File size: 1,631 Bytes
24c2665 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 |
# Convert a directory of solutions to a JSONL file
import json
import os
from tqdm import tqdm
from evalplus.data import get_human_eval_plus, get_mbpp_plus
from evalplus.sanitize import sanitize
def main(directory, sanitize_code: bool = True):
basename = os.path.basename(directory)
parent_dir = os.path.dirname(directory)
target_jsonl_path = os.path.join(parent_dir, f"{basename}.jsonl")
datasets = {**get_human_eval_plus(), **get_mbpp_plus()}
with open(target_jsonl_path, "w") as f:
# iterate directories
for subdir_name in tqdm(os.listdir(directory)):
subdir_path = os.path.join(directory, subdir_name)
if not os.path.isdir(subdir_path):
continue
assert "_" in subdir_name
dataset_name, task_num = subdir_name.split("_")
task_id = f"{dataset_name}/{task_num}"
entrypoint = datasets[task_id]["entry_point"]
for sample_name in os.listdir(subdir_path):
if not sample_name.endswith(".py"):
continue
code = open(os.path.join(subdir_path, sample_name)).read()
if sanitize_code:
try:
code = sanitize(code, entrypoint=entrypoint)
except ValueError as e:
print(f"Failed to sanitize {task_id}/{sample_name}: {e}")
print(code)
continue
f.write(json.dumps({"task_id": task_id, "solution": code}) + "\n")
if __name__ == "__main__":
from fire import Fire
Fire(main)
|