diff --git a/tests/python/test_construct_chat_template_validation.py b/tests/python/test_construct_chat_template_validation.py new file mode 100644 index 0000000000..9ab68639c4 --- /dev/null +++ b/tests/python/test_construct_chat_template_validation.py @@ -0,0 +1,77 @@ +"""Negative-path validation tests for unsloth.chat_templates.construct_chat_template. + +Regression coverage for the str.find() / regex no-match guards added in +PR #5763 follow-up: missing placeholders or unrecoverable two-example +structures must raise RuntimeError with a clear message, not IndexError +or AttributeError, and must never silently drop the last character via +s[:-1]. + +Uses a minimal fake tokenizer so the cases run on CPU-only CI without +HF_TOKEN and without downloading a gated model. The validation paths +exercised here fail before construct_chat_template reaches any heavy +tokenizer interaction, so the stub stays small. +""" + +import pytest + +from unsloth.chat_templates import construct_chat_template + + +class _FakeTokenizer: + """Minimum surface construct_chat_template touches before the + validation guards fire.""" + + name_or_path = "fake/tokenizer" + eos_token = "" + + def get_vocab(self): + return {"": 0} + + +@pytest.mark.parametrize( + "template, expected_in_message", + [ + ("only {INPUT} here, no output marker", "{OUTPUT}"), + ("only {OUTPUT} here, no input marker", "{INPUT}"), + ("neither sentinel here, just literal text", "{INPUT}"), + ("neither sentinel here, just literal text", "{OUTPUT}"), + ], +) +def test_missing_placeholder_in_chat_template_raises(template, expected_in_message): + with pytest.raises(RuntimeError) as exc_info: + construct_chat_template( + tokenizer = _FakeTokenizer(), + chat_template = template, + extra_eos_tokens = [""], + ) + assert expected_in_message in str(exc_info.value) + + +def test_single_pair_template_raises_clear_error_not_attribute_error(): + """One {INPUT}/{OUTPUT} pair (rather than the required two) used to + crash with AttributeError on `found.group(1)` after the for-loop + broke without setting `found`. Must raise RuntimeError now.""" + template = "user: {INPUT}\nassistant: {OUTPUT}\n" + with pytest.raises(RuntimeError): + construct_chat_template( + tokenizer = _FakeTokenizer(), + chat_template = template, + extra_eos_tokens = [""], + ) + + +def test_error_message_excerpt_is_bounded(): + """Error messages must include a bounded excerpt of the offending + template, not dump arbitrarily large content into the traceback.""" + huge = ("garbage " * 5000) + "{INPUT}" # ~40 KB, missing {OUTPUT} + with pytest.raises(RuntimeError) as exc_info: + construct_chat_template( + tokenizer = _FakeTokenizer(), + chat_template = huge, + extra_eos_tokens = [""], + ) + msg = str(exc_info.value) + # Excerpt is repr-quoted and capped; total message should stay well + # under the template length. + assert len(msg) < 1000 + assert "{OUTPUT}" in msg diff --git a/unsloth/chat_templates.py b/unsloth/chat_templates.py index e8a34cbc60..956fcb2392 100644 --- a/unsloth/chat_templates.py +++ b/unsloth/chat_templates.py @@ -2461,17 +2461,40 @@ extra_eos_tokens = None, f"{left_changed}" ) except: - ending = chat_template[chat_template.find("{OUTPUT}") + len("{OUTPUT}"):] + output_pos = chat_template.find("{OUTPUT}") + input_pos = chat_template.find("{INPUT}") + if output_pos == -1 or input_pos == -1: + missing = [] + if input_pos == -1: missing.append("{INPUT}") + if output_pos == -1: missing.append("{OUTPUT}") + raise RuntimeError( + f"Unsloth: chat_template must contain {' and '.join(missing)} " + f"placeholder(s). Got: {chat_template[:200]!r}" + ) + ending = chat_template[output_pos + len("{OUTPUT}"):] ending = re.escape(ending) find_text = "{INPUT}" + ending + "(.+?{OUTPUT}" + ending + ")" response_part = re.findall(find_text, chat_template, flags = re.DOTALL | re.MULTILINE) + if len(response_part) == 0: + raise RuntimeError( + "Unsloth: Could not recover a two-example structure from chat_template. " + "Provide exactly two {INPUT}/{OUTPUT} pairs (and optionally {SYSTEM}). " + f"Got: {chat_template[:200]!r}" + ) response_part = response_part[0] + found = None for j in range(1, len(response_part)): try_find = re.escape(response_part[:j]) try: found = next(re.finditer("(" + try_find + ").+?\\{INPUT\\}", chat_template, flags = re.DOTALL | re.MULTILINE)) except: break + if found is None: + raise RuntimeError( + "Unsloth: Could not locate a separator between examples in chat_template. " + "Provide exactly two {INPUT}/{OUTPUT} pairs (and optionally {SYSTEM}). " + f"Got: {chat_template[:200]!r}" + ) separator = found.group(1) response_start = chat_template.find(response_part) @@ -2607,8 +2630,20 @@ extra_eos_tokens = None, jinja_template = "{{ bos_token }}" + jinja_template # Get instruction and output parts for train_on_inputs = False - input_part = input_part [:input_part .find("{INPUT}")] - output_part = output_part[:output_part.find("{OUTPUT}")] + input_idx = input_part .find("{INPUT}") + output_idx = output_part.find("{OUTPUT}") + if input_idx == -1: + raise RuntimeError( + f"Unsloth: The instruction section of the template must contain the " + f"'{{INPUT}}' placeholder. Section: {input_part[:200]!r}" + ) + if output_idx == -1: + raise RuntimeError( + f"Unsloth: The response section of the template must contain the " + f"'{{OUTPUT}}' placeholder. Section: {output_part[:200]!r}" + ) + input_part = input_part [:input_idx ] + output_part = output_part[:output_idx] return modelfile, jinja_template, input_part, output_part