Add merge_best_of_both.py: picks lowest-cost model from 5743 vs V90 per task (42 tasks improved, 55% size reduction)"

own-solver/merge_best_of_both.py

@@ -0,0 +1,200 @@
+#!/usr/bin/env python3
+"""
+Best-of-both merger: Pick the better model from 5743 and V90 for each task.
+
+Under new formula (score = 25 - ln(memory + params)):
+- memory = sum of ALL intermediate tensor bytes
+- Each node output that isn't 'output' costs memory
+- FEWER NODES = LESS MEMORY = HIGHER SCORE
+
+Strategy: For each task, pick the model with fewer intermediate tensors.
+If tied, pick the smaller file (fewer params/weight bytes).
+Validates each candidate against train+test+arc-gen before accepting.
+
+Usage:
+  python merge_best_of_both.py \
+    --sub_5743 submission-5743.zip \
+    --sub_v90 submission-6043.zip \
+    --data_dir ./tasks \
+    --output_zip submission_merged.zip
+"""
+
+import json
+import math
+import os
+import zipfile
+
+import numpy as np
+import onnx
+import onnxruntime as ort
+from onnx import helper, TensorProto, numpy_helper
+
+
+def encode_grid(grid):
+    arr = np.array(grid, dtype=np.int32)
+    h, w = arr.shape
+    t = np.zeros((1, 10, 30, 30), dtype=np.float32)
+    for r in range(h):
+        for c in range(w):
+            v = int(arr[r, c])
+            if 0 <= v < 10:
+                t[0, v, r, c] = 1.0
+    return t
+
+
+def validate_model(model_bytes, examples, max_check=30):
+    """Validate model produces correct output on examples."""
+    try:
+        opts = ort.SessionOptions()
+        opts.log_severity_level = 3
+        sess = ort.InferenceSession(model_bytes, sess_options=opts, providers=['CPUExecutionProvider'])
+    except Exception:
+        return False
+    for ex in examples[:max_check]:
+        try:
+            inp = encode_grid(ex['input'])
+            out = sess.run(['output'], {'input': inp})[0]
+            expected = encode_grid(ex['output'])
+            if not np.array_equal((out > 0.0).astype(np.float32), expected):
+                return False
+        except Exception:
+            return False
+    return True
+
+
+def count_intermediates(model_bytes):
+    """Count intermediate tensors (proxy for runtime memory cost)."""
+    try:
+        model = onnx.load_from_string(model_bytes)
+        count = 0
+        for node in model.graph.node:
+            for out in node.output:
+                if out and out != 'output':
+                    count += 1
+        return count
+    except Exception:
+        return 999999
+
+
+def estimate_cost(model_bytes):
+    """Estimate cost under new formula: memory + params."""
+    try:
+        model = onnx.load_from_string(model_bytes)
+    except Exception:
+        return float('inf')
+    
+    weight_memory = 0
+    params = 0
+    for init in model.graph.initializer:
+        arr = numpy_helper.to_array(init)
+        weight_memory += arr.nbytes
+        params += arr.size
+    for node in model.graph.node:
+        if node.op_type == 'Constant':
+            for attr in node.attribute:
+                if attr.name == 'value' and attr.t.ByteSize() > 0:
+                    try:
+                        arr = numpy_helper.to_array(attr.t)
+                        weight_memory += arr.nbytes
+                        params += arr.size
+                    except:
+                        params += 1
+
+    intermediates = 0
+    for node in model.graph.node:
+        for out in node.output:
+            if out and out != 'output':
+                intermediates += 1
+    
+    intermediate_memory = intermediates * 20000  # avg estimate
+    cost = weight_memory + intermediate_memory + params
+    return cost
+
+
+def main():
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--sub_5743', required=True, help='5743 submission zip')
+    parser.add_argument('--sub_v90', required=True, help='V90 submission zip')
+    parser.add_argument('--data_dir', required=True, help='Task JSON directory')
+    parser.add_argument('--output_zip', required=True, help='Output merged zip')
+    args = parser.parse_args()
+    
+    models_5743 = {}
+    models_v90 = {}
+    
+    with zipfile.ZipFile(args.sub_5743, 'r') as zf:
+        for tid in range(1, 401):
+            fname = f'task{tid:03d}.onnx'
+            if fname in zf.namelist():
+                models_5743[tid] = zf.read(fname)
+    
+    with zipfile.ZipFile(args.sub_v90, 'r') as zf:
+        for tid in range(1, 401):
+            fname = f'task{tid:03d}.onnx'
+            if fname in zf.namelist():
+                models_v90[tid] = zf.read(fname)
+    
+    print(f"Loaded {len(models_5743)} from 5743, {len(models_v90)} from V90")
+    
+    merged = {}
+    stats = {'5743': 0, 'v90': 0, 'validated_5743': 0}
+    
+    for tid in range(1, 401):
+        b5 = models_5743.get(tid)
+        bv = models_v90.get(tid)
+        
+        if not b5 and not bv:
+            continue
+        if not b5:
+            merged[tid] = bv
+            stats['v90'] += 1
+            continue
+        if not bv:
+            merged[tid] = b5
+            stats['5743'] += 1
+            continue
+        
+        cost5 = estimate_cost(b5)
+        costv = estimate_cost(bv)
+        
+        if cost5 < costv:
+            task_path = os.path.join(args.data_dir, f'task{tid:03d}.json')
+            if os.path.exists(task_path):
+                with open(task_path) as f:
+                    task_data = json.load(f)
+                examples = task_data.get('train', []) + task_data.get('test', [])
+                arcgen = task_data.get('arc-gen', [])[:30]
+                all_ex = examples + arcgen
+                
+                if validate_model(b5, all_ex):
+                    merged[tid] = b5
+                    stats['validated_5743'] += 1
+                    inter5 = count_intermediates(b5)
+                    interv = count_intermediates(bv)
+                    if interv - inter5 > 5:
+                        print(f"  Task {tid:3d}: USE 5743 ({inter5} vs {interv} intermediates)")
+                    continue
+            
+            merged[tid] = bv
+            stats['v90'] += 1
+        else:
+            merged[tid] = bv
+            stats['v90'] += 1
+    
+    print(f"\nUsing 5743: {stats['5743']} + {stats['validated_5743']} validated")
+    print(f"Using V90: {stats['v90']}")
+    print(f"Total: {len(merged)}")
+    
+    with zipfile.ZipFile(args.output_zip, 'w', zipfile.ZIP_DEFLATED) as zf:
+        for tid in range(1, 401):
+            fname = f'task{tid:03d}.onnx'
+            if tid in merged:
+                zf.writestr(fname, merged[tid])
+    
+    total_size = sum(len(v) for v in merged.values())
+    print(f"Output: {args.output_zip} ({total_size:,} bytes)")
+
+
+if __name__ == '__main__':
+    main()