diff options
Diffstat (limited to 'Qwen2.5-Eval/evaluation/math_eval.py')
| -rwxr-xr-x | Qwen2.5-Eval/evaluation/math_eval.py | 437 |
1 files changed, 437 insertions, 0 deletions
diff --git a/Qwen2.5-Eval/evaluation/math_eval.py b/Qwen2.5-Eval/evaluation/math_eval.py new file mode 100755 index 0000000..0a35929 --- /dev/null +++ b/Qwen2.5-Eval/evaluation/math_eval.py @@ -0,0 +1,437 @@ +# rm -rf math_eval.py; vim math_eval.py +import random +import os +import argparse +import time +from vllm import LLM, SamplingParams +from datetime import datetime +from tqdm import tqdm + +import torch +from transformers import AutoTokenizer, AutoModelForCausalLM + +from evaluate import evaluate +from utils import set_seed, load_jsonl, save_jsonl, construct_prompt +from parser import * +from trajectory import * +from data_loader import load_data +from python_executor import PythonExecutor +from model_utils import load_hf_lm_and_tokenizer, generate_completions + + +def parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument("--data_names", default="gsm8k,math", type=str) + parser.add_argument("--data_dir", default="./data", type=str) + parser.add_argument("--model_name_or_path", default="gpt-4", type=str) + parser.add_argument("--output_dir", default="./output", type=str) + parser.add_argument("--prompt_type", default="tool-integrated", type=str) + parser.add_argument("--split", default="test", type=str) + parser.add_argument("--num_test_sample", default=-1, type=int) # -1 for full data + parser.add_argument("--seed", default=0, type=int) + parser.add_argument("--start", default=0, type=int) + parser.add_argument("--end", default=-1, type=int) + parser.add_argument("--temperature", default=0, type=float) + parser.add_argument("--n_sampling", default=1, type=int) + parser.add_argument("--top_p", default=1, type=float) + parser.add_argument("--max_tokens_per_call", default=8192, type=int) + parser.add_argument("--shuffle", action="store_true") + parser.add_argument("--use_vllm", action="store_true") + parser.add_argument("--vllm_batch_size",default=0, type=int) + parser.add_argument("--save_outputs", action="store_true") + parser.add_argument("--overwrite", action="store_true") + parser.add_argument("--use_safetensors", action="store_true") + parser.add_argument("--num_shots", type=int, default=0) + parser.add_argument( + "--apply_chat_template", + action="store_true", + help="Apply chat template to prompt.", + ) + parser.add_argument("--pipeline_parallel_size", type=int, default=1) + parser.add_argument( + "--adapt_few_shot", + action="store_true", + help="Few shot for multiple-choice questions, zero shot for others.", + ) + args = parser.parse_args() + args.top_p = ( + 1 if args.temperature == 0 else args.top_p + ) # top_p must be 1 when using greedy sampling (vllm) + return args + + +def prepare_data(data_name, args): + examples = load_data(data_name, args.split, args.data_dir) + + # sample `num_test_sample` from dataset + if args.num_test_sample > 0: + # examples = random.sample(examples, min(args.num_test_sample, len(examples))) + examples = examples[: args.num_test_sample] + + # shuffle + if args.shuffle: + random.seed(datetime.now().timestamp()) + random.shuffle(examples) + + # select start and end + examples = examples[args.start : len(examples) if args.end == -1 else args.end] + + # get out_file name + dt_string = datetime.now().strftime("%m-%d_%H-%M") + model_name = "/".join(args.model_name_or_path.split("/")[-2:]) + out_file_prefix = f"{args.split}_{args.prompt_type}_{args.num_test_sample}_seed{args.seed}_t{args.temperature}" + output_dir = args.output_dir + # if not os.path.exists(output_dir): + # output_dir = f"outputs/{output_dir}" # don't add this + out_file = f"{output_dir}/{data_name}/{out_file_prefix}_s{args.start}_e{args.end}.jsonl" + os.makedirs(f"{output_dir}/{data_name}", exist_ok=True) + + # load all processed samples + processed_samples = [] + if not args.overwrite: + processed_files = [ + f + for f in os.listdir(f"{output_dir}/{data_name}/") + if f.endswith(".jsonl") and f.startswith(out_file_prefix) + ] + for f in processed_files: + processed_samples.extend( + list(load_jsonl(f"{output_dir}/{data_name}/{f}")) + ) + + # dedepulicate + processed_samples = {sample["idx"]: sample for sample in processed_samples} + processed_idxs = list(processed_samples.keys()) + processed_samples = list(processed_samples.values()) + examples = [example for example in examples if example["idx"] not in processed_idxs] + return examples, processed_samples, out_file + + +def setup(args): + # load model + available_gpus = os.environ["CUDA_VISIBLE_DEVICES"].split(",") + if args.use_vllm: + llm = LLM( + # model="/volume/ailab4sci/ztgao/em/checkpoints/step_20", + model=args.model_name_or_path, + tensor_parallel_size=4, + pipeline_parallel_size=args.pipeline_parallel_size, + trust_remote_code=True, + ) + tokenizer = None + if args.apply_chat_template: + tokenizer = AutoTokenizer.from_pretrained( + args.model_name_or_path, trust_remote_code=True + ) + else: + llm, tokenizer = load_hf_lm_and_tokenizer( + model_name_or_path=args.model_name_or_path, + load_in_half=True, + use_fast_tokenizer=True, + use_safetensors=args.use_safetensors, + ) + + # infer & eval + data_list = args.data_names.split(",") + results = [] + for data_name in data_list: + results.append(main(llm, tokenizer, data_name, args)) + + # add "avg" result to data_list and results + data_list.append("avg") + results.append( + { + "acc": sum([result["acc"] for result in results]) / len(results), + } + ) + + # print all results + pad = max([len(data_name) for data_name in data_list]) + print("\t".join(data_name.ljust(pad, " ") for data_name in data_list)) + print("\t".join([f"{result['acc']:.1f}".ljust(pad, " ") for result in results])) + + +def is_multi_choice(answer): + for c in answer: + if c not in ["A", "B", "C", "D", "E"]: + return False + return True + + +def main(llm, tokenizer, data_name, args): + examples, processed_samples, out_file = prepare_data(data_name, args) + print("=" * 50) + print("data:", data_name, " ,remain samples:", len(examples)) + if len(examples) > 0: + print(examples[0]) + + # init python executor + if "pal" in args.prompt_type: + executor = PythonExecutor(get_answer_expr="solution()") + else: + executor = PythonExecutor(get_answer_from_stdout=True) + + samples = [] + for example in tqdm(examples, total=len(examples)): + idx = example["idx"] + + # parse question and answer + example["question"] = parse_question(example, data_name) + if example["question"] == "": + continue + gt_cot, gt_ans = parse_ground_truth(example, data_name) + example["gt_ans"] = gt_ans + full_prompt = construct_prompt(example, data_name, args) + + if idx == args.start: + print(full_prompt) + + sample = { + "idx": idx, + "question": example["question"], + "gt_cot": gt_cot, + "gt": gt_ans, + "prompt": full_prompt, + } + + # add remain fields + for key in [ + "level", + "type", + "unit", + "solution_type", + "choices", + "solution", + "ques_type", + "ans_type", + "answer_type", + "dataset", + "subfield", + "filed", + "theorem", + "answer", + "difficulty" + ]: + if key in example: + sample[key] = example[key] + samples.append(sample) + + # repeat n times + input_prompts = [ + sample["prompt"] for sample in samples for _ in range(args.n_sampling) + ] + if args.apply_chat_template: + input_prompts = [ + tokenizer.apply_chat_template( + [{"role": "user", "content": prompt.strip()}], + tokenize=False, + add_generation_prompt=True, + ) + for prompt in input_prompts + ] + remain_prompts = input_prompts + remain_prompts = [(i, prompt) for i, prompt in enumerate(remain_prompts)] + end_prompts = [] + + max_func_call = 1 if args.prompt_type in ["cot", "pal"] else 4 + + stop_words = ["</s>", "<|im_end|>", "<|endoftext|>"] + + if args.prompt_type in ["cot"]: + stop_words.append("\n\nQuestion:") + if args.prompt_type in ["pal", "tool-integrated", "jiuzhang_tora"]: + stop_words.extend(["\n\n---", "```output"]) + elif args.prompt_type in ["wizard_zs", "platypus_fs"]: + stop_words.extend(["Instruction", "Response"]) + elif "jiuzhang" in args.prompt_type: + stop_words.append("\n\n## Question") + elif "numina" in args.prompt_type: + stop_words.append("\n### Problem") + elif "pure" in args.prompt_type: + stop_words.append("\n\n\n") + + # start inference + # measure time use + start_time = time.time() + result_prompts = [] + for epoch in range(max_func_call): + print("-" * 20, "Epoch", epoch) + current_prompts = remain_prompts + if len(current_prompts) == 0: + break + + # get all outputs + prompts = [item[1] for item in current_prompts] + result_prompts.extend(prompts) + if args.use_vllm: + if args.vllm_batch_size: + outputs = [] + for i in range(0, len(prompts), args.vllm_batch_size): + batch_prompts = prompts[i:i+args.vllm_batch_size] + print(f"args.temperature: {args.temperature}") + batch_outputs = llm.generate( + batch_prompts, + SamplingParams( + temperature=args.temperature, + max_tokens=args.max_tokens_per_call, + stop=stop_words, + stop_token_ids=( + [151645, 151643] + if "qwen2" in args.model_name_or_path.lower() + else None + ), + ), + ) + + batch_outputs = sorted( + batch_outputs, key=lambda x: int(x.request_id) + ) + batch_outputs = [output.outputs[0].text for output in batch_outputs] + outputs.extend(batch_outputs) + else: + outputs = llm.generate( + prompts, + SamplingParams( + temperature=args.temperature, + # top_p=args.top_p, + max_tokens=args.max_tokens_per_call, + # n=1, + stop=stop_words, + stop_token_ids=( + [151645, 151643] + if "qwen2" in args.model_name_or_path.lower() + else None + ), + ), + ) + + outputs = sorted( + outputs, key=lambda x: int(x.request_id) + ) # sort outputs by request_id + outputs = [output.outputs[0].text for output in outputs] + else: + outputs = generate_completions( + model=llm, + tokenizer=tokenizer, + prompts=prompts, + max_new_tokens=args.max_tokens_per_call, + batch_size=16, + stop_id_sequences=stop_words, + ) + + assert len(outputs) == len(current_prompts) + + # process all outputs + remain_prompts = [] + remain_codes = [] + for (i, query), output in zip(current_prompts, outputs): + output = output.rstrip() + query += output + if args.prompt_type == "pal": + remain_prompts.append((i, query)) + if "```python" in output: + output = extract_program(query) + remain_codes.append(output) + elif args.prompt_type == "cot": + end_prompts.append((i, query)) + elif "boxed" not in output and output.endswith("```"): + program = extract_program(query) + remain_prompts.append((i, query)) + remain_codes.append(program) + else: + end_prompts.append((i, query)) + + # execute the remain prompts + remain_results = executor.batch_apply(remain_codes) + for k in range(len(remain_prompts)): + i, query = remain_prompts[k] + res, report = remain_results[k] + exec_result = res if res else report + if "pal" in args.prompt_type: + exec_result = "\\boxed{" + exec_result + "}" + exec_result = f"\n```output\n{exec_result}\n```\n" + query += exec_result + # not end + if epoch == max_func_call - 1: + query += "\nReach max function call limit." + remain_prompts[k] = (i, query) + + # unsolved samples + print("Unsolved samples:", len(remain_prompts)) + end_prompts.extend(remain_prompts) + # sort by idx + end_prompts = sorted(end_prompts, key=lambda x: x[0]) + + # remove input_prompt from end_prompt + codes = [] + assert len(input_prompts) == len(end_prompts) + for i in range(len(input_prompts)): + _, end_prompt = end_prompts[i] + code = end_prompt.split(input_prompts[i])[-1].strip() + for stop_word in stop_words: + if stop_word in code: + code = code.split(stop_word)[0].strip() + codes.append(code) + + # extract preds + results = [ + run_execute(executor, code, args.prompt_type, data_name) for code in codes + ] + time_use = time.time() - start_time + + # put results back to examples + all_samples = [] + for i, sample in enumerate(samples): + code = codes[i * args.n_sampling : (i + 1) * args.n_sampling] + result = results[i * args.n_sampling : (i + 1) * args.n_sampling] + preds = [item[0] for item in result] + reports = [item[1] for item in result] + for j in range(len(preds)): + if sample["gt"] in ["A", "B", "C", "D", "E"] and preds[j] not in [ + "A", + "B", + "C", + "D", + "E", + ]: + preds[j] = choice_answer_clean(code[j]) + elif is_multi_choice(sample["gt"]) and not is_multi_choice(preds[j]): + # remove any non-choice char + preds[j] = "".join( + [c for c in preds[j] if c in ["A", "B", "C", "D", "E"]] + ) + + sample.pop("prompt") + sample.update({"code": code, "pred": preds, "report": reports}) + all_samples.append(sample) + + # add processed samples + all_samples.extend(processed_samples) + all_samples, result_json = evaluate( + samples=all_samples, + data_name=data_name, + prompt_type=args.prompt_type, + execute=True, + ) + + # save outputs + if len(processed_samples) < len(all_samples) and args.save_outputs: + save_jsonl(all_samples, out_file) + + result_json["time_use_in_second"] = time_use + result_json["time_use_in_minite"] = ( + f"{int(time_use // 60)}:{int(time_use % 60):02d}" + ) + result_json["prompts"] = result_prompts + + with open( + out_file.replace(".jsonl", f"_{args.prompt_type}_metrics.json"), "w" + ) as f: + json.dump(result_json, f, indent=4) + return result_json + + +if __name__ == "__main__": + args = parse_args() + set_seed(args.seed) + setup(args) |