in evaluation/math_eval.py [0:0]
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",
]:
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()
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]
if args.use_vllm:
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}"
)
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