diff --git a/verifiers/v1/__init__.py b/verifiers/v1/__init__.py index 547ad39d2b..d7a479cb15 100644 --- a/verifiers/v1/__init__.py +++ b/verifiers/v1/__init__.py @@ -138,6 +138,7 @@ Tool, ToolCall, ToolMessage, + TopLogprobs, TurnTokens, Usage, UserMessage, @@ -260,6 +261,7 @@ "Branch", "TurnTokens", "SamplingMask", + "TopLogprobs", "Timing", "TimeSpan", "TimeSplit", diff --git a/verifiers/v1/clients/train.py b/verifiers/v1/clients/train.py index 672c9823bb..ca22f5a0bb 100644 --- a/verifiers/v1/clients/train.py +++ b/verifiers/v1/clients/train.py @@ -28,6 +28,7 @@ SamplingMask, Tool, ToolCall, + TopLogprobs, TurnTokens, Usage, WeightVersionSpan, @@ -185,6 +186,13 @@ def response_from_generate( weight_versions=[WeightVersionSpan(**span) for span in spans] if (spans := result.get("weight_versions")) else None, + top_logprobs=( + TopLogprobs.from_top_logprobs( + result["completion_top_ids"], result["completion_top_logprobs"] + ) + if result.get("completion_top_ids") is not None + else None + ), ), ) diff --git a/verifiers/v1/graph.py b/verifiers/v1/graph.py index 8380d49d85..a75a8febdf 100644 --- a/verifiers/v1/graph.py +++ b/verifiers/v1/graph.py @@ -24,7 +24,7 @@ import hashlib import json import time -from dataclasses import dataclass +from dataclasses import dataclass, fields from typing import TYPE_CHECKING, Any import numpy as np @@ -33,6 +33,7 @@ ConfigDict, Field, FieldSerializationInfo, + ValidationInfo, field_serializer, field_validator, ) @@ -49,6 +50,7 @@ TextContentPart, Tool, ToolMessage, + TopLogprobs, WeightVersionSpan, ) @@ -170,6 +172,9 @@ class MessageNode(BaseModel): completion, as vLLM returns them). Assistant nodes only; None when the engine labels none.""" + 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 @@ -210,26 +215,28 @@ 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("sampling_mask") - def serialize_sampling_mask(self, mask: SamplingMask | None) -> dict | None: - if mask is None: + @field_serializer("sampling_mask", "top_logprobs") + def serialize_sampling_arrays( + self, payload: SamplingMask | TopLogprobs | None + ) -> dict | None: + if payload is None: return None return { - "ids": _encode_ndarray(mask.ids), - "counts": _encode_ndarray(mask.counts), + field.name: _encode_ndarray(getattr(payload, field.name)) + for field in fields(payload) } - @field_validator("sampling_mask", mode="before") + @field_validator("sampling_mask", "top_logprobs", mode="before") @classmethod - def deserialize_sampling_mask(cls, value: Any) -> SamplingMask | None: - if value is None or isinstance(value, SamplingMask): + def deserialize_sampling_arrays( + cls, value: Any, info: ValidationInfo + ) -> SamplingMask | TopLogprobs | None: + kind = SamplingMask if info.field_name == "sampling_mask" else TopLogprobs + if value is None or isinstance(value, kind): return value if isinstance(value, dict): - return SamplingMask( - ids=_decode_ndarray(value["ids"]), - counts=_decode_ndarray(value["counts"]), - ) - raise TypeError(f"cannot build SamplingMask from {type(value).__name__}") + return kind(**{key: _decode_ndarray(array) for key, array in value.items()}) + raise TypeError(f"cannot build {kind.__name__} from {type(value).__name__}") def _canonical_tool_arguments(arguments: str) -> str: @@ -823,6 +830,23 @@ 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 (payload.counts < 0).any() + 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 @@ -1020,6 +1044,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 51e432afb1..28540a472b 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,24 @@ 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.""" + heads = [n.top_logprobs for n in self.nodes if n.top_logprobs is not None] + if not heads: + return None + counts = [] + for node in self.nodes: + rows = np.zeros(len(node.mask), dtype=np.int32) + if node.top_logprobs is not None: + rows[np.nonzero(node.mask)[0]] = node.top_logprobs.counts + counts.append(rows) + return TopLogprobs( + ids=np.concatenate([head.ids for head in heads]), + logprobs=np.concatenate([head.logprobs for head in heads]), + counts=np.concatenate(counts), + ) + @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 167c74f9d5..5bd4607a53 100644 --- a/verifiers/v1/types.py +++ b/verifiers/v1/types.py @@ -249,6 +249,41 @@ def _decode_packed(packed: dict[str, Any], dtype: type) -> np.ndarray: return array.astype(dtype, copy=False) +@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") + return cls( + ids=np.fromiter((token for row in ids for token in row), dtype=np.int32), + logprobs=np.fromiter( + (p for row in logprobs for p in row), dtype=np.float32 + ), + counts=counts, + ) + + class TurnTokens(BaseModel): """Training tokens from renderer tokenization or provider-returned token IDs.""" @@ -288,6 +323,10 @@ class TurnTokens(BaseModel): # attributed to the assistant node by the turn's `commit`, then dropped. weight_versions: list[WeightVersionSpan] | 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): id: str