Update modeling.py

modeling.py

@@ -150,20 +150,24 @@ class RecombinationTransformerForCausalLM(PreTrainedModel):
         self.final_rms_norm = RMSNorm(config.embed_dim)
         self.lm_head = nn.Linear(config.embed_dim, config.vocab_size, bias=False)
     
-    def forward(self, input_ids, attention_mask=None, past_key_values=None, return_dict=None):
+    def forward(self, input_ids, attention_mask=None, past_key_values=None, return_dict=None, **kwargs):
         if attention_mask is None:
             attention_mask = torch.ones(input_ids.shape, device=input_ids.device)
         
-        # Create causal mask
         batch_size, seq_length = input_ids.size()
         causal_mask = torch.tril(torch.ones((seq_length, seq_length), device=input_ids.device)).unsqueeze(0).expand(batch_size, -1, -1)
         
+        if past_key_values is None:
+            past_key_values = [None] * len(self.layers)
+        
         # Embedding
         x = self.embed_tokens(input_ids)
         
-        # Apply layers
-        for layer in self.layers:
+        new_past_key_values = []
+        for i, layer in enumerate(self.layers):
+            past_key_value = past_key_values[i]
             x = layer(x, attn_mask=causal_mask)
+            new_past_key_values.append(x)
         
         # Final normalization
         x = self.final_rms_norm(x)
@@ -172,9 +176,10 @@ class RecombinationTransformerForCausalLM(PreTrainedModel):
         logits = self.lm_head(x)
         
         if not return_dict:
-            return (logits,)
+            return (logits, new_past_key_values)
+
+        return CausalLMOutputWithPast(logits=logits, past_key_values=new_past_key_values)
 
-        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: