Spaces:
Runtime error
Runtime error
| import torch | |
| def get_focus_rate(attn, src_padding_mask=None, tgt_padding_mask=None): | |
| ''' | |
| attn: bs x L_t x L_s | |
| ''' | |
| if src_padding_mask is not None: | |
| attn = attn * (1 - src_padding_mask.float())[:, None, :] | |
| if tgt_padding_mask is not None: | |
| attn = attn * (1 - tgt_padding_mask.float())[:, :, None] | |
| focus_rate = attn.max(-1).values.sum(-1) | |
| focus_rate = focus_rate / attn.sum(-1).sum(-1) | |
| return focus_rate | |
| def get_phone_coverage_rate(attn, src_padding_mask=None, src_seg_mask=None, tgt_padding_mask=None): | |
| ''' | |
| attn: bs x L_t x L_s | |
| ''' | |
| src_mask = attn.new(attn.size(0), attn.size(-1)).bool().fill_(False) | |
| if src_padding_mask is not None: | |
| src_mask |= src_padding_mask | |
| if src_seg_mask is not None: | |
| src_mask |= src_seg_mask | |
| attn = attn * (1 - src_mask.float())[:, None, :] | |
| if tgt_padding_mask is not None: | |
| attn = attn * (1 - tgt_padding_mask.float())[:, :, None] | |
| phone_coverage_rate = attn.max(1).values.sum(-1) | |
| # phone_coverage_rate = phone_coverage_rate / attn.sum(-1).sum(-1) | |
| phone_coverage_rate = phone_coverage_rate / (1 - src_mask.float()).sum(-1) | |
| return phone_coverage_rate | |
| def get_diagonal_focus_rate(attn, attn_ks, target_len, src_padding_mask=None, tgt_padding_mask=None, | |
| band_mask_factor=5, band_width=50): | |
| ''' | |
| attn: bx x L_t x L_s | |
| attn_ks: shape: tensor with shape [batch_size], input_lens/output_lens | |
| diagonal: y=k*x (k=attn_ks, x:output, y:input) | |
| 1 0 0 | |
| 0 1 0 | |
| 0 0 1 | |
| y>=k*(x-width) and y<=k*(x+width):1 | |
| else:0 | |
| ''' | |
| # width = min(target_len/band_mask_factor, 50) | |
| width1 = target_len / band_mask_factor | |
| width2 = target_len.new(target_len.size()).fill_(band_width) | |
| width = torch.where(width1 < width2, width1, width2).float() | |
| base = torch.ones(attn.size()).to(attn.device) | |
| zero = torch.zeros(attn.size()).to(attn.device) | |
| x = torch.arange(0, attn.size(1)).to(attn.device)[None, :, None].float() * base | |
| y = torch.arange(0, attn.size(2)).to(attn.device)[None, None, :].float() * base | |
| cond = (y - attn_ks[:, None, None] * x) | |
| cond1 = cond + attn_ks[:, None, None] * width[:, None, None] | |
| cond2 = cond - attn_ks[:, None, None] * width[:, None, None] | |
| mask1 = torch.where(cond1 < 0, zero, base) | |
| mask2 = torch.where(cond2 > 0, zero, base) | |
| mask = mask1 * mask2 | |
| if src_padding_mask is not None: | |
| attn = attn * (1 - src_padding_mask.float())[:, None, :] | |
| if tgt_padding_mask is not None: | |
| attn = attn * (1 - tgt_padding_mask.float())[:, :, None] | |
| diagonal_attn = attn * mask | |
| diagonal_focus_rate = diagonal_attn.sum(-1).sum(-1) / attn.sum(-1).sum(-1) | |
| return diagonal_focus_rate, mask | |