wehub-resource-sync
57 lines
1.9 KiB
Python
57 lines
1.9 KiB
Python
import argparse
|
|
import copy
|
|
import json
|
|
|
|
from datasets import load_dataset
|
|
|
|
_code_judge_template = """I will show you a programming problem. You will first be given the instruction about the program, and then I will show you the program implementation. Please determine if the completed program is correct or not. You can think step by step before determine its correctness.
|
|
|
|
After your reasoning process, **REMEMBER** to use **### The program is {}.** to repeat your judgement, where {} should only be **correct** or **incorrect**.
|
|
|
|
Instruction: {}
|
|
|
|
Program:
|
|
```
|
|
{}
|
|
```
|
|
"""
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description='Completion Judge')
|
|
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))
|
|
task_id2passed = {}
|
|
for task_id, res in results.items():
|
|
attempt = res[0] # Pass@1
|
|
attempt_res = attempt[1]
|
|
passed = attempt_res["passed"]
|
|
task_id2passed[task_id] = passed
|
|
print(len(task_id2passed))
|
|
|
|
predictions = json.load(open(args.prediction_file))
|
|
print(len(predictions))
|
|
outputs = []
|
|
for item in predictions:
|
|
if str(item["task_id"]) not in task_id2passed:
|
|
continue
|
|
|
|
orig_prompt = item["prompt"]
|
|
new_prompt = _code_judge_template.format("{}", "{}", orig_prompt, item["completion"])
|
|
item["orig_prompt"] = orig_prompt
|
|
item["prompt"] = new_prompt
|
|
item["passed"] = task_id2passed[str(item["task_id"])] # FIXME: This is not aligned.
|
|
outputs.append(item)
|
|
|
|
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()
|