from transformers import AutoTokenizer from vllm import LLM, SamplingParams from arguments import get_args from tqdm import tqdm import torch import os import json os.environ["TOKENIZERS_PARALLELISM"] = "false" def load_vllm_model(args): """Load a vLLM model with specified configuration. Args: args: Command-line arguments containing model configuration: - model_folder: Directory containing the model - model_name: Name of the model to load - tokenizer_folder: Directory containing the tokenizer - tokenizer_name: Name of the tokenizer to load - tensor_parallel_size: Number of GPUs for tensor parallelism - yarn_factor: Scaling factor for YaRN (Yet another RoPE extensioN method) - max_output_len: Maximum output length - seed: Random seed for reproducibility Returns: LLM: Initialized vLLM model instance """ tokenizer_path = os.path.join(args.tokenizer_folder, args.tokenizer_name) model_path = os.path.join(args.model_folder, args.model_name) tensor_parallel_size = args.tensor_parallel_size eager_mode = True if "DeepSeek-R1" in model_path else False print("eager_mode:", eager_mode) print("load tokenizer from %s" % tokenizer_path) print("load model from %s" % model_path) print("tensor_parallel_size:", tensor_parallel_size) if args.yarn_factor == 1: rope_scaling = None else: rope_scaling = {"rope_type":"yarn", "factor": args.yarn_factor, "original_max_position_embeddings":32768, "attention_factor": 0.8782488562869419} max_output_len = int(args.max_output_len * args.yarn_factor) model_vllm = LLM(model_path, tokenizer=tokenizer_path, max_model_len=max_output_len, trust_remote_code=True, tensor_parallel_size=tensor_parallel_size, enforce_eager=eager_mode, seed=args.seed, rope_scaling=rope_scaling ) return model_vllm def apply_template(prompt, tokenizer, think=True): """Apply chat template to format the prompt for model input. Args: prompt: Either a string containing a single user message, or a list of chat messages with 'role' and 'content' fields tokenizer: HuggingFace tokenizer with chat template support think: Whether to enable thinking mode (default: True) Returns: str: Formatted prompt string ready for model input Raises: ValueError: If prompt is neither a string nor a list """ if isinstance(prompt, str): chat = [ {"role": "user", "content": prompt}, ] elif isinstance(prompt, list): chat = prompt else: raise ValueError("prompt must be str or list") return tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True, enable_thinking=think) def get_prompt_list(args): """Load and preprocess prompts from the specified evaluation dataset. This function supports multiple benchmark datasets including: - Math: MATH, MATH500, GSM8K, Minerva Math, OmniMath, AIME - Coding: MBPP, HumanEval, LiveCodeBench - Multiple Choice: MMLU, MMLU Pro, GPQA - Instruction Following: IFEval, IFBench, MT-Bench - General: AlpacaEval, Arena-Hard Args: args: Command-line arguments containing: - eval_dataset: Name of the evaluation dataset - benchmark_folder: Root directory containing benchmark data - start_idx: Starting index for subsetting (optional) - end_idx: Ending index for subsetting (optional) - Various dataset-specific paths Returns: tuple: (prompt_list, qid_list) - prompt_list: List of formatted prompts ready for inference - qid_list: List of question IDs (None for some datasets) Raises: ValueError: If eval_dataset is not recognized """ if args.eval_dataset == "mbpp": from data.benchmark import preprocess_mbpp_chatml_template input_datapath = os.path.join(args.benchmark_folder, args.mbpp_path) prompt_list, qid_list = preprocess_mbpp_chatml_template(input_datapath) elif args.eval_dataset == "mbpp_sanitized": from data.benchmark import preprocess_mbpp_chatml_template input_datapath = os.path.join(args.benchmark_folder, args.mbpp_sanitized_path) prompt_list, qid_list = preprocess_mbpp_chatml_template(input_datapath) elif args.eval_dataset == "mbpp_plus": from data.benchmark import preprocess_mbpp_chatml_template input_datapath = os.path.join(args.benchmark_folder, args.mbpp_plus_path) prompt_list, qid_list = preprocess_mbpp_chatml_template(input_datapath) elif args.eval_dataset == "math": from data.benchmark import preprocess_math_zeroshot_chatml_template input_datapath = os.path.join(args.benchmark_folder, args.math_path) prompt_list = preprocess_math_zeroshot_chatml_template(input_datapath) qid_list = None elif args.eval_dataset == "math500": from data.benchmark import preprocess_math500_zeroshot_chatml_template input_datapath = os.path.join(args.benchmark_folder, args.math500_path) prompt_list = preprocess_math500_zeroshot_chatml_template(input_datapath, use_r1=args.use_r1) qid_list = None elif args.eval_dataset == "gsm8k": from data.benchmark import preprocess_gsm8k_zeroshot_raw input_datapath = os.path.join(args.benchmark_folder, args.gsm8k_path) prompt_list = preprocess_gsm8k_zeroshot_raw(input_datapath) qid_list = None elif args.eval_dataset == "humaneval": from data.benchmark import preprocess_humaneval_raw input_datapath = os.path.join(args.benchmark_folder, args.humaneval_path) prompt_list, qid_list = preprocess_humaneval_raw(input_datapath) elif args.eval_dataset == "mmlu": from data.benchmark import preprocess_mmlu_raw_template input_datapath = os.path.join(args.benchmark_folder, args.mmlu_path) prompt_list = preprocess_mmlu_raw_template(input_datapath) qid_list = None elif args.eval_dataset == "mmlu_r1": from data.benchmark import preprocess_mmlu_r1_raw_template_wdai input_datapath = os.path.join(args.benchmark_folder, args.mmlu_path) prompt_list = preprocess_mmlu_r1_raw_template_wdai(input_datapath) qid_list = None elif args.eval_dataset == "alpaca_eval": from data.benchmark import preprocess_alpaca_eval_raw input_datapath = os.path.join(args.benchmark_folder, args.alpaca_eval_path) prompt_list, qid_list = preprocess_alpaca_eval_raw(input_datapath) elif args.eval_dataset == "arena_hard": from data.benchmark import preprocess_arena_hard_raw input_datapath = os.path.join(args.benchmark_folder, args.arena_hard_path) prompt_list, qid_list = preprocess_arena_hard_raw(input_datapath) elif args.eval_dataset == "arena_hard_v2": from data.benchmark import preprocess_arena_hard_v2_raw input_datapath = os.path.join(args.benchmark_folder, args.arena_hard_v2_path) prompt_list, qid_list = preprocess_arena_hard_v2_raw(input_datapath) elif args.eval_dataset == "ifeval": from data.benchmark import preprocess_ifeval_raw input_datapath = os.path.join(args.benchmark_folder, args.ifeval_path) prompt_list, qid_list = preprocess_ifeval_raw(input_datapath) elif args.eval_dataset == "ifeval_training": from data.benchmark import preprocess_ifeval_raw input_datapath = os.path.join(args.benchmark_folder, args.ifeval_training_path) prompt_list, qid_list = preprocess_ifeval_raw(input_datapath) elif args.eval_dataset == "ifbench": from data.benchmark import preprocess_ifbench_raw input_datapath = os.path.join(args.benchmark_folder, args.ifbench_path) prompt_list, qid_list = preprocess_ifbench_raw(input_datapath) elif args.eval_dataset == "mtbench_firstturn": from data.benchmark import preprocess_mtbench_firstturn_raw input_datapath = os.path.join(args.benchmark_folder, args.mtbench_path) prompt_list, qid_list = preprocess_mtbench_firstturn_raw(input_datapath) elif args.eval_dataset == "mtbench_secondturn": from data.benchmark import preprocess_mtbench_secondturn_raw input_datapath = os.path.join(args.benchmark_folder, args.mtbench_path) prompt_list, qid_list = preprocess_mtbench_secondturn_raw(input_datapath, args.model_output_path) elif args.eval_dataset == "lcb5": from data.benchmark import preprocess_livecodebench_raw input_datapath = os.path.join(args.benchmark_folder, args.livecodebench_path) prompt_list, qid_list = preprocess_livecodebench_raw(input_datapath) elif args.eval_dataset == "lcb6": from data.benchmark import preprocess_livecodebench_raw print(args) input_datapath = os.path.join(args.benchmark_folder, args.livecodebench6_path) prompt_list, qid_list = preprocess_livecodebench_raw(input_datapath) elif args.eval_dataset == "minerva_math": from data.benchmark import preprocess_minerva_math_chatml_template input_datapath = os.path.join(args.benchmark_folder, args.minervamath_path) prompt_list = preprocess_minerva_math_chatml_template(input_datapath) qid_list = None elif args.eval_dataset == "gaokao2023en": from data.benchmark import preprocess_gaokao2023en_chatml_template input_datapath = os.path.join(args.benchmark_folder, args.gaokao2023en_path) prompt_list = preprocess_gaokao2023en_chatml_template(input_datapath) qid_list = None elif args.eval_dataset == "olympiadbench": from data.benchmark import preprocess_olympiadbench_chatml_template input_datapath = os.path.join(args.benchmark_folder, args.olympiadbench_path) prompt_list = preprocess_olympiadbench_chatml_template(input_datapath) qid_list = None elif args.eval_dataset == "collegemath": from data.benchmark import preprocess_collegemath_chatml_template input_datapath = os.path.join(args.benchmark_folder, args.collegemath_path) prompt_list = preprocess_collegemath_chatml_template(input_datapath) qid_list = None elif args.eval_dataset == "mmlu_stem": from data.benchmark import preprocess_mmlu_stem_chatml_template input_datapath = os.path.join(args.benchmark_folder, args.mmlustem_path) prompt_list = preprocess_mmlu_stem_chatml_template(input_datapath) qid_list = None elif args.eval_dataset == "amc23": from data.benchmark import preprocess_amc23_chatml_template input_datapath = os.path.join(args.benchmark_folder, args.amc23_path) prompt_list = preprocess_amc23_chatml_template(input_datapath) qid_list = None elif args.eval_dataset == "aime24": from data.benchmark import preprocess_aime24_raw input_datapath = os.path.join(args.benchmark_folder, args.aime24_path) prompt_list = preprocess_aime24_raw(input_datapath) qid_list = None elif args.eval_dataset == "aime25": from data.benchmark import preprocess_aime25_raw input_datapath = os.path.join(args.benchmark_folder, args.aime25_path) prompt_list = preprocess_aime25_raw(input_datapath) qid_list = None elif args.eval_dataset == "omnimath": from data.benchmark import preprocess_omnimath_chatml_template input_datapath = os.path.join(args.benchmark_folder, args.omnimath_path) prompt_list = preprocess_omnimath_chatml_template(input_datapath) qid_list = None elif args.eval_dataset == "gpqa_diamond": from data.benchmark import preprocess_gpqa_raw_template input_datapath = os.path.join(args.benchmark_folder, args.gpqa_diamond_path) prompt_list = preprocess_gpqa_raw_template(input_datapath, use_r1=args.use_r1) qid_list = None elif args.eval_dataset == "mmlu_pro": from data.benchmark import preprocess_mmlu_pro_zero_shot_raw_template input_datapath = os.path.join(args.benchmark_folder, args.mmlupro_path) fewshot_datapath = os.path.join(args.benchmark_folder, args.mmlupro_fewshot_path) prompt_list = preprocess_mmlu_pro_zero_shot_raw_template(input_datapath, fewshot_datapath) qid_list = None else: raise ValueError("please input a correct eval_dataset name!") print("number of total prompt_list:", len(prompt_list)) if args.start_idx != -1 and args.end_idx != -1: print("getting data from %d to %d" % (args.start_idx, args.end_idx)) prompt_list = prompt_list[args.start_idx:args.end_idx] if qid_list: qid_list = qid_list[args.start_idx:args.end_idx] print("number of test samples in the dataset:", len(prompt_list)) return prompt_list, qid_list def main(): """Main function to run inference on evaluation benchmarks. This function: 1. Parses command-line arguments 2. Loads the vLLM model and tokenizer 3. Loads test data from the specified benchmark 4. Runs batched inference with specified sampling parameters 5. Post-processes outputs (extracts reasoning, handles special tokens) 6. Saves results to JSONL format The output directory structure is: {model_folder}/{model_name}/outputs_vllm073[_topp{topp}_seed{seed}]/{eval_dataset}.jsonl """ args = get_args(add_evaluation=True) if args.device_id: os.environ["CUDA_VISIBLE_DEVICES"] = args.device_id for key, value in vars(args).items(): print(f"{key}: {value}") ## load model model_vllm = load_vllm_model(args) tokenizer_path = os.path.join(args.tokenizer_folder, args.tokenizer_name) tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, trust_remote_code=True) ## load test data prompt_list, qid_list = get_prompt_list(args) ## run inference max_output_len = int(args.max_output_len * args.yarn_factor) print("args.max_output_len:", max_output_len) if args.topp < 1: sampling_params = SamplingParams(temperature=args.temperature, top_p=args.topp, max_tokens=max_output_len, seed=args.seed) print("args.seed:", args.seed) print("args.topp:", args.topp) print("args.temperature:", args.temperature) else: sampling_params = SamplingParams(temperature=args.temperature, top_k=args.topk, max_tokens=max_output_len, seed=args.seed) print("Greedy decoding", args.temperature, args.topk) output_list = [] for i in tqdm(range(0, len(prompt_list), args.batch_size)): batch_prompts = prompt_list[i:i + args.batch_size] if qid_list: batch_qids = qid_list[i:i + args.batch_size] if args.eval_dataset in ("ifeval", "ifbench", "alpaca_eval", "arena_hard", "mtbench_secondturn", "mtbench_firstturn", "mmlu", "humaneval", "gsm8k", "mmlu_r1", "aime24", "aime25", "arena_hard_v2", "lcb5", "lcb6", "ifeval_training", "mmlu_pro", "gpqa_diamond"): raw_prompts = batch_prompts batch_prompts = [apply_template(prompt, tokenizer, think=args.think) for prompt in batch_prompts] for i in range(3): print(batch_prompts[i]) outputs = model_vllm.generate(batch_prompts, sampling_params) if torch.distributed.is_initialized() and torch.distributed.get_rank() != 0: continue for j, output in enumerate(outputs): generated_text = output.outputs[0].text if "<|im_end|>" in generated_text: idx = generated_text.index("<|im_end|>") generated_text = generated_text[:idx] if "<|end_of_text|>" in generated_text: idx = generated_text.index("<|end_of_text|>") generated_text = generated_text[:idx] if "<|eot_id|>" in generated_text: idx = generated_text.index("<|eot_id|>") generated_text = generated_text[:idx] reason = False reason_text = '' if "" in generated_text: idx = generated_text.index("") reason_text = generated_text[:idx] generated_text = generated_text[idx + len(""):].strip() reason = True if qid_list: qid = batch_qids[j] if args.eval_dataset in ("ifeval", "ifeval_training", "ifbench"): output_dict = {"task_id": qid, "prompt": raw_prompts[j], "response": generated_text, "reason": reason, "reason_text": reason_text} elif args.eval_dataset == 'arena_hard': output_dict = {"question_id": qid, "model_id": args.model_name, "choices": [{"index": 0, "turns": [{"content": generated_text}]}], "reason": reason, "reason_text": reason_text } elif args.eval_dataset == 'arena_hard_v2': output_dict = {"uid": qid, "model": args.model_name, "messages": [{"role": "user", "content": raw_prompts[j]}, {"role": "assistant", "content": {"answer": generated_text}}], "reason": reason, "reason_text": reason_text } elif args.eval_dataset == 'alpaca_eval': output_dict = {"question_id": qid, "model_id": args.model_name, "instruction": raw_prompts[j], "datasplit": "eval", "output": generated_text, "reason": reason, "reason_text": reason_text} else: output_dict = {"task_id": qid, "output": generated_text, "reason": reason, "reason_text": reason_text} output_list.append(output_dict) else: output_dict = {"output": generated_text, "reason": reason, "reason_text": reason_text} output_list.append(output_dict) if torch.distributed.is_initialized() and torch.distributed.get_rank() != 0: return ## write to output_datapath if args.topp < 1: foldername = "outputs_vllm073_topp{}_seed{}".format(args.topp, args.seed) else: foldername = "outputs_vllm073" if not args.think: foldername = "nothink_" + foldername output_folder = os.path.join(os.path.join(args.model_folder, args.model_name), foldername) if not os.path.exists(output_folder): os.makedirs(output_folder) output_name = "%s_%dto%d" % (args.eval_dataset, args.start_idx, args.end_idx) \ if args.start_idx != -1 and args.end_idx != -1 else args.eval_dataset output_name = output_name + ".jsonl" output_datapath = os.path.join(output_folder, output_name) print("writing to %s" % output_datapath) with open(output_datapath, "w", encoding='utf-8') as f: for output in output_list: if type(output) == dict: f.write(json.dumps(output) + "\n") else: f.write(output + "\n") if __name__ == "__main__": main()