v5 refactor: add main.py entry point (with W&B init)

neurogolf_solver/main.py

@@ -0,0 +1,88 @@
+#!/usr/bin/env python3
+"""
+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 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 .constants import EXCLUDED_TASKS
+
+try:
+    import wandb
+except ImportError:
+    wandb = None
+
+
+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')
+    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}")
+
+    # Load tasks
+    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)")
+    print(f"Excluded tasks: {sorted(EXCLUDED_TASKS)}")
+
+    task_nums = [int(t) for t in args.tasks.split(',')] if args.tasks else sorted(tasks.keys())
+    active_tasks = [t for t in task_nums if t not in EXCLUDED_TASKS]
+    print(f"Solving {len(active_tasks)} active tasks (skipping {len(task_nums) - len(active_tasks)} excluded)")
+    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
+
+    submission_info = generate_submission(args.output_dir, results, costs_dict, active_tasks)
+    print_summary(results, submission_info, elapsed)
+
+
+if __name__ == '__main__':
+    main()