diff --git a/studio/backend/core/inference/__init__.py b/studio/backend/core/inference/__init__.py index 2faf70bb7..ad7815741 100644 --- a/studio/backend/core/inference/__init__.py +++ b/studio/backend/core/inference/__init__.py @@ -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 diff --git a/studio/backend/core/inference/chat_eos.py b/studio/backend/core/inference/chat_eos.py new file mode 100644 index 000000000..2a5d0db22 --- /dev/null +++ b/studio/backend/core/inference/chat_eos.py @@ -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 + "", # Gemma + "", # 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)) diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index 4dca4db76..eaee5a213 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -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, diff --git a/studio/backend/tests/test_chat_eos_template_refresh.py b/studio/backend/tests/test_chat_eos_template_refresh.py new file mode 100644 index 000000000..75d011701 --- /dev/null +++ b/studio/backend/tests/test_chat_eos_template_refresh.py @@ -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 %}{{m.role}}\n{{m.content}}{% 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: + # 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): =107 distinct from + # eos=1, so the load-time cache resolved to [1, 107]. + orig_tok = _FakeTokenizer(1, chat_template = _GEMMA, token_ids = {"": 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: now maps onto eos id 1, so + # resolving on it yields only [1] (drops 107). + swapped_tok = _FakeTokenizer(1, chat_template = _GEMMA, token_ids = {"": 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 =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 = {"": 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] diff --git a/studio/backend/tests/test_chat_turn_end_eos.py b/studio/backend/tests/test_chat_turn_end_eos.py new file mode 100644 index 000000000..c49e39f8f --- /dev/null +++ b/studio/backend/tests/test_chat_turn_end_eos.py @@ -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 while keeping a document eos, so must + # be added as a stop token. + tok = _FakeTokenizer( + 1, chat_template = ".........", token_ids = {"": 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 ). eos synced to =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