初始化项目，由ModelHub XC社区提供模型
Model: ayh015/myLightningOPD Source: Original Platform
This commit is contained in:
ModelHub XC
494
tools/eval_aime2024_vllm.py
Normal file
494
tools/eval_aime2024_vllm.py
Normal file
@@ -0,0 +1,494 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""
|
||||
Evaluate a HF-format model on AIME 2024 using vLLM.
|
||||
|
||||
Example:
|
||||
|
||||
CUDA_VISIBLE_DEVICES=0,1,2,3 python tools/eval_aime_vllm.py \
|
||||
--model /workspace/Lightning-OPD/checkpoints/qwen3-4b-lightning-opd-hf \
|
||||
--num-gpus 4 \
|
||||
--output outputs/eval_aime2024_qwen3_4b_lightning_opd.jsonl
|
||||
|
||||
Paper-style AIME setting:
|
||||
temperature = 0.6
|
||||
top_p = 0.95
|
||||
max_tokens = 32768
|
||||
n_samples = 32
|
||||
metric = average pass@1
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser()
|
||||
|
||||
# Model / hardware
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Path to HF-format model checkpoint.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-gpus",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Tensor parallel size for vLLM.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--gpu-ids",
|
||||
type=str,
|
||||
default=None,
|
||||
help='Optional CUDA_VISIBLE_DEVICES, e.g. "0,1,2,3". Must match --num-gpus.',
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dtype",
|
||||
type=str,
|
||||
default="bfloat16",
|
||||
choices=["auto", "float16", "bfloat16", "float32"],
|
||||
)
|
||||
parser.add_argument(
|
||||
"--gpu-memory-utilization",
|
||||
type=float,
|
||||
default=0.90,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--trust-remote-code",
|
||||
action="store_true",
|
||||
default=True,
|
||||
)
|
||||
|
||||
# Dataset
|
||||
parser.add_argument(
|
||||
"--dataset",
|
||||
type=str,
|
||||
default="AI-MO/aimo-validation-aime",
|
||||
help="HF dataset name for AIME 2024.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--split",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Dataset split. If not set, use the first available split.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--hf-cache",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Optional HuggingFace cache dir.",
|
||||
)
|
||||
|
||||
# Generation settings from paper
|
||||
parser.add_argument("--n-samples", type=int, default=32)
|
||||
parser.add_argument("--temperature", type=float, default=0.6)
|
||||
parser.add_argument("--top-p", type=float, default=0.95)
|
||||
parser.add_argument("--max-tokens", type=int, default=32768)
|
||||
parser.add_argument("--seed", type=int, default=42)
|
||||
|
||||
# Prompt / tokenizer behavior
|
||||
parser.add_argument(
|
||||
"--no-chat-template",
|
||||
action="store_true",
|
||||
help="Do not apply tokenizer chat template; use raw prompt.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable-thinking",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Pass enable_thinking=True to Qwen3 chat template if supported.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--disable-thinking",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Pass enable_thinking=False to Qwen3 chat template if supported.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--assistant-prefix",
|
||||
type=str,
|
||||
default=None,
|
||||
help=(
|
||||
"Optional text appended after the assistant generation prompt. "
|
||||
"Example for no-thinking style: '<think>\\n\\n</think>\\n\\n'"
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt-template",
|
||||
type=str,
|
||||
default="train",
|
||||
choices=["train", "paper", "chat"],
|
||||
help=(
|
||||
"Prompt template to use. "
|
||||
"'train' uses the exact training ChatML prompt; "
|
||||
"'paper' uses the paper-style ChatML prompt; "
|
||||
"'chat' uses build_problem_prompt plus tokenizer.apply_chat_template."
|
||||
),
|
||||
)
|
||||
|
||||
# Output / debug
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
type=str,
|
||||
default="outputs/eval_aime2024_vllm.jsonl",
|
||||
help="Path to save per-sample generations and scores.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--limit",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Limit number of problems for quick debugging.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--print-examples",
|
||||
type=int,
|
||||
default=3,
|
||||
help="Print first N examples with predictions.",
|
||||
)
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def normalize_answer(x: Any) -> Optional[str]:
|
||||
if x is None:
|
||||
return None
|
||||
|
||||
s = str(x).strip()
|
||||
s = s.replace("$", "")
|
||||
s = s.replace(",", "")
|
||||
s = s.replace("\\,", "")
|
||||
s = s.strip()
|
||||
|
||||
# Remove simple latex wrappers.
|
||||
s = s.replace("\\text", "")
|
||||
s = s.replace("\\mathrm", "")
|
||||
s = s.strip("{}").strip()
|
||||
|
||||
# AIME answers are integers from 0 to 999.
|
||||
m = re.search(r"-?\d+", s)
|
||||
if m is None:
|
||||
return None
|
||||
|
||||
try:
|
||||
return str(int(m.group(0)))
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
|
||||
def extract_last_boxed(text: str) -> Optional[str]:
|
||||
"""
|
||||
Extract the last \\boxed{...}. Handles simple nested braces better than regex.
|
||||
"""
|
||||
marker = r"\boxed{"
|
||||
positions = [m.start() for m in re.finditer(re.escape(marker), text)]
|
||||
if not positions:
|
||||
return None
|
||||
|
||||
for start in reversed(positions):
|
||||
i = start + len(marker)
|
||||
depth = 1
|
||||
chars = []
|
||||
|
||||
while i < len(text):
|
||||
ch = text[i]
|
||||
if ch == "{":
|
||||
depth += 1
|
||||
chars.append(ch)
|
||||
elif ch == "}":
|
||||
depth -= 1
|
||||
if depth == 0:
|
||||
return "".join(chars)
|
||||
chars.append(ch)
|
||||
else:
|
||||
chars.append(ch)
|
||||
i += 1
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def extract_answer(text: str) -> Optional[str]:
|
||||
# Prefer boxed answer, because the actual training prompt asks for \boxed{$Answer}.
|
||||
boxed = extract_last_boxed(text)
|
||||
if boxed is not None:
|
||||
ans = normalize_answer(boxed)
|
||||
if ans is not None:
|
||||
return ans
|
||||
|
||||
# Fallback: final Answer: line.
|
||||
matches = re.findall(r"Answer:\s*([^\n]+)", text, flags=re.IGNORECASE)
|
||||
if matches:
|
||||
ans = normalize_answer(matches[-1])
|
||||
if ans is not None:
|
||||
return ans
|
||||
|
||||
# No valid final answer.
|
||||
return None
|
||||
|
||||
def build_exact_training_prompt(problem: str) -> str:
|
||||
return (
|
||||
"<|im_start|>user\n"
|
||||
"Solve the following math problem step by step. "
|
||||
"The last line of your response should be of the form "
|
||||
"Answer: \\boxed{$Answer} where $Answer is the answer to the problem.\n\n"
|
||||
f"{problem}\n\n"
|
||||
"Remember to put your answer on its own line after \"Answer:\"."
|
||||
"<|im_end|>\n"
|
||||
"<|im_start|>assistant\n\n\n\n"
|
||||
)
|
||||
|
||||
def build_problem_prompt(problem: str) -> str:
|
||||
return (
|
||||
f"{problem}\n\n"
|
||||
"Please reason step by step, and put your final answer within \\boxed{}."
|
||||
)
|
||||
|
||||
def build_paper_math_eval_prompt(problem: str) -> str:
|
||||
return (
|
||||
"<|im_start|>user\n"
|
||||
f"Question: {problem}\n"
|
||||
"Please reason step by step, and put your final answer within \\boxed{}.\n"
|
||||
"<|im_end|>\n"
|
||||
"<|im_start|>assistant\n"
|
||||
)
|
||||
|
||||
def get_field(ex: Dict[str, Any], candidates: List[str]) -> Any:
|
||||
for key in candidates:
|
||||
if key in ex and ex[key] is not None:
|
||||
return ex[key]
|
||||
raise KeyError(f"Cannot find any of {candidates}. Example keys: {list(ex.keys())}")
|
||||
|
||||
# def build_problem_prompt(problem: str) -> str:
|
||||
# return (
|
||||
# "Solve the following math problem step by step. "
|
||||
# "The last line of your response should be of the form "
|
||||
# "Answer: $Answer (without quotes) where $Answer is the answer to the problem.\n\n"
|
||||
# f"{problem}\n\n"
|
||||
# 'Remember to put your answer on its own line after "Answer:".'
|
||||
# )
|
||||
|
||||
def apply_chat_template(tokenizer, prompt: str, args) -> str:
|
||||
if args.no_chat_template:
|
||||
return prompt
|
||||
|
||||
messages = [{"role": "user", "content": prompt}]
|
||||
|
||||
kwargs = {
|
||||
"tokenize": False,
|
||||
"add_generation_prompt": True,
|
||||
}
|
||||
|
||||
# Qwen3 tokenizer may support enable_thinking.
|
||||
if args.enable_thinking and args.disable_thinking:
|
||||
raise ValueError("Do not set both --enable-thinking and --disable-thinking.")
|
||||
|
||||
if args.enable_thinking:
|
||||
kwargs["enable_thinking"] = True
|
||||
elif args.disable_thinking:
|
||||
kwargs["enable_thinking"] = False
|
||||
|
||||
try:
|
||||
text = tokenizer.apply_chat_template(messages, **kwargs)
|
||||
except TypeError:
|
||||
# Some tokenizers do not accept enable_thinking.
|
||||
kwargs.pop("enable_thinking", None)
|
||||
text = tokenizer.apply_chat_template(messages, **kwargs)
|
||||
|
||||
if args.assistant_prefix is not None:
|
||||
text += args.assistant_prefix.encode("utf-8").decode("unicode_escape")
|
||||
|
||||
return text
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
|
||||
if args.gpu_ids is not None:
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_ids
|
||||
|
||||
if args.hf_cache is not None:
|
||||
os.environ["HF_HOME"] = args.hf_cache
|
||||
os.environ["HF_DATASETS_CACHE"] = str(Path(args.hf_cache) / "datasets")
|
||||
os.environ["HF_HUB_CACHE"] = str(Path(args.hf_cache) / "hub")
|
||||
|
||||
# Import after CUDA_VISIBLE_DEVICES is set.
|
||||
from datasets import load_dataset
|
||||
from transformers import AutoTokenizer
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
output_path = Path(args.output)
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
print("=" * 80)
|
||||
print("AIME 2024 vLLM Evaluation")
|
||||
print("=" * 80)
|
||||
print(f"model: {args.model}")
|
||||
print(f"dataset: {args.dataset}")
|
||||
print(f"num_gpus / TP size: {args.num_gpus}")
|
||||
print(f"CUDA_VISIBLE_DEVICES: {os.environ.get('CUDA_VISIBLE_DEVICES')}")
|
||||
print(f"n_samples/problem: {args.n_samples}")
|
||||
print(f"temperature: {args.temperature}")
|
||||
print(f"top_p: {args.top_p}")
|
||||
print(f"max_tokens: {args.max_tokens}")
|
||||
print(f"prompt_template: {args.prompt_template}")
|
||||
print(f"assistant_prefix: {repr(args.assistant_prefix)}")
|
||||
print(f"output: {args.output}")
|
||||
print("=" * 80)
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(
|
||||
args.model,
|
||||
trust_remote_code=args.trust_remote_code,
|
||||
)
|
||||
|
||||
dataset_dict = load_dataset(args.dataset)
|
||||
split = args.split or list(dataset_dict.keys())[0]
|
||||
data = dataset_dict[split]
|
||||
|
||||
# AI-MO/aimo-validation-aime contains multiple AIME years.
|
||||
# Keep only AIME 2024 examples.
|
||||
if "url" in data.column_names:
|
||||
data = data.filter(lambda ex: "2024_AIME" in ex["url"])
|
||||
else:
|
||||
raise ValueError(
|
||||
"Expected a `url` column for filtering AIME 2024, "
|
||||
f"but got columns: {data.column_names}"
|
||||
)
|
||||
|
||||
print(f"Filtered AIME 2024 examples: {len(data)}")
|
||||
for i in range(min(3, len(data))):
|
||||
print(i, data[i].get("url", "NO_URL"), data[i].get("answer", "NO_ANSWER"))
|
||||
|
||||
assert len(data) == 30, f"Expected 30 AIME 2024 problems, got {len(data)}"
|
||||
|
||||
if args.limit is not None:
|
||||
data = data.select(range(min(args.limit, len(data))))
|
||||
|
||||
print(f"Loaded split: {split}")
|
||||
print(f"Number of problems: {len(data)}")
|
||||
print(f"First example keys: {list(data[0].keys())}")
|
||||
|
||||
prompts = []
|
||||
examples = []
|
||||
|
||||
for idx, ex in enumerate(data):
|
||||
problem = get_field(ex, ["problem", "question", "prompt"])
|
||||
gold_raw = get_field(ex, ["answer", "final_answer", "target", "solution"])
|
||||
gold = normalize_answer(gold_raw)
|
||||
|
||||
if args.prompt_template == "train":
|
||||
# Raw ChatML prompt matching the OPD training parquet.
|
||||
# Do NOT call apply_chat_template again, otherwise ChatML will be nested.
|
||||
full_prompt = build_exact_training_prompt(problem)
|
||||
elif args.prompt_template == "paper":
|
||||
# Raw ChatML prompt matching the paper-style math evaluation prompt.
|
||||
# Do NOT call apply_chat_template again.
|
||||
full_prompt = build_paper_math_eval_prompt(problem)
|
||||
elif args.prompt_template == "chat":
|
||||
# Normal user-content prompt; tokenizer will add ChatML.
|
||||
raw_prompt = build_problem_prompt(problem)
|
||||
full_prompt = apply_chat_template(tokenizer, raw_prompt, args)
|
||||
else:
|
||||
raise ValueError(f"Unknown prompt_template: {args.prompt_template}")
|
||||
|
||||
# Important: when using raw ChatML prompts, apply_chat_template() is bypassed.
|
||||
# Therefore --assistant-prefix must be appended here, not only inside apply_chat_template().
|
||||
# This is useful for Qwen3 no-thinking style, e.g.:
|
||||
# --assistant-prefix '<think>\\n\\n</think>\\n\\n'
|
||||
if args.prompt_template in {"train", "paper"} and args.assistant_prefix is not None:
|
||||
full_prompt += args.assistant_prefix.encode("utf-8").decode("unicode_escape")
|
||||
|
||||
prompts.append(full_prompt)
|
||||
examples.append(
|
||||
{
|
||||
"idx": idx,
|
||||
"problem": problem,
|
||||
"gold_raw": gold_raw,
|
||||
"gold": gold,
|
||||
"prompt": full_prompt,
|
||||
}
|
||||
)
|
||||
|
||||
llm = LLM(
|
||||
model=args.model,
|
||||
tensor_parallel_size=args.num_gpus,
|
||||
dtype=args.dtype,
|
||||
trust_remote_code=args.trust_remote_code,
|
||||
gpu_memory_utilization=args.gpu_memory_utilization,
|
||||
seed=args.seed,
|
||||
)
|
||||
|
||||
sampling_params = SamplingParams(
|
||||
n=args.n_samples,
|
||||
temperature=args.temperature,
|
||||
top_p=args.top_p,
|
||||
#repetition_penalty=1.05,
|
||||
max_tokens=args.max_tokens,
|
||||
stop=["<|im_end|>"],
|
||||
)
|
||||
|
||||
outputs = llm.generate(prompts, sampling_params)
|
||||
|
||||
total = 0
|
||||
correct = 0
|
||||
per_problem_records = []
|
||||
|
||||
with output_path.open("w", encoding="utf-8") as f:
|
||||
for ex, out in zip(examples, outputs):
|
||||
sample_records = []
|
||||
problem_correct = 0
|
||||
|
||||
for sample_id, completion in enumerate(out.outputs):
|
||||
text = completion.text
|
||||
pred = extract_answer(text)
|
||||
is_correct = pred == ex["gold"]
|
||||
|
||||
total += 1
|
||||
correct += int(is_correct)
|
||||
problem_correct += int(is_correct)
|
||||
|
||||
record = {
|
||||
"idx": ex["idx"],
|
||||
"sample_id": sample_id,
|
||||
"gold": ex["gold"],
|
||||
"gold_raw": str(ex["gold_raw"]),
|
||||
"pred": pred,
|
||||
"correct": is_correct,
|
||||
"completion": text,
|
||||
"finish_reason": completion.finish_reason,
|
||||
"stop_reason": getattr(completion, "stop_reason", None),
|
||||
}
|
||||
sample_records.append(record)
|
||||
f.write(json.dumps(record, ensure_ascii=False) + "\n")
|
||||
|
||||
problem_acc = problem_correct / max(1, len(out.outputs))
|
||||
per_problem_records.append(problem_acc)
|
||||
|
||||
if ex["idx"] < args.print_examples:
|
||||
print("-" * 80)
|
||||
print(f"Problem {ex['idx']}")
|
||||
print(f"Gold: {ex['gold']} | Correct samples: {problem_correct}/{len(out.outputs)}")
|
||||
print(f"First pred: {sample_records[0]['pred']}")
|
||||
print(f"First completion preview:\n{sample_records[0]['completion'][:1000]}")
|
||||
|
||||
avg_pass1_micro = correct / total if total > 0 else 0.0
|
||||
avg_pass1_macro = sum(per_problem_records) / len(per_problem_records) if per_problem_records else 0.0
|
||||
|
||||
print("=" * 80)
|
||||
print("Final results")
|
||||
print("=" * 80)
|
||||
print(f"Problems: {len(examples)}")
|
||||
print(f"Samples per problem: {args.n_samples}")
|
||||
print(f"Total samples: {total}")
|
||||
print(f"Correct samples: {correct}")
|
||||
print(f"Average pass@1 micro: {avg_pass1_micro:.6f} ({100 * avg_pass1_micro:.2f}%)")
|
||||
print(f"Average pass@1 macro: {avg_pass1_macro:.6f} ({100 * avg_pass1_macro:.2f}%)")
|
||||
print(f"Saved generations to: {output_path}")
|
||||
print("=" * 80)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user