diff --git a/tests/unit/test_trace.py b/tests/unit/test_trace.py index b4fd9ed279..cc094d3a93 100644 --- a/tests/unit/test_trace.py +++ b/tests/unit/test_trace.py @@ -241,3 +241,28 @@ def test_transcript(make_full_trace): assert "[tool search]\nTOOL RESULT: Paris is the capital." in transcript assert "It is Paris." in transcript assert "SECRET REASONING" not in transcript # reasoning is excluded + + +def test_sampling_mask_from_packed_matches_lists(): + import base64 + + import numpy as np + + 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: { + "data": base64.b64encode(a.tobytes()).decode(), + "shape": list(a.shape), + "dtype": "int32", + } + for name, a in (("ids", ids), ("counts", counts)) + } + for mask in ( + SamplingMask.from_sampling_mask(rows), + SamplingMask.from_sampling_mask(packed), + ): + np.testing.assert_array_equal(mask.ids, ids) + np.testing.assert_array_equal(mask.counts, counts) diff --git a/verifiers/v1/types.py b/verifiers/v1/types.py index 79a01de60a..0e15ad9f8b 100644 --- a/verifiers/v1/types.py +++ b/verifiers/v1/types.py @@ -1,3 +1,4 @@ +import base64 from collections.abc import Iterable from dataclasses import dataclass from typing import Annotated, Any, Literal @@ -208,7 +209,16 @@ class SamplingMask: counts: Any @classmethod - def from_sampling_mask(cls, sampling_mask: list[list[int]]) -> "SamplingMask": + def from_sampling_mask( + cls, sampling_mask: list[list[int]] | dict[str, Any] + ) -> "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.""" + if isinstance(sampling_mask, dict): + return cls( + ids=_decode_packed(sampling_mask["ids"], np.int32), + counts=_decode_packed(sampling_mask["counts"], np.int32), + ) counts = np.fromiter( (len(row) for row in sampling_mask), dtype=np.int32, @@ -222,6 +232,11 @@ def from_sampling_mask(cls, sampling_mask: list[list[int]]) -> "SamplingMask": return cls(ids=ids, counts=counts) +def _decode_packed(packed: dict[str, Any], dtype: type) -> np.ndarray: + array = np.frombuffer(base64.b64decode(packed["data"]), dtype=packed["dtype"]) + return array.astype(dtype, copy=False) + + class TurnTokens(BaseModel): """Training tokens from renderer tokenization or provider-returned token IDs."""