# MIT License
#
# Copyright (c) 2021 Soohwan Kim and Sangchun Ha and Soyoung Cho
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
import torch
from torch import Tensor
[docs]def get_attn_pad_mask(inputs, input_lengths, expand_length):
""" mask position is set to 1 """
def get_transformer_non_pad_mask(inputs: Tensor, input_lengths: Tensor) -> Tensor:
""" Padding position is set to 0, either use input_lengths or pad_id """
batch_size = inputs.size(0)
if len(inputs.size()) == 2:
non_pad_mask = inputs.new_ones(inputs.size()) # B x T
elif len(inputs.size()) == 3:
non_pad_mask = inputs.new_ones(inputs.size()[:-1]) # B x T
else:
raise ValueError(f"Unsupported input shape {inputs.size()}")
for i in range(batch_size):
non_pad_mask[i, input_lengths[i]:] = 0
return non_pad_mask
non_pad_mask = get_transformer_non_pad_mask(inputs, input_lengths)
pad_mask = non_pad_mask.lt(1)
attn_pad_mask = pad_mask.unsqueeze(1).expand(-1, expand_length, -1)
return attn_pad_mask
def get_attn_subsequent_mask(seq):
assert seq.dim() == 2
attn_shape = [seq.size(0), seq.size(1), seq.size(1)]
subsequent_mask = torch.triu(torch.ones(attn_shape), diagonal=1)
if seq.is_cuda:
subsequent_mask = subsequent_mask.cuda()
return subsequent_mask