Update modeling_axion.py

modeling_axion.py

@@ -69,6 +69,19 @@ class DeepSeekNanoForCausalLM(PreTrainedModel):
                 labels=None, use_cache=False, **kwargs):
         x = self.embed(input_ids)
         new_caches = [] if use_cache else None
+
+        # Compatibilidade com DynamicCache (Transformers >= 4.36)
+        # Converte para lista simples que o nosso MLA entende
+        if past_key_values is not None and not isinstance(past_key_values, list):
+            try:
+                past_key_values = [
+                    (past_key_values.key_cache[i], past_key_values.value_cache[i])
+                    if i < len(past_key_values.key_cache) else None
+                    for i in range(len(self.blocks))
+                ]
+            except Exception:
+                past_key_values = None
+
         for i, block in enumerate(self.blocks):
             cache = past_key_values[i] if past_key_values else None
             x, nc = block(x, kv_cache=cache, use_cache=use_cache)
@@ -82,6 +95,7 @@ class DeepSeekNanoForCausalLM(PreTrainedModel):
         return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=new_caches)
 
     def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):
-        if past_key_values:
+        # Se tem cache (qualquer tipo), usa só o último token
+        if past_key_values is not None:
             input_ids = input_ids[:, -1:]
         return {"input_ids": input_ids, "past_key_values": past_key_values, "use_cache": True}