[eval] GSM8k support for run_eval (#17041)

python/sglang/test/run_eval.py

@@ -129,6 +129,15 @@ def run_eval(args):
         from sglang.test.simple_eval_aime25 import AIME25Eval
 
         eval_obj = AIME25Eval(args.num_examples, args.num_threads)
+    elif args.eval_name == "gsm8k":
+        from sglang.test.simple_eval_gsm8k import GSM8KEval
+
+        eval_obj = GSM8KEval(
+            num_examples=args.num_examples,
+            num_threads=args.num_threads,
+            num_shots=getattr(args, "num_shots", 5),
+            data_path=getattr(args, "gsm8k_data_path", None),
+        )
     else:
         raise ValueError(f"Invalid eval name: {args.eval_name}")
 
@@ -268,6 +277,18 @@ if __name__ == "__main__":
         type=int,
         help="Minimum context length in characters for LongBench-v2",
     )
+    parser.add_argument(
+        "--num-shots",
+        type=int,
+        default=5,
+        help="Number of few-shot examples for GSM8K (default: 5)",
+    )
+    parser.add_argument(
+        "--gsm8k-data-path",
+        type=str,
+        default=None,
+        help="Path to GSM8K data file (e.g., test.jsonl)",
+    )
 
     args = parser.parse_args()
 

python/sglang/test/simple_eval_gsm8k.py

@@ -0,0 +1,99 @@
+# Adapted from https://github.com/openai/simple-evals/
+
+import ast
+import re
+from typing import Optional
+
+from sglang.test import simple_eval_common as common
+from sglang.test.simple_eval_common import (
+    HTML_JINJA,
+    Eval,
+    EvalResult,
+    SamplerBase,
+    SingleEvalResult,
+)
+from sglang.utils import download_and_cache_file, read_jsonl
+
+GSM8K_URL = "https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl"
+INVALID = -9999999
+
+
+def get_one_example(lines, i, include_answer):
+    ret = f"Question: {lines[i]['question']}\nAnswer:"
+    if include_answer:
+        ret += f" {lines[i]['answer']}"
+    return ret
+
+
+def get_few_shot_examples(lines, k):
+    return "".join(get_one_example(lines, i, True) + "\n\n" for i in range(k))
+
+
+def get_answer_value(answer_str):
+    answer_str = answer_str.replace(",", "")
+    numbers = re.findall(r"-?\d+\.?\d*", answer_str)
+    if len(numbers) < 1:
+        return INVALID
+    try:
+        return ast.literal_eval(numbers[-1])
+    except (SyntaxError, ValueError):
+        return INVALID
+
+
+class GSM8KEval(Eval):
+    def __init__(
+        self,
+        num_examples: Optional[int] = None,
+        num_threads: int = 64,
+        num_shots: int = 5,
+        data_path: Optional[str] = None,
+    ):
+        self._num_threads = num_threads
+        self._num_shots = num_shots
+
+        if data_path:
+            filename = data_path
+        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)
+
+        # The evaluation data should not include the few-shot examples to prevent data leakage.
+        self._lines = self._lines[num_shots:]
+        if num_examples is not None:
+            self._lines = self._lines[:num_examples]
+
+    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_messages = [
+                sampler._pack_message(content=prompt_content, role="user")
+            ]
+
+            try:
+                response_text = sampler(prompt_messages)
+            except Exception:
+                response_text = ""
+
+            extracted_answer = get_answer_value(response_text)
+            score = float(extracted_answer == correct_answer)
+
+            html = common.jinja_env.from_string(HTML_JINJA).render(
+                prompt_messages=prompt_messages,
+                next_message=dict(content=response_text, role="assistant"),
+                score=score,
+                correct_answer=correct_answer,
+                extracted_answer=extracted_answer,
+            )
+            convo = prompt_messages + [dict(content=response_text, role="assistant")]
+
+            return SingleEvalResult(html=html, score=score, convo=convo)
+
+        results = common.map_with_progress(
+            fn, list(range(len(self._lines))), num_threads=self._num_threads
+        )
+        return common.aggregate_results(results, default_stats=("mean", "std"))