From b307823b1daf7013632340bceac5b2f70dbc04a8 Mon Sep 17 00:00:00 2001 From: Andrew Chen <48723787+chuenchen309@users.noreply.github.com> Date: Sun, 19 Jul 2026 21:33:48 +0800 Subject: [PATCH] fix(chat_templates): bind loop_messages when default_system_message is None (#7199) * fix(chat_templates): bind loop_messages when default_system_message is None construct_chat_template(default_system_message=None) built a system part that binds loop_messages only inside the `{% if messages[0]['role'] == 'system' %}` arm. The `Fix missing loop_messages` step right below then found no unconditional `{% set loop_messages = messages %}`, concluded loop_messages was missing, and rewrote `{% for message in loop_messages %}` back to `{% for message in messages %}` -- undoing the `messages[1:]` skip. A caller-supplied system message therefore reached the loop and tripped raise_exception: Only user and assistant roles are supported! Add the `{% else %}` arm so loop_messages is always bound, mirroring the default_system_message is not None branch minus the default text. That also stops the rewrite from firing, since the unconditional binding is now present. Renders before / after, same template, same inputs: default_system_message input before after None system msg raise_exception 'Be terse.\n### User: Hi\n' None no system '### User: Hi\n' unchanged 'You are helpful.' system msg 'Be terse.\n### User: Hi\n' unchanged 'You are helpful.' no system 'You are helpful.\n...' unchanged The rewrite still fires for templates with no {SYSTEM} part, which is what it was there for -- verified unchanged. Co-Authored-By: Claude Opus 4.8 * Scope loop_messages binding to {SYSTEM} templates for PR #7199 The None branch now only adds the else arm when system_part contains {SYSTEM}, so a static prefix with no {SYSTEM} placeholder keeps raising on a caller system message instead of silently dropping it. Strengthen the tests: assert the default does not leak when a caller system message is present, and add a regression test for the static prefix case. --------- Co-authored-by: Claude Opus 4.8 Co-authored-by: danielhanchen --- ...test_construct_chat_template_validation.py | 90 +++++++++++++++++++ unsloth/chat_templates.py | 6 ++ 2 files changed, 96 insertions(+) diff --git a/tests/python/test_construct_chat_template_validation.py b/tests/python/test_construct_chat_template_validation.py index 66d3d80920..53d281d435 100644 --- a/tests/python/test_construct_chat_template_validation.py +++ b/tests/python/test_construct_chat_template_validation.py @@ -104,3 +104,93 @@ def test_chat_template_does_not_leak_sentinel_when_section_starts_with_it(chat_t ) assert "{INPUT}" not in jinja_template assert "{OUTPUT}" not in jinja_template + + +_SYSTEM_CHAT_TEMPLATE = ( + "{SYSTEM}\n" + "### User: {INPUT}\n### Assistant: {OUTPUT}" + "### User: {INPUT}\n### Assistant: {OUTPUT}" +) + + +def _render(jinja_template, messages): + from jinja2.sandbox import ImmutableSandboxedEnvironment + + env = ImmutableSandboxedEnvironment() + env.globals["raise_exception"] = lambda message: (_ for _ in ()).throw(RuntimeError(message)) + return env.from_string(jinja_template).render( + messages = messages, + bos_token = "", + eos_token = "", + add_generation_prompt = False, + ) + + +@pytest.mark.parametrize("default_system_message", [None, "You are helpful."]) +def test_system_message_is_consumed_by_the_system_part(default_system_message): + """A caller-supplied system message must be rendered by the system part and + skipped by the message loop, whatever `default_system_message` is. + + With `default_system_message = None` the generated template used to bind + `loop_messages` only inside the `{% if %}` arm. The `Fix missing + loop_messages` step then saw no unconditional binding, rewrote the loop back + to `messages`, and the system message reached the loop and tripped + `raise_exception`. + """ + _, jinja_template, _, _ = construct_chat_template( + tokenizer = _SuccessFakeTokenizer(), + chat_template = _SYSTEM_CHAT_TEMPLATE, + default_system_message = default_system_message, + extra_eos_tokens = [""], + ) + rendered = _render( + jinja_template, + [ + {"role": "system", "content": "Be terse."}, + {"role": "user", "content": "Hi"}, + ], + ) + assert rendered.count("Be terse.") == 1 + assert rendered.count("Hi") == 1 + # A caller system message overrides the default; the default must not leak in. + if default_system_message is not None: + assert default_system_message not in rendered + + +def test_absent_system_message_still_renders_without_default(): + """`default_system_message = None` with no system message in the input must + keep working -- the `{% else %}` arm has to bind `loop_messages = messages`.""" + _, jinja_template, _, _ = construct_chat_template( + tokenizer = _SuccessFakeTokenizer(), + chat_template = _SYSTEM_CHAT_TEMPLATE, + default_system_message = None, + extra_eos_tokens = [""], + ) + rendered = _render(jinja_template, [{"role": "user", "content": "Hi"}]) + assert "Hi" in rendered + + +_NO_SYSTEM_CHAT_TEMPLATE = ( + "PREAMBLE\n" + "### User: {INPUT}\n### Assistant: {OUTPUT}" + "### User: {INPUT}\n### Assistant: {OUTPUT}" +) + + +def test_static_prefix_without_system_still_rejects_system_message(): + """A template with a static prefix but no {SYSTEM} placeholder cannot render a + caller system message, so it must still raise rather than silently drop it.""" + _, jinja_template, _, _ = construct_chat_template( + tokenizer = _SuccessFakeTokenizer(), + chat_template = _NO_SYSTEM_CHAT_TEMPLATE, + default_system_message = None, + extra_eos_tokens = [""], + ) + with pytest.raises(RuntimeError, match = "Only user and assistant roles are supported!"): + _render( + jinja_template, + [ + {"role": "system", "content": "Be terse."}, + {"role": "user", "content": "Hi"}, + ], + ) diff --git a/unsloth/chat_templates.py b/unsloth/chat_templates.py index f47c78ba80..b857c34bcb 100644 --- a/unsloth/chat_templates.py +++ b/unsloth/chat_templates.py @@ -2652,6 +2652,12 @@ extra_eos_tokens = None, "{{ '" + full_system + "' }}"\ "{% set loop_messages = messages %}"\ "{% endif %}" + elif "{SYSTEM}" in system_part: + # Only bind loop_messages when the template can render a caller system + # message. A static prefix with no {SYSTEM} must still raise, not drop it. + partial_system += "{% else %}"\ + "{% set loop_messages = messages %}"\ + "{% endif %}" else: partial_system += "{% endif %}"