88 lines
2.9 KiB
Python
88 lines
2.9 KiB
Python
import os.path
|
|
|
|
from datasets import load_dataset
|
|
from glob import glob
|
|
from tqdm import tqdm
|
|
|
|
|
|
class CodeContestReader:
|
|
def __call__(self, file_path):
|
|
if os.path.exists(file_path):
|
|
data = load_dataset("parquet", data_files=file_path)["train"]
|
|
else:
|
|
data = load_dataset("parquet", data_files=list(glob(file_path)))["train"]
|
|
print(len(data))
|
|
outputs = []
|
|
for item in tqdm(data):
|
|
description = item["description"]
|
|
title = item["name"]
|
|
if item["time_limit"]:
|
|
time_limits = item["time_limit"]["seconds"]
|
|
else:
|
|
time_limits = "N/A"
|
|
if item["memory_limit_bytes"]:
|
|
mem_limits = item["memory_limit_bytes"]
|
|
else:
|
|
mem_limits = "N/A"
|
|
pos_solution = item["solutions"]["solution"]
|
|
neg_solution = item["incorrect_solutions"]["solution"]
|
|
|
|
outputs.append({
|
|
"name": title,
|
|
"description": description,
|
|
"time_limit": time_limits,
|
|
"memory_limit": mem_limits,
|
|
"pos_solution": pos_solution,
|
|
"neg_solution": neg_solution,
|
|
"id": f"{title}",
|
|
})
|
|
|
|
return outputs
|
|
|
|
|
|
class CodeContestFlatReader:
|
|
def __call__(self, file_path):
|
|
if os.path.exists(file_path):
|
|
data = load_dataset("parquet", data_files=file_path)["train"]
|
|
else:
|
|
data = load_dataset("parquet", data_files=list(glob(file_path)))["train"]
|
|
print(len(data))
|
|
outputs = []
|
|
for item in tqdm(data):
|
|
description = item["description"]
|
|
title = item["name"]
|
|
if item["time_limit"]:
|
|
time_limits = item["time_limit"]["seconds"]
|
|
else:
|
|
time_limits = "N/A"
|
|
if item["memory_limit_bytes"]:
|
|
mem_limits = item["memory_limit_bytes"]
|
|
else:
|
|
mem_limits = "N/A"
|
|
pos_solution = item["solutions"]["solution"]
|
|
neg_solution = item["incorrect_solutions"]["solution"]
|
|
|
|
for i, solution in enumerate(pos_solution):
|
|
outputs.append({
|
|
"name": title,
|
|
"description": description,
|
|
"time_limit": time_limits,
|
|
"memory_limit": mem_limits,
|
|
"solution": solution,
|
|
"id": f"{title}_{i}",
|
|
"value": 1,
|
|
})
|
|
|
|
for i, solution in enumerate(neg_solution):
|
|
outputs.append({
|
|
"name": title,
|
|
"description": description,
|
|
"time_limit": time_limits,
|
|
"memory_limit": mem_limits,
|
|
"solution": solution,
|
|
"id": f"{title}_{i}",
|
|
"value": 0,
|
|
})
|
|
|
|
return outputs
|