verifiers 0.2.1.dev12__tar.gz → 0.2.1.dev19__tar.gz

{verifiers-0.2.1.dev12 → verifiers-0.2.1.dev19}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: verifiers
-Version: 0.2.1.dev12
+Version: 0.2.1.dev19
 Summary: Verifiers: Environments for LLM Reinforcement Learning
 Project-URL: Homepage, https://github.com/primeintellect-ai/verifiers
 Project-URL: Documentation, https://github.com/primeintellect-ai/verifiers

{verifiers-0.2.1.dev12 → verifiers-0.2.1.dev19}/verifiers/v1/__init__.py

@@ -114,6 +114,7 @@ from verifiers.v1.types import (
     ID,
     ImageUrlContentPart,
     ImageUrlSource,
+    KeptTokens,
     Message,
     MessageContent,
     Messages,
@@ -166,6 +167,7 @@ __all__ = [
     "MessageNode",
     "Branch",
     "TurnTokens",
+    "KeptTokens",
     "Timing",
     "TimeSpan",
     "Error",

{verifiers-0.2.1.dev12 → verifiers-0.2.1.dev19}/verifiers/v1/clients/train.py

@@ -25,6 +25,7 @@ from verifiers.v1.graph import PendingTurn
 from verifiers.v1.types import (
     AssistantMessage,
     FinishReason,
+    KeptTokens,
     Response,
     SamplingConfig,
     Tool,
@@ -149,6 +150,9 @@ def response_from_generate(
             is_content=attribution.is_content if attribution is not None else None,
             multi_modal_data=result.get("multi_modal_data"),
             routed_experts=result.get("routed_experts"),
+            kept_tokens=KeptTokens(**kept)
+            if (kept := result.get("kept_tokens"))
+            else None,
         ),
     )
 

{verifiers-0.2.1.dev12 → verifiers-0.2.1.dev19}/verifiers/v1/graph.py

@@ -31,6 +31,7 @@ from renderers.base import MultiModalData, PlaceholderRange, RenderedTokens
 from verifiers.v1.types import (
     AssistantMessage,
     FinishReason,
+    KeptTokens,
     Message,
     Response,
     StrictBaseModel,
@@ -118,6 +119,12 @@ class MessageNode(StrictBaseModel):
     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: 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."""
 
     model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)
 
@@ -166,18 +173,40 @@ class MessageNode(StrictBaseModel):
         )
 
     @field_serializer("routed_experts")
-    def serialize_routed_experts(self, re: np.ndarray | None) -> dict | None:
-        """uint8 routing array -> raw-bytes `__nd__` dict so it rides the wire (numpy can't JSON)."""
-        return None if re is None else _encode_ndarray(re)
+    def serialize_ndarray_field(self, arr: np.ndarray | None) -> dict | None:
+        """Integer array -> raw-bytes `__nd__` dict so it rides the wire (numpy can't JSON)."""
+        return None if arr is None else _encode_ndarray(arr)
 
     @field_validator("routed_experts", mode="before")
     @classmethod
-    def deserialize_routed_experts(cls, value: Any) -> np.ndarray | None:
+    def deserialize_ndarray_field(cls, value: Any) -> np.ndarray | None:
         if value is None or isinstance(value, np.ndarray):
             return value
         if isinstance(value, dict) and value.get("__nd__"):
             return _decode_ndarray(value)
-        raise TypeError(f"cannot build routed_experts from {type(value).__name__}")
+        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:
+            return None
+        return {
+            "ids": _encode_ndarray(kept.ids),
+            "counts": _encode_ndarray(kept.counts),
+        }
+
+    @field_validator("kept_tokens", mode="before")
+    @classmethod
+    def deserialize_kept_tokens(cls, value: Any) -> KeptTokens | None:
+        if value is None or isinstance(value, KeptTokens):
+            return value
+        if isinstance(value, dict):
+            return KeptTokens(
+                ids=_decode_ndarray(value["ids"]),
+                counts=_decode_ndarray(value["counts"]),
+            )
+        raise TypeError(f"cannot build KeptTokens from {type(value).__name__}")
 
 
 def _canonical_tool_arguments(arguments: str) -> str:
@@ -434,6 +463,23 @@ def _attribute_routed_experts(
         off = end
 
 
+def _attribute_kept_tokens(
+    trace: Trace, assistant_id: int, payload: KeptTokens | 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."""
+    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):
+        return
+    # Own the buffers — the payload views reference the turn's response bytes.
+    node.kept_tokens = KeptTokens(ids=ids.copy(), counts=counts.copy())
+
+
 def _commit_turn(turn: PendingTurn, response: Response) -> None:
     trace = turn.trace
     prompt = turn.prompt
@@ -540,6 +586,10 @@ def _commit_turn(turn: PendingTurn, response: Response) -> None:
         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)
+
 
 # --- walking the graph (views) ---------------------------------------------------------
 

{verifiers-0.2.1.dev12 → verifiers-0.2.1.dev19}/verifiers/v1/trace.py

@@ -22,6 +22,7 @@ from verifiers.v1.state import State, StateT
 from verifiers.v1.task import DataT, WireTaskData
 from verifiers.v1.types import (
     AssistantMessage,
+    KeptTokens,
     Messages,
     StrictBaseModel,
     Tool,
@@ -139,6 +140,32 @@ class Branch(StrictBaseModel):
         total = sum(len(n.token_ids) for n in nodes)
         return merged if merged.shape[0] == total else None
 
+    @property
+    def kept_tokens(self) -> KeptTokens | None:
+        """The branch's kept-set sampling masks: `counts` is int32 aligned 1:1 with
+        `token_ids` (0 = no mask, safe under partial coverage — unlike `routed_experts`
+        this is not all-or-nothing), `ids` the flat int32 concatenation of the kept
+        sets in position order. None when no node carries kept-set data."""
+        if all(n.kept_tokens 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).
+        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)
+            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))
+
     @property
     def num_total_tokens(self) -> int:
         return sum(len(n.token_ids) for n in self.nodes)
@@ -172,7 +199,13 @@ class Branch(StrictBaseModel):
 
 
 _NODE_DUMP_EXCLUDE: dict = {
-    "nodes": {"__all__": {"multi_modal_data", "routed_experts"}}
+    "nodes": {
+        "__all__": {
+            "multi_modal_data",
+            "routed_experts",
+            "kept_tokens",
+        }
+    }
 }
 """Raw tensor fields kept on the msgpack wire but excluded from JSON records."""
 

{verifiers-0.2.1.dev12 → verifiers-0.2.1.dev19}/verifiers/v1/types.py

@@ -1,4 +1,5 @@
 from collections.abc import Iterable
+from dataclasses import dataclass
 from typing import Annotated, Any, Literal
 
 from pydantic import AfterValidator, AliasChoices, BaseModel, ConfigDict, Field
@@ -178,6 +179,18 @@ class RoutedExperts(TypedDict):
     start: int
 
 
+@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`)."""
+
+    ids: Any
+    counts: Any
+
+
 class TurnTokens(StrictBaseModel):
     """Training tokens from renderer tokenization or provider-returned token IDs."""
 
@@ -200,6 +213,11 @@ class TurnTokens(StrictBaseModel):
     # 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)
 
 
 class Response(StrictBaseModel):