mirror of https://github.com/razor-ai/soup.git
277 lines
9.6 KiB
Python
277 lines
9.6 KiB
Python
"""Tests for chat-template hardening (v0.36.0 Part C).
|
|
|
|
Replaces sft.py's silent ``f"{role}: {content}"`` fallback (which produced
|
|
garbage training data on tokenizers without ``chat_template``) with a hard
|
|
error and an explicit ``DataConfig.chat_template`` field for overrides.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Schema field
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestSchemaField:
|
|
def test_chat_template_default_none(self):
|
|
from soup_cli.config.schema import DataConfig
|
|
|
|
cfg = DataConfig(train="data.jsonl")
|
|
assert cfg.chat_template is None
|
|
|
|
def test_chat_template_accepts_registered_name(self):
|
|
from soup_cli.config.schema import DataConfig
|
|
|
|
cfg = DataConfig(train="data.jsonl", chat_template="chatml")
|
|
assert cfg.chat_template == "chatml"
|
|
|
|
def test_chat_template_accepts_jinja_string(self):
|
|
from soup_cli.config.schema import DataConfig
|
|
|
|
jinja = "{% for m in messages %}{{ m.role }}: {{ m.content }}{% endfor %}"
|
|
cfg = DataConfig(train="data.jsonl", chat_template=jinja)
|
|
assert cfg.chat_template == jinja
|
|
|
|
def test_chat_template_rejects_null_byte(self):
|
|
from soup_cli.config.schema import DataConfig
|
|
|
|
with pytest.raises(ValueError):
|
|
DataConfig(train="data.jsonl", chat_template="bad\x00template")
|
|
|
|
def test_chat_template_rejects_oversize(self):
|
|
from soup_cli.config.schema import DataConfig
|
|
|
|
# Cap at 64KB to prevent template-injection DoS payloads.
|
|
with pytest.raises(ValueError):
|
|
DataConfig(train="data.jsonl", chat_template="x" * 100_000)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Registry
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestRegistry:
|
|
def test_lists_known_templates(self):
|
|
from soup_cli.data.chat_templates import list_template_names
|
|
|
|
names = list_template_names()
|
|
# At minimum: the 7 declared in v0.36.0 Part C.
|
|
for required in (
|
|
"chatml", "llama3", "qwen2.5", "gemma3", "phi4", "deepseek-r1", "mistral",
|
|
):
|
|
assert required in names
|
|
|
|
def test_chatml_is_jinja_string(self):
|
|
from soup_cli.data.chat_templates import get_template
|
|
|
|
tmpl = get_template("chatml")
|
|
assert isinstance(tmpl, str)
|
|
assert "{%" in tmpl
|
|
# ChatML signature markers.
|
|
assert "<|im_start|>" in tmpl
|
|
assert "<|im_end|>" in tmpl
|
|
|
|
def test_unknown_name_raises(self):
|
|
from soup_cli.data.chat_templates import get_template
|
|
|
|
with pytest.raises(KeyError, match="not registered"):
|
|
get_template("not-a-real-template")
|
|
|
|
def test_resolve_returns_jinja_for_known_name(self):
|
|
from soup_cli.data.chat_templates import resolve_chat_template
|
|
|
|
out = resolve_chat_template("chatml")
|
|
assert "<|im_start|>" in out
|
|
|
|
def test_resolve_returns_passthrough_for_jinja(self):
|
|
from soup_cli.data.chat_templates import resolve_chat_template
|
|
|
|
jinja = "{% for m in messages %}<x>{{ m.content }}</x>{% endfor %}"
|
|
out = resolve_chat_template(jinja)
|
|
assert out == jinja
|
|
|
|
def test_resolve_none_returns_none(self):
|
|
from soup_cli.data.chat_templates import resolve_chat_template
|
|
|
|
assert resolve_chat_template(None) is None
|
|
|
|
def test_resolve_empty_returns_none(self):
|
|
from soup_cli.data.chat_templates import resolve_chat_template
|
|
|
|
assert resolve_chat_template("") is None
|
|
|
|
def test_resolve_unknown_name_raises(self):
|
|
"""Public surface: bad name through resolve_chat_template."""
|
|
from soup_cli.data.chat_templates import resolve_chat_template
|
|
|
|
with pytest.raises(KeyError, match="not registered"):
|
|
resolve_chat_template("not-a-real-template")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Schema rejects Jinja directives that touch the filesystem
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestJinjaDirectiveBlocking:
|
|
@pytest.mark.parametrize(
|
|
"bad",
|
|
[
|
|
"{% include 'config.yaml' %}",
|
|
"{%- include 'config.yaml' %}",
|
|
"{% import 'os' as os %}",
|
|
"{% from 'os' import system %}",
|
|
"{% macro evil() %}{% endmacro %}",
|
|
"{% extends 'base.j2' %}",
|
|
],
|
|
)
|
|
def test_schema_rejects_filesystem_directives(self, bad):
|
|
from soup_cli.config.schema import DataConfig
|
|
|
|
with pytest.raises(ValueError, match="directive"):
|
|
DataConfig(train="data.jsonl", chat_template=bad)
|
|
|
|
def test_schema_accepts_for_loop(self):
|
|
"""Standard control flow must still work."""
|
|
from soup_cli.config.schema import DataConfig
|
|
|
|
ok = "{% for m in messages %}{{ m.content }}{% endfor %}"
|
|
cfg = DataConfig(train="data.jsonl", chat_template=ok)
|
|
assert cfg.chat_template == ok
|
|
|
|
def test_schema_accepts_if_else(self):
|
|
from soup_cli.config.schema import DataConfig
|
|
|
|
ok = "{% if true %}x{% else %}y{% endif %}"
|
|
cfg = DataConfig(train="data.jsonl", chat_template=ok)
|
|
assert cfg.chat_template == ok
|
|
|
|
def test_schema_empty_string_normalised_to_none(self):
|
|
from soup_cli.config.schema import DataConfig
|
|
|
|
cfg = DataConfig(train="data.jsonl", chat_template="")
|
|
assert cfg.chat_template is None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# apply_chat_template_override
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestApplyOverride:
|
|
def test_sets_tokenizer_chat_template(self):
|
|
from soup_cli.data.chat_templates import apply_chat_template_override
|
|
|
|
class _T:
|
|
chat_template = None
|
|
|
|
tok = _T()
|
|
apply_chat_template_override(tok, "chatml")
|
|
assert tok.chat_template is not None
|
|
assert "<|im_start|>" in tok.chat_template
|
|
|
|
def test_none_leaves_tokenizer_alone(self):
|
|
from soup_cli.data.chat_templates import apply_chat_template_override
|
|
|
|
class _T:
|
|
chat_template = "existing"
|
|
|
|
tok = _T()
|
|
apply_chat_template_override(tok, None)
|
|
assert tok.chat_template == "existing"
|
|
|
|
def test_empty_leaves_tokenizer_alone(self):
|
|
from soup_cli.data.chat_templates import apply_chat_template_override
|
|
|
|
class _T:
|
|
chat_template = "existing"
|
|
|
|
tok = _T()
|
|
apply_chat_template_override(tok, "")
|
|
assert tok.chat_template == "existing"
|
|
|
|
def test_override_emits_save_pretrained_warning(self):
|
|
"""v0.36.0 review fix: warn when overriding so users know push will
|
|
persist the new template into tokenizer_config.json."""
|
|
from io import StringIO
|
|
|
|
from rich.console import Console
|
|
|
|
from soup_cli.data.chat_templates import apply_chat_template_override
|
|
|
|
class _T:
|
|
chat_template = None
|
|
|
|
buf = StringIO()
|
|
console = Console(file=buf, force_terminal=False)
|
|
applied = apply_chat_template_override(_T(), "chatml", console=console)
|
|
assert applied is True
|
|
assert "save_pretrained" in buf.getvalue() or "soup push" in buf.getvalue()
|
|
|
|
def test_override_returns_false_when_noop(self):
|
|
from soup_cli.data.chat_templates import apply_chat_template_override
|
|
|
|
class _T:
|
|
chat_template = "existing"
|
|
|
|
applied = apply_chat_template_override(_T(), None)
|
|
assert applied is False
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Hard error in sft_format when no chat_template AND no override
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestHardError:
|
|
def test_no_template_no_override_raises(self):
|
|
"""The legacy `f"{role}: {content}"` silent fallback is now an error."""
|
|
from soup_cli.config.schema import DataConfig
|
|
from soup_cli.data.sft_format import build_format_row
|
|
|
|
class _NoTemplate:
|
|
chat_template = None
|
|
|
|
def apply_chat_template(self, *args, **kwargs):
|
|
raise AssertionError("must not be reached")
|
|
|
|
cfg = DataConfig(
|
|
train="data.jsonl",
|
|
train_on_responses_only=False, # legacy path
|
|
train_on_messages_with_train_field=False,
|
|
chat_template=None,
|
|
)
|
|
fn = build_format_row(_NoTemplate(), cfg, console=None)
|
|
with pytest.raises(ValueError, match="chat_template"):
|
|
fn({"messages": [{"role": "user", "content": "hi"}]})
|
|
|
|
def test_override_applies_to_legacy_path(self):
|
|
"""When user passes chat_template override, build_format_row works."""
|
|
from soup_cli.config.schema import DataConfig
|
|
from soup_cli.data.sft_format import build_format_row
|
|
|
|
class _T:
|
|
chat_template = None
|
|
applied = []
|
|
|
|
def apply_chat_template(
|
|
self, messages, tokenize=False, add_generation_prompt=False, **kwargs
|
|
):
|
|
self.applied.append(messages)
|
|
return "RENDERED"
|
|
|
|
tok = _T()
|
|
cfg = DataConfig(
|
|
train="data.jsonl",
|
|
train_on_responses_only=False,
|
|
train_on_messages_with_train_field=False,
|
|
chat_template="chatml",
|
|
)
|
|
fn = build_format_row(tok, cfg, console=None)
|
|
out = fn({"messages": [{"role": "user", "content": "hi"}]})
|
|
assert out["text"] == "RENDERED"
|
|
assert tok.chat_template is not None # override was applied
|