alKoGolik's picture
Upload 169 files
c87c295 verified
if __name__ == "__main__":
import argparse
import json
import os
from tempdir import TempDir
parser = argparse.ArgumentParser()
parser.add_argument("--dataset", default="humaneval", type=str)
parser.add_argument("--plus-input", required=True, type=str)
parser.add_argument("--output", required=True, type=str)
args = parser.parse_args()
assert args.dataset == "humaneval"
assert not os.path.exists(args.output), f"{args.output} already exists!"
with TempDir() as tempdir:
# Generate inputs
plus_input = {}
with open(args.plus_input) as file:
for line in file:
problem = json.loads(line)
plus_input[problem["task_id"]] = problem["inputs"]
tempf = None
if args.dataset == "humaneval":
from evalplus.data import get_human_eval_plus
# Allow it to be incomplete
problems = get_human_eval_plus(err_incomplete=False)
tempf = os.path.join(tempdir, "HumanEvalPlus.jsonl")
with open(tempf, "w") as file:
for problem in problems:
problem["plus_input"] = plus_input[problem["task_id"]]
file.write(json.dumps(problem) + "\n")
# Move to the right place
os.rename(tempf, args.output)