diff --git a/tests/unit/test_graph.py b/tests/unit/test_graph.py index 4c73c34e5c..a80bd495fc 100644 --- a/tests/unit/test_graph.py +++ b/tests/unit/test_graph.py @@ -191,6 +191,84 @@ def test_routed_experts_none_when_absent(): assert trace.branches[-1].routed_experts is None +def _two_turn_sampling_mask_trace(masks: list[vf.SamplingMask]) -> vf.Trace: + trace = vf.Trace( + agent=vf.AgentInfo(config=vf.AgentConfig()), + task=vf.TraceTask(type="Task", data=vf.TaskData(idx=0, prompt="x")), + ) + user = vf.UserMessage(content="u1") + turns = [ + ([user], [10, 11, 12], [20, 21], [(0, 2)]), + ( + [user, vf.AssistantMessage(content="a1"), vf.UserMessage(content="u2")], + [10, 11, 12, 20, 21, 30, 31], + [40, 41], + [(0, 2), None, (5, 7)], + ), + ] + for i, ((messages, prompt_ids, completion_ids, spans), mask) in enumerate( + zip(turns, masks) + ): + graph.prepare_turn(trace, messages).commit( + vf.Response( + id=str(i), + created=0, + model="t", + message=vf.AssistantMessage(content=f"a{i + 1}"), + finish_reason="stop", + tokens=TurnTokens( + prompt_ids=prompt_ids, + completion_ids=completion_ids, + message_spans=spans, + sampling_mask=mask, + ), + ) + ) + return trace + + +def _mask( + rows: list[list[int]], logprobs: list[float] | None = None +) -> vf.SamplingMask: + mask = vf.SamplingMask.from_sampling_mask(rows) + if logprobs is not None: + mask.logprobs = np.array(logprobs, np.float32) + return mask + + +def test_sampling_mask_logprobs_attributed_aggregated_and_round_tripped(): + """Optional sampler logprobs ride along with the mask ids: concatenated on the branch, + kept by the wire round trip, and dropped (None) when any turn lacks them. A mask whose + logprobs do not match its ids is not attributed.""" + first = _mask([[20, 5], [21]], [-0.1, -2.4, 0.0]) + second = _mask([[40], [41, 7, 8]], [0.0, -0.5, -1.0, -3.0]) + trace = _two_turn_sampling_mask_trace([first, second]) + + branch_mask = trace.branches[-1].sampling_mask + assert branch_mask is not None + assert len(branch_mask.counts) == len(trace.branches[-1].token_ids) + assert branch_mask.ids.tolist() == [20, 5, 21, 40, 41, 7, 8] + assert branch_mask.logprobs.tolist() == pytest.approx( + [-0.1, -2.4, 0.0, 0.0, -0.5, -1.0, -3.0] + ) + restored = type(trace).model_validate(trace.model_dump()).branches[-1].sampling_mask + assert restored is not None + for name in ("ids", "counts", "logprobs"): + np.testing.assert_array_equal( + getattr(restored, name), getattr(branch_mask, name) + ) + + mixed = _two_turn_sampling_mask_trace([first, _mask([[40], [41, 7, 8]])]) + assert mixed.branches[-1].sampling_mask.logprobs is None + restored = type(mixed).model_validate(mixed.model_dump()).branches[-1].sampling_mask + assert restored.logprobs is None + assert mixed.branches[-1].sampling_mask.ids.tolist() == branch_mask.ids.tolist() + + mismatched = _mask([[40], [41, 7, 8]], [0.0, -0.5]) + trace = _two_turn_sampling_mask_trace([first, mismatched]) + assert trace.nodes[-1].sampling_mask is None + + def test_tool_call_hash_matches_v0_content_and_arguments_normalization(): left = vf.AssistantMessage( content=None, diff --git a/tests/unit/test_trace.py b/tests/unit/test_trace.py index 131a23b9d9..39d6170742 100644 --- a/tests/unit/test_trace.py +++ b/tests/unit/test_trace.py @@ -251,22 +251,24 @@ def test_sampling_mask_from_packed_matches_lists(): from verifiers.v1.types import SamplingMask - rows = [[4, 7], [], [9]] - ids, counts = np.array([4, 7, 9], np.int32), np.array([2, 0, 1], np.int32) - packed = { - name: { + def pack(a): + return { "data": base64.b64encode(a.tobytes()).decode(), "shape": list(a.shape), - "dtype": "int32", + "dtype": a.dtype.name, } - for name, a in (("ids", ids), ("counts", counts)) - } + + rows = [[4, 7], [], [9]] + ids, counts = np.array([4, 7, 9], np.int32), np.array([2, 0, 1], np.int32) + logprobs = np.array([-0.5, -1.0, 0.0], np.float32) + packed = {"ids": pack(ids), "counts": pack(counts)} for mask in ( SamplingMask.from_sampling_mask(rows), - SamplingMask.from_sampling_mask(packed), + SamplingMask.from_sampling_mask(packed, pack(logprobs)), ): np.testing.assert_array_equal(mask.ids, ids) np.testing.assert_array_equal(mask.counts, counts) + np.testing.assert_array_equal(mask.logprobs, logprobs) def test_max_turn_tokens_limit(): diff --git a/verifiers/v1/clients/train.py b/verifiers/v1/clients/train.py index ca22f5a0bb..0fbb0b3260 100644 --- a/verifiers/v1/clients/train.py +++ b/verifiers/v1/clients/train.py @@ -180,7 +180,9 @@ def response_from_generate( mm_placeholders=mm_placeholders, mm_token_type_id_map=mm_token_type_id_map, routed_experts=result.get("routed_experts"), - sampling_mask=SamplingMask.from_sampling_mask(mask) + sampling_mask=SamplingMask.from_sampling_mask( + mask, result.get("sampling_mask_logprobs") + ) if (mask := result.get("sampling_mask")) else None, weight_versions=[WeightVersionSpan(**span) for span in spans] diff --git a/verifiers/v1/graph.py b/verifiers/v1/graph.py index a75a8febdf..3efd5c4c29 100644 --- a/verifiers/v1/graph.py +++ b/verifiers/v1/graph.py @@ -222,8 +222,9 @@ def serialize_sampling_arrays( if payload is None: return None return { - field.name: _encode_ndarray(getattr(payload, field.name)) + field.name: _encode_ndarray(value) for field in fields(payload) + if (value := getattr(payload, field.name)) is not None } @field_validator("sampling_mask", "top_logprobs", mode="before") @@ -754,8 +755,10 @@ def _attribute_sampling_mask( 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 + if ( + len(payload.counts) != sum(node.mask) + or int(payload.counts.sum()) != len(payload.ids) + or (payload.logprobs is not None and len(payload.logprobs) != len(payload.ids)) ): return node.sampling_mask = payload diff --git a/verifiers/v1/trace.py b/verifiers/v1/trace.py index 34472801de..e20012c244 100644 --- a/verifiers/v1/trace.py +++ b/verifiers/v1/trace.py @@ -335,19 +335,29 @@ def sampling_mask(self) -> SamplingMask | None: return None # Attribution validates each mask against the node's sampled positions. ids_parts: list[np.ndarray] = [] + logprobs_parts: list[np.ndarray | None] = [] counts_parts: list[np.ndarray] = [] for node in self.nodes: counts = np.zeros(len(node.mask), dtype=np.int32) 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) + logprobs_parts.append(node.sampling_mask.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) ) - return SamplingMask(ids=ids, counts=np.concatenate(counts_parts)) + # Logprobs only when every mask on the path carries them. + logprobs = ( + np.concatenate([np.zeros(0, np.float32), *logprobs_parts]) + if all(part is not None for part in logprobs_parts) + else None + ) + return SamplingMask( + ids=ids, counts=np.concatenate(counts_parts), logprobs=logprobs + ) @property def top_logprobs(self) -> TopLogprobs | None: diff --git a/verifiers/v1/types.py b/verifiers/v1/types.py index b91f34d81e..de61aeba36 100644 --- a/verifiers/v1/types.py +++ b/verifiers/v1/types.py @@ -217,24 +217,32 @@ 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`. + token. Row boundaries are recovered from `counts`. The optional float32 `logprobs`, + parallel to `ids`, are the sampler's renormalized logprobs of those ids. """ ids: Any counts: Any + logprobs: Any = None @classmethod def from_sampling_mask( - cls, sampling_mask: list[list[int]] | dict[str, Any] + cls, + sampling_mask: list[list[int]] | dict[str, Any], + logprobs: dict[str, Any] | None = None, ) -> "SamplingMask": """Build from vLLM's one-list-per-token masks, or from prime-rl's packed CSR - form: `{"ids", "counts"}`, each a base64 `{data, shape, dtype}` int32 array.""" + form: `{"ids", "counts"}`, each a base64 `{data, shape, dtype}` int32 array. + `logprobs` is packed float32 parallel to the ids (only prime-rl sends it).""" import numpy as np + if logprobs is not None: + logprobs = _decode_packed(logprobs, np.float32) if isinstance(sampling_mask, dict): return cls( ids=_decode_packed(sampling_mask["ids"], np.int32), counts=_decode_packed(sampling_mask["counts"], np.int32), + logprobs=logprobs, ) counts = np.fromiter( (len(row) for row in sampling_mask), @@ -246,7 +254,7 @@ def from_sampling_mask( if int(counts.sum()) else np.empty(0, dtype=np.int32) ) - return cls(ids=ids, counts=counts) + return cls(ids=ids, counts=counts, logprobs=logprobs) def _decode_packed(packed: dict[str, Any], dtype: type) -> "np.ndarray":