Spaces:
Running
on
Zero
Running
on
Zero
# Copyright (c) 2019 Shigeki Karita | |
# 2020 Mobvoi Inc (Binbin Zhang) | |
# 2024 Alibaba Inc (authors: Xiang Lyu) | |
# | |
# Licensed under the Apache License, Version 2.0 (the "License"); | |
# 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. | |
import math | |
import torch | |
from typing import List | |
def make_pad_mask(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(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 | |