[Auto Sync] Update logits_processor.py, test_logprobs.py (20260124) (#17664)

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: yehu-ux <yehu@x.ai>
2026-01-23 17:57:41 -08:00
parent e1833c4f5a
commit bc6f0b5ce7
2 changed files with 9 additions and 3 deletions
@@ -334,7 +334,8 @@ class TestLogprobsDense(unittest.TestCase):
                print(f"Testing with config: {config}")

                # Sample records for this config
-                test_records = random.sample(records, k=min(NUM_SAMPLES, len(records)))
+                num_samples = config.get("num_samples", NUM_SAMPLES)
+                test_records = random.sample(records, k=min(num_samples, len(records)))
                random.shuffle(test_records)

                # Calculate how many samples should return logprobs