diff --git a/bash/case/GLM52_API_TEST1.sh b/bash/case/GLM52_API_TEST1.sh index 909e09d..2389a09 100755 --- a/bash/case/GLM52_API_TEST1.sh +++ b/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[@]}" diff --git a/bash/case/README_API_TEST.md b/bash/case/README_API_TEST.md index 58cd4c6..7131536 100644 --- a/bash/case/README_API_TEST.md +++ b/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. 查看进度 diff --git a/bash/run.py b/bash/run.py index 2594454..ee3f914 100644 --- a/bash/run.py +++ b/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,