wehub-resource-sync
134 lines
4.5 KiB
Python
134 lines
4.5 KiB
Python
import copy
|
|
import json
|
|
import re
|
|
import sys
|
|
from argparse import ArgumentParser
|
|
from datasets import load_dataset
|
|
from collections import defaultdict
|
|
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 run_test
|
|
|
|
|
|
def extract_content_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 main():
|
|
parser = ArgumentParser()
|
|
parser.add_argument("--split", type=str, default="train")
|
|
parser.add_argument("--test_case_inputs", type=str, required=True, help="The completion file.")
|
|
parser.add_argument("--test_case_outputs", type=str, required=True, help="The completion file.")
|
|
# parser.add_argument("--output_file", type=str, required=True)
|
|
args = parser.parse_args()
|
|
|
|
data = load_dataset("codeparrot/apps", split="train").to_list()
|
|
print(len(data))
|
|
# for item in data: # Currently, we do not require the solutions and input_output fields
|
|
# if item["solutions"]:
|
|
# item["solutions"] = json.loads(item["solutions"])
|
|
# if item["input_output"]:
|
|
# item["input_output"] = json.loads(item["input_output"])
|
|
|
|
case_inputs = [json.loads(line) for line in open(args.test_case_inputs).readlines()]
|
|
case_outputs = [json.loads(line) for line in open(args.test_case_outputs).readlines()]
|
|
|
|
problem_id2case = defaultdict(list)
|
|
for item in case_inputs:
|
|
cases = extract_test_cases(item["completion"])
|
|
if len(cases):
|
|
for i, case in enumerate(cases):
|
|
problem_id2case[item["problem_id"]].append({
|
|
"input": case,
|
|
"case_id": i,
|
|
})
|
|
|
|
for item in case_outputs:
|
|
output = extract_content_between_tags(item["completion"])
|
|
if output:
|
|
assert isinstance(output, str), output
|
|
case_id = int(item["case_id"].split("_")[1])
|
|
assert case_id == problem_id2case[item["problem_id"]][case_id]["case_id"]
|
|
problem_id2case[item["problem_id"]][case_id]["output"] = output
|
|
|
|
# Currently, we can directly verify
|
|
outputs = []
|
|
for item in data:
|
|
if item["solutions"]:
|
|
item["solutions"] = json.loads(item["solutions"])
|
|
|
|
problem_id = item["problem_id"]
|
|
cases = {
|
|
"inputs": [],
|
|
"outputs": [],
|
|
}
|
|
for case in problem_id2case[problem_id]:
|
|
if "output" in case:
|
|
cases["inputs"].append(case["input"])
|
|
cases["outputs"].append(case["output"])
|
|
assert isinstance(case["input"], str), case["input"]
|
|
assert isinstance(case["output"], str), case["output"]
|
|
if cases["inputs"]:
|
|
item["gen_test_cases"] = cases
|
|
outputs.append(item)
|
|
|
|
# print(outputs[0]["gen_test_cases"])
|
|
|
|
cnt = 0
|
|
num_program = 0
|
|
for i, item in tqdm(enumerate(outputs)):
|
|
if i in [7]:
|
|
continue # Some bad cases.
|
|
print(item["gen_test_cases"])
|
|
for solution in item["solutions"]:
|
|
try:
|
|
res = run_test(item["gen_test_cases"], solution)
|
|
pass_num = sum([1 for item in res if item is True])
|
|
cnt += pass_num
|
|
except:
|
|
pass
|
|
num_program += 1
|
|
|
|
print(f"Pass rate: {cnt / num_program}")
|
|
|
|
# print(json.dumps(data[0], indent=2))
|
|
# print(data[0]["question"])
|
|
# print(data[0]["solutions"][0])
|
|
# print(data[0]["input_output"]["inputs"][0])
|
|
# print(data[0]["input_output"]["outputs"][0])
|
|
|
|
# with open(args.output_file, "w") as f:
|
|
# for item in tqdm(outputs, desc="Writing prompts"):
|
|
# f.write(json.dumps(item) + "\n")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|