renderers 0.1.8.dev2__tar.gz → 0.1.9.dev0__tar.gz

{renderers-0.1.8.dev2 → renderers-0.1.9.dev0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: renderers
-Version: 0.1.8.dev2
+Version: 0.1.9.dev0
 Summary: Chat template renderers — deterministic message-to-token conversion for LLM training
 License-Expression: Apache-2.0
 License-File: LICENSE

{renderers-0.1.8.dev2 → renderers-0.1.9.dev0}/renderers/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '0.1.8.dev2'
-__version_tuple__ = version_tuple = (0, 1, 8, 'dev2')
+__version__ = version = '0.1.9.dev0'
+__version_tuple__ = version_tuple = (0, 1, 9, 'dev0')
 
 __commit_id__ = commit_id = None

{renderers-0.1.8.dev2 → renderers-0.1.9.dev0}/renderers/base.py

@@ -177,8 +177,162 @@ class RenderedTokens:
     token_ids: list[int] = field(default_factory=list)
     message_indices: list[int] = field(default_factory=list)
     sampled_mask: list[bool] = field(default_factory=list)
+    message_roles: list[str] = field(default_factory=list)
     multi_modal_data: "MultiModalData | None" = None
 
+    def tokens_per_message(
+        self, n_messages: int | None = None, *, sampled_only: bool = False
+    ) -> list[int]:
+        """Count rendered tokens attributed to each caller-relative message.
+
+        ``out[i]`` is the number of tokens with ``message_indices[k] == i``,
+        i.e. tokens the renderer attributed to ``messages[i]``. This
+        includes template scaffolding the renderer wraps around the
+        message — the ``<|im_start|>role\\n`` opener, the closing
+        ``<|im_end|>\\n``, etc. — because those are the renderer's own
+        attribution decision and are preserved verbatim here. Tokens with
+        ``message_indices[k] == -1`` (scaffolding outside any single
+        message, e.g. the trailing generation prompt) are not counted.
+
+        With ``sampled_only=True``, counts only tokens the model would
+        have emitted at inference (``sampled_mask[k] is True``). For
+        example, length-penalty signals in RL: the template wraps each
+        assistant turn in scaffolding tokens (e.g. ``<|im_start|>assistant\\n``,
+        ``<|im_end|>\\n``) that are constant-size and not chosen by the
+        model, so they shouldn't enter the penalty. For roles the model
+        never samples (``user``, ``tool``, ``system``), the
+        ``sampled_only`` count is zero by construction. Renderers that
+        don't populate ``sampled_mask`` (``DefaultRenderer`` — the Jinja
+        template is opaque) return all zeros under ``sampled_only=True``.
+
+        ``n_messages`` defaults to ``len(self.message_roles)``, which
+        every Renderer populates with the caller-relative message list
+        (caller's ``messages`` for ``render()``; ``new_messages`` for
+        ``bridge_to_next_turn()``). Pass it explicitly only to truncate
+        — indices outside ``[0, n_messages)`` are ignored, so passing a
+        smaller value won't raise; it just drops the tail. Values larger
+        than ``len(self.message_roles)`` are clamped, so the returned
+        list never claims more messages than the renderer attributed.
+
+        Works on results from both :meth:`Renderer.render` and
+        :meth:`Renderer.bridge_to_next_turn`. For a bridge result the
+        indices are relative to the new messages the bridge added, not
+        the full conversation history; the prior portion is uniformly
+        ``-1`` (and ``sampled_mask`` uniformly ``False``), so it
+        contributes nothing to either count.
+        """
+        if n_messages is None:
+            n_messages = len(self.message_roles)
+        else:
+            n_messages = min(n_messages, len(self.message_roles))
+        out = [0] * n_messages
+        if sampled_only:
+            if len(self.sampled_mask) != len(self.token_ids):
+                return out
+            for idx, sampled in zip(self.message_indices, self.sampled_mask):
+                if sampled and 0 <= idx < n_messages:
+                    out[idx] += 1
+        else:
+            for idx in self.message_indices:
+                if 0 <= idx < n_messages:
+                    out[idx] += 1
+        return out
+
+    def message_token_spans(self) -> list[tuple[int, int] | None]:
+        """Per-message ``(start, end)`` slices into :attr:`token_ids`.
+
+        ``out[i]`` is the half-open span ``[start, end)`` such that
+        ``token_ids[start:end]`` are the tokens attributed to
+        ``messages[i]`` (or ``new_messages[i]`` for a bridge result).
+        Messages that contributed no tokens get ``None``. Renderer
+        scaffolding outside any message (``message_indices[k] == -1``)
+        is not represented.
+
+        Hand-coded renderers emit each message's tokens contiguously,
+        so the span is well-defined. The implementation tolerates
+        non-contiguous attribution by returning the outer span
+        ``(first_k, last_k + 1)``; if you suspect interleaving, slice
+        ``message_indices`` yourself to verify.
+
+        Returns ``len(self.message_roles)`` entries when ``message_roles``
+        is populated. Otherwise infers the count from
+        ``max(message_indices) + 1`` — useful for manually-constructed
+        ``RenderedTokens`` in tests but only correct when the last
+        message contributed at least one token.
+
+        Cheap to call: single pass over ``message_indices``. Re-call
+        rather than caching the result if you mutate the dataclass.
+        """
+        if self.message_roles:
+            n_messages = len(self.message_roles)
+        else:
+            max_idx = -1
+            for idx in self.message_indices:
+                if idx > max_idx:
+                    max_idx = idx
+            n_messages = max_idx + 1
+
+        firsts: list[int] = [-1] * n_messages
+        lasts: list[int] = [-1] * n_messages
+        for k, idx in enumerate(self.message_indices):
+            if 0 <= idx < n_messages:
+                if firsts[idx] == -1:
+                    firsts[idx] = k
+                lasts[idx] = k
+
+        out: list[tuple[int, int] | None] = []
+        for i in range(n_messages):
+            if firsts[i] == -1:
+                out.append(None)
+            else:
+                out.append((firsts[i], lasts[i] + 1))
+        return out
+
+    def role_token_spans(self) -> dict[str, list[tuple[int, int]]]:
+        """:meth:`message_token_spans` regrouped by ``message_roles``.
+
+        Maps each role appearing in :attr:`message_roles` to a list of
+        ``(start, end)`` spans — one per occurrence of that role, in
+        message order. Messages with no contributed tokens are skipped.
+        Returns an empty dict if :attr:`message_roles` is empty.
+
+        Intended for per-role statistics that operate on per-token
+        signals — e.g. ``logprobs[start:end]`` for each assistant span
+        to compute per-turn perplexity, or
+        ``attention[start:end]`` for tool-response attention analysis.
+        """
+        spans = self.message_token_spans()
+        out: dict[str, list[tuple[int, int]]] = {}
+        for role, span in zip(self.message_roles, spans):
+            if span is None:
+                out.setdefault(role, [])
+                continue
+            out.setdefault(role, []).append(span)
+        return out
+
+    def tokens_by_role(self, *, sampled_only: bool = False) -> dict[str, int]:
+        """Sum :meth:`tokens_per_message` grouped by ``message_roles``.
+
+        Convenience for length-penalty bookkeeping in RL trainers:
+        ``rendered.tokens_by_role(sampled_only=True)["assistant"]`` is
+        the count of tokens the model actually emitted across all
+        assistant turns — template scaffolding excluded.
+        ``rendered.tokens_by_role()["tool"]`` is the raw count of
+        tool-response tokens (``sampled_only`` is zero for ``tool`` by
+        construction since the model never samples those).
+
+        Roles present in :attr:`message_roles` always appear in the
+        returned dict, even with post-filter count ``0``, so callers
+        can index directly without ``KeyError`` on conversations that
+        happen to lack a role. Returns an empty dict if
+        :attr:`message_roles` is empty.
+        """
+        counts = self.tokens_per_message(sampled_only=sampled_only)
+        out: dict[str, int] = {}
+        for role, n in zip(self.message_roles, counts):
+            out[role] = out.get(role, 0) + n
+        return out
+
 
 class ToolCallParseStatus(str, enum.Enum):
     """Per-attempt outcome of parsing a single ``<tool_call>`` block.
@@ -358,6 +512,25 @@ class Renderer(Protocol):
         list so far with ``add_generation_prompt=True`` — except prev
         sampled tokens are kept verbatim rather than re-rendered).
 
+        Attribution on the returned ``RenderedTokens``:
+
+        - ``message_indices`` is ``-1`` over the entire prior portion
+          (length ``len(previous_ids)`` after :func:`trim_to_turn_close`)
+          because the bridge gets the prior as raw token lists with no
+          attribution. Over the bridge-added portion, indices are
+          relative to ``new_messages``: a token rendered as part of
+          ``new_messages[i]`` carries ``i``, and inter-turn separators /
+          the trailing generation prompt carry ``-1``. So
+          ``bridge.tokens_per_message(len(new_messages))`` gives the
+          per-new-message token count for length-penalty bookkeeping.
+        - ``sampled_mask`` is uniformly ``False`` across the entire
+          returned sequence. The bridge output is consumed as the next
+          turn's prompt; nothing it emits was model-sampled, and the
+          bridge has no way to recover which prior tokens were. If the
+          caller needs that distinction for the prior portion, they
+          have it directly: every token in ``prev_completion_ids`` was
+          sampled; every token in ``prev_prompt_ids`` was not.
+
         Text-only renderers return :class:`RenderedTokens` with
         ``multi_modal_data=None``. Multimodal renderers (see
         :class:`MultimodalRenderer`) populate ``multi_modal_data`` so

{renderers-0.1.8.dev2 → renderers-0.1.9.dev0}/renderers/deepseek_v3.py

@@ -210,7 +210,10 @@ class DeepSeekV3Renderer:
                 emit_text("<think>\n", -1, is_sampled=False)
 
         return RenderedTokens(
-            token_ids=tokens, message_indices=indices, sampled_mask=sampled
+            token_ids=tokens,
+            message_indices=indices,
+            sampled_mask=sampled,
+            message_roles=[m.get("role") or "" for m in messages],
         )
 
     def render_ids(
@@ -271,22 +274,29 @@ class DeepSeekV3Renderer:
             return None
 
         ext: list[int] = []
-
-        # Bridge output is consumed as the next turn's prompt — the
-        # caller blanket-masks it via ``prompt_mask=[False]*N``, so we
-        # don't track sampled_mask here. Local helpers accept the kwarg
-        # for signature compatibility with ``_render_tool`` and ignore
-        # it; the returned ``RenderedTokens`` leaves ``sampled_mask``
-        # empty.
+        ext_indices: list[int] = []
+        ext_sampled: list[bool] = []
+
+        # Bridge populates ``message_indices`` (relative to ``new_messages``)
+        # and ``sampled_mask`` (uniformly ``False`` — every token the
+        # bridge emits is template scaffolding for the next prompt, not
+        # something the model sampled). Downstream consumers can run
+        # :meth:`RenderedTokens.tokens_per_message` on the bridge output
+        # to get per-new-message token counts without re-rendering.
         def emit_special(
-            token_id: int, _msg_idx: int = -1, *, is_sampled: bool = False
+            token_id: int, msg_idx: int = -1, *, is_sampled: bool = False
         ) -> None:
             ext.append(token_id)
+            ext_indices.append(msg_idx)
+            ext_sampled.append(is_sampled)
 
         def emit_text(
-            text: str, _msg_idx: int = -1, *, is_sampled: bool = False
+            text: str, msg_idx: int = -1, *, is_sampled: bool = False
         ) -> None:
-            ext.extend(self._encode(text))
+            ids = self._encode(text)
+            ext.extend(ids)
+            ext_indices.extend([msg_idx] * len(ids))
+            ext_sampled.extend([is_sampled] * len(ids))
 
         for i, msg in enumerate(new_messages):
             role = msg.get("role")
@@ -329,7 +339,13 @@ class DeepSeekV3Renderer:
         if self._enable_thinking:
             emit_text("<think>\n", -1)
 
-        return RenderedTokens(token_ids=previous_ids + ext)
+        total_len = len(previous_ids) + len(ext)
+        return RenderedTokens(
+            token_ids=previous_ids + ext,
+            message_indices=[-1] * len(previous_ids) + ext_indices,
+            sampled_mask=[False] * total_len,
+            message_roles=[m.get("role") or "" for m in new_messages],
+        )
 
     # ------------------------------------------------------------------
     # Assistant rendering

{renderers-0.1.8.dev2 → renderers-0.1.9.dev0}/renderers/default.py

@@ -143,7 +143,12 @@ class DefaultRenderer:
             token_ids = full_ids
             message_indices.extend([-1] * len(gen_tokens))
 
-        return RenderedTokens(token_ids=token_ids, message_indices=message_indices)
+        message_roles = [m.get("role") or "" for m in messages]
+        return RenderedTokens(
+            token_ids=token_ids,
+            message_indices=message_indices,
+            message_roles=message_roles,
+        )
 
     def _apply(self, messages, *, tools=None, add_generation_prompt=False) -> list[int]:
         kwargs = dict(self._chat_template_kwargs)

{renderers-0.1.8.dev2 → renderers-0.1.9.dev0}/renderers/glm45.py

@@ -203,7 +203,10 @@ class GLM45Renderer:
                 emit_special(self._think_end, -1, is_sampled=False)
 
         return RenderedTokens(
-            token_ids=tokens, message_indices=indices, sampled_mask=sampled
+            token_ids=tokens,
+            message_indices=indices,
+            sampled_mask=sampled,
+            message_roles=[m.get("role") or "" for m in messages],
         )
 
     def render_ids(
@@ -271,22 +274,29 @@ class GLM45Renderer:
         last_prev = previous_ids[-1]
 
         ext: list[int] = []
-
-        # Bridge output is consumed as the next turn's prompt — the
-        # caller blanket-masks it via ``prompt_mask=[False]*N``, so we
-        # don't track sampled_mask here. Local helpers accept the kwarg
-        # for signature compatibility with ``_render_tool`` and ignore
-        # it; the returned ``RenderedTokens`` leaves ``sampled_mask``
-        # empty.
+        ext_indices: list[int] = []
+        ext_sampled: list[bool] = []
+
+        # Bridge populates ``message_indices`` (relative to ``new_messages``)
+        # and ``sampled_mask`` (uniformly ``False`` — every token the
+        # bridge emits is template scaffolding for the next prompt, not
+        # something the model sampled). Downstream consumers can run
+        # :meth:`RenderedTokens.tokens_per_message` on the bridge output
+        # to get per-new-message token counts without re-rendering.
         def emit_special(
-            token_id: int, _msg_idx: int = -1, *, is_sampled: bool = False
+            token_id: int, msg_idx: int = -1, *, is_sampled: bool = False
         ) -> None:
             ext.append(token_id)
+            ext_indices.append(msg_idx)
+            ext_sampled.append(is_sampled)
 
         def emit_text(
-            text: str, _msg_idx: int = -1, *, is_sampled: bool = False
+            text: str, msg_idx: int = -1, *, is_sampled: bool = False
         ) -> None:
-            ext.extend(self._encode(text))
+            ids = self._encode(text)
+            ext.extend(ids)
+            ext_indices.extend([msg_idx] * len(ids))
+            ext_sampled.extend([is_sampled] * len(ids))
 
         for i, msg in enumerate(new_messages):
             role = msg.get("role")
@@ -318,7 +328,13 @@ class GLM45Renderer:
             emit_special(self._think, -1)
             emit_special(self._think_end, -1)
 
-        return RenderedTokens(token_ids=previous_ids + ext)
+        total_len = len(previous_ids) + len(ext)
+        return RenderedTokens(
+            token_ids=previous_ids + ext,
+            message_indices=[-1] * len(previous_ids) + ext_indices,
+            sampled_mask=[False] * total_len,
+            message_roles=[m.get("role") or "" for m in new_messages],
+        )
 
     def _render_assistant(
         self,

{renderers-0.1.8.dev2 → renderers-0.1.9.dev0}/renderers/glm5.py

@@ -220,7 +220,10 @@ class GLM5Renderer:
                 emit_special(self._think_end, -1, is_sampled=False)
 
         return RenderedTokens(
-            token_ids=tokens, message_indices=indices, sampled_mask=sampled
+            token_ids=tokens,
+            message_indices=indices,
+            sampled_mask=sampled,
+            message_roles=[m.get("role") or "" for m in messages],
         )
 
     def render_ids(
@@ -292,22 +295,29 @@ class GLM5Renderer:
         last_prev = previous_ids[-1]
 
         ext: list[int] = []
-
-        # Bridge output is consumed as the next turn's prompt — the
-        # caller blanket-masks it via ``prompt_mask=[False]*N``, so we
-        # don't track sampled_mask here. Local helpers accept the kwarg
-        # for signature compatibility with ``_render_assistant`` /
-        # ``_render_tool`` and ignore it; the returned ``RenderedTokens``
-        # leaves ``sampled_mask`` empty.
+        ext_indices: list[int] = []
+        ext_sampled: list[bool] = []
+
+        # Bridge populates ``message_indices`` (relative to ``new_messages``)
+        # and ``sampled_mask`` (uniformly ``False`` — every token the
+        # bridge emits is template scaffolding for the next prompt, not
+        # something the model sampled). Downstream consumers can run
+        # :meth:`RenderedTokens.tokens_per_message` on the bridge output
+        # to get per-new-message token counts without re-rendering.
         def emit_special(
-            token_id: int, _msg_idx: int = -1, *, is_sampled: bool = False
+            token_id: int, msg_idx: int = -1, *, is_sampled: bool = False
         ) -> None:
             ext.append(token_id)
+            ext_indices.append(msg_idx)
+            ext_sampled.append(is_sampled)
 
         def emit_text(
-            text: str, _msg_idx: int = -1, *, is_sampled: bool = False
+            text: str, msg_idx: int = -1, *, is_sampled: bool = False
         ) -> None:
-            ext.extend(self._encode(text))
+            ids = self._encode(text)
+            ext.extend(ids)
+            ext_indices.extend([msg_idx] * len(ids))
+            ext_sampled.extend([is_sampled] * len(ids))
 
         for i, msg in enumerate(new_messages):
             role = msg.get("role")
@@ -340,7 +350,13 @@ class GLM5Renderer:
         else:
             emit_special(self._think_end, -1)
 
-        return RenderedTokens(token_ids=previous_ids + ext)
+        total_len = len(previous_ids) + len(ext)
+        return RenderedTokens(
+            token_ids=previous_ids + ext,
+            message_indices=[-1] * len(previous_ids) + ext_indices,
+            sampled_mask=[False] * total_len,
+            message_roles=[m.get("role") or "" for m in new_messages],
+        )
 
     def _render_assistant(
         self,

{renderers-0.1.8.dev2 → renderers-0.1.9.dev0}/renderers/gpt_oss.py

@@ -333,7 +333,10 @@ class GptOssRenderer:
             emit([self._message], -1, is_sampled=False)
 
         return RenderedTokens(
-            token_ids=tokens, message_indices=indices, sampled_mask=sampled
+            token_ids=tokens,
+            message_indices=indices,
+            sampled_mask=sampled,
+            message_roles=[m.get("role") or "" for m in messages],
         )
 
     def render_ids(
@@ -400,22 +403,38 @@ class GptOssRenderer:
         if previous_ids is None:
             return None
 
+        # Bridge populates ``message_indices`` (relative to ``new_messages``)
+        # and ``sampled_mask`` (uniformly ``False``). The harmony encoder
+        # renders each ``new_messages[i]`` as a single block, so every
+        # token in that block carries index ``i``; the trailing
+        # generation prompt uses ``-1``.
         ext: list[int] = []
-        for msg in new_messages:
+        ext_indices: list[int] = []
+        for i, msg in enumerate(new_messages):
             role = msg.get("role")
             if role not in ("tool", "user", "system", "developer"):
                 return None
             for hm in self._to_harmony_messages(msg):
-                ext.extend(self._enc.render(hm))
+                ids = self._enc.render(hm)
+                ext.extend(ids)
+                ext_indices.extend([i] * len(ids))
 
         # Generation prompt: <|start|>assistant<|channel|>analysis<|message|>
+        gen_before = len(ext)
         ext.append(self._start)
         ext.extend(self._encode("assistant"))
         ext.append(self._channel)
         ext.extend(self._encode("analysis"))
         ext.append(self._message)
+        ext_indices.extend([-1] * (len(ext) - gen_before))
 
-        return RenderedTokens(token_ids=previous_ids + ext)
+        total_len = len(previous_ids) + len(ext)
+        return RenderedTokens(
+            token_ids=previous_ids + ext,
+            message_indices=[-1] * len(previous_ids) + ext_indices,
+            sampled_mask=[False] * total_len,
+            message_roles=[m.get("role") or "" for m in new_messages],
+        )
 
     # ── message conversion ───────────────────────────────────────────────────
 

{renderers-0.1.8.dev2 → renderers-0.1.9.dev0}/renderers/kimi_k2.py

@@ -270,7 +270,10 @@ class KimiK2Renderer:
             emit_special(self._im_middle, -1, is_sampled=False)
 
         return RenderedTokens(
-            token_ids=token_ids, message_indices=indices, sampled_mask=sampled
+            token_ids=token_ids,
+            message_indices=indices,
+            sampled_mask=sampled,
+            message_roles=[m.get("role") or "" for m in messages],
         )
 
     def render_ids(
@@ -331,21 +334,29 @@ class KimiK2Renderer:
             return None
 
         ext: list[int] = []
-
-        # Bridge output is consumed as the next turn's prompt — the caller
-        # blanket-masks it via ``prompt_mask=[False]*N``, so we don't track
-        # sampled_mask here. Local helpers accept the kwarg for signature
-        # compatibility with ``_render_tool`` and ignore it; the returned
-        # ``RenderedTokens`` leaves ``sampled_mask`` empty.
+        ext_indices: list[int] = []
+        ext_sampled: list[bool] = []
+
+        # Bridge populates ``message_indices`` (relative to ``new_messages``)
+        # and ``sampled_mask`` (uniformly ``False`` — every token the
+        # bridge emits is template scaffolding for the next prompt, not
+        # something the model sampled). Downstream consumers can run
+        # :meth:`RenderedTokens.tokens_per_message` on the bridge output
+        # to get per-new-message token counts without re-rendering.
         def emit_special(
-            token_id: int, _msg_idx: int = -1, *, is_sampled: bool = False
+            token_id: int, msg_idx: int = -1, *, is_sampled: bool = False
         ) -> None:
             ext.append(token_id)
+            ext_indices.append(msg_idx)
+            ext_sampled.append(is_sampled)
 
         def emit_text(
-            text: str, _msg_idx: int = -1, *, is_sampled: bool = False
+            text: str, msg_idx: int = -1, *, is_sampled: bool = False
         ) -> None:
-            ext.extend(self._encode(text))
+            ids = self._encode(text)
+            ext.extend(ids)
+            ext_indices.extend([msg_idx] * len(ids))
+            ext_sampled.extend([is_sampled] * len(ids))
 
         for i, msg in enumerate(new_messages):
             role = msg.get("role")
@@ -388,7 +399,13 @@ class KimiK2Renderer:
         emit_text("assistant", -1, is_sampled=False)
         emit_special(self._im_middle, -1, is_sampled=False)
 
-        return RenderedTokens(token_ids=previous_ids + ext)
+        total_len = len(previous_ids) + len(ext)
+        return RenderedTokens(
+            token_ids=previous_ids + ext,
+            message_indices=[-1] * len(previous_ids) + ext_indices,
+            sampled_mask=[False] * total_len,
+            message_roles=[m.get("role") or "" for m in new_messages],
+        )
 
     def _render_assistant(
         self,

{renderers-0.1.8.dev2 → renderers-0.1.9.dev0}/renderers/kimi_k25.py

@@ -906,6 +906,7 @@ class KimiK25Renderer:
             token_ids=tokens,
             message_indices=indices,
             sampled_mask=sampled,
+            message_roles=[m.get("role") or "" for m in messages],
             multi_modal_data=mm_data,
         )
 
@@ -995,44 +996,52 @@ class KimiK25Renderer:
             return None
 
         # Seed combined-token list with prior turn so placeholder offsets
-        # are absolute in the bridged sequence.
+        # are absolute in the bridged sequence. Parallel
+        # ``indices``/``sampled`` are seeded with ``-1``/``False`` for the
+        # prior portion — the bridge has no attribution info for
+        # ``previous_ids``. Bridge-added tokens get proper ``msg_idx``
+        # (relative to ``new_messages``) and uniformly ``False``
+        # ``sampled``: nothing the bridge emits was model-sampled.
         tokens: list[int] = list(previous_ids)
+        indices: list[int] = [-1] * len(previous_ids)
+        sampled: list[bool] = [False] * len(previous_ids)
         new_hashes: dict[str, list[str]] = {}
         new_placeholders: dict[str, list[PlaceholderRange]] = {}
         new_items: dict[str, list[dict[str, Any]]] = {}
 
-        # Bridge output is consumed as the next turn's prompt — the caller
-        # blanket-masks it via ``prompt_mask=[False]*N``, so we don't track
-        # sampled_mask here. Local helpers accept the kwarg for signature
-        # compatibility with ``_render_tool_body`` / ``_emit_content`` and
-        # ignore it; the returned ``RenderedTokens`` leaves ``sampled_mask``
-        # empty.
         def emit_special(
-            token_id: int, _msg_idx: int = -1, *, is_sampled: bool = False
+            token_id: int, msg_idx: int = -1, *, is_sampled: bool = False
         ) -> None:
             tokens.append(token_id)
+            indices.append(msg_idx)
+            sampled.append(is_sampled)
 
         def emit_text(
-            text: str, _msg_idx: int = -1, *, is_sampled: bool = False
+            text: str, msg_idx: int = -1, *, is_sampled: bool = False
         ) -> None:
-            tokens.extend(self._encode(text))
+            ids = self._encode(text)
+            tokens.extend(ids)
+            indices.extend([msg_idx] * len(ids))
+            sampled.extend([is_sampled] * len(ids))
 
         def emit_ids(
-            ids: list[int], _msg_idx: int = -1, *, is_sampled: bool = False
+            ids: list[int], msg_idx: int = -1, *, is_sampled: bool = False
         ) -> None:
             tokens.extend(ids)
+            indices.extend([msg_idx] * len(ids))
+            sampled.extend([is_sampled] * len(ids))
 
         def emit_image(
-            part: dict[str, Any], _msg_idx: int = -1, *, is_sampled: bool = False
+            part: dict[str, Any], msg_idx: int = -1, *, is_sampled: bool = False
         ) -> None:
             _, out, _num_patches, h = self._process_image(part)
-            emit_special(self._media_begin)
-            emit_text("image")
-            emit_special(self._media_content)
+            emit_special(self._media_begin, msg_idx)
+            emit_text("image", msg_idx)
+            emit_special(self._media_content, msg_idx)
             offset = len(tokens)
-            emit_special(self._media_pad)
-            emit_special(self._media_end)
-            emit_text("\n")
+            emit_special(self._media_pad, msg_idx)
+            emit_special(self._media_end, msg_idx)
+            emit_text("\n", msg_idx)
             new_hashes.setdefault("image", []).append(h)
             new_placeholders.setdefault("image", []).append(
                 PlaceholderRange(offset=offset, length=1)
@@ -1113,8 +1122,14 @@ class KimiK25Renderer:
         for modality, vals in new_items.items():
             merged_items.setdefault(modality, []).extend(vals)
 
+        bridge_roles = [m.get("role") or "" for m in new_messages]
         if not (merged_hashes or merged_placeholders or merged_items):
-            return RenderedTokens(token_ids=tokens)
+            return RenderedTokens(
+                token_ids=tokens,
+                message_indices=indices,
+                sampled_mask=sampled,
+                message_roles=bridge_roles,
+            )
 
         mm_data = MultiModalData(
             mm_hashes=merged_hashes,
@@ -1123,7 +1138,9 @@ class KimiK25Renderer:
         )
         return RenderedTokens(
             token_ids=tokens,
-            message_indices=[-1] * len(tokens),
+            message_indices=indices,
+            sampled_mask=sampled,
+            message_roles=bridge_roles,
             multi_modal_data=mm_data,
         )