medal-solvers/swap_and_submit.py

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