Update raven_modeling_minimal.py

raven_modeling_minimal.py

@@ -302,12 +302,13 @@ class RavenForCausalLM(RavenPreTrainedModel):
         cache_position: Optional[torch.Tensor] = None,
         **kwargs,
     ) -> CausalLMOutputRecurrentLatents:
+        
         if position_ids is None and cache_position is None:
             freqs_cis = self.freqs_cis[:, : input_ids.shape[1]]
         elif position_ids is not None:
-            freqs_cis = self.freqs_cis.index_select(1, position_ids)
+            freqs_cis = self.freqs_cis.index_select(1, position_ids.squeeze())
         elif cache_position is not None:  # support HF format
-            freqs_cis = self.freqs_cis[:, cache_position : cache_position + 1]
+            freqs_cis = self.freqs_cis[:, cache_position]
 
         if input_embeds is None:
             input_embeds = self.transformer.wte(input_ids)
@@ -445,10 +446,11 @@ class RavenForCausalLM(RavenPreTrainedModel):
             input_ids = input_ids[:, cache_position]  # type: ignore
         model_inputs["input_ids"] = input_ids.clone(memory_format=torch.contiguous_format)
 
-        position_ids = torch.arange(current_input_length)[None, :].to(input_ids.device)
-        model_inputs["position_ids"] = position_ids[:, -current_input_length:].clone(
-            memory_format=torch.contiguous_format
-        )  # positions_ids is a critical argument for the model to correctly apply rope!
+        if cache_position is None:
+            position_ids = torch.arange(current_input_length)[None, :].to(input_ids.device)
+            model_inputs["position_ids"] = position_ids[:, -current_input_length:].clone(
+                memory_format=torch.contiguous_format
+            )  # some form of position_ids is a critical argument for the model to correctly apply rope!
 
         # forward all other entries
         for key, value in kwargs.items():