diff --git a/tests/python/test_change_system_message.py b/tests/python/test_change_system_message.py index 7edd873a8d..23430b03c5 100644 --- a/tests/python/test_change_system_message.py +++ b/tests/python/test_change_system_message.py @@ -13,7 +13,8 @@ def _load_change_system_message(): funcs = [ node for node in tree.body - if isinstance(node, ast.FunctionDef) and node.name == "_change_system_message" + if isinstance(node, ast.FunctionDef) + and node.name in ("_change_system_message", "_escape_jinja_literal") ] namespace = { "re": re, @@ -62,10 +63,20 @@ def test_custom_template_without_placeholder_unchanged(): def test_predefined_template_uses_default_then_override(): - # Predefined templates with a default are unaffected. fn = _load_change_system_message() t1, u1 = fn("System: {system_message}", "unsloth", None) assert t1 == "System: You are a helpful assistant to the user" t2, u2 = fn("System: {system_message}", "unsloth", "Custom override") assert t2 == "System: Custom override" assert u2 == "Custom override" + + +def test_predefined_template_escapes_but_custom_does_not(): + # Predefined templates hold {system_message} inside a Jinja literal, so it is escaped; + # a custom template's placeholder may be raw text, so it stays verbatim. + fn = _load_change_system_message() + msg = """it's a \\test "x".""" + predefined, used = fn("System: {system_message}", "unsloth", msg) + assert predefined == """System: it\\'s a \\\\test \\"x\\".""" + assert used == msg # returned message is the raw one, not the escaped source + assert fn("System: {system_message}", CUSTOM, msg)[0] == f"System: {msg}" diff --git a/tests/python/test_get_chat_template_escaping.py b/tests/python/test_get_chat_template_escaping.py new file mode 100644 index 0000000000..85239390f0 --- /dev/null +++ b/tests/python/test_get_chat_template_escaping.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. +"""Escaping tests for the predefined chat templates' {system_message} placeholder. + +Every predefined template holds {system_message} inside a Jinja string literal, so a +system message carrying a quote or a backslash has to be escaped on the way in. Without +it a quote closes the literal (TemplateSyntaxError) and a backslash is read as an escape, +which silently rewrites the text (\\boxed -> \\x08oxed). Rendering, not just compiling, +is asserted here: the corrupting cases compile fine. +""" + +import jinja2 +import pytest +from jinja2.sandbox import ImmutableSandboxedEnvironment + +from unsloth.chat_templates import ( + CHAT_TEMPLATES, + DEFAULT_SYSTEM_MESSAGE, + _change_system_message, + _escape_jinja_literal, +) + + +# Discovered, not hardcoded, so a new predefined template is covered on arrival. +TEMPLATES_WITH_SYSTEM_MESSAGE = sorted( + name + for name, entry in CHAT_TEMPLATES.items() + if isinstance(entry[0], str) and "{system_message}" in entry[0] +) + +MESSAGES = [ + ("apostrophe", "Answer the user's question."), + ("double", 'He said "hi" today.'), + ("both_quotes", """mix ' and " here"""), + ("latex", r"Put it in \boxed{\frac{1}{2}}."), + ("windows", r"C:\Users\me"), + ("trailing", "ends with a backslash \\"), + ("crlf", "line one\r\nline two"), + ("group_ref", r"use \1 for the group"), + ("jinja", "{{ 7*6 }} and {% raw %}x{% endraw %}"), +] + + +def _render(template): + # No system turn in `messages`, so the template falls back to the baked-in + # {system_message} literal, the path under test. + environment = ImmutableSandboxedEnvironment(trim_blocks = True, lstrip_blocks = True) + return environment.from_string(template).render( + messages = [{"role": "user", "content": "Hi"}], + bos_token = "", + eos_token = "", + add_generation_prompt = False, + ) + + +def test_templates_with_system_message_were_found(): + # An empty discovery list would make every case below vacuous. + assert len(TEMPLATES_WITH_SYSTEM_MESSAGE) >= 10 + + +@pytest.mark.parametrize("name", TEMPLATES_WITH_SYSTEM_MESSAGE) +@pytest.mark.parametrize("label, system_message", MESSAGES, ids = [m[0] for m in MESSAGES]) +def test_system_message_survives_the_jinja_literal(name, label, system_message): + template, used = _change_system_message(CHAT_TEMPLATES[name][0], name, system_message) + assert used == system_message, "the returned message must be the raw one" + assert system_message in _render(template) + + +@pytest.mark.parametrize("name", TEMPLATES_WITH_SYSTEM_MESSAGE) +def test_default_system_message_renders_verbatim(name): + # Defaults are no longer hand-escaped in the source; escaping them twice would + # surface a literal backslash to the user. + default = DEFAULT_SYSTEM_MESSAGE[name] + template, _ = _change_system_message(CHAT_TEMPLATES[name][0], name, None) + assert default in _render(template) + + +@pytest.mark.parametrize("name", ["vicuna", "vicuna_old", "vicuna old"]) +def test_vicuna_default_has_a_plain_apostrophe(name): + assert "\\" not in DEFAULT_SYSTEM_MESSAGE[name] + assert "'s questions." in _render( + _change_system_message(CHAT_TEMPLATES[name][0], name, None)[0] + ) + + +@pytest.mark.parametrize("quote", ["'", '"']) +@pytest.mark.parametrize("label, text", MESSAGES, ids = [m[0] for m in MESSAGES]) +def test_escape_round_trips_in_either_quote_style(quote, label, text): + # get_chat_template also splices ShareGPT `mapping` values into literals, and + # llama-3.1 uses "..." where the rest use '...', so one escaper must cover both. + template = "{{ " + quote + _escape_jinja_literal(text) + quote + " }}" + assert jinja2.Environment().from_string(template).render() == text diff --git a/unsloth/chat_templates.py b/unsloth/chat_templates.py index c4c559a413..9d3a9f0be9 100644 --- a/unsloth/chat_templates.py +++ b/unsloth/chat_templates.py @@ -218,7 +218,7 @@ vicuna_ollama = _ollama_template("vicuna") vicuna_eos_token = "eos_token" CHAT_TEMPLATES["vicuna"] = (vicuna_template, vicuna_eos_token, False, vicuna_ollama,) -DEFAULT_SYSTEM_MESSAGE["vicuna"] = "A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user\\'s questions." +DEFAULT_SYSTEM_MESSAGE["vicuna"] = "A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions." # =========================================== Vicuna Old # https://github.com/lm-sys/FastChat/blob/main/docs/vicuna_weights_version.md#prompt-template @@ -248,7 +248,7 @@ vicuna_old_ollama = _ollama_template("vicuna_old") vicuna_old_eos_token = "eos_token" CHAT_TEMPLATES["vicuna_old"] = (vicuna_old_template, vicuna_old_eos_token, False, vicuna_old_ollama,) -DEFAULT_SYSTEM_MESSAGE["vicuna_old"] = "A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human\\'s questions." +DEFAULT_SYSTEM_MESSAGE["vicuna_old"] = "A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions." CHAT_TEMPLATES["vicuna old"] = CHAT_TEMPLATES["vicuna_old"] DEFAULT_SYSTEM_MESSAGE["vicuna old"] = DEFAULT_SYSTEM_MESSAGE["vicuna_old"] @@ -1810,6 +1810,13 @@ yi_chat_template_eos_token = "<|endoftext|>" CHAT_TEMPLATES["yi-chat"] = (yi_chat_template, yi_chat_template_eos_token, False, yi_chat_ollama) DEFAULT_SYSTEM_MESSAGE["yi-chat"] = None +# Caller text is spliced into Jinja '...' / "..." literals: a bare quote closes the +# literal, a backslash reads as an escape, and \r becomes \n before unescaping. Backslash +# first, else it re-escapes the rest; escaping both quotes is safe in either style. +def _escape_jinja_literal(text): + return text.replace("\\", "\\\\").replace("'", "\\'").replace('"', '\\"').replace("\r", "\\r") + + def _change_system_message(template: str, type_chat_template: str, system_message: str = None): # For predefined templates, check if default system message exists default_system_message = DEFAULT_SYSTEM_MESSAGE.get(f"{type_chat_template}", None) @@ -1833,7 +1840,9 @@ def _change_system_message(template: str, type_chat_template: str, system_messag # For predefined templates with default system message message_to_use = system_message if system_message is not None else default_system_message - new_template = template.replace("{system_message}", message_to_use) + # Predefined templates hold {system_message} inside a Jinja literal, so escape it. + # Custom ones (above) are filled verbatim: their placeholder may sit in raw text. + new_template = template.replace("{system_message}", _escape_jinja_literal(message_to_use)) return new_template, message_to_use @@ -2053,11 +2062,12 @@ def get_chat_template( chat_template = "{{ bos_token }}" + chat_template # For ShareGPT role -> from and content -> value + # The spliced values land inside Jinja literals, so escape them. new_chat_template = chat_template\ - .replace("'role'", "'" + mapping["role"] + "'")\ - .replace("'content'", "'" + mapping["content"] + "'")\ - .replace("'user'", "'" + mapping["user"] + "'")\ - .replace("'assistant'", "'" + mapping["assistant"] + "'") + .replace("'role'", "'" + _escape_jinja_literal(mapping["role"]) + "'")\ + .replace("'content'", "'" + _escape_jinja_literal(mapping["content"]) + "'")\ + .replace("'user'", "'" + _escape_jinja_literal(mapping["user"]) + "'")\ + .replace("'assistant'", "'" + _escape_jinja_literal(mapping["assistant"]) + "'") if use_zoo_tokenizer_patch: # Unsloth MLX avoids the model-utils tokenizer wrapper because that @@ -2617,15 +2627,9 @@ extra_eos_tokens = None, part + '\n\n' + ollama_eos # HF Jinja Chat template - # Caller text is spliced into Jinja '...' literals, where a bare quote closes - # the literal and a backslash is read as an escape, and \r is normalised to \n - # before unescaping. Backslash first, else it re-escapes the ones added after. - def escape_jinja_literal(text): - return text.replace("\\", "\\\\").replace("'", "\\'").replace("\r", "\\r") - def process(part, which, content = "message['content']"): # Escape the literal pieces only; the placeholder becomes Jinja syntax below. - part = which.join(escape_jinja_literal(piece) for piece in part.split(which)) + part = which.join(_escape_jinja_literal(piece) for piece in part.split(which)) if part.endswith(which): part = "'" + part[:part.find(which)] + f"' + {content}" elif part.startswith(which): @@ -2648,7 +2652,7 @@ extra_eos_tokens = None, "{% endif %}"\ "{% endfor %}"\ "{% if add_generation_prompt %}"\ - "{{ '" + escape_jinja_literal(output_part[:output_part.find("{OUTPUT}")]) + "' }}"\ + "{{ '" + _escape_jinja_literal(output_part[:output_part.find("{OUTPUT}")]) + "' }}"\ "{% endif %}" # Now add system prompt to jinja @@ -2668,7 +2672,7 @@ extra_eos_tokens = None, "{{ " + partial_system + " }}"\ "{% set loop_messages = messages[1:] %}" if default_system_message is not None: - full_system = escape_jinja_literal(system_part.replace("{SYSTEM}", default_system_message)) + full_system = _escape_jinja_literal(system_part.replace("{SYSTEM}", default_system_message)) if "{SYSTEM}" in system_part: modelfile += '\nSYSTEM "' + default_system_message + '"' partial_system += "{% else %}"\