f38672da65* Studio: stop chat generation on the assistant-turn-end token A small chat model (e.g. Qwen3.5-0.8B) looped on the safetensors path: it emitted a valid response or tool call, then ran past its turn and re-emitted the call, hallucinating <|im_start|>user turns. Root cause: the model's tokenizer.eos_token is synced to the config document terminator (<|endoftext|>, 248044) while chat turns actually end with <|im_end|> (248046), so generate_stream's single eos_token_id never stopped at the turn boundary. Stop on every assistant-turn-end marker the vocab defines (tokenizer.eos plus <|im_end|>, <|eot_id|>, <end_of_turn>, ...). Verified on the real weights: the single-eos control loops (400 tokens) while the fixed set yields a clean 38-token tool call and a clean answer from the tool result. No-op when eos is already the turn-ender (the id just dedups). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: repair chat generation_config.eos_token_id at load time Qwen3.5 / Qwen3.6 small chat checkpoints declare the chat turn-end as tokenizer.eos_token (<|im_end|>) but ship config.eos_token_id = <|endoftext|> and no generation_config.json (upstream shipped generation_config only on the large chat models). So every .generate() path that reads generation_config -- the vision path and tool loops, not just generate_stream -- never stops at the turn boundary and loops. At load time, when the tokenizer's own eos is a chat turn-end marker but generation_config.eos_token_id omits it, add it. This fixes the config once for all generation paths and complements the generate_stream turn-end stop. No-op for base models (eos is a plain document terminator) and already-correct configs. Verified on unsloth/Qwen3.5-0.8B: 248044 -> [248044, 248046]. * Studio: derive chat turn-end eos from the template, resolve once at load Address PR review of the turn-end stop handling: - Do not call tokenizer.get_vocab() per generation request (serializes the whole 100k+ vocab). Resolve the turn-end tokens once at load and cache them on model_info; generate_stream reads the cache. - Derive turn-end markers from the chat_template the model actually uses, not raw vocab membership, so a base/coder model that merely carries ChatML control tokens in a shared vocab is not stopped early, and a loader that synced tokenizer.eos to the document terminator is still covered. - Skip harmony/gpt-oss templates: <|end|> there is an intra-message channel delimiter, not the turn end (dropped <|return|> from the marker list too). - Move the logic to a dependency-light module (core.inference.chat_eos) so the unit test does not import the full unsloth/torch inference stack. Verified on unsloth/Qwen3.5-0.8B (gen_config 248044 -> [248044, 248046], clean 38-token tool call with generation_config-only stopping), Phi-3.5 (adds <|end|>), Llama-3 / Qwen3 (unchanged), and a harmony template (left untouched). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: refresh turn-end eos after the mapper installs its template For a MODEL_TO_TEMPLATE_MAPPER model whose own tokenizer ships no chat_template, the effective template is applied at generate time via get_chat_template, but the turn-end eos ids were resolved once at load when the template was still empty, so only the document eos was cached. Qwen2.5 / Yi base checkpoints (eos <|endoftext|>, ChatML turns end with <|im_end|>) then run past the assistant boundary in generate_stream and loop. Re-resolve the turn-end eos from the now-templated tokenizer and refresh the cached ids right after applying the mapper template, so generate_stream stops at the ChatML turn end. Add a regression test. * Studio: union turn-end eos refresh into load-time cache instead of overwriting get_chat_template can return a different tokenizer whose vocab was remapped (Gemma folds <end_of_turn> onto the eos id), while generate_stream re-reads the original model_info tokenizer. Overwriting the cache with the refreshed set dropped a valid load-time id (e.g. <end_of_turn>=107) and let generation run past the real turn marker. Union the refresh into the existing cache so it can only add ids, never drop a valid one. Add a regression test covering the destructive-swap case the prior test missed. * Studio: resolve refreshed turn-end ids on the generation tokenizer, add Gemma-4 marker Two residual gaps in the turn-end eos refresh: - For map_eos_token=True mapped templates (e.g. chatml on a Yi-6B base), get_chat_template returns a tokenizer whose vocab folds the turn-end token onto the document eos id, while generate_stream re-reads the original tokenizer. The refresh resolved ids on the returned tokenizer, so it stored the doc eos and missed the real turn-end id, and generation ran past the boundary. Read the turn-end marker strings from the mapped template but resolve their ids on the original generation tokenizer (new resolve_chat_turn_end_eos_ids_using). - Add Gemma-4's <turn|> turn terminator to the marker allowlist; those templates keep a document eos so resolve otherwise missed the real turn marker. Add regression tests for both. * Fix turn-end detection for Starling, multi-variant and vision templates; keep tests collectable The turn-end marker set missed OpenChat/Starling's barred <|end_of_turn|> (distinct from Gemma's unbarred form), so Starling generations ran past the assistant boundary. A dict/list chat_template (Hermes-3 style default+tool_use variants) hit an early non-string return and skipped detection; flatten and scan every variant. Vision models carry the chat_template on the ProcessorMixin, not the unwrapped inner tokenizer, so read markers from the template-carrying container while resolving ids on the generation tokenizer. The refresh test constructs the real backend, so it is guarded with a module-level skip when unsloth/unsloth_zoo is absent (the lightweight pytest matrix), and core.inference package init is made lazy so the dependency-light chat_eos tests collect without the heavy stack. * Studio: tighten chat turn-end eos comments * Studio: condense chat turn-end eos comments --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
150 lines
6 KiB
Python
150 lines
6 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""chat_eos: resolve assistant-turn-end stop tokens from the chat_template and
|
|
repair generation_config so a chat model whose eos is a bare document terminator
|
|
(Qwen3.5: config eos <|endoftext|>, turns end with <|im_end|>) stops at the turn
|
|
boundary instead of running past it and looping. Dependency-light: imported here
|
|
without the full inference stack.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
_BACKEND = Path(__file__).resolve().parent.parent
|
|
if str(_BACKEND) not in sys.path:
|
|
sys.path.insert(0, str(_BACKEND))
|
|
|
|
from core.inference.chat_eos import ( # noqa: E402
|
|
chat_eos_repair,
|
|
resolve_chat_turn_end_eos_ids,
|
|
resolve_chat_turn_end_eos_ids_using,
|
|
)
|
|
|
|
|
|
class _FakeTokenizer:
|
|
def __init__(
|
|
self,
|
|
eos_id,
|
|
chat_template = "",
|
|
token_ids = None,
|
|
unk_token_id = None,
|
|
):
|
|
self.eos_token_id = eos_id
|
|
self.chat_template = chat_template
|
|
self.unk_token_id = unk_token_id
|
|
self._ids = dict(token_ids or {})
|
|
|
|
def convert_tokens_to_ids(self, tok):
|
|
return self._ids.get(tok, self.unk_token_id)
|
|
|
|
|
|
# ---- resolve_chat_turn_end_eos_ids ---------------------------------------
|
|
|
|
_CHATML = "{% for m in messages %}<|im_start|>{{m.role}}\n{{m.content}}<|im_end|>{% endfor %}"
|
|
|
|
|
|
def test_qwen35_adds_im_end_from_template():
|
|
# eos synced to <|endoftext|> (248044); template uses <|im_end|> (248046).
|
|
tok = _FakeTokenizer(248044, chat_template = _CHATML, token_ids = {"<|im_end|>": 248046})
|
|
assert resolve_chat_turn_end_eos_ids(tok) == [248044, 248046]
|
|
|
|
|
|
def test_marker_in_vocab_but_not_in_template_is_ignored():
|
|
# Base/coder model: <|im_end|> is in the vocab but the template does not use
|
|
# it, so it must not become a stop token.
|
|
tok = _FakeTokenizer(248044, chat_template = "{{ messages }}", token_ids = {"<|im_end|>": 248046})
|
|
assert resolve_chat_turn_end_eos_ids(tok) == [248044]
|
|
|
|
|
|
def test_harmony_template_is_left_untouched():
|
|
# gpt-oss/harmony: <|end|> is a channel delimiter, not the turn end.
|
|
harmony = "<|start|>assistant<|channel|>analysis<|message|>...<|end|>"
|
|
tok = _FakeTokenizer(200002, chat_template = harmony, token_ids = {"<|end|>": 200007})
|
|
assert resolve_chat_turn_end_eos_ids(tok) == [200002]
|
|
|
|
|
|
def test_llama3_eot_id_from_template():
|
|
tok = _FakeTokenizer(128001, chat_template = "...<|eot_id|>...", token_ids = {"<|eot_id|>": 128009})
|
|
assert resolve_chat_turn_end_eos_ids(tok) == [128001, 128009]
|
|
|
|
|
|
def test_gemma4_turn_marker_from_template():
|
|
# Gemma-4 ends turns with <turn|> while keeping a document eos, so <turn|> must
|
|
# be added as a stop token.
|
|
tok = _FakeTokenizer(
|
|
1, chat_template = "...<start_of_turn>...<turn|>...", token_ids = {"<turn|>": 106}
|
|
)
|
|
assert resolve_chat_turn_end_eos_ids(tok) == [1, 106]
|
|
|
|
|
|
def test_resolve_using_reads_markers_from_template_but_ids_from_generation_tokenizer():
|
|
# map_eos_token=True: the mapped template remaps <|im_end|> onto the doc-eos id,
|
|
# but the original keeps it atomic. Reading marker STRINGS from the template but
|
|
# IDS on the original recovers the real turn-end id (7), not the doc-eos id (2).
|
|
template_tok = _FakeTokenizer(2, chat_template = _CHATML, token_ids = {"<|im_end|>": 2})
|
|
id_tok = _FakeTokenizer(2, chat_template = "", token_ids = {"<|im_end|>": 7})
|
|
assert resolve_chat_turn_end_eos_ids_using(template_tok, id_tok) == [2, 7]
|
|
# Same tokenizer for both reproduces the plain resolve (load-time behaviour).
|
|
assert resolve_chat_turn_end_eos_ids_using(template_tok, template_tok) == [2]
|
|
|
|
|
|
def test_list_eos_preserved():
|
|
tok = _FakeTokenizer([1, 2], chat_template = _CHATML, token_ids = {"<|im_end|>": 2})
|
|
assert resolve_chat_turn_end_eos_ids(tok) == [1, 2]
|
|
|
|
|
|
def test_missing_marker_maps_to_unk_and_is_skipped():
|
|
tok = _FakeTokenizer(7, chat_template = _CHATML, token_ids = {}, unk_token_id = 0)
|
|
assert resolve_chat_turn_end_eos_ids(tok) == [7]
|
|
|
|
|
|
def test_starling_barred_end_of_turn_from_template():
|
|
# OpenChat/Starling end turns with the BARRED <|end_of_turn|> (distinct from
|
|
# Gemma's <end_of_turn>). eos synced to </s>=2, turn marker at 32000.
|
|
starling = "GPT4 Correct Assistant: hi<|end_of_turn|>"
|
|
tok = _FakeTokenizer(2, chat_template = starling, token_ids = {"<|end_of_turn|>": 32000})
|
|
assert resolve_chat_turn_end_eos_ids(tok) == [2, 32000]
|
|
|
|
|
|
def test_dict_chat_template_scans_all_variants():
|
|
# Hermes-3 style: chat_template is a {name: template} dict. Detection must scan
|
|
# every variant, not bail because the container is not a plain str.
|
|
tmpl = {"default": "{{ messages }}", "tool_use": _CHATML}
|
|
tok = _FakeTokenizer(2, chat_template = tmpl, token_ids = {"<|im_end|>": 5})
|
|
assert resolve_chat_turn_end_eos_ids(tok) == [2, 5]
|
|
|
|
|
|
def test_list_of_dicts_chat_template_scans_all_variants():
|
|
# tokenizer_config.json stores multi-templates as a list of {name, template}.
|
|
tmpl = [{"name": "default", "template": _CHATML}]
|
|
tok = _FakeTokenizer(2, chat_template = tmpl, token_ids = {"<|im_end|>": 5})
|
|
assert resolve_chat_turn_end_eos_ids(tok) == [2, 5]
|
|
|
|
|
|
def test_dict_harmony_template_left_untouched():
|
|
# A multi-variant container whose variant is harmony must still be left alone.
|
|
tmpl = {"default": "<|start|>assistant<|channel|>analysis<|message|>...<|end|>"}
|
|
tok = _FakeTokenizer(200002, chat_template = tmpl, token_ids = {"<|end|>": 200007})
|
|
assert resolve_chat_turn_end_eos_ids(tok) == [200002]
|
|
|
|
|
|
# ---- chat_eos_repair ------------------------------------------------------
|
|
|
|
|
|
def test_repair_adds_missing_turn_end():
|
|
assert chat_eos_repair(248044, [248044, 248046]) == [248044, 248046]
|
|
|
|
|
|
def test_repair_from_missing_generation_config_eos():
|
|
assert chat_eos_repair(None, [248046]) == [248046]
|
|
|
|
|
|
def test_repair_noop_when_already_covered():
|
|
assert chat_eos_repair([248046, 248044], [248046]) is None
|
|
|
|
|
|
def test_repair_noop_when_no_turn_end_ids():
|
|
assert chat_eos_repair(248044, []) is None
|