Add task277 ONNX build script (266/266 PASS, score 14.85 vs base 13.30, gain +1.55)

medal-solvers/build_task277_onnx.py

@@ -0,0 +1,195 @@
+"""Build ONNX model for Task 277 — Smallest CC recoloring. V2 compact.
+
+Key facts: Always exactly 3 CCs, gap≥2, no ties.
+
+Approach:
+1. MaxPool label propagation → 3 unique labels
+2. Find 3 labels: max, second max, third max
+3. Count cells per label
+4. Min count → color 2, others → color 1
+"""
+import numpy as np
+import onnx
+from onnx import helper, numpy_helper, TensorProto
+import math
+import os
+
+
+def build_task277():
+    nodes, inits, vis = [], [], []
+    counter = [0]
+
+    def nm():
+        counter[0] += 1
+        return f"t{counter[0]}"
+
+    def const(name, val, dtype='f'):
+        arr = np.array(val, dtype=np.float32 if dtype == 'f' else np.int64)
+        inits.append(numpy_helper.from_array(arr, name))
+
+    def vi(name, shape, dt=TensorProto.FLOAT):
+        vis.append(helper.make_tensor_value_info(name, dt, shape))
+
+    def nd(op, ins, outs_shapes, **kwargs):
+        out_names = []
+        for sd in outs_shapes:
+            if isinstance(sd, tuple):
+                shape, dt = sd
+            else:
+                shape, dt = sd, TensorProto.FLOAT
+            n = nm()
+            vi(n, shape, dt)
+            out_names.append(n)
+        nodes.append(helper.make_node(op, ins, out_names, **kwargs))
+        return out_names[0] if len(out_names) == 1 else out_names
+
+    # === CONSTANTS ===
+    const('c_one', [1.0])
+    const('c_zero', [0.0])
+    const('c_half', [0.5])
+    const('c_big', [200.0])
+
+    const('sl_start', [0, 0, 0, 0], 'i')
+    const('sl_end', [1, 10, 10, 10], 'i')
+    const('sl_axes', [0, 1, 2, 3], 'i')
+    const('sl_ch8_start', [0, 8, 0, 0], 'i')
+    const('sl_ch8_end', [1, 9, 10, 10], 'i')
+
+    # Label grid: unique per cell (1-100)
+    const('label_init', np.arange(1, 101, dtype=np.float32).reshape(1, 1, 10, 10))
+    const('axes23', [2, 3], 'i')
+    
+    const('pad_10_to_30', [0, 0, 0, 0, 0, 0, 20, 20], 'i')
+    const('pad_val_zero', [0.0])
+    const('zeros_7ch', np.zeros((1, 7, 10, 10), dtype=np.float32))
+
+    # === STEP 1: Get mask ===
+    inp10 = nd('Slice', ['input', 'sl_start', 'sl_end', 'sl_axes'], [[1, 10, 10, 10]])
+    mask8 = nd('Slice', [inp10, 'sl_ch8_start', 'sl_ch8_end', 'sl_axes'], [[1, 1, 10, 10]])
+
+    # === STEP 2: Label propagation ===
+    labels = nd('Mul', [mask8, 'label_init'], [[1, 1, 10, 10]])
+    for _ in range(10):
+        pooled = nd('MaxPool', [labels], [[1, 1, 10, 10]], kernel_shape=[3, 3], pads=[1, 1, 1, 1])
+        labels = nd('Mul', [pooled, mask8], [[1, 1, 10, 10]])
+
+    # === STEP 3: Find 3 unique labels ===
+    # L1 = max of all labels
+    L1 = nd('ReduceMax', [labels, 'axes23'], [[1, 1, 1, 1]], keepdims=1)
+    
+    # Mask out L1: set cells with label==L1 to 0
+    diff_L1 = nd('Abs', [nd('Sub', [labels, L1], [[1, 1, 10, 10]])], [[1, 1, 10, 10]])
+    is_L1_b = nd('Less', [diff_L1, 'c_half'], [([1, 1, 10, 10], TensorProto.BOOL)])
+    is_L1 = nd('Cast', [is_L1_b], [[1, 1, 10, 10]], to=1)
+    not_L1 = nd('Sub', ['c_one', is_L1], [[1, 1, 10, 10]])
+    labels_no_L1 = nd('Mul', [labels, not_L1], [[1, 1, 10, 10]])
+    
+    # L2 = max of remaining labels
+    L2 = nd('ReduceMax', [labels_no_L1, 'axes23'], [[1, 1, 1, 1]], keepdims=1)
+    
+    # Mask out L2
+    diff_L2 = nd('Abs', [nd('Sub', [labels, L2], [[1, 1, 10, 10]])], [[1, 1, 10, 10]])
+    is_L2_b = nd('Less', [diff_L2, 'c_half'], [([1, 1, 10, 10], TensorProto.BOOL)])
+    is_L2 = nd('Cast', [is_L2_b], [[1, 1, 10, 10]], to=1)
+    not_L2 = nd('Sub', ['c_one', is_L2], [[1, 1, 10, 10]])
+    labels_no_L1_L2 = nd('Mul', [labels_no_L1, not_L2], [[1, 1, 10, 10]])
+    
+    # L3 mask is just: mask8 - is_L1 - is_L2
+    is_L3 = nd('Sub', [mask8, nd('Add', [is_L1, is_L2], [[1, 1, 10, 10]])], [[1, 1, 10, 10]])
+
+    # === STEP 4: Count cells per CC ===
+    count_L1 = nd('ReduceSum', [is_L1, 'axes23'], [[1, 1, 1, 1]], keepdims=1)
+    count_L2 = nd('ReduceSum', [is_L2, 'axes23'], [[1, 1, 1, 1]], keepdims=1)
+    count_L3 = nd('ReduceSum', [is_L3, 'axes23'], [[1, 1, 1, 1]], keepdims=1)
+
+    # === STEP 5: Find minimum count ===
+    min_12 = nd('Min', [count_L1, count_L2], [[1, 1, 1, 1]])
+    min_count = nd('Min', [min_12, count_L3], [[1, 1, 1, 1]])
+
+    # Which CC has this count? (no ties guaranteed)
+    diff_c1 = nd('Abs', [nd('Sub', [count_L1, min_count], [[1, 1, 1, 1]])], [[1, 1, 1, 1]])
+    smallest_is_L1_b = nd('Less', [diff_c1, 'c_half'], [([1, 1, 1, 1], TensorProto.BOOL)])
+    smallest_is_L1 = nd('Cast', [smallest_is_L1_b], [[1, 1, 1, 1]], to=1)
+    
+    diff_c2 = nd('Abs', [nd('Sub', [count_L2, min_count], [[1, 1, 1, 1]])], [[1, 1, 1, 1]])
+    smallest_is_L2_b = nd('Less', [diff_c2, 'c_half'], [([1, 1, 1, 1], TensorProto.BOOL)])
+    smallest_is_L2 = nd('Cast', [smallest_is_L2_b], [[1, 1, 1, 1]], to=1)
+    
+    diff_c3 = nd('Abs', [nd('Sub', [count_L3, min_count], [[1, 1, 1, 1]])], [[1, 1, 1, 1]])
+    smallest_is_L3_b = nd('Less', [diff_c3, 'c_half'], [([1, 1, 1, 1], TensorProto.BOOL)])
+    smallest_is_L3 = nd('Cast', [smallest_is_L3_b], [[1, 1, 1, 1]], to=1)
+
+    # === STEP 6: Build smallest mask and other mask ===
+    sm1 = nd('Mul', [is_L1, smallest_is_L1], [[1, 1, 10, 10]])
+    sm2 = nd('Mul', [is_L2, smallest_is_L2], [[1, 1, 10, 10]])
+    sm3 = nd('Mul', [is_L3, smallest_is_L3], [[1, 1, 10, 10]])
+    smallest_mask = nd('Add', [sm1, nd('Add', [sm2, sm3], [[1, 1, 10, 10]])], [[1, 1, 10, 10]])
+    other_mask = nd('Sub', [mask8, smallest_mask], [[1, 1, 10, 10]])
+
+    # === STEP 7: Build output ===
+    bg_mask = nd('Sub', ['c_one', mask8], [[1, 1, 10, 10]])
+    out_10 = nd('Concat', [bg_mask, other_mask, smallest_mask, 'zeros_7ch'], [[1, 10, 10, 10]], axis=1)
+
+    # === STEP 8: Pad — output directly ===
+    nodes.append(helper.make_node('Pad', [out_10, 'pad_10_to_30', 'pad_val_zero'], ['output']))
+
+    x = helper.make_tensor_value_info('input', TensorProto.FLOAT, [1, 10, 30, 30])
+    y = helper.make_tensor_value_info('output', TensorProto.FLOAT, [1, 10, 30, 30])
+
+    graph = helper.make_graph(nodes, 'task277', [x], [y], initializer=inits, value_info=vis)
+    model = helper.make_model(graph, ir_version=10, opset_imports=[helper.make_opsetid('', 18)])
+    return model
+
+
+if __name__ == '__main__':
+    import onnxruntime as ort
+    import json
+
+    print("Building Task 277 ONNX model v2...")
+    model = build_task277()
+
+    output_path = '/app/task277.onnx'
+    del model.graph.value_info[:]
+    model = onnx.shape_inference.infer_shapes(model, strict_mode=True)
+    onnx.save(model, output_path)
+
+    print(f"  Nodes: {len(model.graph.node)}")
+    print(f"  File size: {os.path.getsize(output_path):,} bytes")
+
+    sess = ort.InferenceSession(output_path)
+    with open('/app/task-data/task277.json') as f:
+        data = json.load(f)
+
+    all_examples = data['train'] + data['test'] + data.get('arc-gen', [])
+    right_count, wrong_count = 0, 0
+    for i, ex in enumerate(all_examples):
+        inp = np.zeros((1, 10, 30, 30), dtype=np.float32)
+        for r, row in enumerate(ex['input']):
+            for c, v in enumerate(row):
+                if r < 30 and c < 30:
+                    inp[0][v][r][c] = 1.0
+        result = sess.run(['output'], {'input': inp})
+        out = (result[0] > 0.0).astype(float)
+        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
+        if np.array_equal(out, exp):
+            right_count += 1
+        else:
+            wrong_count += 1
+            if wrong_count <= 5:
+                print(f"  FAIL {i}: {len(np.where(out != exp)[0])} diffs")
+
+    print(f"  Results: {right_count}/{right_count+wrong_count} pass ({wrong_count} fail)")
+
+    if wrong_count == 0:
+        params = sum(int(np.prod(init.dims)) for init in model.graph.initializer)
+        mem_est = sum(int(np.prod([d.dim_value for d in vi.type.tensor_type.shape.dim])) *
+                      np.dtype(onnx.helper.tensor_dtype_to_np_dtype(vi.type.tensor_type.elem_type)).itemsize
+                      for vi in model.graph.value_info
+                      if vi.type.HasField('tensor_type') and vi.type.tensor_type.HasField('shape'))
+        score_est = max(1.0, 25.0 - math.log(max(1.0, mem_est + params)))
+        print(f"  Params: {params:,}, Memory: {mem_est:,}, Score: {score_est:.3f}")
+        print(f"  Base: 13.30. Gain: {score_est - 13.30:+.3f}")