Update modeling.py

modeling.py

@@ -172,7 +172,13 @@ class RecombinationTransformerForCausalLM(PreTrainedModel):
         logits = self.lm_head(x)
         
         return CausalLMOutputWithPast(logits=logits, past_key_values=past_key_values)
-
+        
+    def prepare_inputs_for_generation(self, input_ids, past=None, attention_mask=None, **kwargs):
+        if past:
+            input_ids = input_ids[:, -1].unsqueeze(-1)
+        
+        return {"input_ids": input_ids, "attention_mask": attention_mask, "past_key_values": past}
+        
     def generate(self, input_ids, attention_mask=None, max_length=512, min_length=None, num_return_sequences=1):
         logits_processor = LogitsProcessorList()