feat(client): keep compact sampler-draw rows with their repeats

renderers/client.py

@@ -198,10 +198,15 @@ 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")
+    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,

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(