# Copyright (c) 2021 Graphcore Ltd. All rights reserved. # Copyright 2019 Shigeki Karita # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # Copyright 2020 Johns Hopkins University (Shinji Watanabe) # Northwestern Polytechnical University (Pengcheng Guo) # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # This file has been modified by Graphcore Ltd. # Main changes are removing to_device function and so on. # Copyright 2021 Mobvoi Inc. All Rights Reserved. # -*- coding: utf-8 -*- # Copyright 2019 Shigeki Karita # Apache 2.0 (http://www.apache.org/licenses/LICENSE-2.0) """ This script has been adapted from some of the original EspNet and WeNet repo found here: [ https://github.com/espnet/espnet/blob/master/espnet/nets/pytorch_backend/nets_utils.py https://github.com/wenet-e2e/wenet/blob/main/wenet/utils/mask.py ] """ import torch def make_pad_mask(lengths, maxlen): """Make mask tensor containing indices of padded part. Args: lengths (LongTensor): Batch of lengths (B,). maxlen (int): length of batch Returns: Tensor: Mask tensor containing indices of padded part. dtype=torch.uint8 in PyTorch 1.2- dtype=torch.bool in PyTorch 1.2+ (including 1.2) Examples: With only lengths. >>> lengths = [5, 3, 2] >>> make_non_pad_mask(lengths) masks = [[0, 0, 0, 0 ,0], [0, 0, 0, 1, 1], [0, 0, 1, 1, 1]] With the reference tensor. >>> xs = torch.zeros((3, 2, 4)) >>> make_pad_mask(lengths, xs) tensor([[[0, 0, 0, 0], [0, 0, 0, 0]], [[0, 0, 0, 1], [0, 0, 0, 1]], [[0, 0, 1, 1], [0, 0, 1, 1]]], dtype=torch.uint8) >>> xs = torch.zeros((3, 2, 6)) >>> make_pad_mask(lengths, xs) tensor([[[0, 0, 0, 0, 0, 1], [0, 0, 0, 0, 0, 1]], [[0, 0, 0, 1, 1, 1], [0, 0, 0, 1, 1, 1]], [[0, 0, 1, 1, 1, 1], [0, 0, 1, 1, 1, 1]]], dtype=torch.uint8) With the reference tensor and dimension indicator. >>> xs = torch.zeros((3, 6, 6)) >>> make_pad_mask(lengths, xs, 1) tensor([[[0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1]], [[0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1]], [[0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1]]], dtype=torch.uint8) >>> make_pad_mask(lengths, xs, 2) tensor([[[0, 0, 0, 0, 0, 1], [0, 0, 0, 0, 0, 1], [0, 0, 0, 0, 0, 1], [0, 0, 0, 0, 0, 1], [0, 0, 0, 0, 0, 1], [0, 0, 0, 0, 0, 1]], [[0, 0, 0, 1, 1, 1], [0, 0, 0, 1, 1, 1], [0, 0, 0, 1, 1, 1], [0, 0, 0, 1, 1, 1], [0, 0, 0, 1, 1, 1], [0, 0, 0, 1, 1, 1]], [[0, 0, 1, 1, 1, 1], [0, 0, 1, 1, 1, 1], [0, 0, 1, 1, 1, 1], [0, 0, 1, 1, 1, 1], [0, 0, 1, 1, 1, 1], [0, 0, 1, 1, 1, 1]]], dtype=torch.uint8) """ bs = lengths.size(0) seq_range = torch.arange(0, maxlen, dtype=lengths.dtype) seq_range_expand = seq_range.unsqueeze(0).expand(bs, maxlen).contiguous() mask = seq_range_expand >= lengths.unsqueeze(-1) return mask def subsequent_mask(size, dtype=torch.bool): """Create mask for subsequent steps (size, size). :param int size: size of mask :param torch.dtype dtype: result dtype :rtype: torch.Tensor >>> subsequent_mask(3) [[1, 0, 0], [1, 1, 0], [1, 1, 1]] """ ret = torch.ones(size, size, dtype=dtype) return torch.tril(ret, out=ret) def target_mask(ys_in_pad, ignore_id): """Create mask for decoder self-attention. :param torch.Tensor ys_pad: batch of padded target sequences (B, Lmax) :param int ignore_id: index of padding :param torch.dtype dtype: result dtype :rtype: torch.Tensor (B, Lmax, Lmax) """ ys_mask = ys_in_pad != ignore_id m = subsequent_mask(ys_mask.size(-1)).unsqueeze(0) return ys_mask.unsqueeze(-2) & m def subsequent_mask( size: int, device: torch.device = None, ) -> torch.Tensor: """Create mask for subsequent steps (size, size). This mask is used only in decoder which works in an auto-regressive mode. This means the current step could only do attention with its left steps. In encoder, fully attention is used when streaming is not necessary and the sequence is not long. In this case, no attention mask is needed. When streaming is need, chunk-based attention is used in encoder. See subsequent_chunk_mask for the chunk-based attention mask. Args: size (int): size of mask str device (str): "cpu" or "cuda" or torch.Tensor.device dtype (torch.device): result dtype Returns: torch.Tensor: mask Examples: >>> subsequent_mask(3) [[1, 0, 0], [1, 1, 0], [1, 1, 1]] """ arange = torch.arange(size, device=device) mask = arange.expand(size, size) arange = arange.unsqueeze(-1) mask = mask <= arange return mask def make_pad_mask_decoder(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor: """Make mask tensor containing indices of padded part. See description of make_non_pad_mask. Args: lengths (torch.Tensor): Batch of lengths (B,). Returns: torch.Tensor: Mask tensor containing indices of padded part. Examples: >>> lengths = [5, 3, 2] >>> make_pad_mask_decoder(lengths) masks = [[0, 0, 0, 0 ,0], [0, 0, 0, 1, 1], [0, 0, 1, 1, 1]] """ batch_size = lengths.size(0) max_len = max_len if max_len > 0 else lengths.max().item() seq_range = torch.arange(0, max_len, dtype=torch.int64, device=lengths.device) seq_range_expand = seq_range.unsqueeze(0).expand(batch_size, max_len) seq_length_expand = lengths.unsqueeze(-1) mask = seq_range_expand >= seq_length_expand return mask def mask_finished_scores(score: torch.Tensor, flag: torch.Tensor) -> torch.Tensor: """ If a sequence is finished, we only allow one alive branch. This function aims to give one branch a zero score and the rest -inf score. Args: score (torch.Tensor): A real value array with shape (batch_size * beam_size, beam_size). flag (torch.Tensor): A bool array with shape (batch_size * beam_size, 1). Returns: torch.Tensor: (batch_size * beam_size, beam_size). """ beam_size = score.size(-1) zero_mask = torch.zeros_like(flag, dtype=torch.bool) if beam_size > 1: unfinished = torch.cat((zero_mask, flag.repeat([1, beam_size - 1])), dim=1) finished = torch.cat((flag, zero_mask.repeat([1, beam_size - 1])), dim=1) else: unfinished = zero_mask finished = flag score.masked_fill_(unfinished, -float("inf")) score.masked_fill_(finished, 0) return score def mask_finished_preds(pred: torch.Tensor, flag: torch.Tensor, eos: int) -> torch.Tensor: """ If a sequence is finished, all of its branch should be Args: pred (torch.Tensor): A int array with shape (batch_size * beam_size, beam_size). flag (torch.Tensor): A bool array with shape (batch_size * beam_size, 1). Returns: torch.Tensor: (batch_size * beam_size). """ beam_size = pred.size(-1) finished = flag.repeat([1, beam_size]) return pred.masked_fill_(finished, eos)