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