From 609f5f549c3c6a4e837fc03ba0c02bc33eeb8eaa Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Tue, 9 Jun 2026 20:17:41 +0800 Subject: [PATCH] Add mixed-prefix gsm8k eval and its CPU unit test (#27502) --- python/sglang/test/run_eval.py | 23 +++ python/sglang/test/simple_eval_gsm8k.py | 19 +- .../test/simple_eval_mixed_prefix_gsm8k.py | 51 +++++ .../unit/bench/test_mixed_prefix_gsm8k.py | 177 ++++++++++++++++++ 4 files changed, 265 insertions(+), 5 deletions(-) create mode 100644 python/sglang/test/simple_eval_mixed_prefix_gsm8k.py create mode 100644 test/registered/unit/bench/test_mixed_prefix_gsm8k.py diff --git a/python/sglang/test/run_eval.py b/python/sglang/test/run_eval.py index d872966e7..c80b09540 100644 --- a/python/sglang/test/run_eval.py +++ b/python/sglang/test/run_eval.py @@ -178,6 +178,17 @@ def run_eval(args): num_shots=getattr(args, "num_shots", 5), data_path=getattr(args, "gsm8k_data_path", None), ) + elif args.eval_name == "mixed_prefix_gsm8k": + from sglang.test.simple_eval_mixed_prefix_gsm8k import MixedPrefixGSM8KEval + + eval_obj = MixedPrefixGSM8KEval( + num_examples=args.num_examples, + num_threads=args.num_threads, + num_shots=args.num_shots, + secondary_pool_size=args.mixed_prefix_gsm8k_secondary_pool_size, + data_path=args.gsm8k_data_path, + seed=args.mixed_prefix_gsm8k_seed, + ) else: raise ValueError(f"Invalid eval name: {args.eval_name}") @@ -367,6 +378,18 @@ if __name__ == "__main__": default=None, help="Path to GSM8K data file (e.g., test.jsonl)", ) + parser.add_argument( + "--mixed-prefix-gsm8k-secondary-pool-size", + type=int, + default=15, + help="Size of secondary example pool for eval_name=mixed_prefix_gsm8k (default: 15)", + ) + parser.add_argument( + "--mixed-prefix-gsm8k-seed", + type=int, + default=42, + help="Seed for per-question random sampling in mixed_prefix_gsm8k (default: 42)", + ) args = parser.parse_args() diff --git a/python/sglang/test/simple_eval_gsm8k.py b/python/sglang/test/simple_eval_gsm8k.py index 88665d9c7..9408d7006 100644 --- a/python/sglang/test/simple_eval_gsm8k.py +++ b/python/sglang/test/simple_eval_gsm8k.py @@ -56,20 +56,29 @@ class GSM8KEval(Eval): else: filename = download_and_cache_file(GSM8K_URL) - self._lines = list(read_jsonl(filename)) - self._few_shot_prompt = get_few_shot_examples(self._lines, num_shots) - + all_lines = list(read_jsonl(filename)) + pool_size = self._setup_prefix_pool(all_lines, num_shots) # The evaluation data should not include the few-shot examples to prevent data leakage. - self._lines = self._lines[num_shots:] + self._lines = all_lines[pool_size:] if num_examples is not None: + # Slice caps silently when num_examples exceeds the available lines, + # matching upstream: callers like test_basic_sanity_eagle3 pass a + # num_examples larger than the dataset on purpose. self._lines = self._lines[:num_examples] + def _setup_prefix_pool(self, all_lines: list, num_shots: int) -> int: + self._few_shot_prompt = get_few_shot_examples(all_lines, num_shots) + return num_shots + + def _build_prefix(self, idx: int) -> str: + return self._few_shot_prompt + def __call__(self, sampler: SamplerBase) -> EvalResult: def fn(idx: int) -> SingleEvalResult: question = get_one_example(self._lines, idx, include_answer=False) correct_answer = get_answer_value(self._lines[idx]["answer"]) - prompt_content = self._few_shot_prompt + question + prompt_content = self._build_prefix(idx) + question prompt_messages = [ sampler._pack_message(content=prompt_content, role="user") ] diff --git a/python/sglang/test/simple_eval_mixed_prefix_gsm8k.py b/python/sglang/test/simple_eval_mixed_prefix_gsm8k.py new file mode 100644 index 000000000..1a62d2c61 --- /dev/null +++ b/python/sglang/test/simple_eval_mixed_prefix_gsm8k.py @@ -0,0 +1,51 @@ +import random +from typing import Optional + +from sglang.test.simple_eval_gsm8k import GSM8KEval, get_one_example + + +class MixedPrefixGSM8KEval(GSM8KEval): + + def __init__( + self, + num_examples: Optional[int], + num_threads: int, + num_shots: int, + secondary_pool_size: int, + data_path: Optional[str], + seed: int, + ): + self._secondary_pool_size = secondary_pool_size + self._seed = seed + super().__init__( + num_examples=num_examples, + num_threads=num_threads, + num_shots=num_shots, + data_path=data_path, + ) + + def _setup_prefix_pool(self, all_lines: list, num_shots: int) -> int: + overall_pool_size = num_shots + self._secondary_pool_size + if len(all_lines) < overall_pool_size + 1: + raise ValueError( + f"GSM8K dataset has {len(all_lines)} examples but mixed-prefix " + f"eval needs at least {overall_pool_size + 1} " + f"(num_shots {num_shots} + secondary " + f"{self._secondary_pool_size} + 1 test)." + ) + self._primary_shots = all_lines[:num_shots] + self._secondary_pool = all_lines[num_shots:overall_pool_size] + return overall_pool_size + + def _build_prefix(self, idx: int) -> str: + rng = random.Random(self._seed + idx) + num_primary = rng.randint(0, self._num_shots) + secondary_size = rng.randint(0, self._secondary_pool_size) + secondary_indices = rng.sample(range(len(self._secondary_pool)), secondary_size) + primary = self._primary_shots[:num_primary] + secondary = [self._secondary_pool[i] for i in secondary_indices] + combined = primary + secondary + return "".join( + get_one_example(combined, i, include_answer=True) + "\n\n" + for i in range(len(combined)) + ) diff --git a/test/registered/unit/bench/test_mixed_prefix_gsm8k.py b/test/registered/unit/bench/test_mixed_prefix_gsm8k.py new file mode 100644 index 000000000..fb4b39bfb --- /dev/null +++ b/test/registered/unit/bench/test_mixed_prefix_gsm8k.py @@ -0,0 +1,177 @@ +import json +import os +import tempfile +import unittest +from typing import List, Tuple + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.simple_eval_gsm8k import get_one_example +from sglang.test.simple_eval_mixed_prefix_gsm8k import MixedPrefixGSM8KEval +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-b-test-cpu") + + +def _write_synthetic_dataset(path: str, n: int) -> None: + with open(path, "w") as f: + for i in range(n): + f.write( + json.dumps( + { + "question": f"Synthetic question {i}: what is {i} + {i}?", + "answer": f"The answer is {2 * i}. #### {2 * i}", + } + ) + + "\n" + ) + + +class TestMixedPrefixGSM8KEval(CustomTestCase): + NUM_SHOTS = 4 + SECONDARY_POOL_SIZE = 12 + NUM_EXAMPLES = 40 + + @classmethod + def setUpClass(cls): + cls._tmpdir = tempfile.TemporaryDirectory() + cls._data_path = os.path.join(cls._tmpdir.name, "synthetic.jsonl") + _write_synthetic_dataset(cls._data_path, 100) + + @classmethod + def tearDownClass(cls): + cls._tmpdir.cleanup() + + def _make_eval(self, seed: int = 42, num_examples=None): + return MixedPrefixGSM8KEval( + num_examples=( + num_examples if num_examples is not None else self.NUM_EXAMPLES + ), + num_threads=1, + num_shots=self.NUM_SHOTS, + secondary_pool_size=self.SECONDARY_POOL_SIZE, + data_path=self._data_path, + seed=seed, + ) + + def _primary_lines(self, evaluator) -> List[str]: + return [ + get_one_example(evaluator._primary_shots, j, include_answer=True) + "\n\n" + for j in range(self.NUM_SHOTS) + ] + + def _decompose(self, evaluator, prefix: str) -> Tuple[int, List[str]]: + k = 0 + for line in self._primary_lines(evaluator): + if prefix.startswith(line): + prefix = prefix[len(line) :] + k += 1 + else: + break + remainder_questions: List[str] = [] + if prefix: + chunks = prefix.split("\n\n") + for chunk in chunks: + if chunk.startswith("Question: "): + q_text = chunk[len("Question: ") :].split("\nAnswer:")[0] + remainder_questions.append(q_text) + return k, remainder_questions + + def test_primary_segment_is_strict_prefix_of_primary_shots(self): + e = self._make_eval() + for i in range(self.NUM_EXAMPLES): + k, _ = self._decompose(e, e._build_prefix(i)) + self.assertGreaterEqual(k, 0) + self.assertLessEqual(k, self.NUM_SHOTS) + + def test_remainder_questions_come_from_secondary_pool(self): + e = self._make_eval() + secondary_qs = {item["question"] for item in e._secondary_pool} + for i in range(self.NUM_EXAMPLES): + _, remainder = self._decompose(e, e._build_prefix(i)) + for q in remainder: + self.assertIn(q, secondary_qs) + + def test_remainder_no_duplicates_within_one_query(self): + e = self._make_eval() + for i in range(self.NUM_EXAMPLES): + _, remainder = self._decompose(e, e._build_prefix(i)) + self.assertEqual( + len(remainder), + len(set(remainder)), + f"query {i} has duplicate secondary samples", + ) + + def test_remainder_size_within_secondary_pool_bound(self): + e = self._make_eval() + for i in range(self.NUM_EXAMPLES): + _, remainder = self._decompose(e, e._build_prefix(i)) + self.assertGreaterEqual(len(remainder), 0) + self.assertLessEqual(len(remainder), self.SECONDARY_POOL_SIZE) + + def test_primary_depth_takes_multiple_values(self): + e = self._make_eval() + ks = { + self._decompose(e, e._build_prefix(i))[0] for i in range(self.NUM_EXAMPLES) + } + self.assertGreater(len(ks), 2, f"k values seen: {ks}") + + def test_secondary_size_takes_multiple_values(self): + e = self._make_eval() + sizes = { + len(self._decompose(e, e._build_prefix(i))[1]) + for i in range(self.NUM_EXAMPLES) + } + self.assertGreater(len(sizes), 2, f"sizes seen: {sizes}") + + def test_two_queries_share_min_primary_prefix(self): + e = self._make_eval() + lines = self._primary_lines(e) + prefixes = [e._build_prefix(i) for i in range(self.NUM_EXAMPLES)] + ks = [self._decompose(e, p)[0] for p in prefixes] + for i in range(self.NUM_EXAMPLES): + for j in range(i + 1, self.NUM_EXAMPLES): + shared = "".join(lines[: min(ks[i], ks[j])]) + self.assertTrue(prefixes[i].startswith(shared)) + self.assertTrue(prefixes[j].startswith(shared)) + + def test_build_prefix_is_deterministic(self): + a = self._make_eval(seed=42) + b = self._make_eval(seed=42) + for i in range(self.NUM_EXAMPLES): + self.assertEqual(a._build_prefix(i), b._build_prefix(i)) + + def test_seed_actually_matters(self): + a = self._make_eval(seed=42) + b = self._make_eval(seed=43) + differences = sum( + 1 + for i in range(self.NUM_EXAMPLES) + if a._build_prefix(i) != b._build_prefix(i) + ) + self.assertGreater(differences, self.NUM_EXAMPLES // 2) + + def test_pools_and_test_lines_pairwise_disjoint(self): + e = self._make_eval(num_examples=None) + primary_qs = {item["question"] for item in e._primary_shots} + secondary_qs = {item["question"] for item in e._secondary_pool} + test_qs = {item["question"] for item in e._lines} + self.assertEqual(primary_qs & secondary_qs, set()) + self.assertEqual(primary_qs & test_qs, set()) + self.assertEqual(secondary_qs & test_qs, set()) + + def test_insufficient_dataset_raises(self): + tiny = os.path.join(self._tmpdir.name, "tiny.jsonl") + _write_synthetic_dataset(tiny, n=5) + with self.assertRaises(ValueError): + MixedPrefixGSM8KEval( + num_examples=1, + num_threads=1, + num_shots=self.NUM_SHOTS, + secondary_pool_size=self.SECONDARY_POOL_SIZE, + data_path=tiny, + seed=42, + ) + + +if __name__ == "__main__": + unittest.main()