Method beam_search_decode

axlearn/audio/decoder_asr.py:1016–1099

Transducer label-synchronous search. Each hypothesis in the beam has the same length of tokens, including both blank and label tokens. Args: input_batch: A dict containing: inputs: A Tensor of shape [batch_size, num_frames, dim] from encoder

(  # pytype: disable=signature-mismatch
        self,
        input_batch: Nested[Tensor],
        num_decodes: int,
        max_decode_len: int,
        prefix_merger: Optional[PrefixMerger] = None,
    )

        return tokens_to_scores

    def beam_search_decode(  # pytype: disable=signature-mismatch
        self,
        input_batch: Nested[Tensor],
        num_decodes: int,
        max_decode_len: int,
        prefix_merger: Optional[PrefixMerger] = None,
    ) -> DecodeOutputs:
        """Transducer label-synchronous search.

        Each hypothesis in the beam has the same length of tokens, including
            both blank and label tokens.

        Args:
            input_batch: A dict containing:
                inputs: A Tensor of shape [batch_size, num_frames, dim] from encoder outputs.
                paddings: A 0/1 Tensor of shape [batch_size, num_frames]. 1's represent paddings.
            num_decodes: Beam size.
            max_decode_len: maximum number of decode steps to run beam search.
                Decoding terminates if an eos token is not emitted after max_decode_steps
                steps. This value can depend on the tokenization.
            prefix_merger: An optional PrefixMerger to apply during decoding.

        Returns:
            DecodeOutputs, containing
                raw_sequences: An int Tensor of shape [batch_size, num_decodes, max_decode_len].
                sequences: An int Tensor of shape [batch_size, num_decodes, max_decode_len].
                paddings: A 0/1 Tensor of shape [batch_size, num_decodes, max_decode_len].
                scores: A Tensor of shape [batch_size, num_decodes].

        Raises:
            ValueError: If max_decode_len <= src_max_len.
        """
        paddings: Tensor = input_batch["paddings"]
        batch_size, src_max_len = paddings.shape
        if max_decode_len <= src_max_len:
            raise ValueError(f"{max_decode_len=} is smaller than {src_max_len=}.")

        cfg = self.config
        blank_id, eos_id, bos_id = cfg.blank_id, cfg.eos_id, cfg.bos_id

        # Starts decoding with [BOS] token.
        inputs = jnp.zeros((batch_size, max_decode_len), dtype=jnp.int32)
        inputs = inputs.at[:, 0].set(bos_id)

        init_states = {
            "am_step": jnp.zeros(batch_size),
            "lm_states": self.prediction_network.init_states(batch_size=batch_size),
            "lm_data": jnp.zeros((batch_size, 1, self.config.joint_dim)),
            "decode_step": jnp.array(0),
        }

        beam_search_outputs = beam_search_decode(
            inputs=inputs,
            time_step=infer_initial_time_step(inputs, pad_id=0),
            cache=init_states,
            tokens_to_scores=self._tokens_to_scores(
                input_batch, num_decodes=num_decodes, max_decode_len=max_decode_len
            ),