5c2e53606e* Studio: render thinking blocks for safetensors inference with prefilled <think> templates Reasoning templates like Qwen3.6 end the generation prompt with an open <think> tag. skip_prompt streaming drops it, so the frontend never sees the opening tag and shows reasoning as plain text. Detect the prefill and re-emit it at the start of the stream on the transformers and MLX paths. Also stop stripping think tags in _clean_generated_text when a tokenizer marks them special. * Studio: guard think re-emit for special close tags, yield prefill early Address review feedback: - Guard: skip re-emitting the open <think> when the tokenizer marks </think> as a special token, since skip_special_tokens would strip the model's close tag and leave an unclosed block that swallows the answer. Falls back to plain text (pre-fix behaviour) for those tokenizers. - Yield the prefilled <think> before the first token so the thinking block renders during prompt prefill instead of after the first generated token. - Drop the now-unnecessary _clean_generated_text think-tag exemption; the guard handles the special-token case at the source. No mainstream reasoning model (Qwen3.6, Qwen3, DeepSeek-R1, QwQ, GLM-4.6) marks think tags special, so behaviour is unchanged for them. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Lyxot <longyixing331@gmail.com>
89 lines
3.2 KiB
Python
89 lines
3.2 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
|
|
|
|
"""
|
|
Unit tests for detect_think_prefill.
|
|
|
|
Reasoning templates (Qwen3.6-style) end the generation prompt with an open
|
|
``<think>\\n`` so the model starts reasoning immediately. skip_prompt
|
|
streaming drops that opening tag, so the safetensors/MLX paths must re-emit
|
|
it for the frontend's <think> parser to render a thinking block.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
|
|
_backend = os.path.join(os.path.dirname(__file__), "..")
|
|
sys.path.insert(0, _backend)
|
|
|
|
from core.inference.chat_template_helpers import detect_think_prefill
|
|
|
|
|
|
QWEN_PROMPT = "<|im_start|>user\nHi!<|im_end|>\n<|im_start|>assistant\n"
|
|
|
|
|
|
def test_open_think_prefill_reemitted():
|
|
"""Qwen3.6-style enable_thinking=True prompt tail: <think>\\n."""
|
|
assert detect_think_prefill(QWEN_PROMPT + "<think>\n") == "<think>\n"
|
|
|
|
|
|
def test_bare_open_think_prefill_reemitted():
|
|
"""Prefill without trailing newline still detected."""
|
|
assert detect_think_prefill(QWEN_PROMPT + "<think>") == "<think>"
|
|
|
|
|
|
def test_closed_think_prefill_not_reemitted():
|
|
"""enable_thinking=False prefills a closed, empty think block."""
|
|
assert detect_think_prefill(QWEN_PROMPT + "<think>\n\n</think>\n\n") == ""
|
|
|
|
|
|
def test_prompt_without_think_untouched():
|
|
"""Non-reasoning templates produce no prefix."""
|
|
assert detect_think_prefill(QWEN_PROMPT) == ""
|
|
|
|
|
|
def test_historical_think_blocks_ignored():
|
|
"""A closed think block in a prior assistant turn (preserve_thinking)
|
|
must not trigger re-emission when the generation tail is plain."""
|
|
prompt = (
|
|
"<|im_start|>user\nHi!<|im_end|>\n"
|
|
"<|im_start|>assistant\n<think>\nprior reasoning\n</think>\n\nHello!<|im_end|>\n"
|
|
"<|im_start|>user\nAgain?<|im_end|>\n<|im_start|>assistant\n"
|
|
)
|
|
assert detect_think_prefill(prompt) == ""
|
|
|
|
|
|
def test_historical_blocks_plus_open_prefill():
|
|
"""Prior closed blocks plus a fresh open prefill: only the tail matters."""
|
|
prompt = (
|
|
"<|im_start|>assistant\n<think>\nprior\n</think>\n\nHello!<|im_end|>\n"
|
|
"<|im_start|>assistant\n<think>\n"
|
|
)
|
|
assert detect_think_prefill(prompt) == "<think>\n"
|
|
|
|
|
|
def test_content_after_open_tag_not_reemitted():
|
|
"""If non-whitespace follows the tag it is not a plain prefill."""
|
|
assert detect_think_prefill(QWEN_PROMPT + "<think>\npartial reasoning") == ""
|
|
|
|
|
|
def test_empty_and_none_prompts():
|
|
assert detect_think_prefill("") == ""
|
|
assert detect_think_prefill(None) == ""
|
|
|
|
|
|
def test_guard_suppresses_when_close_tag_is_special():
|
|
"""If </think> is a special token, skip_special_tokens strips the model's
|
|
close tag, so re-emitting the open would leave an unclosed block. Guard off."""
|
|
specials = ["<|im_end|>", "<think>", "</think>"]
|
|
assert detect_think_prefill(QWEN_PROMPT + "<think>\n", specials) == ""
|
|
|
|
|
|
def test_guard_emits_when_think_not_special():
|
|
specials = ["<|im_end|>", "<|endoftext|>"]
|
|
assert detect_think_prefill(QWEN_PROMPT + "<think>\n", specials) == "<think>\n"
|
|
|
|
|
|
def test_guard_default_and_empty_keep_emitting():
|
|
assert detect_think_prefill(QWEN_PROMPT + "<think>\n", None) == "<think>\n"
|
|
assert detect_think_prefill(QWEN_PROMPT + "<think>\n", []) == "<think>\n"
|