wehub-resource-sync
55 lines
1.9 KiB
Python
55 lines
1.9 KiB
Python
import json
|
|
import sys
|
|
from argparse import ArgumentParser
|
|
|
|
from datasets import load_dataset
|
|
from tqdm import tqdm
|
|
|
|
sys.set_int_max_str_digits(0)
|
|
|
|
|
|
def main():
|
|
parser = ArgumentParser()
|
|
parser.add_argument("--split", type=str, default="train")
|
|
parser.add_argument("--prompt_file", type=str, default="prompts/apps/test_input_gen_2shot_v2.0.txt")
|
|
parser.add_argument("--output_file", type=str, required=True)
|
|
parser.add_argument("--function_only", default=False, action="store_true")
|
|
args = parser.parse_args()
|
|
|
|
data = load_dataset("codeparrot/apps", split="train").to_list()
|
|
print(len(data))
|
|
|
|
prompt_template = open(args.prompt_file).read()
|
|
|
|
with open(args.output_file, "w") as f:
|
|
for item in tqdm(data, desc="Writing prompts"):
|
|
question = item["question"]
|
|
if item["input_output"]:
|
|
input_output = json.loads(item["input_output"])
|
|
if "fn_name" in input_output:
|
|
question += f"\n\n{item['starter_code']}"
|
|
else:
|
|
input_output = {}
|
|
|
|
prompt = prompt_template.replace("[[Question]]", question)
|
|
item["prompt"] = prompt
|
|
if args.function_only:
|
|
if "fn_name" in input_output:
|
|
f.write(json.dumps(item) + "\n")
|
|
else:
|
|
f.write(json.dumps(item) + "\n")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|
|
|
|
"""
|
|
>>> python scripts/apps/pp_test_case_gen_inputs_v2.0.py \
|
|
--split train --prompt_file prompts/apps/test_input_gen_2shot_v2.0.txt \
|
|
--output outputs/apps/test_case_inputs_gen/apps.train.test_case_inputs.gen.v2.0.inputs.jsonl
|
|
|
|
>>> python scripts/apps/pp_test_case_gen_inputs_v2.0.py \
|
|
--split train --prompt_file prompts/apps/test_input_gen_2shot_v2.1.txt \
|
|
--output outputs/apps/test_case_inputs_gen/apps.train.test_case_inputs.gen.v2.1.func_only.inputs.jsonl --function_only
|
|
"""
|