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:]