perf(v1): accept packed CSR sampling masks from prime-rl's /generate

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)

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."""