Files
microsoft--unilm/PFPO/scripts/apps/pp_test_case_gen_public_outputs.py
2026-07-13 13:24:13 +08:00

117 lines
2.9 KiB
Python

import copy
import json
import re
import sys
from argparse import ArgumentParser
import json
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:
{}
------------------------
And here is one group of test case inputs:
```
{}
```
------------------------
Please carefully comprehend the process described in the problem, and simulate it **step-by-step** to obtain the expected outputs the above test case inputs.
**REMEMBER** to follow the format below to provide the outputs:
[BEGIN]
outputs
[END]
"""
_vanilla_prompt_2 = """You are an expert programmer. Here is a programming problem:
{}
------------------------
And here is one group of test case inputs:
```
{}
```
------------------------
Please carefully comprehend the process described in the problem, and derive the expected outputs the above test case inputs. You should follow the steps below:
(1) Simulate the algorithm/process described in the problem with the given inputs **step-by-step**, and record the step-level intermediate results.
(2) Derive the correct outcome.
(3) Put the final outputs in the format below:
[BEGIN]
*put your outputs here*
[END]
and return back it to me.
"""
_vanilla_prompt_3 = """You are an expert programmer. Here is a programming problem:
{}
------------------------
And here is one group of test case inputs:
```
{}
```
------------------------
Please carefully comprehend the process described in the problem, and derive the expected outputs the above test case inputs. You should follow the steps below:
(1) Simulate the algorithm/process described in the problem with the given inputs **step-by-step**, and record the step-level intermediate results.
(2) Derive the correct outcome.
Response format:
[BEGIN]
*The expected outputs*
[END]
"""
PROMPTS = {
"vanilla": _vanilla_prompt,
"vanilla_2": _vanilla_prompt_2,
"vanilla_3": _vanilla_prompt_3,
}
def main():
parser = ArgumentParser()
parser.add_argument("--split", type=str, required=True)
parser.add_argument("--output_file", type=str, required=True)
parser.add_argument("--prompt_type", type=str, default="vanilla")
args = parser.parse_args()
data = load_dataset("codeparrot/apps", split=args.split).to_list()
print(len(data))
inputs = []
for item in data:
if item["input_output"]:
item["input_output"] = json.loads(item["input_output"])
inputs.append(item)
with open(args.output_file, "w") as f:
for item in tqdm(inputs, total=len(inputs)):
problem = item["question"]
for case in item["input_output"]["inputs"]:
new_item = copy.deepcopy(item)
prompt = PROMPTS[args.prompt_type].format(problem, case)
new_item["prompt"] = prompt
f.write(json.dumps(new_item) + "\n")
if __name__ == '__main__':
main()
"""
"""