diff --git a/renderers/client.py b/renderers/client.py index 656f424a..2e56f532 100644 --- a/renderers/client.py +++ b/renderers/client.py @@ -198,10 +198,15 @@ def _parse_compact_logprobs( values = np.frombuffer(raw, dtype=" 1e-6).any(): raise MalformedGenerateResponseError("Invalid compact logprob values") + samples = packed.get("samples") is True 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 + supports = ( + _parse_sampling_mask(choice, completion_ids) + if k is not None and not samples + else None + ) for index, (token_id, start, end) in enumerate( zip(completion_ids, offsets, offsets[1:]) ): @@ -215,7 +220,13 @@ def _parse_compact_logprobs( "Missing compact sampled-token evidence" ) sampled.append(row[token_id]) - if k is not None: + if k is not None and samples: + # Rows of i.i.d. sampler draws: the sampled token, then the draws with repeats kept. + if ids[start] != token_id or end - start < 2: + raise MalformedGenerateResponseError("Invalid compact sampler draws") + top_ids.append(ids[start + 1 : end].tolist()) + top_logprobs.append(values[start + 1 : end].tolist()) + elif k is not None: head = dict(islice(row.items(), max(num_top, 1))) selected_ids, selected_logprobs = _select_sampler_head( head, diff --git a/tests/test_client.py b/tests/test_client.py index 3b53fe25..46f0b1db 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -132,6 +132,30 @@ def test_parse_generate_response_preserves_multiple_compact_choices(separators): assert sampled.tolist() == [] and top_ids == [] and top_logprobs == [] +def test_compact_sampler_draws_keep_repeats(): + choice = _compact_choice( + { + "logprobs": { + "content": [ + { + "token": "token_id:7", + "logprob": -2.0, + "top_logprobs": [ + {"token": f"token_id:{i}", "logprob": p} + for i, p in [(1, -0.1), (1, -0.1), (7, -2.0)] + ], + } + ] + } + } + ) + choice["compact_logprobs"]["samples"] = True + sampled, ids, logprobs = _parse_compact_logprobs(choice, [7], 3) + assert sampled.tolist() == [-2.0] + assert ids == [[1, 1, 7]] + assert np.allclose(logprobs, [[-0.1, -0.1, -2.0]]) + + @pytest.mark.parametrize("sampled", [1, 3, 5]) @pytest.mark.parametrize("compact", [False, True]) @pytest.mark.parametrize(