From 90b6ca0f094a77977a6fff5e37d1c39a3308b3c1 Mon Sep 17 00:00:00 2001 From: Phi Date: Thu, 8 Oct 2026 00:46:04 +0000 Subject: [PATCH] feat(trace): carry sampler probability heads to training branches --- verifiers/v1/__init__.py | 2 ++ verifiers/v1/clients/train.py | 8 ++++++ verifiers/v1/graph.py | 46 +++++++++++++++++++++++++++++++++++ verifiers/v1/trace.py | 32 ++++++++++++++++++++++++ verifiers/v1/types.py | 43 ++++++++++++++++++++++++++++++++ 5 files changed, 131 insertions(+) diff --git a/verifiers/v1/__init__.py b/verifiers/v1/__init__.py index 78120a9f41..4e71a5376f 100644 --- a/verifiers/v1/__init__.py +++ b/verifiers/v1/__init__.py @@ -131,6 +131,7 @@ SamplingConfig, SamplingMask, SystemMessage, + TopLogprobs, TextContentPart, Tool, ToolCall, @@ -257,6 +258,7 @@ "Branch", "TurnTokens", "SamplingMask", + "TopLogprobs", "Timing", "TimeSpan", "TimeSplit", diff --git a/verifiers/v1/clients/train.py b/verifiers/v1/clients/train.py index 025679711f..bc17464caf 100644 --- a/verifiers/v1/clients/train.py +++ b/verifiers/v1/clients/train.py @@ -28,6 +28,7 @@ SamplingMask, Tool, ToolCall, + TopLogprobs, TurnTokens, Usage, ) @@ -181,6 +182,13 @@ def response_from_generate( sampling_mask=SamplingMask.from_sampling_mask(mask) if (mask := result.get("sampling_mask")) else None, + top_logprobs=( + TopLogprobs.from_top_logprobs( + result["completion_top_ids"], result["completion_top_logprobs"] + ) + if result.get("completion_top_ids") + else None + ), ), ) diff --git a/verifiers/v1/graph.py b/verifiers/v1/graph.py index 270f087d3e..651c13b78a 100644 --- a/verifiers/v1/graph.py +++ b/verifiers/v1/graph.py @@ -48,6 +48,7 @@ TextContentPart, Tool, ToolMessage, + TopLogprobs, ) if TYPE_CHECKING: @@ -164,6 +165,9 @@ class MessageNode(BaseModel): nodes only. The arrays serialize as raw-byte `__nd__` dictionaries. """ + top_logprobs: SkipJsonSchema[TopLogprobs | None] = None + """Sampler top-k probabilities for sampled tokens; unnormalized over the head.""" + model_config = ConfigDict(arbitrary_types_allowed=True) @property @@ -225,6 +229,29 @@ def deserialize_sampling_mask(cls, value: Any) -> SamplingMask | None: ) raise TypeError(f"cannot build SamplingMask from {type(value).__name__}") + @field_serializer("top_logprobs") + def serialize_top_logprobs(self, head: TopLogprobs | None) -> dict | None: + if head is None: + return None + return { + "ids": _encode_ndarray(head.ids), + "logprobs": _encode_ndarray(head.logprobs), + "counts": _encode_ndarray(head.counts), + } + + @field_validator("top_logprobs", mode="before") + @classmethod + def deserialize_top_logprobs(cls, value: Any) -> TopLogprobs | None: + if value is None or isinstance(value, TopLogprobs): + return value + if isinstance(value, dict): + return TopLogprobs( + ids=_decode_ndarray(value["ids"]), + logprobs=_decode_ndarray(value["logprobs"]), + counts=_decode_ndarray(value["counts"]), + ) + raise TypeError(f"cannot build TopLogprobs from {type(value).__name__}") + def _canonical_tool_arguments(arguments: str) -> str: # Ignore JSON key order and whitespace when hashing equivalent tool calls. @@ -775,6 +802,22 @@ def _project_prompt_attribution( return projected_spans, projected_is_content +def _attribute_top_logprobs( + trace: Trace, assistant_id: int, payload: TopLogprobs | None +) -> None: + """Attach a completion-aligned top-k sampling head to the assistant node.""" + if payload is None: + return + node = trace.nodes[assistant_id] + if ( + len(payload.counts) != sum(node.mask) + or int(payload.counts.sum()) != len(payload.ids) + or len(payload.ids) != len(payload.logprobs) + ): + raise ValueError("Top logprobs must align with every sampled node token") + node.top_logprobs = payload + + def _commit_turn(turn: PendingTurn, response: Response) -> int: trace = turn.trace prompt = turn.prompt @@ -970,6 +1013,9 @@ def _commit_turn(turn: PendingTurn, response: Response) -> int: _attribute_sampling_mask( trace, assistant_id, tokens.sampling_mask if tokens else None ) + _attribute_top_logprobs( + trace, assistant_id, tokens.top_logprobs if tokens else None + ) return assistant_id diff --git a/verifiers/v1/trace.py b/verifiers/v1/trace.py index c3b4f6c58d..71f01c19c0 100644 --- a/verifiers/v1/trace.py +++ b/verifiers/v1/trace.py @@ -31,6 +31,7 @@ SamplingMask, Tool, ToolMessage, + TopLogprobs, Usage, content_text, ) @@ -44,6 +45,7 @@ "__all__": { "routed_experts", "sampling_mask", + "top_logprobs", } } } @@ -347,6 +349,36 @@ def sampling_mask(self) -> SamplingMask | None: ) return SamplingMask(ids=ids, counts=np.concatenate(counts_parts)) + @property + def top_logprobs(self) -> TopLogprobs | None: + """Top-k sampling heads aligned to this branch's token ids.""" + if all(n.top_logprobs is None for n in self.nodes): + return None + # Attribution validates each head against the node's sampled positions. + ids_parts: list[np.ndarray] = [] + logprobs_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.top_logprobs is not None and len(node.top_logprobs.counts): + counts[np.nonzero(node.mask)[0]] = node.top_logprobs.counts + ids_parts.append(node.top_logprobs.ids) + logprobs_parts.append(node.top_logprobs.logprobs) + counts_parts.append(counts) + ids = ( + np.concatenate(ids_parts).astype(np.int32, copy=False) + if ids_parts + else np.zeros(0, dtype=np.int32) + ) + logprobs = ( + np.concatenate(logprobs_parts).astype(np.float32, copy=False) + if logprobs_parts + else np.zeros(0, dtype=np.float32) + ) + return TopLogprobs( + ids=ids, logprobs=logprobs, counts=np.concatenate(counts_parts) + ) + @property def usage(self) -> Usage | None: return Usage.aggregate(c.usage for c in self.calls if c.usage is not None) diff --git a/verifiers/v1/types.py b/verifiers/v1/types.py index f9edd7d0cf..5e96e2f6e7 100644 --- a/verifiers/v1/types.py +++ b/verifiers/v1/types.py @@ -216,6 +216,46 @@ def from_sampling_mask(cls, sampling_mask: list[list[int]]) -> "SamplingMask": return cls(ids=ids, counts=counts) +@dataclass +class TopLogprobs: + """Top-k sampling head stored as flat `ids`, `logprobs` and `counts` arrays. + + Each row holds one completion token's top-k sampler candidates in descending + probability order, independent of the sampled action. Row boundaries come from + `counts`; `ids` and `logprobs` share the same flat layout. + """ + + ids: Any + logprobs: Any + counts: Any + + @classmethod + def from_top_logprobs( + cls, ids: list[list[int]], logprobs: list[list[float]] + ) -> "TopLogprobs": + if len(ids) != len(logprobs): + raise ValueError("top ids and logprobs must have one row per token") + counts = np.fromiter( + (len(row) for row in ids), + dtype=np.int32, + count=len(ids), + ) + if any(len(ids[i]) != len(logprobs[i]) for i in range(len(ids))): + raise ValueError("top ids and logprobs rows must be the same length") + flat = int(counts.sum()) + ids_flat = ( + np.concatenate([np.asarray(row, dtype=np.int32) for row in ids]) + if flat + else np.empty(0, dtype=np.int32) + ) + logprobs_flat = ( + np.concatenate([np.asarray(row, dtype=np.float32) for row in logprobs]) + if flat + else np.empty(0, dtype=np.float32) + ) + return cls(ids=ids_flat, logprobs=logprobs_flat, counts=counts) + + class TurnTokens(BaseModel): """Training tokens from renderer tokenization or provider-returned token IDs.""" @@ -251,6 +291,9 @@ class TurnTokens(BaseModel): # 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) + # Transient carrier (excluded): per-completion-token top-k sampling heads, + # attributed to the assistant node by the turn's `commit`, then dropped. + top_logprobs: TopLogprobs | None = Field(default=None, exclude=True) class Response(BaseModel):