diff --git a/pyproject.toml b/pyproject.toml index 80a9dd0f3f..f770af755f 100644 --- a/pyproject.toml +++ b/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", ] diff --git a/uv.lock b/uv.lock index 3238932633..f051199b71 100644 --- a/uv.lock +++ b/uv.lock @@ -32,7 +32,7 @@ ruff = "2026-07-28T00:00:00Z" name = "nemo-gym" version = "0.4.0" requires-dist = ["aiohttp>=3.14.1", "fastapi", "gprof2dot", "hydra-core", "itsdangerous", "mcp>=1.27,<2", "omegaconf", "openai>=2.9", "orjson", "pandas", "pydantic", "pydot", "ray>=2.55.1", "requests", "rich", "typing-extensions", "uvicorn", "wandb", "yappi"] -requires-python = ">=3.11, <3.14" +requires-python = ">=3.12" [[package]] name = "aiofile" @@ -3838,7 +3838,7 @@ wheels = [ [[package]] name = "renderers" -version = "0.1.11" +version = "0.1.12.dev2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "jinja2" }, @@ -3849,9 +3849,9 @@ dependencies = [ { name = "prime-pydantic-config" }, { name = "tiktoken" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/6f/37/3441c75b45fc818ad1881d7837ff073f92ce4e3ab3bfd1640ad3062d6bb0/renderers-0.1.11.tar.gz", hash = "sha256:575f02e44200b525dc5c4fac207148d18bacadc083eca783068d4195d21f8960", size = 392707, upload-time = "2026-08-31T17:36:02.706Z" } +sdist = { url = "https://files.pythonhosted.org/packages/09/86/4016c7ea6cbc2c3258265b11766d3018b805563f3b7fbcf7edd6df38690d/renderers-0.1.12.dev2.tar.gz", hash = "sha256:7cc764dec4d0c5c48a20a610df52ec3be2de3d4f0c65cb5d709487e0bad761a8", size = 394034, upload-time = "2026-09-01T21:31:06.753Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/c4/bb/4e2d84f717a0266dcb5c26c6cd1f1beae46582de41cf323ee29e57fcb45a/renderers-0.1.11-py3-none-any.whl", hash = "sha256:8fac795d32f8aeacffe6dbb863d60f1b3437a06389c27593ebabe66cec9bdc45", size = 230033, upload-time = "2026-08-31T17:36:00.652Z" }, + { url = "https://files.pythonhosted.org/packages/7b/89/4be1a4ce35e85ecbcccfd8bc7e869fb23424c7b93696e29149e555bddae1/renderers-0.1.12.dev2-py3-none-any.whl", hash = "sha256:9e4043d87982a9f51ca151071baac7f30053e9c896356957e228fa4791cfd793", size = 231100, upload-time = "2026-09-01T21:31:04.936Z" }, ] [package.optional-dependencies] diff --git a/verifiers/v1/__init__.py b/verifiers/v1/__init__.py index d3012c8c5e..733fbeab03 100644 --- a/verifiers/v1/__init__.py +++ b/verifiers/v1/__init__.py @@ -116,7 +116,6 @@ ContentPart, ImageUrlContentPart, ImageUrlSource, - KeptTokens, Message, MessageContent, Messages, @@ -124,6 +123,7 @@ Response, Sampling, SamplingConfig, + SamplingMask, SystemMessage, TextContentPart, Tool, @@ -249,7 +249,7 @@ "MessageNode", "Branch", "TurnTokens", - "KeptTokens", + "SamplingMask", "Timing", "TimeSpan", "TimeSplit", diff --git a/verifiers/v1/clients/train.py b/verifiers/v1/clients/train.py index fd49504153..e803610231 100644 --- a/verifiers/v1/clients/train.py +++ b/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, ), ) diff --git a/verifiers/v1/graph.py b/verifiers/v1/graph.py index 8c787adf25..af1dcee658 100644 --- a/verifiers/v1/graph.py +++ b/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 diff --git a/verifiers/v1/trace.py b/verifiers/v1/trace.py index e8b559b6ad..66022215f5 100644 --- a/verifiers/v1/trace.py +++ b/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: diff --git a/verifiers/v1/types.py b/verifiers/v1/types.py index 03cf38b059..4f97e95b63 100644 --- a/verifiers/v1/types.py +++ b/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):