Upload 11 files

README.md

@@ -0,0 +1,48 @@
+---
+library_name: mlx
+tags:
+  - mlx
+  - model-routing
+  - synthetic-data
+---
+
+# Succinct Router 14M
+
+A 14M-parameter decoder-only routing model trained from random initialization. It predicts an
+independent pass probability for each of these candidate configurations:
+
+- `gpt-5.6-luna-none`
+- `gpt-5.6-terra-low`
+- `gpt-5.6-sol-medium`
+
+This model does **not** answer prompts. It selects the cheapest candidate configuration expected to
+pass a calibrated quality threshold.
+
+## Evaluation
+
+```json
+{
+  "exact_route_accuracy": 0.9508599508599509,
+  "unsafe_downroute_rate": 0.005528255528255528,
+  "selected_model_pass_rate": 0.9944717444717445,
+  "abstention_rate": 0.0,
+  "savings_vs_always_large": 0.8486732186732187
+}
+```
+
+## Run with MLX
+
+```bash
+pip install "mlx>=0.32,<0.33" tokenizers
+python route_mlx.py --model-dir . --prompt "Extract the invoice number from INV-4821."
+```
+
+The result contains calibrated pass probabilities and the selected candidate configuration. If no
+candidate clears the threshold, `abstained_to_largest` is true and the route falls back to the
+largest configuration.
+
+## Limitations
+
+Training and evaluation data are synthetic and mechanically graded. The artifact is a learning
+prototype, not a production-ready router. Validate it on anonymized real traffic before making
+product or cost claims. MLX runtime parity and latency must be measured on Apple Silicon.

calibration.json

@@ -0,0 +1,40 @@
+{
+  "schema_version": "1",
+  "model_ids": [
+    "gpt-5.6-luna-none",
+    "gpt-5.6-terra-low",
+    "gpt-5.6-sol-medium"
+  ],
+  "temperatures": [
+    1.0598043203353882,
+    1.106358528137207,
+    0.3136223554611206
+  ],
+  "selection_threshold": 0.58,
+  "constraints": {
+    "max_unsafe_downroute_rate": 0.005,
+    "min_selected_model_pass_rate": 0.99,
+    "satisfied": true
+  },
+  "validation_route_metrics": {
+    "exact_route_accuracy": 0.9509283819628647,
+    "unsafe_downroute_rate": 0.004641909814323607,
+    "selected_model_pass_rate": 0.9953580901856764,
+    "abstention_rate": 0.0,
+    "savings_vs_always_large": 0.8416445623342175
+  },
+  "uncalibrated": {
+    "pass_vector_accuracy": 0.9244031830238727,
+    "pass_bit_accuracy": 0.9748010610079576,
+    "macro_f1": 0.985090243264452,
+    "brier_score": 0.01777769201095797,
+    "expected_calibration_error": 0.018437407725585803
+  },
+  "calibrated": {
+    "pass_vector_accuracy": 0.9244031830238727,
+    "pass_bit_accuracy": 0.9748010610079576,
+    "macro_f1": 0.985090243264452,
+    "brier_score": 0.017453370615839958,
+    "expected_calibration_error": 0.01830040662266169
+  }
+}

config.json

@@ -0,0 +1,16 @@
+{
+  "schema_version": "1",
+  "architecture": "succinct-router-decoder-v1",
+  "model": {
+    "vocab_size": 8192,
+    "max_sequence_length": 1024,
+    "layers": 6,
+    "model_width": 384,
+    "attention_heads": 6,
+    "ffn_width": 1024,
+    "dropout": 0.1,
+    "candidate_count": 3
+  },
+  "route_state": "final_<ROUTE>_token",
+  "tied_lm_head": true
+}

example_inputs.jsonl

@@ -0,0 +1,8 @@
+{"example_id": "sha256:5db7b3e355519960ac61e93ac9f7cd537f5010bb5142e8d2f0b85d9e801162e5", "input_ids": [2, 26, 200, 28, 94, 219, 125, 222, 158, 13, 177, 180, 220, 141, 221, 133, 188, 215, 11, 217, 11, 115, 168, 223, 13, 214, 198, 115, 218, 224, 159, 201, 199, 13, 94, 94, 26, 211, 28, 94, 492, 352, 11, 351, 11, 133, 612, 159, 141, 13, 94, 94, 560, 11, 859, 530, 605, 13, 332, 700, 526, 13, 1036, 238, 540, 664, 13, 686, 608, 702, 99, 973, 59, 974, 29, 4], "reference_pass_logits": [7.015040397644043, 6.347684383392334, 6.796907901763916], "reference_scores": [0.989065408706665, 0.9848109483718872, 0.990003764629364]}
+{"example_id": "sha256:780d0b0600848b763ea1d247395a6abc9295344d2a44d8cd7ba33641927bec67", "input_ids": [2, 26, 200, 28, 94, 219, 125, 222, 158, 13, 177, 180, 220, 141, 221, 133, 188, 215, 11, 217, 11, 115, 168, 223, 13, 214, 198, 115, 218, 224, 159, 201, 199, 13, 94, 94, 26, 211, 28, 94, 492, 352, 11, 351, 11, 133, 612, 159, 141, 13, 94, 94, 560, 11, 859, 530, 605, 13, 332, 700, 526, 13, 1036, 238, 540, 664, 13, 686, 608, 702, 99, 973, 59, 974, 29, 4], "reference_pass_logits": [7.015040397644043, 6.347684383392334, 6.796907901763916], "reference_scores": [0.989065408706665, 0.9848109483718872, 0.990003764629364]}
+{"example_id": "sha256:2f773ba7a80b8d601866a7d93b6ea025b2fe8dfc9140365dac1e8e60ab9e26d1", "input_ids": [2, 26, 200, 28, 94, 219, 125, 222, 158, 13, 177, 180, 220, 141, 221, 133, 188, 215, 11, 217, 11, 115, 168, 223, 13, 214, 198, 115, 218, 224, 159, 201, 199, 13, 94, 94, 26, 211, 28, 94, 814, 180, 617, 133, 240, 158, 148, 595, 651, 345, 704, 683, 537, 59, 270, 11, 639, 644, 267, 158, 110, 571, 6, 348, 116, 483, 629, 238, 13, 712, 24, 349, 234, 107, 110, 678, 59, 270, 181, 110, 345, 107, 148, 270, 59, 144, 107, 110, 340, 641, 148, 234, 107, 110, 445, 59, 270, 181, 110, 345, 107, 148, 270, 59, 144, 107, 110, 340, 181, 110, 618, 107, 110, 340, 642, 94, 94, 558, 434, 12, 1216, 820, 25, 295, 682, 491, 541, 13, 4], "reference_pass_logits": [6.926111698150635, 6.366804122924805, 6.842347621917725], "reference_scores": [0.99000084400177, 0.9857044219970703, 0.9898501634597778]}
+{"example_id": "sha256:2c7baa36710166560457211f804dd37cf02b40470e395dda9a0c64c72b528a91", "input_ids": [2, 26, 200, 28, 94, 219, 125, 222, 158, 13, 177, 180, 220, 141, 221, 133, 188, 215, 11, 217, 11, 115, 168, 223, 13, 214, 198, 115, 218, 224, 159, 201, 199, 13, 94, 94, 26, 211, 28, 94, 814, 180, 617, 133, 240, 158, 148, 595, 651, 345, 704, 683, 537, 59, 270, 11, 639, 644, 267, 158, 110, 571, 6, 348, 116, 483, 629, 238, 13, 712, 24, 349, 234, 107, 110, 678, 59, 270, 181, 110, 345, 107, 148, 270, 59, 144, 107, 110, 340, 641, 148, 234, 107, 110, 445, 59, 270, 181, 110, 345, 107, 148, 270, 59, 144, 107, 110, 340, 181, 110, 618, 107, 110, 340, 642, 94, 94, 558, 434, 12, 1216, 820, 25, 295, 682, 491, 541, 13, 4], "reference_pass_logits": [6.926111698150635, 6.366804122924805, 6.842347621917725], "reference_scores": [0.99000084400177, 0.9857044219970703, 0.9898501634597778]}
+{"example_id": "sha256:1746c9cf5738c6f562730b79c8a48454e0457b259ee27e32fb1ac9b0b8eba4f9", "input_ids": [2, 26, 200, 28, 94, 219, 125, 222, 158, 13, 177, 180, 220, 141, 221, 133, 188, 215, 11, 217, 11, 115, 168, 223, 13, 214, 198, 115, 218, 224, 159, 201, 199, 13, 94, 94, 26, 211, 28, 94, 166, 125, 464, 459, 699, 548, 677, 717, 706, 159, 148, 436, 107, 481, 456, 671, 684, 819, 11, 551, 99, 554, 7, 654, 11, 755, 551, 333, 7, 834, 548, 125, 991, 992, 13, 670, 24, 349, 284, 107, 1114, 13, 1062, 11, 110, 286, 107, 335, 285, 148, 284, 107, 1095, 13, 17, 11, 110, 286, 107, 288, 285, 148, 284, 107, 1150, 13, 1049, 11, 110, 286, 107, 288, 285, 148, 284, 107, 1156, 13, 1009, 11, 110, 286, 107, 230, 457, 4], "reference_pass_logits": [-7.627079963684082, 2.626823663711548, 4.990095615386963], "reference_scores": [0.018578819930553436, 0.9573246240615845, 0.9830689430236816]}
+{"example_id": "sha256:40396bbb4a9130ac5fdeb93980af7a9f98f12406d55c20b9c48a05ad6e65ece7", "input_ids": [2, 26, 200, 28, 94, 219, 125, 222, 158, 13, 177, 180, 220, 141, 221, 133, 188, 215, 11, 217, 11, 115, 168, 223, 13, 214, 198, 115, 218, 224, 159, 201, 199, 13, 94, 94, 26, 211, 28, 94, 833, 506, 464, 459, 699, 548, 677, 717, 706, 159, 148, 436, 107, 481, 456, 671, 684, 819, 11, 551, 99, 554, 7, 654, 11, 755, 551, 333, 7, 834, 548, 506, 991, 992, 13, 670, 24, 349, 284, 107, 1114, 13, 1062, 11, 110, 286, 107, 335, 285, 148, 284, 107, 1095, 13, 17, 11, 110, 286, 107, 288, 285, 148, 284, 107, 1150, 13, 1049, 11, 110, 286, 107, 288, 285, 148, 284, 107, 1156, 13, 1009, 11, 110, 286, 107, 230, 457, 4], "reference_pass_logits": [-7.622464656829834, 2.636568307876587, 4.985684394836426], "reference_scores": [0.018651440739631653, 0.9573564529418945, 0.9829824566841125]}
+{"example_id": "sha256:62fafb8ee1842978bd6c2b7a324561b09a2b546db5c40264b20d33c6fa2de220", "input_ids": [2, 26, 200, 28, 94, 219, 125, 222, 158, 13, 177, 180, 220, 141, 221, 133, 188, 215, 11, 217, 11, 115, 168, 223, 13, 214, 198, 115, 218, 224, 159, 201, 199, 13, 94, 94, 26, 211, 28, 94, 817, 125, 715, 315, 342, 59, 144, 133, 714, 459, 59, 633, 13, 177, 141, 271, 13, 94, 94, 132, 568, 107, 148, 144, 107, 110, 691, 12, 1221, 865, 273, 110, 688, 107, 110, 695, 181, 110, 436, 59, 574, 107, 1120, 13, 1067, 87, 4], "reference_pass_logits": [6.730648040771484, 6.451633453369141, 6.800724506378174], "reference_scores": [0.9886912107467651, 0.9871037602424622, 0.9892838597297668]}
+{"example_id": "sha256:3c9e6ec5cf888091fd117958d1f87455b125ac6985f38d2adfa72fcd66815b50", "input_ids": [2, 26, 200, 28, 94, 219, 125, 222, 158, 13, 177, 180, 220, 141, 221, 133, 188, 215, 11, 217, 11, 115, 168, 223, 13, 214, 198, 115, 218, 224, 159, 201, 199, 13, 94, 94, 26, 211, 28, 94, 817, 506, 715, 315, 342, 59, 144, 133, 714, 459, 59, 633, 13, 240, 141, 271, 13, 94, 94, 132, 568, 107, 148, 144, 107, 110, 691, 12, 1221, 865, 273, 110, 688, 107, 110, 695, 181, 110, 436, 59, 574, 107, 1120, 13, 1067, 87, 4], "reference_pass_logits": [6.713947296142578, 6.455345153808594, 6.808040618896484], "reference_scores": [0.9886158108711243, 0.9869857430458069, 0.9892330169677734]}

labels.json

@@ -0,0 +1,7 @@
+{
+  "candidate_model_config_ids": [
+    "gpt-5.6-luna-none",
+    "gpt-5.6-terra-low",
+    "gpt-5.6-sol-medium"
+  ]
+}

manifest.json

@@ -0,0 +1,391 @@
+{
+  "schema_version": "1",
+  "format": "succinct-router-mlx-v1",
+  "linux_validation": {
+    "tensor_names_shapes_dtypes": "passed",
+    "tokenizer_round_trip": "passed",
+    "reference_pytorch_logits": 8
+  },
+  "mlx_runtime_validation": {
+    "status": "passed",
+    "device": "Device(cpu, 0)",
+    "examples": 8,
+    "max_abs_pass_logit_error": 3.814697265625e-06,
+    "max_abs_score_error": 5.960464477539063e-08,
+    "tolerance": 0.0001
+  },
+  "apple_silicon_metal_latency_validation": "required",
+  "weights": {
+    "tensor_count": 48,
+    "sha256": "35e54bfcffe5ade95e32bee903c708ca30554170acb9c7b82dc7db6e3d4eebe4",
+    "bytes": 55084128,
+    "tied_weight": {
+      "lm_head.weight": "embedding.weight"
+    },
+    "tensors": {
+      "blocks.0.attention.output.weight": {
+        "shape": [
+          384,
+          384
+        ],
+        "dtype": "float32"
+      },
+      "blocks.0.attention.qkv.weight": {
+        "shape": [
+          1152,
+          384
+        ],
+        "dtype": "float32"
+      },
+      "blocks.0.attention_norm.weight": {
+        "shape": [
+          384
+        ],
+        "dtype": "float32"
+      },
+      "blocks.0.ffn.down.weight": {
+        "shape": [
+          384,
+          1024
+        ],
+        "dtype": "float32"
+      },
+      "blocks.0.ffn.gate.weight": {
+        "shape": [
+          1024,
+          384
+        ],
+        "dtype": "float32"
+      },
+      "blocks.0.ffn.up.weight": {
+        "shape": [
+          1024,
+          384
+        ],
+        "dtype": "float32"
+      },
+      "blocks.0.ffn_norm.weight": {
+        "shape": [
+          384
+        ],
+        "dtype": "float32"
+      },
+      "blocks.1.attention.output.weight": {
+        "shape": [
+          384,
+          384
+        ],
+        "dtype": "float32"
+      },
+      "blocks.1.attention.qkv.weight": {
+        "shape": [
+          1152,
+          384
+        ],
+        "dtype": "float32"
+      },
+      "blocks.1.attention_norm.weight": {
+        "shape": [
+          384
+        ],
+        "dtype": "float32"
+      },
+      "blocks.1.ffn.down.weight": {
+        "shape": [
+          384,
+          1024
+        ],
+        "dtype": "float32"
+      },
+      "blocks.1.ffn.gate.weight": {
+        "shape": [
+          1024,
+          384
+        ],
+        "dtype": "float32"
+      },
+      "blocks.1.ffn.up.weight": {
+        "shape": [
+          1024,
+          384
+        ],
+        "dtype": "float32"
+      },
+      "blocks.1.ffn_norm.weight": {
+        "shape": [
+          384
+        ],
+        "dtype": "float32"
+      },
+      "blocks.2.attention.output.weight": {
+        "shape": [
+          384,
+          384
+        ],
+        "dtype": "float32"
+      },
+      "blocks.2.attention.qkv.weight": {
+        "shape": [
+          1152,
+          384
+        ],
+        "dtype": "float32"
+      },
+      "blocks.2.attention_norm.weight": {
+        "shape": [
+          384
+        ],
+        "dtype": "float32"
+      },
+      "blocks.2.ffn.down.weight": {
+        "shape": [
+          384,
+          1024
+        ],
+        "dtype": "float32"
+      },
+      "blocks.2.ffn.gate.weight": {
+        "shape": [
+          1024,
+          384
+        ],
+        "dtype": "float32"
+      },
+      "blocks.2.ffn.up.weight": {
+        "shape": [
+          1024,
+          384
+        ],
+        "dtype": "float32"
+      },
+      "blocks.2.ffn_norm.weight": {
+        "shape": [
+          384
+        ],
+        "dtype": "float32"
+      },
+      "blocks.3.attention.output.weight": {
+        "shape": [
+          384,
+          384
+        ],
+        "dtype": "float32"
+      },
+      "blocks.3.attention.qkv.weight": {
+        "shape": [
+          1152,
+          384
+        ],
+        "dtype": "float32"
+      },
+      "blocks.3.attention_norm.weight": {
+        "shape": [
+          384
+        ],
+        "dtype": "float32"
+      },
+      "blocks.3.ffn.down.weight": {
+        "shape": [
+          384,
+          1024
+        ],
+        "dtype": "float32"
+      },
+      "blocks.3.ffn.gate.weight": {
+        "shape": [
+          1024,
+          384
+        ],
+        "dtype": "float32"
+      },
+      "blocks.3.ffn.up.weight": {
+        "shape": [
+          1024,
+          384
+        ],
+        "dtype": "float32"
+      },
+      "blocks.3.ffn_norm.weight": {
+        "shape": [
+          384
+        ],
+        "dtype": "float32"
+      },
+      "blocks.4.attention.output.weight": {
+        "shape": [
+          384,
+          384
+        ],
+        "dtype": "float32"
+      },
+      "blocks.4.attention.qkv.weight": {
+        "shape": [
+          1152,
+          384
+        ],
+        "dtype": "float32"
+      },
+      "blocks.4.attention_norm.weight": {
+        "shape": [
+          384
+        ],
+        "dtype": "float32"
+      },
+      "blocks.4.ffn.down.weight": {
+        "shape": [
+          384,
+          1024
+        ],
+        "dtype": "float32"
+      },
+      "blocks.4.ffn.gate.weight": {
+        "shape": [
+          1024,
+          384
+        ],
+        "dtype": "float32"
+      },
+      "blocks.4.ffn.up.weight": {
+        "shape": [
+          1024,
+          384
+        ],
+        "dtype": "float32"
+      },
+      "blocks.4.ffn_norm.weight": {
+        "shape": [
+          384
+        ],
+        "dtype": "float32"
+      },
+      "blocks.5.attention.output.weight": {
+        "shape": [
+          384,
+          384
+        ],
+        "dtype": "float32"
+      },
+      "blocks.5.attention.qkv.weight": {
+        "shape": [
+          1152,
+          384
+        ],
+        "dtype": "float32"
+      },
+      "blocks.5.attention_norm.weight": {
+        "shape": [
+          384
+        ],
+        "dtype": "float32"
+      },
+      "blocks.5.ffn.down.weight": {
+        "shape": [
+          384,
+          1024
+        ],
+        "dtype": "float32"
+      },
+      "blocks.5.ffn.gate.weight": {
+        "shape": [
+          1024,
+          384
+        ],
+        "dtype": "float32"
+      },
+      "blocks.5.ffn.up.weight": {
+        "shape": [
+          1024,
+          384
+        ],
+        "dtype": "float32"
+      },
+      "blocks.5.ffn_norm.weight": {
+        "shape": [
+          384
+        ],
+        "dtype": "float32"
+      },
+      "embedding.weight": {
+        "shape": [
+          8192,
+          384
+        ],
+        "dtype": "float32"
+      },
+      "norm.weight": {
+        "shape": [
+          384
+        ],
+        "dtype": "float32"
+      },
+      "pass_head.bias": {
+        "shape": [
+          3
+        ],
+        "dtype": "float32"
+      },
+      "pass_head.weight": {
+        "shape": [
+          3,
+          384
+        ],
+        "dtype": "float32"
+      },
+      "score_head.bias": {
+        "shape": [
+          3
+        ],
+        "dtype": "float32"
+      },
+      "score_head.weight": {
+        "shape": [
+          3,
+          384
+        ],
+        "dtype": "float32"
+      }
+    }
+  },
+  "files": {
+    "README.md": {
+      "bytes": 1397,
+      "sha256": "a1c485e0e037508e1ea8b7fe6c6422d24ed1ddaa8c60579ba552eeae4a673c29"
+    },
+    "calibration.json": {
+      "bytes": 1139,
+      "sha256": "416b2d6d5d609608c405a615646cab6e7ee9022d24586b9245053a75e78daf74"
+    },
+    "config.json": {
+      "bytes": 350,
+      "sha256": "68c3fd9f8ae9c8376cfdbba58acb642c76645c145705f6bb7e68e86c3de4367a"
+    },
+    "example_inputs.jsonl": {
+      "bytes": 6124,
+      "sha256": "223879c190c301466279ddd177fd46af91e448b6b8e1be44a4aa89c2ccc43ca0"
+    },
+    "labels.json": {
+      "bytes": 117,
+      "sha256": "c4ea84d5fcde04d423647df6843a07c21ee680ec95983ac485370f07697d495a"
+    },
+    "mlx_model.py": {
+      "bytes": 5486,
+      "sha256": "5fbb97e004f7c9e613ae94144f3d04966d64d1a2a0bbf152664296876fd28901"
+    },
+    "model.safetensors": {
+      "bytes": 55084128,
+      "sha256": "35e54bfcffe5ade95e32bee903c708ca30554170acb9c7b82dc7db6e3d4eebe4"
+    },
+    "route_mlx.py": {
+      "bytes": 3265,
+      "sha256": "20b7520c89c876d9ae2e964bf49679294a1f50e67567538f719c8107db9a4d6c"
+    },
+    "tokenizer.json": {
+      "bytes": 532269,
+      "sha256": "f417f5abf0de1c4a475620c4b16b487834115555688d379b29202c4edc5d7481"
+    },
+    "tokenizer_config.json": {
+      "bytes": 183,
+      "sha256": "568172f0cc44b6d1d080d53807dce5537a0e58cdbe7fb9d7e6678f97c1f3641c"
+    }
+  }
+}

mlx_model.py

@@ -0,0 +1,132 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any
+
+try:
+    import mlx.core as mx
+    import mlx.nn as nn
+except ImportError:  # Linux export validation intentionally has no MLX runtime.
+    mx = None
+    nn = None
+
+
+if nn is not None and mx is not None:
+
+    class RMSNorm(nn.Module):
+        def __init__(self, dimension: int, epsilon: float = 1e-6) -> None:
+            super().__init__()
+            self.weight = mx.ones((dimension,))
+            self.epsilon = epsilon
+
+        def __call__(self, values: Any) -> Any:
+            normalized = values * mx.rsqrt(
+                mx.mean(mx.square(values), axis=-1, keepdims=True) + self.epsilon
+            )
+            return normalized * self.weight
+
+    def _rotate_half(values: Any) -> Any:
+        first, second = mx.split(values, 2, axis=-1)
+        return mx.concatenate((-second, first), axis=-1)
+
+    def apply_rope(query: Any, key: Any) -> tuple[Any, Any]:
+        sequence_length = query.shape[-2]
+        dimension = query.shape[-1]
+        positions = mx.arange(sequence_length, dtype=mx.float32)
+        frequencies = 1.0 / (10000 ** (mx.arange(0, dimension, 2, dtype=mx.float32) / dimension))
+        angles = positions[:, None] * frequencies[None, :]
+        angles = mx.concatenate((angles, angles), axis=-1)[None, None, :, :]
+        cosine = mx.cos(angles)
+        sine = mx.sin(angles)
+        return query * cosine + _rotate_half(query) * sine, key * cosine + _rotate_half(key) * sine
+
+    class CausalSelfAttention(nn.Module):
+        def __init__(self, config: dict[str, Any]) -> None:
+            super().__init__()
+            self.heads = int(config["attention_heads"])
+            self.head_dimension = int(config["model_width"]) // self.heads
+            width = int(config["model_width"])
+            self.qkv = nn.Linear(width, 3 * width, bias=False)
+            self.output = nn.Linear(width, width, bias=False)
+
+        def __call__(self, values: Any) -> Any:
+            batch, sequence, width = values.shape
+            qkv = self.qkv(values).reshape(batch, sequence, 3, self.heads, self.head_dimension)
+            qkv = qkv.transpose(2, 0, 3, 1, 4)
+            query, key, value = qkv[0], qkv[1], qkv[2]
+            query, key = apply_rope(query, key)
+            mask = nn.MultiHeadAttention.create_additive_causal_mask(sequence)
+            attended = mx.fast.scaled_dot_product_attention(
+                query,
+                key,
+                value,
+                scale=self.head_dimension**-0.5,
+                mask=mask,
+            )
+            return self.output(attended.transpose(0, 2, 1, 3).reshape(batch, sequence, width))
+
+    class SwiGLU(nn.Module):
+        def __init__(self, config: dict[str, Any]) -> None:
+            super().__init__()
+            width = int(config["model_width"])
+            ffn_width = int(config["ffn_width"])
+            self.gate = nn.Linear(width, ffn_width, bias=False)
+            self.up = nn.Linear(width, ffn_width, bias=False)
+            self.down = nn.Linear(ffn_width, width, bias=False)
+
+        def __call__(self, values: Any) -> Any:
+            return self.down(nn.silu(self.gate(values)) * self.up(values))
+
+    class TransformerBlock(nn.Module):
+        def __init__(self, config: dict[str, Any]) -> None:
+            super().__init__()
+            width = int(config["model_width"])
+            self.attention_norm = RMSNorm(width)
+            self.attention = CausalSelfAttention(config)
+            self.ffn_norm = RMSNorm(width)
+            self.ffn = SwiGLU(config)
+
+        def __call__(self, values: Any) -> Any:
+            values = values + self.attention(self.attention_norm(values))
+            return values + self.ffn(self.ffn_norm(values))
+
+    class TransformerRouterMLX(nn.Module):
+        def __init__(self, config: dict[str, Any]) -> None:
+            super().__init__()
+            width = int(config["model_width"])
+            self.config = config
+            self.embedding = nn.Embedding(int(config["vocab_size"]), width)
+            self.blocks = [TransformerBlock(config) for _ in range(int(config["layers"]))]
+            self.norm = RMSNorm(width)
+            self.pass_head = nn.Linear(width, int(config["candidate_count"]))
+            self.score_head = nn.Linear(width, int(config["candidate_count"]))
+
+        def __call__(self, input_ids: Any) -> dict[str, Any]:
+            hidden = self.embedding(input_ids)
+            for block in self.blocks:
+                hidden = block(hidden)
+            hidden = self.norm(hidden)
+            route_state = hidden[:, -1]
+            return {
+                "pass_logits": self.pass_head(route_state),
+                "scores": mx.sigmoid(self.score_head(route_state)),
+            }
+
+    def load_mlx_router(directory: str | Path) -> Any:
+        root = Path(directory)
+        config = json.loads((root / "config.json").read_text(encoding="utf-8"))["model"]
+        model = TransformerRouterMLX(config)
+        weights = mx.load(str(root / "model.safetensors"))
+        model.load_weights(list(weights.items()), strict=True)
+        mx.eval(model.parameters())
+        return model
+
+else:
+
+    class TransformerRouterMLX:  # type: ignore[no-redef]
+        def __init__(self, config: dict[str, Any]) -> None:
+            raise RuntimeError("MLX requires macOS on Apple Silicon")
+
+    def load_mlx_router(directory: str | Path) -> Any:
+        raise RuntimeError("MLX requires macOS on Apple Silicon")

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:35e54bfcffe5ade95e32bee903c708ca30554170acb9c7b82dc7db6e3d4eebe4
+size 55084128

route_mlx.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+from typing import Any
+
+import mlx.core as mx
+from mlx_model import load_mlx_router
+from tokenizers import Tokenizer
+
+
+def _encode_prompt(
+    tokenizer: Tokenizer,
+    prompt: str,
+    *,
+    system: str | None,
+    max_length: int,
+) -> list[int]:
+    text = f"<user>\n{prompt}" if system is None else f"<system>\n{system}\n\n<user>\n{prompt}"
+    bos_id = tokenizer.token_to_id("<BOS>")
+    route_id = tokenizer.token_to_id("<ROUTE>")
+    head_tail_id = tokenizer.token_to_id("<HEAD_TAIL>")
+    if bos_id is None or route_id is None or head_tail_id is None:
+        raise ValueError("Tokenizer is missing required routing tokens")
+    body = tokenizer.encode(text, add_special_tokens=False).ids
+    available = max_length - 2
+    if len(body) > available:
+        head_length = (available - 1 + 1) // 2
+        tail_length = available - 1 - head_length
+        body = body[:head_length] + [head_tail_id] + (body[-tail_length:] if tail_length else [])
+    return [bos_id, *body, route_id]
+
+
+def route_prompt(
+    model_dir: str | Path,
+    prompt: str,
+    *,
+    system: str | None = None,
+) -> dict[str, Any]:
+    root = Path(model_dir)
+    config = json.loads((root / "config.json").read_text(encoding="utf-8"))
+    calibration = json.loads((root / "calibration.json").read_text(encoding="utf-8"))
+    labels = json.loads((root / "labels.json").read_text(encoding="utf-8"))[
+        "candidate_model_config_ids"
+    ]
+    tokenizer = Tokenizer.from_file(str(root / "tokenizer.json"))
+    input_ids = _encode_prompt(
+        tokenizer,
+        prompt,
+        system=system,
+        max_length=int(config["model"]["max_sequence_length"]),
+    )
+    model = load_mlx_router(root)
+    prediction = model(mx.array([input_ids]))
+    temperatures = mx.array(calibration["temperatures"])
+    probabilities = mx.sigmoid(prediction["pass_logits"][0] / temperatures)
+    mx.eval(probabilities, prediction["scores"])
+    pass_probabilities = probabilities.tolist()
+    threshold = float(calibration["selection_threshold"])
+    selected_index = next(
+        (index for index, probability in enumerate(pass_probabilities) if probability >= threshold),
+        None,
+    )
+    abstained = selected_index is None
+    fallback_index = len(labels) - 1
+    route_index = fallback_index if abstained else selected_index
+    return {
+        "route_config_id": labels[route_index],
+        "abstained_to_largest": abstained,
+        "selection_threshold": threshold,
+        "pass_probabilities": dict(zip(labels, pass_probabilities, strict=True)),
+        "predicted_scores": dict(zip(labels, prediction["scores"][0].tolist(), strict=True)),
+        "input_tokens": len(input_ids),
+    }
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Route one prompt with the exported MLX model")
+    parser.add_argument("--model-dir", default=".")
+    parser.add_argument("--prompt", required=True)
+    parser.add_argument("--system")
+    args = parser.parse_args()
+    print(
+        json.dumps(
+            route_prompt(args.model_dir, args.prompt, system=args.system),
+            indent=2,
+            sort_keys=True,
+        )
+    )
+
+
+if __name__ == "__main__":
+    main()

tokenizer_config.json

@@ -0,0 +1,12 @@
+{
+  "vocab_size": 8192,
+  "model_vocab_size": 8192,
+  "special_tokens": {
+    "<PAD>": 0,
+    "<UNK>": 1,
+    "<BOS>": 2,
+    "<EOS>": 3,
+    "<ROUTE>": 4,
+    "<HEAD_TAIL>": 5
+  }
+}