electron-rare/mlx-lm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

custom dsv32 chat template (#693)

Browse Source 
* custom dsv32 chat template

* use has_chat_template
This commit is contained in:
Awni Hannun
2025-12-22 13:52:58 -08:00
committed by GitHub
parent 1b2d11b5c7
commit 9fe5f43abf
5 changed files with 368 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mlx_lm/cache_prompt.py
+1 -16
View File
@@ -41,16 +41,6 @@ def setup_arg_parser():
default=None,
help="End of sequence token for tokenizer",
)
parser.add_argument(
"--ignore-chat-template",
action="store_true",
help="Use the raw prompt without the tokenizer's chat template.",
)
parser.add_argument(
"--use-default-chat-template",
action="store_true",
help="Use the default chat template",
)
parser.add_argument(
"--max-kv-size",
type=int,
@@ -107,11 +97,7 @@ def main():
args.prompt = sys.stdin.read() if args.prompt == "-" else args.prompt
if args.use_default_chat_template:
if tokenizer.chat_template is None:
tokenizer.chat_template = tokenizer.default_chat_template
if not args.ignore_chat_template and tokenizer.chat_template is not None:
if tokenizer.has_chat_template:
messages = [{"role": "user", "content": args.prompt}]
prompt = tokenizer.apply_chat_template(
messages,
@@ -155,7 +141,6 @@ def main():
print("Saving...")
metadata = {}
metadata["model"] = args.model
metadata["chat_template"] = json.dumps(tokenizer.chat_template)
metadata["tokenizer_config"] = json.dumps(tokenizer_config)
save_prompt_cache(args.prompt_cache_file, cache, metadata)
mlx_lm/generate.py
+1 -7
View File
@@ -1302,15 +1302,9 @@ def main():
if args.chat_template_config is not None:
template_kwargs = json.loads(args.chat_template_config)
if args.use_default_chat_template:
if tokenizer.chat_template is None:
tokenizer.chat_template = tokenizer.default_chat_template
elif using_cache:
tokenizer.chat_template = json.loads(metadata["chat_template"])
prompt = args.prompt.replace("\\n", "\n").replace("\\t", "\t")
prompt = sys.stdin.read() if prompt == "-" else prompt
if not args.ignore_chat_template and tokenizer.chat_template is not None:
if not args.ignore_chat_template and tokenizer.has_chat_template:
if args.system_prompt is not None:
messages = [{"role": "system", "content": args.system_prompt}]
else:
mlx_lm/tokenizer_utils.py
+34 -3
View File
@@ -1,3 +1,4 @@
import importlib
import json
from functools import partial
from json import JSONDecodeError
@@ -248,7 +249,11 @@ class TokenizerWrapper:
"""
def __init__(
self, tokenizer, detokenizer_class=NaiveStreamingDetokenizer, eos_token_ids=None
self,
tokenizer,
detokenizer_class=NaiveStreamingDetokenizer,
eos_token_ids=None,
chat_template=None,
):
self._tokenizer = tokenizer
self._detokenizer_class = detokenizer_class
@@ -261,6 +266,10 @@ class TokenizerWrapper:
self._think_end = None
self._tool_call_start = None
self._tool_call_end = None
self._chat_template = chat_template
self.has_chat_template = (
tokenizer.chat_template is not None or chat_template is not None
)
THINK_TOKENS = [("<think>", "</think>")]
TOOL_CALL_TOKENS = [("<tool_call>", "</tool_call>")]
@@ -278,9 +287,15 @@ class TokenizerWrapper:
self._tool_call_end = tool_call_end
break
def apply_chat_template(self, *args, **kwargs):
def apply_chat_template(self, *args, tokenize=True, **kwargs):
if self._chat_template is not None:
out = self._chat_template(*args, **kwargs)
if tokenize:
out = self._tokenizer.encode(out, add_special_tokens=False)
return out
kwargs["return_dict"] = False
return self._tokenizer.apply_chat_template(*args, **kwargs)
return self._tokenizer.apply_chat_template(*args, tokenize=tokenize, **kwargs)
def add_eos_token(self, token: str):
token_id = None
@@ -450,12 +465,28 @@ def load(
if isinstance(eos_token_ids, int):
eos_token_ids = [eos_token_ids]
tokenizer_config_file = model_path / "tokenizer_config.json"
custom_tokenizer = None
if tokenizer_config_file.exists():
with open(tokenizer_config_file, "r", encoding="utf-8") as fid:
try:
tokenizer_config = json.load(fid)
except JSONDecodeError as e:
raise JSONDecodeError(
"Failed to parse tokenizer_config.json", e.doc, e.pos
)
if tokenizer_type := tokenizer_config.get("tokenizer_type", False):
custom_tokenizer = importlib.import_module(
f"mlx_lm.tokenizers.{tokenizer_type}"
)
if return_tokenizer:
kwargs = tokenizer_config_extra or {}
return TokenizerWrapper(
AutoTokenizer.from_pretrained(model_path, **kwargs),
detokenizer_class,
eos_token_ids=eos_token_ids,
chat_template=getattr(custom_tokenizer, "apply_chat_template", None),
)
else:
return detokenizer_class
mlx_lm/tokenizers/deepseek_v32.py
+330
View File
@@ -0,0 +1,330 @@
import copy
import json
import re
from typing import Any, Dict, List, Optional, Tuple, Union
TOOLS_SYSTEM_TEMPLATE = """## Tools
You have access to a set of tools you can use to answer the user's question.
You can invoke functions by writing a "<{dsml_token}function_calls>" block like the following as part of your reply to the user:
<{dsml_token}function_calls>
<{dsml_token}invoke name="$FUNCTION_NAME">
<{dsml_token}parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE</{dsml_token}parameter>
...
</{dsml_token}invoke>
<{dsml_token}invoke name="$FUNCTION_NAME2">
...
</{dsml_token}invoke>
</{dsml_token}function_calls>
String and scalar parameters should be specified as is without any escaping or quotes, while lists and objects should use JSON format. The "string" attribute should be set to "true" for string type parameters and "false" for other types (numbers, booleans, arrays, objects).
If the thinking_mode is enabled, then after function results you should strongly consider outputting a thinking block. Here is an example:
<{dsml_token}function_calls>
...
</{dsml_token}function_calls>
<function_results>
...
</function_results>
{thinking_start_token}...thinking about results{thinking_end_token}
Here are the functions available in JSONSchema format:
<functions>
{tool_schemas}
</functions>
"""
bos_token: str = "<｜begin▁of▁sentence｜>"
eos_token: str = "<｜end▁of▁sentence｜>"
thinking_start_token: str = "<think>"
thinking_end_token: str = "</think>"
dsml_token: str = "｜DSML｜"
system_msg_template: str = "{content}"
user_msg_template: str = "<｜User｜>{content}<｜Assistant｜>"
assistant_msg_template: str = "{reasoning}{content}{tool_calls}<｜end▁of▁sentence｜>"
thinking_template = "{reasoning_content}"
response_format_template: str = (
"## Response Format:\n\nYou MUST strictly adhere to the following schema to reply:\n{schema}"
)
tool_call_template: str = (
'<{dsml_token}invoke name="{name}">\n{arguments}\n</{dsml_token}invoke>'
)
tool_calls_template = (
"<{dsml_token}function_calls>\n{tool_calls}\n</{dsml_token}function_calls>"
)
tool_output_template: str = "\n<result>{content}</result>"
def to_json(value: Any) -> str:
try:
return json.dumps(value, ensure_ascii=False)
except:
return json.dumps(value, ensure_ascii=True)
def tools_from_openai_format(tools):
return [tool["function"] for tool in tools]
def tool_calls_from_openai_format(tool_calls):
return [
{
"name": tool_call["function"]["name"],
"arguments": tool_call["function"]["arguments"],
}
for tool_call in tool_calls
]
def encode_arguments_to_dsml(tool_call: Dict[str, str]) -> str:
p_dsml_template = """<{dsml_token}parameter name="{key}" string="{is_str}">{value}</{dsml_token}parameter>"""
P_dsml_strs = []
arguments = json.loads(tool_call["arguments"])
for k, v in arguments.items():
p_dsml_str = p_dsml_template.format(
dsml_token=dsml_token,
key=k,
is_str="true" if isinstance(v, str) else "false",
value=v if isinstance(v, str) else to_json(v),
)
P_dsml_strs.append(p_dsml_str)
return "\n".join(P_dsml_strs)
def decode_dsml_to_arguments(
tool_name: str, tool_args: Dict[str, Tuple[str, str]]
) -> Dict[str, str]:
def _decode_value(key: str, value: str, string: str):
if string == "true":
value = to_json(value)
return f"{to_json(key)}: {value}"
tool_args_json = (
"{"
+ ", ".join(
[_decode_value(k, v, string=is_str) for k, (v, is_str) in tool_args.items()]
)
+ "}"
)
return dict(name=tool_name, arguments=tool_args_json)
def render_tools(tools: List[Dict[str, Union[str, Dict[str, Any]]]]) -> str:
tools_json = [to_json(t) for t in tools]
return TOOLS_SYSTEM_TEMPLATE.format(
tool_schemas="\n".join(tools_json),
dsml_token=dsml_token,
thinking_start_token=thinking_start_token,
thinking_end_token=thinking_end_token,
)
def find_last_user_index(messages: List[Dict[str, Any]]) -> int:
last_user_index = -1
for idx in range(len(messages) - 1, -1, -1):
if messages[idx].get("role") in ["user", "developer"]:
last_user_index = idx
break
return last_user_index
def render_message(
index: int, messages: List[Dict[str, Any]], thinking_mode: str
) -> str:
assert 0 <= index < len(messages)
assert thinking_mode in [
"chat",
"thinking",
], f"Invalid thinking_mode `{thinking_mode}`"
prompt = ""
msg = messages[index]
last_user_idx = find_last_user_index(messages)
role = msg.get("role")
content = msg.get("content")
tools = msg.get("tools")
response_format = msg.get("response_format")
tool_calls = msg.get("tool_calls")
reasoning_content = msg.get("reasoning_content")
if tools:
tools = tools_from_openai_format(tools)
if tool_calls:
tool_calls = tool_calls_from_openai_format(tool_calls)
if role == "system":
prompt += system_msg_template.format(content=content or "")
if tools:
prompt += "\n\n" + render_tools(tools)
if response_format:
prompt += "\n\n" + response_format_template.format(
schema=to_json(response_format)
)
elif role == "developer":
assert content, f"Invalid message for role `{role}`: {msg}"
content_developer = ""
if tools:
content_developer += "\n\n" + render_tools(tools)
if response_format:
content_developer += "\n\n" + response_format_template.format(
schema=to_json(response_format)
)
content_developer += "\n\n# The user's message is: {}".format(content)
prompt += user_msg_template.format(content=content_developer)
if index == last_user_idx and thinking_mode == "thinking":
prompt += thinking_start_token
else:
prompt += thinking_end_token
elif role == "user":
prompt += user_msg_template.format(content=content)
if index == last_user_idx and thinking_mode == "thinking":
prompt += thinking_start_token
else:
prompt += thinking_end_token
elif role == "tool":
prev_assistant_idx = index - 1
assistant_msg = messages[prev_assistant_idx]
while prev_assistant_idx >= 0 and assistant_msg.get("role") == "tool":
prev_assistant_idx -= 1
assistant_msg = messages[prev_assistant_idx]
assert (
index == 0
or prev_assistant_idx >= 0
and assistant_msg.get("role") == "assistant"
), f"Invalid messages at {index}:\n{assistant_msg}"
tool_call_order = index - prev_assistant_idx
assistant_tool_calls = assistant_msg.get("tool_calls")
assert (
assistant_tool_calls and len(assistant_tool_calls) >= tool_call_order
), "No tool calls but found tool output"
if tool_call_order == 1:
prompt += "\n\n<function_results>"
prompt += tool_output_template.format(content=content)
if tool_call_order == len(assistant_tool_calls):
prompt += "\n</function_results>"
if index >= last_user_idx and thinking_mode == "thinking":
prompt += "\n\n" + thinking_start_token
else:
prompt += "\n\n" + thinking_end_token
elif role == "assistant":
prev_assistant_idx = index
thinking_part = ""
tool_calls_content = ""
if tool_calls:
tool_calls = [
tool_call_template.format(
dsml_token=dsml_token,
name=tool_call.get("name"),
arguments=encode_arguments_to_dsml(tool_call),
)
for tool_call in tool_calls
]
tool_calls_content += "\n\n" + tool_calls_template.format(
dsml_token=dsml_token, tool_calls="\n".join(tool_calls)
)
summary_content = content or ""
if thinking_mode == "thinking" and index > last_user_idx:
assert (
reasoning_content or tool_calls
), f"ThinkingMode: {thinking_mode}, invalid message without reasoning_content/tool_calls `{msg}` after last user message"
thinking_part = (
thinking_template.format(reasoning_content=reasoning_content or "")
+ thinking_end_token
)
prompt += assistant_msg_template.format(
reasoning=thinking_part,
content=summary_content,
tool_calls=tool_calls_content,
)
else:
raise NotImplementedError(f"Unknown role: {role}")
return prompt
def drop_thinking_messages(
messages: List[Dict[str, Any]], last_user_idx: Optional[int] = None
) -> List[Dict[str, Any]]:
messages_wo_thinking: List[Dict[str, Any]] = []
last_user_idx = (
find_last_user_index(messages) if last_user_idx is None else last_user_idx
)
for idx, msg in enumerate(messages):
role = msg.get("role")
if role in ["user", "system", "tool"] or idx >= last_user_idx:
messages_wo_thinking.append(msg)
continue
elif role == "assistant":
msg_wo_thinking = copy.copy(msg)
msg_wo_thinking.pop("reasoning_content", None)
messages_wo_thinking.append(msg_wo_thinking)
return messages_wo_thinking
def encode_messages(
messages: List[Dict[str, Any]],
thinking_mode: str = "thinking",
context: Optional[List[Dict[str, Any]]] = None,
drop_thinking: bool = True,
add_default_bos_token: bool = True,
) -> str:
context = context if context else []
full_messages = context + messages
prompt = bos_token if add_default_bos_token and len(context) == 0 else ""
if thinking_mode == "thinking" and drop_thinking:
full_messages = drop_thinking_messages(full_messages)
for idx in range(len(messages)):
prompt += render_message(
idx + len(context), full_messages, thinking_mode=thinking_mode
)
return prompt
def apply_chat_template(
messages, continue_final_message=False, add_generation_prompt=False, **kwargs
):
out = encode_messages(messages, **kwargs)
if continue_final_message and add_generation_prompt:
raise ValueError(
"Only one of continue_final_message or add_generation_prompt can be True"
)
if not add_generation_prompt and messages[-1]["role"] == "user":
out = out.removesuffix("<｜Assistant｜><think>")
if continue_final_message and messages[-1]["role"] == "assistant":
out = out.removesuffix(eos_token)
return out
mlx_lm/utils.py
+2 -9
View File
@@ -5,7 +5,6 @@ import glob
import importlib
import inspect
import json
import logging
import os
import resource
import shutil
@@ -71,7 +70,6 @@ def _get_classes(config: dict):
arch = importlib.import_module(f"mlx_lm.models.{model_type}")
except ImportError:
msg = f"Model type {model_type} not supported."
logging.error(msg)
raise ValueError(msg)
return arch.Model, arch.ModelArgs
@@ -145,13 +143,8 @@ def hf_repo_to_path(hf_repo):
def load_config(model_path: Path) -> dict:
try:
with open(model_path / "config.json", "r") as f:
config = json.load(f)
except FileNotFoundError:
logging.error(f"Config file not found in {model_path}")
raise
return config
with open(model_path / "config.json", "r") as f:
return json.load(f)
def load_model(