#!/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()