    def take_search_step(
        self, last_predictions: torch.Tensor, state: Dict[str, torch.Tensor], step: int
    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
        """
        Take step during beam search.

        This function is what gets passed to the `BeamSearch.search` method. It takes
        predictions from the last timestep and the current state and outputs
        the log probabilities assigned to tokens for the next timestep, as well as the updated
        state.
        """
        decoder_cache: Optional[List[KeyValueStates]] = None
        decoder_cache_dict = {
            k: state[k].contiguous() for k in state if k.startswith("decoder_cache_")
        }
        if decoder_cache_dict:
            decoder_cache = self._dict_to_decoder_cache(decoder_cache_dict)

        if len(last_predictions.shape) == 1:
            last_predictions = last_predictions.unsqueeze(-1)

        decoder_outputs: T5StackOutput = self.decoder(
            input_ids=last_predictions,
            past_key_values=decoder_cache,
            encoder_hidden_states=state["encoder_hidden_states"],
            encoder_attention_mask=state["encoder_attention_mask"],
            use_cache=True,
        )

        # Shape: (group_size, 2, vocab_size)
        lm_logits = self._get_lm_logits(decoder_outputs.last_hidden_state)

        # Shape: (group_size, vocab_size)
        logits = lm_logits[:, -1, :]

        # Shape: (group_size, vocab_size)
        log_probabilities = F.log_softmax(logits, dim=-1)

        # Update state with decoder cache.
        decoder_cache = decoder_outputs.past_key_values
        assert decoder_cache is not None
        decoder_cache_dict = self._decoder_cache_to_dict(decoder_cache)
        state.update(decoder_cache_dict)

        return log_probabilities, state
Take step during beam search.

This function is what gets passed to the `BeamSearch.search` method. It takes
predictions from the last timestep and the current state and outputs
the log probabilities assigned to tokens for the next timestep, as well as the updated
state.
allennlp/modules/transformer/t5.py