Files
2026-07-13 13:06:23 +08:00

121 lines
3.5 KiB
Python

import torch
from datasets import load_dataset
from tqdm import tqdm
import multiprocessing
import random
import requests
from functools import partial
import argparse
from pathlib import Path
import yaml
from metachain.agents.math.math_solver_agent import get_math_solver_agent
from metachain import MetaChain
from metachain.workflows.math_solver_workflow_flow import majority_voting
import importlib
import os
import asyncio
from evaluation.math500.prompts import MATH_COT_PROMPT
def save_yaml(path: Path, data, sort_keys=True):
with open(path, "w") as f:
yaml.dump(data, f, sort_keys=sort_keys)
async def run_inference(item, save_dir, workflow):
outpath = save_dir / f"{item['id']}.yaml"
if outpath.exists():
return
prompt = MATH_COT_PROMPT + f"\n\nProblem:\n{item['problem']}\n\nYour task is to solve this problem."
prompt += "Please given your final answer (answer ONLY) within the format of `Final Answer: The final answer is <answer>. I hope it is correct.` after your reasoning \n"
prompt += "For example: According to ...\nFinal Answer: The final answer is $24$. I hope it is correct.\n"
if workflow == "majority_voting":
answer = await majority_voting(prompt)
elif workflow == None:
agent = get_math_solver_agent(model="deepseek/deepseek-chat")
client = MetaChain()
messages = [
{"role": "user", "content": prompt},
]
context_variables = {
}
response = await client.run_async(agent, messages, context_variables)
answer = response.messages[-1]['content']
else: raise ValueError(f"Unknown workflow: {workflow}")
out = {
"prompt": prompt,
"question": item["problem"],
"answer": answer,
"gt_answer": item["answer"],
}
save_yaml(outpath, out)
async def main(args):
test_dataset = list(
load_dataset(
"HuggingFaceH4/MATH-500", "default", split="test", trust_remote_code=True
)
)
print(f"Number of test items: {len(test_dataset)}")
random.seed(12345)
for i, data in enumerate(test_dataset):
data["id"] = i
random.shuffle(test_dataset)
if args.limit is not None:
limit = args.limit
else:
limit = len(test_dataset)
if args.stride is not None:
stride = args.stride
else:
stride = 1
if args.offset is not None:
offset = args.offset
else:
offset = 0
test_dataset = test_dataset[offset:limit:stride]
print(f"Total number of items to process: {len(test_dataset)}")
if args.workflow == None:
save_dir = os.path.join(args.save_dir, "math_solver")
save_dir = Path(save_dir)
save_dir.mkdir(parents=True, exist_ok=True)
else:
save_dir = os.path.join(args.save_dir, args.workflow)
save_dir = Path(save_dir)
save_dir.mkdir(parents=True, exist_ok=True)
predictions = []
for item in tqdm(test_dataset):
predictions.append(await run_inference(item, save_dir, args.workflow))
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--num_few_shot", type=int, default=2)
parser.add_argument("--limit", type=int, default=3)
parser.add_argument("--stride", type=int, default=1)
parser.add_argument("--offset", type=int, default=0)
parser.add_argument("--save_dir", type=str, default="evaluation_results/math500")
parser.add_argument("--workflow", type=str, default=None)
args = parser.parse_args()
asyncio.run(main(args))