Method beam_search_decode

axlearn/common/decoder.py:218–285

Perform beam search decoding. Args: input_batch: A dict containing: prefix: The prefix to use for prompting of shape [batch, max_prefix_length]. The prefix for each example in the batch should begin with a prompt token (e.g.

(
        self,
        *,
        input_batch: Nested[Tensor],
        max_sequence_length: int,
        num_decodes: int,
        cross_attention_data: Optional[Tensor] = None,
        cross_attention_logit_biases: Optional[Tensor] = None,
        brevity_penalty: Optional[BrevityPenaltyFn] = None,
    )

    def beam_search_decode(
        self,
        *,
        input_batch: Nested[Tensor],
        max_sequence_length: int,
        num_decodes: int,
        cross_attention_data: Optional[Tensor] = None,
        cross_attention_logit_biases: Optional[Tensor] = None,
        brevity_penalty: Optional[BrevityPenaltyFn] = None,
    ) -> BeamSearchOutputs:
        """Perform beam search decoding.

        Args:
            input_batch: A dict containing:
                prefix: The prefix to use for prompting of shape [batch, max_prefix_length].
                    The prefix for each example in the batch should begin with a prompt token (e.g.
                    BOS).
                    The prefix will be padded with `cfg.pad_token_id` to `max_sequence_length`, thus
                    it is expected that `max_prefix_length <= max_sequence_length`.
            max_sequence_length: The maximum sequence length of tokens to generate.
            num_decodes: The number of decoded sequences to return. These are the number of
                hypotheses per batch example.
            cross_attention_data: A float Tensor of shape [batch_size, source_len, hidden_dim].
            cross_attention_logit_biases: A Tensor of shape [batch_size, target_len, source_len].
                A -inf represents a disconnected position pair.
                `target_len` should be broadcastable to `max_sequence_length`.
            brevity_penalty: Brevity penalty function for length normalization during beam search.

        Returns:
            The beam search outputs.

        Raises:
            ValueError: If pad_token_id is non-zero.
        """
        validate_contains_paths(input_batch, paths=["prefix"])
        prefix = input_batch["prefix"]

        cfg = self.config
        tokens_to_scores_fn = self._tokens_to_scores(
            num_decodes=num_decodes,
            cross_attention_data=cross_attention_data,
            cross_attention_logit_biases=cross_attention_logit_biases,
            logits_modifier=None,
        )
        input_ids = self._pad(
            prefix, max_sequence_length=max_sequence_length, pad_id=cfg.pad_token_id
        )
        time_step = infer_initial_time_step(prefix, pad_id=cfg.pad_token_id)
        prefill_batch = {**input_batch}
        prefill_batch["input_ids"] = input_ids
        # Note: it prefills `k-1` tokens and used the last prefix token in first generation loop.
        # TODO(axlearn-dev): prefill all prefix tokens in one shot.
        init_states, _ = self._decoder.prefill_states(
            time_step=time_step,
            input_batch=prefill_batch,
            cross_attention_data=cross_attention_data,
            cross_attention_logit_biases=cross_attention_logit_biases,
        )
        return beam_search_decode(