mirror of
https://github.com/unslothai/unsloth.git
synced 2026-08-21 06:43:53 +00:00
44f113cf0d
* Escape the system message spliced into predefined chat templates
get_chat_template(..., system_message = ...) substitutes the message into a
{system_message} placeholder that sits inside a Jinja string literal in all 15
predefined templates that carry one, so a quote closes the literal and a
backslash is read as an escape:
vicuna "Answer the user's question." -> TemplateSyntaxError
vicuna r"Put it in \boxed{}." -> renders '\x08oxed{}'
vicuna r"C:\Users\me" -> TemplateSyntaxError
Reuse the escaper PR #7731 added for construct_chat_template, promoted to a
module-level _escape_jinja_literal and extended to escape double quotes so the
one helper covers llama-3.1's "..." literal as well as the '...' the rest use.
Apply it to the predefined branch of _change_system_message and to the ShareGPT
mapping values, and drop the hand-escaping from the two vicuna defaults, which
would otherwise be escaped twice.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Tighten the escaping comments
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
92 lines
3.7 KiB
Python
92 lines
3.7 KiB
Python
# 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 = "<s>",
|
|
eos_token = "</s>",
|
|
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
|