fix

Author: alien-0119

mindone/transformers/generation/utils.py

@@ -1942,7 +1942,7 @@ def _prepare_generation_config(
 
         return generation_config, model_kwargs
 
-    def _get_initial_cache_position(self, input_ids, model_kwargs):
+    def _get_initial_cache_position(self, seq_length, model_kwargs):
         """Calculates `cache_position` for the pre-fill stage based on `input_ids` and optionally past length"""
         if "cache_position" in model_kwargs and model_kwargs["cache_position"]:
             return model_kwargs
@@ -1951,10 +1951,10 @@ def _get_initial_cache_position(self, input_ids, model_kwargs):
             cache_position = mint.ones_like(model_kwargs["inputs_embeds"][0, :, 0], dtype=ms.int32).cumsum(0) - 1
         elif "decoder_inputs_embeds" in model_kwargs and self.config.is_encoder_decoder:
             cache_position = (
-                mint.ones_like(model_kwargs["decoder_inputs_embeds"][0, :, 0], dtype=ms.int32).cumsum(0) - 1
+                mint.ones(model_kwargs["decoder_inputs_embeds"][0, :, 0], dtype=ms.int32).cumsum(0) - 1
             )
         else:
-            cache_position = mint.ones_like(input_ids[0, :], dtype=ms.int32).cumsum(0) - 1
+            cache_position = mint.ones(seq_length, dtype=ms.int32).cumsum(0) - 1
 
         if model_kwargs.get("past_key_values") is not None:
             cache = model_kwargs["past_key_values"]
@@ -3721,7 +3721,7 @@ def _sample(
         batch_size, cur_len = input_ids.shape
         this_peer_finished = False
         unfinished_sequences = mint.ones(batch_size, dtype=ms.int32)
-        model_kwargs = self._get_initial_cache_position(input_ids, model_kwargs)
+        model_kwargs = self._get_initial_cache_position(cur_len, model_kwargs)
 
         model_forward = self.__call__
         compile_forward = self._valid_auto_compile_criteria(model_kwargs, generation_config)