diff --git a/mlx_lm/tokenizer_utils.py b/mlx_lm/tokenizer_utils.py index 4bb44eab1..f5dc5e118 100644 --- a/mlx_lm/tokenizer_utils.py +++ b/mlx_lm/tokenizer_utils.py @@ -253,6 +253,9 @@ def make_byte_decoder(cls): def _infer_thinking(tokenizer): vocab = tokenizer.get_vocab() THINK_TOKENS = [ + # Apertus has / in its vocabulary but never emits them, + # so its own markers must be checked first. + ("<|inner_prefix|>", "<|inner_suffix|>"), ("", ""), ("", ""), ] @@ -566,6 +569,8 @@ def _infer_tool_parser(chat_template): return "qwen3_coder" elif "<|tool_calls_section_begin|>" in chat_template: return "kimi_k2" + elif "<|tools_prefix|>" in chat_template: + return "apertus" elif "[TOOL_CALLS]" in chat_template: return "mistral" elif "" in chat_template and "tool_call.name" in chat_template: diff --git a/mlx_lm/tool_parsers/apertus.py b/mlx_lm/tool_parsers/apertus.py new file mode 100644 index 000000000..05796d422 --- /dev/null +++ b/mlx_lm/tool_parsers/apertus.py @@ -0,0 +1,35 @@ +# Copyright © 2026 Apple Inc. +""" +Modified from: +https://github.com/vllm-project/vllm/blob/main/vllm/tool_parsers/apertus_tool_parser.py +""" + +import json +from typing import Any + +tool_call_start = "<|tools_prefix|>" +tool_call_end = "<|tools_suffix|>" + + +def parse_tool_call(text: str, tools: Any | None = None): + # Apertus emits an array of single key objects mapping the function name to + # its arguments, e.g. + # [{"get_weather": {"location": "London"}}, {"get_time": {}}] + calls = json.loads(text) + if not isinstance(calls, list): + calls = [calls] + + tool_calls = [] + for call in calls: + if not isinstance(call, dict) or not call: + continue + name, arguments = next(iter(call.items())) + # A call with no arguments can come back as null, but clients expect an + # object. + if arguments is None: + arguments = {} + tool_calls.append(dict(name=name, arguments=arguments)) + + if not tool_calls: + raise ValueError(f"Could not parse tool call from: {text}") + return tool_calls diff --git a/tests/test_tokenizers.py b/tests/test_tokenizers.py index d5c81e34a..0797e2d9d 100644 --- a/tests/test_tokenizers.py +++ b/tests/test_tokenizers.py @@ -10,6 +10,7 @@ NaiveStreamingDetokenizer, SPMStreamingDetokenizer, TokenizerWrapper, + _infer_thinking, ) from mlx_lm.utils import load_tokenizer @@ -110,6 +111,35 @@ def test_thinking(self): self.assertIsNone(tokenizer.think_start_id) self.assertIsNone(tokenizer.think_end_id) + def test_thinking_marker_precedence(self): + # Apertus carries unused tokens alongside the markers it + # actually emits, so its own markers have to win. + class Tokenizer: + def __init__(self, vocab): + self._vocab = vocab + + def get_vocab(self): + return self._vocab + + apertus_vocab = { + "<|inner_prefix|>": 32, + "<|inner_suffix|>": 33, + "": 69, + "": 70, + } + think_start, think_end, start_ids, end_ids = _infer_thinking( + Tokenizer(apertus_vocab) + ) + self.assertEqual(think_start, "<|inner_prefix|>") + self.assertEqual(think_end, "<|inner_suffix|>") + self.assertEqual(start_ids, (32,)) + self.assertEqual(end_ids, (33,)) + + think_start, _, _, _ = _infer_thinking(Tokenizer({"": 1, "": 2})) + self.assertEqual(think_start, "") + + self.assertEqual(_infer_thinking(Tokenizer({})), (None, None, None, None)) + def test_find_token(self): # Check that _find returns a valid index when # searching for a think token in short system prompts diff --git a/tests/test_tool_parsing.py b/tests/test_tool_parsing.py index 52892b7ff..3ecdb2fba 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 ( + apertus, function_gemma, gemma4, glm47, @@ -313,6 +314,61 @@ def test_kimi_k2(self): ] self.assertEqual(tool_calls, expected) + def test_apertus(self): + # Single tool call + test_case = '[{"multiply": {"a": 12234585, "b": 48838483920}}]' + tool_calls = apertus.parse_tool_call(test_case, None) + expected = [ + {"name": "multiply", "arguments": {"a": 12234585, "b": 48838483920}} + ] + self.assertEqual(tool_calls, expected) + + # Multiple tool calls + test_case = ( + '[{"get_weather": {"location": "London"}}, ' + '{"get_time": {"location": "London"}}]' + ) + tool_calls = apertus.parse_tool_call(test_case, None) + expected = [ + {"name": "get_weather", "arguments": {"location": "London"}}, + {"name": "get_time", "arguments": {"location": "London"}}, + ] + self.assertEqual(tool_calls, expected) + + # Nested arguments + test_case = ( + '[{"complex_function": {"nested": {"inner": "value"}, "list": ["a", "b"]}}]' + ) + tool_calls = apertus.parse_tool_call(test_case, None) + expected = [ + { + "name": "complex_function", + "arguments": {"nested": {"inner": "value"}, "list": ["a", "b"]}, + } + ] + self.assertEqual(tool_calls, expected) + + # Null arguments are normalized so clients receive an empty object + test_case = '[{"get_time": null}]' + tool_calls = apertus.parse_tool_call(test_case, None) + self.assertEqual(tool_calls, [{"name": "get_time", "arguments": {}}]) + + # A bare object rather than an array + test_case = '{"get_time": {}}' + tool_calls = apertus.parse_tool_call(test_case, None) + self.assertEqual(tool_calls, [{"name": "get_time", "arguments": {}}]) + + # Entries which cannot be a tool call are skipped + test_case = '[{"get_time": {}}, {}, "junk"]' + tool_calls = apertus.parse_tool_call(test_case, None) + self.assertEqual(tool_calls, [{"name": "get_time", "arguments": {}}]) + + # Nothing parseable raises, as does a call truncated mid generation + with self.assertRaises(ValueError): + apertus.parse_tool_call("[]", None) + with self.assertRaises(ValueError): + apertus.parse_tool_call('[{"get_weather": {"location": "London"}', None) + def test_minimax_m2(self): test_case = ( '\n'