|
if __name__ == "__main__": |
|
import argparse |
|
import os |
|
import re |
|
from zipfile import ZipFile |
|
|
|
parser = argparse.ArgumentParser( |
|
description="Package and zip LLM provided solutions" |
|
) |
|
parser.add_argument("--root", type=str, help="Root directory of solutions") |
|
parser.add_argument( |
|
"--output", type=str, required=True, help="Output directory of zip files" |
|
) |
|
args = parser.parse_args() |
|
|
|
|
|
directory_pattern = re.compile(r"(.*)_temp_(?:\d*\.?\d+)$") |
|
|
|
assert os.path.isdir(args.root) |
|
os.makedirs(args.output, exist_ok=True) |
|
|
|
for directory in os.listdir(args.root): |
|
match = directory_pattern.match(directory) |
|
if not match: |
|
continue |
|
directory_name = match.group(0) |
|
full_dir_path = os.path.join(args.root, directory_name) |
|
assert os.path.isdir(full_dir_path) |
|
print(f"Processing {full_dir_path}") |
|
|
|
zip_file_path = os.path.join(args.output, f"{directory_name}.zip") |
|
if os.path.exists(zip_file_path): |
|
print(f"Skipping -- {zip_file_path} already exists") |
|
continue |
|
|
|
with ZipFile(zip_file_path, "w") as zip_file: |
|
|
|
for task_id in os.listdir(full_dir_path): |
|
task_dir = os.path.join(full_dir_path, task_id) |
|
if not os.path.isdir(task_dir): |
|
continue |
|
for sample_id in os.listdir(task_dir): |
|
sample_file = os.path.join(task_dir, sample_id) |
|
if not sample_file.endswith(".py"): |
|
continue |
|
zip_file.write( |
|
sample_file, os.path.join(directory_name, task_id, sample_id) |
|
) |
|
|
