Update neurogolf_solver.py

neurogolf_solver.py

@@ -681,6 +681,43 @@ def solve_task(tn, td, outdir, conv_budget=30.0):
     
     return False, None, None, time.time() - t_start, path
 
+def run_tasks(task_nums, tasks, output_dir, conv_budget, use_wandb):
+    results = {}
+
+    for tn in task_nums:
+        if tn not in tasks:
+            continue
+
+        td = tasks[tn]['data']
+        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
+            score = macs + memory + params
+
+            results[tn] = (sname, t_task, sz)
+            print(f"Task {tn:3d}: {sname:20s}  {score} {t_task:7.3f}s  ({sz:>8,} bytes)")
+        else:
+            print(f"Task {tn:3d}: UNSOLVED  {t_task:7.3f}s")
+            macs, memory, params, score = 0, 0, 0, 0
+
+        if use_wandb:
+            wandb.log({
+                "task_id": tn,
+                "solver": sname if ok else "unsolved",
+                "onnx_bytes": sz if ok else 0,
+                "task_time_sec": t_task,
+                "macs": macs,
+                "memory": memory,
+                "params": params,
+                "score": score,
+            })
+
+    return results
+
+
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument('--data_dir', default='ARC-AGI/data/training/')
@@ -689,7 +726,7 @@ def main():
     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('--wandb_run_id', type=str, default=None)
+    parser.add_argument('--use_wandb', action='store_true')
     args = parser.parse_args()
     global ORT_PROVIDERS
     config = {
@@ -712,38 +749,30 @@ def main():
     print("=" * 70)
     t0 = time.time()
     results = {}
-    with wandb.init(
-        project="neurogolf",
-        name="solver_run",
-        config=config,
-    ) as run:
-        
-        for tn in task_nums:
-            if tn not in tasks: continue
-            td = tasks[tn]['data']
-            ok, sname, sz, t_task, model_path  = solve_task(tn, td, args.output_dir, args.conv_budget)
-            if ok:
-                macs, memory, params = score_network(model_path)
-                score = macs + memory + params
-                results[tn] = (sname, t_task, sz)
-                print(f"Task {tn:3d}: {sname:20s}  {score} {t_task:7.3f}s  ({sz:>8,} bytes)")
-                if macs is None:  
-                    macs, memory, params = 0, 0, 0
-            else:
-                print(f"Task {tn:3d}: UNSOLVED  {t_task:7.3f}s")
-                macs, memory, params, score = 0, 0, 0, 0
-            wandb.log({
-                "task_id": tn,
-                "solver": sname if ok else "unsolved",
-                "onnx_bytes": sz if ok else 0,
-                "task_time_sec": t_task,
-                "macs": macs,
-                "memory": memory,
-                "params": params,
-                "score": score,
-            })
+    
+    if args.use_wandb:
+        with wandb.init(
+            project="neurogolf",
+            name="solver_run",
+            config=config,
+        ):
+            results = run_tasks(
+                task_nums,
+                tasks,
+                args.output_dir,
+                args.conv_budget,
+                use_wandb=True
+            )
+    
+    else:
+        results = 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(task_nums)} in {elapsed:.0f}s")