src/transformers/models/xglm/modeling_tf_xglm.py/TFXGLMForCausalLM/prepare_inputs_for_generation
class TFXGLMForCausalLM:    def prepare_inputs_for_generation(self, inputs, past_key_values=None, use_cache=None, **kwargs):
        # only last token for inputs_ids if past is defined in kwargs
        if past_key_values:
            inputs = tf.expand_dims(inputs[:, -1], -1)

        attention_mask = kwargs.get("attention_mask", None)

        return {
            "input_ids": inputs,
            "attention_mask": attention_mask,
            "past_key_values": past_key_values,
            "use_cache": use_cache,
        }