renderers 0.2.1.dev8__tar.gz → 0.2.1.dev10__tar.gz

{renderers-0.2.1.dev8 → renderers-0.2.1.dev10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: renderers
-Version: 0.2.1.dev8
+Version: 0.2.1.dev10
 Summary: Chat template renderers — deterministic message-to-token conversion for LLM training
 License-Expression: Apache-2.0
 License-File: LICENSE
@@ -12,7 +12,7 @@ Requires-Dist: openai>=1.108.1
 Requires-Dist: prime-pydantic-config>=0.3.0.dev83
 Requires-Dist: tiktoken
 Provides-Extra: multimodal
-Requires-Dist: pillow>=12.2.0; extra == 'multimodal'
+Requires-Dist: pillow>=13.0.0; extra == 'multimodal'
 Requires-Dist: transformers>=4.50.0; extra == 'multimodal'
 Provides-Extra: transformers
 Requires-Dist: transformers>=4.50.0; extra == 'transformers'

{renderers-0.2.1.dev8 → renderers-0.2.1.dev10}/pyproject.toml

@@ -52,7 +52,7 @@ transformers = [
 # Image/audio renderers also need Pillow to resolve media inputs. Keep this as
 # a separate convenience extra so tokenizer-only users do not pull it in.
 multimodal = [
-    "pillow>=12.2.0",
+    "pillow>=13.0.0",
     "transformers>=4.50.0",
 ]
 
@@ -96,7 +96,7 @@ packages = ["renderers"]
 
 [dependency-groups]
 dev = [
-    "pillow>=12.2.0",
+    "pillow>=13.0.0",
     "pre-commit",
     "pytest>=7.0.0",
     "pytest-asyncio>=0.21.0",

{renderers-0.2.1.dev8 → renderers-0.2.1.dev10}/renderers/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '0.2.1.dev8'
-__version_tuple__ = version_tuple = (0, 2, 1, 'dev8')
+__version__ = version = '0.2.1.dev10'
+__version_tuple__ = version_tuple = (0, 2, 1, 'dev10')
 
 __commit_id__ = commit_id = None

{renderers-0.2.1.dev8 → renderers-0.2.1.dev10}/renderers/base.py

@@ -2230,7 +2230,7 @@ def build_trajectory_step(
         "completion_mask": [True] * len(completion_ids),
         "completion_logprobs": [0.0] * len(completion_ids),
         "routed_experts": None,
-        "sampling_mask": None,
+        "sampler_support": None,
     }
     if (
         full_rendered.multi_modal_data is not None

{renderers-0.2.1.dev8 → renderers-0.2.1.dev10}/renderers/client.py

@@ -169,9 +169,10 @@ def _parse_packed_completion_logprobs(
     return values.astype(np.float32, copy=False)
 
 
-def _parse_compact_logprobs(
-    choice: Mapping[str, Any], completion_ids: list[int], k: int | None
-) -> tuple[np.ndarray, list[list[int]] | None, list[list[float]] | None]:
+def _decode_compact_logprobs(
+    choice: Mapping[str, Any], completion_ids: list[int]
+) -> tuple[np.ndarray, np.ndarray, list[int], int]:
+    """Flat ids and logprobs of every compact row, the row offsets and the head size."""
     packed = choice.get("compact_logprobs")
     if not isinstance(packed, Mapping) or packed.get("format") != "u32-f32-le-v1":
         raise MalformedGenerateResponseError("Unsupported compact logprob format")
@@ -198,13 +199,115 @@ def _parse_compact_logprobs(
     values = np.frombuffer(raw, dtype="<f4", count=count, offset=count * 4)
     if not np.isfinite(values).all() or (values > 1e-6).any():
         raise MalformedGenerateResponseError("Invalid compact logprob values")
+    return ids, values, offsets, num_top
+
+
+def _sampler_support(
+    ids: np.ndarray, counts: np.ndarray, logprobs: np.ndarray | None, replayed: bool
+) -> dict[str, Any]:
+    """The sampler's per-token support in CSR form: completion token ``i`` owns ``counts[i]`` entries of
+    the flat ``ids`` (and ``logprobs``, the sampler's logprob of each id). ``replayed`` rows are the exact
+    set the token was sampled from; other rows are a head of the sampler's distribution."""
+    return {
+        "ids": ids.astype(np.int32, copy=False),
+        "counts": counts.astype(np.int32, copy=False),
+        "logprobs": None
+        if logprobs is None
+        else logprobs.astype(np.float32, copy=False),
+        "replayed": replayed,
+    }
+
+
+def _pack_sampler_rows(
+    ids: list[list[int]], logprobs: list[list[float]], replayed: bool
+) -> dict[str, Any]:
+    return _sampler_support(
+        np.fromiter((i for row in ids for i in row), dtype=np.int32),
+        np.fromiter((len(row) for row in ids), dtype=np.int32, count=len(ids)),
+        np.fromiter((p for row in logprobs for p in row), dtype=np.float32),
+        replayed,
+    )
+
+
+def _sampled_entries(
+    ids: np.ndarray, counts: np.ndarray, completion_ids: list[int]
+) -> np.ndarray:
+    """Which entry of each replayed row is the sampled token. Every row must be non-empty, hold distinct
+    ids and contain its sampled token."""
+    if (
+        len(counts) != len(completion_ids)
+        or int(counts.sum()) != len(ids)
+        or (counts <= 0).any()
+    ):
+        raise MalformedGenerateResponseError("Sampling mask token count mismatch")
+    ids = ids.astype(np.int64)
+    rows = np.repeat(np.arange(len(counts)), counts)
+    order = np.lexsort((ids, rows))
+    sorted_ids, sorted_rows = ids[order], rows[order]
+    repeated = (sorted_ids[1:] == sorted_ids[:-1]) & (
+        sorted_rows[1:] == sorted_rows[:-1]
+    )
+    sampled = ids == np.asarray(completion_ids, dtype=np.int64)[rows]
+    if repeated.any() or np.count_nonzero(sampled) != len(counts):
+        raise MalformedGenerateResponseError("Invalid recorded sampling support")
+    return sampled
+
+
+def _parse_replayed_support(
+    choice: Mapping[str, Any], completion_ids: list[int]
+) -> tuple[np.ndarray, dict[str, Any]]:
+    """A replayed request's sampled logprobs and support: the recorded sampling mask with the sampler's
+    logprob of every kept id, which must normalize over it."""
+    ids, values, offsets, _ = _decode_compact_logprobs(choice, completion_ids)
+    mask = choice.get("sampling_mask")
+    if not isinstance(mask, Mapping):
+        raise MalformedGenerateResponseError("Invalid packed sampling mask")
+    mask_ids, counts = _decode_packed_sampling_mask(mask)
+    if not np.array_equal(counts, np.diff(offsets)) or not np.array_equal(
+        mask_ids, ids
+    ):
+        raise MalformedGenerateResponseError(
+            "Sampler logprobs must cover exactly the recorded sampling support"
+        )
+    if (values == VLLM_LOGPROB_SENTINEL).any():
+        raise MalformedGenerateResponseError("Missing sampling-support probabilities")
+    sampled = _sampled_entries(mask_ids, counts, completion_ids)
+    if len(counts):
+        mass = np.add.reduceat(np.exp(values.astype(np.float64)), offsets[:-1])
+        if not np.allclose(mass, 1.0, rtol=0, atol=1e-4):
+            raise MalformedGenerateResponseError(
+                "Sampler probabilities must normalize over the recorded support"
+            )
+    return values[sampled], _sampler_support(mask_ids, counts, values, replayed=True)
+
+
+def _parse_mask_support(
+    choice: Mapping[str, Any], completion_ids: list[int]
+) -> dict[str, Any]:
+    """The recorded sampling mask of a request that returned no sampler logprobs for it."""
+    mask = choice.get("sampling_mask")
+    if isinstance(mask, Mapping):
+        ids, counts = _decode_packed_sampling_mask(mask)
+        _sampled_entries(ids, counts, completion_ids)
+        return _sampler_support(ids, counts, None, replayed=True)
+    rows = _parse_sampling_mask(choice, completion_ids)
+    assert rows is not None
+    return _sampler_support(
+        np.fromiter((i for row in rows for i in row), dtype=np.int32),
+        np.fromiter((len(row) for row in rows), dtype=np.int32, count=len(rows)),
+        None,
+        replayed=True,
+    )
+
+
+def _parse_compact_logprobs(
+    choice: Mapping[str, Any], completion_ids: list[int], k: int | None
+) -> tuple[np.ndarray, list[list[int]] | None, list[list[float]] | None]:
+    ids, values, offsets, num_top = _decode_compact_logprobs(choice, completion_ids)
     sampled = []
     top_ids: list[list[int]] = []
     top_logprobs: list[list[float]] = []
-    supports = _parse_sampling_mask(choice, completion_ids) if k is not None else None
-    for index, (token_id, start, end) in enumerate(
-        zip(completion_ids, offsets, offsets[1:])
-    ):
+    for token_id, start, end in zip(completion_ids, offsets, offsets[1:]):
         # FlatLogprobs can repeat the sampled ID. Match vLLM's dict order/value semantics.
         items = list(zip(ids[start:end].tolist(), values[start:end].tolist()))
         row = dict(items)
@@ -218,11 +321,7 @@ def _parse_compact_logprobs(
         if k is not None:
             head = dict(islice(row.items(), max(num_top, 1)))
             selected_ids, selected_logprobs = _select_sampler_head(
-                head,
-                token_id,
-                row[token_id],
-                k,
-                supports[index] if supports is not None else None,
+                head, token_id, row[token_id], k
             )
             top_ids.append(selected_ids)
             top_logprobs.append(selected_logprobs)
@@ -300,6 +399,28 @@ def _parse_completion_logprobs(
     return np.asarray(completion_logprobs, dtype=np.float32)
 
 
+def _decode_packed_sampling_mask(
+    mask: Mapping[str, Any],
+) -> tuple[np.ndarray, np.ndarray]:
+    arrays = {}
+    for name in ("ids", "counts"):
+        packed = mask.get(name)
+        if not isinstance(packed, Mapping) or packed.get("dtype") != "int32":
+            raise MalformedGenerateResponseError("Invalid packed sampling mask")
+        try:
+            values = np.frombuffer(
+                base64.b64decode(packed["data"], validate=True), dtype=np.int32
+            )
+        except (KeyError, TypeError, ValueError, binascii.Error) as exc:
+            raise MalformedGenerateResponseError(
+                "Invalid packed sampling mask"
+            ) from exc
+        if packed.get("shape") != [len(values)] or (values < 0).any():
+            raise MalformedGenerateResponseError("Invalid packed sampling mask")
+        arrays[name] = values
+    return arrays["ids"], arrays["counts"]
+
+
 def _parse_sampling_mask(
     choice: Mapping[str, Any], completion_ids: list[int]
 ) -> list[list[int]] | None:
@@ -307,23 +428,7 @@ def _parse_sampling_mask(
     if mask is None:
         return None
     if isinstance(mask, Mapping):
-        arrays = {}
-        for name in ("ids", "counts"):
-            packed = mask.get(name)
-            if not isinstance(packed, Mapping) or packed.get("dtype") != "int32":
-                raise MalformedGenerateResponseError("Invalid packed sampling mask")
-            try:
-                values = np.frombuffer(
-                    base64.b64decode(packed["data"], validate=True), dtype=np.int32
-                )
-            except (KeyError, TypeError, ValueError, binascii.Error) as exc:
-                raise MalformedGenerateResponseError(
-                    "Invalid packed sampling mask"
-                ) from exc
-            if packed.get("shape") != [len(values)] or (values < 0).any():
-                raise MalformedGenerateResponseError("Invalid packed sampling mask")
-            arrays[name] = values
-        counts, ids = arrays["counts"], arrays["ids"]
+        ids, counts = _decode_packed_sampling_mask(mask)
         if len(counts) != len(completion_ids) or int(counts.sum()) != len(ids):
             raise MalformedGenerateResponseError("Sampling mask token count mismatch")
         mask = [row.tolist() for row in np.split(ids, np.cumsum(counts)[:-1])]
@@ -500,7 +605,7 @@ async def generate(
     Returns a dict with: request_id, prompt_ids, renderer_prompt_ids,
     mm_placeholders, completion_ids, completion_logprobs, content,
     reasoning_content, tool_calls, finish_reason, routed_experts,
-    weight_versions, multi_modal_data, prompt_attribution.
+    sampler_support, weight_versions, multi_modal_data, prompt_attribution.
     ``completion_logprobs`` is a float32 numpy array, one value per completion
     token: vLLM samples in float32, so the packed form decodes without a copy
     and the JSON forms narrow losslessly.
@@ -509,9 +614,13 @@ async def generate(
     is vLLM's (>= 0.32) list of ``{version, start, end}`` spans over
     completion-token indexes, or ``None`` when the engine labels none.
 
-    When ``sampling_params["logprobs"]`` is a positive int, also
-    ``completion_top_ids`` / ``completion_top_logprobs``
-    (original sampler probabilities, retaining full-vocabulary normalization).
+    ``sampler_support`` is the sampler's per-token support in CSR form, or ``None`` when
+    the engine returned none: ``counts[i]`` entries of the flat int32 ``ids`` (and float32
+    ``logprobs``, or ``None``) per completion token. With ``replayed`` the rows are the
+    recorded sampling masks of truncated sampling, with logprobs normalized over each
+    mask when the engine returned them. Otherwise, when ``sampling_params["logprobs"]``
+    is a positive int, the rows are that many top candidates of the sampler's
+    distribution, independent of the sampled token and full-vocabulary normalized.
 
     ``prompt_attribution`` is the renderer's :class:`RenderedTokens` for
     the prompt — either the one this call computed via
@@ -646,18 +755,27 @@ async def generate(
             "Engine response must include mm_placeholders when process_multimodal=False."
         )
 
-    completion_top_ids = None
-    completion_top_logprobs = None
-    if choice.get("compact_logprobs") is not None:
-        completion_logprobs, completion_top_ids, completion_top_logprobs = (
-            _parse_compact_logprobs(choice, completion_ids, head_size)
+    sampler_support = None
+    replayed = choice.get("sampling_mask") is not None
+    if choice.get("compact_logprobs") is not None and replayed:
+        completion_logprobs, sampler_support = _parse_replayed_support(
+            choice, completion_ids
+        )
+    elif choice.get("compact_logprobs") is not None:
+        completion_logprobs, top_ids, top_logprobs = _parse_compact_logprobs(
+            choice, completion_ids, head_size
         )
+        if top_ids is not None and top_logprobs is not None:
+            sampler_support = _pack_sampler_rows(top_ids, top_logprobs, replayed=False)
     else:
         completion_logprobs = _parse_completion_logprobs(choice, completion_ids)
         if head_size is not None:
-            completion_top_ids, completion_top_logprobs = (
-                _parse_completion_top_logprobs(choice, completion_ids, head_size)
+            sampler_support = _pack_sampler_rows(
+                *_parse_completion_top_logprobs(choice, completion_ids, head_size),
+                replayed=replayed,
             )
+        elif replayed:
+            sampler_support = _parse_mask_support(choice, completion_ids)
 
     parsed = renderer.parse_response(
         completion_ids,
@@ -666,10 +784,6 @@ async def generate(
     )
 
     routed_experts = choice.get("routed_experts")
-    # vLLM's native kept-set sampling masks (``--return-sampling-mask``):
-    # one list of surviving vocab ids per completion token, or (prime-rl's
-    # server) packed CSR ``{"ids", "counts"}`` base64 int32 arrays.
-    sampling_mask = choice.get("sampling_mask")
 
     # /inference/v1/generate returns finish_reason in {"stop","length",...} —
     # never "tool_calls" (a chat-completions concept). Promote stop→tool_calls
@@ -694,15 +808,13 @@ async def generate(
         "mm_placeholders": mm_placeholders,
         "completion_ids": list(completion_ids),
         "completion_logprobs": completion_logprobs,
-        "completion_top_ids": completion_top_ids,
-        "completion_top_logprobs": completion_top_logprobs,
+        "sampler_support": sampler_support,
         "content": parsed.content,
         "reasoning_content": parsed.reasoning_content,
         "tool_calls": parsed.tool_calls,
         "finish_reason": finish_reason,
         "reasoning_complete": parsed.reasoning_complete,
         "routed_experts": routed_experts,
-        "sampling_mask": sampling_mask,
         "weight_versions": choice.get("weight_versions"),
         # The mm sidecar consumed on the request side, surfaced back so
         # callers can persist it on the trajectory step for downstream

{renderers-0.2.1.dev8 → renderers-0.2.1.dev10}/tests/replay.py

@@ -45,4 +45,7 @@ def replay_client(completions, requests):
 def assert_evidence(result, ids):
     assert result["completion_ids"] == ids
     assert result["completion_logprobs"].tolist() == [-0.25] * len(ids)
-    assert result["sampling_mask"] == [[tid] for tid in ids]
+    support = result["sampler_support"]
+    assert support["replayed"] and support["logprobs"] is None
+    assert support["ids"].tolist() == ids
+    assert support["counts"].tolist() == [1] * len(ids)

{renderers-0.2.1.dev8 → renderers-0.2.1.dev10}/tests/test_client.py

@@ -16,6 +16,7 @@ from renderers.base import (
 from renderers.client import (
     _parse_compact_logprobs,
     _parse_completion_top_logprobs,
+    _parse_replayed_support,
     generate,
     parse_generate_response,
 )
@@ -59,6 +60,64 @@ def _parse_test_head(choice, sampled, k, compact):
     return _parse_completion_top_logprobs(choice, sampled, k)
 
 
+def _replayed_choice(rows, sampled):
+    """A paired replay response: each row's kept ids and their sampler probabilities."""
+    ids = np.asarray([i for row in rows for i, _ in row], dtype=np.int32)
+    counts = np.asarray([len(row) for row in rows], dtype=np.int32)
+    values = np.log(np.asarray([p for row in rows for _, p in row], dtype=np.float32))
+    return {
+        "token_ids": sampled,
+        "compact_logprobs": {
+            "data": base64.b64encode(
+                ids.astype("<u4").tobytes() + values.astype("<f4").tobytes()
+            ).decode("ascii"),
+            "format": "u32-f32-le-v1",
+            "offsets": [0, *np.cumsum(counts).tolist()],
+            "num_top_logprobs": int(counts.max(initial=0)),
+        },
+        "sampling_mask": {
+            name: {
+                "data": base64.b64encode(array.tobytes()).decode(),
+                "shape": [len(array)],
+                "dtype": "int32",
+            }
+            for name, array in (("ids", ids), ("counts", counts))
+        },
+    }
+
+
+@pytest.mark.parametrize(
+    "error", [None, "mass", "unsampled", "duplicate", "mask", "empty"]
+)
+def test_replayed_support_keeps_the_recorded_mask_with_its_probabilities(error):
+    rows = [[(3, 0.2), (1, 0.6), (2, 0.2)], [(9, 1.0)]]
+    sampled = [3, 9]
+    if error == "mass":
+        rows[0][1] = (1, 0.5)
+    elif error == "unsampled":
+        sampled = [4, 9]
+    elif error == "duplicate":
+        rows[0] = [(3, 0.2), (1, 0.6), (3, 0.2)]
+    choice = _replayed_choice(rows, sampled)
+    if error == "mask":
+        choice["sampling_mask"] = _replayed_choice(
+            [[(3, 0.2), (1, 0.6), (5, 0.2)], [(9, 1.0)]], sampled
+        )["sampling_mask"]
+    elif error == "empty":
+        choice = _replayed_choice([], [])
+    if error not in (None, "empty"):
+        with pytest.raises(MalformedGenerateResponseError):
+            _parse_replayed_support(choice, sampled)
+        return
+    logprobs, support = _parse_replayed_support(choice, choice["token_ids"])
+    expected = [] if error == "empty" else [0.2, 1.0]
+    np.testing.assert_allclose(np.exp(logprobs), expected, rtol=1e-6)
+    assert support["replayed"]
+    assert support["ids"].tolist() == [i for row in rows for i, _ in row] * (not error)
+    assert support["counts"].tolist() == [len(row) for row in rows] * (not error)
+    assert len(support["logprobs"]) == len(support["ids"])
+
+
 @pytest.mark.parametrize(
     "update",
     [
@@ -168,10 +227,9 @@ def test_top_logprobs_head_is_independent_of_sampled_action(
         _parse_test_head({"logprobs": {"content": [entry]}}, [sampled], 2, compact)
 
 
-@pytest.mark.parametrize("compact", [False, True])
 @pytest.mark.parametrize("packed", [False, True])
 @pytest.mark.parametrize("error", [None, "missing", "mass", "mask"])
-def test_top_logprobs_preserve_complete_replay_support(compact, packed, error):
+def test_top_logprobs_preserve_complete_replay_support(packed, error):
     import math
 
     probabilities = {1: 0.6, 2: 0.2, 3: 0.2}
@@ -200,9 +258,9 @@ def test_top_logprobs_preserve_complete_replay_support(compact, packed, error):
     choice = {"logprobs": {"content": [entry]}, "sampling_mask": mask}
     if error is not None:
         with pytest.raises(MalformedGenerateResponseError):
-            _parse_test_head(choice, [3], 2, compact)
+            _parse_completion_top_logprobs(choice, [3], 2)
     else:
-        ids, logprobs = _parse_test_head(choice, [3], 2, compact)
+        ids, logprobs = _parse_completion_top_logprobs(choice, [3], 2)
         assert ids == [[1, 2, 3]]
         np.testing.assert_allclose(np.exp(logprobs), [[0.6, 0.2, 0.2]])
 
@@ -425,13 +483,13 @@ def test_generate_builds_request_body_and_parses_response(
     assert result["completion_ids"] == [7, 8]
     np.testing.assert_allclose(result["completion_logprobs"], [-0.1, -0.2])
     if head_size is not None:
-        assert result["completion_top_ids"] == [[7, 20], [8, 20]]
-        np.testing.assert_allclose(
-            result["completion_top_logprobs"], [[-0.1, -4.0], [-0.2, -4.0]]
-        )
+        support = result["sampler_support"]
+        assert not support["replayed"]
+        assert support["ids"].tolist() == [7, 20, 8, 20]
+        assert support["counts"].tolist() == [2, 2]
+        np.testing.assert_allclose(support["logprobs"], [-0.1, -4.0, -0.2, -4.0])
     else:
-        assert result["completion_top_ids"] is None
-        assert result["completion_top_logprobs"] is None
+        assert result["sampler_support"] is None
     if extra_args is not None:
         assert extra_args == {"prl_compact_logprobs": False, "custom_arg": 7}
     assert result["routed_experts"]["shape"] == [2, 1, 1]

{renderers-0.2.1.dev8 → renderers-0.2.1.dev10}/uv.lock

@@ -1385,7 +1385,7 @@ requires-dist = [
     { name = "numpy" },
     { name = "openai", specifier = ">=1.108.1" },
     { name = "openai-harmony", specifier = ">=0.0.4" },
-    { name = "pillow", marker = "extra == 'multimodal'", specifier = ">=12.2.0" },
+    { name = "pillow", marker = "extra == 'multimodal'", specifier = ">=13.0.0" },
     { name = "prime-pydantic-config", specifier = ">=0.3.0.dev83" },
     { name = "tiktoken" },
     { name = "transformers", marker = "extra == 'multimodal'", specifier = ">=4.50.0" },
@@ -1395,7 +1395,7 @@ provides-extras = ["multimodal", "transformers"]
 
 [package.metadata.requires-dev]
 dev = [
-    { name = "pillow", specifier = ">=12.2.0" },
+    { name = "pillow", specifier = ">=13.0.0" },
     { name = "pre-commit" },
     { name = "pytest", specifier = ">=7.0.0" },
     { name = "pytest-asyncio", specifier = ">=0.21.0" },