Method sample_decode

axlearn/audio/decoder_asr.py:517–558

CTC sample decoding. The output hypotheses will have blanks and repeats removed (via `_map_label_sequences`). To perform greedy decoding, provide `top_k_logits(1)` as the logits modifier. Args: input_batch: See `beam_search_decode`. num_decodes: See

(
        self,
        input_batch: Nested[Tensor],
        *,
        num_decodes: int = 1,
        logits_modifier: Optional[ConfigOr[LogitsToLogitsFn]] = None,
    )

        )

    def sample_decode(
        self,
        input_batch: Nested[Tensor],
        *,
        num_decodes: int = 1,
        logits_modifier: Optional[ConfigOr[LogitsToLogitsFn]] = None,
    ) -> DecodeOutputs:
        """CTC sample decoding.

        The output hypotheses will have blanks and repeats removed (via `_map_label_sequences`).
        To perform greedy decoding, provide `top_k_logits(1)` as the logits modifier.

        Args:
            input_batch: See `beam_search_decode`.
            num_decodes: See `beam_search_decode`.
            logits_modifier: An optional logits modifier to apply prior to softmax.
                If None, do not modify the logits.

        Returns:
            See `beam_search_decode`.
        """
        cfg: CTCDecoderModel.Config = self.config
        paddings: Tensor = input_batch["paddings"]
        # Add 1 so we can drop EOS while ensuring decodes can be up to `num_frames`.
        max_decode_len = paddings.shape[-1] + 1
        sample_decode_outputs = sample_decode(
            inputs=jnp.zeros_like(paddings),
            time_step=jnp.zeros(paddings.shape[0], dtype=jnp.int32),
            cache={"time_step": jnp.array(0)},
            tokens_to_scores=self._tokens_to_scores(
                input_batch, num_decodes=num_decodes, logits_modifier=logits_modifier
            ),
            num_decodes=num_decodes,
            prng_key=self.prng_key,
            max_decode_len=max_decode_len,
            stop_decoding_condition=StopOnSubsequence([[cfg.vocab_size]]),  # Dummy EOS token.
        )
        return self._postprocess_outputs(
            sequences=sample_decode_outputs.sequences,
            paddings=paddings,
            scores=sample_decode_outputs.token_scores,
        )

    def greedy_decode(self, input_batch: Nested[Tensor]) -> DecodeOutputs:
        """CTC greedy decoding.