Add build_submission.py - single script to generate submission from pre-built models

medal-solvers/build_submission.py

@@ -0,0 +1,342 @@
+#!/usr/bin/env python3
+"""
+build_submission.py — Generates the full submission zip for NeuroGolf.
+
+Usage (Kaggle notebook):
+    python build_submission.py \
+        --base /kaggle/input/competitions/neurogolf-2026/submission-6043.zip \
+        --task-data-dir /kaggle/input/competitions/neurogolf-2026 \
+        --wave22 /path/to/wave22.py \
+        --optimized-dir /path/to/medal-solvers/optimized \
+        --output /kaggle/working/submission.zip
+
+What it does:
+    1. Loads pre-built .onnx files from --optimized-dir (the 7 hand-crafted + 16 wave22)
+    2. Optionally builds MORE wave22 models on-the-fly (if --wave22 given)
+    3. Validates every replacement model (100% pass required)
+    4. Creates submission zip: base + all valid replacements
+
+Arguments:
+    --base            Path to the base submission zip (submission-6043.zip)
+    --task-data-dir   Directory containing taskNNN.json files
+    --optimized-dir   Directory with pre-built .onnx files to swap in
+    --wave22          Path to wave22.py to build additional models (optional)
+    --output          Output zip path (default: submission.zip)
+    --skip-validation Skip example validation (not recommended)
+
+Minimal Kaggle usage (just swap pre-built models):
+    python build_submission.py \
+        --base submission-6043.zip \
+        --task-data-dir /kaggle/input/competitions/neurogolf-2026 \
+        --optimized-dir optimized \
+        --output /kaggle/working/submission.zip
+"""
+import argparse
+import os
+import sys
+import json
+import zipfile
+import math
+import numpy as np
+import onnx
+from onnx import helper as oh, numpy_helper as onh, TensorProto
+import onnxruntime as ort
+import zlib
+import base64
+import re
+
+
+def _d(b64, shape, dtype):
+    """Decode base64+zlib compressed numpy array (used by wave22)."""
+    return np.frombuffer(zlib.decompress(base64.b64decode(b64)), dtype=dtype).reshape(shape)
+
+
+def validate_model(model_path, task_data_dir, task_num):
+    """Validate model passes ALL examples. Returns (right, wrong)."""
+    task_file = os.path.join(task_data_dir, f'task{task_num:03d}.json')
+    if not os.path.exists(task_file):
+        return None, None
+
+    with open(task_file) as f:
+        data = json.load(f)
+
+    sess = ort.InferenceSession(model_path)
+    all_ex = data.get('train', []) + data.get('test', []) + data.get('arc-gen', [])
+
+    right, wrong = 0, 0
+    for ex in all_ex:
+        inp_grid = ex['input']
+        if max(len(inp_grid), max((len(r) for r in inp_grid), default=0)) > 30:
+            continue
+        inp = np.zeros((1, 10, 30, 30), dtype=np.float32)
+        for r, row in enumerate(inp_grid):
+            for c, v in enumerate(row):
+                if r < 30 and c < 30:
+                    inp[0][v][r][c] = 1.0
+        exp = np.zeros((1, 10, 30, 30), dtype=np.float32)
+        for r, row in enumerate(ex['output']):
+            for c, v in enumerate(row):
+                if r < 30 and c < 30:
+                    exp[0][v][r][c] = 1.0
+
+        result = sess.run(['output'], {'input': inp})
+        out = (result[0] > 0.0).astype(float)
+        if np.array_equal(out, exp):
+            right += 1
+        else:
+            wrong += 1
+    return right, wrong
+
+
+def score_model_static(model):
+    """Static score approximation from value_info + initializers."""
+    try:
+        mi = onnx.shape_inference.infer_shapes(model, strict_mode=False)
+    except:
+        mi = model
+    g = mi.graph
+    ini = {i.name for i in g.initializer}
+    mem = 0
+    for vi in g.value_info:
+        if vi.name in ini:
+            continue
+        if not vi.type.HasField('tensor_type'):
+            continue
+        tt = vi.type.tensor_type
+        if not tt.HasField('shape'):
+            continue
+        n = 1
+        ok = True
+        for d in tt.shape.dim:
+            if d.HasField('dim_value') and d.dim_value > 0:
+                n *= d.dim_value
+            else:
+                ok = False
+                break
+        if not ok:
+            continue
+        dt = onnx.helper.tensor_dtype_to_np_dtype(tt.elem_type)
+        mem += n * np.dtype(dt).itemsize
+    par = 0
+    for i in g.initializer:
+        if any(d <= 0 for d in i.dims):
+            return None
+        par += math.prod(i.dims) if i.dims else 1
+    for nd in g.node:
+        if nd.op_type == 'Constant':
+            for a in nd.attribute:
+                if a.name == 'value':
+                    par += math.prod(a.t.dims) if a.t.dims else 1
+                elif a.name == 'value_floats':
+                    par += len(a.floats)
+                elif a.name == 'value_ints':
+                    par += len(a.ints)
+    if mem + par == 0:
+        return None
+    return max(1.0, 25.0 - math.log(max(1.0, mem + par)))
+
+
+def load_task_data(task_data_dir, task_num):
+    """Load task JSON data."""
+    path = os.path.join(task_data_dir, f'task{task_num:03d}.json')
+    if not os.path.exists(path):
+        return None
+    with open(path) as f:
+        data = json.load(f)
+    if 'arc-gen' not in data:
+        data['arc-gen'] = []
+    return data
+
+
+def build_wave22_models(wave22_path, task_data_dir, output_dir, base_zip_path):
+    """Build wave22 models that improve over base. Returns dict {task_num: path}."""
+    if not os.path.exists(wave22_path):
+        print(f"  WARNING: wave22.py not found at {wave22_path}, skipping")
+        return {}
+
+    with open(wave22_path) as f:
+        wave22_code = f.read()
+
+    wave22_ns = {
+        'np': np, 'zlib': zlib, 'base64': base64,
+        'oh': oh, 'onh': onh, 'TensorProto': TensorProto,
+        '_d': _d
+    }
+    exec(wave22_code, wave22_ns)
+
+    funcs = re.findall(r'def (s_\w+)\(td\):\s*\n\s*"""Task (\d+)', wave22_code)
+    print(f"  Found {len(funcs)} wave22 functions")
+
+    # Score base models
+    base_scores = {}
+    with zipfile.ZipFile(base_zip_path, 'r') as zf:
+        for name in zf.namelist():
+            if name.startswith('task') and name.endswith('.onnx'):
+                try:
+                    tn = int(name[4:7])
+                    data = zf.read(name)
+                    tmp = os.path.join(output_dir, f'_base_{name}')
+                    with open(tmp, 'wb') as f:
+                        f.write(data)
+                    m = onnx.load(tmp)
+                    s = score_model_static(m)
+                    if s:
+                        base_scores[tn] = s
+                    os.remove(tmp)
+                except:
+                    pass
+
+    skip_tasks = {53, 77, 98, 100, 291, 307, 398}
+    models = {}
+    file_limit = 1.44 * 1024 * 1024
+
+    for func_name, task_str in funcs:
+        tn = int(task_str)
+        if tn in skip_tasks:
+            continue
+
+        td = load_task_data(task_data_dir, tn)
+        if td is None:
+            continue
+
+        try:
+            model = wave22_ns[func_name](td)
+            if model is None:
+                continue
+
+            out_path = os.path.join(output_dir, f'task{tn:03d}.onnx')
+            onnx.save(model, out_path)
+            fsize = os.path.getsize(out_path)
+
+            if fsize > file_limit:
+                os.remove(out_path)
+                continue
+
+            wave_score = score_model_static(model)
+            base_score = base_scores.get(tn)
+            if wave_score is None or base_score is None:
+                os.remove(out_path)
+                continue
+
+            if wave_score <= base_score:
+                os.remove(out_path)
+                continue
+
+            right, wrong = validate_model(out_path, task_data_dir, tn)
+            if wrong != 0 or right is None:
+                os.remove(out_path)
+                continue
+
+            gain = wave_score - base_score
+            models[tn] = out_path
+            print(f"    task{tn:03d}: +{gain:.3f} ({base_score:.3f} -> {wave_score:.3f}), "
+                  f"{len(model.graph.node)} nodes, {right}/{right+wrong} PASS")
+
+        except Exception:
+            p = os.path.join(output_dir, f'task{tn:03d}.onnx')
+            if os.path.exists(p):
+                os.remove(p)
+
+    return models
+
+
+def create_submission(base_zip_path, replacements, output_path):
+    """Create submission zip with replacements."""
+    os.makedirs(os.path.dirname(output_path) or '.', exist_ok=True)
+
+    with zipfile.ZipFile(base_zip_path, 'r') as base_zip:
+        with zipfile.ZipFile(output_path, 'w', zipfile.ZIP_DEFLATED) as out_zip:
+            replaced = []
+            for item in base_zip.namelist():
+                basename = os.path.basename(item)
+                if basename.startswith('task') and basename.endswith('.onnx'):
+                    try:
+                        tn = int(basename[4:7])
+                        if tn in replacements:
+                            out_zip.write(replacements[tn], basename)
+                            replaced.append(tn)
+                            continue
+                    except ValueError:
+                        pass
+                data = base_zip.read(item)
+                out_zip.writestr(basename, data)
+
+    print(f"\n  Submission: {output_path}")
+    print(f"  Size: {os.path.getsize(output_path):,} bytes")
+    print(f"  Replaced {len(replaced)} models: {sorted(replaced)}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Build NeuroGolf submission")
+    parser.add_argument('--base', required=True,
+                        help='Path to base submission zip (submission-6043.zip)')
+    parser.add_argument('--task-data-dir', required=True,
+                        help='Directory with taskNNN.json files')
+    parser.add_argument('--wave22', default=None,
+                        help='Path to wave22.py (optional, builds additional models)')
+    parser.add_argument('--optimized-dir', default=None,
+                        help='Directory with pre-built optimized .onnx files')
+    parser.add_argument('--output', default='submission.zip',
+                        help='Output submission zip path')
+    parser.add_argument('--skip-validation', action='store_true',
+                        help='Skip validation (not recommended)')
+    args = parser.parse_args()
+
+    if not os.path.exists(args.base):
+        print(f"ERROR: Base zip not found: {args.base}")
+        sys.exit(1)
+    if not os.path.exists(args.task_data_dir):
+        print(f"ERROR: Task data dir not found: {args.task_data_dir}")
+        sys.exit(1)
+
+    replacements = {}
+    file_limit = 1.44 * 1024 * 1024
+
+    # --- Step 1: Load pre-built optimized models ---
+    if args.optimized_dir and os.path.isdir(args.optimized_dir):
+        print(f"[1] Loading pre-built models from {args.optimized_dir}")
+        for f in sorted(os.listdir(args.optimized_dir)):
+            if f.startswith('task') and f.endswith('.onnx'):
+                tn = int(f[4:7])
+                path = os.path.join(args.optimized_dir, f)
+                fsize = os.path.getsize(path)
+                if fsize > file_limit:
+                    print(f"  SKIP {f}: exceeds size limit ({fsize:,} bytes)")
+                    continue
+                if not args.skip_validation:
+                    right, wrong = validate_model(path, args.task_data_dir, tn)
+                    if wrong != 0 or right is None:
+                        print(f"  SKIP {f}: validation failed ({wrong} wrong)")
+                        continue
+                    print(f"  {f}: {right}/{right+wrong} PASS")
+                else:
+                    print(f"  {f}: loaded (validation skipped)")
+                replacements[tn] = path
+        print(f"  Total from optimized-dir: {len(replacements)}")
+    else:
+        print("[1] No --optimized-dir provided, skipping pre-built models")
+
+    # --- Step 2: Build wave22 models ---
+    if args.wave22:
+        print(f"\n[2] Building wave22 models from {args.wave22}")
+        tmp_dir = os.path.join(os.path.dirname(args.output) or '.', '_wave22_tmp')
+        os.makedirs(tmp_dir, exist_ok=True)
+        wave22_models = build_wave22_models(
+            args.wave22, args.task_data_dir, tmp_dir, args.base)
+        added = 0
+        for tn, path in wave22_models.items():
+            if tn not in replacements:
+                replacements[tn] = path
+                added += 1
+        print(f"  Added from wave22: {added} new models")
+    else:
+        print("\n[2] No --wave22 provided, skipping wave22 extraction")
+
+    # --- Step 3: Create submission ---
+    print(f"\n[3] Creating submission zip")
+    create_submission(args.base, replacements, args.output)
+    print("\nDone!")
+
+
+if __name__ == '__main__':
+    main()