Fix json_tools: match "<tool_call" prefix to survive ">" token merge - #1336
Closed
snagnever wants to merge 1 commit into
Closed
Fix json_tools: match "<tool_call" prefix to survive ">" token merge#1336snagnever wants to merge 1 commit into
snagnever wants to merge 1 commit into
Conversation
mlx-lm matches tool-call markers as exact token-id sequences (Aho-Corasick over the generated tokens). json_tools's start marker "<tool_call>" encodes to a run ending in a standalone ">" token, but many tokenizers merge that ">" with the next byte -- "<tool_call>\n" becomes a single ">\n" token -- so the precomputed marker is never a contiguous subsequence of the generated stream. The state machine never enters the tool-capture state and a valid tool call is returned as assistant content with tool_calls=null. Match the stable "<tool_call" prefix (the "<", "tool", "_call" tokens, which don't merge) and extract the first brace-balanced JSON object, tolerating the leftover ">" / newline before the JSON and a trailing "</tool_call>" if the end marker likewise merges. Output is identical to json.loads(text.strip()) for clean segments, and streaming is unaffected (server.py accumulates the whole tool segment before parsing). Verified end-to-end on mlx-community/DeepSeek-V4-Flash-2bit-DQ: tool-call parse rate 0/4 -> 4/4 on a probe and 8/40 -> 33/40 on a tool-calling benchmark. Adds tests for the merge cases; existing json_tools cases (clean JSON) still pass. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Member
|
Closed due to #1501. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #1335.
Cause
Tool-call detection matches markers as exact token-id sequences (
SequenceStateMachine, Aho-Corasick).json_tools.tool_call_start = "<tool_call>"encodes to a run ending in a standalone>token — but many tokenizers merge that>with the following byte (<tool_call>\n→ a single>\ntoken), so the precomputed marker is never a contiguous subsequence of the generated stream. The state machine never enters thetoolstate and a valid tool call is returned as assistant content withtool_calls = null. (Token IDs + a deterministic reproducer in the issue.)Change
mlx_lm/tool_parsers/json_tools.py:"<tool_call"(drop the volatile>; the<,tool,_calltokens are stable).parse_tool_callextracts the first brace-balanced JSON object, tolerating the leftover>/newline that the wider capture now includes and a trailing</tool_call>if that marker likewise merges.Why it's safe
<tool_callis a subsequence of what currently-matching models emit, and brace extraction returns output identical tojson.loads(text.strip())for clean segments. The existingjson_toolscases intests/test_tool_parsing.py(clean JSON, nested braces) still pass.server.pyaccumulates the entiretoolsegment (tool_text) while in the tool state and only callsparse_tool_callafter the transition back tonormal, so the leftover>never leaks into a streamed delta.Tests
Adds
test_json_tools_marker_merge(clean JSON, leftover->, trailing-</tool_call>, nested-brace/}-in-string).python -m unittest tests.test_tool_parsing→ all pass.Verification (end-to-end)
On
mlx-community/DeepSeek-V4-Flash-2bit-DQ(markers not special tokens;>merges): tool-call parse rate 0/4 → 4/4 on a probe and 8/40 → 33/40 on jdhodges tool-calling — same checkpoint, no other change.Note for reviewers
This changes the shared
json_toolsparser. I validated output-identity for clean input and the streaming path, but couldn't test against a live Qwen2.5/Hermes model — happy to gate this behind a separate parser (e.g.json_tools_lenient) or adjust the approach if you'd prefer. Full analysis: writeup.