Function _map_label_sequences

axlearn/audio/decoder_asr.py:650–698

Removes blanks, paddings, and repeats from the input sequences, as used in CTC or RNN-T. Note that unless pad_id is the same as blank_id, pad_id should not be in the range of the vocab. We use pad_id to infer padding positions in the inputs. Args: inputs: An int Tensor of shape

(
    inputs: Tensor, *, remove_repeats: bool, blank_id: int = 0, pad_id: int = 0
)

def _map_label_sequences(
    inputs: Tensor, *, remove_repeats: bool, blank_id: int = 0, pad_id: int = 0
) -> Nested[Tensor]:
    """Removes blanks, paddings, and repeats from the input sequences, as used in CTC or RNN-T.

    Note that unless pad_id is the same as blank_id, pad_id should not be in the range of the vocab.
    We use pad_id to infer padding positions in the inputs.

    Args:
        inputs: An int Tensor of shape [..., max_decode_len] containing decoded sequences.
        remove_repeats: A boolean indicating whether we remove repeats or not. It is True for CTC,
            False for RNN-T.
        blank_id: Token ID corresponding to blanks.
        pad_id: Token ID corresponding to paddings.

    Returns:
        A dict containing:
            sequences: A Tensor of shape [..., max_decode_len] containing label sequences.
            paddings: A 0/1 Tensor of shape [..., max_decode_len]. 1's represent paddings.
            lengths: A Tensor of shape [..., 1] containing the length of each sequence.
    """
    max_decode_len = inputs.shape[-1]
    # Identify points at which the token is a valid label token to keep.
    # `indicators` has shape [batch_size, num_decodes, max_decode_len], and has a value of 1
    # in positions corresponding to inputs we intend to keep,
    # i.e., the token is not blank or padding.
    indicators = (inputs != blank_id) & (inputs != pad_id)
    if remove_repeats:
        # Identify points at which curr != prev.
        # I.e., the token is not blank or padding, and is different from the previous token.
        y = jnp.concatenate([jnp.full(inputs.shape[:-1] + (1,), pad_id), inputs], axis=-1)
        indicators = (y[..., 1:] != y[..., :-1]) & indicators

    # Compute lengths of final sequences. [..., 1].
    lens = jnp.sum(indicators, axis=-1, keepdims=True, dtype=inputs.dtype)

    # Compute sequences by left-justifying the tokens-to-keep. Under jit, we use a dispatch matrix
    # of shape [batch_size, num_decodes, max_decode_len, max_decode_len].
    # dispatch[..., from, to] == 1 means inputs[:, :, from] is put at sequences[:, :, to].
    # dispatch[..., i, :] == 0 means we drop token i in the inputs.
    # [batch_size, num_decodes, max_decode_len, max_decode_len].
    dispatch = jax.nn.one_hot(
        jnp.cumsum(indicators, axis=-1) * indicators - 1, max_decode_len, dtype=inputs.dtype
    )
    sequences = jnp.einsum("...nm,...n->...m", dispatch, inputs)
    paddings = jnp.arange(max_decode_len) >= lens
    if pad_id != 0:
        sequences = jnp.where(paddings, pad_id, sequences)
    return dict(sequences=sequences, paddings=paddings, lengths=lens)


class RNNPredictionNetwork(BaseLayer):