mirror of
https://github.com/unslothai/unsloth.git
synced 2026-07-09 15:58:41 +00:00
Studio: stop chat generation on the assistant-turn-end token (fixes Qwen3.5 loop) (#6804)
* 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>
This commit is contained in:
Daniel Han
GitHub
parent
f0a5c52821
commit
f38672da65
5 changed files with 558 additions and 6 deletions
|
|
@ -7,13 +7,16 @@ Inference submodule - backend for model loading and generation.
|
|||
The default get_inference_backend() returns an InferenceOrchestrator that
|
||||
delegates to a subprocess. The original InferenceBackend runs inside the
|
||||
subprocess and can be imported directly from .inference when needed.
|
||||
|
||||
Public names are resolved lazily (PEP 562): importing this package -- or a
|
||||
dependency-light leaf like ``core.inference.chat_eos`` -- must NOT eagerly pull
|
||||
the orchestrator / llama_cpp import chain (httpx, subprocess plumbing, the ML
|
||||
backend and its Studio dependencies). Those load only when a public name is
|
||||
actually accessed, so standalone helpers stay unit-testable without the full
|
||||
inference stack.
|
||||
"""
|
||||
|
||||
from .orchestrator import InferenceOrchestrator, get_inference_backend
|
||||
from .llama_cpp import LlamaCppBackend
|
||||
|
||||
# Expose InferenceOrchestrator as InferenceBackend for backward compat.
|
||||
InferenceBackend = InferenceOrchestrator
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
__all__ = [
|
||||
"InferenceBackend",
|
||||
|
|
@ -21,3 +24,33 @@ __all__ = [
|
|||
"get_inference_backend",
|
||||
"LlamaCppBackend",
|
||||
]
|
||||
|
||||
# name -> (submodule, attribute); InferenceBackend aliases InferenceOrchestrator.
|
||||
_LAZY_ATTRS = {
|
||||
"InferenceOrchestrator": ("orchestrator", "InferenceOrchestrator"),
|
||||
"InferenceBackend": ("orchestrator", "InferenceOrchestrator"),
|
||||
"get_inference_backend": ("orchestrator", "get_inference_backend"),
|
||||
"LlamaCppBackend": ("llama_cpp", "LlamaCppBackend"),
|
||||
}
|
||||
|
||||
|
||||
def __getattr__(name):
|
||||
try:
|
||||
submodule, attr = _LAZY_ATTRS[name]
|
||||
except KeyError:
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from None
|
||||
from importlib import import_module
|
||||
|
||||
value = getattr(import_module(f"{__name__}.{submodule}"), attr)
|
||||
globals()[name] = value # cache so later access skips __getattr__
|
||||
return value
|
||||
|
||||
|
||||
def __dir__():
|
||||
return sorted(set(globals()) | set(__all__))
|
||||
|
||||
|
||||
if TYPE_CHECKING: # keep static analysers / IDEs aware of the lazy names
|
||||
from .llama_cpp import LlamaCppBackend
|
||||
from .orchestrator import InferenceOrchestrator, get_inference_backend
|
||||
InferenceBackend = InferenceOrchestrator
|
||||
|
|
|
|||
109
studio/backend/core/inference/chat_eos.py
Normal file
109
studio/backend/core/inference/chat_eos.py
Normal file
|
|
@ -0,0 +1,109 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Resolve a chat model's assistant-turn-end stop tokens.
|
||||
|
||||
Some checkpoints set eos_token_id to a bare document terminator (Qwen3.5 ships
|
||||
config eos ``<|endoftext|>`` though chat turns end with ``<|im_end|>``, and its
|
||||
small chat variants ship no generation_config), so generation runs past the turn
|
||||
and loops -- re-emitting tool calls or hallucinating ``<|im_start|>`` turns.
|
||||
|
||||
Turn-end markers are derived from the tokenizer's ``chat_template`` (the tokens it
|
||||
actually uses to end a turn), not raw vocab membership: a base/coder model can
|
||||
carry ChatML control tokens in a shared vocab without using them, and a loader
|
||||
may have synced ``eos_token`` to the document terminator. Dependency-light (no
|
||||
torch / unsloth) so it is unit-testable without the full inference stack.
|
||||
"""
|
||||
|
||||
from typing import Optional
|
||||
|
||||
# Canonical assistant-turn-end markers per chat family.
|
||||
_CHAT_TURN_END_TOKENS = (
|
||||
"<|im_end|>", # ChatML: Qwen, Yi
|
||||
"<|eot_id|>", # Llama 3.x
|
||||
"<|eom_id|>", # Llama 3.x tool turns
|
||||
"<end_of_turn>", # Gemma
|
||||
"<turn|>", # Gemma-4
|
||||
"<|end|>", # Phi
|
||||
"<|end_of_turn|>", # OpenChat / Starling (barred, distinct from Gemma's)
|
||||
)
|
||||
# harmony/gpt-oss uses <|end|> as a channel delimiter, not the turn end, and has
|
||||
# its own streamer, so its eos is left untouched.
|
||||
_HARMONY_MARKERS = ("<|channel|>", "<|constrain|>")
|
||||
|
||||
|
||||
def _eos_id_set(eos_token_id) -> set:
|
||||
if isinstance(eos_token_id, (list, tuple)):
|
||||
return {int(t) for t in eos_token_id if t is not None}
|
||||
if eos_token_id is not None:
|
||||
return {int(eos_token_id)}
|
||||
return set()
|
||||
|
||||
|
||||
def _collect_template_text(chat_template) -> str:
|
||||
"""Flatten a tokenizer ``chat_template`` into one scannable string.
|
||||
|
||||
Usually the template is a single jinja string, but multi-variant models
|
||||
(e.g. Hermes-3: a ``default`` plus a ``tool_use`` template) expose it as a
|
||||
``{name: template}`` dict -- or, as stored in tokenizer_config.json, a list
|
||||
of ``{"name": ..., "template": ...}`` dicts. Scanning only the ``str`` case
|
||||
would skip turn-end detection for those valid models, so gather every string
|
||||
leaf (variant names are harmless: they never contain the markers).
|
||||
"""
|
||||
if isinstance(chat_template, str):
|
||||
return chat_template
|
||||
if isinstance(chat_template, dict):
|
||||
values = chat_template.values()
|
||||
elif isinstance(chat_template, (list, tuple)):
|
||||
values = chat_template
|
||||
else:
|
||||
return ""
|
||||
parts = [_collect_template_text(v) for v in values]
|
||||
return "\n".join(p for p in parts if p)
|
||||
|
||||
|
||||
def resolve_chat_turn_end_eos_ids_using(template_tokenizer, id_tokenizer) -> list:
|
||||
"""eos of ``id_tokenizer`` plus any canonical turn-end marker the
|
||||
``template_tokenizer``'s chat_template uses, resolved to ids on ``id_tokenizer`` --
|
||||
the tokenizer generation actually uses.
|
||||
|
||||
Pass the same tokenizer for both at load time. After a mapped ``get_chat_template``
|
||||
pass the MAPPED tokenizer as ``template_tokenizer`` (it carries the effective
|
||||
template) and the ORIGINAL generation tokenizer as ``id_tokenizer``: a mapped
|
||||
template registered ``map_eos_token=True`` can hand back a tokenizer whose vocab
|
||||
folds the turn-end token onto the doc-eos id, and generate_stream re-reads the
|
||||
original tokenizer, so resolving ids on the mapped tokenizer would store the wrong
|
||||
(doc-eos) id and let generation run past the real turn marker."""
|
||||
ids = _eos_id_set(getattr(id_tokenizer, "eos_token_id", None))
|
||||
template = _collect_template_text(getattr(template_tokenizer, "chat_template", None))
|
||||
if not template or any(h in template for h in _HARMONY_MARKERS):
|
||||
return sorted(ids)
|
||||
unk = getattr(id_tokenizer, "unk_token_id", None)
|
||||
for marker in _CHAT_TURN_END_TOKENS:
|
||||
if marker in template:
|
||||
try:
|
||||
tid = id_tokenizer.convert_tokens_to_ids(marker)
|
||||
except Exception:
|
||||
tid = None
|
||||
if tid is not None and tid != unk and int(tid) >= 0:
|
||||
ids.add(int(tid))
|
||||
return sorted(ids)
|
||||
|
||||
|
||||
def resolve_chat_turn_end_eos_ids(tokenizer) -> list:
|
||||
"""tokenizer.eos plus any canonical turn-end marker the model's chat_template
|
||||
actually uses. Cheap (convert_tokens_to_ids per marker, no get_vocab); intended
|
||||
to be resolved once at load. Returns eos unchanged for harmony templates."""
|
||||
return resolve_chat_turn_end_eos_ids_using(tokenizer, tokenizer)
|
||||
|
||||
|
||||
def chat_eos_repair(current_eos, turn_end_ids) -> Optional[list]:
|
||||
"""Merged eos_token_id list, or None if ``current_eos`` already covers every
|
||||
resolved turn-end id. Used to repair a model's generation_config at load so
|
||||
every ``.generate()`` path (vision, tool loops) stops at the turn boundary."""
|
||||
if not turn_end_ids:
|
||||
return None
|
||||
current_set = _eos_id_set(current_eos)
|
||||
if set(turn_end_ids) <= current_set:
|
||||
return None
|
||||
return sorted(current_set | set(turn_end_ids))
|
||||
|
|
@ -27,6 +27,10 @@ from utils.hardware import (
|
|||
from core.inference.audio_codecs import AudioCodecManager
|
||||
from core.inference.runtime_context import runtime_context_length
|
||||
from core.inference.message_content import content_to_text
|
||||
from core.inference.chat_eos import (
|
||||
chat_eos_repair,
|
||||
resolve_chat_turn_end_eos_ids_using,
|
||||
)
|
||||
from io import StringIO
|
||||
import structlog
|
||||
from loggers import get_logger
|
||||
|
|
@ -210,6 +214,50 @@ class InferenceBackend:
|
|||
# API uses -1 to disable top-k; transformers uses 0.
|
||||
return 0 if top_k < 0 else top_k
|
||||
|
||||
def _resolve_chat_eos(self, model_name: str) -> None:
|
||||
"""Resolve this chat model's assistant-turn-end stop tokens once at load,
|
||||
cache them in model_info, and repair generation_config so every
|
||||
``.generate()`` path stops at the turn boundary.
|
||||
|
||||
Some checkpoints (e.g. Qwen3.5 / Qwen3.6 small chat models) end turns with
|
||||
``<|im_end|>`` but ship ``config.eos_token_id = <|endoftext|>`` and no
|
||||
``generation_config.json``, so paths that read ``generation_config`` (the
|
||||
vision path, tool loops) run past the turn and loop. Turn-end markers are
|
||||
derived from the chat_template (see chat_eos.resolve_chat_turn_end_eos_ids),
|
||||
so base/coder models and harmony templates are left untouched.
|
||||
"""
|
||||
info = self.models.get(model_name) or {}
|
||||
model = info.get("model")
|
||||
container = info.get("tokenizer")
|
||||
tokenizer = getattr(container, "tokenizer", container) # unwrap processors
|
||||
if model is None or tokenizer is None:
|
||||
return
|
||||
# Vision models carry the chat_template on the processor, not the inner
|
||||
# tokenizer. Read markers from whichever has one, but resolve ids on the
|
||||
# generation tokenizer, else the vision path misses the turn-end token.
|
||||
template_source = container if getattr(container, "chat_template", None) else tokenizer
|
||||
try:
|
||||
turn_end_ids = resolve_chat_turn_end_eos_ids_using(template_source, tokenizer)
|
||||
except Exception as e: # never block a load on eos resolution
|
||||
logger.warning("Chat turn-end eos resolution failed for %s: %s", model_name, e)
|
||||
return
|
||||
info["chat_turn_end_eos_ids"] = turn_end_ids
|
||||
|
||||
gen = getattr(model, "generation_config", None)
|
||||
if gen is None:
|
||||
return
|
||||
repaired = chat_eos_repair(gen.eos_token_id, turn_end_ids)
|
||||
if repaired is None:
|
||||
return
|
||||
previous = gen.eos_token_id
|
||||
gen.eos_token_id = repaired
|
||||
logger.info(
|
||||
"Repaired generation_config.eos_token_id for %s: %s -> %s",
|
||||
model_name,
|
||||
previous,
|
||||
repaired,
|
||||
)
|
||||
|
||||
def load_model(
|
||||
self,
|
||||
config: ModelConfig,
|
||||
|
|
@ -496,6 +544,7 @@ class InferenceBackend:
|
|||
max_seq_length,
|
||||
)
|
||||
|
||||
self._resolve_chat_eos(model_name)
|
||||
self._load_chat_template_info(model_name)
|
||||
|
||||
self.active_model_name = model_name
|
||||
|
|
@ -946,6 +995,22 @@ class InferenceBackend:
|
|||
tokenizer,
|
||||
chat_template = template_name,
|
||||
)
|
||||
# The mapper installs the effective template only now, at generate
|
||||
# time, so re-resolve and UNION into the load-time cache (never
|
||||
# overwrite). get_chat_template can return a remapped tokenizer
|
||||
# (turn-end folded onto doc-eos) while generate_stream reads the
|
||||
# original, so take marker strings from the mapped template but
|
||||
# resolve their ids on the original.
|
||||
try:
|
||||
_gen_tok = model_info.get("tokenizer") or tokenizer
|
||||
refreshed = resolve_chat_turn_end_eos_ids_using(
|
||||
getattr(tokenizer, "tokenizer", tokenizer),
|
||||
getattr(_gen_tok, "tokenizer", _gen_tok),
|
||||
)
|
||||
existing = model_info.get("chat_turn_end_eos_ids") or []
|
||||
model_info["chat_turn_end_eos_ids"] = sorted(set(existing) | set(refreshed))
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not refresh chat turn-end eos after template: {e}")
|
||||
else:
|
||||
logger.info(
|
||||
f"No registered Unsloth template for {self.active_model_name}, using tokenizer default"
|
||||
|
|
@ -1382,7 +1447,8 @@ class InferenceBackend:
|
|||
min_p = min_p,
|
||||
repetition_penalty = repetition_penalty,
|
||||
do_sample = temperature > 0,
|
||||
eos_token_id = tokenizer.eos_token_id,
|
||||
# Resolved once at load (chat_template-derived turn-end tokens).
|
||||
eos_token_id = model_info.get("chat_turn_end_eos_ids") or tokenizer.eos_token_id,
|
||||
pad_token_id = tokenizer.eos_token_id
|
||||
if tokenizer.pad_token_id is None
|
||||
else tokenizer.pad_token_id,
|
||||
|
|
|
|||
194
studio/backend/tests/test_chat_eos_template_refresh.py
Normal file
194
studio/backend/tests/test_chat_eos_template_refresh.py
Normal file
|
|
@ -0,0 +1,194 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Mapper models whose own tokenizer ships no chat_template have their turn-end
|
||||
eos resolved at LOAD from an empty template (document eos only). The effective
|
||||
template is installed later, at generate time, via get_chat_template, so the
|
||||
turn-end-eos cache must be refreshed then; otherwise generate_stream runs past
|
||||
the ChatML <|im_end|> boundary and loops (the exact bug this PR fixes).
|
||||
"""
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
_BACKEND = Path(__file__).resolve().parent.parent
|
||||
if str(_BACKEND) not in sys.path:
|
||||
sys.path.insert(0, str(_BACKEND))
|
||||
|
||||
# These tests construct InferenceBackend, pulling the full stack. CI may lack
|
||||
# unsloth/unsloth_zoo (ImportError) or have a broken CUDA/bitsandbytes setup
|
||||
# (RuntimeError); skip at module level so collection is not aborted (exit 2).
|
||||
try:
|
||||
from core.inference import inference as inf_mod # noqa: E402
|
||||
from core.inference.inference import InferenceBackend # noqa: E402
|
||||
except (ImportError, RuntimeError) as exc: # pragma: no cover - env-dependent
|
||||
pytest.skip(
|
||||
f"full inference backend unavailable ({type(exc).__name__}: {exc})",
|
||||
allow_module_level = True,
|
||||
)
|
||||
|
||||
_CHATML = "{% for m in messages %}<|im_start|>{{m.role}}\n{{m.content}}<|im_end|>{% endfor %}"
|
||||
_GEMMA = "{% for m in messages %}<start_of_turn>{{m.role}}\n{{m.content}}<end_of_turn>{% endfor %}"
|
||||
|
||||
|
||||
class _FakeTokenizer:
|
||||
def __init__(
|
||||
self,
|
||||
eos_id,
|
||||
chat_template = "",
|
||||
token_ids = None,
|
||||
):
|
||||
self.eos_token_id = eos_id
|
||||
self.chat_template = chat_template
|
||||
self.pad_token_id = eos_id
|
||||
self.unk_token_id = None
|
||||
self._ids = dict(token_ids or {})
|
||||
|
||||
def convert_tokens_to_ids(self, tok):
|
||||
return self._ids.get(tok)
|
||||
|
||||
|
||||
def test_turn_end_eos_refreshed_after_generate_time_template(monkeypatch):
|
||||
import utils.datasets as ds
|
||||
|
||||
backend = InferenceBackend.__new__(InferenceBackend)
|
||||
backend.active_model_name = "unsloth/qwen2.5-0.5b"
|
||||
|
||||
# No chat_template at load, so the cache stored only the document eos, though
|
||||
# <|im_end|> is atomic in the vocab (unused until the mapper installs a template).
|
||||
bare_tok = _FakeTokenizer(151643, chat_template = "", token_ids = {"<|im_end|>": 151645})
|
||||
model_info = {
|
||||
"tokenizer": bare_tok,
|
||||
"is_vision": False,
|
||||
"chat_turn_end_eos_ids": [151643],
|
||||
}
|
||||
backend.models = {backend.active_model_name: model_info}
|
||||
|
||||
# The mapper installs a ChatML template (turns end with <|im_end|>) at generate time.
|
||||
templated_tok = _FakeTokenizer(151643, chat_template = _CHATML, token_ids = {"<|im_end|>": 151645})
|
||||
monkeypatch.setattr(inf_mod, "get_chat_template", lambda tok, chat_template = None: templated_tok)
|
||||
monkeypatch.setattr(
|
||||
ds, "MODEL_TO_TEMPLATE_MAPPER", {backend.active_model_name: "qwen-2.5"}, raising = False
|
||||
)
|
||||
|
||||
# Stub the tail so the generator runs through the refresh without a real model.
|
||||
monkeypatch.setattr(backend, "_normalize_top_k", lambda k: k, raising = False)
|
||||
monkeypatch.setattr(
|
||||
backend, "_apply_chat_template_for_generation", lambda *a, **k: "PROMPT", raising = False
|
||||
)
|
||||
monkeypatch.setattr(backend, "generate_stream", lambda *a, **k: iter(()), raising = False)
|
||||
|
||||
list(backend._generate_chat_response_inner(messages = [{"role": "user", "content": "hi"}]))
|
||||
|
||||
# After the template is applied the cache must include the ChatML turn-end id.
|
||||
assert model_info["chat_turn_end_eos_ids"] == [151643, 151645]
|
||||
|
||||
|
||||
def test_turn_end_eos_refresh_preserves_load_time_ids_on_destructive_swap(monkeypatch):
|
||||
# Regression: get_chat_template can return a remapped tokenizer (Gemma: <end_of_turn>
|
||||
# folded onto the eos id) while generate_stream re-reads the original. Resolving on
|
||||
# the swap yields a narrower set, so the refresh must UNION, never overwrite.
|
||||
import utils.datasets as ds
|
||||
|
||||
backend = InferenceBackend.__new__(InferenceBackend)
|
||||
backend.active_model_name = "unsloth/gemma-2b-it"
|
||||
|
||||
# Original tokenizer (used by generate_stream): <end_of_turn>=107 distinct from
|
||||
# eos=1, so the load-time cache resolved to [1, 107].
|
||||
orig_tok = _FakeTokenizer(1, chat_template = _GEMMA, token_ids = {"<end_of_turn>": 107})
|
||||
model_info = {
|
||||
"tokenizer": orig_tok,
|
||||
"is_vision": False,
|
||||
"chat_turn_end_eos_ids": [1, 107],
|
||||
}
|
||||
backend.models = {backend.active_model_name: model_info}
|
||||
|
||||
# Destructively-swapped tokenizer: <end_of_turn> now maps onto eos id 1, so
|
||||
# resolving on it yields only [1] (drops 107).
|
||||
swapped_tok = _FakeTokenizer(1, chat_template = _GEMMA, token_ids = {"<end_of_turn>": 1})
|
||||
monkeypatch.setattr(inf_mod, "get_chat_template", lambda tok, chat_template = None: swapped_tok)
|
||||
monkeypatch.setattr(
|
||||
ds, "MODEL_TO_TEMPLATE_MAPPER", {backend.active_model_name: "gemma-3"}, raising = False
|
||||
)
|
||||
|
||||
monkeypatch.setattr(backend, "_normalize_top_k", lambda k: k, raising = False)
|
||||
monkeypatch.setattr(
|
||||
backend, "_apply_chat_template_for_generation", lambda *a, **k: "PROMPT", raising = False
|
||||
)
|
||||
monkeypatch.setattr(backend, "generate_stream", lambda *a, **k: iter(()), raising = False)
|
||||
|
||||
list(backend._generate_chat_response_inner(messages = [{"role": "user", "content": "hi"}]))
|
||||
|
||||
# The load-time <end_of_turn>=107 must survive: overwriting with the swapped
|
||||
# [1] would regress and loop past the turn.
|
||||
assert model_info["chat_turn_end_eos_ids"] == [1, 107]
|
||||
|
||||
|
||||
def test_turn_end_eos_refresh_resolves_marker_id_on_original_not_remapped(monkeypatch):
|
||||
# Yi-style map_eos_token=True: the original carries <|im_end|> at its own id, but
|
||||
# get_chat_template folds it onto the doc-eos id. generate_stream uses the original,
|
||||
# so read marker strings from the mapped template but ids from the original.
|
||||
import utils.datasets as ds
|
||||
|
||||
backend = InferenceBackend.__new__(InferenceBackend)
|
||||
backend.active_model_name = "01-ai/yi-6b"
|
||||
|
||||
# Original: no template of its own, doc eos = 2, <|im_end|> atomic = 7.
|
||||
orig_tok = _FakeTokenizer(2, chat_template = "", token_ids = {"<|im_end|>": 7})
|
||||
model_info = {
|
||||
"tokenizer": orig_tok,
|
||||
"is_vision": False,
|
||||
"chat_turn_end_eos_ids": [2],
|
||||
}
|
||||
backend.models = {backend.active_model_name: model_info}
|
||||
|
||||
# Remapped tokenizer: ChatML template, but <|im_end|> folded onto doc-eos id 2.
|
||||
remapped_tok = _FakeTokenizer(2, chat_template = _CHATML, token_ids = {"<|im_end|>": 2})
|
||||
monkeypatch.setattr(inf_mod, "get_chat_template", lambda tok, chat_template = None: remapped_tok)
|
||||
monkeypatch.setattr(
|
||||
ds, "MODEL_TO_TEMPLATE_MAPPER", {backend.active_model_name: "chatml"}, raising = False
|
||||
)
|
||||
|
||||
monkeypatch.setattr(backend, "_normalize_top_k", lambda k: k, raising = False)
|
||||
monkeypatch.setattr(
|
||||
backend, "_apply_chat_template_for_generation", lambda *a, **k: "PROMPT", raising = False
|
||||
)
|
||||
monkeypatch.setattr(backend, "generate_stream", lambda *a, **k: iter(()), raising = False)
|
||||
|
||||
list(backend._generate_chat_response_inner(messages = [{"role": "user", "content": "hi"}]))
|
||||
|
||||
# The real <|im_end|>=7 (original vocab) must be recovered, not the remapped 2.
|
||||
assert model_info["chat_turn_end_eos_ids"] == [2, 7]
|
||||
|
||||
|
||||
class _FakeProcessor:
|
||||
"""A ProcessorMixin-like container: carries the chat_template itself and
|
||||
wraps the real text tokenizer as ``.tokenizer`` (the vision layout)."""
|
||||
|
||||
def __init__(self, chat_template, tokenizer):
|
||||
self.chat_template = chat_template
|
||||
self.tokenizer = tokenizer
|
||||
|
||||
|
||||
def test_resolve_chat_eos_reads_vision_processor_template():
|
||||
# Vision model: the chat_template lives on the processor while the inner tokenizer
|
||||
# ships none. _resolve_chat_eos must read the marker from the processor but resolve
|
||||
# its id on the inner tokenizer, and repair generation_config.
|
||||
from types import SimpleNamespace
|
||||
|
||||
inner_tok = _FakeTokenizer(1, chat_template = "", token_ids = {"<end_of_turn>": 107})
|
||||
processor = _FakeProcessor(_GEMMA, inner_tok)
|
||||
model = SimpleNamespace(generation_config = SimpleNamespace(eos_token_id = 1))
|
||||
|
||||
backend = InferenceBackend.__new__(InferenceBackend)
|
||||
backend.active_model_name = "unsloth/gemma-3-4b-it"
|
||||
model_info = {"model": model, "tokenizer": processor, "processor": processor, "is_vision": True}
|
||||
backend.models = {backend.active_model_name: model_info}
|
||||
|
||||
backend._resolve_chat_eos(backend.active_model_name)
|
||||
|
||||
assert model_info["chat_turn_end_eos_ids"] == [1, 107]
|
||||
# generation_config repaired so the vision .generate() path stops at the turn.
|
||||
assert model.generation_config.eos_token_id == [1, 107]
|
||||
150
studio/backend/tests/test_chat_turn_end_eos.py
Normal file
150
studio/backend/tests/test_chat_turn_end_eos.py
Normal file
|
|
@ -0,0 +1,150 @@
|
|||
# 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
|
||||
Loading…
Add table
Add a link
Reference in a new issue