Upload own-solver/neurogolf_solver/main.py

own-solver/neurogolf_solver/main.py

@@ -0,0 +1,103 @@
+#!/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 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()