Skip to content

Normalize ChatML content once, instead of per render site #159

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ItzikVa merged 1 commit into from
Nov 2, 2026
Merged

Normalize ChatML content once, instead of per render site #159

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
233 changes: 118 additions & 115 deletions src/granite_switch/composer/tokenizer_setup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -725,55 +725,85 @@ 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
``<tool_response>``, 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
upgrade it in place, rather than stack a second competing block.
"""
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"
)


Expand All @@ -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"
Expand All @@ -815,76 +857,39 @@ def _inject_chatml_system_message_flatten(template: str) -> str:
)


# The ChatML tool branch emits message.content directly inside a
# <tool_response> 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 = (
" {{- '<tool_response>\\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
``<tool_response>`` 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 = (
" {{- '<tool_response>\\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:
Expand All @@ -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`` /
Expand All @@ -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
Expand Down Expand Up @@ -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:
Expand Down
8 changes: 4 additions & 4 deletions tests/composer/test_chat_template.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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)


Expand Down
Loading