src/transformers/generation/tf_utils.py/TFGenerationMixin/adjust_logits_during_generation
class TFGenerationMixin:    def adjust_logits_during_generation(
        self, logits, cur_len, max_length, forced_bos_token_id, forced_eos_token_id, **kwargs
    ):
        """
        Implement in subclasses of [`PreTrainedModel`] for custom behavior to adjust the logits in the generate method.
        """
        vocab_size = getattr(self.config, "vocab_size", None)
        if vocab_size is None and self.config.is_encoder_decoder:
            decoder_config = getattr(self.config, "decoder", None)
            if decoder_config is not None:
                vocab_size = getattr(self.config.decoder, "vocab_size", None)

        if cur_len == 1 and forced_bos_token_id is not None:
            vocab_range = tf.constant(range(vocab_size))
            return tf.where(vocab_range != forced_bos_token_id, -1e8, logits)
        elif cur_len == max_length - 1 and forced_eos_token_id is not None:
            vocab_range = tf.constant(range(vocab_size))
            return tf.where(vocab_range != forced_eos_token_id, -1e8, logits)
        else:
            return logits