diff --git a/studio/backend/tests/test_response_template_markers.py b/studio/backend/tests/test_response_template_markers.py new file mode 100644 index 0000000000..8c813e62f2 --- /dev/null +++ b/studio/backend/tests/test_response_template_markers.py @@ -0,0 +1,216 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""TEMPLATE_TO_RESPONSES_MAPPER markers must match what the templates render. + +The manual instruction/response markers are the fallback for +train_on_completions when auto-detection is unavailable, so a marker that +never matches the rendered chat template masks every assistant token and the +run dies on the all-labels-masked safety net. Six template families shipped +such markers: + + mistral - "[INST] " / " [/INST]": the surrounding spaces fold into + the neighbouring tokens ("[INST]" is a single special + token in Mistral v0.3), so the padded strings never match. + llama - same space folding, plus llama-2 tokenizes [INST] after + as bare "[" on transformers 5.x while the standalone + encoding gives "▁[", so the marker must anchor on . + starling - trailing space after "GPT4 Correct Assistant:" folds + into the next content token ("▁Hello"). + glm - "[gMASK]" renders once at text start, never before + later user turns; "" is generation scaffolding + that non-final turns render as a lone "". + qwen3-thinking - "" is stripped from non-final assistant turns + (Qwen3-Thinking-2507) or never rendered (QwQ). + zephyr - role tags are plain text, and SentencePiece tokenizes + "<|assistant|>" differently at text start than after + "\\n" mid-conversation; the markers need the leading + newline anchor to tokenize like a real turn boundary. + +Literal assertions run everywhere; the token-level masking checks need the +representative tokenizers plus unsloth_zoo and skip when either is +unavailable (offline CI). +""" + +from __future__ import annotations + +import importlib.util +import sys +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# model_mappings is dependency-free: load it directly so these tests run +# without the studio venv / package import side effects. +_MM_PATH = Path(_BACKEND_DIR) / "utils" / "datasets" / "model_mappings.py" +_mm_spec = importlib.util.spec_from_file_location("_marker_test_mm", _MM_PATH) +model_mappings = importlib.util.module_from_spec(_mm_spec) +_mm_spec.loader.exec_module(model_mappings) + +T2R = model_mappings.TEMPLATE_TO_RESPONSES_MAPPER + + +# ── Fixed entries: markers derived from what each representative tokenizer +# actually renders (see PR for the token-level derivation). ── +EXPECTED_FIXED = { + "mistral": {"instruction": "[INST]", "response": "[/INST]"}, + "llama": {"instruction": "[INST]", "response": "[/INST]"}, + "starling": {"instruction": "GPT4 Correct User:", "response": "GPT4 Correct Assistant:"}, + "glm": {"instruction": "<|user|>", "response": "<|assistant|>"}, + "qwen3-thinking": {"instruction": "<|im_start|>user\n", "response": "<|im_start|>assistant\n"}, + "zephyr": {"instruction": "\n<|user|>\n", "response": "\n<|assistant|>\n"}, +} + +# Spot-pin some known-good entries so a refactor cannot silently change them. +EXPECTED_UNCHANGED = { + "qwen3": {"instruction": "<|im_start|>user\n", "response": "<|im_start|>assistant\n"}, + "llama-3.1": { + "instruction": "<|start_header_id|>user<|end_header_id|>\n\n", + "response": "<|start_header_id|>assistant<|end_header_id|>\n\n", + }, + "phi-4": { + "instruction": "<|im_start|>user<|im_sep|>", + "response": "<|im_start|>assistant<|im_sep|>", + }, + "gemma-3": {"instruction": "user\n", "response": "model\n"}, + "gpt-oss": { + "instruction": "<|start|>user<|message|>", + "response": "<|start|>assistant<|channel|>final<|message|>", + }, +} + + +@pytest.mark.parametrize("template", sorted(EXPECTED_FIXED)) +def test_fixed_marker_literals(template): + assert T2R[template] == EXPECTED_FIXED[template] + + +@pytest.mark.parametrize("template", sorted(EXPECTED_UNCHANGED)) +def test_unchanged_marker_literals(template): + assert T2R[template] == EXPECTED_UNCHANGED[template] + + +def test_no_marker_is_empty_or_whitespace(): + for template, parts in T2R.items(): + assert parts["instruction"].strip(), template + assert parts["response"].strip(), template + + +# ── Token-level checks: markers must select exactly the assistant turns on a +# rendered two-turn fixture, and the final EOS label must never be -100. ── + +REPRESENTATIVES = { + "mistral": ["unsloth/mistral-7b-instruct-v0.3"], + "llama": ["unsloth/llama-2-7b-chat"], + "starling": ["unsloth/Starling-LM-7B-beta"], + "glm": ["unsloth/GLM-4.7-Flash"], + "qwen3-thinking": ["unsloth/Qwen3-4B-Thinking-2507", "Qwen/QwQ-32B"], + "zephyr": ["unsloth/zephyr-sft"], +} + +FIXTURE = [ + {"role": "user", "content": "zebra alpha question one?"}, + {"role": "assistant", "content": "grape reply number one."}, + {"role": "user", "content": "zebra beta question two?"}, + {"role": "assistant", "content": "grape reply number two."}, +] + + +def _load_tokenizer(repo): + try: + from transformers import AutoTokenizer + except Exception as e: # pragma: no cover + pytest.skip(f"transformers unavailable: {e}") + try: + return AutoTokenizer.from_pretrained(repo) + except OSError as e: + pytest.skip(f"tokenizer {repo} unavailable (offline?): {e}") + except Exception: + # Tokenizer class newer than this transformers (e.g. GLM-4.7's + # TokenizersBackend): build directly from tokenizer.json. + try: + import json as _json + from huggingface_hub import hf_hub_download + from transformers import PreTrainedTokenizerFast + + with open(hf_hub_download(repo, "tokenizer_config.json"), encoding = "utf-8") as f: + cfg = _json.load(f) + tok_file = hf_hub_download(repo, "tokenizer.json") + + def _tokval(v): + return v["content"] if isinstance(v, dict) else v + + return PreTrainedTokenizerFast( + tokenizer_file = tok_file, + chat_template = cfg.get("chat_template"), + **{ + k: _tokval(cfg[k]) + for k in ("bos_token", "eos_token", "pad_token", "unk_token") + if cfg.get(k) is not None + }, + ) + except Exception as e: + pytest.skip(f"tokenizer {repo} unavailable (offline?): {e}") + + +def _train_on_responses_only(): + try: + from unsloth_zoo.dataset_utils import train_on_responses_only + except Exception as e: + pytest.skip(f"unsloth_zoo unavailable: {e}") + return train_on_responses_only + + +@pytest.mark.parametrize( + "template,repo", + [(t, r) for t, repos in sorted(REPRESENTATIVES.items()) for r in repos], +) +def test_fixed_markers_token_level(template, repo): + tor = _train_on_responses_only() + tok = _load_tokenizer(repo) + parts = T2R[template] + + msgs = [{"role": "system", "content": "You are a terse assistant."}] + FIXTURE + try: + ids = tok.apply_chat_template(msgs, tokenize = True, add_generation_prompt = False) + if hasattr(ids, "keys"): + ids = ids["input_ids"] # transformers 5.x returns a BatchEncoding + except Exception: + ids = tok.apply_chat_template(FIXTURE, tokenize = True, add_generation_prompt = False) + if hasattr(ids, "keys"): + ids = ids["input_ids"] + + fn = tor( + None, + instruction_part = parts["instruction"], + response_part = parts["response"], + tokenizer = tok, + return_function = True, + ) + labels = fn({"input_ids": [list(ids)]})["labels"][0] + + n = len(ids) + trained = tok.decode([ids[i] for i in range(n) if labels[i] != -100]) + masked = tok.decode([ids[i] for i in range(n) if labels[i] == -100]) + + # User and system content fully masked + assert "question one" not in trained and "question one" in masked + assert "question two" not in trained and "question two" in masked + assert "terse assistant" not in trained + # EVERY assistant turn trained, not just the last + assert "reply number one" in trained + assert "reply number two" in trained + # The final EOS (last non-whitespace token) must never be -100, or the + # fine-tuned model never learns to stop generating. + i = n - 1 + while i > 0 and tok.decode([ids[i]]).strip() == "": + i -= 1 + assert labels[i] != -100, f"final token {tok.convert_ids_to_tokens(int(ids[i]))!r} is masked" + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v"])) diff --git a/studio/backend/utils/datasets/model_mappings.py b/studio/backend/utils/datasets/model_mappings.py index 9d2c983aed..65ba4b4688 100644 --- a/studio/backend/utils/datasets/model_mappings.py +++ b/studio/backend/utils/datasets/model_mappings.py @@ -485,9 +485,11 @@ TEMPLATE_TO_RESPONSES_MAPPER = { "instruction": "<|im_start|>user\n", "response": "<|im_start|>assistant\n", }, + # No "" suffix: Qwen3-Thinking-2507 strips it from non-final turns + # and QwQ renders none, so a marker holding it masks those responses. "qwen3-thinking": { "instruction": "<|im_start|>user\n", - "response": "<|im_start|>assistant\n", + "response": "<|im_start|>assistant\n", }, "qwen3": { "instruction": "<|im_start|>user\n", @@ -525,29 +527,39 @@ TEMPLATE_TO_RESPONSES_MAPPER = { "instruction": "<|im_start|>user<|im_sep|>", "response": "<|im_start|>assistant<|im_sep|>", }, + # No surrounding spaces: in Mistral v0.3 they fold into neighbouring text + # tokens ("[INST]"/"[/INST]" are single special tokens), so padded strings + # never match and everything masks. Same for Llama-2's SentencePiece. "mistral": { - "instruction": "[INST] ", - "response": " [/INST]", + "instruction": "[INST]", + "response": "[/INST]", }, "llama": { - "instruction": "[INST] ", - "response": " [/INST]", + # -anchored: llama-2 tokenizes [INST] after as bare "[" on + # transformers 5.x (standalone gives space-prefixed "▁["), so an + # unanchored marker misses every turn boundary there. + "instruction": "[INST]", + "response": "[/INST]", }, "chatml": { "instruction": "<|im_start|>user\n", "response": "<|im_start|>assistant\n", }, + # Leading newline required: Zephyr's role tags are plain text, and + # SentencePiece tokenizes "<|assistant|>" differently at text start than + # after "\n". Without the "\n" anchor the markers never match real + # turns, so every assistant token masks. "zephyr": { - "instruction": "<|user|>\n", - "response": "<|assistant|>\n", + "instruction": "\n<|user|>\n", + "response": "\n<|assistant|>\n", }, "unsloth": { - "instruction": ">>> User: ", - "response": ">>> Assistant: ", + "instruction": ">>> User:", + "response": ">>> Assistant:", }, "vicuna": { - "instruction": "USER: ", - "response": "ASSISTANT: ", + "instruction": "USER:", + "response": "ASSISTANT:", }, "alpaca": { "instruction": "### Instruction:\n", @@ -573,16 +585,21 @@ TEMPLATE_TO_RESPONSES_MAPPER = { "instruction": "<|im_start|>user\n", "response": "<|im_start|>assistant\n", }, + # No trailing space: SentencePiece folds it into the next content token + # ("▁Hello"), so the padded marker never matches and masks everything. "starling": { - "instruction": "GPT4 Correct User: ", - "response": "GPT4 Correct Assistant: ", + "instruction": "GPT4 Correct User:", + "response": "GPT4 Correct Assistant:", }, "yi-chat": { "instruction": "<|im_start|>user\n", "response": "<|im_start|>assistant\n", }, + # "[gMASK]" appears once at text start, so a marker holding it matches + # no later user turn; "" is scaffolding GLM-4.x renders as a lone + # "" on non-final turns, so "<|assistant|>" never matches. "glm": { - "instruction": "[gMASK]<|user|>", - "response": "<|assistant|>", + "instruction": "<|user|>", + "response": "<|assistant|>", }, }