compatible with transformers>=4.42.0

modeling_kangaroo.py

@@ -1346,7 +1346,15 @@ class KangarooForCausalLM(LlamaPreTrainedModel):
                 position_ids = position_ids[:, -input_ids.shape[1] :]
 
         # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
-        if inputs_embeds is not None and past_key_values is None:
+        set_inputs_embeds = False
+        if inputs_embeds is not None:
+            if isinstance(past_key_values, Cache):
+                if past_key_values.get_seq_length() == 0:
+                    set_inputs_embeds = True
+            else:
+                if past_key_values is None:
+                    set_inputs_embeds = True
+        if set_inputs_embeds:
             model_inputs = {"inputs_embeds": inputs_embeds}
         else:
             # The `contiguous()` here is necessary to have a static stride during decoding. torchdynamo otherwise