Function sample_decode

axlearn/common/decoding.py:978–1151

Performs sampling decoding. Args: inputs: An int32 Tensor of shape [batch_size, length] containing a sequence of tokens. Please refer to `beam_search_decode` for more information on `inputs`. time_step: Initial time steps for decoding of shape [batch_size].

(
    *,
    inputs: Tensor,
    time_step: Tensor,
    cache: NestedTensor,
    tokens_to_scores: Callable[[Tensor, NestedTensor], tuple[Tensor, NestedTensor]],
    stop_decoding_condition: StopDecodingCondition,
    num_decodes: int,
    prng_key: Tensor,
    max_decode_len: Optional[int] = None,
    loop: Literal["lax", "python"] = "lax",
    pad_id: int = 0,
    input_token_scores: Optional[Tensor] = None,
)

def sample_decode(
    *,
    inputs: Tensor,
    time_step: Tensor,
    cache: NestedTensor,
    tokens_to_scores: Callable[[Tensor, NestedTensor], tuple[Tensor, NestedTensor]],
    stop_decoding_condition: StopDecodingCondition,
    num_decodes: int,
    prng_key: Tensor,
    max_decode_len: Optional[int] = None,
    loop: Literal["lax", "python"] = "lax",
    pad_id: int = 0,
    input_token_scores: Optional[Tensor] = None,
) -> SampleOutputs:
    """Performs sampling decoding.

    Args:
        inputs: An int32 Tensor of shape [batch_size, length] containing a sequence of tokens.
            Please refer to `beam_search_decode` for more information on `inputs`.
        time_step: Initial time steps for decoding of shape [batch_size].
            Please refer to `beam_search_decode` for more information on `time_step`.
        cache: State of the decoder model.
        tokens_to_scores: Fast autoregressive decoder function taking single token
            slices and cache and returning next-token scores and updated cache.
            [batch*num_decodes, vocab], {} = tokens_to_scores([batch*num_decodes, 1], {}).
            NestedTensor usually has batch*num_decodes as the leading dim.
            The scores represents logits for sampling the next token, i.e.,
            the sampling probabilities will be softmax(scores, axis=-1).
            The caller can implement temperature-based, top-k, or top-p sampling as part
            of `tokens_to_scores`.
        stop_decoding_condition: StopDecodingCondition instance which given index, current sequences
            and prompt mask returns a bool tensor indicating if a sequence is complete.
        num_decodes: Number of decoded sequences to be returned for each input sequence.
        prng_key: The random key.
        max_decode_len: An optional maximum length of decoded sequence. If
            None, it uses `inputs.shape[1]` as `max_decode_len`.
        loop: "lax" or "python". The latter should only be used for debugging.
        pad_id: Token ID associated with padded input.
        input_token_scores: Optional initial scores of shape [batch_size, length] where
            length < max_decode_len, e.g. as produced by prefilling. Note that length should be
            strictly less than max_decode_len, as we exclude the scores for the dummy prompt token.
            In other words, input_token_scores[i, j - 1] represents the score for inputs[i, j],
            since the token at inputs[i, 0] does not have a score. Defaults to all zeros.

    Returns:
        SampleOutputs, containing the finished or live sequences and their corresponding scores.

    Raises:
        NotImplementedError: If an unsupported loop is provided.
    """
    batch_size = inputs.shape[0]
    if max_decode_len is None:
        max_decode_len = inputs.shape[1]
    # Account for conditioning input token.
    max_decode_len += 1

    # Initialize state.
    sample_decode_init_state = _decode_init(