From f3f0676baa854caf6eb511d00d2930d31b66e238 Mon Sep 17 00:00:00 2001 From: noaa Date: Sun, 25 Oct 2026 18:30:08 +0200 Subject: [PATCH 1/2] Add history_adapters chat-template option to keep SR control tokens in history Text-based multi-turn callers re-render history without control tokens, so an SR turn's header renders its site token instead of the control token that was sent and the prefix cache misses from there on. history_adapters (one entry per assistant message) restores them in place of the site token, for both the 4.1 role-marker and 4.2 ChatML templates. Adds tests and documents the option and the caller rules it needs in SR_ARCHITECTURE.md. Signed-off-by: noaa --- docs/SR_ARCHITECTURE.md | 41 ++ .../composer/tokenizer_setup.py | 195 +++++++- tests/composer/test_chat_template.py | 428 ++++++++++++++++++ 3 files changed, 663 insertions(+), 1 deletion(-) diff --git a/docs/SR_ARCHITECTURE.md b/docs/SR_ARCHITECTURE.md index 0e3c3b92..e605b37f 100644 --- a/docs/SR_ARCHITECTURE.md +++ b/docs/SR_ARCHITECTURE.md @@ -144,6 +144,47 @@ No manual configuration is needed — pass SR adapter paths to the composer and `validate_cross_stream_population` then asserts every adapter in a `dual_stream` checkpoint has non-zero `cross_stream` weights, catching a silent load failure. +## Keeping Control Tokens in History + +A text-based multi-turn caller re-renders history from `messages`, which carry no +control tokens, so an SR turn's header renders with its site token instead of the +control token that was sent. For SR this changes no KV and no used output, but the +ids no longer match what was sent and the prefix cache misses from there on. + +Pass `history_adapters`, exactly one entry per assistant message (the adapter +that produced it, or `None`), and the template puts the tokens back. The control +token replaces the site token; nothing is added: + +```python +tokenizer.apply_chat_template(messages, add_generation_prompt=True, + adapter_name="sr2", history_adapters=[None, "sr1"]) +# vLLM /v1/chat/completions and mellea: same keys inside chat_template_kwargs +``` + +``` +4.1 assistant<|end_of_role|>A1 -> assistant<|sr1|>A1 +4.2 enable_thinking=F assistant\nA1 -> assistant\n<|sr1|>A1 +``` + +A list of the wrong length, or an entry that is not an SR adapter, raises. No +effect when the current turn is base. + +For the cache to hit, the caller must: + +* store the exact answer text the model produced (no post-processing), and send + `skip_special_tokens: False` so a control token generated mid-answer is not + stripped from it; +* on Granite 4.2 with `enable_thinking=True`, store `"\n"` + the returned + text and pass `truncate_history_thinking=False`; otherwise history thinking is + rewritten and the match breaks right after the control token; +* serve Granite 4.2 with `--chat-template-content-format string`, or vLLM rejects + every chat request with a 400. + +Even exact text can re-encode to different ids (the model may emit `"\n"` + `"\n"` +where the tokenizer picks `"\n\n"`), so the cache misses from that point on that +turn. Only a caller that keeps the generated ids and sends them to `/v1/completions` +avoids this. + ## Config Parameters | Parameter | Type | Default | Description | diff --git a/src/granite_switch/composer/tokenizer_setup.py b/src/granite_switch/composer/tokenizer_setup.py index a127835d..d681b831 100644 --- a/src/granite_switch/composer/tokenizer_setup.py +++ b/src/granite_switch/composer/tokenizer_setup.py @@ -9,6 +9,8 @@ import re from dataclasses import dataclass +import jinja2 + from .adapter_loader import load_adapter_config # Name of the optional return-to-base control token (``<|base_reset|>``). It is @@ -556,6 +558,185 @@ def _rewrite(emission: "re.Match") -> str: return head + region, site, count +#: A ``for`` / ``endfor`` statement tag, for depth-matching the message loop. +#: Accepts every whitespace-control spelling: ``{%``, ``{%-`` and ``{%+``. +_FOR_TAG_RE = re.compile(r"\{%[-+]?\s*(for|endfor)\b") + + +def _jinja_string_literal(text: str) -> str: + """*text* as a single-quoted Jinja string literal that renders back to it.""" + escaped = text.replace("\\", "\\\\").replace("'", "\\'").replace("\n", "\\n") + return "'" + escaped + "'" + + +def _sr_history_header(template: str, site: str) -> str: + """The rendered assistant header, through *site*, that history turns repeat. + + Taken from the base template's generation prompt: every branch opens with + the same assistant header, and a history assistant turn renders that same + header before its content (``<|start_of_role|>assistant<|end_of_role|>`` on + 4.1, ``<|im_start|>assistant\\n`` on 4.2, where the history turn + continues ````). + + Raises: + ValueError: If the generation prompt's branches share no prefix that + contains *site*. + """ + _, region = _split_at_generation_prompt(template) + literals = [_unescape_literal(body) for body in _generation_prompt_literals(region)] + prefix = os.path.commonprefix(literals) if literals else "" + cut = prefix.rfind(site) + if cut == -1: + raise ValueError( + f"The generation prompt's branches share no prefix containing the " + f"Shadow Residual substitution site {site!r} ({literals!r}), so the " + f"history_adapters rewrite has no header to match." + ) + return prefix[: cut + len(site)] + + +def _wrap_message_loop_for_history_adapters( + template: str, loop_var_source: str, header: str, site: str +) -> str: + """Let ``history_adapters`` put SR control tokens back into history turns. + + Rendering history from ``messages`` emits every assistant header with the + SR *site* (the token the control token replaced at generation time), so the + re-rendered ids diverge from the ones that were sent and the prefix cache + misses. ``history_adapters[k]`` names the adapter that produced the k-th + assistant message; this rewrites that message's header back so its *site* + is the control token. + + The match is on the whole *header* (e.g. ``<|im_start|>assistant\\n``), + which ends in *site*, not on *site* alone. Granite 4.2 writes its own + ```` only when the answer has no think tags, so for an answer + that mentions ```` the first bare ```` is the answer's text; the + header text cannot occur there. + + The main loop body is wrapped in ``{% filter replace(header, new, 1) %}``, + so the replacement applies to the body's rendered text, rather than + rewriting each assistant emission in the source: the 4.2 template emits the + assistant header from several branches and reads ```` in its own + history-truncation logic, so a token placed in the source would be seen + (and stripped, or misread) by that logic. Rendered-text replacement runs + after it and is indifferent to which branch ran. + + Only applied when the current turn is SR (``ns.adapter_type == 'sr'``): a + base turn after a kept SR token would generate under SR (latest-wins). + + ``history_adapters`` must have exactly one entry per assistant message + (``None`` for a turn that had no adapter); any other length raises. + + Raises: + ValueError: If *header* does not end in *site*, the template has no + main message loop with a matching ``endfor``, or the wrapped + template no longer parses (e.g. an ``endfor`` inside a comment + fooled the tag scan). + """ + if not header.endswith(site): + raise ValueError( + f"history_adapters header {header!r} does not end in the Shadow " + f"Residual substitution site {site!r}." + ) + open_tag = re.search( + r"\{%[-+]?\s*for message in " + re.escape(loop_var_source) + r"\s*[-+]?%\}", + template, + ) + if open_tag is None: + raise ValueError( + f"No '{{%- for message in {loop_var_source} %}}' loop in the chat " + f"template, so history_adapters has nothing to wrap." + ) + depth = 1 + for tag in _FOR_TAG_RE.finditer(template, open_tag.end()): + depth += 1 if tag.group(1) == "for" else -1 + if depth == 0: + close_at = tag.start() + break + else: + raise ValueError( + "The chat template's message loop has no matching endfor, so " + "history_adapters has nothing to wrap." + ) + + # Checked once, before the loop: a list that does not line up one-to-one + # with the assistant messages is a caller bug, and silently padding or + # truncating it would put tokens on the wrong turns. + precheck = ( + "{%- if ns.adapter_type == '" + ANCHOR_MODE_SR + "'" + " and history_adapters is defined and history_adapters is not none %}\n" + "{%- if history_adapters is string or history_adapters is mapping" + " or history_adapters is not iterable %}\n" + "{{- raise_exception('history_adapters must be a list of adapter names " + "(or none), one per assistant message') }}\n" + "{%- endif %}\n" + "{%- set _gs_expected = " + loop_var_source + " | selectattr('role', " + "'equalto', 'assistant') | list | length %}\n" + "{%- if history_adapters | length != _gs_expected %}\n" + "{{- raise_exception('history_adapters has ' ~ (history_adapters | length)" + " ~ ' entries but the conversation has ' ~ _gs_expected ~ ' assistant " + "messages; pass exactly one entry (an adapter name or none) per " + "assistant message') }}\n" + "{%- endif %}\n" + "{%- for _gs_entry in history_adapters %}\n" + # ``is not string`` first: ``or`` short-circuits, so an unhashable entry + # never reaches the ``in`` lookup and gets this message, not a TypeError. + "{%- if _gs_entry is not none and (_gs_entry is not string" + " or _gs_entry not in adapter_map" + " or adapter_map[_gs_entry]['type'] != '" + ANCHOR_MODE_SR + "') %}\n" + "{{- raise_exception('history_adapters names ' ~ _gs_entry ~ ', which is " + "not a Shadow Residual adapter of this checkpoint') }}\n" + "{%- endif %}\n" + "{%- endfor %}\n" + "{%- endif %}\n" + ) + # Picks this message's replacement before its body renders; the default + # (header -> header) leaves the body untouched. + choose = ( + "\n{%- set _gs_new = " + _jinja_string_literal(header) + " %}\n" + "{%- if message.role == 'assistant' %}\n" + "{%- if ns.adapter_type == '" + ANCHOR_MODE_SR + "'" + " and history_adapters is defined and history_adapters is not none" + " and history_adapters[ns.assistant_idx] is not none %}\n" + "{%- set _gs_new = " + + _jinja_string_literal(header[: len(header) - len(site)]) + + " ~ adapter_map[history_adapters[ns.assistant_idx]]['token'] %}\n" + "{%- endif %}\n" + "{%- set ns.assistant_idx = ns.assistant_idx + 1 %}\n" + "{%- endif %}\n" + "{%- filter replace(" + _jinja_string_literal(header) + ", _gs_new, 1) %}" + ) + wrapped = ( + template[: open_tag.start()] + + precheck + + template[open_tag.start() : open_tag.end()] + + choose + + template[open_tag.end() : close_at] + + "{%- endfilter %}\n" + + template[close_at:] + ) + + # The tag scan is textual, so a template it misreads composes fine and then + # fails at render time. Parse both versions to catch that here instead; a + # template that does not parse even before the wrap (an extension tag this + # environment lacks) is left for the real renderer to judge. + env = jinja2.Environment(extensions=["jinja2.ext.loopcontrols"]) + try: + env.parse(template) + except jinja2.TemplateSyntaxError: + return wrapped + try: + env.parse(wrapped) + except jinja2.TemplateSyntaxError as exc: + raise ValueError( + f"Wrapping the chat template's message loop for history_adapters " + f"produced a template that does not parse ({exc}). The loop's " + f"endfor was probably mis-identified, e.g. one inside a comment or " + f"a string." + ) from exc + return wrapped + + def add_control_tokens( tokenizer, discovered_adapters: list[tuple[str | None, str, str, str | None]], @@ -1214,7 +1395,9 @@ def configure_chat_template( "\n adapter_invocation_text=adapter_invocation_text," "\n adapter_invocation_drop_chars=adapter_invocation_drop_chars," "\n alora_target_idx=-1," - "\n skip_next_role_marker=false" + # history_adapters state (SR only): the assistant message being visited. + + ("\n assistant_idx=0," if sr_anchors else "") + + "\n skip_next_role_marker=false" "\n )" ) modified_chat_template = ( @@ -1374,6 +1557,12 @@ def configure_chat_template( for name, info in adapter_mapping.items() if info["type"] == ANCHOR_MODE_SR } + modified_chat_template = _wrap_message_loop_for_history_adapters( + modified_chat_template, + fmt.loop_var_source, + _sr_history_header(base_chat_template, sr_site), + sr_site, + ) # Skip-once wrapper for every role-open marker emission in the template. # ns.skip_next_role_marker is set to true immediately after a LoRA or @@ -1470,4 +1659,8 @@ def _split_merged(match: "re.Match") -> str: f" - SR tokens: substituted for {sr_site!r} in all " f"{sr_replacements} generation-prompt render path(s)" ) + print( + f" - SR history: history_adapters=[...] restores {sr_site!r} -> " + f"control token in earlier assistant turns" + ) return fmt.name, sr_substitute_token_ids diff --git a/tests/composer/test_chat_template.py b/tests/composer/test_chat_template.py index 5cbabc0c..32744bde 100644 --- a/tests/composer/test_chat_template.py +++ b/tests/composer/test_chat_template.py @@ -23,6 +23,7 @@ - End-to-end: adapter_config.json → render (no patching) """ +import json import os from types import SimpleNamespace from unittest.mock import patch @@ -31,6 +32,8 @@ from jinja2 import Environment from granite_switch.composer.tokenizer_setup import ( + _wrap_message_loop_for_history_adapters, + add_control_tokens, configure_audio_chat_template, configure_chat_template, detect_template_format, @@ -1423,3 +1426,428 @@ def test_sr_anchor_id_must_match_the_tokenizer(self): configure_chat_template( tokenizer, [(self._SR_GRANITE, "answerability", "sr")] ) + + +class TestShadowResidualHistoryAdapters: + """``history_adapters`` restores SR control tokens in earlier assistant turns. + + For SR the token at a history position changes no KV (K/V come from the base + stream) and no used output; what it buys is id equality with what was sent, + which is what vLLM's prefix cache keys on. So the central property is a prefix + relation: the turn-N prompt plus the model's answer must be a prefix of the + turn-N+1 render. + + Renders use HF's whitespace flags (``trim_blocks``/``lstrip_blocks``), as + ``apply_chat_template`` and vLLM do: under plain Jinja the 4.2 template's own + history-truncation path leaks whitespace, so the prefix relation fails there + even without this feature. + """ + + _ENV = Environment(trim_blocks=True, lstrip_blocks=True) + + _SR_GRANITE = os.path.join(_FIXTURES, "sr_answerability_adapter") + _SR_CHATML = os.path.join(_FIXTURES, "sr_answerability_adapter_chatml") + + _MESSAGES = [ + {"role": "user", "content": "Q1"}, + {"role": "assistant", "content": "A1"}, + {"role": "user", "content": "Q2"}, + {"role": "assistant", "content": "A2"}, + {"role": "user", "content": "Q3"}, + ] + + @staticmethod + def _raise(message): + raise ValueError(message) + + def _granite(self): + tokenizer = _SRTokenizer(_GRANITE_TEMPLATE, {"<|end_of_role|>": [49153]}) + configure_chat_template( + tokenizer, + [(self._SR_GRANITE, "sr1", "sr"), (self._SR_GRANITE, "sr2", "sr")], + ) + return tokenizer + + def _chatml(self): + tokenizer = _SRTokenizer(_CHATML_TEMPLATE, {"": [49154]}) + configure_chat_template( + tokenizer, + [(self._SR_CHATML, "sr1", "sr"), (self._SR_CHATML, "sr2", "sr")], + ) + return tokenizer + + def _render(self, tokenizer, messages, **kwargs): + return self._ENV.from_string(tokenizer.chat_template).render( + messages=messages, + add_generation_prompt=True, + raise_exception=self._raise, + **kwargs, + ) + + def test_granite_history_headers_carry_their_control_tokens(self): + result = self._render( + self._granite(), + self._MESSAGES, + adapter_name="sr1", + history_adapters=["sr1", "sr2"], + ) + assert "<|start_of_role|>assistant<|sr1|>A1<|end_of_text|>" in result + assert "<|start_of_role|>assistant<|sr2|>A2<|end_of_text|>" in result + assert result.endswith("<|start_of_role|>assistant<|sr1|>") + + def test_granite_sent_prompt_is_prefix_of_next_turn(self): + tokenizer = self._granite() + sent = self._render( + tokenizer, self._MESSAGES[:3], adapter_name="sr2", history_adapters=["sr1"] + ) + nxt = self._render( + tokenizer, + self._MESSAGES, + adapter_name="sr1", + history_adapters=["sr1", "sr2"], + ) + assert nxt.startswith(sent + "A2<|end_of_text|>\n") + + def test_chatml_no_thinking_sent_prompt_is_prefix_of_next_turn(self): + tokenizer = self._chatml() + sent = self._render( + tokenizer, + self._MESSAGES[:3], + enable_thinking=False, + adapter_name="sr2", + history_adapters=["sr1"], + ) + nxt = self._render( + tokenizer, + self._MESSAGES, + enable_thinking=False, + adapter_name="sr1", + history_adapters=["sr1", "sr2"], + ) + assert "<|im_start|>assistant\n<|sr1|>A1" in nxt + assert nxt.startswith(sent + "A2<|im_end|>\n") + + def test_chatml_thinking_places_tokens_but_breaks_the_prefix(self): + """With thinking on and the answer stored as returned, the sent prompt + ends ``<|sr2|>\\n`` but history re-renders as ``<|sr2|>A``, so no + prefix relation holds. The next test shows the caller-side fix. + """ + tokenizer = self._chatml() + sent = self._render( + tokenizer, + self._MESSAGES[:3], + enable_thinking=True, + adapter_name="sr2", + history_adapters=["sr1"], + ) + nxt = self._render( + tokenizer, + self._MESSAGES, + enable_thinking=True, + adapter_name="sr1", + history_adapters=["sr1", "sr2"], + ) + assert "<|im_start|>assistant\n<|sr2|>A2" in nxt + assert not nxt.startswith(sent + "A2") + + @pytest.mark.parametrize("truncate", [False, True]) + def test_chatml_thinking_prefix_holds_with_think_restored(self, truncate): + """With thinking on, the control token takes ````'s place in the sent + prompt, so the model's text starts after it. Storing ``"\\n"`` + that + text gives the template its header back and the prefix holds -- but only + with ``truncate_history_thinking=False``: truncation strips the reasoning + out of every answer before the last user turn. + """ + tokenizer = self._chatml() + answers = ["R1\n\nA1", "R2\n\nA2"] + messages = [ + {"role": "user", "content": "Q1"}, + {"role": "assistant", "content": "\n" + answers[0]}, + {"role": "user", "content": "Q2"}, + {"role": "assistant", "content": "\n" + answers[1]}, + {"role": "user", "content": "Q3"}, + ] + kwargs = {"enable_thinking": True, "truncate_history_thinking": truncate} + sent = self._render( + tokenizer, + messages[:3], + adapter_name="sr2", + history_adapters=["sr1"], + **kwargs, + ) + nxt = self._render( + tokenizer, + messages, + adapter_name="sr1", + history_adapters=["sr1", "sr2"], + **kwargs, + ) + assert sent.endswith("<|im_start|>assistant\n<|sr2|>\n") + assert nxt.startswith(sent + answers[1] + "<|im_end|>\n") is not truncate + + @pytest.mark.parametrize("fmt", ["granite", "chatml"]) + def test_base_turn_emits_no_history_tokens(self, fmt): + tokenizer = self._granite() if fmt == "granite" else self._chatml() + result = self._render( + tokenizer, self._MESSAGES, history_adapters=["sr1", "sr2"] + ) + assert "<|sr" not in result + + @pytest.mark.parametrize("fmt", ["granite", "chatml"]) + def test_history_adapters_none_is_a_no_op(self, fmt): + tokenizer = self._granite() if fmt == "granite" else self._chatml() + absent = self._render(tokenizer, self._MESSAGES, adapter_name="sr1") + explicit_none = self._render( + tokenizer, self._MESSAGES, adapter_name="sr1", history_adapters=None + ) + all_none = self._render( + tokenizer, + self._MESSAGES, + adapter_name="sr1", + history_adapters=[None, None], + ) + assert absent == explicit_none == all_none + assert absent.count("<|sr1|>") == 1 + + def test_absent_history_adapters_matches_the_unpatched_template(self): + """The wrap is invisible when unused: undoing the generation-prompt + substitution reproduces the original template's render exactly.""" + for enable_thinking in (False, True): + base = self._ENV.from_string(_CHATML_TEMPLATE).render( + messages=self._MESSAGES, + add_generation_prompt=True, + enable_thinking=enable_thinking, + ) + patched = self._render( + self._chatml(), + self._MESSAGES, + enable_thinking=enable_thinking, + adapter_name="sr1", + ) + assert patched.replace("<|sr1|>", "") == base + base = self._ENV.from_string(_GRANITE_TEMPLATE).render( + messages=self._MESSAGES, add_generation_prompt=True + ) + patched = self._render(self._granite(), self._MESSAGES, adapter_name="sr1") + assert patched.replace("<|sr1|>", "<|end_of_role|>") == base + + def test_render_identical_under_plain_jinja(self): + """The class renders with HF's flags; the filter block must not shift + whitespace under plain Jinja either.""" + env = Environment() + for tokenizer, site, base_template in ( + (self._granite(), "<|end_of_role|>", _GRANITE_TEMPLATE), + (self._chatml(), "", _CHATML_TEMPLATE), + ): + base = env.from_string(base_template).render( + messages=self._MESSAGES, add_generation_prompt=True + ) + patched = env.from_string(tokenizer.chat_template).render( + messages=self._MESSAGES, + add_generation_prompt=True, + adapter_name="sr1", + raise_exception=self._raise, + ) + assert patched.replace("<|sr1|>", site) == base + + def test_index_counts_assistant_messages_only(self): + messages = [ + {"role": "system", "content": "S"}, + {"role": "user", "content": "Q1"}, + {"role": "assistant", "content": "A1"}, + {"role": "tool", "content": "T1"}, + {"role": "user", "content": "Q2"}, + {"role": "assistant", "content": "A2"}, + {"role": "user", "content": "Q3"}, + ] + result = self._render( + self._granite(), + messages, + adapter_name="sr1", + history_adapters=[None, "sr2"], + ) + assert "assistant<|end_of_role|>A1" in result + assert "assistant<|sr2|>A2" in result + + @pytest.mark.parametrize( + "history_adapters", + [["sr2"], ["sr2", "sr2", "sr2"], []], + ids=["too-short", "too-long", "empty"], + ) + def test_length_mismatch_raises(self, history_adapters): + """One entry per assistant message, no more and no fewer: a list that + does not line up would put tokens on the wrong turns.""" + with pytest.raises(ValueError, match="has 2 assistant messages"): + self._render( + self._granite(), + self._MESSAGES, + adapter_name="sr1", + history_adapters=history_adapters, + ) + + def test_length_is_not_checked_on_a_base_turn(self): + """A base turn does not read the list, so a wrong length is not an error.""" + result = self._render(self._granite(), self._MESSAGES, history_adapters=["sr2"]) + assert "<|sr" not in result + + def test_unknown_adapter_raises(self): + with pytest.raises(ValueError, match="not a Shadow Residual adapter"): + self._render( + self._granite(), + self._MESSAGES, + adapter_name="sr1", + history_adapters=["nope", None], + ) + + def test_non_sr_checkpoint_is_unaffected(self): + with patch(_PATCH_TARGET, return_value=""): + tokenizer = _make_tokenizer() + configure_chat_template(tokenizer, [("/path/a", "req_check", "alora")]) + assert "history_adapters" not in tokenizer.chat_template + assert "assistant_idx" not in tokenizer.chat_template + + def test_chatml_answer_containing_the_site_is_not_rewritten(self): + """4.2 writes its own ```` header only when the answer has + no think tags, so for this answer the first ```` is the answer's + own text. Matching the full header leaves it alone.""" + messages = [ + {"role": "user", "content": "Q1"}, + {"role": "assistant", "content": "See tag"}, + {"role": "user", "content": "Q2"}, + ] + result = self._render( + self._chatml(), + messages, + enable_thinking=False, + adapter_name="sr1", + history_adapters=["sr2"], + ) + assert "<|im_start|>assistant\nSee tag<|im_end|>" in result + assert "<|sr2|>" not in result + + def test_bare_string_raises(self): + """A string is a sequence of characters, not a list of names.""" + with pytest.raises(ValueError, match="must be a list"): + self._render( + self._granite(), + self._MESSAGES, + adapter_name="sr1", + history_adapters="sr1", + ) + + def test_non_string_entry_raises_the_named_error(self): + with pytest.raises(ValueError, match="not a Shadow Residual adapter"): + self._render( + self._granite(), + self._MESSAGES, + adapter_name="sr1", + history_adapters=[["sr1"], None], + ) + + def test_wrap_rejects_a_result_that_does_not_parse(self): + """An ``endfor`` inside a comment fools the tag scan; the wrapped + template must fail at compose time, not at render time.""" + template = ( + "{%- set ns = namespace(adapter_type='') %}" + "{%- for message in messages %}{# a {% endfor %} b #}" + "{{ message.content }}{%- endfor %}" + ) + with pytest.raises(ValueError, match="does not parse"): + _wrap_message_loop_for_history_adapters(template, "messages", "", "") + + def test_wrap_counts_plus_whitespace_control_tags(self): + """``{%+ for`` opens a loop like ``{%- for`` does.""" + template = ( + "{%- set ns = namespace(adapter_type='', assistant_idx=0) %}" + "{%- for message in messages %}" + "{%+ for x in [1] %}{{ message.content }}{% endfor %}" + "{%- endfor %}" + ) + wrapped = _wrap_message_loop_for_history_adapters( + template, "messages", "", "" + ) + rendered = self._ENV.from_string(wrapped).render( + messages=[{"role": "user", "content": "a"}], adapter_map={} + ) + assert rendered == "a" + + +class TestShadowResidualHistoryAdaptersTokenIds: + """The prefix relation on token ids, with the real Granite tokenizers. + + vLLM's prefix cache keys on ids, not text; matching strings can still + tokenize differently at a boundary such as the one right after a control + token. Skipped when the tokenizer cannot be fetched. + """ + + _CASES = [ + ( + "ibm-granite/granite-4.1-3b", + "sr_answerability_adapter", + {}, + "<|end_of_text|>", + ), + ( + "ibm-granite/granite-4.2-3b", + "sr_answerability_adapter_chatml", + {"enable_thinking": False}, + "<|im_end|>", + ), + ] + + @staticmethod + def _ids(out): + # transformers 5.x returns a BatchEncoding here; 4.x returns a plain list. + return list(out["input_ids"]) if hasattr(out, "keys") else list(out) + + @pytest.mark.parametrize( + "model,fixture,kwargs,turn_end", _CASES, ids=["granite-4.1", "granite-4.2"] + ) + def test_sent_ids_are_prefix_of_next_turn( + self, model, fixture, kwargs, turn_end, tmp_path + ): + from transformers import AutoTokenizer + + try: + tokenizer = AutoTokenizer.from_pretrained(model) + except Exception as e: + pytest.skip(f"could not fetch {model} tokenizer: {e}") + # The fixture pins the fake tokenizer's site id; drop it so the composer + # resolves the id against the real one. + with open(os.path.join(_FIXTURES, fixture, "adapter_config.json")) as f: + adapter_config = json.load(f) + adapter_config.pop("last_context_token_id", None) + adapter_path = str(tmp_path / fixture) + os.makedirs(adapter_path) + with open(os.path.join(adapter_path, "adapter_config.json"), "w") as f: + json.dump(adapter_config, f) + adapters = [ + (adapter_path, "sr1", "sr", None), + (adapter_path, "sr2", "sr", None), + ] + add_control_tokens(tokenizer, adapters) + configure_chat_template(tokenizer, [a[:3] for a in adapters]) + + messages = TestShadowResidualHistoryAdapters._MESSAGES + sent = self._ids( + tokenizer.apply_chat_template( + messages[:3], + add_generation_prompt=True, + adapter_name="sr2", + history_adapters=[None], + **kwargs, + ) + ) + answer = tokenizer.encode( + messages[3]["content"] + turn_end, add_special_tokens=False + ) + nxt = self._ids( + tokenizer.apply_chat_template( + messages, + add_generation_prompt=True, + adapter_name="sr1", + history_adapters=[None, "sr2"], + **kwargs, + ) + ) + assert nxt[: len(sent) + len(answer)] == sent + answer From fdc436a50d6c0dd410dabd3f6595fc5b98a52033 Mon Sep 17 00:00:00 2001 From: noaa Date: Tue, 3 Nov 2026 16:24:31 +0200 Subject: [PATCH 2/2] Match full SR headers and validate the history_adapters loop Fixes the review on #154: 4.2 answers starting with are no longer rewritten, the wrap keeps the template's whitespace, and the composer raises on templates it can't wrap correctly. Declares jinja2. Signed-off-by: noaa --- docs/SR_ARCHITECTURE.md | 5 +- pyproject.toml | 1 + .../composer/tokenizer_setup.py | 311 +++++++++++++----- tests/composer/test_chat_template.py | 157 +++++++++ uv.lock | 2 + 5 files changed, 397 insertions(+), 79 deletions(-) diff --git a/docs/SR_ARCHITECTURE.md b/docs/SR_ARCHITECTURE.md index e605b37f..1e3bec74 100644 --- a/docs/SR_ARCHITECTURE.md +++ b/docs/SR_ARCHITECTURE.md @@ -167,7 +167,10 @@ tokenizer.apply_chat_template(messages, add_generation_prompt=True, ``` A list of the wrong length, or an entry that is not an SR adapter, raises. No -effect when the current turn is base. +effect when the current turn is base. A turn whose render has no assistant header +(e.g. a 4.2 thinking answer stored without its leading `\n`) is left +unchanged, with no control token and no error: the base render already differs +from what was sent there, so no token could make the cache hit. For the cache to hit, the caller must: diff --git a/pyproject.toml b/pyproject.toml index ee201ae3..d5bb9487 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,6 +12,7 @@ requires-python = ">=3.11,<3.14" dependencies = [ "torch>=2.11.0", "transformers>=5.16,<=5.17.0", + "jinja2>=3.1", ] [project.urls] diff --git a/src/granite_switch/composer/tokenizer_setup.py b/src/granite_switch/composer/tokenizer_setup.py index d681b831..476b020a 100644 --- a/src/granite_switch/composer/tokenizer_setup.py +++ b/src/granite_switch/composer/tokenizer_setup.py @@ -10,6 +10,8 @@ from dataclasses import dataclass import jinja2 +import jinja2.ext +import jinja2.nodes from .adapter_loader import load_adapter_config @@ -465,7 +467,7 @@ def _resolve_sr_substitution_site(region: str, anchor: str, tokenizer) -> str: def _replace_anchor_in_generation_prompt( template: str, anchor: str, tokenizer -) -> tuple[str, str, int]: +) -> tuple[str, str, int, str, list[str]]: """Make every generation-prompt branch emit the control token in place of one token. @@ -484,9 +486,12 @@ def _replace_anchor_in_generation_prompt( tokenizer: Used to check candidate sites are a single token. Returns: - ``(patched_template, site, num_replacements)`` where *site* is the - literal actually replaced — the anchor itself unless the branches forced - an earlier common token (see :func:`_resolve_sr_substitution_site`). + ``(patched_template, site, num_replacements, header, literals)`` where + *site* is the literal actually replaced — the anchor itself unless the + branches forced an earlier common token (see + :func:`_resolve_sr_substitution_site`) — and *header* / *literals* are + what the ``history_adapters`` wrap matches (see + :func:`_sr_history_header`), taken from the same branches. Raises: ValueError: If there is no ``add_generation_prompt`` block, if the site @@ -555,7 +560,10 @@ def _rewrite(emission: "re.Match") -> str: f"control token, otherwise the ones that miss it render an adapter " f"that never activates." ) - return head + region, site, count + header, history_literals = _sr_history_header( + [_unescape_literal(body) for body in literals], site + ) + return head + region, site, count, header, history_literals #: A ``for`` / ``endfor`` statement tag, for depth-matching the message loop. @@ -563,27 +571,91 @@ def _rewrite(emission: "re.Match") -> str: _FOR_TAG_RE = re.compile(r"\{%[-+]?\s*(for|endfor)\b") +class _GenerationTag(jinja2.ext.Extension): + """Parses HF's ``{% generation %}`` block (assistant-token masking) so that + templates using it parse here too; the block is marked so the + ``history_adapters`` wrap can refuse one inside the message loop.""" + + tags = {"generation"} + + def parse(self, parser): + lineno = next(parser.stream).lineno + body = parser.parse_statements(("name:endgeneration",), drop_needle=True) + # Jinja forbids custom node types, so the block is an ordinary scope + # carrying an extension attribute that marks it. + return jinja2.nodes.Scope( + [jinja2.nodes.ExtensionAttribute(self.identifier, "generation"), *body], + lineno=lineno, + ) + + +def _is_generation_block(node: "jinja2.nodes.Node") -> bool: + return ( + isinstance(node, jinja2.nodes.Scope) + and bool(node.body) + and isinstance(node.body[0], jinja2.nodes.ExtensionAttribute) + and node.body[0].identifier == _GenerationTag.identifier + ) + + +#: Parses chat templates the way HF's renderer accepts them. +_TEMPLATE_PARSE_ENV = jinja2.Environment( + extensions=["jinja2.ext.loopcontrols", _GenerationTag] +) + + +def _message_loop_node(tree: "jinja2.nodes.Template", loop_var_source: str): + """The first ``for message in `` loop in *tree*, or None.""" + for node in tree.find_all(jinja2.nodes.For): + if ( + isinstance(node.target, jinja2.nodes.Name) + and node.target.name == "message" + and isinstance(node.iter, jinja2.nodes.Name) + and node.iter.name == loop_var_source + ): + return node + return None + + +def _has_own_loop_control(nodes_: list) -> bool: + """Whether *nodes_* hold a ``continue``/``break`` that belongs to their own + loop, i.e. one not inside a nested ``for``.""" + for node in nodes_: + if isinstance(node, (jinja2.nodes.Continue, jinja2.nodes.Break)): + return True + if isinstance(node, jinja2.nodes.For): + # Its body's controls are the inner loop's; only its else-block + # runs in the outer loop's iteration. + if _has_own_loop_control(node.else_): + return True + continue + if _has_own_loop_control(list(node.iter_child_nodes())): + return True + return False + + def _jinja_string_literal(text: str) -> str: """*text* as a single-quoted Jinja string literal that renders back to it.""" escaped = text.replace("\\", "\\\\").replace("'", "\\'").replace("\n", "\\n") return "'" + escaped + "'" -def _sr_history_header(template: str, site: str) -> str: - """The rendered assistant header, through *site*, that history turns repeat. +def _sr_history_header(literals: list[str], site: str) -> tuple[str, list[str]]: + """The rendered assistant header, through *site*, that history turns repeat, + and the generation prompt's full *literals* (rendered text) that start with it. - Taken from the base template's generation prompt: every branch opens with - the same assistant header, and a history assistant turn renders that same - header before its content (``<|start_of_role|>assistant<|end_of_role|>`` on - 4.1, ``<|im_start|>assistant\\n`` on 4.2, where the history turn - continues ````). + Every generation-prompt branch opens with the same assistant header, and a + history assistant turn renders that same header before its content + (``<|start_of_role|>assistant<|end_of_role|>`` on 4.1, + ``<|im_start|>assistant\\n`` on 4.2, where the history turn continues + ````). The full literals (4.2: ``...\\n`` and + ``...``) are what a history turn renders when the template + wrote the header itself, so matching them rather than the bare header keeps + an answer that merely starts with *site* untouched. Raises: - ValueError: If the generation prompt's branches share no prefix that - contains *site*. + ValueError: If the branches share no prefix that contains *site*. """ - _, region = _split_at_generation_prompt(template) - literals = [_unescape_literal(body) for body in _generation_prompt_literals(region)] prefix = os.path.commonprefix(literals) if literals else "" cut = prefix.rfind(site) if cut == -1: @@ -592,11 +664,17 @@ def _sr_history_header(template: str, site: str) -> str: f"Shadow Residual substitution site {site!r} ({literals!r}), so the " f"history_adapters rewrite has no header to match." ) - return prefix[: cut + len(site)] + header = prefix[: cut + len(site)] + # Longest first, so a literal that extends another is tried before it. + return header, sorted(set(literals), key=len, reverse=True) def _wrap_message_loop_for_history_adapters( - template: str, loop_var_source: str, header: str, site: str + template: str, + loop_var_source: str, + header: str, + site: str, + literals: list[str] | None = None, ) -> str: """Let ``history_adapters`` put SR control tokens back into history turns. @@ -607,20 +685,33 @@ def _wrap_message_loop_for_history_adapters( assistant message; this rewrites that message's header back so its *site* is the control token. - The match is on the whole *header* (e.g. ``<|im_start|>assistant\\n``), - which ends in *site*, not on *site* alone. Granite 4.2 writes its own - ```` only when the answer has no think tags, so for an answer - that mentions ```` the first bare ```` is the answer's text; the - header text cannot occur there. - - The main loop body is wrapped in ``{% filter replace(header, new, 1) %}``, - so the replacement applies to the body's rendered text, rather than - rewriting each assistant emission in the source: the 4.2 template emits the - assistant header from several branches and reads ```` in its own + The match is on each full generation-prompt *literal* (e.g. 4.2's + ``<|im_start|>assistant\\n`` and ``...\\n``), each of + which starts with *header*, not on *site* or *header* alone. Granite 4.2 + writes its own ```` only when the answer has no think tags, + so for an answer that mentions ```` -- even at its very start -- the + first ```` is the answer's text, and only the full literal tells the + two apart. The thinking-on form a caller stores (``"\\n"`` + answer) + renders the ``...\\n`` literal, so it still gets its token. Only an + answer that itself starts with a full literal's tail (``\\n`` or + ````) is indistinguishable from the header. *literals* + defaults to ``[header]``. + + The replacement applies to each message's rendered text rather than to + each assistant emission in the source: the 4.2 template emits the assistant + header from several branches and reads ```` in its own history-truncation logic, so a token placed in the source would be seen (and stripped, or misread) by that logic. Rendered-text replacement runs after it and is indifferent to which branch ran. + Each message's body is wrapped in a ``{% filter %}`` that does the + replacement. A ``{% continue %}`` / ``{% break %}`` would leave that filter + before it prints, dropping the message's output, and HF's + ``{% generation %}`` would take its assistant-mask offsets from text the + filter has not printed yet. A message loop that uses either (a nested + loop's own continue/break are fine), or has a ``{% else %}``, is rejected + at compose time. + Only applied when the current turn is SR (``ns.adapter_type == 'sr'``): a base turn after a kept SR token would generate under SR (latest-wins). @@ -628,16 +719,49 @@ def _wrap_message_loop_for_history_adapters( (``None`` for a turn that had no adapter); any other length raises. Raises: - ValueError: If *header* does not end in *site*, the template has no - main message loop with a matching ``endfor``, or the wrapped - template no longer parses (e.g. an ``endfor`` inside a comment - fooled the tag scan). + ValueError: If *header* does not end in *site*, the template does not + parse, has no main message loop with a matching ``endfor``, uses + ``continue``/``break``, ``{% generation %}`` or ``{% else %}`` in + that loop, or the tag scan misread a + comment or string as the loop (the result does not parse, or the + real loop is left unwrapped). """ if not header.endswith(site): raise ValueError( f"history_adapters header {header!r} does not end in the Shadow " f"Residual substitution site {site!r}." ) + literals = literals or [header] + try: + tree = _TEMPLATE_PARSE_ENV.parse(template) + except jinja2.TemplateSyntaxError as exc: + raise ValueError( + f"The chat template does not parse ({exc}), so history_adapters " + f"cannot be added to it." + ) from exc + loop = _message_loop_node(tree, loop_var_source) + if loop is not None: + # Each would misbehave inside the per-message filter: continue/break + # leave it before it prints, dropping the message's output; HF's + # generation block takes its assistant-mask offsets from text the + # filter has not printed yet; an else block would land inside it. + if loop.else_: + unsupported = "an else block" + elif any( + _is_generation_block(node) + for body_node in loop.body + for node in [body_node, *body_node.find_all(jinja2.nodes.Scope)] + ): + unsupported = "a {% generation %} block" + elif _has_own_loop_control(loop.body): + unsupported = "continue/break" + else: + unsupported = None + if unsupported is not None: + raise ValueError( + f"The chat template's message loop uses {unsupported}, which " + f"history_adapters does not support (line {loop.lineno})." + ) open_tag = re.search( r"\{%[-+]?\s*for message in " + re.escape(loop_var_source) + r"\s*[-+]?%\}", template, @@ -659,52 +783,74 @@ def _wrap_message_loop_for_history_adapters( "history_adapters has nothing to wrap." ) + # Each inserted tag copies the whitespace control of the loop tag it sits + # next to, so the wrap neither strips nor adds whitespace when unused. + # ``+`` is copied too: under ``lstrip_blocks`` it keeps the indentation + # before the tag, which a plain inserted tag would strip. + def _control(char: str) -> str: + return char if char in "-+" else "" + + lead = _control(open_tag.group(0)[2]) + trail = _control(open_tag.group(0)[-3]) + close_lead = _control(template[close_at + 2]) # Checked once, before the loop: a list that does not line up one-to-one # with the assistant messages is a caller bug, and silently padding or # truncating it would put tokens on the wrong turns. precheck = ( - "{%- if ns.adapter_type == '" + ANCHOR_MODE_SR + "'" - " and history_adapters is defined and history_adapters is not none %}\n" - "{%- if history_adapters is string or history_adapters is mapping" - " or history_adapters is not iterable %}\n" - "{{- raise_exception('history_adapters must be a list of adapter names " - "(or none), one per assistant message') }}\n" - "{%- endif %}\n" - "{%- set _gs_expected = " + loop_var_source + " | selectattr('role', " - "'equalto', 'assistant') | list | length %}\n" - "{%- if history_adapters | length != _gs_expected %}\n" - "{{- raise_exception('history_adapters has ' ~ (history_adapters | length)" + "{%" + lead + " if ns.adapter_type == '" + ANCHOR_MODE_SR + "'" + " and history_adapters is defined and history_adapters is not none %}" + "{% if history_adapters is string or history_adapters is mapping" + " or history_adapters is not iterable %}" + "{{ raise_exception('history_adapters must be a list of adapter names " + "(or none), one per assistant message') }}" + "{% endif %}" + "{% set _gs_expected = " + loop_var_source + " | selectattr('role', " + "'equalto', 'assistant') | list | length %}" + "{% if history_adapters | length != _gs_expected %}" + "{{ raise_exception('history_adapters has ' ~ (history_adapters | length)" " ~ ' entries but the conversation has ' ~ _gs_expected ~ ' assistant " "messages; pass exactly one entry (an adapter name or none) per " - "assistant message') }}\n" - "{%- endif %}\n" - "{%- for _gs_entry in history_adapters %}\n" + "assistant message') }}" + "{% endif %}" + "{% for _gs_entry in history_adapters %}" # ``is not string`` first: ``or`` short-circuits, so an unhashable entry # never reaches the ``in`` lookup and gets this message, not a TypeError. - "{%- if _gs_entry is not none and (_gs_entry is not string" + "{% if _gs_entry is not none and (_gs_entry is not string" " or _gs_entry not in adapter_map" - " or adapter_map[_gs_entry]['type'] != '" + ANCHOR_MODE_SR + "') %}\n" - "{{- raise_exception('history_adapters names ' ~ _gs_entry ~ ', which is " - "not a Shadow Residual adapter of this checkpoint') }}\n" - "{%- endif %}\n" - "{%- endfor %}\n" - "{%- endif %}\n" + " or adapter_map[_gs_entry]['type'] != '" + ANCHOR_MODE_SR + "') %}" + "{{ raise_exception('history_adapters names ' ~ _gs_entry ~ ', which is " + "not a Shadow Residual adapter of this checkpoint') }}" + "{% endif %}" + "{% endfor %}" + "{% endif %}" ) - # Picks this message's replacement before its body renders; the default - # (header -> header) leaves the body untouched. + # Every literal gets a replace, but at most one matches the header: the + # rest would need the full literal, role marker included, inside the answer. + # With ``_gs_tok`` left at *site* every replace is a no-op. + head = header[: len(header) - len(site)] + replaces = " | ".join( + "replace(" + + _jinja_string_literal(literal) + + ", " + + _jinja_string_literal(head) + + " ~ _gs_tok ~ " + + _jinja_string_literal(literal[len(header) :]) + + ", 1)" + for literal in literals + ) + # Picks this message's token before its body renders. The last tag copies + # the loop tag's trailing whitespace control, since the body now follows it. choose = ( - "\n{%- set _gs_new = " + _jinja_string_literal(header) + " %}\n" - "{%- if message.role == 'assistant' %}\n" - "{%- if ns.adapter_type == '" + ANCHOR_MODE_SR + "'" + "{% set _gs_tok = " + _jinja_string_literal(site) + " %}" + "{% if message.role == 'assistant' %}" + "{% if ns.adapter_type == '" + ANCHOR_MODE_SR + "'" " and history_adapters is defined and history_adapters is not none" - " and history_adapters[ns.assistant_idx] is not none %}\n" - "{%- set _gs_new = " - + _jinja_string_literal(header[: len(header) - len(site)]) - + " ~ adapter_map[history_adapters[ns.assistant_idx]]['token'] %}\n" - "{%- endif %}\n" - "{%- set ns.assistant_idx = ns.assistant_idx + 1 %}\n" - "{%- endif %}\n" - "{%- filter replace(" + _jinja_string_literal(header) + ", _gs_new, 1) %}" + " and history_adapters[ns.assistant_idx] is not none %}" + "{% set _gs_tok = adapter_map[history_adapters[ns.assistant_idx]]['token'] %}" + "{% endif %}" + "{% set ns.assistant_idx = ns.assistant_idx + 1 %}" + "{% endif %}" + "{% filter " + replaces + " " + trail + "%}" ) wrapped = ( template[: open_tag.start()] @@ -712,21 +858,18 @@ def _wrap_message_loop_for_history_adapters( + template[open_tag.start() : open_tag.end()] + choose + template[open_tag.end() : close_at] - + "{%- endfilter %}\n" + + "{%" + + close_lead + + " endfilter %}" + template[close_at:] ) - # The tag scan is textual, so a template it misreads composes fine and then - # fails at render time. Parse both versions to catch that here instead; a - # template that does not parse even before the wrap (an extension tag this - # environment lacks) is left for the real renderer to judge. - env = jinja2.Environment(extensions=["jinja2.ext.loopcontrols"]) - try: - env.parse(template) - except jinja2.TemplateSyntaxError: - return wrapped + # The tag scan is textual, so it can misread a comment or a string as the + # loop. Check the parsed result instead: a misread either leaves an inserted + # tag unbalanced, or wraps text that is not the real loop, whose body then + # has no filter. try: - env.parse(wrapped) + wrapped_tree = _TEMPLATE_PARSE_ENV.parse(wrapped) except jinja2.TemplateSyntaxError as exc: raise ValueError( f"Wrapping the chat template's message loop for history_adapters " @@ -734,6 +877,15 @@ def _wrap_message_loop_for_history_adapters( f"endfor was probably mis-identified, e.g. one inside a comment or " f"a string." ) from exc + wrapped_loop = _message_loop_node(wrapped_tree, loop_var_source) + if wrapped_loop is None or not any( + isinstance(node, jinja2.nodes.FilterBlock) for node in wrapped_loop.body + ): + raise ValueError( + "Wrapping the chat template's message loop for history_adapters " + "wrapped text that is not the loop itself, e.g. a copy of it inside " + "a comment or a string." + ) return wrapped @@ -1546,6 +1698,8 @@ def configure_chat_template( modified_chat_template, sr_site, sr_replacements, + sr_header, + sr_literals, ) = _replace_anchor_in_generation_prompt( modified_chat_template, next(iter(sr_anchors)), tokenizer ) @@ -1560,8 +1714,9 @@ def configure_chat_template( modified_chat_template = _wrap_message_loop_for_history_adapters( modified_chat_template, fmt.loop_var_source, - _sr_history_header(base_chat_template, sr_site), + sr_header, sr_site, + sr_literals, ) # Skip-once wrapper for every role-open marker emission in the template. diff --git a/tests/composer/test_chat_template.py b/tests/composer/test_chat_template.py index 32744bde..61174497 100644 --- a/tests/composer/test_chat_template.py +++ b/tests/composer/test_chat_template.py @@ -1771,6 +1771,163 @@ def test_wrap_counts_plus_whitespace_control_tags(self): ) assert rendered == "a" + @pytest.mark.parametrize( + "body", + [ + "{{ message.content }}" + "{%- if message.role == 'user' %}{% continue %}{% endif %}x", + "{{ message.content }}" + "{%- if message.content == 'b' %}{% break %}{% endif %}x", + "{%- for x in [1] %}{% else %}{% continue %}{% endfor %}", + "{% generation %}{{ message.content }}{% endgeneration %}", + "{%- for x in [1] %}{% generation %}x{% endgeneration %}{% endfor %}", + "{{ message.content }}{%- else %}empty", + ], + ids=[ + "continue", + "break", + "in-inner-else", + "generation", + "nested-generation", + "for-else", + ], + ) + def test_wrap_rejects_what_the_filter_would_break(self, body): + """Inside the per-message filter, continue/break drop the message's + output, a ``{% generation %}`` block gets the wrong assistant-mask + offsets, and an else block lands inside the filter, so compose refuses + each.""" + template = ( + "{%- set ns = namespace(adapter_type='', assistant_idx=0) %}" + "{%- for message in messages %}" + body + "{%- endfor %}END" + ) + with pytest.raises(ValueError, match="does not support"): + _wrap_message_loop_for_history_adapters(template, "messages", "", "") + + def test_wrap_rejects_a_template_that_does_not_parse(self): + template = "{%- for message in messages %}{{ message.content }}" + with pytest.raises(ValueError, match="does not parse"): + _wrap_message_loop_for_history_adapters(template, "messages", "", "") + + def test_wrap_rejects_a_loop_copy_inside_a_comment(self): + """A whole loop inside a comment fools the tag scan but still parses + once wrapped; the real loop is left unwrapped, so compose refuses it.""" + template = ( + "{%- set ns = namespace(adapter_type='', assistant_idx=0) %}" + "{# e.g. {% for message in messages %}x{% endfor %} #}" + "{%- for message in messages %}{{ message.content }}{%- endfor %}" + ) + with pytest.raises(ValueError, match="not the loop itself"): + _wrap_message_loop_for_history_adapters(template, "messages", "", "") + + def test_wrap_parses_generation_blocks_outside_the_loop(self): + """HF's ``{% generation %}`` tag parses here, so a template that uses it + outside the message loop is still validated and wrapped.""" + template = ( + "{%- set ns = namespace(adapter_type='', assistant_idx=0) %}" + "{%- for message in messages %}{{ message.content }}{%- endfor %}" + "{% generation %}G{% endgeneration %}" + ) + wrapped = _wrap_message_loop_for_history_adapters( + template, "messages", "", "" + ) + assert "{% filter " in wrapped + + @pytest.mark.parametrize("enable_thinking", [False, True]) + def test_chatml_answer_without_header_is_left_unchanged(self, enable_thinking): + """A 4.2 answer stored mid-thought (no leading ````) renders no + header, so its turn gets no control token and nothing raises: the + render matches the one without ``history_adapters``.""" + messages = [ + {"role": "user", "content": "Q1"}, + {"role": "assistant", "content": "R1\n\nA1"}, + {"role": "user", "content": "Q2"}, + ] + kwargs = {"enable_thinking": enable_thinking, "adapter_name": "sr1"} + tokenizer = self._chatml() + with_history = self._render( + tokenizer, messages, history_adapters=["sr2"], **kwargs + ) + without = self._render(tokenizer, messages, **kwargs) + assert with_history == without + assert "<|sr2|>" not in with_history + + def test_wrap_allows_loop_controls_in_a_nested_loop(self): + """A loop control in an inner loop stays inside that loop.""" + template = ( + "{%- set ns = namespace(adapter_type='', assistant_idx=0) %}" + "{%- for message in messages %}" + "{%- for x in [1, 2] %}{% if x == 2 %}{% break %}{% endif %}" + "{{ message.content }}{% endfor %}" + "{%- endfor %}" + ) + wrapped = _wrap_message_loop_for_history_adapters( + template, "messages", "", "" + ) + env = Environment(extensions=["jinja2.ext.loopcontrols"]) + messages = [{"role": "user", "content": "a"}, {"role": "user", "content": "b"}] + assert env.from_string(wrapped).render(messages=messages) == "ab" + + @pytest.mark.parametrize( + "template", + [ + "X {% for message in messages %}{{ message.content }} {% endfor %}END", + "X {%- for message in messages %}{{ message.content }}{% endfor %}END", + "X {% for message in messages -%} {{ message.content }}{%- endfor %} END", + "X {%- for message in messages -%}{{ message.content }}{%- endfor -%} END", + "X\n{% for message in messages %}\n{{ message.content }}\n{% endfor %}\nEND", + "X\n {%+ for message in messages %}{{ message.content }}\n {%+ endfor %}END", + "X\n{% for message in messages +%}\n{{ message.content }}{% endfor %}END", + ], + ids=[ + "no-dash", + "lead-dash", + "trail-dash", + "all-dash", + "newlines", + "plus", + "trail-plus", + ], + ) + @pytest.mark.parametrize("hf_flags", [False, True], ids=["plain", "hf-flags"]) + def test_wrap_preserves_whitespace_when_unused(self, template, hf_flags): + """The inserted tags copy the loop tags' whitespace control, so without + ``history_adapters`` the render matches the original for any style.""" + template = ( + "{%- set ns = namespace(adapter_type='', assistant_idx=0) -%}" + template + ) + wrapped = _wrap_message_loop_for_history_adapters( + template, "messages", "", "" + ) + env = Environment(trim_blocks=hf_flags, lstrip_blocks=hf_flags) + messages = [{"role": "user", "content": "a"}, {"role": "user", "content": "b"}] + assert env.from_string(wrapped).render(messages=messages) == env.from_string( + template + ).render(messages=messages) + + @pytest.mark.parametrize("enable_thinking", [False, True]) + def test_chatml_answer_starting_with_the_site_is_not_rewritten( + self, enable_thinking + ): + """An answer that *starts* with ```` completes the bare header + text, so only matching the full generation-prompt literal leaves it + alone. The restored thinking-on form still gets its token + (``test_chatml_thinking_prefix_holds_with_think_restored``).""" + messages = [ + {"role": "user", "content": "Q1"}, + {"role": "assistant", "content": " is the tag"}, + {"role": "user", "content": "Q2"}, + ] + result = self._render( + self._chatml(), + messages, + enable_thinking=enable_thinking, + adapter_name="sr1", + history_adapters=["sr2"], + ) + assert "<|im_start|>assistant\n is the tag<|im_end|>" in result + assert "<|sr2|>" not in result + class TestShadowResidualHistoryAdaptersTokenIds: """The prefix relation on token ids, with the real Granite tokenizers. diff --git a/uv.lock b/uv.lock index 8e5f16b9..4f85b46e 100644 --- a/uv.lock +++ b/uv.lock @@ -1578,6 +1578,7 @@ name = "granite-switch" version = "0.1.1" source = { editable = "." } dependencies = [ + { name = "jinja2" }, { name = "torch", version = "2.11.0", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'group-14-granite-switch-dev' or extra == 'group-14-granite-switch-test' or extra == 'group-14-granite-switch-vllm26' or (extra == 'group-14-granite-switch-dev-latest' and extra == 'group-14-granite-switch-dev-vllm27') or (extra == 'group-14-granite-switch-dev-latest' and extra == 'group-14-granite-switch-vllm27') or (extra == 'group-14-granite-switch-dev-vllm27' and extra == 'group-14-granite-switch-dev-vllm30') or (extra == 'group-14-granite-switch-dev-vllm27' and extra == 'group-14-granite-switch-vllm30') or (extra == 'group-14-granite-switch-dev-vllm30' and extra == 'group-14-granite-switch-vllm27') or (extra == 'group-14-granite-switch-vllm27' and extra == 'group-14-granite-switch-vllm30')" }, { name = "torch", version = "2.13.0", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'group-14-granite-switch-dev-latest' or extra == 'group-14-granite-switch-dev-vllm27' or extra == 'group-14-granite-switch-dev-vllm30' or (extra == 'group-14-granite-switch-dev' and extra == 'group-14-granite-switch-vllm27') or (extra == 'group-14-granite-switch-dev' and extra == 'group-14-granite-switch-vllm30') or (extra == 'group-14-granite-switch-test' and extra == 'group-14-granite-switch-vllm27') or (extra == 'group-14-granite-switch-test' and extra == 'group-14-granite-switch-vllm30') or (extra == 'group-14-granite-switch-vllm26' and extra == 'group-14-granite-switch-vllm27') or (extra == 'group-14-granite-switch-vllm26' and extra == 'group-14-granite-switch-vllm30') or (extra != 'group-14-granite-switch-dev' and extra != 'group-14-granite-switch-test' and extra != 'group-14-granite-switch-vllm26')" }, { name = "transformers" }, @@ -1682,6 +1683,7 @@ requires-dist = [ { name = "huggingface-hub", marker = "extra == 'build'" }, { name = "huggingface-hub", marker = "extra == 'compose'" }, { name = "ipython", marker = "extra == 'tutorials'", specifier = ">=8.10.0" }, + { name = "jinja2", specifier = ">=3.1" }, { name = "mellea", marker = "extra == 'tutorials'", specifier = "==0.6.0" }, { name = "python-dotenv", marker = "extra == 'tutorials'", specifier = ">=1.0.0" }, { name = "pyyaml", marker = "extra == 'build'" },