Move own-solver/neurogolf_solver/solvers/solver_registry.py to own-solver/

own-solver/neurogolf_solver/solvers/solver_registry.py

@@ -0,0 +1,163 @@
+#!/usr/bin/env python3
+"""Solver registry and task solving orchestration."""
+
+import os
+import time
+import onnx
+from .analytical import s_identity, s_constant, s_color_map, s_transpose
+from .geometric import s_flip, s_rotate, s_shift, s_fixed_crop, s_gravity
+from .tiling import (s_tile, s_upscale, s_kronecker, s_nonuniform_scale, s_diagonal_tile,
+                     s_mirror_h, s_mirror_v, s_quad_mirror, s_concat, s_concat_enhanced,
+                     s_spatial_gather, s_varshape_spatial_gather)
+from .gravity import s_gravity_unrolled
+from .edge import s_edge_detect
+from .mode import s_mode_fill
+from .wave1 import (s_downsample_stride, s_symmetry_complete, s_extract_inner,
+                    s_add_border, s_sparse_fill, s_channel_filter)
+from .conv import solve_conv_fixed, solve_conv_variable, solve_conv_diffshape, solve_conv_var_diff
+from ..data_loader import get_exs, fixed_shapes
+from ..validators import validate
+from ..profiler import score_network
+from ..constants import EXCLUDED_TASKS, MAX_ONNX_FILESIZE
+
+# Analytical solvers registry — order matters (cheaper first)
+ANALYTICAL_SOLVERS = [
+    ('identity', s_identity),
+    ('constant', s_constant),
+    ('color_map', s_color_map),
+    ('transpose', s_transpose),
+    ('flip', s_flip),
+    ('rotate', s_rotate),
+    ('shift', s_shift),
+    ('tile', s_tile),
+    ('upscale', s_upscale),
+    ('kronecker', s_kronecker),
+    ('nonuniform_scale', s_nonuniform_scale),
+    ('mirror_h', s_mirror_h),
+    ('mirror_v', s_mirror_v),
+    ('quad_mirror', s_quad_mirror),
+    ('concat', s_concat),
+    ('concat_enhanced', s_concat_enhanced),
+    ('diagonal_tile', s_diagonal_tile),
+    ('fixed_crop', s_fixed_crop),
+    ('spatial_gather', s_spatial_gather),
+    ('varshape_spatial_gather', s_varshape_spatial_gather),
+    ('gravity_unrolled', s_gravity_unrolled),
+    ('edge_detect', s_edge_detect),
+    ('mode_fill', s_mode_fill),
+    ('downsample_stride', s_downsample_stride),
+    ('symmetry_complete', s_symmetry_complete),
+    ('extract_inner', s_extract_inner),
+    ('add_border', s_add_border),
+    ('sparse_fill', s_sparse_fill),
+    ('channel_filter', s_channel_filter),
+]
+
+
+def _check_size(path):
+    """Return True if file is within 1.44MB limit."""
+    try:
+        return os.path.getsize(path) <= MAX_ONNX_FILESIZE
+    except OSError:
+        return False
+
+
+def _check_scoreable(path):
+    """Return True if score_network returns valid scores (not None).
+    A model that can't be scored will be REJECTED by Kaggle."""
+    macs, memory, params = score_network(path)
+    if macs is None or memory is None or params is None:
+        return False
+    return True
+
+
+def _accept_model(path, td, providers):
+    """Full acceptance check: size + validate (outputs) + scoreable.
+    Returns True only if model would be accepted by Kaggle."""
+    if not _check_size(path):
+        return False
+    if not validate(path, td, providers):
+        return False
+    if not _check_scoreable(path):
+        return False
+    return True
+
+
+def _cleanup_failed(path):
+    """Delete leftover .onnx file from failed solve attempts.
+    Prevents bad files from ending up in submission zip."""
+    try:
+        if os.path.exists(path):
+            os.remove(path)
+    except OSError:
+        pass
+
+
+def solve_task(tn, td, outdir, providers, conv_budget=30.0, excluded_tasks=None):
+    """Solve a single ARC-AGI task.
+    
+    Returns: (ok, solver_name, file_size, elapsed, model_path)
+    If unsolved, deletes any leftover .onnx file.
+    """
+    if excluded_tasks is None:
+        excluded_tasks = EXCLUDED_TASKS
+    
+    t_start = time.time()
+    os.makedirs(outdir, exist_ok=True)
+    path = os.path.join(outdir, f"task{tn:03d}.onnx")
+
+    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 = sfn(td)
+            if model is None:
+                continue
+            onnx.save(model, path)
+            if _accept_model(path, td, providers):
+                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, providers, time_budget=conv_time / 2)
+            if result is not None:
+                if _check_size(path) and _check_scoreable(path):
+                    sname, model = result
+                    return True, sname, os.path.getsize(path), time.time() - t_start, path
+        result = solve_conv_variable(td, path, providers, time_budget=conv_time)
+        if result is not None:
+            if _check_size(path) and _check_scoreable(path):
+                sname, model = result
+                return True, sname, os.path.getsize(path), time.time() - t_start, path
+    else:
+        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, providers, time_budget=conv_time)
+                if result is not None:
+                    if _check_size(path) and _check_scoreable(path):
+                        sname, model = result
+                        return True, sname, os.path.getsize(path), time.time() - t_start, path
+
+        result = solve_conv_var_diff(td, path, providers, time_budget=conv_time)
+        if result is not None:
+            if _check_size(path) and _check_scoreable(path):
+                sname, model = result
+                return True, sname, os.path.getsize(path), time.time() - t_start, path
+
+    # All solvers failed — delete leftover .onnx file
+    _cleanup_failed(path)
+    return False, None, None, time.time() - t_start, path