feat: sampling mask pass-through

pyproject.toml

@@ -46,7 +46,7 @@ dependencies = [
     "uvloop>=0.21.0; sys_platform != 'win32' and sys_platform != 'cygwin' and platform_python_implementation != 'PyPy'",
     "loguru>=0.7.0",
     "tomli-w>=1.0.0",
-    "renderers[multimodal]>=0.1.11",
+    "renderers[multimodal]>=0.1.12.dev2",
     "typing_extensions>=4.12.2",
     "mcp==2.0.0",
 ]

verifiers/v1/__init__.py

@@ -116,14 +116,14 @@
     ContentPart,
     ImageUrlContentPart,
     ImageUrlSource,
-    KeptTokens,
     Message,
     MessageContent,
     Messages,
     Request,
     Response,
     Sampling,
     SamplingConfig,
+    SamplingMask,
     SystemMessage,
     TextContentPart,
     Tool,
@@ -249,7 +249,7 @@
     "MessageNode",
     "Branch",
     "TurnTokens",
-    "KeptTokens",
+    "SamplingMask",
     "Timing",
     "TimeSpan",
     "TimeSplit",

verifiers/v1/clients/train.py

@@ -23,9 +23,9 @@
 from verifiers.v1.types import (
     AssistantMessage,
     FinishReason,
-    KeptTokens,
     Response,
     SamplingConfig,
+    SamplingMask,
     Tool,
     ToolCall,
     TurnTokens,
@@ -161,8 +161,8 @@ def response_from_generate(
             multi_modal_data=result.get("multi_modal_data"),
             mm_token_type_id_map=mm_token_type_id_map,
             routed_experts=result.get("routed_experts"),
-            kept_tokens=KeptTokens(**kept)
-            if (kept := result.get("kept_tokens"))
+            sampling_mask=SamplingMask.from_sampling_mask(mask)
+            if (mask := result.get("sampling_mask"))
             else None,
         ),
     )

verifiers/v1/graph.py

@@ -35,9 +35,9 @@
 from verifiers.v1.semantic import ParentLink
 from verifiers.v1.types import (
     AssistantMessage,
-    KeptTokens,
     Message,
     Response,
+    SamplingMask,
     TextContentPart,
     Tool,
     ToolMessage,
@@ -141,12 +141,12 @@ class MessageNode(BaseModel):
     the turn's `generate` payload by `_attribute_routed_experts`; `Branch.routed_experts`
     concatenates these along the path into the trainer's router-replay input. Rides the wire as
     a raw-bytes `__nd__` dict; kept off disk by the dump-site `exclude` in prime-rl."""
-    kept_tokens: SkipJsonSchema[KeptTokens | None] = None
-    """Kept-set sampling masks for this node's sampled tokens, decoded: `ids` flat int32
-    in position order, `counts` the per-token kept-set sizes (aligned with `logprobs`;
-    0 = no mask). Assistant nodes only; consumed via `Branch.kept_tokens` for
-    sampling-replay training. Rides the wire as raw-bytes `__nd__` dicts; kept off disk
-    by the dump-site `exclude` in prime-rl."""
+    sampling_mask: SkipJsonSchema[SamplingMask | None] = None
+    """Sampling masks for this node's sampled tokens.
+
+    `ids` stores the flat token ids and `counts` stores each token's row size. Assistant
+    nodes only. The arrays serialize as raw-byte `__nd__` dictionaries.
+    """
 
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
@@ -208,27 +208,26 @@ def deserialize_ndarray_field(cls, value: Any) -> np.ndarray | None:
             return _decode_ndarray(value)
         raise TypeError(f"cannot build ndarray field from {type(value).__name__}")
 
-    @field_serializer("kept_tokens")
-    def serialize_kept_tokens(self, kept: KeptTokens | None) -> dict | None:
-        """`KeptTokens` -> dict of raw-bytes `__nd__` entries so the arrays ride the wire."""
-        if kept is None:
+    @field_serializer("sampling_mask")
+    def serialize_sampling_mask(self, mask: SamplingMask | None) -> dict | None:
+        if mask is None:
             return None
         return {
-            "ids": _encode_ndarray(kept.ids),
-            "counts": _encode_ndarray(kept.counts),
+            "ids": _encode_ndarray(mask.ids),
+            "counts": _encode_ndarray(mask.counts),
         }
 
-    @field_validator("kept_tokens", mode="before")
+    @field_validator("sampling_mask", mode="before")
     @classmethod
-    def deserialize_kept_tokens(cls, value: Any) -> KeptTokens | None:
-        if value is None or isinstance(value, KeptTokens):
+    def deserialize_sampling_mask(cls, value: Any) -> SamplingMask | None:
+        if value is None or isinstance(value, SamplingMask):
             return value
         if isinstance(value, dict):
-            return KeptTokens(
+            return SamplingMask(
                 ids=_decode_ndarray(value["ids"]),
                 counts=_decode_ndarray(value["counts"]),
             )
-        raise TypeError(f"cannot build KeptTokens from {type(value).__name__}")
+        raise TypeError(f"cannot build SamplingMask from {type(value).__name__}")
 
 
 def _canonical_tool_arguments(arguments: str) -> str:
@@ -612,21 +611,18 @@ def _attribute_routed_experts(
         off = end
 
 
-def _attribute_kept_tokens(
-    trace: Trace, assistant_id: int, payload: KeptTokens | None
+def _attribute_sampling_mask(
+    trace: Trace, assistant_id: int, payload: SamplingMask | None
 ) -> None:
-    """Attach this turn's kept-set sampling masks to the assistant node (the payload
-    covers exactly the turn's completion tokens, so no path arithmetic). A payload
-    that doesn't line up with the node's sampled tokens is dropped, not misaligned."""
+    """Attach a completion-aligned sampling mask to the assistant node."""
     if payload is None:
         return
-    counts = np.frombuffer(binascii.a2b_base64(payload.counts), dtype=np.int32)
-    ids = np.frombuffer(binascii.a2b_base64(payload.ids), dtype=np.int32)
     node = trace.nodes[assistant_id]
-    if len(counts) != sum(node.mask) or int(counts.sum()) != len(ids):
+    if len(payload.counts) != sum(node.mask) or int(payload.counts.sum()) != len(
+        payload.ids
+    ):
         return
-    # Own the buffers — the payload views reference the turn's response bytes.
-    node.kept_tokens = KeptTokens(ids=ids.copy(), counts=counts.copy())
+    node.sampling_mask = payload
 
 
 def _commit_turn(turn: PendingTurn, response: Response) -> int:
@@ -775,9 +771,10 @@ def _commit_turn(turn: PendingTurn, response: Response) -> int:
         trace, new_node_ids, path_len, tokens.routed_experts if tokens else None
     )
 
-    # Attribute this turn's kept-set sampling masks onto the assistant node (they are
-    # completion-aligned, so only the sampled node carries them).
-    _attribute_kept_tokens(trace, assistant_id, tokens.kept_tokens if tokens else None)
+    # Sampling masks are completion-aligned, so only the sampled node carries them.
+    _attribute_sampling_mask(
+        trace, assistant_id, tokens.sampling_mask if tokens else None
+    )
 
     return assistant_id
 

verifiers/v1/trace.py

@@ -29,10 +29,10 @@
 from verifiers.v1.types import (
     AssistantMessage,
     FinishReason,
-    KeptTokens,
     Message,
     Messages,
     Sampling,
+    SamplingMask,
     Tool,
     ToolMessage,
     Usage,
@@ -48,7 +48,7 @@
         "__all__": {
             "multi_modal_data",
             "routed_experts",
-            "kept_tokens",
+            "sampling_mask",
         }
     }
 }
@@ -343,28 +343,25 @@ def routed_experts(self) -> np.ndarray | None:
         return merged if merged.shape[0] == total else None
 
     @property
-    def kept_tokens(self) -> KeptTokens | None:
-        """int32 kept-set `counts` aligned 1:1 with `token_ids` plus flat `ids` in
-        position order; partial data scatters as 0 counts, no data returns None."""
-        if all(n.kept_tokens is None for n in self.nodes):
+    def sampling_mask(self) -> SamplingMask | None:
+        """Sampling masks aligned to this branch's token ids."""
+        if all(n.sampling_mask is None for n in self.nodes):
             return None
-        # `_attribute_kept_tokens` validates counts/ids against the node's sampled
-        # tokens before setting the field, so this is a straight scatter+concat
-        # (a corrupted node would fail loudly on the scatter shape mismatch).
+        # Attribution validates each mask against the node's sampled positions.
         ids_parts: list[np.ndarray] = []
         counts_parts: list[np.ndarray] = []
         for node in self.nodes:
             counts = np.zeros(len(node.mask), dtype=np.int32)
-            if node.kept_tokens is not None and len(node.kept_tokens.counts):
-                counts[np.nonzero(node.mask)[0]] = node.kept_tokens.counts
-                ids_parts.append(node.kept_tokens.ids)
+            if node.sampling_mask is not None and len(node.sampling_mask.counts):
+                counts[np.nonzero(node.mask)[0]] = node.sampling_mask.counts
+                ids_parts.append(node.sampling_mask.ids)
             counts_parts.append(counts)
         ids = (
             np.concatenate(ids_parts).astype(np.int32, copy=False)
             if ids_parts
             else np.zeros(0, dtype=np.int32)
         )
-        return KeptTokens(ids=ids, counts=np.concatenate(counts_parts))
+        return SamplingMask(ids=ids, counts=np.concatenate(counts_parts))
 
     @property
     def usage(self) -> Usage | None:

verifiers/v1/types.py

@@ -2,6 +2,7 @@
 from dataclasses import dataclass
 from typing import Annotated, Any, Literal
 
+import numpy as np
 from pydantic import AliasChoices, BaseModel, ConfigDict, Field
 from renderers.base import MultiModalData
 from typing_extensions import TypedDict
@@ -184,16 +185,30 @@ class RoutedExperts(TypedDict):
 
 
 @dataclass
-class KeptTokens:
-    """Kept-set sampling masks for sampling replay: `ids` (every kept set concatenated
-    in position order) and `counts` (kept-set size per completion token; 0 = no usable
-    mask). Base64 blobs straight off the `generate` response on the `TurnTokens`
-    carrier; decoded to flat int32 arrays on `MessageNode` (`len(ids) == sum(counts)`,
-    row boundaries recovered from `counts`)."""
+class SamplingMask:
+    """Sampling masks stored as flat int32 `ids` and `counts` arrays.
+
+    Each row contains the token ids that survived sampling filters for one completion
+    token. Row boundaries are recovered from `counts`.
+    """
 
     ids: Any
     counts: Any
 
+    @classmethod
+    def from_sampling_mask(cls, sampling_mask: list[list[int]]) -> "SamplingMask":
+        counts = np.fromiter(
+            (len(row) for row in sampling_mask),
+            dtype=np.int32,
+            count=len(sampling_mask),
+        )
+        ids = (
+            np.concatenate([np.asarray(row, dtype=np.int32) for row in sampling_mask])
+            if int(counts.sum())
+            else np.empty(0, dtype=np.int32)
+        )
+        return cls(ids=ids, counts=counts)
+
 
 class TurnTokens(BaseModel):
     """Training tokens from renderer tokenization or provider-returned token IDs."""
@@ -221,11 +236,9 @@ class TurnTokens(BaseModel):
     # per token), attributed per node by the turn's `commit` into `MessageNode.routed_experts`,
     # then dropped. None unless the engine ran with `enable_return_routed_experts`.
     routed_experts: RoutedExperts | None = Field(default=None, exclude=True)
-    # Transient carrier (excluded): the kept-set sampling masks from `generate` (token ids
-    # surviving top-p/top-k truncation, per completion token), attributed to the assistant
-    # node by the turn's `commit`, then dropped. None unless the engine ran with
-    # `enable_return_kept_tokens`.
-    kept_tokens: KeptTokens | None = Field(default=None, exclude=True)
+    # Transient carrier (excluded): per-completion-token sampling masks,
+    # attributed to the assistant node by the turn's `commit`, then dropped.
+    sampling_mask: SamplingMask | None = Field(default=None, exclude=True)
 
 
 class Response(BaseModel):