jundot--omlx
wehub-resource-sync
836 行
28 KiB
Python
836 行
28 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
||
"""Tests for protocol-specific output parser sessions."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import sys
|
||
import types
|
||
from types import SimpleNamespace
|
||
|
||
from omlx.adapter.gemma4 import Gemma4OutputParserSession
|
||
from omlx.adapter.harmony import load_harmony_gpt_oss_encoding
|
||
from omlx.adapter.output_parser import detect_output_parser
|
||
|
||
|
||
class FakeDetokenizer:
|
||
def __init__(self, decode_one):
|
||
self._decode_one = decode_one
|
||
self.last_segment = ""
|
||
|
||
def reset(self):
|
||
self.last_segment = ""
|
||
|
||
def add_token(self, token_id: int):
|
||
self.last_segment = self._decode_one(token_id)
|
||
|
||
def finalize(self):
|
||
self.last_segment = ""
|
||
|
||
|
||
class GemmaTokenizer:
|
||
def __init__(self, token_map: dict[int, str]):
|
||
self._token_map = token_map
|
||
|
||
@property
|
||
def detokenizer(self):
|
||
return FakeDetokenizer(lambda token_id: self._token_map[token_id])
|
||
|
||
def decode(self, token_ids, skip_special_tokens: bool = True):
|
||
return "".join(self._token_map[token_id] for token_id in token_ids)
|
||
|
||
|
||
class HarmonyTokenizer:
|
||
def __init__(self, encoding):
|
||
self._encoding = encoding
|
||
|
||
def convert_tokens_to_ids(self, token: str) -> int:
|
||
ids = self._encoding.encode(token, allowed_special="all")
|
||
return ids[0] if ids else -1
|
||
|
||
def decode(self, token_ids, skip_special_tokens: bool = True):
|
||
return self._encoding.decode(token_ids)
|
||
|
||
@property
|
||
def detokenizer(self):
|
||
return FakeDetokenizer(lambda token_id: self._encoding.decode([token_id]))
|
||
|
||
|
||
class CohereTokenizer:
|
||
def __init__(self, token_map: dict[int, str]):
|
||
self._token_map = token_map
|
||
|
||
@property
|
||
def detokenizer(self):
|
||
return FakeDetokenizer(lambda token_id: self._token_map[token_id])
|
||
|
||
def decode(self, token_ids, skip_special_tokens: bool = True):
|
||
return "".join(self._token_map[token_id] for token_id in token_ids)
|
||
|
||
|
||
class DeepSeekV4Tokenizer(CohereTokenizer):
|
||
has_tool_calling = True
|
||
tool_call_start = "<｜DSML｜tool_calls>"
|
||
tool_call_end = "</｜DSML｜tool_calls>"
|
||
|
||
def tool_parser(self, text: str, tools=None):
|
||
from omlx.patches.deepseek_v4.tool_parser_v4 import parse_tool_call
|
||
|
||
return parse_tool_call(text, tools)
|
||
|
||
|
||
class _FakeMelodyOptions:
|
||
def cmd4(self):
|
||
return self
|
||
|
||
def stream_tool_actions(self):
|
||
return self
|
||
|
||
|
||
class _FakeMelodyFilter:
|
||
def __init__(self, options):
|
||
self.options = options
|
||
|
||
def write_decoded(self, decoded_text: str):
|
||
if decoded_text.startswith("R:"):
|
||
return SimpleNamespace(
|
||
content=None,
|
||
reasoning=decoded_text[2:],
|
||
tool_calls=[],
|
||
)
|
||
if decoded_text.startswith("C:"):
|
||
return SimpleNamespace(
|
||
content=decoded_text[2:],
|
||
reasoning=None,
|
||
tool_calls=[],
|
||
)
|
||
if decoded_text.startswith("T1"):
|
||
tool_call = SimpleNamespace(
|
||
index=0,
|
||
id="call_",
|
||
name="look",
|
||
arguments='{"q"',
|
||
)
|
||
return SimpleNamespace(content=None, reasoning=None, tool_calls=[tool_call])
|
||
if decoded_text.startswith("T2"):
|
||
tool_call = SimpleNamespace(
|
||
index=0,
|
||
id="1",
|
||
name="up",
|
||
arguments=':"x"}',
|
||
)
|
||
return SimpleNamespace(content=None, reasoning=None, tool_calls=[tool_call])
|
||
return SimpleNamespace(content=None, reasoning=None, tool_calls=[])
|
||
|
||
def flush_partials(self):
|
||
return SimpleNamespace(content=None, reasoning=None, tool_calls=[])
|
||
|
||
|
||
def _install_fake_melody(monkeypatch):
|
||
module = types.ModuleType("cohere_melody")
|
||
module.PyFilter = _FakeMelodyFilter
|
||
module.PyFilterOptions = _FakeMelodyOptions
|
||
monkeypatch.setitem(sys.modules, "cohere_melody", module)
|
||
|
||
|
||
def _write_json(path, data):
|
||
path.write_text(json.dumps(data))
|
||
|
||
|
||
def _spm_decoder():
|
||
return {
|
||
"type": "Sequence",
|
||
"decoders": [
|
||
{
|
||
"type": "Replace",
|
||
"pattern": {"String": "\u2581"},
|
||
"content": " ",
|
||
},
|
||
{"type": "ByteFallback"},
|
||
{"type": "Fuse"},
|
||
{"type": "Strip", "content": " ", "start": 1, "stop": 0},
|
||
],
|
||
}
|
||
|
||
|
||
class ByteFallbackTokenizer:
|
||
clean_up_tokenization_spaces = False
|
||
vocab = {
|
||
"<pad>": 0,
|
||
"<0xEC>": 1,
|
||
"<0x9E>": 2,
|
||
"<0xA0>": 3,
|
||
}
|
||
|
||
def decode(self, token_ids, skip_special_tokens: bool = True):
|
||
table = {
|
||
0: b"",
|
||
1: bytes([0xEC]),
|
||
2: bytes([0x9E]),
|
||
3: bytes([0xA0]),
|
||
}
|
||
raw = b"".join(table[token_id] for token_id in token_ids)
|
||
if not raw:
|
||
return ""
|
||
if raw == bytes([0xEC, 0x9E, 0xA0]):
|
||
return "\uc7a0"
|
||
return "\ufffd" * sum(1 for token_id in token_ids if token_id != 0)
|
||
|
||
|
||
class TestCohere2MoeOutputParserSession:
|
||
def test_detects_cohere2_moe_from_model_config(self, monkeypatch):
|
||
_install_fake_melody(monkeypatch)
|
||
tokenizer = CohereTokenizer({1: "C:hello"})
|
||
|
||
factory = detect_output_parser(
|
||
"North-Mini-Code",
|
||
tokenizer,
|
||
{"model_type": "cohere2_moe"},
|
||
)
|
||
|
||
assert factory is not None
|
||
assert factory.kind == "cohere2_moe"
|
||
|
||
def test_streams_reasoning_as_think_block_and_visible_content(self, monkeypatch):
|
||
_install_fake_melody(monkeypatch)
|
||
tokenizer = CohereTokenizer(
|
||
{
|
||
1: "R:reasoning",
|
||
2: "C:answer",
|
||
}
|
||
)
|
||
factory = detect_output_parser(
|
||
"North-Mini-Code",
|
||
tokenizer,
|
||
{"model_type": "cohere2_moe"},
|
||
)
|
||
session = factory.create_session(tokenizer)
|
||
|
||
parts = []
|
||
visible = []
|
||
for token_id in [1, 2]:
|
||
result = session.process_token(token_id)
|
||
parts.append(result.stream_text)
|
||
visible.append(result.visible_text)
|
||
final = session.finalize()
|
||
parts.append(final.stream_text)
|
||
visible.append(final.visible_text)
|
||
|
||
assert "".join(parts) == "<think>\nreasoning</think>\nanswer"
|
||
assert "".join(visible) == "<think>\nreasoning</think>\nanswer"
|
||
assert final.tool_calls == []
|
||
assert final.finish_reason is None
|
||
|
||
def test_accumulates_streamed_tool_call_deltas(self, monkeypatch):
|
||
_install_fake_melody(monkeypatch)
|
||
tokenizer = CohereTokenizer({1: "T1", 2: "T2"})
|
||
factory = detect_output_parser(
|
||
"North-Mini-Code",
|
||
tokenizer,
|
||
{"model_type": "cohere2_moe"},
|
||
)
|
||
session = factory.create_session(tokenizer)
|
||
|
||
assert session.process_token(1).stream_text == ""
|
||
assert session.process_token(2).stream_text == ""
|
||
final = session.finalize()
|
||
|
||
assert final.tool_calls == [
|
||
{
|
||
"id": "call_1",
|
||
"name": "lookup",
|
||
"arguments": '{"q":"x"}',
|
||
}
|
||
]
|
||
assert final.finish_reason == "tool_calls"
|
||
|
||
def test_literal_newline_in_arguments_is_reescaped(self, monkeypatch):
|
||
"""Melody may stream literal control chars when the model emits them inside
|
||
JSON string values (e.g. newlines inside code arguments). finalize() must
|
||
re-serialize the accumulated arguments so they are valid JSON."""
|
||
# Build a fake Melody that returns arguments containing a literal newline
|
||
# (U+000A) inside the JSON string value, as the real model sometimes does.
|
||
literal_newline_args = '{"path":"f.py","code":"line1\nline2"}' # literal \n
|
||
|
||
class _FakeMelodyFilterLiteralNewline:
|
||
def __init__(self, options):
|
||
pass
|
||
|
||
def write_decoded(self, decoded_text: str):
|
||
if decoded_text == "TC":
|
||
tc = SimpleNamespace(
|
||
index=0,
|
||
id="call_1",
|
||
name="edit",
|
||
arguments=literal_newline_args,
|
||
)
|
||
return SimpleNamespace(
|
||
content=None, reasoning=None, tool_calls=[tc]
|
||
)
|
||
return SimpleNamespace(content=None, reasoning=None, tool_calls=[])
|
||
|
||
def flush_partials(self):
|
||
return SimpleNamespace(content=None, reasoning=None, tool_calls=[])
|
||
|
||
import types, json as _json
|
||
|
||
module = types.ModuleType("cohere_melody")
|
||
module.PyFilter = _FakeMelodyFilterLiteralNewline
|
||
module.PyFilterOptions = _FakeMelodyOptions
|
||
monkeypatch.setitem(__import__("sys").modules, "cohere_melody", module)
|
||
|
||
tokenizer = CohereTokenizer({"TC": "TC"})
|
||
from omlx.adapter.output_parser import Cohere2MoeOutputParserSession
|
||
|
||
session = Cohere2MoeOutputParserSession.__new__(Cohere2MoeOutputParserSession)
|
||
session._tokenizer = tokenizer
|
||
session._melody = _FakeMelodyFilterLiteralNewline(None)
|
||
session._detokenizer = None
|
||
session._thinking_started = False
|
||
session._thinking_closed = False
|
||
session._tool_calls = {}
|
||
|
||
session.process_token("TC")
|
||
final = session.finalize()
|
||
|
||
assert len(final.tool_calls) == 1
|
||
args_str = final.tool_calls[0]["arguments"]
|
||
# Must be valid strict JSON (no literal control characters)
|
||
parsed = _json.loads(args_str)
|
||
assert parsed["code"] == "line1\nline2"
|
||
# The literal newline must have been escaped
|
||
assert "\n" not in args_str or "\\n" in args_str
|
||
|
||
|
||
class TestGemma4OutputParserSession:
|
||
def test_normal_reasoning_block(self):
|
||
token_map = {
|
||
1: "<|channel>",
|
||
2: "thought\n",
|
||
3: "reasoning",
|
||
4: "<channel|>",
|
||
5: "final answer",
|
||
}
|
||
tokenizer = GemmaTokenizer(token_map)
|
||
session = Gemma4OutputParserSession(tokenizer)
|
||
|
||
stream = []
|
||
visible = []
|
||
for token_id in [1, 2, 3, 4, 5]:
|
||
result = session.process_token(token_id)
|
||
stream.append(result.stream_text)
|
||
visible.append(result.visible_text)
|
||
final = session.finalize()
|
||
stream.append(final.stream_text)
|
||
visible.append(final.visible_text)
|
||
|
||
full_stream = "".join(stream)
|
||
full_visible = "".join(visible)
|
||
|
||
assert full_stream == "<think>\nreasoning</think>\nfinal answer"
|
||
assert full_visible == full_stream
|
||
assert "<|channel>" not in full_stream
|
||
assert "<channel|>" not in full_stream
|
||
|
||
def test_empty_thought_block(self):
|
||
token_map = {
|
||
1: "<|channel>thought\n",
|
||
2: "<channel|>",
|
||
3: "answer",
|
||
}
|
||
tokenizer = GemmaTokenizer(token_map)
|
||
session = Gemma4OutputParserSession(tokenizer)
|
||
|
||
parts = []
|
||
for token_id in [1, 2, 3]:
|
||
parts.append(session.process_token(token_id).stream_text)
|
||
parts.append(session.finalize().stream_text)
|
||
|
||
assert "".join(parts) == "<think>\n</think>\nanswer"
|
||
|
||
def test_partial_marker_across_tokens(self):
|
||
token_map = {
|
||
1: "<|chan",
|
||
2: "nel>thought\nstep 1",
|
||
3: " and step 2<chan",
|
||
4: "nel|>",
|
||
5: "done",
|
||
}
|
||
tokenizer = GemmaTokenizer(token_map)
|
||
session = Gemma4OutputParserSession(tokenizer)
|
||
|
||
parts = []
|
||
for token_id in [1, 2, 3, 4, 5]:
|
||
parts.append(session.process_token(token_id).stream_text)
|
||
parts.append(session.finalize().stream_text)
|
||
|
||
text = "".join(parts)
|
||
assert text == "<think>\nstep 1 and step 2</think>\ndone"
|
||
assert "<|channel>thought" not in text
|
||
assert "<channel|>" not in text
|
||
|
||
def test_suppresses_turn_end_marker(self):
|
||
token_map = {
|
||
1: "<|channel>thought\n",
|
||
2: "reasoning",
|
||
3: "<channel|>",
|
||
4: "answer",
|
||
5: "<turn|>",
|
||
}
|
||
tokenizer = GemmaTokenizer(token_map)
|
||
session = Gemma4OutputParserSession(tokenizer)
|
||
|
||
parts = []
|
||
for token_id in [1, 2, 3, 4, 5]:
|
||
result = session.process_token(token_id)
|
||
parts.append(result.stream_text)
|
||
assert "<turn|>" not in result.stream_text
|
||
assert "<turn|>" not in result.visible_text
|
||
parts.append(session.finalize().stream_text)
|
||
|
||
text = "".join(parts)
|
||
assert text == "<think>\nreasoning</think>\nanswer"
|
||
assert "<turn|>" not in text
|
||
|
||
def test_stray_close_marker_outside_thought_dropped(self):
|
||
"""A bare ``<channel|>`` after the thought block already closed must
|
||
not leak into visible content. Models occasionally emit one in long
|
||
multi-turn contexts and the SDK rejects it as raw markup."""
|
||
token_map = {
|
||
1: "<|channel>thought\n",
|
||
2: "reasoning",
|
||
3: "<channel|>",
|
||
4: "answer",
|
||
5: "<channel|>",
|
||
6: "more",
|
||
}
|
||
tokenizer = GemmaTokenizer(token_map)
|
||
session = Gemma4OutputParserSession(tokenizer)
|
||
|
||
parts = []
|
||
for token_id in [1, 2, 3, 4, 5, 6]:
|
||
parts.append(session.process_token(token_id).stream_text)
|
||
parts.append(session.finalize().stream_text)
|
||
|
||
text = "".join(parts)
|
||
assert text == "<think>\nreasoning</think>\nanswermore"
|
||
assert "<channel|>" not in text
|
||
|
||
def test_stray_open_marker_inside_thought_dropped(self):
|
||
"""A nested ``<|channel>thought\\n`` while already inside a thought
|
||
block must not re-emit ``<think>``. The block stays open until the
|
||
first matching close marker."""
|
||
token_map = {
|
||
1: "<|channel>thought\n",
|
||
2: "step 1",
|
||
3: "<|channel>thought\n",
|
||
4: "step 2",
|
||
5: "<channel|>",
|
||
6: "answer",
|
||
}
|
||
tokenizer = GemmaTokenizer(token_map)
|
||
session = Gemma4OutputParserSession(tokenizer)
|
||
|
||
parts = []
|
||
for token_id in [1, 2, 3, 4, 5, 6]:
|
||
parts.append(session.process_token(token_id).stream_text)
|
||
parts.append(session.finalize().stream_text)
|
||
|
||
text = "".join(parts)
|
||
assert text == "<think>\nstep 1step 2</think>\nanswer"
|
||
assert text.count("<think>\n") == 1
|
||
assert text.count("</think>\n") == 1
|
||
|
||
def test_tool_call_markers_pass_through(self):
|
||
"""Tool-call markup must reach the buffered output text untouched so
|
||
``parse_tool_calls`` can extract the call. ``ToolCallStreamFilter``
|
||
downstream is responsible for removing it from stream deltas."""
|
||
token_map = {
|
||
1: "<|channel>thought\n",
|
||
2: "calling",
|
||
3: "<channel|>",
|
||
4: "<|tool_call>",
|
||
5: "call:bash{cmd:ls}",
|
||
6: "<tool_call|>",
|
||
7: "done",
|
||
}
|
||
tokenizer = GemmaTokenizer(token_map)
|
||
session = Gemma4OutputParserSession(tokenizer)
|
||
|
||
stream_parts = []
|
||
visible_parts = []
|
||
for token_id in [1, 2, 3, 4, 5, 6, 7]:
|
||
result = session.process_token(token_id)
|
||
stream_parts.append(result.stream_text)
|
||
visible_parts.append(result.visible_text)
|
||
final = session.finalize()
|
||
stream_parts.append(final.stream_text)
|
||
visible_parts.append(final.visible_text)
|
||
|
||
stream_text = "".join(stream_parts)
|
||
visible_text = "".join(visible_parts)
|
||
assert stream_text == visible_text
|
||
assert "<|tool_call>" in stream_text
|
||
assert "<tool_call|>" in stream_text
|
||
assert "call:bash{cmd:ls}" in stream_text
|
||
|
||
def test_spm_fallback_buffers_split_utf8(self, tmp_path):
|
||
_write_json(tmp_path / "tokenizer.json", {"decoder": _spm_decoder()})
|
||
session = Gemma4OutputParserSession(
|
||
ByteFallbackTokenizer(),
|
||
model_path=tmp_path,
|
||
)
|
||
|
||
parts = []
|
||
for token_id in [1, 2, 3]:
|
||
parts.append(session.process_token(token_id).stream_text)
|
||
parts.append(session.finalize().stream_text)
|
||
|
||
text = "".join(parts)
|
||
assert text == "\uc7a0"
|
||
assert "\ufffd" not in text
|
||
|
||
|
||
class TestOutputParserFactory:
|
||
def test_detects_deepseek_v4_by_config(self):
|
||
tokenizer = DeepSeekV4Tokenizer({1: "x"})
|
||
factory = detect_output_parser(
|
||
"DeepSeek-V4-Flash-oQ4e",
|
||
tokenizer,
|
||
{"model_type": "deepseek_v4"},
|
||
)
|
||
|
||
assert factory is not None
|
||
assert factory.kind == "deepseek_v4"
|
||
|
||
def test_deepseek_v4_stops_at_first_dsml_tool_block(self):
|
||
tokenizer = DeepSeekV4Tokenizer(
|
||
{
|
||
1: "Before ",
|
||
2: "<｜DSML｜tool",
|
||
3: '_calls>\n<｜DSML｜invoke name="Bash">\n',
|
||
4: '<｜DSML｜parameter name="command" string="true">ls</｜DSML｜parameter>\n'
|
||
"</｜DSML｜invoke>\n",
|
||
5: "</｜DSML｜tool_calls>",
|
||
}
|
||
)
|
||
factory = detect_output_parser(
|
||
"DeepSeek-V4-Flash-oQ4e",
|
||
tokenizer,
|
||
{"model_type": "deepseek_v4"},
|
||
)
|
||
session = factory.create_session(tokenizer)
|
||
|
||
stream = []
|
||
visible = []
|
||
stop_seen = False
|
||
for token_id in [1, 2, 3, 4, 5]:
|
||
result = session.process_token(token_id)
|
||
stream.append(result.stream_text)
|
||
visible.append(result.visible_text)
|
||
stop_seen = stop_seen or result.is_stop
|
||
assert result.record_token is True
|
||
|
||
final = session.finalize()
|
||
stream.append(final.stream_text)
|
||
visible.append(final.visible_text)
|
||
|
||
assert stop_seen is True
|
||
assert "".join(stream) == "Before "
|
||
assert "".join(visible) == "Before "
|
||
assert final.finish_reason == "tool_calls"
|
||
assert len(final.tool_calls) == 1
|
||
assert final.tool_calls[0]["name"] == "Bash"
|
||
assert json.loads(final.tool_calls[0]["arguments"]) == {"command": "ls"}
|
||
|
||
def test_deepseek_v4_drops_text_after_tool_end_in_same_token(self):
|
||
tokenizer = DeepSeekV4Tokenizer(
|
||
{
|
||
1: '<｜DSML｜tool_calls>\n<｜DSML｜invoke name="Bash">\n',
|
||
2: '<｜DSML｜parameter name="command" string="true">ls</｜DSML｜parameter>\n'
|
||
"</｜DSML｜invoke>\n",
|
||
3: "</｜DSML｜tool_calls>\n"
|
||
'<｜DSML｜parameter name="command" string="true">pwd</｜DSML｜parameter>',
|
||
}
|
||
)
|
||
factory = detect_output_parser(
|
||
"DeepSeek-V4-Flash-oQ4e",
|
||
tokenizer,
|
||
{"model_type": "deepseek_v4"},
|
||
)
|
||
session = factory.create_session(tokenizer)
|
||
|
||
stream = []
|
||
stop_seen = False
|
||
for token_id in [1, 2, 3]:
|
||
result = session.process_token(token_id)
|
||
stream.append(result.stream_text)
|
||
stop_seen = stop_seen or result.is_stop
|
||
|
||
final = session.finalize()
|
||
stream.append(final.stream_text)
|
||
|
||
assert stop_seen is True
|
||
assert "".join(stream) == ""
|
||
assert final.finish_reason == "tool_calls"
|
||
assert len(final.tool_calls) == 1
|
||
assert json.loads(final.tool_calls[0]["arguments"]) == {"command": "ls"}
|
||
|
||
def test_detects_minimax_m3_by_config(self):
|
||
tokenizer = CohereTokenizer({1: "x"})
|
||
factory = detect_output_parser(
|
||
"MiniMax-M3-4bit",
|
||
tokenizer,
|
||
{"model_type": "minimax_m3_vl"},
|
||
)
|
||
|
||
assert factory is not None
|
||
assert factory.kind == "minimax_m3"
|
||
|
||
def test_minimax_m3_parser_extracts_tool_calls(self, monkeypatch):
|
||
module = types.ModuleType("mlx_vlm.tool_parsers.minimax_m3")
|
||
|
||
def parse_tool_call(text):
|
||
assert "lookup" in text
|
||
return {"name": "lookup", "arguments": {"query": "mlx"}}
|
||
|
||
module.parse_tool_call = parse_tool_call
|
||
monkeypatch.setitem(sys.modules, "mlx_vlm.tool_parsers.minimax_m3", module)
|
||
|
||
start = "]<]minimax[>[<tool_call>"
|
||
end = "]<]minimax[>[</tool_call>"
|
||
tokenizer = CohereTokenizer(
|
||
{
|
||
1: "before ",
|
||
2: start,
|
||
3: ']<]minimax[>[<invoke name="lookup">',
|
||
4: "]<]minimax[>[</invoke>",
|
||
5: end,
|
||
6: " after",
|
||
}
|
||
)
|
||
factory = detect_output_parser(
|
||
"MiniMax-M3-4bit",
|
||
tokenizer,
|
||
{"model_type": "minimax_m3_vl"},
|
||
)
|
||
session = factory.create_session(tokenizer)
|
||
|
||
visible = []
|
||
stream = []
|
||
for token_id in [1, 2, 3, 4, 5, 6]:
|
||
result = session.process_token(token_id)
|
||
stream.append(result.stream_text)
|
||
visible.append(result.visible_text)
|
||
final = session.finalize()
|
||
|
||
assert "".join(stream) == "before after"
|
||
assert start not in "".join(stream)
|
||
assert "".join(visible) + final.visible_text == "before after"
|
||
assert final.tool_calls == [{"name": "lookup", "arguments": '{"query":"mlx"}'}]
|
||
assert final.finish_reason == "tool_calls"
|
||
|
||
def test_minimax_m3_parser_normalizes_thinking_and_strips_eos(self):
|
||
tokenizer = CohereTokenizer(
|
||
{
|
||
1: "<mm:think>",
|
||
2: "reasoning",
|
||
3: "</mm:think>",
|
||
4: "Answer",
|
||
5: "[e~[",
|
||
6: "]!d~[",
|
||
}
|
||
)
|
||
factory = detect_output_parser(
|
||
"MiniMax-M3-4bit",
|
||
tokenizer,
|
||
{"model_type": "minimax_m3_vl"},
|
||
)
|
||
session = factory.create_session(tokenizer)
|
||
|
||
stream = []
|
||
visible = []
|
||
stop_seen = False
|
||
record_flags = []
|
||
for token_id in [1, 2, 3, 4, 6, 5]:
|
||
result = session.process_token(token_id)
|
||
stream.append(result.stream_text)
|
||
visible.append(result.visible_text)
|
||
stop_seen = stop_seen or result.is_stop
|
||
record_flags.append(result.record_token)
|
||
final = session.finalize()
|
||
stream.append(final.stream_text)
|
||
visible.append(final.visible_text)
|
||
|
||
assert "".join(stream) == "<think>reasoning</think>Answer"
|
||
assert "".join(visible) == "<think>reasoning</think>Answer"
|
||
assert stop_seen is True
|
||
assert record_flags[-1] is False
|
||
|
||
def test_minimax_m3_factory_exposes_native_thinking_markers(self):
|
||
tokenizer = CohereTokenizer({})
|
||
tokenizer.convert_tokens_to_ids = lambda text: {
|
||
"[e~[": 200020,
|
||
"<mm:think>": 200059,
|
||
"</mm:think>": 200060,
|
||
}.get(text, -1)
|
||
tokenizer.unk_token_id = -1
|
||
|
||
factory = detect_output_parser(
|
||
"MiniMax-M3-4bit",
|
||
tokenizer,
|
||
{"model_type": "minimax_m3_vl"},
|
||
)
|
||
|
||
assert factory.thinking_start_text == "<mm:think>"
|
||
assert factory.thinking_start_output_text == "<think>\n"
|
||
assert factory.thinking_end_text == "</mm:think>"
|
||
assert factory.stop_token_ids == {200020}
|
||
|
||
def test_detects_gemma4(self):
|
||
tokenizer = GemmaTokenizer({1: "x"})
|
||
factory = detect_output_parser(
|
||
"google/gemma-4b",
|
||
tokenizer,
|
||
{"model_type": "gemma4"},
|
||
)
|
||
|
||
assert factory is not None
|
||
assert factory.kind == "gemma4"
|
||
|
||
def test_session_receives_model_path_when_provided(self, monkeypatch):
|
||
"""Since #2178 the scheduler's model_name is a display id, so the
|
||
filesystem path must reach parser sessions via model_path."""
|
||
import omlx.adapter.output_parser as output_parser_module
|
||
|
||
seen = {}
|
||
|
||
class RecordingSession:
|
||
def __init__(self, tokenizer, model_path=None):
|
||
seen["model_path"] = model_path
|
||
|
||
monkeypatch.setattr(
|
||
output_parser_module, "MiniMaxM3OutputParserSession", RecordingSession
|
||
)
|
||
tokenizer = CohereTokenizer({})
|
||
tokenizer.convert_tokens_to_ids = lambda text: -1
|
||
tokenizer.unk_token_id = -1
|
||
|
||
factory = detect_output_parser(
|
||
"MiniMax-M3-4bit",
|
||
tokenizer,
|
||
{"model_type": "minimax_m3_vl"},
|
||
model_path="/models/minimax-m3",
|
||
)
|
||
factory.create_session(tokenizer)
|
||
assert seen["model_path"] == "/models/minimax-m3"
|
||
|
||
def test_session_falls_back_to_model_name_without_model_path(self, monkeypatch):
|
||
"""dflash/vlm engines pass their filesystem path as model_name and no
|
||
model_path, so the session fallback must keep using model_name."""
|
||
import omlx.adapter.output_parser as output_parser_module
|
||
|
||
seen = {}
|
||
|
||
class RecordingSession:
|
||
def __init__(self, tokenizer, model_path=None):
|
||
seen["model_path"] = model_path
|
||
|
||
monkeypatch.setattr(
|
||
output_parser_module, "MiniMaxM3OutputParserSession", RecordingSession
|
||
)
|
||
tokenizer = CohereTokenizer({})
|
||
tokenizer.convert_tokens_to_ids = lambda text: -1
|
||
tokenizer.unk_token_id = -1
|
||
|
||
factory = detect_output_parser(
|
||
"/models/MiniMax-M3-4bit",
|
||
tokenizer,
|
||
{"model_type": "minimax_m3_vl"},
|
||
)
|
||
factory.create_session(tokenizer)
|
||
assert seen["model_path"] == "/models/MiniMax-M3-4bit"
|
||
|
||
def test_detects_gemma4_unified_by_config(self):
|
||
tokenizer = GemmaTokenizer({1: "x"})
|
||
factory = detect_output_parser(
|
||
"some-model",
|
||
tokenizer,
|
||
{"model_type": "gemma4_unified"},
|
||
)
|
||
|
||
assert factory is not None
|
||
assert factory.kind == "gemma4"
|
||
|
||
def test_harmony_wrapper_regression(self):
|
||
encoding = load_harmony_gpt_oss_encoding()
|
||
tokenizer = HarmonyTokenizer(encoding)
|
||
factory = detect_output_parser(
|
||
"gpt-oss-20b",
|
||
tokenizer,
|
||
{"model_type": "gpt_oss"},
|
||
)
|
||
|
||
assert factory is not None
|
||
assert factory.kind == "harmony"
|
||
|
||
session = factory.create_session(tokenizer)
|
||
tokens = encoding.encode(
|
||
"<|channel|>analysis<|message|>thinking<|end|>"
|
||
"<|start|>assistant<|channel|>final<|message|>Answer<|return|>",
|
||
allowed_special="all",
|
||
)
|
||
|
||
stream = []
|
||
visible = []
|
||
saw_stop = False
|
||
for token in tokens:
|
||
result = session.process_token(token)
|
||
stream.append(result.stream_text)
|
||
visible.append(result.visible_text)
|
||
saw_stop = saw_stop or result.is_stop
|
||
final = session.finalize()
|
||
stream.append(final.stream_text)
|
||
visible.append(final.visible_text)
|
||
|
||
assert saw_stop is True
|
||
assert "<think>\n" in "".join(stream)
|
||
assert "</think>\n" in "".join(stream)
|
||
assert "".join(visible) == "Answer"
|
||
|
||
def test_harmony_non_streaming_preserves_reasoning(self):
|
||
"""Non-streaming output_text retains analysis-channel reasoning."""
|
||
from omlx.api.thinking import extract_thinking
|
||
|
||
encoding = load_harmony_gpt_oss_encoding()
|
||
tokenizer = HarmonyTokenizer(encoding)
|
||
factory = detect_output_parser(
|
||
"gpt-oss-20b",
|
||
tokenizer,
|
||
{"model_type": "gpt_oss"},
|
||
)
|
||
session = factory.create_session(tokenizer)
|
||
|
||
tokens = encoding.encode(
|
||
"<|channel|>analysis<|message|>Let me think about this<|end|>"
|
||
"<|start|>assistant<|channel|>final<|message|>Four<|return|>",
|
||
allowed_special="all",
|
||
)
|
||
|
||
visible_parts = []
|
||
for token in tokens:
|
||
result = session.process_token(token)
|
||
visible_parts.append(result.visible_text)
|
||
|
||
final = session.finalize()
|
||
visible_parts.append(final.visible_text)
|
||
|
||
# Mirror scheduler aggregation: prepend any parser-provided prefix
|
||
# to the accumulated visible_text before exposing as output_text.
|
||
prefix = getattr(final, "output_text_prefix", "")
|
||
output_text = prefix + "".join(visible_parts)
|
||
|
||
thinking, content = extract_thinking(output_text)
|
||
assert thinking == "Let me think about this"
|
||
assert content == "Four"
|