From a5fbe0ff97019e6973379f8fb7935cf06d3bab67 Mon Sep 17 00:00:00 2001 From: noaa Date: Sun, 25 Oct 2026 18:30:08 +0200 Subject: [PATCH] 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 29be363b..5b3815f4 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 @@ -552,6 +554,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]], @@ -1165,7 +1346,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 = ( @@ -1318,6 +1501,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 @@ -1412,4 +1601,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