Fix deepseek indexer rope argument (#1431)

Author: pcuenca

mlx_lm/chat_templates/deepseek_v32.py

@@ -333,6 +333,10 @@ def encode_messages(
 def apply_chat_template(
     messages, continue_final_message=False, add_generation_prompt=False, **kwargs
 ):
+    if "enable_thinking" in kwargs:
+        kwargs["thinking_mode"] = (
+            "thinking" if kwargs.pop("enable_thinking") else "chat"
+        )
     out = encode_messages(messages, **kwargs)
     if continue_final_message and add_generation_prompt:
         raise ValueError(

mlx_lm/models/deepseek_v32.py

@@ -50,6 +50,7 @@ class ModelArgs(BaseModelArgs):
     rope_theta: float = 10000.0
     rope_scaling: Dict = None
     attention_bias: bool = False
+    indexer_rope_interleave: bool = False
 
 
 class Indexer(nn.Module):
@@ -71,7 +72,7 @@ def __init__(self, args: ModelArgs):
         self.rope = initialize_rope(
             dims=args.qk_rope_head_dim,
             base=args.rope_theta,
-            traditional=True,
+            traditional=args.indexer_rope_interleave,
             max_position_embeddings=args.max_position_embeddings,
             scaling_config=args.rope_scaling,
         )

mlx_lm/models/glm_moe_dsa.py

@@ -42,6 +42,7 @@ class ModelArgs(BaseModelArgs):
     attention_bias: bool
     rope_scaling: Dict = None
     rope_theta: Optional[float] = None
+    indexer_rope_interleave: bool = True
 
     def __post_init__(self):
         self.rope_scaling = self.rope_parameters