Files
microsoft--unilm/PFPO/scripts/apps/execute_gold_sol_on_test_case.py
2026-07-13 13:24:13 +08:00

154 lines
5.0 KiB
Python

import copy
import json
import re
import sys
from argparse import ArgumentParser
from datasets import load_dataset
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed
from glob import glob
import os
from tqdm import tqdm
sys.set_int_max_str_digits(0)
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
from scripts.apps.utils_execute import check_correctness
def extract_solution_between_tags(text):
# Regular expression to match content between [BEGIN] and [END]
pattern = r'<BEGIN>(.*?)<END>'
match = re.search(pattern, text, re.DOTALL)
if match:
return match.group(1).strip()
else:
return None
def extract_test_cases(text):
# Regular expression to match the contents between <TEST INPUT X> and </TEST INPUT X>
pattern = r'<TEST INPUT (\d+)>(.*?)</TEST INPUT \1>'
# Find all matches in the text, the re.DOTALL flag allows . to match newline characters
test_cases = re.findall(pattern, text, re.DOTALL)
# Extract only the content part from the matches
test_cases = [case[1].strip() for case in test_cases]
return test_cases
def _worker(item):
solution = item["solutions"][0]
all_results = check_correctness(item["input_output"], solution, timeout=10, debug=False, return_output=True)
res, outputs, errors = all_results
# res = all_results
for tmp in res:
if (not isinstance(tmp, bool)) and (not isinstance(tmp, int)):
print(tmp, tmp.__class__.__name__)
res = [bool(tmp) if (not isinstance(tmp, bool)) and (not isinstance(tmp, int)) else tmp for tmp in res]
if all(item is True for item in res) is True:
result = True
else:
result = False
return result, res
def main():
parser = ArgumentParser()
parser.add_argument("--test_case_file", type=str)
parser.add_argument("--test_case_field", type=str)
parser.add_argument("--test_case_id_field", type=str, default="problem_id")
parser.add_argument("--output_file", type=str)
parser.add_argument("--num_workers", type=int, default=4)
args = parser.parse_args()
data = load_dataset("codeparrot/apps", split="train", trust_remote_code=True).to_list()
print(f"Total number of items before: {len(data)}")
train_sub_val_ids = set(json.load(open("apps_train_sub_val_ids.json")))
data = [item for item in data if item["problem_id"] not in train_sub_val_ids]
print(len(data))
test_cases = json.load(open(args.test_case_file))
# id2test_cases = {item[args.test_case_id_field]: item[args.test_case_field] for item in test_cases}
missing = 0
repeat = 0
id2test_cases = {}
for item in test_cases:
if item[args.test_case_id_field] in id2test_cases:
repeat += 1
if args.test_case_field not in item:
missing += 1
continue
id2test_cases[item[args.test_case_id_field]] = item[args.test_case_field]
print(f"Missing: {missing}")
print(f"Repeat: {repeat}")
missing_solutions = 0
missing_test_cases = 0
new_data = []
for item in data:
if item["solutions"]:
item["solutions"] = json.loads(item["solutions"])
assert isinstance(item["solutions"], list)
if len(item["solutions"]) > 0:
if item["problem_id"] in id2test_cases:
item["input_output"] = id2test_cases[item["problem_id"]]
new_data.append(item)
elif f"apps-train-{item['problem_id']}" in id2test_cases:
item["input_output"] = id2test_cases[f"apps-train-{item['problem_id']}"]
new_data.append(item)
else:
missing_test_cases += 1
else:
missing_solutions += 1
else:
missing_solutions += 1
print(f"Missing solutions: {missing_solutions}")
print(f"Missing test cases: {missing_test_cases}")
data = new_data
print(f"Total number of items: {len(data)}")
missing = 0
corr = 0
micro_corr = 0
micro_all = 0
pbar = tqdm(data)
outputs = []
with ThreadPoolExecutor(max_workers=args.num_workers) as executor:
futures = []
for _input in pbar:
future = executor.submit(_worker, _input)
futures.append(future)
pbar.update()
for future in tqdm(as_completed(futures), total=len(futures), desc="Collecting results"):
outputs.append(future.result())
for item in outputs:
result, full_res = item
if result:
corr += 1
for r in full_res:
if r:
micro_corr += 1
micro_all += 1
print(f"Missing: {missing} / {len(outputs)} = {missing / len(outputs)}")
print(f"Correct: {corr} / {len(outputs)} = {corr / len(outputs)}")
print(f"Micro Correct: {micro_corr} / {micro_all} = {micro_corr / micro_all}")
json.dump(outputs, open(args.output_file, "w"), ensure_ascii=False, indent=2)
if __name__ == '__main__':
main()