wehub-resource-sync
69 lines
2.1 KiB
Python
69 lines
2.1 KiB
Python
import os.path
|
|
|
|
from datasets import load_dataset
|
|
import argparse
|
|
import json
|
|
from tqdm import tqdm
|
|
from glob import glob
|
|
|
|
import re
|
|
|
|
|
|
def extract_test_case_inputs(text):
|
|
# Regular expression to match content between <TEST CASE INPUTS> and </TEST CASE INPUTS>
|
|
pattern = r'<TEST CASE INPUTS>(.*?)</TEST CASE INPUTS>'
|
|
|
|
match = re.search(pattern, text, re.DOTALL)
|
|
if match:
|
|
return match.group(1).strip()
|
|
else:
|
|
return None
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--input_file", type=str)
|
|
parser.add_argument("--output_file", type=str)
|
|
args = parser.parse_args()
|
|
|
|
data = []
|
|
if os.path.exists(args.input_file):
|
|
data.extend([json.loads(line) for line in open(args.input_file).readlines()])
|
|
else:
|
|
for file in glob(args.input_file):
|
|
data.extend([json.loads(line) for line in open(file).readlines()])
|
|
|
|
outputs = {}
|
|
for item in data:
|
|
cases = extract_test_case_inputs(item["completion"])
|
|
if cases is not None:
|
|
cases = cases.split("\n")
|
|
cases = [case.strip() for case in cases if case.strip()]
|
|
cases = [case for case in cases if "==" not in case]
|
|
cases = [case for case in cases if "assert " in case]
|
|
|
|
if item["task_id"] not in outputs:
|
|
outputs[item["task_id"]] = {
|
|
"source_file": item["source_file"],
|
|
"task_id": item["task_id"],
|
|
"code": item["code"],
|
|
"test_imports": item["test_imports"],
|
|
"test_list": item["test_list"],
|
|
}
|
|
outputs[item["task_id"]]["aug_cases"] = set()
|
|
|
|
outputs[item["task_id"]]["aug_cases"].update(cases)
|
|
|
|
cnt = 0
|
|
for item_id in outputs:
|
|
outputs[item_id]["aug_cases"] = list(outputs[item_id]["aug_cases"])
|
|
cnt += len(outputs[item_id]["aug_cases"])
|
|
|
|
print(f"Averaged {cnt/len(outputs)} test cases per task.")
|
|
|
|
json.dump(outputs, open(args.output_file, "w"), indent=2)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|