diff --git a/src/granite_switch/composer/tokenizer_setup.py b/src/granite_switch/composer/tokenizer_setup.py index bc1bc345..42fdacba 100644 --- a/src/granite_switch/composer/tokenizer_setup.py +++ b/src/granite_switch/composer/tokenizer_setup.py @@ -725,23 +725,47 @@ def _inject_audio_granite_format(template: str, marker: str) -> str: return template.replace(_GRANITE_TEXT_PART_BRANCH, new, 1) -# Marker that identifies the ChatML content-flatten block (so the audio pass can -# find and upgrade a text-only block rather than inject a competing second one). -_CHATML_FLATTEN_SENTINEL = "{#- gs:chatml-content-flatten -#}" +# Marker that identifies the ChatML content-normalize block (so the audio pass +# can find and upgrade a text-only block rather than inject a competing second +# one). +_CHATML_FLATTEN_SENTINEL = "{#- gs:chatml-content-normalize -#}" +# The ChatML message loop. The normalize pass is injected immediately before it, +# so every role-dispatch branch inside the loop sees a plain-string `content`. +_CHATML_MESSAGE_LOOP = "{%- for message in loop_messages %}" -def _chatml_content_flatten_block(marker: str | None) -> str: - """The Jinja block that rebuilds ChatML ``content`` from a parts list. - ChatML's user/system branch does ``{%- set content = message.content | - string %}``, so when vLLM feeds OpenAI-style structured content (a parts - *list*) that statement renders the list's Python repr — ``[{'type': 'text', - ...}]`` — instead of the text. This block flattens the parts back into a - string: text parts are newline-joined in order; everything else is dropped. +def _chatml_content_flatten_block(marker: str | None) -> str: + """The Jinja block that normalizes every ChatML message's ``content``. + + ChatML reads ``message.content`` at *seven* sites with different shapes of + assumption: the user/system branch stringifies it (rendering a parts list as + its Python repr ``[{'type': 'text', ...}]``), the assistant branch gates on + ``content is string`` (silently DROPPING a parts list's text when the turn + also carries tool calls), the tool branch emits it straight into + ````, and the catch-all ``else`` concatenates it to a string + (``TypeError`` on a list). Under vLLM's ``openai`` content format *every* + message arrives as a parts list, including assistant history, because vLLM + wraps plain strings too. + + Patching those sites individually is only ever as complete as the list of + sites someone remembered — three were missed on the first pass. So instead + this rebuilds ``loop_messages`` with flattened content ONCE, before the loop, + which is what the granite_format (4.1) family already does natively. Text + parts are newline-joined in order; everything else is dropped. Role-agnostic + by construction, so a role added to a future base template is covered + without touching this code. + + ``dict(_m, content=...)`` makes a flattened *copy* — the messages are not + mutated, which also keeps this working under the immutable Jinja sandbox + ``transformers`` renders with. When *marker* is given (audio enabled) an extra branch emits it for audio parts, interleaved in part order so the ASR transcript lands where the clip - sat. When *marker* is ``None`` the block is text-only. + sat. When *marker* is ``None`` the block is text-only. Note that normalizing + centrally widens audio to *any* role, where the previous per-site injection + reached user/system only; emitting a marker too widely is the safe direction, + since a DROPPED audio part is what crashes vLLM's prompt replacement. The sentinel comment lets :func:`configure_audio_chat_template` recognize a text-only block already injected by :func:`configure_chat_template` and @@ -749,31 +773,37 @@ def _chatml_content_flatten_block(marker: str | None) -> str: """ audio_branch = ( ( - " {%- if entry.type is defined and 'audio' in entry.type %}\n" - " {%- set _c.val = _c.val + '" + marker + "' %}\n" - " {%- elif entry.type is defined and entry.type == 'text' %}\n" + " {%- if _e.type is defined and 'audio' in _e.type %}\n" + " {%- set _c.val = _c.val + '" + marker + "' %}\n" + " {%- elif _e.type is defined and _e.type == 'text' %}\n" ) if marker is not None - else " {%- if entry.type is defined and entry.type == 'text' %}\n" + else " {%- if _e.type is defined and _e.type == 'text' %}\n" ) return ( - "\n" - " " + _CHATML_FLATTEN_SENTINEL + "\n" - " {#- ChatML stringifies message.content, so rebuild it from the\n" - " parts when it is a list — otherwise a multimodal / OpenAI-format\n" - " turn renders as a Python repr. -#}\n" - " {%- if message.content is not string and message.content is iterable %}\n" - " {%- set _c = namespace(val='') %}\n" - " {%- for entry in message.content %}\n" + _CHATML_FLATTEN_SENTINEL + "\n" + "{#- Normalize every message's content to a plain string before the loop:\n" + " under vLLM's 'openai' content format each content arrives as a parts\n" + " list, and the branches below variously repr it, drop it, or raise. -#}\n" + "{%- set _norm = namespace(msgs=[]) %}\n" + "{%- for _m in loop_messages %}\n" + " {%- if _m.content is not string and _m.content is iterable " + "and _m.content is not mapping %}\n" + " {%- set _c = namespace(val='') %}\n" + " {%- for _e in _m.content %}\n" + audio_branch - + " {%- if _c.val != '' %}\n" - " {%- set _c.val = _c.val + '\\n' %}\n" - " {%- endif %}\n" - " {%- set _c.val = _c.val + entry.text %}\n" + + " {%- if _c.val != '' %}\n" + " {%- set _c.val = _c.val + '\\n' %}\n" " {%- endif %}\n" - " {%- endfor %}\n" - " {%- set content = _c.val %}\n" - " {%- endif %}" + " {%- set _c.val = _c.val + _e.text %}\n" + " {%- endif %}\n" + " {%- endfor %}\n" + " {%- set _norm.msgs = _norm.msgs + [dict(_m, content=_c.val)] %}\n" + " {%- else %}\n" + " {%- set _norm.msgs = _norm.msgs + [_m] %}\n" + " {%- endif %}\n" + "{%- endfor %}\n" + "{%- set loop_messages = _norm.msgs %}\n" ) @@ -791,10 +821,22 @@ def _inject_chatml_system_message_flatten(template: str) -> str: that, when that content is a non-string iterable, rebuilds it as the newline- joined text parts. Text-only: an audio part in the *leading system* message is still dropped (matching the per-message audio scope and the fact that no - caller places a clip in the system turn). No-op if the anchor is absent. + caller places a clip in the system turn). + + This site is NOT covered by the pre-loop normalize pass: the preamble reads + ``messages[0]["content"]`` directly, before ``loop_messages`` exists. + + Raises: + ValueError: if the anchor is absent. A silent no-op here would compose a + checkpoint that renders a parts-list system turn as a ``TypeError``, + with nothing failing until a prompt is served. """ if _CHATML_SYSTEM_MESSAGE_SET not in template: - return template + raise ValueError( + "Could not find the ChatML leading-system-message assignment " + f"({_CHATML_SYSTEM_MESSAGE_SET!r}) to inject content flattening; " + "the base chat template may have changed." + ) block = ( "\n" "{%- if system_message is not string and system_message is iterable %}\n" @@ -815,76 +857,39 @@ def _inject_chatml_system_message_flatten(template: str) -> str: ) -# The ChatML tool branch emits message.content directly inside a -# wrapper. Tool content is external-data-in (like user content), -# so vLLM wraps it into a parts list in 'openai' mode too; without a flatten the -# tool response renders as the Python list-repr. Anchored on the emission in its -# tool-response context so only the tool branch is rewritten. -_CHATML_TOOL_RESPONSE_EMIT = ( - " {{- '\\n' }}\n {{- message.content }}\n" -) - - -def _inject_chatml_tool_content_flatten(template: str) -> str: - """Flatten a parts-list ``tool``-role message in the ChatML tool branch. - - Replaces the direct ``{{- message.content }}`` emission (inside the - ```` wrapper) with a flatten-then-emit: a non-string iterable - is rebuilt as its newline-joined text parts, a string passes through. Text - parts only — a tool response carrying audio is not a shape any caller sends. - No-op if the anchor is absent (template shape changed). - """ - if _CHATML_TOOL_RESPONSE_EMIT not in template: - return template - replacement = ( - " {{- '\\n' }}\n" - " {%- if message.content is not string and message.content is iterable %}\n" - " {%- set _tool = namespace(val='') %}\n" - " {%- for entry in message.content %}\n" - " {%- if entry.type is defined and entry.type == 'text' %}\n" - " {%- if _tool.val != '' %}\n" - " {%- set _tool.val = _tool.val + '\\n' %}\n" - " {%- endif %}\n" - " {%- set _tool.val = _tool.val + entry.text %}\n" - " {%- endif %}\n" - " {%- endfor %}\n" - " {{- _tool.val }}\n" - " {%- else %}\n" - " {{- message.content }}\n" - " {%- endif %}\n" - ) - return template.replace(_CHATML_TOOL_RESPONSE_EMIT, replacement, 1) - - def _inject_chatml_content_flatten(template: str, marker: str | None) -> str: - """Insert (or upgrade) the ChatML content-flatten block. + """Insert (or upgrade) the ChatML content-normalize block. + + Inserted immediately *before* the message loop, so every role-dispatch + branch inside the loop reads an already-flattened string. ALoRA Pass 1 is + injected at the same anchor and must still see the ORIGINAL content shape — + it scans for the invocation text in both string and parts form — so this + block goes in first and Pass 1 lands before it. - Appended immediately *after* ``{%- set content = message.content | string - %}`` rather than replacing it, so that statement stays intact as - :func:`configure_chat_template`'s ALoRA Pass 2 anchor. + If a text-only block (marked by :data:`_CHATML_FLATTEN_SENTINEL`) is already + present — injected by :func:`configure_chat_template` — and *marker* is + given, the existing block is swapped for the audio-aware variant instead of + stacking a second block that would re-normalize ``loop_messages``. - If a text-only flatten block (marked by :data:`_CHATML_FLATTEN_SENTINEL`) is - already present — injected by :func:`configure_chat_template` — and *marker* - is given, the existing block is swapped for the audio-aware variant instead - of stacking a second block that would re-assign ``content``. + Raises: + ValueError: if the message-loop anchor is absent. """ - if _CHATML_CONTENT_SET not in template: - raise ValueError( - "Could not find the ChatML per-message content assignment " - f"({_CHATML_CONTENT_SET!r}) to inject content flattening; the base " - "chat template may have changed." - ) desired = _chatml_content_flatten_block(marker) if _CHATML_FLATTEN_SENTINEL in template: - # A block is already installed; replace it wholesale (upgrade text-only - # -> audio, or re-inject identically). Bound the slice by the block's - # own structure: from the sentinel's leading newanchor to its {%- endif %}. + # A block is already installed; swap it wholesale (upgrade text-only -> + # audio, or re-inject identically). existing = _chatml_content_flatten_block(None) if existing in template: return template.replace(existing, desired, 1) # Already the audio variant (or an unexpected shape): leave as-is. return template - return template.replace(_CHATML_CONTENT_SET, _CHATML_CONTENT_SET + desired, 1) + if _CHATML_MESSAGE_LOOP not in template: + raise ValueError( + "Could not find the ChatML message loop " + f"({_CHATML_MESSAGE_LOOP!r}) to inject content normalization; the " + "base chat template may have changed." + ) + return template.replace(_CHATML_MESSAGE_LOOP, desired + _CHATML_MESSAGE_LOOP, 1) def configure_audio_chat_template(tokenizer, marker: str = "<|audio|>") -> None: @@ -900,10 +905,10 @@ def configure_audio_chat_template(tokenizer, marker: str = "<|audio|>") -> None: * **granite_format** (3.x / 4.0 / 4.1) — one ``elif`` added to the existing content-part loop. See :func:`_inject_audio_granite_format`. - * **chatml** (Granite 4.2) — no content-part loop exists, so a flattening - block rebuilds ``content`` from the parts. A text-only flatten is already - installed by :func:`configure_chat_template`; this upgrades it in place to - also emit the marker for audio parts. See + * **chatml** (Granite 4.2) — no content-part loop exists, so a normalize + pass before the message loop rebuilds each ``content`` from its parts. A + text-only pass is already installed by :func:`configure_chat_template`; + this upgrades it in place to also emit the marker for audio parts. See :func:`_inject_chatml_content_flatten`. Matching on ``'audio' in entry.type`` covers ``audio`` / ``input_audio`` / @@ -912,9 +917,13 @@ def configure_audio_chat_template(tokenizer, marker: str = "<|audio|>") -> None: Call *after* :func:`configure_chat_template`, gated on audio being enabled; the ChatML path depends on that order. - Scope is the per-message loop. An audio part in the leading system message is - dropped by both families — matching the pre-existing granite_format - behaviour, and not a shape any caller sends. + An audio part in the leading *system* message is dropped by both families — + matching the pre-existing granite_format behaviour, and not a shape any + caller sends. Otherwise ChatML's scope is now every role, because the + normalize pass it upgrades is role-agnostic (granite_format remains + per-message-loop). Widening the marker is the safe direction: a *dropped* + audio part is what fails vLLM's prompt replacement, an extra marker in a + role no caller uses is inert. Raises: ValueError: if the template family is unrecognized or its injection @@ -1388,29 +1397,23 @@ def configure_chat_template( + alora_pass2 + modified_chat_template[insertion_point:] ) - # Flatten an OpenAI-style content parts list into a string. vLLM's - # content-format auto-detector resolves to 'openai' whenever a template - # iterates message.content (which the ALoRA Pass 1 scan above does), and - # then feeds content as a parts list. ChatML's base template stringifies - # content directly, so without this the list renders as its Python repr. - # Inserted AFTER Pass 2 so that at render time it runs BEFORE Pass 2 — - # both anchor on the content assignment, and Pass 2 rsplits `content` - # and so needs a string. Text-only here; the audio pass upgrades it. + # Normalize every message's content to a plain string BEFORE the loop. + # vLLM's content-format auto-detector resolves to 'openai' whenever a + # template iterates message.content (which the ALoRA Pass 1 scan above + # does), and then feeds every content as a parts list — assistant history + # included, since it wraps plain strings too. ChatML's branches variously + # repr that list, silently drop it, or raise on it, so normalizing once + # up front is what makes all of them correct; patching each branch is + # only ever as complete as the list of branches someone remembered. + # Text-only here; the audio pass upgrades this same block. modified_chat_template = _inject_chatml_content_flatten( modified_chat_template, marker=None ) - # The leading system message is read into a scalar in the preamble before - # the loop; flatten a parts list there too or the render crashes on a - # str+list concat before the per-message block runs. + # Not covered by the pass above: the preamble reads messages[0]["content"] + # directly, before loop_messages exists, and concatenates it to a string. modified_chat_template = _inject_chatml_system_message_flatten( modified_chat_template ) - # The tool branch emits message.content directly; flatten a parts-list - # tool response the same way (tool content is external-data-in, like user - # content, so vLLM wraps it under 'openai' too). - modified_chat_template = _inject_chatml_tool_content_flatten( - modified_chat_template - ) else: match = re.search(fmt.user_role_anchor_re, modified_chat_template) if match: diff --git a/tests/composer/test_chat_template.py b/tests/composer/test_chat_template.py index efdc5e6e..e07400d3 100644 --- a/tests/composer/test_chat_template.py +++ b/tests/composer/test_chat_template.py @@ -1099,9 +1099,9 @@ def test_part_order_and_multiple_clips_preserved(self): def test_pass2_anchor_survives_injection(self): """The ALoRA Pass 2 anchor must remain findable after audio injection. - The flattening block is appended after the content assignment rather - than replacing it precisely so configure_chat_template can still anchor - on it, in either call order. + The normalize pass is inserted before the message loop rather than at the + content assignment, so that assignment stays intact for + configure_chat_template to anchor Pass 2 on, in either call order. """ tokenizer = _make_chatml_audio_tokenizer() configure_audio_chat_template(tokenizer) @@ -1113,7 +1113,7 @@ def test_missing_chatml_anchor_raises(self): tokenizer = SimpleNamespace( chat_template="{{- '<|im_start|>user\\n' }}{{ messages }}" ) - with pytest.raises(ValueError, match="ChatML per-message content assignment"): + with pytest.raises(ValueError, match="ChatML message loop"): configure_audio_chat_template(tokenizer) diff --git a/tests/composer/test_chat_template_real_tokenizer.py b/tests/composer/test_chat_template_real_tokenizer.py index b6eeb717..41f49ae8 100644 --- a/tests/composer/test_chat_template_real_tokenizer.py +++ b/tests/composer/test_chat_template_real_tokenizer.py @@ -18,12 +18,20 @@ * multiple text parts (join behaviour), * parts-list across multiple turns and in a system message, * parts-list while an adapter is active (the flatten must run before the ALoRA - Pass 2 ``rsplit``). + Pass 2 ``rsplit``), +* parts-list on every role vLLM wraps — assistant history, assistant-with-tool- + calls, tool, and a role outside the known branches. Break criterion: the rendered prompt must never contain the Python list-repr marker ``[{'type'`` — that means a parts list leaked into the prompt instead of being flattened. +**That criterion alone is not sufficient**, and relying on it is how the +assistant-with-tool-calls case slipped through a green 2886-test GPU run: the +ChatML assistant branch gates on ``content is string``, so a parts list there is +DROPPED rather than repr'd. An absence cannot be detected by looking for a +marker. Tests covering a role must therefore also assert the text is PRESENT. + A separate class pins the *trigger*: that the composed template resolves to the vLLM content format the template can actually honour, using a local reimplementation of vLLM's AST-based detector (and the real vLLM detector when @@ -273,6 +281,131 @@ def test_tool_role_string_unchanged( assert "sunny, 72F" in out assert _REPR_MARKER not in out + def test_assistant_history_parts_flattened( + self, model_id, family, adapter_dir, name, atype + ): + """An assistant turn sent as a parts list must flatten. + + Under 'openai' vLLM wraps *every* message's content, assistant history + included — it is vLLM that produces this shape, not the client. Patching + only the roles thought of as "incoming" (user/system/tool) left each + earlier assistant turn rendering as [{'type': 'text', 'text': 'a1'}]. + """ + tok = _configure(model_id, adapter_dir, name, atype) + out = _render( + tok, + [ + {"role": "user", "content": _text_part("q1")}, + {"role": "assistant", "content": _text_part("a1-assistant-reply")}, + {"role": "user", "content": _text_part("q2")}, + ], + ) + assert "a1-assistant-reply" in out, ( + f"assistant text missing ({family}): {out!r}" + ) + assert _REPR_MARKER not in out, f"list repr leaked ({family}): {out!r}" + + def test_assistant_text_kept_with_tool_calls( + self, model_id, family, adapter_dir, name, atype + ): + """A parts-list assistant turn carrying tool calls must keep its text. + + This asserts PRESENCE, not the absence of the list-repr, which is the + whole point: the ChatML assistant branch gates on ``content is string``, + so a parts list fails the check and the text is DROPPED rather than + repr'd. A no-list-repr assertion reports this case as passing while the + text silently vanishes from the history — so only a positive assertion + catches it. + """ + tok = _configure(model_id, adapter_dir, name, atype) + out = _render( + tok, + [ + {"role": "user", "content": _text_part("what is the weather?")}, + { + "role": "assistant", + "content": _text_part("Let me check the weather"), + "tool_calls": [ + { + "type": "function", + "function": { + "name": "get_weather", + "arguments": {"city": "Paris"}, + }, + } + ], + }, + {"role": "tool", "content": _text_part("sunny, 72F")}, + {"role": "user", "content": _text_part("thanks")}, + ], + ) + assert "Let me check the weather" in out, ( + f"assistant text dropped next to tool_calls ({family}): {out!r}" + ) + assert "get_weather" in out, f"tool call missing ({family}): {out!r}" + assert _REPR_MARKER not in out, f"list repr leaked ({family}): {out!r}" + + def test_unknown_role_parts_renders_like_string( + self, model_id, family, adapter_dir, name, atype + ): + """A role outside the known branches must render, not raise. + + ChatML's catch-all ``else`` concatenates message.content straight onto a + string, so a parts list there raised TypeError before the normalize pass + ran. Base 4.2 only escapes this because vLLM detects 'string' for it and + flattens upstream — composing removes that protection, so this is a + composed-only regression. + + Asserted as equality against the equivalent plain-string render rather + than as text presence, because what a template DOES with an unknown role + differs by family: 4.2 has a catch-all that emits it, 4.1 has none and + drops it (true of a plain string there too, so not a bug). Equality is + the family-agnostic claim and the stronger one — flattening must make a + parts list behave exactly like the string it carries, whatever that is. + """ + tok = _configure(model_id, adapter_dir, name, atype) + as_parts = _render( + tok, + [ + {"role": "user", "content": _text_part("summarize")}, + {"role": "document", "content": _text_part("doc-body-text")}, + ], + ) + as_string = _render( + tok, + [ + {"role": "user", "content": "summarize"}, + {"role": "document", "content": "doc-body-text"}, + ], + ) + assert as_parts == as_string, ( + f"parts-list render diverged from the string render ({family}):\n" + f"parts={as_parts!r}\nstring={as_string!r}" + ) + assert _REPR_MARKER not in as_parts, ( + f"list repr leaked ({family}): {as_parts!r}" + ) + + +def test_chatml_injections_raise_on_missing_anchor(): + """Every ChatML injection must fail loudly when its anchor is absent. + + Two of the three originally returned the template unchanged, so a base-template + revision that moved an anchor would compose GREEN and then render wrongly at + serve time. A composer cannot no-op a correctness fix: the checkpoint is + written and published before anyone renders a prompt. + """ + from granite_switch.composer.tokenizer_setup import ( + _inject_chatml_content_flatten, + _inject_chatml_system_message_flatten, + ) + + anchorless = "{{- '<|im_start|>user\\n' }}{{ messages }}" + with pytest.raises(ValueError, match="ChatML message loop"): + _inject_chatml_content_flatten(anchorless, marker=None) + with pytest.raises(ValueError, match="leading-system-message"): + _inject_chatml_system_message_flatten(anchorless) + # ──────────────────────────────────────────────────────────────────── # Detector-level: the composed template must resolve to a content format