From d42256a5c590858605978bfe05f1300270c11bfb Mon Sep 17 00:00:00 2001 From: Vineeth Sai Date: Wed, 24 Jun 2026 00:40:53 -0700 Subject: [PATCH] Fix construct_chat_template leaking {INPUT}/{OUTPUT} sentinel into the chat template (#6531) * Fix construct_chat_template leaking {INPUT}/{OUTPUT} sentinel into the template In construct_chat_template's inner process() helper, the branch handling a section that starts with the {INPUT}/{OUTPUT} sentinel sliced the part from part.find(which) (which is 0 in that branch), so the literal sentinel was re-included in the generated Jinja chat template. The endswith branch already slices correctly with part[:part.find(which)]; this slices past the sentinel with part[len(which):], so a template whose input or output section begins with the sentinel (for example a user turn that starts with {INPUT}) renders correctly instead of emitting a literal {INPUT}/{OUTPUT}. Added a regression test covering {INPUT}-leading and {OUTPUT}-leading sections. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- ...test_construct_chat_template_validation.py | 37 +++++++++++++++++++ unsloth/chat_templates.py | 2 +- 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/tests/python/test_construct_chat_template_validation.py b/tests/python/test_construct_chat_template_validation.py index 2b1d012d14..66d3d80920 100644 --- a/tests/python/test_construct_chat_template_validation.py +++ b/tests/python/test_construct_chat_template_validation.py @@ -7,6 +7,8 @@ not silently drop the last char via s[:-1]. A minimal fake tokenizer keeps the cases CPU-only (no HF_TOKEN, no gated download). """ +from types import SimpleNamespace + import pytest from unsloth.chat_templates import construct_chat_template @@ -67,3 +69,38 @@ def test_error_message_excerpt_is_bounded(): # Excerpt is capped well under the template length. assert len(msg) < 1000 assert "{OUTPUT}" in msg + + +class _SuccessFakeTokenizer(_FakeTokenizer): + """Adds the surface construct_chat_template touches on the success path.""" + + bos_token = "" + bos_token_id = 1 + added_tokens_decoder: dict = {} + + def __call__(self, text): + # input_ids[0] must differ from bos_token_id so the BOS-handling branch is skipped. + return SimpleNamespace(input_ids = [5]) + + +@pytest.mark.parametrize( + "chat_template", + [ + # User turn begins with {INPUT} (no prefix before the sentinel). + "{INPUT} [/INST] {OUTPUT}{INPUT} [/INST] {OUTPUT}", + # Assistant turn begins with {OUTPUT} (no prefix before the sentinel). + "User: {INPUT}\n{OUTPUT}User: {INPUT}\n{OUTPUT}", + ], +) +def test_chat_template_does_not_leak_sentinel_when_section_starts_with_it(chat_template): + """When an input/output section begins with the {INPUT}/{OUTPUT} sentinel, the + generated Jinja template must not keep the literal sentinel text. The `startswith` + branch in the internal `process()` helper used to slice from `find()` (which is 0 + here) instead of past the sentinel, re-including the literal `{INPUT}`/`{OUTPUT}`.""" + _, jinja_template, _, _ = construct_chat_template( + tokenizer = _SuccessFakeTokenizer(), + chat_template = chat_template, + extra_eos_tokens = [""], + ) + assert "{INPUT}" not in jinja_template + assert "{OUTPUT}" not in jinja_template diff --git a/unsloth/chat_templates.py b/unsloth/chat_templates.py index fe078c5025..7f453bff82 100644 --- a/unsloth/chat_templates.py +++ b/unsloth/chat_templates.py @@ -2559,7 +2559,7 @@ extra_eos_tokens = None, if part.endswith(which): part = "'" + part[:part.find(which)] + f"' + {content}" elif part.startswith(which): - part = f"{content} + '" + part[part.find(which):] + "'" + part = f"{content} + '" + part[len(which):] + "'" else: part = "'" + part.replace(which, f"' + {content} + '") + "'" if part.startswith("'' + "): part = part[5:]