wehub-resource-sync
101 lines
3.8 KiB
Python
101 lines
3.8 KiB
Python
import argparse
|
|
import copy
|
|
import json
|
|
|
|
from datasets import load_dataset
|
|
|
|
_decompose_judge_template = """I will show you a competitive programming problem as well as its one solution. Note that the solution can be correct or incorrect. Please:
|
|
(1) Use **natural language** to clarify the thoughts behind the solution. You may consider including: the limitation about running time and memory, potential data structure, potential algorithms, etc. Please put your thoughts in <THOUGHTS></THOUGHTS>.
|
|
(2) After that, please split the above program solution into several segments. Each segment should follow the following format:
|
|
<SEGMENT 1>
|
|
<CODE SEGMENT>
|
|
*One segment of code from the original solution. It has to be a continuous part of the original solution, and should implement ONE basic and meaningful function.*
|
|
</CODE SEGMENT>
|
|
<NL DESCRIPTION>
|
|
*The corresponding description or interpretation of the above code snippet/segment. You should also include the clarification about the definition of the involved variables.*
|
|
</NL DESCRIPTION>
|
|
<WRAPPED CODE SEGMENT>
|
|
*According to the original code solution, the split code snippet, and the corresponding sub-clarification about the code segment, please wrap the code snippet into a function. The function should be a standalone function that can be executed independently. You should also include necessary docstring/comments in the function to define the input variables and return values.*
|
|
</WRAPPED CODE SEGMENT>
|
|
</SEGMENT 1>
|
|
<SEGMENT 2>
|
|
*Same structure as SEGMENT 1*
|
|
</SEGMENT 2>
|
|
|
|
Now, let's get started:
|
|
|
|
Problem Title: {}
|
|
|
|
Description:
|
|
{}
|
|
|
|
Time Limits:
|
|
{} Seconds
|
|
Memory Limits:
|
|
{} Bytes
|
|
|
|
Solution Program:
|
|
```
|
|
{}
|
|
```
|
|
"""
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description='Completion Judge')
|
|
parser.add_argument("--file_path", type=str, required=True)
|
|
parser.add_argument("--output_file", type=str, required=True)
|
|
parser.add_argument("--sample_num", default=-1, type=int, help="Number of samples to generate")
|
|
args = parser.parse_args()
|
|
|
|
data = load_dataset("parquet", data_files=args.file_path)["train"]
|
|
outputs = []
|
|
if args.sample_num > 0:
|
|
sample_num = args.sample_num
|
|
else:
|
|
sample_num = len(data)
|
|
for item in data:
|
|
description = item["description"]
|
|
title = item["name"]
|
|
time_limits = item["time_limit"]["seconds"]
|
|
mem_limits = item["memory_limit_bytes"]
|
|
if len(item["solutions"]["solution"]) == 0 or len(item["incorrect_solutions"]["solution"]) == 0:
|
|
continue
|
|
|
|
pos_solution = item["solutions"]["solution"][0]
|
|
neg_solution = item["incorrect_solutions"]["solution"][0]
|
|
|
|
item = {
|
|
"title": title,
|
|
"description": description,
|
|
"time_limits": time_limits,
|
|
"mem_limits": mem_limits,
|
|
}
|
|
pos_item = copy.deepcopy(item)
|
|
pos_item["solution"] = pos_solution
|
|
pos_item["prompt"] = _decompose_judge_template.format(title, description, time_limits, mem_limits,
|
|
pos_solution)
|
|
pos_item["passed"] = True
|
|
outputs.append(pos_item)
|
|
|
|
neg_item = copy.deepcopy(item)
|
|
neg_item["solution"] = neg_solution
|
|
neg_item["prompt"] = _decompose_judge_template.format(title, description, time_limits, mem_limits,
|
|
neg_solution)
|
|
neg_item["passed"] = False
|
|
outputs.append(neg_item)
|
|
|
|
if len(outputs) >= sample_num:
|
|
break
|
|
|
|
if args.sample_num > 0:
|
|
args.output_file = args.output_file.replace(".jsonl", f"_{sample_num}.jsonl")
|
|
|
|
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()
|