From a9d59776cc9c75b4bd286b0e8b0ed49b65c68390 Mon Sep 17 00:00:00 2001
From: Ke Bao <ispobaoke@gmail.com>
Date: Fri, 13 Feb 2026 18:08:57 +0800
Subject: [PATCH] Enhence gsm8k test (#18791)

---
 benchmark/gsm8k/bench_sglang.py | 46 ++++++++++++++++++++++++++++++---
 1 file changed, 42 insertions(+), 4 deletions(-)

diff --git a/benchmark/gsm8k/bench_sglang.py b/benchmark/gsm8k/bench_sglang.py
index 98c28b39b..be766cd9a 100644
--- a/benchmark/gsm8k/bench_sglang.py
+++ b/benchmark/gsm8k/bench_sglang.py
@@ -48,6 +48,18 @@ def main(args):
     # Select backend
     set_default_backend(select_sglang_backend(args))
 
+    # Load tokenizer if enable_thinking is set
+    tokenizer = None
+    if args.enable_thinking:
+        from transformers import AutoTokenizer
+
+        assert (
+            args.tokenizer_path is not None
+        ), "--tokenizer-path is required when --enable-thinking is set"
+        tokenizer = AutoTokenizer.from_pretrained(
+            args.tokenizer_path, trust_remote_code=True
+        )
+
     # Read data
     if args.platinum:
         print("Loading GSM8K Platinum dataset from HuggingFace...")
@@ -70,7 +82,16 @@ def main(args):
     questions = []
     labels = []
     for i in range(len(lines[:num_questions])):
-        questions.append(get_one_example(lines, i, False))
+        raw_question = few_shot_examples + get_one_example(lines, i, False)
+        if tokenizer is not None:
+            messages = [{"role": "user", "content": raw_question}]
+            raw_question = tokenizer.apply_chat_template(
+                messages,
+                tokenize=False,
+                add_generation_prompt=True,
+                enable_thinking=True,
+            )
+        questions.append(raw_question)
         labels.append(get_answer_value(lines[i]["answer"]))
     assert all(l != INVALID for l in labels)
     arguments = [{"question": q} for q in questions]
@@ -83,9 +104,11 @@ def main(args):
 
     @sgl.function
     def few_shot_gsm8k(s, question):
-        s += few_shot_examples + question
+        s += question
         s += sgl.gen(
-            "answer", max_tokens=512, stop=["Question", "Assistant:", "<|separator|>"]
+            "answer",
+            max_tokens=args.max_new_tokens,
+            stop=["Question", "Assistant:", "<|separator|>"],
         )
 
     #####################################
@@ -96,7 +119,8 @@ def main(args):
     tic = time.perf_counter()
     states = few_shot_gsm8k.run_batch(
         arguments,
-        temperature=0,
+        temperature=args.temperature,
+        top_p=args.top_p,
         num_threads=args.parallel,
         progress_bar=True,
     )
@@ -152,6 +176,20 @@ if __name__ == "__main__":
     parser.add_argument("--num-shots", type=int, default=5)
     parser.add_argument("--data-path", type=str, default="test.jsonl")
     parser.add_argument("--num-questions", type=int, default=200)
+    parser.add_argument("--max-new-tokens", type=int, default=512)
+    parser.add_argument("--temperature", type=float, default=0.0)
+    parser.add_argument("--top-p", type=float, default=1.0)
+    parser.add_argument(
+        "--enable-thinking",
+        action="store_true",
+        help="Enable thinking mode by wrapping prompts with chat template",
+    )
+    parser.add_argument(
+        "--tokenizer-path",
+        type=str,
+        default=None,
+        help="Path to tokenizer (required when --enable-thinking is set)",
+    )
     parser.add_argument(
         "--platinum",
         action="store_true",