From ff6b319779a511715c513beb5b7692322a04aa9c Mon Sep 17 00:00:00 2001 From: Vitor de Araujo Date: Sun, 31 May 2026 12:55:10 -0300 Subject: [PATCH] Add deepseek_dsml tool parser for DeepSeek-V4's native DSML format MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit DeepSeek-V4 (Flash/Pro) emits tool calls in its native DSML format: <|DSML|tool_calls> <|DSML|invoke name="get_weather"> <|DSML|parameter name="city" string="true">Paris ... with multiple <|DSML|invoke> per block (native parallel calls). string="true" means a literal string value, string="false" a JSON value. Adds mlx_lm/tool_parsers/deepseek_dsml.py (modeled on minimax_m2) and an _infer_tool_parser entry so the official DSML chat template auto-selects it. The start/end markers use the "<|DSML|tool_calls" prefix (dropping the trailing ">"): mlx-lm matches markers by token-id sequence and the ">" merges with the following byte on this tokenizer (same class of issue as #1335); the parser extracts the invokes/parameters regardless of the leftover ">". Verified on mlx-community/DeepSeek-V4-Flash-2bit-DQ: 0 -> 39/40 (98%) on a jdhodges-style tool suite (the no-tool-template baseline scored 0), 8/8 on parallel multi-tool cases. Adds tests (single, parallel, mixed string/JSON). Co-Authored-By: Claude Opus 4.8 (1M context) --- mlx_lm/tokenizer_utils.py | 2 + mlx_lm/tool_parsers/deepseek_dsml.py | 79 ++++++++++++++++++++++++++++ tests/test_tool_parsing.py | 39 ++++++++++++++ 3 files changed, 120 insertions(+) create mode 100644 mlx_lm/tool_parsers/deepseek_dsml.py diff --git a/mlx_lm/tokenizer_utils.py b/mlx_lm/tokenizer_utils.py index c7e50fbe7..ac0926791 100644 --- a/mlx_lm/tokenizer_utils.py +++ b/mlx_lm/tokenizer_utils.py @@ -549,6 +549,8 @@ def _infer_tool_parser(chat_template): """Attempt to auto-infer a tool parser from the chat template.""" if not isinstance(chat_template, str): return None + elif "<|DSML|tool_calls>" in chat_template: + return "deepseek_dsml" elif "" in chat_template: return "minimax_m2" elif "<|tool_call>" in chat_template and "" in chat_template: diff --git a/mlx_lm/tool_parsers/deepseek_dsml.py b/mlx_lm/tool_parsers/deepseek_dsml.py new file mode 100644 index 000000000..17adb8e40 --- /dev/null +++ b/mlx_lm/tool_parsers/deepseek_dsml.py @@ -0,0 +1,79 @@ +# Tool parser for DeepSeek-V4's native DSML tool-call format. Deploy into +# venv: mlx_lm/tool_parsers/deepseek_dsml.py, set tokenizer_config +# "tool_parser_type": "deepseek_dsml", and use the official DSML chat template. +# +# Format (multiple invokes per tool_calls block => native parallel calls): +# <|DSML|tool_calls> +# <|DSML|invoke name="get_weather"> +# <|DSML|parameter name="city" string="true">Paris +# <|DSML|parameter name="days" string="false">3 +# +# ... +# +# +# string="true" -> value is a literal string +# string="false" -> value is JSON (numbers, booleans, arrays, objects) +# +# Start/end markers drop the trailing ">": mlx-lm matches markers by token-id +# sequence and BPE merges ">" with the next byte ("...calls>\n" -> one token), +# so the precomputed full marker never matches. "<|DSML|tool_calls" (the +# <,|DSML|,tool,_c,alls tokens) is stable; the parser tolerates the leftover ">". +import json + +import regex as re + +tool_call_start = "<|DSML|tool_calls" +tool_call_end = "", re.DOTALL) +_PARAM = re.compile( + r'<|DSML|parameter\s+name=(?P"[^"]*"|\'[^\']*\'|[^\s>]+)' + r'\s+string="(?Ptrue|false)"\s*>(?P.*?)', + re.DOTALL, +) +_NAME_BODY = re.compile(r'\s*(?P"[^"]*"|\'[^\']*\'|[^\s>]+)\s*>(?P.*)', re.DOTALL) + + +def _unquote(s): + s = s.strip() + if len(s) >= 2 and s[0] in "\"'" and s[-1] == s[0]: + return s[1:-1] + return s + + +def _strip_one_newline(v): + if v.startswith("\n"): + v = v[1:] + if v.endswith("\n"): + v = v[:-1] + return v + + +def parse_tool_call(text, tools=None): + calls = [] + for invoke in _INVOKE.findall(text): + m = _NAME_BODY.match(invoke) + if not m: + continue + name = _unquote(m.group("name")) + args = {} + for pm in _PARAM.finditer(m.group("body")): + pname = _unquote(pm.group("name")) + val = _strip_one_newline(pm.group("val")) + if pm.group("is_str") == "true": + args[pname] = val + else: + try: + args[pname] = json.loads(val) + except (json.JSONDecodeError, ValueError): + args[pname] = val + # template fallback: whole arg blob passed as a single "arguments" param + if set(args) == {"arguments"} and isinstance(args["arguments"], str): + try: + args = json.loads(args["arguments"]) + except (json.JSONDecodeError, ValueError): + pass + calls.append({"name": name, "arguments": args}) + if not calls: + raise ValueError("no DSML tool call found") + return calls[0] if len(calls) == 1 else calls diff --git a/tests/test_tool_parsing.py b/tests/test_tool_parsing.py index 52892b7ff..e720b751c 100644 --- a/tests/test_tool_parsing.py +++ b/tests/test_tool_parsing.py @@ -2,6 +2,7 @@ from pathlib import Path from mlx_lm.tool_parsers import ( + deepseek_dsml, function_gemma, gemma4, glm47, @@ -33,6 +34,10 @@ def test_parsers(self): '{"name": "multiply", "arguments": {"a": 12234585, "b": 48838483920}}', json_tools, ), + ( + '>\n<|DSML|invoke name="multiply">\n<|DSML|parameter name="a" string="false">12234585\n<|DSML|parameter name="b" string="false">48838483920\n\n', + deepseek_dsml, + ), ( '\n12234585\n48838483920\n', minimax_m2, @@ -103,6 +108,10 @@ def test_parsers(self): '{"name": "get_current_temperature", "arguments": {"location": "London"}}', json_tools, ), + ( + '>\n<|DSML|invoke name="get_current_temperature">\n<|DSML|parameter name="location" string="true">London\n\n', + deepseek_dsml, + ), ( '\nLondon\n', minimax_m2, @@ -329,6 +338,36 @@ def test_minimax_m2(self): tool_calls = minimax_m2.parse_tool_call(test_case, None) self.assertEqual(expected, tool_calls) + def test_deepseek_dsml_parallel(self): + # DSML puts multiple <|DSML|invoke> in one <|DSML|tool_calls> block = native + # parallel calls. The leading ">" is the leftover the server captures after the + # "<|DSML|tool_calls" prefix start marker (which drops the ">" that BPE merges). + test_case = ( + ">\n" + '<|DSML|invoke name="get_weather">\n' + '<|DSML|parameter name="city" string="true">Tokyo\n' + "\n" + '<|DSML|invoke name="get_weather">\n' + '<|DSML|parameter name="city" string="true">Paris\n' + "\n" + ) + expected = [ + {"name": "get_weather", "arguments": {"city": "Tokyo"}}, + {"name": "get_weather", "arguments": {"city": "Paris"}}, + ] + self.assertEqual(deepseek_dsml.parse_tool_call(test_case, None), expected) + + # string="true" -> literal string; string="false" -> JSON-decoded value + test_case = ( + '<|DSML|invoke name="calc">\n' + '<|DSML|parameter name="expr" string="true">2+2\n' + '<|DSML|parameter name="n" string="false">3\n' + '<|DSML|parameter name="tags" string="false">["a", "b"]\n' + "" + ) + tc = deepseek_dsml.parse_tool_call(test_case, None) + self.assertEqual(tc["arguments"], {"expr": "2+2", "n": 3, "tags": ["a", "b"]}) + if __name__ == "__main__": unittest.main()