src/transformers/models/xlnet/modeling_tf_xlnet.py/TFXLNetLMHeadModel/prepare_inputs_for_generation
class TFXLNetLMHeadModel:    def prepare_inputs_for_generation(self, inputs, past, use_mems=None, **kwargs):
        # Add dummy token at the end (no attention on this one)

        # At every pass, the attention values for the new token and the two last generated tokens
        # are computed, the rest is reloaded from the `past` cache. A purely auto-regressive model would have
        # offset = 1; offset = 2 seems to have slightly better computation.
        offset = 2

        effective_batch_size = inputs.shape[0]
        dummy_token = tf.zeros((effective_batch_size, 1), dtype=inputs.dtype)

        if past:
            inputs = tf.concat([inputs[:, -offset:], dummy_token], axis=1)
        else:
            inputs = tf.concat([inputs, dummy_token], axis=1)

        # Build permutation mask so that previous tokens don't see last token
        sequence_length = inputs.shape[1]
        perm_mask = tf.zeros((effective_batch_size, sequence_length, sequence_length - 1))
        perm_mask_seq_end = tf.ones((effective_batch_size, sequence_length, 1))
        perm_mask = tf.concat([perm_mask, perm_mask_seq_end], axis=-1)

        # We'll only predict the last token
        target_mapping = tf.zeros((effective_batch_size, 1, sequence_length - 1))
        target_mapping_seq_end = tf.ones((effective_batch_size, 1, 1))
        target_mapping = tf.concat([target_mapping, target_mapping_seq_end], axis=-1)

        inputs = {
            "input_ids": inputs,
            "perm_mask": perm_mask,
            "target_mapping": target_mapping,
            "use_mems": kwargs.get("use_mems"),
        }

        # if past is defined in model kwargs then use it for faster decoding
        if past:
            inputs["mems"] = tuple(layer_past[:-offset, :, :] for layer_past in past)

        return inputs