unsloth/studio/backend
Daniel Han f38672da65
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>
2026-07-06 10:07:56 -07:00
..
assets Studio: require signed capability tokens for /p preview links (#6666) 2026-06-25 21:40:48 -07:00
auth Studio: opt-in OpenAI /v1 model auto-switch and idle keep-warm (#6392) 2026-07-01 06:42:23 -07:00
core Studio: stop chat generation on the assistant-turn-end token (fixes Qwen3.5 loop) (#6804) 2026-07-06 10:07:56 -07:00
hub Fix Studio custom folders on Linux external drives (#6799) 2026-07-03 19:10:04 +01:00
loggers Studio: trim serving-log noise and surface llama-server engine stats (#6377) 2026-06-17 05:37:57 -07:00
models Studio: multi-select export formats, portable FP8/INT8, GGUF LoRA, and source parity (#6767) 2026-07-03 08:25:10 -07:00
plugins Reduce and tighten code comments and docstrings repo-wide (#6095) 2026-06-08 23:09:51 -07:00
requirements Whole-document context for RAG chat attachments (#6693) 2026-06-30 15:55:23 +02:00
routes studio: tool calling + healing parity for Llama-3, Mistral, Gemma 4 on safetensors + MLX (#5620) 2026-07-06 10:06:06 -07:00
state Studio: shareable per-checkpoint preview links (#6486) 2026-06-24 06:31:53 -07:00
storage Fix Studio custom folders on Linux external drives (#6799) 2026-07-03 19:10:04 +01:00
tests Studio: stop chat generation on the assistant-turn-end token (fixes Qwen3.5 loop) (#6804) 2026-07-06 10:07:56 -07:00
utils Auto Xet to HTTP download fallback in from_pretrained; share Studio's fallback via unsloth_zoo (#6638) 2026-07-06 05:13:25 -07:00
__init__.py Final cleanup 2026-03-12 18:28:04 +00:00
_platform_compat.py Reduce and tighten code comments and docstrings repo-wide (#6095) 2026-06-08 23:09:51 -07:00
cloudflare_tunnel.py Reap Studio child processes when the parent dies abnormally (#6425) 2026-06-18 05:51:22 -07:00
colab.py Studio Colab: opt-in shareable Cloudflare tunnel link (#6684) 2026-06-26 00:56:23 -07:00
main.py Studio: multi-select export formats, portable FP8/INT8, GGUF LoRA, and source parity (#6767) 2026-07-03 08:25:10 -07:00
run.py studio: explicit Cloudflare tunnel notice and public-exposure warning at startup (#6515) 2026-06-30 17:47:48 +02:00
startup_banner.py Fix Windows Studio UTF-8 startup handling (#6614) 2026-07-01 13:47:33 +01:00