Function make_source_mask

funasr/models/transformer/utils/nets_utils.py:648–665

Create source mask for given lengths. Reference: https://github.com/k2-fsa/icefall/blob/master/icefall/utils.py Args: lengths: Sequence lengths. (B,) Returns: : Mask for the sequence lengths. (B, max_len)

(lengths: torch.Tensor)

def make_source_mask(lengths: torch.Tensor) -> torch.Tensor:
    """Create source mask for given lengths.

    Reference: https://github.com/k2-fsa/icefall/blob/master/icefall/utils.py

    Args:
        lengths: Sequence lengths. (B,)

    Returns:
        : Mask for the sequence lengths. (B, max_len)

    """
    max_len = lengths.max()
    batch_size = lengths.size(0)

    expanded_lengths = torch.arange(max_len).expand(batch_size, max_len).to(lengths)

    return expanded_lengths >= lengths.unsqueeze(1)


def get_transducer_task_io(