diff --git a/mlx_lm/chat_templates/deepseek_v32.py b/mlx_lm/chat_templates/deepseek_v32.py index 4bb7c8fe7..021efc5a9 100644 --- a/mlx_lm/chat_templates/deepseek_v32.py +++ b/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( diff --git a/mlx_lm/models/deepseek_v32.py b/mlx_lm/models/deepseek_v32.py index 7c97682e7..d696c02c5 100644 --- a/mlx_lm/models/deepseek_v32.py +++ b/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, ) diff --git a/mlx_lm/models/glm_moe_dsa.py b/mlx_lm/models/glm_moe_dsa.py index 14e96e365..dbaabd79f 100644 --- a/mlx_lm/models/glm_moe_dsa.py +++ b/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