| |
| """ |
| ARC-AGI NeuroGolf Championship - Main Entry Point |
| |
| Usage: |
| python -m neurogolf_solver.main --data_dir ARC-AGI/data/training/ --output_dir submission |
| python -m neurogolf_solver.main --kaggle --output_dir /kaggle/working/submission |
| python -m neurogolf_solver.main --data_dir ARC-AGI/data/training/ --arcgen_dir ARC-GEN-100K/ --use_wandb |
| """ |
|
|
| import argparse |
| import os |
| import sys |
| import time |
| import onnxruntime as ort |
| from .config import get_providers |
| from .data_loader import load_tasks_dir, load_tasks_kaggle |
| from .submission import run_tasks, generate_submission, print_summary |
| from .profiler import score_network |
| from .constants import EXCLUDED_TASKS, MAX_ONNX_FILESIZE |
|
|
| try: |
| import wandb |
| except ImportError: |
| wandb = None |
|
|
|
|
| def check_all_models(output_dir, strict_size, strict_score): |
| """Check all .onnx files for size limit and scoreability.""" |
| size_problems = [] |
| score_problems = [] |
| for f in sorted(os.listdir(output_dir)): |
| if not f.endswith('.onnx'): |
| continue |
| fpath = os.path.join(output_dir, f) |
| fsize = os.path.getsize(fpath) |
| if fsize > MAX_ONNX_FILESIZE: |
| size_problems.append((f, fsize)) |
| macs, memory, params = score_network(fpath) |
| if macs is None or memory is None or params is None: |
| score_problems.append(f) |
| if size_problems: |
| print(f"\n{'!'*70}") |
| print(f"FATAL: {len(size_problems)} .onnx files exceed 1.44MB limit:") |
| for f, sz in size_problems: |
| print(f" {f}: {sz:,} bytes ({sz/1024:.1f} KB)") |
| print(f"{'!'*70}") |
| if strict_size: |
| sys.exit(1) |
| if score_problems: |
| print(f"\nWARNING: {len(score_problems)} .onnx files unscorable by onnx_tool:") |
| for f in score_problems: |
| print(f" {f}") |
| if strict_score: |
| print("Stopping (--strict_score is on).") |
| sys.exit(1) |
| if not size_problems and not score_problems: |
| print(f"\nAll .onnx files pass size and score checks.") |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description='NeuroGolf Solver v5') |
| parser.add_argument('--data_dir', default='ARC-AGI/data/training/') |
| parser.add_argument('--arcgen_dir', default='', help='Path to ARC-GEN-100K/ directory') |
| parser.add_argument('--output_dir', default='/kaggle/working/submission') |
| parser.add_argument('--kaggle', action='store_true', help='Use Kaggle task format') |
| parser.add_argument('--conv_budget', type=float, default=30.0, help='Seconds per conv solver per task') |
| parser.add_argument('--tasks', type=str, default='', help='Comma-separated task numbers') |
| parser.add_argument('--device', type=str, default='auto', choices=['auto', 'cpu', 'cuda']) |
| parser.add_argument('--use_wandb', action='store_true', help='Enable W&B logging') |
| parser.add_argument('--strict_size', type=bool, default=True, help='Halt if any .onnx > 1.44MB (default: True)') |
| parser.add_argument('--strict_score', type=bool, default=False, help='Halt if any model unscorable (default: False)') |
| args = parser.parse_args() |
| providers = get_providers(args.device) |
| config = {"device": args.device, "conv_budget": args.conv_budget, "data_dir": args.data_dir, "arcgen_dir": args.arcgen_dir, "tasks": args.tasks} |
| ort.set_default_logger_severity(3) |
| print(f"Using providers: {providers}") |
| print(f"Strict size: {args.strict_size} | Strict score: {args.strict_score}") |
| print(f"Max .onnx file size: {MAX_ONNX_FILESIZE:,} bytes") |
| if args.kaggle: |
| tasks = load_tasks_kaggle(args.data_dir) |
| else: |
| arcgen = args.arcgen_dir if args.arcgen_dir else None |
| tasks = load_tasks_dir(args.data_dir, arcgen_dir=arcgen) |
| total_arcgen = sum(len(t['data'].get('arc-gen', [])) for t in tasks.values()) |
| print(f"Loaded {len(tasks)} tasks ({total_arcgen} ARC-GEN examples)") |
| task_nums = [int(t) for t in args.tasks.split(',')] if args.tasks else sorted(tasks.keys()) |
| print(f"Solving {len(task_nums)} tasks") |
| print(f"Conv budget: {args.conv_budget}s per task") |
| print("=" * 70) |
| t0 = time.time() |
| if args.use_wandb and wandb is not None: |
| with wandb.init(project="neurogolf", name="solver_run", config=config): |
| results, costs_dict, total_score = run_tasks(task_nums, tasks, args.output_dir, providers, args.conv_budget, EXCLUDED_TASKS, use_wandb=True) |
| else: |
| results, costs_dict, total_score = run_tasks(task_nums, tasks, args.output_dir, providers, args.conv_budget, EXCLUDED_TASKS, use_wandb=False) |
| elapsed = time.time() - t0 |
| check_all_models(args.output_dir, args.strict_size, args.strict_score) |
| submission_info = generate_submission(args.output_dir, results, costs_dict, task_nums) |
| print_summary(results, submission_info, elapsed) |
|
|
| if __name__ == '__main__': |
| main() |
|
|
