| 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) |
| ) |
|
|
