Put W&B code back into run_task

neurogolf_solver.py

@@ -2,7 +2,6 @@
 """
 ARC-AGI NeuroGolf Championship - Complete Solver v5
 Format: [1,10,30,30] one-hot input/output, opset 17, IR version 8.
-
 v5 CHANGES (from v4):
   - Opset 10 → 17, IR 10 → 8
   - s_flip: Slice(step=-1) replaces Gather — 0 MACs (was ~165K)
@@ -11,7 +10,6 @@ v5 CHANGES (from v4):
   - All Pad nodes: tensor-based pads input (opset 17 requirement)
   - All ReduceSum nodes: axes as tensor input (opset 13+ requirement)
   - All other solvers unchanged from v4
-
 Solvers:
   - Analytical: identity, constant, color_map, transpose, flip, rotate, tile, upscale,
                 concat, concat_enhanced, spatial_gather, varshape_spatial_gather,
@@ -19,7 +17,6 @@ Solvers:
   - Conv (fixed shape): Slice -> Conv -> ArgMax -> Equal+Cast -> Pad
   - Conv (variable shape): Conv(30x30) -> ArgMax -> Equal+Cast -> Mul(mask)
   - Conv (diff shape): Slice -> Conv -> Slice(crop) -> ArgMax -> Equal+Cast -> Pad
-
 Usage:
   python neurogolf_solver.py --data_dir ARC-AGI/data/training/ --output_dir submission
   python neurogolf_solver.py --data_dir ARC-AGI/data/training/ --output_dir submission --conv_budget 60 --arcgen_dir ARC-GEN-100K/
@@ -1170,130 +1167,198 @@ ANALYTICAL_SOLVERS = [
     ('varshape_spatial_gather', s_varshape_spatial_gather),
 ]
 
-def solve_task(tn, td, output_dir, conv_budget=30.0, verbose=True):
-    path = os.path.join(output_dir, f"task{tn:03d}.onnx")
-    for name, solver in ANALYTICAL_SOLVERS:
+
+def solve_task(tn, td, outdir, conv_budget=30.0):
+    t_start = time.time()
+    os.makedirs(outdir, exist_ok=True)
+    path = os.path.join(outdir, f"task{tn:03d}.onnx")
+    
+    # Skip excluded tasks
+    if tn in EXCLUDED_TASKS:
+        return False, 'excluded', None, time.time() - t_start, path
+    
+    # 1. Try analytical solvers (fast, tiny models)
+    for sname, sfn in ANALYTICAL_SOLVERS:
         try:
-            model = solver(td)
-        except Exception as e:
-            if verbose: print(f"  {name}: ERROR {e}")
-            continue
-        if model is not None:
+            model = sfn(td)
+            if model is None: continue
             onnx.save(model, path)
-            if validate(path, td):
-                macs, mem, par = score_network(path)
-                if macs is not None:
-                    cost = macs + mem + par
-                    score = max(1.0, 25.0 - math.log(cost)) if cost > 0 else 25.0
-                    if verbose: print(f"  {name}: PASS  cost={cost}  score={score:.2f}")
-                    return name, score
-            else:
-                if verbose: print(f"  {name}: model built but FAILED validation")
-    conv_solvers = [
-        ('conv_fixed', solve_conv_fixed),
-        ('conv_variable', solve_conv_variable),
-        ('conv_diffshape', solve_conv_diffshape),
-        ('conv_var_diff', solve_conv_var_diff),
-    ]
-    for name, solver in conv_solvers:
-        try:
-            result = solver(td, path, time_budget=conv_budget)
-        except Exception as e:
-            if verbose: print(f"  {name}: ERROR {e}")
-            continue
+            if validate(path, td): 
+                return True, sname, os.path.getsize(path), time.time() - t_start, path
+        except: pass
+    
+    # 2. Determine task shape category and try conv solvers
+    exs = get_exs(td)
+    same_shape = all(inp.shape == out.shape for inp, out in exs)
+    shapes = set(inp.shape for inp, _ in exs)
+    fixed_in = len(shapes) == 1
+    
+    conv_time = conv_budget
+    
+    if same_shape:
+        if fixed_in:
+            result = solve_conv_fixed(td, path, time_budget=conv_time/2)
+            if result is not None:
+                sname, model = result
+                return True, sname, os.path.getsize(path), time.time() - t_start, path
+        result = solve_conv_variable(td, path, time_budget=conv_time)
         if result is not None:
-            solver_type, model = result
-            onnx.save(model, path)
-            macs, mem, par = score_network(path)
-            if macs is not None:
-                cost = macs + mem + par
-                score = max(1.0, 25.0 - math.log(cost)) if cost > 0 else 25.0
-                if verbose: print(f"  {solver_type}: PASS  cost={cost}  score={score:.2f}")
-                return solver_type, score
-    return None
-
-def main():
-    parser = argparse.ArgumentParser(description='NeuroGolf Solver v5')
-    parser.add_argument('--data_dir', type=str, default=None)
-    parser.add_argument('--kaggle_dir', type=str, default=None)
-    parser.add_argument('--arcgen_dir', type=str, default=None)
-    parser.add_argument('--output_dir', type=str, default='submission')
-    parser.add_argument('--conv_budget', type=float, default=30.0)
-    parser.add_argument('--task', type=int, default=None)
-    parser.add_argument('--verbose', action='store_true', default=True)
-    parser.add_argument('--quiet', action='store_true', default=False)
-    args = parser.parse_args()
-    if args.quiet:
-        args.verbose = False
-    os.makedirs(args.output_dir, exist_ok=True)
-    if args.kaggle_dir:
-        tasks = load_tasks_kaggle(args.kaggle_dir)
-    elif args.data_dir:
-        tasks = load_tasks_dir(args.data_dir, args.arcgen_dir)
+            sname, model = result
+            return True, sname, os.path.getsize(path), time.time() - t_start, path
     else:
-        for p in ['/kaggle/input/competitions/neurogolf-2026/', 'ARC-AGI/data/training/']:
-            if os.path.exists(p):
-                if 'kaggle' in p:
-                    tasks = load_tasks_kaggle(p)
-                else:
-                    tasks = load_tasks_dir(p, args.arcgen_dir)
-                break
-        else:
-            print("ERROR: No data directory found. Use --data_dir or --kaggle_dir")
-            sys.exit(1)
+        sp = fixed_shapes(td)
+        if sp is not None:
+            (IH,IW),(OH,OW) = sp
+            if OH <= IH and OW <= IW:
+                result = solve_conv_diffshape(td, path, time_budget=conv_time)
+                if result is not None:
+                    sname, model = result
+                    return True, sname, os.path.getsize(path), time.time() - t_start, path
+        
+        # Try variable diff-shape conv (output within input bounds)
+        result = solve_conv_var_diff(td, path, time_budget=conv_time)
+        if result is not None:
+            sname, model = result
+            return True, sname, os.path.getsize(path), time.time() - t_start, path
+    
+    return False, None, None, time.time() - t_start, path
+
+def run_tasks(task_nums, tasks, output_dir, conv_budget, use_wandb):
     results = {}
-    total_score = 0.0
-    solved = 0
-    t_total = time.time()
-    task_nums = [args.task] if args.task else sorted(tasks.keys())
+    costs_dict = {}
+    total_score = 0
     for tn in task_nums:
-        if tn in EXCLUDED_TASKS:
-            if args.verbose: print(f"Task {tn:3d}: EXCLUDED")
-            continue
         if tn not in tasks:
-            if args.verbose: print(f"Task {tn:3d}: NOT FOUND")
             continue
+        if tn in EXCLUDED_TASKS:
+            print(f"Task {tn:3d}: EXCLUDED (officially)")
+            continue
+
         td = tasks[tn]['data']
-        hex_id = tasks[tn]['hex']
-        if args.verbose: print(f"\nTask {tn:3d} ({hex_id}):")
-        result = solve_task(tn, td, args.output_dir, args.conv_budget, args.verbose)
-        if result is not None:
-            solver_type, score = result
-            results[tn] = {'solver': solver_type, 'score': score, 'hex': hex_id}
+        ok, sname, sz, t_task, model_path = solve_task(tn, td, output_dir, conv_budget)
+
+        if ok:
+            macs, memory, params = score_network(model_path)
+            if macs is None:
+                macs, memory, params = 0, 0, 0
+            cost = macs + memory + params
+            score = max(1.0, 25.0 - math.log(max(1, cost)))
             total_score += score
-            solved += 1
+
+            results[tn] = (sname, t_task, sz)
+            costs_dict[tn] = cost
+            print(f"Task {tn:3d}: {sname:25s} {score:7.3f} {cost:>12} {t_task:7.3f}s  ({sz:>8,} bytes)")
         else:
-            total_score += 1.0
-            if args.verbose: print(f"  UNSOLVED")
-    elapsed = time.time() - t_total
-    print(f"\n{'='*60}")
-    print(f"RESULTS: {solved}/{len(task_nums)} tasks solved")
-    print(f"Total score: {total_score:.1f}")
-    print(f"Time: {elapsed:.1f}s")
-    print(f"{'='*60}")
-    solver_counts = Counter(r['solver'] for r in results.values())
-    solver_scores = {}
-    for tn, r in results.items():
-        st = r['solver']
-        solver_scores[st] = solver_scores.get(st, 0) + r['score']
-    print("\nSolver breakdown:")
-    for st in sorted(solver_counts.keys()):
-        print(f"  {st}: {solver_counts[st]} tasks, total score {solver_scores[st]:.1f}, avg {solver_scores[st]/solver_counts[st]:.2f}")
-    csv_path = os.path.join(args.output_dir, 'submission.csv')
+            print(f"Task {tn:3d}: UNSOLVED  {t_task:7.3f}s")
+            cost = 0
+
+        if use_wandb and wandb is not None:
+            wandb.log({
+                "task_id": tn,
+                "solver": sname if ok else "unsolved",
+                "onnx_bytes": sz if ok else 0,
+                "task_time_sec": t_task,
+                "cost": cost,
+                "score": score if ok else 0,
+            })
+
+    return results, costs_dict, total_score
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    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')
+    parser.add_argument('--conv_budget', type=float, default=30.0)
+    parser.add_argument('--tasks', type=str, default='')
+    parser.add_argument('--device', type=str, default='auto', choices=['auto','cpu','cuda'])
+    parser.add_argument('--use_wandb', action='store_true')
+    args = parser.parse_args()
+    global ORT_PROVIDERS
+    config = {
+        "device": args.device,
+        "conv_budget": args.conv_budget,
+        "data_dir": args.data_dir,
+        "arcgen_dir": args.arcgen_dir,
+        "tasks": args.tasks,
+    }
+ 
+    if args.device == 'cuda':
+        ORT_PROVIDERS = ['CUDAExecutionProvider', 'CPUExecutionProvider']
+    elif args.device == 'cpu':
+        ORT_PROVIDERS = ['CPUExecutionProvider']
+    
+    ort.set_default_logger_severity(3)
+    print(f"Using providers: {ORT_PROVIDERS}")
+    
+    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)
+    
+    # Count arc-gen examples
+    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, args.conv_budget, use_wandb=True)
+    else:
+        results, costs_dict, total_score = run_tasks(task_nums, tasks, args.output_dir, args.conv_budget, use_wandb=False)
+
+    elapsed = time.time() - t0
+    print(f"\n{'='*70}")
+    print(f"Solved: {len(results)}/{len(active_tasks)} active tasks in {elapsed:.0f}s")
+    solver_names = [v[0] for v in results.values()]
+    sc = Counter(solver_names)
+    for s, c in sc.most_common(): print(f"  {s}: {c}")
+    
+    # Generate submission
+    outdir = args.output_dir
+    n_files = len([f for f in os.listdir(outdir) if f.endswith('.onnx')])
+    total_size = sum(os.path.getsize(os.path.join(outdir, f))
+                     for f in os.listdir(outdir) if f.endswith('.onnx'))
+    
+    # Create submission.zip
+    zip_path = os.path.join(os.path.dirname(outdir) or '/kaggle/working/', 'submission.zip')
+    buf = io.BytesIO()
+    with zipfile.ZipFile(buf, 'w', zipfile.ZIP_DEFLATED) as zf:
+        for f in sorted(os.listdir(outdir)):
+            if f.endswith('.onnx'):
+                zf.write(os.path.join(outdir, f), f)
+    zip_bytes = buf.getvalue()
+    with open(zip_path, 'wb') as f:
+        f.write(zip_bytes)
+    zip_size = len(zip_bytes)
+    
+    # Create submission.csv
+    csv_path = os.path.join(os.path.dirname(outdir) or '.', 'submission.csv')
     with open(csv_path, 'w', newline='') as f:
         w = csv.writer(f)
-        w.writerow(['task_num', 'hex_id', 'solver', 'score', 'onnx_file'])
-        for tn in sorted(results.keys()):
-            r = results[tn]
-            w.writerow([tn, r['hex'], r['solver'], f"{r['score']:.3f}", f"task{tn:03d}.onnx"])
-    zip_path = os.path.join(args.output_dir, 'submission.zip')
-    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zf:
-        for tn in sorted(results.keys()):
-            onnx_path = os.path.join(args.output_dir, f"task{tn:03d}.onnx")
-            if os.path.exists(onnx_path):
-                zf.write(onnx_path, f"task{tn:03d}.onnx")
-    print(f"\nSubmission files: {csv_path}, {zip_path}")
-    print(f"Models in zip: {len(results)}")
+        w.writerow(['task_id', 'total_cost'])
+        for tn in sorted(costs_dict.keys()):
+            w.writerow([f'task{tn:03d}', costs_dict[tn]])
+    
+    # Estimate LB score: solved tasks get their score, unsolved get 1.0
+    unsolved_count = len(active_tasks) - len(results)
+    est_lb = total_score + unsolved_count * 1.0
+    
+    print(f"\n{n_files} ONNX files, {total_size/1024:.1f} KB uncompressed")
+    print(f"ZIP size: {zip_size/1024:.1f} KB / {MAX_FILESIZE/1024:.0f} KB limit {'OK' if zip_size <= MAX_FILESIZE else 'OVER!'}")
+    print(f"Estimated LB score: {est_lb:.1f} (solved: {total_score:.1f} + unsolved: {unsolved_count}×1.0)")
+    print(f"Written: {zip_path} | {csv_path}")
 
 if __name__ == '__main__':
     main()
+