wehub-resource-sync
55 lines
1.6 KiB
Python
55 lines
1.6 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)
|
|
|
|
_vanilla_prompt = """You are an expert programmer. Here is a programming problem:
|
|
--------
|
|
{}
|
|
--------
|
|
Please carefully comprehend the requirements in the problem, and write down the solution program to pass it under the given time and memory constraints.
|
|
|
|
**REMEMBER** to strictly follow the steps below to help reduce the potential flaws:
|
|
(1) According to the input scale and the time/memory constraints, think about the time complexity and space complexity of your solution.
|
|
(2) Think **step-by-step** to design the algorithm.
|
|
(3) Translate your thoughts into Python program to solve it.
|
|
|
|
Besides, your Python solution program should be located between <BEGIN> and <END> tags. For example:
|
|
<BEGIN>
|
|
t = int(input())
|
|
...
|
|
print(ans)
|
|
<END>
|
|
|
|
Now, remember my instruction and let's get started:
|
|
"""
|
|
|
|
PROMPTS = {
|
|
"vanilla": _vanilla_prompt,
|
|
}
|
|
|
|
|
|
def main():
|
|
parser = ArgumentParser()
|
|
parser.add_argument("--split", type=str, default="train")
|
|
parser.add_argument("--prompt_type", type=str, default="vanilla")
|
|
parser.add_argument("--output_file", type=str, required=True)
|
|
args = parser.parse_args()
|
|
|
|
data = load_dataset("codeparrot/apps", split="train").to_list()
|
|
print(len(data))
|
|
|
|
with open(args.output_file, "w") as f:
|
|
for item in tqdm(data, desc="Writing prompts"):
|
|
prompt = PROMPTS[args.prompt_type].format(item["question"])
|
|
item["prompt"] = prompt
|
|
f.write(json.dumps(item) + "\n")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|