Files
microsoft--unilm/PFPO/scripts/math/merge_dp_multi_solution.py
T
2026-07-13 13:24:13 +08:00

69 lines
2.1 KiB
Python

import argparse
import json
import os
import sys
from tqdm import tqdm
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
from post_processors.openai_api_callback import majority_voting_predict
from data.deepseek_math_utils import eval_script
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--input_file_format", type=str)
parser.add_argument("--seed_list", type=str)
parser.add_argument("--output_file", type=str)
args = parser.parse_args()
data = []
for seed in eval(args.seed_list):
_file = args.input_file_format.format(seed)
if os.path.exists(_file):
print(f"Loading: {_file}")
data.extend(json.load(open(_file)))
else:
print(f"File not found: {_file}")
id2data = {}
for item in tqdm(data):
if item["id"] not in id2data:
id2data[item["id"]] = item
id2data[item["id"]]["response"] = [id2data[item["id"]]["response"]]
id2data[item["id"]]["pred"] = [id2data[item["id"]]["pred"]]
id2data[item["id"]]["res"] = [id2data[item["id"]]["res"]]
else:
assert item["text"] == id2data[item["id"]]["text"]
id2data[item["id"]]["response"].append(item["response"])
id2data[item["id"]]["pred"].append(item["pred"])
id2data[item["id"]]["res"].append(item["res"])
data = list(id2data.values())
pass_at_k = 0
maj_at_k = 0
for item in tqdm(data):
sc_pred = majority_voting_predict(item["pred"])
if sc_pred != "":
sc_res = eval_script.eval_math({"prediction": sc_pred, "answer": item["label"]})
else:
sc_res = False
item["sc_pred"] = sc_pred
item["sc_res"] = sc_res
if any(item["res"]):
pass_at_k += 1
if sc_res:
maj_at_k += 1
print("Pass at k: ", pass_at_k / len(data))
print("Maj at k: ", maj_at_k / len(data))
json.dump(data, open(args.output_file, "w"), indent=2)
if __name__ == '__main__':
main()