Update modeling_c_cubed.py

modeling_c_cubed.py

@@ -707,4 +707,30 @@ class CcubedForConditionalGeneration(CcubedPreTrainedModel):
             hidden_states=outputs.hidden_states,
             attentions=outputs.attentions,
             context_hidden_states=context_features,
-        )
+        )
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        inputs_embeds=None,
+        past_key_values=None,
+        attention_mask=None,
+        context_attention_mask=None,
+        **kwargs
+    ):
+        if past_key_values:
+            input_ids = input_ids[:, -1:]
+
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids}
+
+        model_inputs.update({
+            "past_key_values": past_key_values,
+            "use_cache": kwargs.get("use_cache"),
+            "attention_mask": attention_mask,
+            "context_attention_mask": context_attention_mask
+        })
+        return model_inputs