Method tokens_to_scores

axlearn/common/decoder.py:417–456

Maps current token IDs and model state to next logits and updated state. Args: token_ids: An int Tensor of shape [batch*num_decodes, 1]. cache: A NestedTensor of cached states. Returns: (log_probs, updated_cache), where log_pr

(token_ids: Tensor, cache: NestedTensor)

                (log_probs, updated_cache), where log_probs has shape
                [token_ids.shape[0], vocab_size] and represents log probabilities of the next
                tokens; and updated_cache is the updated cache.
            """
            time_step = cache["time_step"]
            assert time_step.ndim == 1

            # Select attention biases corresponding to the current time steps.
            # We alias the nonlocal variable.
            cross_attention_biases = cross_attention_logit_biases
            if cross_attention_biases is not None:
                # Note: the target_len dimension can be 1 during decoding.
                # When indexing, we clip the indices instead of producing NaNs.
                # TODO(markblee): Consider removing `take_along_axis` entirely if we restrict
                # target_len to always be 1 during decoding.
                # [batch*num_decodes, num_heads, 1, source_len].
                cross_attention_biases = jnp.take_along_axis(
                    cross_attention_biases, time_step[:, None, None, None], mode="clip", axis=2
                )

            # Use a temporary output collection to avoid tracer leaks during extend_step.
            with _temporary_output_collection():
                updated_state, outputs = self._decoder.extend_step(
                    cached_states=cache,
                    input_batch={"input_ids": token_ids},
                    cross_attention_data=cross_attention_data,
                    cross_attention_logit_biases=cross_attention_biases,
                )

            logits = outputs["logits"]
            log_probs = log_probs_from_logits(logits[:, -1, :], logits_modifier=logits_modifier)
            return log_probs, updated_state

        return tokens_to_scores

    @staticmethod
    def _pad(prefix: Tensor, *, max_sequence_length: int, pad_id: int) -> Tensor:
        """Accept token IDs input tensor and pad if necessary to max_sequence_length."""
        return jnp.concatenate(
            [
                prefix,
                jnp.full(
                    (prefix.shape[0], max_sequence_length - prefix.shape[1]),
                    pad_id,
                    dtype=prefix.dtype,