wehub-resource-sync
66 lines
1.8 KiB
Python
66 lines
1.8 KiB
Python
from datasets import load_dataset
|
|
import argparse
|
|
import json
|
|
from tqdm import tqdm
|
|
|
|
|
|
_prompt_w_test_cases = """You are an expert Python programmer, and here is your task:
|
|
|
|
{}
|
|
|
|
And here are some test cases in assertion format:
|
|
|
|
<EXAMPLE TEST CASE INPUT>
|
|
{}
|
|
</EXAMPLE TEST CASE INPUT>
|
|
|
|
where only the function name and test cases inputs are included, and the expected outputs are omitted to avoid any confusion.
|
|
|
|
Now, please **STRICTLY** following the above assertion format to generate **50** more test case inputs for me. Organize your results between <TEST CASE INPUTS> and </TEST CASE INPUTS> tags:
|
|
|
|
Here is the return format:
|
|
|
|
<TEST CASE INPUTS>
|
|
*assertion inputs 1*
|
|
*assertion inputs 2*
|
|
...
|
|
*assertion inputs 50*
|
|
</TEST CASE INPUTS>
|
|
|
|
Remember my requirements and now let's get started:
|
|
|
|
"""
|
|
|
|
|
|
PROMPTS = {
|
|
"w_test_cases": _prompt_w_test_cases,
|
|
}
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--sanitized", default=False, action="store_true")
|
|
parser.add_argument("--output_file", type=str, required=True)
|
|
parser.add_argument("--prompt_type", type=str, default="vanilla")
|
|
args = parser.parse_args()
|
|
|
|
if args.sanitized:
|
|
dataset = load_dataset("mbpp", "sanitized", split="test").to_list()
|
|
else:
|
|
dataset = load_dataset("mbpp", split="test").to_list()
|
|
prompt_key = "prompt" if args.sanitized else "text"
|
|
|
|
with open(args.output_file, "w") as f:
|
|
for item in tqdm(dataset):
|
|
query = item[prompt_key]
|
|
test_cases = "\n".join([case.split(" == ")[0] for case in item["test_list"]])
|
|
assert len(item["test_list"]), item
|
|
prompt = PROMPTS[args.prompt_type].format(query, test_cases)
|
|
|
|
item["prompt"] = prompt
|
|
f.write(json.dumps(item) + "\n")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|