| |
| import io |
| import json |
| import sys |
| import tempfile |
| import unittest |
| import urllib.error |
| from pathlib import Path |
| from unittest import mock |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parents[1])) |
| import public_bench as pb |
|
|
|
|
| class Response: |
| def __init__(self, status, body): |
| self.status = status |
| self.body = body |
|
|
| def __enter__(self): |
| return self |
|
|
| def __exit__(self, *args): |
| return False |
|
|
| def getcode(self): |
| return self.status |
|
|
| def read(self): |
| return self.body |
|
|
|
|
| def success_result(content="A", error=None): |
| return { |
| "content": content, "reasoning_chars": 0, "elapsed_sec": 0.01, |
| "timings": {}, "usage": None, "model": "test", "error": error, |
| "error_type": "test" if error else None, "http_status": 500 if error else 200, |
| "error_body": "failure" if error else "", "attempts": [], "retry_count": 0, |
| "prompt_hash": "hash", |
| } |
|
|
|
|
| class PublicBenchTests(unittest.TestCase): |
| def test_balanced_sample_is_exact_and_stable(self): |
| rows = [{"id": str(i), "domain": domain} |
| for i, domain in enumerate("aaabbbcccc")] |
| first = pb.balanced_sample(rows, "domain", 7, 123, |
| lambda row: row["id"]) |
| second = pb.balanced_sample(list(reversed(rows)), "domain", 7, 123, |
| lambda row: row["id"]) |
| self.assertEqual([x["id"] for x in first], [x["id"] for x in second]) |
| self.assertEqual(len(first), 7) |
| counts = {domain: sum(x["domain"] == domain for x in first) |
| for domain in "abc"} |
| self.assertEqual(sum(counts.values()), 7) |
| self.assertLessEqual(max(counts.values()) - min(counts.values()), 1) |
|
|
| def test_client_retries_and_records_http_body(self): |
| calls = [] |
|
|
| def opener(request, timeout): |
| calls.append(request) |
| if len(calls) == 1: |
| raise urllib.error.HTTPError( |
| request.full_url, 500, "server", {}, io.BytesIO(b"overloaded")) |
| return Response(200, b'{"choices":[{"message":{"content":"A"}}]}') |
|
|
| client = pb.LocalClient("http://example.test/v1/chat/completions", False, |
| 8, 1, retries=2, retry_backoff=0, opener=opener) |
| result = client.call("question", "item-1") |
| self.assertEqual(result["content"], "A") |
| self.assertEqual(len(calls), 2) |
| self.assertEqual(result["retry_count"], 1) |
| self.assertEqual(result["attempts"][0]["status"], 500) |
| self.assertEqual(result["attempts"][0]["body_excerpt"], "overloaded") |
| self.assertIsNone(result["error"]) |
|
|
| def test_client_does_not_repeat_reasoning_only_truncation(self): |
| calls = [] |
|
|
| def opener(request, timeout): |
| calls.append(request) |
| return Response( |
| 200, |
| b'{"choices":[{"finish_reason":"length","message":' |
| b'{"content":"","reasoning_content":"long thought"}}]}', |
| ) |
|
|
| client = pb.LocalClient("http://example.test/v1/chat/completions", True, |
| 8, 1, retries=3, retry_backoff=0, opener=opener) |
| result = client.call("question", "item-1") |
| self.assertEqual(len(calls), 1) |
| self.assertEqual(result["error_type"], "OutputTruncatedError") |
| self.assertEqual(result["finish_reason"], "length") |
| self.assertEqual(result["reasoning_chars"], len("long thought")) |
|
|
| def test_adaptive_client_votes_and_assigns_max_profile_to_aime(self): |
| class FakeClient: |
| seed = 100 |
|
|
| def __init__(self): |
| self.calls = [] |
|
|
| def call(self, prompt, item_id, **kwargs): |
| self.calls.append((prompt, item_id, kwargs)) |
| answers = {"sample-1": "7", "sample-2": "8", "sample-3": "7"} |
| content = next((value for key, value in answers.items() |
| if key in item_id), "7") |
| return { |
| "content": content, "error": None, "error_type": None, |
| "http_status": 200, "error_body": "", "attempts": [], |
| "retry_count": 0, "elapsed_sec": 0.01, "timings": {}, |
| "usage": None, "model": "test", "reasoning_chars": 10, |
| "request": kwargs, |
| } |
|
|
| profiles = { |
| "fixed": pb.RequestProfile("fixed", False, 8, 1), |
| "fast": pb.RequestProfile("fast", False, 8, 1), |
| "deep": pb.RequestProfile("deep", True, 16, 2), |
| "max": pb.RequestProfile("max", True, 32, 3), |
| } |
| fake = FakeClient() |
| client = pb.AdaptiveClient( |
| fake, "aime", "adaptive", profiles, pb.extract_integer, |
| samples=3, sample_temperature=0.35, |
| ) |
| result = client.call("Solve this problem", "aime/1") |
| self.assertEqual(result["winner"], "7") |
| self.assertEqual(result["vote_counts"], {"7": 2, "8": 1}) |
| self.assertEqual(result["profile"], "max") |
| self.assertEqual(result["request_count"], 3) |
| self.assertTrue(all(call[2]["think"] for call in fake.calls)) |
| self.assertTrue(all(call[2]["max_tokens"] == 32 for call in fake.calls)) |
|
|
| def test_adaptive_client_retries_fast_when_answer_format_is_invalid(self): |
| class FakeClient: |
| seed = 0 |
|
|
| def __init__(self): |
| self.calls = [] |
|
|
| def call(self, prompt, item_id, **kwargs): |
| self.calls.append((item_id, kwargs)) |
| content = "not an answer" if "fallback" not in item_id else "B" |
| return { |
| "content": content, "error": None, "error_type": None, |
| "http_status": 200, "error_body": "", "attempts": [], |
| "retry_count": 0, "elapsed_sec": 0.01, "timings": {}, |
| "usage": None, "model": "test", "reasoning_chars": 0, |
| "request": kwargs, |
| } |
|
|
| profiles = { |
| "fixed": pb.RequestProfile("fixed", False, 8, 1), |
| "fast": pb.RequestProfile("fast", False, 8, 1), |
| "deep": pb.RequestProfile("deep", True, 16, 2), |
| } |
| profiles["max"] = pb.RequestProfile("max", True, 32, 3) |
| fake = FakeClient() |
| client = pb.AdaptiveClient( |
| fake, "mmlu_pro", "adaptive", profiles, |
| lambda text: pb.extract_letter(text, "ABCD"), |
| ) |
| result = client.call("Choose one", "mmlu/1") |
| self.assertEqual(result["content"], "B") |
| self.assertEqual(len(fake.calls), 2) |
| self.assertEqual(fake.calls[-1][1]["think"], False) |
| self.assertEqual(fake.calls[-1][1]["max_tokens"], 264) |
|
|
| def test_max_client_rescues_reasoning_only_response_with_fast_call(self): |
| class FakeClient: |
| seed = 0 |
|
|
| def __init__(self): |
| self.calls = [] |
|
|
| def call(self, prompt, item_id, **kwargs): |
| self.calls.append((item_id, kwargs)) |
| if "fallback-fast" not in item_id: |
| return { |
| "item_id": item_id, "content": "", "error": "cut", |
| "error_type": "OutputTruncatedError", "elapsed_sec": 1, |
| "timings": {}, "request": kwargs, |
| } |
| return { |
| "item_id": item_id, "content": "C", "error": None, |
| "error_type": None, "elapsed_sec": 0.1, "timings": {}, |
| "request": kwargs, |
| } |
|
|
| profiles = { |
| "fixed": pb.RequestProfile("fixed", False, 8, 1), |
| "fast": pb.RequestProfile("fast", False, 8, 1), |
| "deep": pb.RequestProfile("deep", True, 16, 2), |
| "max": pb.RequestProfile("max", True, 32, 3), |
| } |
| fake = FakeClient() |
| client = pb.AdaptiveClient( |
| fake, "gpqa_diamond", "max", profiles, |
| lambda text: pb.extract_letter(text, "ABCD"), |
| ) |
| result = client.call("Choose one", "gpqa/1") |
| self.assertEqual(result["content"], "C") |
| self.assertTrue(result["fallback_used"]) |
| self.assertEqual([kwargs["think"] for _, kwargs in fake.calls], [True, False]) |
|
|
| def test_corrupt_cache_is_quarantined_and_refetched(self): |
| with tempfile.TemporaryDirectory() as directory: |
| out = Path(directory) |
| cache = out / "datasets" / "x__c__test.jsonl" |
| cache.parent.mkdir() |
| cache.write_text('{broken\n', encoding="utf-8") |
| with mock.patch.object(pb, "get_json", return_value={ |
| "num_rows_total": 1, "rows": [{"row": {"id": 1}}] |
| }): |
| rows = pb.fetch_rows("x", "c", "test", out) |
| self.assertEqual(rows, [{"id": 1}]) |
| self.assertTrue(list(cache.parent.glob(cache.name + ".corrupt-*"))) |
| self.assertEqual(pb.read_jsonl(cache), [{"id": 1}]) |
|
|
| def test_run_rows_retries_failed_existing_item_only(self): |
| class FakeClient: |
| def __init__(self, result): |
| self.result = result |
| self.calls = [] |
|
|
| def call(self, prompt, item_id): |
| self.calls.append(item_id) |
| return dict(self.result) |
|
|
| rows = [{"id": "one", "answer": "A"}, {"id": "two", "answer": "A"}] |
| with tempfile.TemporaryDirectory() as directory: |
| out = Path(directory) |
| first_client = FakeClient(success_result("", "HTTPError")) |
| first = pb.run_rows("toy", rows, lambda row: row["id"], |
| lambda content: content or None, first_client, out, |
| lambda row: row["id"], gold_fn=lambda row: "A") |
| self.assertEqual(first["failed_n"], 2) |
| self.assertEqual(first_client.calls, ["one", "two"]) |
|
|
| second_client = FakeClient(success_result("A")) |
| second = pb.run_rows("toy", rows, lambda row: row["id"], |
| lambda content: content, second_client, out, |
| lambda row: row["id"], gold_fn=lambda row: "A") |
| self.assertEqual(second["failed_n"], 0) |
| self.assertEqual(second["correct"], 2) |
| self.assertEqual(second_client.calls, ["one", "two"]) |
|
|
| items = pb.read_jsonl(out / "toy.jsonl") |
| self.assertEqual(len(items), 2) |
| self.assertTrue(all(item["row_sha256"] for item in items)) |
|
|
| def test_atomic_jsonl_is_valid_after_write(self): |
| with tempfile.TemporaryDirectory() as directory: |
| path = Path(directory) / "rows.jsonl" |
| pb.jsonl_dump(path, [{"x": 1}, {"x": 2}]) |
| self.assertEqual(pb.read_jsonl(path), [{"x": 1}, {"x": 2}]) |
| self.assertFalse(list(Path(directory).glob("*.tmp"))) |
|
|
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|
