diff --git a/docs/SR_ARCHITECTURE.md b/docs/SR_ARCHITECTURE.md index 0e3c3b92..1e3bec74 100644 --- a/docs/SR_ARCHITECTURE.md +++ b/docs/SR_ARCHITECTURE.md @@ -144,6 +144,50 @@ 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. 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: + +* 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/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 a127835d..476b020a 100644 --- a/src/granite_switch/composer/tokenizer_setup.py +++ b/src/granite_switch/composer/tokenizer_setup.py @@ -9,6 +9,10 @@ import re from dataclasses import dataclass +import jinja2 +import jinja2.ext +import jinja2.nodes + from .adapter_loader import load_adapter_config # Name of the optional return-to-base control token (``<|base_reset|>``). It is @@ -463,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. @@ -482,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 @@ -553,7 +560,333 @@ 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. +#: Accepts every whitespace-control spelling: ``{%``, ``{%-`` and ``{%+``. +_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(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. + + 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 branches share no prefix that contains *site*. + """ + 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." + ) + 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, + literals: list[str] | None = None, +) -> 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 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). + + ``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 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, + ) + 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." + ) + + # 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 = ( + "{%" + 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') }}" + "{% 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" + " or _gs_entry not in adapter_map" + " 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 %}" + ) + # 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 = ( + "{% 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 %}" + "{% 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()] + + precheck + + template[open_tag.start() : open_tag.end()] + + choose + + template[open_tag.end() : close_at] + + "{%" + + close_lead + + " endfilter %}" + + template[close_at:] + ) + + # 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: + 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 " + 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 + 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 def add_control_tokens( @@ -1214,7 +1547,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 = ( @@ -1363,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 ) @@ -1374,6 +1711,13 @@ 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_header, + sr_site, + sr_literals, + ) # 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 +1814,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..61174497 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,585 @@ 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" + + @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. + + 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 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'" },