""" onnx_builder.py — Shared ONNX graph building utilities for medal-solvers. Provides: - OnnxBuilder: DSL for building ONNX graphs (const, nd, vi helpers) - validate_model(): Quick ORT-based validation against task JSON - validate_official(): Full validation with neurogolf_utils.py - build_and_validate(): CLI entry point for build scripts Usage in a build script: from onnx_builder import OnnxBuilder, build_and_validate def build_taskNNN(): b = OnnxBuilder() # Use b.const(), b.nd(), etc. to build the graph ... return b.finish('taskNNN', input_shape=[1, 10, 30, 30], output_shape=[1, 10, 30, 30]) if __name__ == '__main__': build_and_validate(build_taskNNN, task_num=NNN) """ import onnx from onnx import helper, TensorProto, numpy_helper import numpy as np import os import sys import json class OnnxBuilder: """DSL for building ONNX graphs with minimal boilerplate.""" def __init__(self): self.nodes = [] self.inits = [] self.vis = [] self._counter = 0 def _nm(self): """Generate unique tensor name.""" self._counter += 1 return f"t{self._counter}" def const(self, name, val, dtype='f'): """Add a constant initializer. Args: name: Initializer name (must be unique across the model) val: Value (scalar, list, or numpy array) dtype: 'f' for float32, 'i' for int64 """ if dtype == 'f': arr = np.array(val, dtype=np.float32) else: arr = np.array(val, dtype=np.int64) self.inits.append(numpy_helper.from_array(arr, name)) def vi(self, name, shape, dt=TensorProto.FLOAT): """Register a value_info (intermediate tensor shape declaration).""" self.vis.append(helper.make_tensor_value_info(name, dt, shape)) def nd(self, op, ins, outs_shapes, **kwargs): """Add a node to the graph. Args: op: ONNX operator name (e.g. 'Add', 'Conv', 'Gather') ins: List of input tensor names outs_shapes: List of output shape specs. Each is either: - a list (shape, assumes FLOAT dtype) - a tuple (shape, dtype) for non-float outputs **kwargs: Additional node attributes (e.g. axis=1, pads=[...]) Returns: Single output name if 1 output, else list of output names. """ out_names = [] for sd in outs_shapes: if isinstance(sd, tuple): shape, dt = sd else: shape, dt = sd, TensorProto.FLOAT n = self._nm() self.vi(n, shape, dt) out_names.append(n) self.nodes.append(helper.make_node(op, ins, out_names, **kwargs)) return out_names[0] if len(out_names) == 1 else out_names def finish(self, graph_name, input_shape=None, output_shape=None, input_name='input', output_name='output', last_tensor=None, opset=18): """Finalize and return the ONNX ModelProto. Args: graph_name: Name for the ONNX graph input_shape: Shape of the input tensor (default [1, 10, 30, 30]) output_shape: Shape of the output tensor (default [1, 10, 30, 30]) input_name: Name of input tensor (default 'input') output_name: Name of output tensor (default 'output') last_tensor: If provided, adds an Identity node connecting this to output opset: ONNX opset version (default 18) Returns: onnx.ModelProto """ if input_shape is None: input_shape = [1, 10, 30, 30] if output_shape is None: output_shape = [1, 10, 30, 30] if last_tensor is not None: self.nodes.append(helper.make_node('Identity', [last_tensor], [output_name])) x = helper.make_tensor_value_info(input_name, TensorProto.FLOAT, input_shape) y = helper.make_tensor_value_info(output_name, TensorProto.FLOAT, output_shape) graph = helper.make_graph( self.nodes, graph_name, [x], [y], initializer=self.inits, value_info=self.vis ) model = helper.make_model( graph, ir_version=10, opset_imports=[helper.make_opsetid('', opset)] ) return model def validate_model(model_path, task_json_path, task_num=None): """Quick validation: run model on all examples and check correctness. Returns: (pass_count, fail_count) """ import onnxruntime as ort with open(task_json_path) as f: data = json.load(f) sess = ort.InferenceSession(model_path) right, wrong = 0, 0 all_ex = data.get('train', []) + data.get('test', []) + data.get('arc-gen', []) for ex in all_ex: inp_grid = ex['input'] if max(len(inp_grid), max((len(r) for r in inp_grid), default=0)) > 30: continue 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 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 result = sess.run(['output'], {'input': inp}) out = (result[0] > 0.0).astype(float) if np.array_equal(out, exp): right += 1 else: wrong += 1 return right, wrong def validate_official(model_path, task_num, neurogolf_utils_path, task_data_dir): """Run official neurogolf_utils.verify_network() for scoring. Returns True on success. """ import types # Mock IPython.display (not available outside notebooks) mock_ipython = types.ModuleType('IPython') mock_display = types.ModuleType('IPython.display') mock_display.display = lambda *a, **k: None mock_display.FileLink = lambda x: x mock_ipython.display = mock_display sys.modules['IPython'] = mock_ipython sys.modules['IPython.display'] = mock_display import matplotlib matplotlib.use('Agg') utils_dir = os.path.dirname(os.path.abspath(neurogolf_utils_path)) if utils_dir not in sys.path: sys.path.insert(0, utils_dir) import neurogolf_utils neurogolf_utils._NEUROGOLF_DIR = task_data_dir.rstrip('/') + '/' examples = neurogolf_utils.load_examples(task_num) print(f" Loaded: train={len(examples['train'])}, " f"test={len(examples['test'])}, arc-gen={len(examples['arc-gen'])}") network = onnx.load(model_path) print(f"\n{'='*60}") print(f" Official neurogolf_utils.verify_network():") print(f"{'='*60}\n") neurogolf_utils.verify_network(network, task_num, examples) print(f"\n{'='*60}") return True def find_path(filename, candidates): """Find first existing path from candidates list.""" if os.path.exists(filename): return filename for c in candidates: if os.path.exists(c): return c return filename def build_and_validate(build_fn, task_num, default_output=None): """CLI entry point for build scripts. Handles argument parsing, building, saving, and validation. Args: build_fn: Function that returns an onnx.ModelProto task_num: Task number (e.g. 319) default_output: Default output path (default: optimized/taskNNN.onnx) """ import argparse if default_output is None: default_output = f'/app/repo/medal-solvers/optimized/task{task_num:03d}.onnx' parser = argparse.ArgumentParser( description=f"Build and validate optimized Task {task_num} model" ) parser.add_argument('--neurogolf-utils', default='neurogolf_utils.py', help='Path to official neurogolf_utils.py') parser.add_argument('--task-data-dir', default='../task-data', help='Directory containing taskNNN.json files') parser.add_argument('--output', default=default_output, help=f'Output path (default: {default_output})') parser.add_argument('--skip-official', action='store_true', help='Skip official neurogolf_utils validation') args = parser.parse_args() # Auto-detect paths task_data_dir = find_path( args.task_data_dir, ['task-data', '../task-data', '.', '/kaggle/input/competitions/neurogolf-2026'] ) neurogolf_path = find_path( args.neurogolf_utils, ['neurogolf_utils.py', '../own-solver/neurogolf_utils.py', 'own-solver/neurogolf_utils.py'] ) # Build print(f"Building Task {task_num} optimized model...") model = build_fn() os.makedirs(os.path.dirname(args.output) or '.', exist_ok=True) onnx.save(model, args.output) fsize = os.path.getsize(args.output) print(f" Nodes: {len(model.graph.node)}") print(f" File size: {fsize:,} bytes (limit: {int(1.44*1024*1024):,})") print(f" Saved to: {args.output}") # Quick validation task_json = os.path.join(task_data_dir, f'task{task_num:03d}.json') if os.path.exists(task_json): print(f"\n Quick validation against {task_json}...") right, wrong = validate_model(args.output, task_json, task_num) print(f" Results: {right} pass, {wrong} fail out of {right + wrong}") if wrong > 0: print(f" \u26a0\ufe0f FAILURES DETECTED \u2014 model is NOT ready for submission!") return False else: print(f" \u2713 All examples pass!") else: print(f"\n \u26a0\ufe0f Task data not found at {task_json} \u2014 skipping validation") # Official validation if not args.skip_official and os.path.exists(neurogolf_path) and os.path.exists(task_json): print(f"\n Official validation with {neurogolf_path}...") validate_official(args.output, task_num, neurogolf_path, task_data_dir) return True