Update recastmlp_llama/modeling_recastmlp_llama.py

recastmlp_llama/modeling_recastmlp_llama.py

@@ -281,21 +281,22 @@ class RECASTMLP_llamaModel(PreTrainedModel):
 
         if inputs_embeds is None:
             inputs_embeds = self.embed_tokens(input_ids)
-
+        # Set up cache position if not provided
+        if cache_position is None:
+            past_seen_tokens = 0 if past_key_values is None else (
+                past_key_values.get_seq_length() if isinstance(past_key_values, Cache) 
+                else past_key_values[0][0].size(-2) if past_key_values 
+                else 0
+            )
+            cache_position = torch.arange(
+                past_seen_tokens, 
+                past_seen_tokens + inputs_embeds.shape[1], 
+                device=inputs_embeds.device
+            )
         # Create position embeddings to be shared across the decoder layers
+        # Set up position IDs if not provided
         if position_ids is None:
-            past_seen_tokens = (
-                past_key_values.get_seq_length() if past_key_values is not None else 0
-            )
-            position_ids = torch.arange(
-                past_seen_tokens,
-                past_seen_tokens + inputs_embeds.shape[1],
-                device=inputs_embeds.device,
-            ).unsqueeze(0)
-
-        position_embeddings = self.rotary_emb(inputs_embeds, position_ids)
-        hidden_states = inputs_embeds
-
+            position_ids = cache_position.unsqueeze(0)
         # Get updated causal mask
         causal_mask = self._update_causal_mask(
             attention_mask,
@@ -304,6 +305,9 @@ class RECASTMLP_llamaModel(PreTrainedModel):
             past_key_values,
             output_attentions,
         )
+        hidden_states = inputs_embeds
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+        
 
         # Initialize outputs
         all_hidden_states = () if output_hidden_states else None