diff --git a/renderers/client.py b/renderers/client.py index 2e56f532..fcc2953e 100644 --- a/renderers/client.py +++ b/renderers/client.py @@ -199,6 +199,14 @@ def _parse_compact_logprobs( if not np.isfinite(values).all() or (values > 1e-6).any(): raise MalformedGenerateResponseError("Invalid compact logprob values") samples = packed.get("samples") is True + # Replayed supports wider than the server's width return only part of their ids, with + # probabilities still normalized over the whole support. + partial_rows = packed.get("partial_rows", []) + if not isinstance(partial_rows, list) or any( + type(row) is not int or not 0 <= row < len(completion_ids) for row in partial_rows + ): + raise MalformedGenerateResponseError("Invalid compact logprob partial rows") + partial_rows = set(partial_rows) sampled = [] top_ids: list[list[int]] = [] top_logprobs: list[list[float]] = [] @@ -234,6 +242,7 @@ def _parse_compact_logprobs( row[token_id], k, supports[index] if supports is not None else None, + partial=index in partial_rows, ) top_ids.append(selected_ids) top_logprobs.append(selected_logprobs) @@ -360,8 +369,10 @@ def _select_sampler_head( sampled_logprob: float, k: int, support: list[int] | None = None, + partial: bool = False, ) -> tuple[list[int], list[float]]: - """Keep an action-independent head, or the complete recorded replay support.""" + """Keep an action-independent head, or the recorded replay support. A ``partial`` + support holds part of the ids, so its probabilities may sum to less than one.""" if sampled_id in head and not math.isclose( head[sampled_id], sampled_logprob, rel_tol=1e-6, abs_tol=1e-5 ): @@ -380,7 +391,7 @@ def _select_sampler_head( "Missing sampling-support probabilities" ) ranked = [(i, p) for i, p in ranked if i in support_set] - if not math.isclose( + if not partial and not math.isclose( math.fsum(math.exp(p) for _, p in ranked), 1.0, rel_tol=0, abs_tol=1e-4 ): raise MalformedGenerateResponseError( diff --git a/tests/test_client.py b/tests/test_client.py index 46f0b1db..a70361df 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -76,6 +76,9 @@ def _parse_test_head(choice, sampled, k, compact): {"data": "?"}, {"data": None}, {"data": ""}, + {"partial_rows": "0"}, + {"partial_rows": [1]}, + {"partial_rows": [True]}, ], ) def test_compact_logprobs_reject_malformed_payload(update): @@ -194,12 +197,12 @@ def test_top_logprobs_head_is_independent_of_sampled_action( @pytest.mark.parametrize("compact", [False, True]) @pytest.mark.parametrize("packed", [False, True]) -@pytest.mark.parametrize("error", [None, "missing", "mass", "mask"]) +@pytest.mark.parametrize("error", [None, "missing", "mass", "mask", "partial"]) def test_top_logprobs_preserve_complete_replay_support(compact, packed, error): import math probabilities = {1: 0.6, 2: 0.2, 3: 0.2} - if error == "mass": + if error in ("mass", "partial"): probabilities[1] = 0.5 entry = { "token": "token_id:3", @@ -222,7 +225,14 @@ def test_top_logprobs_preserve_complete_replay_support(compact, packed, error): for name, values in (("ids", support), ("counts", [len(support)])) } choice = {"logprobs": {"content": [entry]}, "sampling_mask": mask} - if error is not None: + if error == "partial" and compact: + # Only the server declares a support partial: its mass may then fall short of one. + choice = _compact_choice(choice) + choice["compact_logprobs"]["partial_rows"] = [0] + _, ids, logprobs = _parse_compact_logprobs(choice, [3], 2) + assert ids == [[1, 2, 3]] + np.testing.assert_allclose(np.exp(logprobs), [[0.5, 0.2, 0.2]]) + elif error is not None: with pytest.raises(MalformedGenerateResponseError): _parse_test_head(choice, [3], 2, compact) else: