src/transformers/models/roc_bert/modeling_roc_bert.py/RoCBertForCausalLM/prepare_inputs_for_generation
class RoCBertForCausalLM:    def prepare_inputs_for_generation(
        self,
        input_ids,
        input_shape_ids=None,
        input_pronunciation_ids=None,
        past_key_values=None,
        attention_mask=None,
        **model_kwargs,
    ):
        input_shape = input_ids.shape

        # if model is used as a decoder in encoder-decoder model, the decoder attention mask is created on the fly
        if attention_mask is None:
            attention_mask = input_ids.new_ones(input_shape)

        # cut decoder_input_ids if past is used
        if past_key_values is not None:
            input_ids = input_ids[:, -1:]
            if input_shape_ids is not None:
                input_shape_ids = input_shape_ids[:, -1:]
            if input_pronunciation_ids is not None:
                input_pronunciation_ids = input_pronunciation_ids[:, -1:]

        return {
            "input_ids": input_ids,
            "input_shape_ids": input_shape_ids,
            "input_pronunciation_ids": input_pronunciation_ids,
            "attention_mask": attention_mask,
            "past_key_values": past_key_values,
        }