wehub-resource-sync
142 lines
3.8 KiB
Python
142 lines
3.8 KiB
Python
from datasets import load_dataset
|
|
import argparse
|
|
import json
|
|
|
|
_vanilla_prompt = """You are an expert Python programmer, and here is your task: {}
|
|
Your code should pass these tests:
|
|
|
|
{}
|
|
|
|
Please put your code into [BEGIN]<Your Code>[END] block.
|
|
"""
|
|
|
|
_vanilla_prompt_emp = """You are an expert Python programmer, and here is your task: {}
|
|
Your code should pass these tests:
|
|
|
|
{}
|
|
|
|
**REMEMBER** to tag your solution program with the following format:
|
|
[BEGIN]
|
|
<Your Program>
|
|
[END]
|
|
and then return back it to me.
|
|
"""
|
|
|
|
_description_prompt = """You are an expert Python programmer, and here is your task: {}
|
|
Your code should pass these tests:
|
|
|
|
{}
|
|
|
|
**DO NOT** show me any code. Please use **NATURAL LANGUAGE** to describe your thoughts and implementation in details. Take care if your algorithm can pass the test cases by simulation.
|
|
|
|
**REMEMBER** to tag your description with the following format:
|
|
[BEGIN]
|
|
<Your Description>
|
|
[END]
|
|
and then return back it to me.
|
|
"""
|
|
|
|
_desc2code_no_task_prompt = """You are an expert Python programmer, and here is a description of an algorithm:
|
|
{}
|
|
|
|
Your task is to implement the algorithm in Python. Your code should pass these test cases:
|
|
|
|
{}
|
|
|
|
**REMEMBER** to tag your solution program with the following format:
|
|
[BEGIN]
|
|
<Your Program>
|
|
[END]
|
|
and then return back it to me.
|
|
"""
|
|
|
|
_desc2code_prompt = """You are an expert Python programmer, and here is your task: {}
|
|
|
|
I have helped you to describe how to implement the algorithm in Python:
|
|
{}
|
|
|
|
Your task is to implement the algorithm in Python. Your code should pass these test cases:
|
|
|
|
{}
|
|
|
|
**REMEMBER** to tag your solution program with the following format:
|
|
|
|
[BEGIN]
|
|
<Your Program>
|
|
[END]
|
|
|
|
and then return back it to me.
|
|
"""
|
|
|
|
|
|
PROMPTS = {
|
|
"vanilla": _vanilla_prompt,
|
|
"vanilla_emp": _vanilla_prompt_emp,
|
|
"description": _description_prompt,
|
|
"desc2code_no_task": _desc2code_no_task_prompt,
|
|
"desc2code": _desc2code_prompt,
|
|
}
|
|
|
|
|
|
def load_descriptions(desc_file):
|
|
raw_completions = [json.loads(line) for line in open(desc_file).readlines()]
|
|
|
|
id2desc = {}
|
|
missing = 0
|
|
for item in raw_completions:
|
|
if "[BEGIN]" not in item["completion"] or "[END]" not in item["completion"]:
|
|
missing += 1
|
|
continue
|
|
|
|
s = item["completion"].index("[BEGIN]") + len("[BEGIN]")
|
|
e = item["completion"].index("[END]")
|
|
desc = item["completion"][s:e].strip()
|
|
id2desc[item["task_id"]] = desc
|
|
print(f"Missing {missing} descriptions.")
|
|
return id2desc
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--desc_file", type=str, required=True)
|
|
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="desc2code_no_task")
|
|
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"
|
|
|
|
id2desc = load_descriptions(args.desc_file)
|
|
|
|
outputs = []
|
|
for item in dataset:
|
|
task_id = item["task_id"]
|
|
query = item[prompt_key]
|
|
test_cases = "\n".join(item["test_list"])
|
|
|
|
if task_id in id2desc:
|
|
desc = id2desc[task_id]
|
|
if args.prompt_type == "desc2code":
|
|
prompt = PROMPTS[args.prompt_type].format(query, desc, test_cases)
|
|
else:
|
|
prompt = PROMPTS[args.prompt_type].format(desc, test_cases)
|
|
else:
|
|
prompt = PROMPTS["_vanilla_prompt_emp"].format(query, test_cases)
|
|
|
|
outputs.append({
|
|
"prompt": prompt,
|
|
"task_id": task_id,
|
|
})
|
|
|
|
with open(args.output_file, "w") as f:
|
|
for item in outputs:
|
|
f.write(json.dumps(item) + "\n")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|