#!/usr/bin/env python3 """ swap_and_submit.py — Replace task models in submission-6043.zip with optimized versions. Usage: python swap_and_submit.py [--base submission-6043.zip] [--output submission.zip] [--models task285.onnx ...] If no --models specified, it looks for optimized .onnx files in ./optimized/ directory. Prerequisites: - submission-6043.zip in the repo (or specify with --base) - Optimized .onnx files named taskNNN.onnx - pip install onnx onnxruntime numpy The script will: 1. Extract submission-6043.zip 2. Replace specified task files with optimized versions 3. Validate each replacement with neurogolf_utils.py logic (if task data available) 4. Create new submission.zip ready for Kaggle upload """ import argparse import os import sys import zipfile import json import math def main(): parser = argparse.ArgumentParser(description="Swap optimized models into submission zip") parser.add_argument("--base", default="submission-6043.zip", help="Base submission zip") parser.add_argument("--output", default="/kaggle/working/submission.zip", help="Output submission zip") parser.add_argument("--models", nargs="+", help="Optimized .onnx files to swap in (e.g. task285.onnx)") parser.add_argument("--optimized-dir", default="optimized", help="Directory with optimized models") parser.add_argument("--task-data-dir", default=None, help="Directory with taskNNN.json files for validation") parser.add_argument("--skip-validation", action="store_true", help="Skip validation (not recommended)") parser.add_argument("--validate-only", action="store_true", help="Only validate, don't create zip") args = parser.parse_args() # Find optimized models models_to_swap = {} if args.models: for m in args.models: if not os.path.exists(m): print(f"ERROR: Model file not found: {m}") sys.exit(1) basename = os.path.basename(m) if basename.startswith("task") and basename.endswith(".onnx"): task_num = int(basename[4:7]) models_to_swap[task_num] = m else: print(f"ERROR: Model filename must be taskNNN.onnx, got: {basename}") sys.exit(1) elif os.path.isdir(args.optimized_dir): for f in sorted(os.listdir(args.optimized_dir)): if f.startswith("task") and f.endswith(".onnx"): task_num = int(f[4:7]) models_to_swap[task_num] = os.path.join(args.optimized_dir, f) if not models_to_swap: print("ERROR: No optimized models found.") print(f" Either pass --models task285.onnx or place files in ./{args.optimized_dir}/") sys.exit(1) print(f"Found {len(models_to_swap)} models to swap: {sorted(models_to_swap.keys())}") for tn in sorted(models_to_swap.keys()): path = models_to_swap[tn] size = os.path.getsize(path) print(f" task{tn:03d}.onnx ({size:,} bytes) <- {path}") # Validate each model (if task data available) if not args.skip_validation: try: import onnx import onnxruntime import numpy as np except ImportError: print("WARNING: onnx/onnxruntime not installed, skipping validation") args.skip_validation = True if not args.skip_validation: task_data_dir = args.task_data_dir if task_data_dir is None: for candidate in ["task-data", ".", "/kaggle/input/competitions/neurogolf-2026"]: if os.path.exists(candidate) and os.path.exists(os.path.join(candidate, "task001.json")): task_data_dir = candidate break if task_data_dir: print(f"\nValidating with task data from: {task_data_dir}") all_pass = True for task_num, model_path in sorted(models_to_swap.items()): task_file = os.path.join(task_data_dir, f"task{task_num:03d}.json") if not os.path.exists(task_file): print(f" Task {task_num}: SKIP (no task data at {task_file})") continue ok = validate_model(model_path, task_file, task_num) if not ok: all_pass = False print(f" Task {task_num}: FAIL -- model does NOT pass all examples!") print(f" DO NOT SUBMIT -- this will score WORSE than base submission") else: print(f" Task {task_num}: PASS (all train+test+arc-gen)") if not all_pass: print("\nERROR: One or more models failed validation.") print("Fix the models or use --skip-validation to override (NOT recommended).") if not args.validate_only: sys.exit(1) else: print("WARNING: No task data found for validation. Use --task-data-dir to specify.") if args.validate_only: return # Check base zip exists if not os.path.exists(args.base): print(f"ERROR: Base submission not found: {args.base}") sys.exit(1) # Check file sizes (Kaggle limit: 1.44 * 1024 * 1024 bytes per model) LIMIT = 1.44 * 1024 * 1024 for task_num, model_path in models_to_swap.items(): size = os.path.getsize(model_path) if size > LIMIT: print(f"ERROR: Task {task_num} model exceeds size limit: {size:,} > {int(LIMIT):,} bytes") sys.exit(1) # Create output zip print(f"\nCreating {args.output}...") replaced_tasks = set() with zipfile.ZipFile(args.base, 'r') as base_zip: base_names = set(base_zip.namelist()) print(f" Base zip contains {len(base_names)} files") # Pre-check: verify all target tasks exist in base zip for tn in sorted(models_to_swap.keys()): expected_name = f"task{tn:03d}.onnx" if expected_name not in base_names: print(f" WARNING: {expected_name} not found in base zip!") with zipfile.ZipFile(args.output, 'w', zipfile.ZIP_DEFLATED) as out_zip: for item in base_zip.namelist(): basename = os.path.basename(item) replaced = False if basename.startswith("task") and basename.endswith(".onnx"): try: task_num = int(basename[4:7]) if task_num in models_to_swap: out_zip.write(models_to_swap[task_num], basename) replaced = True replaced_tasks.add(task_num) except ValueError: pass if not replaced: data = base_zip.read(item) out_zip.writestr(item, data) # Final summary with clear pass/fail print(f"\n{'='*60}") print(f" REPLACEMENT SUMMARY") print(f"{'='*60}") for tn in sorted(models_to_swap.keys()): status = "REPLACED" if tn in replaced_tasks else "MISSING!" marker = "+" if tn in replaced_tasks else "X" print(f" [{marker}] task{tn:03d}.onnx: {status}") missing = set(models_to_swap.keys()) - replaced_tasks if missing: print(f"\n ERROR: {len(missing)} models were NOT replaced!") print(f" These tasks are missing from the base zip or have naming issues:") for tn in sorted(missing): print(f" task{tn:03d}.onnx -- source was: {models_to_swap[tn]}") else: print(f"\n All {len(replaced_tasks)} models successfully replaced.") out_size = os.path.getsize(args.output) print(f" Output: {args.output} ({out_size:,} bytes)") print(f"{'='*60}") def validate_model(model_path, task_file, task_num): """Validate model against all task examples. Returns True if all pass.""" import onnx import onnxruntime import numpy as np try: with open(task_file, 'r') as f: task_data = json.load(f) session = onnxruntime.InferenceSession(model_path) all_examples = task_data.get("train", []) + task_data.get("test", []) + task_data.get("arc-gen", []) right, wrong = 0, 0 for example in all_examples: benchmark_input = np.zeros((1, 10, 30, 30), dtype=np.float32) grid = example["input"] if max(len(grid), len(grid[0]) if grid else 0) > 30: continue for r, row in enumerate(grid): for c, color in enumerate(row): benchmark_input[0][color][r][c] = 1.0 benchmark_output = np.zeros((1, 10, 30, 30), dtype=np.float32) grid_out = example["output"] for r, row in enumerate(grid_out): for c, color in enumerate(row): if r < 30 and c < 30: benchmark_output[0][color][r][c] = 1.0 try: result = session.run(["output"], {"input": benchmark_input}) user_output = (result[0] > 0.0).astype(float) if np.array_equal(user_output, benchmark_output): right += 1 else: wrong += 1 except Exception: wrong += 1 if wrong == 0 and right > 0: return True else: print(f" {right} pass, {wrong} fail (out of {right + wrong})") return False except Exception as e: print(f" Validation error: {e}") return False if __name__ == "__main__": main()