renderers 0.1.12.dev16__tar.gz → 0.1.12.dev17__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (90) hide show
  1. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/PKG-INFO +3 -2
  2. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/README.md +2 -1
  3. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/renderers/_version.py +2 -2
  4. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/renderers/base.py +45 -8
  5. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/renderers/glm45.py +11 -35
  6. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/renderers/glm5.py +21 -32
  7. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/tests/test_build_helpers.py +104 -0
  8. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/tests/test_message_indices.py +63 -0
  9. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/tests/test_sampled_mask.py +2 -0
  10. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/.github/workflows/publish-dev.yml +0 -0
  11. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/.github/workflows/publish.yml +0 -0
  12. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/.github/workflows/style.yml +0 -0
  13. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/.github/workflows/test.yml +0 -0
  14. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/.gitignore +0 -0
  15. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/.pre-commit-config.yaml +0 -0
  16. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/LICENSE +0 -0
  17. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/docs/renderer-config.md +0 -0
  18. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/examples/README.md +0 -0
  19. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/examples/sglang/multiturn_generate_sglang.py +0 -0
  20. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/examples/sglang/online_multiturn_sglang.py +0 -0
  21. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/examples/tinker/multiturn_generate_tinker.py +0 -0
  22. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/examples/transformers/multiturn_generate_transformers.py +0 -0
  23. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/examples/vllm/multiturn_generate_vllm.py +0 -0
  24. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/pyproject.toml +0 -0
  25. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/renderers/__init__.py +0 -0
  26. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/renderers/client.py +0 -0
  27. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/renderers/configs.py +0 -0
  28. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/renderers/custom.py +0 -0
  29. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/renderers/deepseek_r1.py +0 -0
  30. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/renderers/deepseek_v3.py +0 -0
  31. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/renderers/deepseek_v4.py +0 -0
  32. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/renderers/default.py +0 -0
  33. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/renderers/gemma4.py +0 -0
  34. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/renderers/gpt_oss.py +0 -0
  35. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/renderers/hy3.py +0 -0
  36. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/renderers/inkling.py +0 -0
  37. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/renderers/kimi_k2.py +0 -0
  38. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/renderers/kimi_k25.py +0 -0
  39. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/renderers/laguna_s21.py +0 -0
  40. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/renderers/laguna_xs2.py +0 -0
  41. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/renderers/llama_3.py +0 -0
  42. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/renderers/minimax_m2.py +0 -0
  43. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/renderers/nemotron3.py +0 -0
  44. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/renderers/parsers.py +0 -0
  45. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/renderers/parsing.py +0 -0
  46. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/renderers/prime_qwen3.py +0 -0
  47. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/renderers/qwen3.py +0 -0
  48. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/renderers/qwen35.py +0 -0
  49. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/renderers/qwen36.py +0 -0
  50. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/renderers/qwen38.py +0 -0
  51. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/renderers/qwen3_vl.py +0 -0
  52. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/renderers/reasoning.py +0 -0
  53. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/tests/conftest.py +0 -0
  54. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/tests/parity.py +0 -0
  55. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/tests/reference_rendering.py +0 -0
  56. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/tests/test_bridge.py +0 -0
  57. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/tests/test_client.py +0 -0
  58. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/tests/test_custom.py +0 -0
  59. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/tests/test_deepseek_r1.py +0 -0
  60. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/tests/test_deepseek_v4.py +0 -0
  61. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/tests/test_disabled_thinking_stability.py +0 -0
  62. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/tests/test_gemma4.py +0 -0
  63. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/tests/test_glm_tool_name_validation.py +0 -0
  64. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/tests/test_hy3.py +0 -0
  65. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/tests/test_incremental.py +0 -0
  66. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/tests/test_inkling.py +0 -0
  67. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/tests/test_is_content.py +0 -0
  68. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/tests/test_kimi_k25_tool_schema.py +0 -0
  69. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/tests/test_laguna_m1.py +0 -0
  70. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/tests/test_laguna_xs21.py +0 -0
  71. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/tests/test_llama_3.py +0 -0
  72. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/tests/test_load_tokenizer.py +0 -0
  73. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/tests/test_message_tool_names.py +0 -0
  74. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/tests/test_multimodal.py +0 -0
  75. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/tests/test_nemotron35_vlm.py +0 -0
  76. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/tests/test_nemotron3_ultra.py +0 -0
  77. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/tests/test_offsetless_tokenizers.py +0 -0
  78. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/tests/test_parity.py +0 -0
  79. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/tests/test_parse_response.py +0 -0
  80. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/tests/test_parse_response_robustness.py +0 -0
  81. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/tests/test_parsers.py +0 -0
  82. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/tests/test_preserve_thinking.py +0 -0
  83. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/tests/test_qwen35_size_coverage.py +0 -0
  84. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/tests/test_qwen38.py +0 -0
  85. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/tests/test_reasoning_boundaries.py +0 -0
  86. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/tests/test_renderer_config.py +0 -0
  87. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/tests/test_roundtrip.py +0 -0
  88. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/tests/test_tokens_per_message.py +0 -0
  89. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/tests/test_tool_arg_type_preservation.py +0 -0
  90. {renderers-0.1.12.dev16 → renderers-0.1.12.dev17}/uv.lock +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: renderers
3
- Version: 0.1.12.dev16
3
+ Version: 0.1.12.dev17
4
4
  Summary: Chat template renderers — deterministic message-to-token conversion for LLM training
5
5
  License-Expression: Apache-2.0
6
6
  License-File: LICENSE
@@ -127,7 +127,8 @@ class Renderer(Protocol):
127
127
  def bridge_to_next_turn(prev_prompt_ids, prev_completion_ids, new_messages, *, tools=None) -> list[int] | None: ...
128
128
  ```
129
129
 
130
- - `RenderedTokens` carries `token_ids` **and** `message_indices` — one entry per token attributing each to its source message (`-1` for structural scaffolding). Lets `build_training_sample` build a per-token loss mask in one render.
130
+ - `RenderedTokens` carries `token_ids` **and** `message_indices` — one entry per token attributing each to its owning message (`-1` for structural scaffolding). Sampled closing markers belong to their assistant, including GLM markers emitted at the next user/tool boundary; user/tool bodies keep their own ownership.
131
+ - `build_training_sample(..., message_loss_mask=[0, 0, 0, 1])` keeps all messages as context but restricts loss to selected messages. Supply one boolean or integer `0`/`1` per message, or `None` for the normal mask. Excluding an assistant also excludes its closing token. Selection never enables targets disabled by the existing role/content settings or changes token IDs.
131
132
  - `ParsedResponse` is `(content, reasoning_content, tool_calls)`. It scans token ids for special-token boundaries (e.g. id `151657` for `<tool_call>` on Qwen3) — a literal `"<tool_call>"` in user content tokenizes to ordinary text ids and never matches.
132
133
  - Round-trip: rendering `[user, assistant(content, reasoning, tool_calls)]`, slicing the assistant completion, and feeding it through `parse_response` returns an equivalent structured message. Tested per-renderer in `tests/test_roundtrip.py`.
133
134
 
@@ -107,7 +107,8 @@ class Renderer(Protocol):
107
107
  def bridge_to_next_turn(prev_prompt_ids, prev_completion_ids, new_messages, *, tools=None) -> list[int] | None: ...
108
108
  ```
109
109
 
110
- - `RenderedTokens` carries `token_ids` **and** `message_indices` — one entry per token attributing each to its source message (`-1` for structural scaffolding). Lets `build_training_sample` build a per-token loss mask in one render.
110
+ - `RenderedTokens` carries `token_ids` **and** `message_indices` — one entry per token attributing each to its owning message (`-1` for structural scaffolding). Sampled closing markers belong to their assistant, including GLM markers emitted at the next user/tool boundary; user/tool bodies keep their own ownership.
111
+ - `build_training_sample(..., message_loss_mask=[0, 0, 0, 1])` keeps all messages as context but restricts loss to selected messages. Supply one boolean or integer `0`/`1` per message, or `None` for the normal mask. Excluding an assistant also excludes its closing token. Selection never enables targets disabled by the existing role/content settings or changes token IDs.
111
112
  - `ParsedResponse` is `(content, reasoning_content, tool_calls)`. It scans token ids for special-token boundaries (e.g. id `151657` for `<tool_call>` on Qwen3) — a literal `"<tool_call>"` in user content tokenizes to ordinary text ids and never matches.
112
113
  - Round-trip: rendering `[user, assistant(content, reasoning, tool_calls)]`, slicing the assistant completion, and feeding it through `parse_response` returns an equivalent structured message. Tested per-renderer in `tests/test_roundtrip.py`.
113
114
 
@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
18
18
  commit_id: str | None
19
19
  __commit_id__: str | None
20
20
 
21
- __version__ = version = '0.1.12.dev16'
22
- __version_tuple__ = version_tuple = (0, 1, 12, 'dev16')
21
+ __version__ = version = '0.1.12.dev17'
22
+ __version_tuple__ = version_tuple = (0, 1, 12, 'dev17')
23
23
 
24
24
  __commit_id__ = commit_id = None
@@ -224,8 +224,11 @@ class MultiModalData:
224
224
  class RenderedTokens:
225
225
  """Result of rendering messages to tokens.
226
226
 
227
- Each token carries an index into the original message list so callers can
228
- build per-token loss masks without re-rendering. Tokens from structural
227
+ Each token carries its owning message's index into the original message
228
+ list so callers can build per-token loss masks without re-rendering.
229
+ Sampled closing tokens belong to the assistant that generated them, even
230
+ when emitted while rendering the next message's role marker. User/tool
231
+ bodies retain their own message indices. Tokens from structural
229
232
  scaffolding the renderer adds outside any single message (e.g. the
230
233
  trailing generation prompt) carry index ``-1``.
231
234
 
@@ -383,7 +386,10 @@ class RenderedTokens:
383
386
  is not represented.
384
387
 
385
388
  Hand-coded renderers emit each message's tokens contiguously,
386
- so the span is well-defined. The implementation tolerates
389
+ so the span is well-defined. Spans follow input-message indices,
390
+ not necessarily token order: templates can reorder tool results.
391
+ Empty messages whose role marker closes the preceding assistant
392
+ may have no tokens of their own. The implementation tolerates
387
393
  non-contiguous attribution by returning the outer span
388
394
  ``(first_k, last_k + 1)``; if you suspect interleaving, slice
389
395
  ``message_indices`` yourself to verify.
@@ -1640,6 +1646,7 @@ def build_training_sample(
1640
1646
  tools: list[ToolSpec] | None = None,
1641
1647
  content_sft_roles: "set[str] | frozenset[str] | None" = None,
1642
1648
  ensure_final_stop: bool = False,
1649
+ message_loss_mask: list[bool | int] | None = None,
1643
1650
  ) -> RenderedTrainingSample:
1644
1651
  """Build a :class:`RenderedTrainingSample` for supervised training.
1645
1652
 
@@ -1656,11 +1663,8 @@ def build_training_sample(
1656
1663
  it's attributed to. This is the recommended default for renderer
1657
1664
  callers — the renderer owns the per-token "is this model output"
1658
1665
  signal, so role-level filtering becomes a downstream constraint
1659
- rather than a precondition. (Some role markers — e.g. GLM
1660
- ``<|user|>`` / ``<|observation|>`` after a tool-calling assistant
1661
- turn — *are* sampled by the model at inference and live inside the
1662
- next message's span; ``sampled_mask`` captures that, but a
1663
- naive role filter would mask them out.)
1666
+ rather than a precondition. Sampled closing markers, including GLM's
1667
+ ``<|user|>`` / ``<|observation|>``, are attributed to their assistant.
1664
1668
 
1665
1669
  When ``role_to_mask`` is provided, ``loss_mask`` is the AND of the
1666
1670
  role-based attribution and the sampled signal: only tokens the
@@ -1706,7 +1710,27 @@ def build_training_sample(
1706
1710
  fires, the output intentionally diverges from ``apply_chat_template``.
1707
1711
  Ignored for renderers without ``sampled_mask`` (``DefaultRenderer``) —
1708
1712
  the close of an opaque template can't be located reliably.
1713
+
1714
+ ``message_loss_mask`` optionally selects messages for supervision with
1715
+ one boolean or integer 0/1 per message. It only restricts the existing
1716
+ mask: selecting a message never enables targets disabled by the role,
1717
+ sampled, or content settings. Unselected messages remain context,
1718
+ including prefilled assistant turns at arbitrary positions. Sampled
1719
+ closing tokens follow their owning assistant, even when a template
1720
+ emits them at the next message boundary. A synthesized final stop
1721
+ follows the final message's selection. Token ids are unaffected.
1722
+ ``None`` preserves the existing mask.
1709
1723
  """
1724
+ if message_loss_mask is not None:
1725
+ if len(message_loss_mask) != len(messages):
1726
+ raise ValueError("message_loss_mask must have one entry per message")
1727
+ if any(
1728
+ not isinstance(value, (bool, int)) or value not in (0, 1)
1729
+ for value in message_loss_mask
1730
+ ):
1731
+ raise ValueError(
1732
+ "message_loss_mask entries must be booleans or integers 0/1"
1733
+ )
1710
1734
  rendered = renderer.render(messages, tools=tools)
1711
1735
  has_sampled_info = len(rendered.sampled_mask) == len(rendered.token_ids)
1712
1736
  has_content_info = len(rendered.is_content) == len(rendered.token_ids)
@@ -1765,6 +1789,19 @@ def build_training_sample(
1765
1789
  # stop is a training target, like any sampled token.
1766
1790
  loss_mask.append(True)
1767
1791
 
1792
+ if message_loss_mask is not None:
1793
+ for k in range(len(loss_mask)):
1794
+ # Any synthesized stop after the rendered sequence belongs to
1795
+ # the final assistant, just like an in-message closing token.
1796
+ msg_idx = (
1797
+ rendered.message_indices[k]
1798
+ if k < len(rendered.message_indices)
1799
+ else len(messages) - 1
1800
+ )
1801
+ loss_mask[k] = (
1802
+ loss_mask[k] and msg_idx >= 0 and bool(message_loss_mask[msg_idx])
1803
+ )
1804
+
1768
1805
  # Surface the multimodal payload for VLM renderers. ``None`` for text
1769
1806
  # renderers and for text-only samples (empty media) so downstream
1770
1807
  # ``multi_modal_data is not None`` is a reliable "has media" check.
@@ -192,20 +192,8 @@ class GLM45Renderer:
192
192
  role = msg["role"]
193
193
  content = self._visible_text(msg.get("content"))
194
194
 
195
- # When the previous message is an assistant, this message's
196
- # role-opening token (``<|user|>`` / ``<|observation|>``) is
197
- # the inference-time stop signal that closes the assistant's
198
- # turn (see ``get_stop_token_ids``). Mark it
199
- # ``is_sampled=True`` so the loss-mask pipeline trains the
200
- # model to emit it after ``</tool_call>`` (instead of
201
- # continuing with another ``<tool_call>`` block). The token
202
- # stays attributed to this message (msg_idx=i) and remains
203
- # ``is_content=False`` — it's a role-marker / scaffold, not
204
- # body bytes, so ``content_mask_for_roles({"tool"})`` and
205
- # ``content_token_spans_by_role()`` correctly exclude it
206
- # from "tool body" views. Byte stream is unchanged.
207
- # ``system`` only appears at the start of a GLM conversation,
208
- # so its opener is never the closer of an assistant turn.
195
+ # A user/tool role marker that closes generation belongs to
196
+ # the preceding assistant's sampled output, not the next body.
209
197
  closes_assistant_turn = i > 0 and messages[i - 1]["role"] == "assistant"
210
198
 
211
199
  if role == "system":
@@ -219,9 +207,9 @@ class GLM45Renderer:
219
207
  elif role == "user":
220
208
  emit_special(
221
209
  self._user,
222
- i,
210
+ i - 1 if closes_assistant_turn else i,
223
211
  is_sampled=closes_assistant_turn,
224
- is_content=False,
212
+ is_content=closes_assistant_turn,
225
213
  )
226
214
  # ``\n`` is scaffold; ``content`` is body; the optional
227
215
  # ``/nothink`` suffix is scaffold the renderer injects
@@ -508,15 +496,9 @@ class GLM45Renderer:
508
496
  # (think block + content + tool calls) is the model-sampled
509
497
  # portion.
510
498
  #
511
- # GLM-4.5 does NOT emit an explicit per-turn close token inside
512
- # the assistant message; the next message's role marker
513
- # (``<|user|>`` / ``<|observation|>`` / ``<|endoftext|>``) acts
514
- # as the stop signal at inference, and those tokens are
515
- # attributed to the *next* message (or are absent on the final
516
- # turn). So no sampled stop-signal token lives inside this
517
- # assistant span — content / think / tool_calls carry the
518
- # is_sampled=True signal.
519
- #
499
+ # The next user/tool role marker closes this assistant turn and
500
+ # is attributed here when that later message is rendered.
501
+
520
502
  # Invariant on assistant tokens: ``is_content == sampled_mask``.
521
503
  # Every scaffold token here gets ``is_sampled=False/is_content=False``;
522
504
  # every model-sampled emit gets both True.
@@ -598,23 +580,17 @@ class GLM45Renderer:
598
580
  emit_text,
599
581
  emit_text_segments,
600
582
  ) -> None:
601
- # Tool body bytes get ``is_content=True``; the wraps are
602
- # scaffold. The ``<|observation|>`` role tag is scaffold too
603
- # (``is_content=False`` so ``content_mask_for_roles({"tool"})``
604
- # excludes it). When the previous message is an assistant it
605
- # doubles as the inference stop signal for that assistant's
606
- # turn — mark it ``is_sampled=True`` so SFT trains the model to
607
- # emit it after ``</tool_call>``. The token stays attributed to
608
- # this tool message; byte stream is unchanged.
583
+ # Tool bodies keep their own attribution; a sampled observation
584
+ # marker belongs to the assistant that generated it.
609
585
  prev_role = messages[msg_idx - 1]["role"] if msg_idx > 0 else None
610
586
  closes_assistant_turn = prev_role == "assistant"
611
587
 
612
588
  if prev_role != "tool":
613
589
  emit_special(
614
590
  self._observation,
615
- msg_idx,
591
+ msg_idx - 1 if closes_assistant_turn else msg_idx,
616
592
  is_sampled=closes_assistant_turn,
617
- is_content=False,
593
+ is_content=closes_assistant_turn,
618
594
  )
619
595
 
620
596
  emit_text_segments(
@@ -225,20 +225,8 @@ class GLM5Renderer:
225
225
  role = msg["role"]
226
226
  content = self._visible_text(msg.get("content"))
227
227
 
228
- # When the previous message is an assistant, this message's
229
- # role-opening token (``<|user|>`` / ``<|observation|>``) is
230
- # the inference-time stop signal that closes the assistant's
231
- # turn (see ``get_stop_token_ids``). Mark it
232
- # ``is_sampled=True`` so the loss-mask pipeline trains the
233
- # model to emit it after ``</tool_call>`` (instead of
234
- # continuing with another ``<tool_call>`` block). The token
235
- # stays attributed to this message (msg_idx=i) and remains
236
- # ``is_content=False`` — it's a role-marker / scaffold, not
237
- # body bytes, so ``content_mask_for_roles({"tool"})`` and
238
- # ``content_token_spans_by_role()`` correctly exclude it
239
- # from "tool body" views. Byte stream is unchanged.
240
- # ``system`` only appears at the start of a GLM conversation,
241
- # so its opener is never the closer of an assistant turn.
228
+ # A user/tool role marker that closes generation belongs to
229
+ # the preceding assistant's sampled output, not the next body.
242
230
  closes_assistant_turn = i > 0 and messages[i - 1]["role"] == "assistant"
243
231
 
244
232
  if role == "system":
@@ -248,9 +236,9 @@ class GLM5Renderer:
248
236
  elif role == "user":
249
237
  emit_special(
250
238
  self._user,
251
- i,
239
+ i - 1 if closes_assistant_turn else i,
252
240
  is_sampled=closes_assistant_turn,
253
- is_content=False,
241
+ is_content=closes_assistant_turn,
254
242
  )
255
243
  emit_text(content, i, is_sampled=False, is_content=True)
256
244
 
@@ -272,13 +260,17 @@ class GLM5Renderer:
272
260
  ):
273
261
  tool_message = messages[tool_idx]
274
262
  self._render_tool(
275
- messages,
276
263
  tool_idx,
277
264
  self._visible_text(tool_message.get("content")),
278
265
  emit_special=emit_special,
279
266
  emit_text=emit_text,
280
267
  emit_text_segments=emit_text_segments,
281
268
  starts_block=position == 0,
269
+ closing_assistant_index=(
270
+ i - 1
271
+ if position == 0 and closes_assistant_turn
272
+ else None
273
+ ),
282
274
  )
283
275
 
284
276
  # ── Generation prompt ───────────────────────────────────────
@@ -649,7 +641,6 @@ class GLM5Renderer:
649
641
 
650
642
  def _render_tool(
651
643
  self,
652
- messages: list[Message],
653
644
  msg_idx: int,
654
645
  content: str,
655
646
  *,
@@ -657,24 +648,22 @@ class GLM5Renderer:
657
648
  emit_text,
658
649
  emit_text_segments,
659
650
  starts_block: bool = False,
651
+ closing_assistant_index: int | None = None,
660
652
  ) -> None:
661
- # Tool body bytes get ``is_content=True``; the wraps are
662
- # scaffold. The ``<|observation|>`` role tag is scaffold too
663
- # (``is_content=False`` so ``content_mask_for_roles({"tool"})``
664
- # excludes it). When the previous message is an assistant it
665
- # doubles as the inference stop signal for that assistant's
666
- # turn — mark it ``is_sampled=True`` so SFT trains the model to
667
- # emit it after ``</tool_call>``. The token stays attributed to
668
- # this tool message; byte stream is unchanged.
669
- prev_role = messages[msg_idx - 1]["role"] if msg_idx > 0 else None
670
- closes_assistant_turn = prev_role == "assistant"
671
-
672
- if starts_block or prev_role != "tool":
653
+ # Tool bodies keep their own attribution; a sampled observation
654
+ # marker belongs to the assistant that generated it.
655
+ # The caller supplies the block's assistant because results may
656
+ # be reordered relative to their input-message indices.
657
+ closes_assistant_turn = closing_assistant_index is not None
658
+
659
+ if starts_block:
673
660
  emit_special(
674
661
  self._observation,
675
- msg_idx,
662
+ closing_assistant_index
663
+ if closing_assistant_index is not None
664
+ else msg_idx,
676
665
  is_sampled=closes_assistant_turn,
677
- is_content=False,
666
+ is_content=closes_assistant_turn,
678
667
  )
679
668
 
680
669
  emit_special(
@@ -95,6 +95,110 @@ def test_build_training_sample_ensures_final_stop(model_name, tokenizer, rendere
95
95
  assert sample.token_ids == baseline.token_ids
96
96
 
97
97
 
98
+ @pytest.mark.parametrize(
99
+ "selection",
100
+ [[0] * 6, [1] * 6, [0, 1, 0, 0, 1, 0], [False, False, True, True, False, True]],
101
+ )
102
+ @pytest.mark.parametrize("ensure_final_stop", [False, True])
103
+ @pytest.mark.parametrize("body_roles", [None, {"user", "tool"}])
104
+ @pytest.mark.parametrize("tool_turn", [False, True])
105
+ def test_build_training_sample_message_loss_mask(
106
+ model_name, tokenizer, renderer, selection, ensure_final_stop, body_roles, tool_turn
107
+ ):
108
+ """Select arbitrary messages without changing tokens or enabling new targets."""
109
+ messages = [
110
+ {"role": "user", "content": "Hi"},
111
+ {"role": "assistant", "content": "Earlier answer."},
112
+ {"role": "user", "content": "Next question"},
113
+ {"role": "assistant", "content": "New answer."},
114
+ {"role": "user", "content": "One more question"},
115
+ {"role": "assistant", "content": "Final answer."},
116
+ ]
117
+ if tool_turn:
118
+ messages[1]["tool_calls"] = [
119
+ {
120
+ "id": "call_1",
121
+ "type": "function",
122
+ "function": {"name": "lookup", "arguments": '{"q":"x"}'},
123
+ }
124
+ ]
125
+ messages[2] = {
126
+ "role": "tool",
127
+ "content": "Lookup result",
128
+ "tool_call_id": "call_1",
129
+ "name": "lookup",
130
+ }
131
+ rendered = renderer.render(messages)
132
+ kwargs = {
133
+ "ensure_final_stop": ensure_final_stop,
134
+ "content_sft_roles": body_roles,
135
+ }
136
+ if not rendered.sampled_mask:
137
+ kwargs["role_to_mask"] = lambda m: m["role"] == "assistant"
138
+ baseline = build_training_sample(renderer, messages, **kwargs)
139
+ sample = build_training_sample(
140
+ renderer, messages, message_loss_mask=selection, **kwargs
141
+ )
142
+ assert sample.token_ids == baseline.token_ids
143
+ assert len(sample.loss_mask) == len(sample.token_ids)
144
+ for k, index in enumerate(rendered.message_indices):
145
+ if index < 0 or not selection[index]:
146
+ assert not sample.loss_mask[k]
147
+ else:
148
+ assert sample.loss_mask[k] == baseline.loss_mask[k]
149
+ # A synthesized stop is controlled by the final assistant's flag.
150
+ assert sample.loss_mask[len(rendered.token_ids) :] == [
151
+ keep and bool(selection[-1])
152
+ for keep in baseline.loss_mask[len(rendered.token_ids) :]
153
+ ]
154
+ disabled = build_training_sample(
155
+ renderer,
156
+ messages,
157
+ message_loss_mask=selection,
158
+ role_to_mask=lambda _: False,
159
+ ensure_final_stop=ensure_final_stop,
160
+ )
161
+ assert not any(disabled.loss_mask)
162
+ trained = tokenizer.decode(
163
+ [t for t, keep in zip(sample.token_ids[1:], sample.loss_mask[1:]) if keep]
164
+ )
165
+ for index, text in [
166
+ (1, "Earlier answer."),
167
+ (3, "New answer."),
168
+ (5, "Final answer."),
169
+ ]:
170
+ assert (text in trained) == bool(selection[index])
171
+ # Closing tokens are assistant output, including markers emitted while
172
+ # rendering the next user/tool message. Its body keeps its own owner.
173
+ for index, sampled in zip(rendered.message_indices, rendered.sampled_mask):
174
+ if sampled:
175
+ assert index >= 0 and messages[index]["role"] == "assistant"
176
+
177
+
178
+ @pytest.mark.parametrize(
179
+ "selection",
180
+ [
181
+ [],
182
+ [1],
183
+ [1] * 5,
184
+ [1, 0, 2, 1],
185
+ [1, -1, 0, 1],
186
+ [1, "0", 0, 1],
187
+ [1, 1.0, 0, 1],
188
+ [1, None, 0, 1],
189
+ ],
190
+ )
191
+ def test_build_training_sample_rejects_invalid_message_loss_mask(selection):
192
+ messages = [
193
+ {"role": "user", "content": "Hi"},
194
+ {"role": "assistant", "content": "Answer"},
195
+ {"role": "tool", "content": "Result"},
196
+ {"role": "assistant", "content": "Final"},
197
+ ]
198
+ with pytest.raises(ValueError, match="message_loss_mask"):
199
+ build_training_sample(None, messages, message_loss_mask=selection)
200
+
201
+
98
202
  def test_build_trajectory_step_reconstructs_full(model_name, tokenizer, renderer):
99
203
  """prompt_ids + completion_ids must equal the full rendered sequence."""
100
204
  prompt = [
@@ -40,6 +40,69 @@ Tests in this file:
40
40
 
41
41
  from __future__ import annotations
42
42
 
43
+ import pytest
44
+
45
+ from renderers import build_training_sample
46
+ from renderers.glm45 import GLM45Renderer
47
+ from renderers.glm5 import GLM5Renderer
48
+ from tests.reference_rendering import render_reference
49
+
50
+
51
+ def test_glm_tool_stop_ownership(model_name, tokenizer, renderer):
52
+ """The first emitted observation closes the assistant, even after reordering."""
53
+ if not isinstance(renderer, (GLM5Renderer, GLM45Renderer)):
54
+ pytest.skip("GLM turn-closing role markers")
55
+ messages = [
56
+ {"role": "user", "content": "Look up both"},
57
+ {
58
+ "role": "assistant",
59
+ "content": "",
60
+ "tool_calls": [
61
+ {
62
+ "id": f"c{i}",
63
+ "type": "function",
64
+ "function": {"name": f"f{i}", "arguments": {}},
65
+ }
66
+ for i in (1, 2)
67
+ ],
68
+ },
69
+ {
70
+ "role": "tool",
71
+ "content": "Second result",
72
+ "tool_call_id": "c2",
73
+ "name": "f2",
74
+ },
75
+ {"role": "tool", "content": "First result", "tool_call_id": "c1", "name": "f1"},
76
+ {"role": "assistant", "content": "Done"},
77
+ ]
78
+ rendered = renderer.render(messages)
79
+ assert rendered.token_ids == render_reference(tokenizer, messages)
80
+ observation = tokenizer.encode("<|observation|>", add_special_tokens=False)[0]
81
+ positions = [
82
+ k for k, token in enumerate(rendered.token_ids) if token == observation
83
+ ]
84
+ first, *rest = positions
85
+ assert rendered.message_indices[first] == 1
86
+ assert rendered.sampled_mask[first] and rendered.is_content[first]
87
+ assert not any(rendered.sampled_mask[k] for k in rest)
88
+ assert not rendered.content_mask_for_roles({"tool"})[first]
89
+ span = rendered.message_token_spans()[1]
90
+ assert span is not None
91
+ assert all(index == 1 for index in rendered.message_indices[span[0] : span[1]])
92
+ for selected in (False, True):
93
+ sample = build_training_sample(
94
+ renderer,
95
+ messages,
96
+ message_loss_mask=[0, selected, not selected, not selected, 0],
97
+ content_sft_roles={"tool"},
98
+ ensure_final_stop=True,
99
+ )
100
+ assert sample.loss_mask[first] == selected
101
+ for k, is_tool_body in enumerate(rendered.content_mask_for_roles({"tool"})):
102
+ if is_tool_body:
103
+ assert sample.loss_mask[k] == (not selected)
104
+ assert not sample.loss_mask[-1]
105
+
43
106
 
44
107
  def test_message_indices_in_range(model_name, renderer):
45
108
  """Every emitted token's ``message_indices`` must be in
@@ -41,6 +41,8 @@ def test_sampled_mask_excludes_user_and_system(model_name, renderer):
41
41
  {"role": "system", "content": "You are helpful."},
42
42
  {"role": "user", "content": "Hi"},
43
43
  {"role": "assistant", "content": "Hello!"},
44
+ {"role": "user", "content": "Another question"},
45
+ {"role": "assistant", "content": "Another answer"},
44
46
  ]
45
47
  rendered = renderer.render(msgs)
46
48
  if not rendered.sampled_mask: