Add build script for task028 optimized ONNX

medal-solvers/build_task028_onnx.py

@@ -0,0 +1,325 @@
+"""Build optimized ONNX model for Task 028.
+
+Rule (265/265 verified):
+1. Input is 10x10 with exactly 2 colored dots on a bg of 0
+2. Upper dot (closer to row 0) → fills upper zone with its color
+3. Lower dot (closer to row 9) → fills lower zone with its color  
+4. Each zone pattern: full row at dot position, full row at border (0 or 9),
+   columns 0 and 9 filled in between. Interior stays 0.
+5. Boundary between zones is midpoint of the two dot rows.
+
+ONNX approach:
+- Slice to 10x10
+- Detect non-bg channels, find dot row positions via ArgMax on row sums
+- Generate masks using row/col grids and comparisons
+- Combine with color one-hots
+- Pad back to 30x30
+
+Base: 10 nodes, 12610 params, score 12.83
+Target: ~40 nodes, ~1500 params → score ~15+ (gain +2.5)
+"""
+import sys, os
+sys.path.insert(0, '/app/repo/medal-solvers')
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+from onnx import TensorProto
+import numpy as np
+import onnx
+from onnx import helper, numpy_helper
+import math
+
+
+def build_task028():
+    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_half', [0.5])
+    const('c_one', [1.0])
+    const('c_zero', [0.0])
+    const('c_two', [2.0])
+    const('c_9f', [9.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('axes3', [3], 'i')  # cols
+    const('axes23', [2, 3], 'i')
+    const('axes1', [1], 'i')
+    
+    # Row grid [1,1,10,1] - row index values
+    const('row_grid', np.arange(10, dtype=np.float32).reshape(1, 1, 10, 1))
+    # Col grid [1,1,1,10]
+    const('col_grid', np.arange(10, dtype=np.float32).reshape(1, 1, 1, 10))
+    
+    const('shape_1_1_1_1', [1, 1, 1, 1], 'i')
+    const('shape_1_10_1_1', [1, 10, 1, 1], 'i')
+    
+    const('depth_10', [10.0])
+    const('oh_vals', [0.0, 1.0])
+    
+    # Pad to get back to 30x30
+    const('pad_10_to_30', [0, 0, 0, 0, 0, 0, 20, 20], 'i')
+    const('pad_val_zero', [0.0])
+
+    # Indices
+    const('idx_0', [0], 'i')
+    const('idx_1', [1], 'i')
+    const('idx_bg', [0], 'i')
+    
+    # Slice to get non-bg channels [1, 9, 10, 10]
+    const('sl_ch_start', [0, 1, 0, 0], 'i')
+    const('sl_ch_end', [1, 10, 10, 10], 'i')
+
+    # === STEP 1: Slice input to 10x10 ===
+    inp10 = nd('Slice', ['input', 'sl_start', 'sl_end', 'sl_axes'], [[1, 10, 10, 10]])
+
+    # === STEP 2: Find the two dots ===
+    # Sum non-bg channels to get overall non-zero mask [1,1,10,10]
+    nonbg = nd('Slice', [inp10, 'sl_ch_start', 'sl_ch_end', 'sl_axes'], [[1, 9, 10, 10]])
+    active = nd('ReduceSum', [nonbg, 'axes1'], [[1, 1, 10, 10]], keepdims=1)
+    
+    # Sum per row to find which rows have dots [1,1,10,1]
+    row_sums = nd('ReduceSum', [active, 'axes3'], [[1, 1, 10, 1]], keepdims=1)
+    
+    # Find the two dot rows: first dot = ArgMax of row_sums
+    # But we have 2 dots. Use: first non-zero row from top = upper dot
+    # For upper dot: row_sums * row_grid gives weighted; we want min non-zero row
+    # Trick: add large value where row_sum==0, then ArgMin
+    const('c_big', [100.0])
+    has_dot_b = nd('Greater', [row_sums, 'c_half'], [([1, 1, 10, 1], TensorProto.BOOL)])
+    has_dot = nd('Cast', [has_dot_b], [[1, 1, 10, 1]], to=1)
+    no_dot = nd('Sub', ['c_one', has_dot], [[1, 1, 10, 1]])
+    
+    # For upper dot (minimum row with a dot):
+    row_masked_upper = nd('Add', [nd('Mul', ['row_grid', has_dot], [[1, 1, 10, 1]]),
+                                   nd('Mul', ['c_big', no_dot], [[1, 1, 10, 1]])], [[1, 1, 10, 1]])
+    const('axes2', [2], 'i')
+    upper_row_idx = nd('ArgMin', [row_masked_upper], [([1, 1, 1, 1], TensorProto.INT64)], axis=2, keepdims=1)
+    upper_row_f = nd('Cast', [upper_row_idx], [[1, 1, 1, 1]], to=1)
+    
+    # For lower dot (maximum row with a dot):
+    row_masked_lower = nd('Mul', ['row_grid', has_dot], [[1, 1, 10, 1]])
+    lower_row_idx = nd('ArgMax', [row_masked_lower], [([1, 1, 1, 1], TensorProto.INT64)], axis=2, keepdims=1)
+    lower_row_f = nd('Cast', [lower_row_idx], [[1, 1, 1, 1]], to=1)
+    
+    # Midpoint between dots
+    sum_rows = nd('Add', [upper_row_f, lower_row_f], [[1, 1, 1, 1]])
+    mid_row = nd('Div', [sum_rows, 'c_two'], [[1, 1, 1, 1]])
+    
+    # === STEP 3: Find dot colors ===
+    # Channel sums per row: for upper dot's row, which channel has the dot?
+    # Gather the input at upper_row → [1, 10, 1, 10] then sum over cols → [1, 10, 1, 1]
+    # Simpler: total channel sums weighted by position
+    
+    # Upper dot color: extract channels at upper dot row
+    # Use: for each channel, check if it has a 1 at the upper dot row
+    # channel_at_upper[c] = sum over cols of inp10[0, c, upper_row, :]
+    # Since there's exactly one dot per row, the channel with max sum IS the color
+    
+    # Sum each channel over all spatial positions → [1, 10, 1, 1]
+    # This gives total pixel count per channel. But we have 2 dots of possibly different colors.
+    # Instead: create mask for upper zone and lower zone, then find color per zone.
+    
+    # Upper zone mask: rows <= mid_row [1,1,10,1]
+    in_upper_b = nd('Less', ['row_grid', nd('Add', [mid_row, 'c_half'], [[1, 1, 1, 1]])],
+                    [([1, 1, 10, 1], TensorProto.BOOL)])
+    in_upper = nd('Cast', [in_upper_b], [[1, 1, 10, 1]], to=1)  # [1,1,10,1]
+    in_lower = nd('Sub', ['c_one', in_upper], [[1, 1, 10, 1]])
+    
+    # Upper color: which channel has pixels in the upper zone?
+    # Mask input by upper zone: inp10 * in_upper [1,10,10,10] * [1,1,10,1] → [1,10,10,10]
+    upper_masked = nd('Mul', [inp10, in_upper], [[1, 10, 10, 10]])
+    upper_ch_sums = nd('ReduceSum', [upper_masked, 'axes23'], [[1, 10, 1, 1]], keepdims=1)
+    # Skip channel 0 (bg)
+    const('bg_mask_10', np.array([[[[0, 1, 1, 1, 1, 1, 1, 1, 1, 1]]]], dtype=np.float32).reshape(1, 10, 1, 1))
+    upper_ch_masked = nd('Mul', [upper_ch_sums, 'bg_mask_10'], [[1, 10, 1, 1]])
+    upper_color_idx = nd('ArgMax', [upper_ch_masked], [([1, 1, 1, 1], TensorProto.INT64)], axis=1, keepdims=1)
+    upper_color_1d = nd('Reshape', [upper_color_idx, 'idx_1'], [([1], TensorProto.INT64)])
+    
+    # Lower color
+    lower_masked = nd('Mul', [inp10, in_lower], [[1, 10, 10, 10]])
+    lower_ch_sums = nd('ReduceSum', [lower_masked, 'axes23'], [[1, 10, 1, 1]], keepdims=1)
+    lower_ch_masked = nd('Mul', [lower_ch_sums, 'bg_mask_10'], [[1, 10, 1, 1]])
+    lower_color_idx = nd('ArgMax', [lower_ch_masked], [([1, 1, 1, 1], TensorProto.INT64)], axis=1, keepdims=1)
+    lower_color_1d = nd('Reshape', [lower_color_idx, 'idx_1'], [([1], TensorProto.INT64)])
+    
+    # === STEP 4: Build output masks ===
+    # Upper zone pattern [1,1,10,10]:
+    # Full row at row 0, full row at upper_dot_row, cols 0&9 in between
+    const('c_0f', [0.0])
+    
+    # Row == 0
+    at_row0_b = nd('Less', ['row_grid', 'c_half'], [([1, 1, 10, 1], TensorProto.BOOL)])
+    at_row0 = nd('Cast', [at_row0_b], [[1, 1, 10, 1]], to=1)
+    
+    # Row == upper_dot_row
+    diff_upper = nd('Abs', [nd('Sub', ['row_grid', upper_row_f], [[1, 1, 10, 1]])], [[1, 1, 10, 1]])
+    at_upper_b = nd('Less', [diff_upper, 'c_half'], [([1, 1, 10, 1], TensorProto.BOOL)])
+    at_upper = nd('Cast', [at_upper_b], [[1, 1, 10, 1]], to=1)
+    
+    # Full rows in upper zone: row 0 OR dot row
+    full_upper_rows = nd('Max', [at_row0, at_upper], [[1, 1, 10, 1]])
+    
+    # Edge cols: col==0 or col==9
+    const('c_8_5', [8.5])
+    at_col0_b = nd('Less', ['col_grid', 'c_half'], [([1, 1, 1, 10], TensorProto.BOOL)])
+    at_col0 = nd('Cast', [at_col0_b], [[1, 1, 1, 10]], to=1)
+    at_col9_b = nd('Greater', ['col_grid', 'c_8_5'], [([1, 1, 1, 10], TensorProto.BOOL)])
+    at_col9 = nd('Cast', [at_col9_b], [[1, 1, 1, 10]], to=1)
+    edge_cols = nd('Max', [at_col0, at_col9], [[1, 1, 1, 10]])
+    
+    # Upper mask: (full_rows broadcast to 10x10) OR (edge_cols * in_upper_zone)
+    # full_rows [1,1,10,1] → broadcast with ones [1,1,1,10] → [1,1,10,10]
+    # But in ONNX: Max already broadcasts
+    upper_full = full_upper_rows  # [1,1,10,1] will broadcast
+    upper_edges = nd('Mul', [edge_cols, in_upper], [[1, 1, 10, 10]])  # [1,1,1,10]*[1,1,10,1]→[1,1,10,10]
+    
+    upper_mask_raw = nd('Max', [upper_full, upper_edges], [[1, 1, 10, 10]])
+    upper_mask = nd('Mul', [upper_mask_raw, in_upper], [[1, 1, 10, 10]])  # clip to upper zone
+    
+    # Lower zone pattern:
+    # Full row at row 9, full row at lower_dot_row, cols 0&9 in between
+    at_row9_b = nd('Greater', ['row_grid', 'c_8_5'], [([1, 1, 10, 1], TensorProto.BOOL)])
+    at_row9 = nd('Cast', [at_row9_b], [[1, 1, 10, 1]], to=1)
+    
+    diff_lower = nd('Abs', [nd('Sub', ['row_grid', lower_row_f], [[1, 1, 10, 1]])], [[1, 1, 10, 1]])
+    at_lower_b = nd('Less', [diff_lower, 'c_half'], [([1, 1, 10, 1], TensorProto.BOOL)])
+    at_lower = nd('Cast', [at_lower_b], [[1, 1, 10, 1]], to=1)
+    
+    full_lower_rows = nd('Max', [at_row9, at_lower], [[1, 1, 10, 1]])
+    
+    lower_full = full_lower_rows
+    lower_edges = nd('Mul', [edge_cols, in_lower], [[1, 1, 10, 10]])
+    
+    lower_mask_raw = nd('Max', [lower_full, lower_edges], [[1, 1, 10, 10]])
+    lower_mask = nd('Mul', [lower_mask_raw, in_lower], [[1, 1, 10, 10]])
+    
+    # === STEP 5: Apply colors ===
+    # Upper color one-hot [1, 10, 1, 1]
+    upper_oh = nd('OneHot', [upper_color_1d, 'depth_10', 'oh_vals'], [[1, 10]], axis=1)
+    upper_oh_4d = nd('Reshape', [upper_oh, 'shape_1_10_1_1'], [[1, 10, 1, 1]])
+    upper_out = nd('Mul', [upper_oh_4d, upper_mask], [[1, 10, 10, 10]])
+    
+    # Lower color one-hot [1, 10, 1, 1]
+    lower_oh = nd('OneHot', [lower_color_1d, 'depth_10', 'oh_vals'], [[1, 10]], axis=1)
+    lower_oh_4d = nd('Reshape', [lower_oh, 'shape_1_10_1_1'], [[1, 10, 1, 1]])
+    lower_out = nd('Mul', [lower_oh_4d, lower_mask], [[1, 10, 10, 10]])
+    
+    # Combine colored channels
+    colored_10 = nd('Add', [upper_out, lower_out], [[1, 10, 10, 10]])
+    
+    # === STEP 6: Add background channel ===
+    # ch0 = 1 where no other channel is active (inside 10x10)
+    # Slice channels 1-9, take max, subtract from 1
+    nonbg_10 = nd('Slice', [colored_10, 'sl_ch_start', 'sl_ch_end', 'sl_axes'], [[1, 9, 10, 10]])
+    any_color = nd('ReduceMax', [nonbg_10, 'axes1'], [[1, 1, 10, 10]], keepdims=1)
+    bg_ch = nd('Sub', ['c_one', any_color], [[1, 1, 10, 10]])
+    
+    # Build ch0 one-hot and add
+    const('ch0_idx', [0], 'i')
+    ch0_oh = nd('OneHot', ['ch0_idx', 'depth_10', 'oh_vals'], [[1, 10]], axis=1)
+    ch0_oh_4d = nd('Reshape', [ch0_oh, 'shape_1_10_1_1'], [[1, 10, 1, 1]])
+    bg_full = nd('Mul', [ch0_oh_4d, bg_ch], [[1, 10, 10, 10]])
+    
+    out_10 = nd('Add', [colored_10, bg_full], [[1, 10, 10, 10]])
+    
+    # === STEP 7: Pad to 30x30 ===
+    final = nd('Pad', [out_10, 'pad_10_to_30', 'pad_val_zero'], [[1, 10, 30, 30]])
+    
+    # Build model
+    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])
+    nodes.append(helper.make_node('Identity', [final], ['output']))
+    
+    graph = helper.make_graph(nodes, 'task028', [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
+    from huggingface_hub import hf_hub_download
+    import zipfile
+    
+    print("Building Task 028 ONNX model...")
+    model = build_task028()
+    
+    output_path = '/app/task028.onnx'
+    onnx.save(model, output_path)
+    fsize = os.path.getsize(output_path)
+    print(f"  Nodes: {len(model.graph.node)}")
+    print(f"  File size: {fsize:,} bytes")
+    
+    # Validate
+    sess = ort.InferenceSession(output_path)
+    
+    task_path = hf_hub_download('rogermt/neurogolf-solver', 'own-solver/neurogolf-2026.zip')
+    with zipfile.ZipFile(task_path, 'r') as zf:
+        data = json.loads(zf.read('task028.json'))
+    
+    all_examples = data['train'] + data['test'] + data.get('arc-gen', [])
+    right_count, wrong_count = 0, 0
+    
+    for i, ex in enumerate(all_examples):
+        inp_grid = ex['input']
+        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
+        
+        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 <= 3:
+                diff_locs = np.where(out != exp)
+                print(f"  FAIL {i}: {len(diff_locs[0])} diffs")
+    
+    print(f"\nResults: {right_count} pass, {wrong_count} fail out of {len(all_examples)}")
+    
+    if wrong_count == 0:
+        # Score estimate
+        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])) * 4
+                      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 est: {mem_est:,}")
+        print(f"  Score est: {score_est:.3f} (base: 12.831)")
+        print(f"  Gain est: +{score_est - 12.831:.3f}")