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

99 lines
3.1 KiB
Python

import argparse
import copy
import json
from datasets import load_dataset
_print_inter_template = """I have a program with the instruction or comments to indicate the usage of it. I want to know the intermediate running process of the program, so please help me:
(1) Add several `print` clauses after each line of code to print the values of the involved variables. Please follow the format: `print(f"x: ", x), where `x` is one of the intermediate values. Each print clause is used for checking exactly one variable. Return the completed program using <NEW PROGRAM></NEW PROGRAM>.
(2) Give me 5 test cases according to the docstring or instruction of the program. Use <TEST CASE></TEST CASE> to return them to me. Each pair of <TEST CASE></TEST CASE> should include only one case. The cases should also in python code format so that I can use `eval(test_case)` to obtain the results directly.
I will first give you an example:
Instruction: Write a function to sort a given matrix in ascending order according to the sum of its rows.
Program:
def sort_matrix(matrix):
matrix.sort(key=sum)
return matrix
Here's the program with added `print` statements to track the intermediate values:
<NEW PROGRAM>
def sort_matrix(matrix):
# Print the original matrix
print("Original matrix:", matrix)
# Print the sum of each row before sorting
for row in matrix:
print(f"Sum of {row}: ", sum(row))
matrix.sort(key=sum)
# Print the matrix after sorting
print("Sorted matrix:", matrix)
return matrix
</NEW PROGRAM>
Next, I'll provide you with 5 test cases based on the instruction:
<TEST CASE>
sort_matrix([[3, 2, 1], [1, 1, 1], [2, 2, 2]])
</TEST CASE>
<TEST CASE>
sort_matrix([[10, 5, 3], [0, 0, 0], [2, 2, 6]])
</TEST CASE>
<TEST CASE>
sort_matrix([[1, 2, 3], [3, 2, 1], [2, 2, 2]])
</TEST CASE>
<TEST CASE>
sort_matrix([[7], [3, 5], [1, 1, 1, 1]])
</TEST CASE>
<TEST CASE>
sort_matrix([[-1, -1, -1], [1, 0, -1], [0, 0, 0]])
</TEST CASE>
Now, I will give a new program and let's get started to process it:
Instruction: {}
Program:
{}
"""
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--prediction_file", type=str, required=True)
parser.add_argument("--result_file", type=str, required=True)
parser.add_argument("--output_file", type=str, required=True)
args = parser.parse_args()
results = json.load(open(args.result_file))
print(len(results))
predictions = json.load(open(args.prediction_file))
print(len(predictions))
outputs = []
for res, pred in zip(results.values(), predictions):
orig_prompt = pred["prompt"]
orig_completion = pred["completion"]
new_prompt = _print_inter_template.format(orig_prompt, pred["completion"], row="row")
pred["orig_prompt"] = orig_prompt
pred["orig_completion"] = orig_completion
pred["prompt"] = new_prompt
pred["passed"] = res[0][1]["passed"]
outputs.append(pred)
with open(args.output_file, "w", encoding="utf8") as f:
for item in outputs:
f.write(json.dumps(item, ensure_ascii=False) + "\n")
if __name__ == '__main__':
main()