feat: support --thinking-budget-tokens to control reasoning budget via API

bash/case/GLM52_API_TEST1.sh

@@ -41,6 +41,7 @@ BATCH_SIZE="${EVAL_BATCH_SIZE:-4}"
 LIMIT="${EVAL_LIMIT:-none}"
 SEED="${EVAL_SEED:-42}"
 THINKING="${EVAL_THINKING:-false}"
+THINKING_BUDGET_TOKENS="${EVAL_THINKING_BUDGET_TOKENS:-}"
 MODE="${EVAL_MODE:-custom}"
 
 # --------------------------------------------------
@@ -59,6 +60,7 @@ while [[ $# -gt 0 ]]; do
         --seed)          SEED="$2"; shift 2 ;;
         --thinking)      THINKING="true"; shift ;;
         --no-thinking)   THINKING="false"; shift ;;
+        --thinking-budget-tokens) THINKING_BUDGET_TOKENS="$2"; shift 2 ;;
         --mode)          MODE="$2"; shift 2 ;;
         -h|--help)
             grep '^# ' "$0" | sed 's/^# //'
@@ -143,6 +145,10 @@ if [[ "$THINKING" == "true" ]]; then
     ARGS+=(--thinking)
 fi
 
+if [[ -n "$THINKING_BUDGET_TOKENS" ]]; then
+    ARGS+=(--thinking-budget-tokens "$THINKING_BUDGET_TOKENS")
+fi
+
 echo "============================================================"
 echo "API 评测启动"
 echo "Mode:         $MODE"
@@ -154,6 +160,9 @@ echo "Config:       $CONFIG"
 echo "Batch size:   $BATCH_SIZE"
 echo "Limit:        $LIMIT"
 echo "Thinking:     $THINKING"
+if [[ -n "$THINKING_BUDGET_TOKENS" ]]; then
+    echo "Thinking budget_tokens: $THINKING_BUDGET_TOKENS"
+fi
 echo "============================================================"
 
 python bash/run.py "${ARGS[@]}"

bash/case/README_API_TEST.md

@@ -63,6 +63,27 @@ export EVAL_FOLDER_NAME="GLM52-API-Test"
 bash bash/case/GLM52_API_TEST1.sh
 ```
 
+### 2.4 控制 thinking budget
+
+如果模型 API 支持 `thinking.budget_tokens`，可以传入 `--thinking-budget-tokens`：
+
+```bash
+# 尝试用 budget_tokens=0 关闭 thinking
+bash bash/case/GLM52_API_TEST1.sh \
+    --thinking \
+    --thinking-budget-tokens 0 \
+    --datasets aime24 \
+    --folder-name no-thinking-test
+```
+
+等价于在请求体里加入：
+
+```json
+"extra_body": {
+  "thinking": {"type": "enabled", "budget_tokens": 0}
+}
+```
+
 ## 3. 常用命令行参数
 
 | 参数 | 说明 |
@@ -78,6 +99,7 @@ bash bash/case/GLM52_API_TEST1.sh
 | `--limit` | 每个 benchmark 最多测多少条，`none` 表示全量 |
 | `--thinking` | 启用 thinking 模式 |
 | `--no-thinking` | 关闭 thinking 模式 |
+| `--thinking-budget-tokens` | 控制 thinking budget，例如 `0` 尝试关闭 thinking |
 
 ## 4. 查看进度
 

bash/run.py

@@ -262,6 +262,9 @@ def build_parser():
                         help='Disable thinking mode (default)')
     parser.add_argument('--thinking-max-tokens-scale', type=float, default=1.0,
                         help='Scale max_tokens by this factor when --thinking is enabled (default: %(default)s)')
+    parser.add_argument('--thinking-budget-tokens', type=int, default=None,
+                        help='Add extra_body.thinking={"type": "enabled", "budget_tokens": N} to control thinking budget. '
+                             'Set to 0 to attempt disabling thinking via API.')
     parser.add_argument('--max-tokens-add', type=int, default=0,
                         help='Add this many tokens to every benchmark max_tokens (applied after scale)')
 
@@ -420,7 +423,7 @@ def load_dataset_configs(config_path: str):
         return yaml.safe_load(f)
 
 
-def configure_thinking(generation_config: dict, enable: bool) -> dict:
+def configure_thinking(generation_config: dict, enable: bool, thinking_budget_tokens: int = None) -> dict:
     extra_body = generation_config.get('extra_body', {})
     chat_template_kwargs = extra_body.get('chat_template_kwargs', {})
     if enable:
@@ -429,6 +432,13 @@ def configure_thinking(generation_config: dict, enable: bool) -> dict:
         chat_template_kwargs.pop('thinking', None)
     if chat_template_kwargs:
         extra_body['chat_template_kwargs'] = chat_template_kwargs
+
+    # 支持 API 级别的 thinking budget 控制（如 {"type": "enabled", "budget_tokens": 0}）
+    if thinking_budget_tokens is not None:
+        extra_body['thinking'] = {'type': 'enabled', 'budget_tokens': thinking_budget_tokens}
+    else:
+        extra_body.pop('thinking', None)
+
     if extra_body:
         generation_config['extra_body'] = extra_body
     return generation_config
@@ -461,6 +471,7 @@ def build_task_config(
     run_idx: int = 0,
     thinking_max_tokens_scale: float = 1.0,
     max_tokens_add: int = 0,
+    thinking_budget_tokens: int = None,
 ) -> TaskConfig:
     if run_idx > 0:
         work_dir = Path(output_dir) / dataset_name / f'seed_{seed}_run_{run_idx}'
@@ -469,7 +480,11 @@ def build_task_config(
     work_dir.mkdir(parents=True, exist_ok=True)
     work_dir = str(work_dir)
 
-    generation_config = configure_thinking(deepcopy(ds_cfg['generation_config']), enable_thinking)
+    generation_config = configure_thinking(
+        deepcopy(ds_cfg['generation_config']), enable_thinking, thinking_budget_tokens
+    )
+    if thinking_budget_tokens is not None:
+        print(f'   [thinking] budget_tokens: {thinking_budget_tokens}')
     if enable_thinking and thinking_max_tokens_scale != 1.0:
         original_max_tokens = generation_config.get('max_tokens', 32768)
         scaled = int(original_max_tokens * thinking_max_tokens_scale)
@@ -709,6 +724,8 @@ def main():
     print(f'Thinking: {enable_thinking}')
     if enable_thinking:
         print(f'Thinking max_tokens scale: {args.thinking_max_tokens_scale}')
+    if args.thinking_budget_tokens is not None:
+        print(f'Thinking budget_tokens: {args.thinking_budget_tokens}')
     print(f'Seed: {args.seed}')
     print(f'Batch Size: {args.batch_size}')
     print(f'Tokenizer Path: {args.tokenizer_path}')
@@ -738,6 +755,7 @@ def main():
             run_idx=run_idx,
             thinking_max_tokens_scale=args.thinking_max_tokens_scale,
             max_tokens_add=args.max_tokens_add,
+            thinking_budget_tokens=args.thinking_budget_tokens,
         )
         try:
             run_and_summarize(task_cfg, write_summary_flag, str(model_output_dir), args.model,