GLiNER2.5-Decide Core ML multifunction package (L64-L512) with bucket router

.gitignore

@@ -0,0 +1,13 @@
+.venv/
+__pycache__/
+.pytest_cache/
+*.egg-info/
+bench/benchbuckets
+bench/plan
+
+# Compiled Core ML model: rebuilt from the .mlpackage on first load.
+models/*.mlmodelc/
+
+# The .mlpackage is ~900 MB. Keep it out of plain git; use Git LFS (or rebuild it with
+# scripts/build) if this folder ever becomes a repository.
+models/*.mlpackage/

LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

NOTICE

@@ -0,0 +1,15 @@
+GLiNER2.5-Decide-CoreML-FP16-MultiFn-L64-512
+
+This product includes a Core ML conversion of GLiNER2.5-Decide
+(https://huggingface.co/fastino/GLiNER2.5-Decide), created by Fastino and
+licensed under the Apache License, Version 2.0.
+
+It includes files derived from gliner2-5-decide-coreml
+(https://huggingface.co/FluidInference/gliner2-5-decide-coreml) by Fluid
+Inference, licensed under the Apache License, Version 2.0:
+  scripts/build/export_model.py, convert_names.py, preprocessing.py and runtime.py (unmodified)
+  scripts/build/convert_bucket.py (modified from convert-coreml.py)
+  src/gliner_decide_coreml/encoding.py and decoding.py (adapted from preprocessing.py and runtime.py)
+
+The model packages in models/ were converted from the Fastino checkpoint at
+revision 65624f1a0265b3f612bae66a2685a06b94a68a9d.

README.md

@@ -0,0 +1,176 @@
+---
+license: apache-2.0
+library_name: coremltools
+pipeline_tag: zero-shot-classification
+base_model: fastino/GLiNER2.5-Decide
+language:
+- en
+tags:
+- coreml
+- gliner
+- gliner2
+- deberta-v3
+- apple-silicon
+- fp16
+- multifunction
+- zero-shot-classification
+---
+
+# GLiNER2.5-Decide-CoreML-FP16-MultiFn-L64-512
+
+[fastino/GLiNER2.5-Decide](https://huggingface.co/fastino/GLiNER2.5-Decide) (340M, DeBERTa-v3-large, zero-shot
+classification) packaged for fast local inference on Apple Silicon:
+
+- **Core ML, FP16**: the classification path (encoder + label classifier), converted with the tooling from
+  [FluidInference/gliner2-5-decide-coreml](https://huggingface.co/FluidInference/gliner2-5-decide-coreml).
+- **MultiFn**: one multifunction package with four fixed-length graphs (`L64`, `L128`, `L256`, `L512`) that
+  share a single copy of the weights: **0.9 GB instead of 3.5 GB** for four separate packages.
+- **Routing**: every request is tokenized and sent to the smallest bucket that fits it, so short requests
+  pay short-request latency. Text longer than 512 tokens is split into chunks automatically.
+
+## Results (Apple M1 Pro, macOS 27.2, `CPU_AND_GPU`)
+
+| Bucket | Model only (Swift) | End to end (Python router) | `ALL` compute units |
+|---:|---:|---:|---:|
+| 64 | 19.3 ms | 20.0 ms | 70.7 ms |
+| 128 | 31.6 ms | 32.2 ms | 197.6 ms |
+| 256 | 63.3 ms | 60.9 ms | 79.6 ms |
+| 512 | 133.9 ms | 135.3 ms | 208.9 ms |
+
+p50 latency. "End to end" includes tokenization, routing, prediction and decoding; Python preprocessing costs
+only 0.2–0.7 ms, so a native Swift runtime would not be meaningfully faster. **Use the GPU**: on this machine
+`ALL` (GPU + Neural Engine) is slower at every size, and the Neural Engine alone is ~10x slower (see
+[Research notes](#research-notes)).
+
+**Fidelity.** Every bucket was verified against the native PyTorch model at conversion time (logit error
+≤ 4.8e-7 for the export wrapper; identical labels on every example that fits). The test suite checks the
+router against native answers on six requests covering all four buckets (1–4 heads, multi-label,
+label descriptions, prompts): identical labels, largest confidence difference **0.0007**.
+
+## Quick start
+
+Requires Apple Silicon, **macOS 15+** (multifunction models), Python 3.10–3.12 and [uv](https://docs.astral.sh/uv/).
+
+```bash
+hf download augustoFranke/GLiNER2.5-Decide-CoreML-FP16-MultiFn-L64-512 --local-dir GLiNER2.5-Decide-CoreML-FP16-MultiFn-L64-512
+cd GLiNER2.5-Decide-CoreML-FP16-MultiFn-L64-512
+uv sync
+```
+
+```python
+from gliner_decide_coreml import DecideRouter
+
+router = DecideRouter()  # loads all four buckets once; do this at startup
+
+router.classify(
+    "My subscription renewed after the service was already down. Can I get that charge refunded?",
+    {
+        "intent": ["order_status", "refund_request", "cancel_subscription", "other"],
+        "urgency": ["low", "normal", "high"],
+        "topics": {"labels": ["billing", "outage", "account"], "multi_label": True, "cls_threshold": 0.4},
+    },
+)
+# {"intent": {"label": "refund_request", "confidence": 0.99}, "urgency": {...}, "topics": [{...}, ...]}
+
+result, route = router.classify_with_route(text, tasks)
+# Route(tokens=63, bucket=64, chunks=1, calls=1)
+```
+
+The task format is the same as native `classify_text`: a list of labels, or a dict with `labels` (list, or
+`{label: description}`), `multi_label`, `cls_threshold` and `prompt`.
+
+The first load of each function compiles it for the GPU (can take a minute or more on a new machine);
+macOS caches the result, and later loads take about 10 s for all four.
+
+## How it works
+
+```
+text + tasks ──► GLiNER2 schema tokenizer ──► n tokens ──► smallest bucket ≥ n ──► Core ML function ──► decode
+                 "( [P] intent ( [L] refund …"             64 / 128 / 256 / 512     (GPU, fp16)          softmax or
+                  [SEP_TEXT] my subscription …"             └─ n > 512: word chunks,                    sigmoid +
+                                                               merged per head                          threshold
+```
+
+- **Input budget.** The schema (questions and labels) shares the sequence with the text. With 2 questions and
+  8 labels (38 tokens) the buckets leave room for about 26 / 90 / 218 / 474 text tokens. Label descriptions
+  and prompts cost tokens too.
+- **Heads.** Each call scores up to 4 heads with up to 32 labels each. More than 4 heads are split across calls.
+  Heads in one call see each other's labels, so results can shift slightly from a single native pass.
+- **Chunking.** For text over 512 tokens: overlapping word chunks (32 words of overlap), each classified,
+  then merged: single-label heads take the most confident chunk, multi-label heads take the per-label
+  maximum. This is a heuristic; no chunk sees the whole text.
+
+## Layout
+
+```
+models/
+  GLiNER2.5-Decide-MultiFn-fp16.mlpackage   functions L64 / L128 / L256 / L512 (default L128)
+  tokenizer/                                DeBERTa-v3 tokenizer + GLiNER2 special tokens
+src/gliner_decide_coreml/
+  router.py      DecideRouter: bucket selection, head grouping, chunking
+  encoding.py    GLiNER2 schema preprocessing and bucket padding
+  decoding.py    activations, thresholds, chunk merging
+tests/           router vs native answers (golden.json), chunking, head splitting
+scripts/
+  make_golden.py             record native PyTorch answers for the tests
+  build/convert_bucket.py    convert + verify one bucket (from the Fluid Inference tooling)
+  build/build_multifunction.py
+bench/           Swift latency and compute-plan tools, Python router benchmark
+research/ane-ceiling/        Neural Engine feasibility experiment (see below)
+```
+
+## Tests and benchmarks
+
+```bash
+uv run pytest                          # router vs native answers, all buckets
+uv run python bench/router_latency.py  # end-to-end latency per bucket
+
+cd bench
+swiftc -O benchbuckets.swift -o benchbuckets   # model-only latency of compiled packages
+swiftc -O plan.swift -o plan                   # per-op device placement (MLComputePlan)
+```
+
+`scripts/make_golden.py` regenerates `tests/golden.json` from the native model (downloads the pinned
+checkpoint, ~1.7 GB).
+
+## Rebuilding the package
+
+```bash
+cd scripts/build
+for L in 64 128 256 512; do
+  uv run python convert_bucket.py --length $L --max-heads 4 --max-options 32 --output-dir build
+done
+uv run python build_multifunction.py --build-dir build
+```
+
+Source checkpoint: `fastino/GLiNER2.5-Decide` at revision `65624f1a0265b3f612bae66a2685a06b94a68a9d`.
+Pinned toolchain: coremltools 9.0, torch 2.7.0, transformers 4.57.6, gliner2 2.0.0.
+
+## Research notes
+
+Measured on an M1 Pro while building this package.
+
+- **Compute units.** Under `ALL`, Core ML splits every layer between the GPU and the Neural Engine (48
+  switches per request), which is slower than the GPU alone. Under `CPU_AND_NEURAL_ENGINE` the model takes
+  ~600 ms: DeBERTa's disentangled attention uses `gather_along_axis` (49 ops: the c2p and p2c
+  relative-position lookups in all 24 layers, plus the label-marker gather), which the Neural Engine cannot run.
+- **Neural Engine ceiling** (`research/ane-ceiling`, random-weight encoders with this model's shape). A plain
+  encoder runs in **32–36 ms** on the Neural Engine, and adding DeBERTa's extra position matmuls costs only
+  +6 ms. Replacing the gathers exactly with gather-free "relative shift" code keeps every op on the Neural
+  Engine but is slow: reshape-based shift **807 ms**, barrel shifter (static slices + blends) **187 ms**.
+  An exact ANE rewrite would only beat the GPU if the shift cost less than ~1 ms per layer.
+- **GPU headroom.** The model needs roughly 175 GFLOP per 256-token request; at the M1 Pro GPU's peak that is a
+  ~35 ms floor, versus 61–63 ms achieved. Bigger wins come from shorter inputs (this package) or a smaller,
+  distilled model, not from kernel tuning.
+
+## Limitations
+
+- Classification only: GLiNER2's entity, relation and structure extraction heads are not included.
+- Preprocessing uses the GLiNER2 Python package, which imports PyTorch (no model weights are loaded).
+- Latency figures are from one machine (M1 Pro); newer chips are faster and may favour different compute units.
+- Multifunction packages need macOS 15 / iOS 18 or later.
+
+## License and attribution
+
+Apache-2.0 (see `LICENSE` and `NOTICE`). Model by [Fastino](https://huggingface.co/fastino/GLiNER2.5-Decide);
+original Core ML conversion tooling by [Fluid Inference](https://huggingface.co/FluidInference/gliner2-5-decide-coreml).

bench/benchbuckets.swift

@@ -0,0 +1,59 @@
+// Latency of each GLiNER2.5-Decide bucket package with valid, full-length inputs.
+// Usage: benchbuckets <units,...> <iterations> <pkg.mlmodelc>...
+import CoreML
+import Foundation
+
+let args = CommandLine.arguments
+let unitsByName: [String: MLComputeUnits] = [
+    "all": .all, "cpuAndGPU": .cpuAndGPU, "cpuOnly": .cpuOnly, "cpuAndNeuralEngine": .cpuAndNeuralEngine,
+]
+let unitNames = args[1].split(separator: ",").map(String.init)
+let iters = Int(args[2])!
+let packages = args.dropFirst(3).map { URL(fileURLWithPath: $0) }
+
+func filled(_ shape: [NSNumber], _ type: MLMultiArrayDataType, _ value: (Int) -> Double) throws -> MLMultiArray {
+    let a = try MLMultiArray(shape: shape, dataType: type)
+    for i in 0..<a.count { a[i] = NSNumber(value: value(i)) }
+    return a
+}
+
+// Every position is a real (unmasked) token: the worst case for a bucket. Padding does not
+// change the compute, so this is also the cost of any request routed to this bucket.
+func inputs(for desc: MLModelDescription) throws -> MLFeatureProvider {
+    let ids = desc.inputDescriptionsByName["input_ids"]!.multiArrayConstraint!.shape
+    let grid = desc.inputDescriptionsByName["marker_indices"]!.multiArrayConstraint!.shape
+    let L = ids[1].intValue, K = grid[2].intValue
+    let used = [5, 3, 0, 0]  // two questions with 5 and 3 labels, like the example request
+    return try MLDictionaryFeatureProvider(dictionary: [
+        "input_ids": try filled(ids, .int32) { _ in Double(Int.random(in: 1000..<100_000)) },
+        "attention_mask": try filled(ids, .int32) { _ in 1 },
+        "marker_indices": try filled(grid, .int32) { i in
+            let h = i / K, k = i % K
+            return k < used[h] ? Double(min(L - 1, 2 + h * 12 + k * 2)) : 0
+        },
+        "marker_mask": try filled(grid, .float32) { i in i % K < used[i / K] ? 1 : 0 },
+    ])
+}
+
+print("bucket   units        load     p50       p90")
+for url in packages {
+    for label in unitNames {
+        let config = MLModelConfiguration()
+        config.computeUnits = unitsByName[label]!
+        let t0 = Date()
+        let model = try MLModel(contentsOf: url, configuration: config)
+        let load = Date().timeIntervalSince(t0)
+        let input = try inputs(for: model.modelDescription)
+        let L = model.modelDescription.inputDescriptionsByName["input_ids"]!.multiArrayConstraint!.shape[1]
+        for _ in 0..<5 { _ = try model.prediction(from: input) }
+        var t: [Double] = []
+        for _ in 0..<iters {
+            let s = Date()
+            _ = try model.prediction(from: input)
+            t.append(Date().timeIntervalSince(s) * 1000)
+        }
+        t.sort()
+        print(String(format: "L%-6@  %-10@  %5.1fs  %6.1f ms  %6.1f ms",
+                     L.stringValue as NSString, label as NSString, load, t[t.count / 2], t[Int(Double(t.count) * 0.9)]))
+    }
+}

bench/plan.swift

@@ -0,0 +1,71 @@
+// Step 0 diagnostic: where does CoreML place each op of the GLiNER2.5-Decide package?
+// Usage: plan <path.mlpackage|.mlmodelc> <out.json>
+import CoreML
+import Foundation
+
+func deviceName(_ d: MLComputeDevice) -> String {
+    switch d {
+    case .cpu: return "CPU"
+    case .gpu: return "GPU"
+    case .neuralEngine: return "ANE"
+    @unknown default: return "?"
+    }
+}
+
+struct OpRecord: Codable {
+    let index: Int
+    let op: String
+    let outputs: [String]
+    let preferred: String
+    let supported: [String]
+    let cost: Double
+}
+
+func walk(_ block: MLModelStructure.Program.Block, plan: MLComputePlan, into records: inout [OpRecord]) {
+    for op in block.operations {
+        let usage = plan.deviceUsage(for: op)
+        let cost = plan.estimatedCost(of: op)?.weight ?? 0
+        records.append(OpRecord(
+            index: records.count,
+            op: op.operatorName,
+            outputs: op.outputs.map { $0.name },
+            preferred: usage.map { deviceName($0.preferred) } ?? "none",
+            supported: usage.map { $0.supported.map(deviceName) } ?? [],
+            cost: cost))
+        for inner in op.blocks { walk(inner, plan: plan, into: &records) }
+    }
+}
+
+let args = CommandLine.arguments
+let modelURL = URL(fileURLWithPath: args[1])
+let outURL = URL(fileURLWithPath: args[2])
+
+let compiledURL: URL
+if modelURL.pathExtension == "mlmodelc" {
+    compiledURL = modelURL
+} else {
+    let t0 = Date()
+    let tmp = try await MLModel.compileModel(at: modelURL)
+    compiledURL = modelURL.deletingPathExtension().appendingPathExtension("mlmodelc")
+    try? FileManager.default.removeItem(at: compiledURL)
+    try FileManager.default.moveItem(at: tmp, to: compiledURL)
+    print("compiled in \(String(format: "%.1f", Date().timeIntervalSince(t0)))s -> \(compiledURL.lastPathComponent)")
+}
+
+var result: [String: [OpRecord]] = [:]
+for (label, units) in [("cpuAndNeuralEngine", MLComputeUnits.cpuAndNeuralEngine), ("all", MLComputeUnits.all)] {
+    let config = MLModelConfiguration()
+    config.computeUnits = units
+    let plan = try await MLComputePlan.load(contentsOf: compiledURL, configuration: config)
+    guard case let .program(program) = plan.modelStructure, let main = program.functions["main"] else {
+        fatalError("not an ML program")
+    }
+    var records: [OpRecord] = []
+    walk(main.block, plan: plan, into: &records)
+    result[label] = records
+    print("\(label): \(records.count) ops")
+}
+let enc = JSONEncoder()
+enc.outputFormatting = [.prettyPrinted, .sortedKeys]
+try enc.encode(result).write(to: outURL)
+print("wrote \(outURL.path)")

bench/router_latency.py

@@ -0,0 +1,33 @@
+"""End-to-end router latency (tokenize + route + Core ML predict + decode) per bucket, from Python.
+
+    uv run python bench/router_latency.py
+"""
+import statistics
+import sys
+import time
+import warnings
+from pathlib import Path
+
+warnings.filterwarnings("ignore")
+sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "tests"))
+
+from cases import CASES  # noqa: E402
+from gliner_decide_coreml import DecideRouter  # noqa: E402
+
+router = DecideRouter()
+seen = set()
+print(f"{'case':<28}{'tokens':>7}{'bucket':>8}{'p50':>10}{'p90':>10}")
+for case in CASES:
+    if case["bucket"] in seen:
+        continue
+    seen.add(case["bucket"])
+    for _ in range(5):
+        router.classify(case["text"], case["tasks"])
+    times = []
+    for _ in range(30):
+        start = time.perf_counter()
+        _, route = router.classify_with_route(case["text"], case["tasks"])
+        times.append((time.perf_counter() - start) * 1000)
+    times.sort()
+    print(f"{case['id']:<28}{route.tokens:>7}{route.bucket:>8}{statistics.median(times):>8.1f}ms"
+          f"{times[int(len(times) * 0.9)]:>8.1f}ms")

bench/time_prep.py

@@ -0,0 +1,25 @@
+"""Time the Python preprocessing (tokenize + build schema + pad) for requests of different sizes."""
+import statistics
+import time
+import warnings
+
+warnings.filterwarnings("ignore")
+from preprocessing import load_processor, prepare_decision  # noqa: E402
+
+proc = load_processor(".")
+tasks = {"intent": ["order_status", "refund_request", "cancel_subscription", "update_payment", "other"],
+         "urgency": ["low", "normal", "high"]}
+sentence = ("My subscription renewed on April 15 for 5,400 yen after the service was already down. "
+            "Can I get that charge refunded? ")
+
+for bucket, repeats in [(64, 1), (128, 3), (256, 8), (512, 18)]:
+    text = (sentence * repeats).strip()
+    for _ in range(5):
+        arrays = prepare_decision(proc, text, tasks, bucket, 4, 32)
+    times = []
+    for _ in range(50):
+        t = time.perf_counter()
+        arrays = prepare_decision(proc, text, tasks, bucket, 4, 32)
+        times.append((time.perf_counter() - t) * 1000)
+    n = int(arrays["attention_mask"].sum())
+    print(f"L{bucket:<4} {n:4d} real tokens   prep p50 {statistics.median(times):5.2f} ms")

models/GLiNER2.5-Decide-MultiFn-fp16.mlpackage/Data/com.apple.CoreML/model.mlmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2af62288520359e51612dcde770266a181de4922dd637b9e33bafebbf2ed734e
+size 1585697

models/GLiNER2.5-Decide-MultiFn-fp16.mlpackage/Data/com.apple.CoreML/weights/weight.bin

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

models/GLiNER2.5-Decide-MultiFn-fp16.mlpackage/Manifest.json

@@ -0,0 +1,18 @@
+{
+    "fileFormatVersion": "1.0.0",
+    "itemInfoEntries": {
+        "1F0D1536-6E76-4AF5-A08C-54426203B6AD": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Weights",
+            "name": "weights",
+            "path": "com.apple.CoreML/weights"
+        },
+        "D12A5651-2419-4D92-A578-C314C16AE33B": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Specification",
+            "name": "model.mlmodel",
+            "path": "com.apple.CoreML/model.mlmodel"
+        }
+    },
+    "rootModelIdentifier": "D12A5651-2419-4D92-A578-C314C16AE33B"
+}

models/tokenizer/config.json

@@ -0,0 +1,24 @@
+{
+  "_attn_implementation_autoset": true,
+  "_name_or_path": "/home/urchadezaratiana/checkpoints/checkpoint-45500",
+  "architecture": "span",
+  "architecture_version": 1,
+  "architectures": [
+    "SpanExtractor"
+  ],
+  "attn_implementation": "sdpa",
+  "config_version": 3,
+  "counting_layer": "count_lstm",
+  "max_len": null,
+  "max_width": 8,
+  "model_name": "microsoft/deberta-v3-large",
+  "model_type": "extractor",
+  "span_head": {
+    "dropout": 0.1,
+    "max_width": 8,
+    "span_mode": "markerV0"
+  },
+  "token_pooling": "first",
+  "transformers_version": "4.48.1",
+  "use_moe": false
+}

models/tokenizer/special_tokens_map.json

@@ -0,0 +1,123 @@
+{
+  "additional_special_tokens": [
+    {
+      "content": "[SEP_STRUCT]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "[SEP_TEXT]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "[P]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "[C]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "[E]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "[R]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "[L]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "[EXAMPLE]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "[OUTPUT]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "[DESCRIPTION]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    }
+  ],
+  "bos_token": {
+    "content": "[CLS]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "cls_token": {
+    "content": "[CLS]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "[SEP]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "mask_token": {
+    "content": "[MASK]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "[PAD]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "sep_token": {
+    "content": "[SEP]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "[UNK]",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  }
+}

models/tokenizer/tokenizer_config.json

@@ -0,0 +1,156 @@
+{
+  "add_prefix_space": true,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "[PAD]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "[CLS]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "[SEP]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "[UNK]",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128000": {
+      "content": "[MASK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128001": {
+      "content": "[SEP_STRUCT]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128002": {
+      "content": "[SEP_TEXT]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128003": {
+      "content": "[P]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128004": {
+      "content": "[C]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128005": {
+      "content": "[E]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128006": {
+      "content": "[R]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128007": {
+      "content": "[L]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128008": {
+      "content": "[EXAMPLE]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128009": {
+      "content": "[OUTPUT]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128010": {
+      "content": "[DESCRIPTION]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [
+    "[SEP_STRUCT]",
+    "[SEP_TEXT]",
+    "[P]",
+    "[C]",
+    "[E]",
+    "[R]",
+    "[L]",
+    "[EXAMPLE]",
+    "[OUTPUT]",
+    "[DESCRIPTION]"
+  ],
+  "backend": "tokenizers",
+  "bos_token": "[CLS]",
+  "clean_up_tokenization_spaces": false,
+  "cls_token": "[CLS]",
+  "do_lower_case": false,
+  "eos_token": "[SEP]",
+  "extra_special_tokens": {},
+  "is_local": false,
+  "local_files_only": false,
+  "mask_token": "[MASK]",
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "sp_model_kwargs": {},
+  "split_by_punct": false,
+  "tokenizer_class": "DebertaV2Tokenizer",
+  "unk_id": 3,
+  "unk_token": "[UNK]",
+  "vocab_type": "spm"
+}

pyproject.toml

@@ -0,0 +1,34 @@
+[project]
+name = "gliner-decide-coreml"
+version = "0.1.0"
+description = "GLiNER2.5-Decide as one Core ML multifunction package (64/128/256/512-token buckets) with automatic bucket routing, tuned for Apple Silicon GPUs."
+readme = "README.md"
+license = "Apache-2.0"
+requires-python = ">=3.10,<3.13"
+# Pinned to the versions the Core ML buckets were converted and verified with.
+# torch/transformers are needed only because GLiNER2's schema preprocessing imports them.
+dependencies = [
+    "coremltools==9.0",
+    "gliner2[local]==2.0.0",
+    "huggingface-hub>=0.34,<1",
+    "numpy<2.3",
+    "sentencepiece>=0.2,<0.3",
+    "torch==2.7.0",
+    "transformers==4.57.6",
+]
+
+[dependency-groups]
+dev = [
+    "pytest>=8",
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/gliner_decide_coreml"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+filterwarnings = ["ignore::DeprecationWarning", "ignore::UserWarning", "ignore::FutureWarning"]

research/ane-ceiling/models.py

@@ -0,0 +1,250 @@
+"""ANE ceiling test: DeBERTa-v3-large-shaped encoders with random weights.
+
+Variants
+  a-std : plain post-LN encoder, standard (B, S, C) layout with nn.Linear
+  a-ane : same math, Apple's ANE layout (B, C, 1, S), 1x1 Conv2d, per-head attention
+  b-ane : a-ane + DeBERTa's c2p / p2c terms, computed gather-free:
+          constant per-layer position tables (one row per possible distance)
+          + the "relative shift" (reshape / slice) instead of gather_along_axis
+Speed does not depend on weight values, so random weights give real timings.
+"""
+import argparse
+import math
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+C, H, D, FF = 1024, 16, 64, 4096  # DeBERTa-v3-large dims (from the port's encoder_config)
+
+
+def skew(t: torch.Tensor, S: int) -> torch.Tensor:
+    """(B, S, 1, 2S-1) scores against every distance -> (B, S, 1, S).
+
+    out[b, i, 0, j] = t[b, i, 0, j - i + S - 1]; done with reshape/slice only.
+    Row i of the flattened tensor starts at i*(2S-1); the wanted entry is at
+    i*(2S-2) + (S-1) + j, i.e. rows of stride 2S-2 once we drop the first S-1.
+    """
+    B = t.shape[0]
+    f = t.reshape(B, S * (2 * S - 1))
+    f = f[:, S - 1 : S - 1 + S * (2 * S - 2)]
+    f = f.reshape(B, S, 2 * S - 2)[:, :, :S]
+    return f.reshape(B, S, 1, S)
+
+
+def shift_last(x: torch.Tensor, S: int, N: int) -> torch.Tensor:
+    """Barrel shifter along the last axis, no reshapes.
+
+    x: (B, N*S, 1, 2S-1), rows grouped per head.  out[.., n*S+i, 0, j] = x[.., n*S+i, 0, j + S-1-i].
+    The per-row shift (0..S-1) is applied in log2(S) steps of 1, 2, 4, ...; each step is two
+    static slices and a select with a constant row mask.
+    """
+    w = 2 * S - 1  # static ints only: traced shapes would become dynamic ops
+    for b, (m, inv) in enumerate(barrel_masks(S, N, "rows")):
+        step = 1 << b
+        w -= step
+        x = x[..., step:step + w] * m + x[..., :w] * inv
+    return x
+
+
+def shift_chan(x: torch.Tensor, S: int) -> torch.Tensor:
+    """Barrel shifter along the channel axis (lands p2c in query-major order, no transpose).
+
+    x: (N, 2S-1, 1, S).  out[n, i, 0, j] = x[n, i + S-1-j, 0, j].
+    """
+    c = 2 * S - 1
+    for b, (m, inv) in enumerate(barrel_masks(S, 1, "cols")):
+        step = 1 << b
+        c -= step
+        x = x[:, step:step + c] * m + x[:, :c] * inv
+    return x
+
+
+_MASKS = {}
+
+
+def barrel_masks(S: int, N: int, kind: str):
+    """Constant 0/1 masks (and their complements) for each barrel step.
+
+    Built once outside tracing so they enter the graph as plain constants; blending with
+    m * hi + (1 - m) * lo is exact because one of the two products is always zero.
+    """
+    key = (S, N, kind)
+    if key not in _MASKS:
+        shift = S - 1 - torch.arange(S)
+        shape = (1, N * S, 1, 1) if kind == "rows" else (1, 1, 1, S)
+        steps = []
+        b = 0
+        while (1 << b) < S:
+            m = ((shift >> b) & 1).repeat(N).float().view(shape)
+            steps.append((m, 1.0 - m))
+            b += 1
+        _MASKS[key] = steps
+    return _MASKS[key]
+
+
+class LayerNormANE(nn.Module):
+    """LayerNorm over the channel axis of a (B, C, 1, S) tensor."""
+
+    def __init__(self, c, eps=1e-5):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(c))
+        self.bias = nn.Parameter(torch.zeros(c))
+
+    def forward(self, x):
+        zc = x - x.mean(dim=1, keepdim=True)
+        var = (zc * zc).mean(dim=1, keepdim=True)
+        out = zc * torch.rsqrt(var + self.eps)
+        return out * self.weight.view(1, -1, 1, 1) + self.bias.view(1, -1, 1, 1)
+
+
+class LayerANE(nn.Module):
+    """rel: None (plain), "skew" (reshape shift), "barrel" (slice/select shift),
+    "noshift" (extra matmuls only, shift skipped: wrong math, timing control)."""
+
+    def __init__(self, S: int, rel=None):
+        super().__init__()
+        self.S, self.rel = S, rel
+        self.q, self.k, self.v, self.o = (nn.Conv2d(C, C, 1) for _ in range(4))
+        self.ln1, self.ln2 = LayerNormANE(C), LayerNormANE(C)
+        self.ff1, self.ff2 = nn.Conv2d(C, FF, 1), nn.Conv2d(FF, C, 1)
+        # DeBERTa scales by sqrt(d * (1 + number of position terms)).
+        self.scale = 1.0 / math.sqrt(D * (3 if rel else 1))
+        if rel:
+            # Precomputed LN(rel_embeddings) @ W_k / W_q, expanded to every distance.
+            self.register_buffer("pk", torch.randn(1, C, 1, 2 * S - 1) * 0.02)
+            self.register_buffer("pq", torch.randn(1, C, 1, 2 * S - 1) * 0.02)
+            self.register_buffer("pqT", torch.randn(1, 2 * S - 1, 1, C) * 0.02)
+        if rel == "barrel":
+            # Build the masks now, so tracing captures them as constants rather than ops.
+            barrel_masks(S, H, "rows")
+            barrel_masks(S, 1, "cols")
+
+    def rel_terms(self, qT, ks, kT):
+        """Per-head (c2p, p2c) score tables, each (B, Sq, 1, Sk)."""
+        S = self.S
+        pks = self.pk.split(D, dim=1)
+        if self.rel == "skew":
+            pqs = self.pq.split(D, dim=1)
+            c2p = [skew(torch.einsum("bchr,bqhc->bqhr", pks[h], qT[h]), S) for h in range(H)]
+            p2c = [skew(torch.einsum("bchr,bkhc->bkhr", pqs[h], kT[h]), S).transpose(1, 3)
+                   for h in range(H)]
+            return c2p, p2c
+        # Heads batched so each shift runs once per layer, not once per head.
+        a = torch.cat([torch.einsum("bchr,bqhc->bqhr", pks[h], qT[h]) for h in range(H)], dim=1)
+        pqT = self.pqT.split(D, dim=3)
+        m = torch.cat([torch.einsum("bchk,brhc->brhk", ks[h], pqT[h]) for h in range(H)], dim=0)
+        if self.rel == "barrel":
+            a, m = shift_last(a, S, H), shift_chan(m, S)
+        else:  # noshift
+            a, m = a[..., :S], m[:, :S]
+        return a.split(S, dim=1), m.split(1, dim=0)
+
+    def forward(self, x, mask):  # x: (B, C, 1, S); mask: (B, 1, 1, S) additive
+        q, k, v = self.q(x), self.k(x), self.v(x)
+        qT = q.transpose(1, 3).split(D, dim=3)  # H x (B, S, 1, D)
+        ks, vs = k.split(D, dim=1), v.split(D, dim=1)  # H x (B, D, 1, S)
+        if self.rel:
+            kT = k.transpose(1, 3).split(D, dim=3) if self.rel == "skew" else None
+            c2p, p2c = self.rel_terms(qT, ks, kT)
+        heads = []
+        for h in range(H):
+            s = torch.einsum("bchk,bqhc->bqhk", ks[h], qT[h])  # (B, Sq, 1, Sk)
+            if self.rel:
+                s = s + c2p[h] + p2c[h]
+            w = (s * self.scale + mask).softmax(dim=3)
+            heads.append(torch.einsum("bqhk,bchk->bchq", w, vs[h]))  # (B, D, 1, Sq)
+        x = self.ln1(x + self.o(torch.cat(heads, dim=1)))
+        return self.ln2(x + self.ff2(F.gelu(self.ff1(x))))
+
+
+class LayerStd(nn.Module):
+    def __init__(self, S: int):
+        super().__init__()
+        self.q, self.k, self.v, self.o = (nn.Linear(C, C) for _ in range(4))
+        self.ln1, self.ln2 = nn.LayerNorm(C), nn.LayerNorm(C)
+        self.ff1, self.ff2 = nn.Linear(C, FF), nn.Linear(FF, C)
+        self.scale = 1.0 / math.sqrt(D)
+
+    def forward(self, x, mask):  # x: (B, S, C); mask: (B, 1, 1, S)
+        B, S, _ = x.shape
+        q = self.q(x).view(B, S, H, D).transpose(1, 2)
+        k = self.k(x).view(B, S, H, D).transpose(1, 2)
+        v = self.v(x).view(B, S, H, D).transpose(1, 2)
+        w = (q @ k.transpose(-1, -2) * self.scale + mask).softmax(-1)
+        a = (w @ v).transpose(1, 2).reshape(B, S, C)
+        x = self.ln1(x + self.o(a))
+        return self.ln2(x + self.ff2(F.gelu(self.ff1(x))))
+
+
+class Encoder(nn.Module):
+    def __init__(self, variant: str, S: int, layers: int):
+        super().__init__()
+        make = {"a-std": lambda: LayerStd(S),
+                "a-ane": lambda: LayerANE(S),
+                "b-ane": lambda: LayerANE(S, "skew"),
+                "c-ane": lambda: LayerANE(S, "barrel"),
+                "n-ane": lambda: LayerANE(S, "noshift")}[variant]
+        self.layers = nn.ModuleList(make() for _ in range(layers))
+
+    def forward(self, x, mask):
+        for layer in self.layers:
+            x = layer(x, mask)
+        return x
+
+
+def check_skew():
+    """The reshape/slice shift must equal the gather it replaces."""
+    S = 7
+    t = torch.randn(2, S, 1, 2 * S - 1)
+    i = torch.arange(S).view(S, 1)
+    j = torch.arange(S).view(1, S)
+    idx = (j - i + S - 1).expand(2, S, S).unsqueeze(2)  # (B, S, 1, S)
+    ref = torch.gather(t, 3, idx)
+    assert torch.equal(skew(t, S), ref), "skew != gather"
+
+    S, N = 8, 3  # barrel shifters need S to be a power of two
+    t = torch.randn(1, N * S, 1, 2 * S - 1)
+    i = torch.arange(S).view(S, 1)
+    j = torch.arange(S).view(1, S)
+    idx = (j - i + S - 1).repeat(N, 1).view(1, N * S, 1, S)
+    assert torch.equal(shift_last(t, S, N), torch.gather(t, 3, idx)), "shift_last != gather"
+    t = torch.randn(N, 2 * S - 1, 1, S)
+    idx = (i + S - 1 - j).view(1, S, 1, S).expand(N, S, 1, S)
+    assert torch.equal(shift_chan(t, S), torch.gather(t, 1, idx)), "shift_chan != gather"
+    print("skew, shift_last, shift_chan == gather: OK")
+
+
+if __name__ == "__main__":
+    import coremltools as ct
+
+    p = argparse.ArgumentParser()
+    p.add_argument("variant", choices=["a-std", "a-ane", "b-ane", "c-ane", "n-ane"])
+    p.add_argument("--seq", type=int, default=256)
+    p.add_argument("--layers", type=int, default=24)
+    p.add_argument("--out", required=True)
+    a = p.parse_args()
+
+    check_skew()
+    torch.manual_seed(0)
+    model = Encoder(a.variant, a.seq, a.layers).eval()
+    n = sum(p.numel() for p in model.parameters()) + sum(b.numel() for b in model.buffers())
+    print(f"{a.variant}: {a.layers} layers, seq {a.seq}, {n/1e6:.1f}M params+constants")
+
+    x = torch.randn(1, a.seq, C) if a.variant == "a-std" else torch.randn(1, C, 1, a.seq)
+    mask = torch.zeros(1, 1, 1, a.seq)
+    with torch.no_grad():
+        traced = torch.jit.trace(model, (x, mask))
+    ml = ct.convert(
+        traced,
+        inputs=[ct.TensorType(name="hidden", shape=x.shape, dtype=np.float16),
+                ct.TensorType(name="mask", shape=mask.shape, dtype=np.float16)],
+        outputs=[ct.TensorType(name="out", dtype=np.float16)],
+        convert_to="mlprogram",
+        compute_precision=ct.precision.FLOAT16,
+        minimum_deployment_target=ct.target.macOS14,
+    )
+    ml.save(a.out)
+    print("saved", a.out)

research/ane-ceiling/summarize.py

@@ -0,0 +1,12 @@
+import json, sys, collections
+d = json.load(open(sys.argv[1]))
+for mode in ("cpuAndNeuralEngine", "all"):
+    ops = [o for o in d[mode] if o["op"] != "const"]
+    dev = collections.Counter(o["preferred"] for o in ops)
+    seq = [o["preferred"] for o in ops]
+    sw = sum(1 for a, b in zip(seq, seq[1:]) if a != b)
+    off = collections.Counter(o["op"] for o in ops if o["preferred"] != "ANE")
+    nosup = collections.Counter(o["op"] for o in ops if "ANE" not in o["supported"])
+    print(f"  {mode:19} ops={len(ops):5} {dict(dev)} switches={sw}")
+    if off: print(f"    not on ANE: {dict(off)}")
+    if nosup: print(f"    ANE can't run: {dict(nosup)}")

scripts/build/build_multifunction.py

@@ -0,0 +1,40 @@
+"""Merge the four bucket packages into one multifunction package with shared weights.
+
+Run convert_bucket.py for each length first, e.g.
+
+    cd scripts/build
+    for L in 64 128 256 512; do
+        uv run python convert_bucket.py --length $L --max-heads 4 --max-options 32 --output-dir build
+    done
+    uv run python build_multifunction.py --build-dir build
+
+Each function (L64, L128, L256, L512) keeps its own fixed-shape graph; identical weight
+tensors are stored once (~0.9 GB instead of ~3.5 GB for four separate packages).
+"""
+import argparse
+from pathlib import Path
+
+import coremltools as ct
+
+from runtime import package_name
+
+BUCKETS = (64, 128, 256, 512)
+ROOT = Path(__file__).resolve().parents[2]
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--build-dir", default="build")
+    parser.add_argument("--output", default=str(ROOT / "models" / "GLiNER2.5-Decide-MultiFn-fp16.mlpackage"))
+    args = parser.parse_args()
+    desc = ct.utils.MultiFunctionDescriptor()
+    for length in BUCKETS:
+        source = Path(args.build_dir) / package_name("fp16", length, 4, 32)
+        desc.add_function(str(source), src_function_name="main", target_function_name=f"L{length}")
+    desc.default_function_name = "L128"
+    ct.utils.save_multifunction(desc, args.output)
+    print("saved", args.output)
+
+
+if __name__ == "__main__":
+    main()

scripts/build/convert_bucket.py

@@ -0,0 +1,189 @@
+"""Convert and verify the pinned GLiNER2.5-Decide multi-head classification path.
+
+Modified from Fluid Inference's gliner2-5-decide-coreml `convert-coreml.py` (Apache-2.0):
+traces with an example that fits the requested length and verifies every example that fits,
+so buckets shorter than the original examples (L64) can be built and checked.
+"""
+import argparse
+import json
+import math
+from pathlib import Path
+
+import coremltools as ct
+import numpy as np
+import torch
+from gliner2 import AutoExtractor
+from huggingface_hub import snapshot_download
+from transformers.models.deberta_v2 import modeling_deberta_v2
+
+from convert_names import MODEL_ID, MODEL_REVISION
+from export_model import GLiNER2DecideExport, coreml_safe_attention_forward
+from preprocessing import classification_schema, prepare_decision
+from runtime import decode, package_name
+
+EXAMPLES = [
+    ("My subscription renewed on April 15 for ¥5,400 after the service was already down. "
+     "Can I get that charge refunded?",
+     {"intent": ["order_status", "refund_request", "cancel_subscription", "update_payment", "login_problem",
+                 "shipping_delay", "bug_report", "speak_to_human"]}),
+    ("Battery dies before lunch, but the keyboard and the screen are the best I have used on a laptop.",
+     {"sentiment": ["positive", "negative", "mixed", "neutral"],
+      "aspects": {"labels": ["battery", "keyboard", "screen", "camera", "price", "support"],
+                  "multi_label": True, "cls_threshold": 0.4}}),
+    ("From: compliance@group.example\nSubject: Protocol update — action required today\n\n"
+     "Please confirm the new retention rule is applied before Friday's audit.",
+     {"intent": ["fyi", "request", "approval", "complaint", "newsletter", "security_alert"],
+      "urgency": ["low", "normal", "high", "critical"],
+      "route": ["support", "billing", "legal", "security", "finance", "archive"]}),
+    ("My subscription renewed on April 15 for 5,400 yen after the service was already down. "
+     "Can I get that charge refunded?",
+     {"intent": ["order_status", "refund_request", "cancel_subscription", "update_payment", "other"],
+      "urgency": ["low", "normal", "high"]}),
+]
+
+
+def fits(processor, text, tasks, bucket):
+    try:
+        prepare_decision(processor, text, tasks, *bucket)
+        return True
+    except ValueError:
+        return False
+
+
+def native_logits(native, text, tasks):
+    """Per-head native logits via the upstream span collator, encoder and shared classifier."""
+    from gliner2.training.trainer import ExtractorCollator
+
+    collator = ExtractorCollator(native.processor, is_training=False, max_len=None, architecture="span")
+    batch = collator([(text, classification_schema(tasks).build())])
+    _, schema_embs = native._encode_batch(batch)
+    return [native.classifier(torch.stack(embs[1:])).squeeze(-1) for embs in schema_embs[0]]
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--output-dir", default="build")
+    parser.add_argument("--length", type=int, default=128)
+    parser.add_argument("--max-heads", type=int, default=4)
+    parser.add_argument("--max-options", type=int, default=8)
+    parser.add_argument("--precision", choices=["fp16", "fp32"], default="fp16")
+    args = parser.parse_args()
+    torch.set_num_threads(4)
+    source = snapshot_download(
+        MODEL_ID, revision=MODEL_REVISION,
+        allow_patterns=[
+            "config.json", "encoder_config/*", "model.safetensors", "tokenizer.json", "tokenizer_config.json",
+            "special_tokens_map.json",
+        ],
+    )
+    native = AutoExtractor.from_pretrained(source, map_location="cpu").eval()
+    wrapper = GLiNER2DecideExport(native).eval()
+    bucket = (args.length, args.max_heads, args.max_options)
+    fitting = [e for e in EXAMPLES if fits(native.processor, *e, bucket)]
+    if not fitting:
+        raise RuntimeError(f"no example fits bucket {bucket}")
+    print(f"examples fitting L{args.length}: {len(fitting)}/{len(EXAMPLES)}")
+    text, tasks = fitting[-1] if args.length < 128 else (EXAMPLES[2] if EXAMPLES[2] in fitting else fitting[-1])
+    arrays = prepare_decision(native.processor, text, tasks, *bucket)
+    tensors = tuple(torch.from_numpy(value) for value in arrays.values())
+
+    def wrapper_error():
+        with torch.no_grad():
+            expected = native_logits(native, text, tasks)
+            actual = wrapper(*tensors)[0][0]
+        return max(float((row - actual[h, : len(row)]).abs().max()) for h, row in enumerate(expected))
+
+    error = wrapper_error()
+    if error > 1e-4:
+        raise RuntimeError(f"Wrapper/native logit mismatch: {error}")
+    # The upstream scale is a constant for a fixed DeBERTa attention head width.
+    # Its traced int32 sqrt is rejected by Core ML; freeze the identical float32
+    # value while tracing, and restore the upstream implementation immediately.
+    original_scale = modeling_deberta_v2.scaled_size_sqrt
+    original_rpos = modeling_deberta_v2.build_rpos
+    original_attention = modeling_deberta_v2.DisentangledSelfAttention.forward
+
+    def static_scale(query_layer, scale_factor):
+        value = math.sqrt(float(query_layer.shape[-1] * scale_factor))
+        return torch.tensor(value, dtype=torch.float32, device=query_layer.device)
+
+    modeling_deberta_v2.scaled_size_sqrt = static_scale
+    # The encoder only uses self-attention: query and key sequence lengths are
+    # identical, so the scripted build_rpos returns relative_pos unchanged.
+    modeling_deberta_v2.build_rpos = lambda query, key, relative_pos, buckets, max_pos: relative_pos
+    modeling_deberta_v2.DisentangledSelfAttention.forward = coreml_safe_attention_forward
+    try:
+        frozen_error = wrapper_error()
+        if frozen_error > 1e-4:
+            raise RuntimeError(f"Frozen attention scale changed native logits: {frozen_error}")
+        with torch.no_grad():
+            traced = torch.jit.trace(wrapper, tensors, check_trace=False)
+    finally:
+        modeling_deberta_v2.scaled_size_sqrt = original_scale
+        modeling_deberta_v2.build_rpos = original_rpos
+        modeling_deberta_v2.DisentangledSelfAttention.forward = original_attention
+    grid = (1, args.max_heads, args.max_options)
+    converted = ct.convert(
+        traced, convert_to="mlprogram", minimum_deployment_target=ct.target.iOS17,
+        compute_precision=ct.precision.FLOAT16 if args.precision == "fp16" else ct.precision.FLOAT32,
+        compute_units=ct.ComputeUnit.CPU_ONLY,
+        inputs=[
+            ct.TensorType(name="input_ids", shape=(1, args.length), dtype=np.int32),
+            ct.TensorType(name="attention_mask", shape=(1, args.length), dtype=np.int32),
+            ct.TensorType(name="marker_indices", shape=grid, dtype=np.int32),
+            ct.TensorType(name="marker_mask", shape=grid, dtype=np.float32),
+        ],
+        outputs=[ct.TensorType(name="logits", dtype=np.float32), ct.TensorType(name="probabilities", dtype=np.float32)],
+    )
+    converted.short_description = "GLiNER2.5-Decide multi-head schema classification path"
+    converted.author = "Fastino (original); Fluid Inference (Core ML conversion)"
+    converted.license = "Apache-2.0"
+    converted.user_defined_metadata.update({
+        "source_model": MODEL_ID, "source_revision": MODEL_REVISION,
+        "scope": "classification only; span and count heads not exported",
+        "length": str(args.length), "max_heads": str(args.max_heads), "max_options": str(args.max_options),
+    })
+    out = Path(args.output_dir)
+    out.mkdir(parents=True, exist_ok=True)
+    package = out / package_name(args.precision, *bucket)
+    converted.save(str(package))
+    runtime = ct.models.MLModel(str(package), compute_units=ct.ComputeUnit.ALL)
+    cases = []
+    for text, tasks in fitting:
+        native_output = native.classify_text(text, tasks, include_confidence=True)
+        logits = np.asarray(runtime.predict(prepare_decision(native.processor, text, tasks, *bucket))["logits"])[0]
+        coreml_output = decode(tasks, logits)
+        if json.dumps(labels_only(native_output)) != json.dumps(labels_only(coreml_output)):
+            raise RuntimeError(f"Core ML/native mismatch: {coreml_output} != {native_output}")
+        cases.append({"text": text, "native": native_output, "coreml": coreml_output,
+                      "max_confidence_error": confidence_error(native_output, coreml_output)})
+    report = {
+        "source_model": MODEL_ID, "source_revision": MODEL_REVISION, "package": str(package),
+        "package_bytes": sum(f.stat().st_size for f in package.rglob("*") if f.is_file()),
+        "native_total_parameters": sum(p.numel() for p in native.parameters()),
+        "exported_parameters": sum(p.numel() for p in wrapper.parameters()),
+        "wrapper_max_logit_error": error, "coremltools": ct.__version__, "torch": torch.__version__,
+        "cases": cases,
+    }
+    (out / f"conversion-{package.stem}.json").write_text(json.dumps(report, indent=2, ensure_ascii=False) + "\n")
+    print(json.dumps(report, indent=2, ensure_ascii=False))
+
+
+def _entries(value):
+    return value if isinstance(value, list) else [value]
+
+
+def labels_only(result: dict) -> dict:
+    return {task: sorted(entry["label"] for entry in _entries(value)) for task, value in result.items()}
+
+
+def confidence_error(native: dict, coreml: dict) -> float:
+    errors = []
+    for task, value in native.items():
+        predicted = {entry["label"]: entry["confidence"] for entry in _entries(coreml[task])}
+        errors += [abs(entry["confidence"] - predicted[entry["label"]]) for entry in _entries(value)]
+    return max(errors)
+
+
+if __name__ == "__main__":
+    main()

scripts/build/convert_names.py

@@ -0,0 +1,3 @@
+"""Pinned source checkpoint shared by the conversion, verification and scoring scripts."""
+MODEL_ID = "fastino/GLiNER2.5-Decide"
+MODEL_REVISION = "65624f1a0265b3f612bae66a2685a06b94a68a9d"

scripts/build/export_model.py

@@ -0,0 +1,54 @@
+"""Native GLiNER2.5-Decide multi-head classification path with explicit marker routing."""
+import torch
+from torch import nn
+from transformers.models.deberta_v2 import modeling_deberta_v2
+
+
+def coreml_safe_attention_forward(
+    self, hidden_states, attention_mask, output_attentions=False,
+    query_states=None, relative_pos=None, rel_embeddings=None,
+):
+    """Native DeBERTa attention with a finite mask sentinel for FP16 Core ML."""
+    if query_states is None:
+        query_states = hidden_states
+    query = self.transpose_for_scores(self.query_proj(query_states), self.num_attention_heads)
+    key = self.transpose_for_scores(self.key_proj(hidden_states), self.num_attention_heads)
+    value = self.transpose_for_scores(self.value_proj(hidden_states), self.num_attention_heads)
+    factor = 1 + int("c2p" in self.pos_att_type) + int("p2c" in self.pos_att_type)
+    scale = modeling_deberta_v2.scaled_size_sqrt(query, factor)
+    scores = torch.bmm(query, key.transpose(-1, -2) / scale.to(dtype=query.dtype))
+    if self.relative_attention:
+        relative = self.disentangled_attention_bias(
+            query, key, relative_pos, self.pos_dropout(rel_embeddings), factor
+        )
+        scores = scores + relative
+    scores = scores.view(-1, self.num_attention_heads, scores.size(-2), scores.size(-1))
+    scores = scores.masked_fill(~attention_mask.bool(), -1e4)
+    probabilities = self.dropout(torch.softmax(scores, dim=-1))
+    context = torch.bmm(probabilities.view(-1, probabilities.size(-2), probabilities.size(-1)), value)
+    context = context.view(-1, self.num_attention_heads, context.size(-2), context.size(-1))
+    context = context.permute(0, 2, 1, 3).contiguous()
+    context = context.view(context.size()[:-2] + (-1,))
+    return (context, probabilities) if output_attentions else (context, None)
+
+
+class GLiNER2DecideExport(nn.Module):
+    """Encoder plus the shared label classifier, scored for every (head, label) marker at once.
+
+    Logits are returned per head so the host can apply softmax (single-label) or sigmoid
+    (multi-label) exactly as the native runtime does. Probabilities are the per-head softmax.
+    """
+
+    def __init__(self, native: nn.Module):
+        super().__init__()
+        self.encoder = native.encoder
+        self.classifier = native.classifier
+
+    def forward(self, input_ids, attention_mask, marker_indices, marker_mask):
+        hidden = self.encoder(input_ids=input_ids.long(), attention_mask=attention_mask.long()).last_hidden_state
+        heads, options = marker_indices.shape[1], marker_indices.shape[2]
+        flat = marker_indices.long().reshape(1, heads * options, 1).expand(-1, -1, hidden.shape[-1])
+        states = hidden.gather(1, flat)
+        logits = self.classifier(states).reshape(1, heads, options)
+        logits = torch.where(marker_mask > 0.5, logits, torch.full_like(logits, -1e4))
+        return logits, torch.softmax(logits, dim=-1)

scripts/build/preprocessing.py

@@ -0,0 +1,67 @@
+"""Native GLiNER2 span-architecture schema preprocessing for a fixed Core ML bucket."""
+import numpy as np
+from gliner2 import Schema
+from gliner2.models.base import load_extractor_tokenizer
+from gliner2.processor import SchemaTransformer
+from gliner2.training.trainer import ExtractorCollator
+
+
+def load_processor(tokenizer_dir: str):
+    """Load only the tokenizer and schema formatter needed by the Core ML model."""
+    return SchemaTransformer(tokenizer=load_extractor_tokenizer(tokenizer_dir), token_pooling="first")
+
+
+def classification_schema(tasks: dict) -> Schema:
+    """Same task-dict handling as native ``classify_text``."""
+    schema = Schema()
+    for name, config in tasks.items():
+        if isinstance(config, dict) and "labels" in config:
+            cfg = config.copy()
+            labels = cfg.pop("labels")
+            schema.classification(name, labels, **cfg)
+        else:
+            schema.classification(name, config)
+    return schema
+
+
+def task_labels(tasks: dict) -> dict[str, list[str]]:
+    result = {}
+    for name, config in tasks.items():
+        labels = config["labels"] if isinstance(config, dict) and "labels" in config else config
+        result[name] = list(labels.keys()) if isinstance(labels, dict) else list(labels)
+    return result
+
+
+def prepare_decision(processor, text: str, tasks: dict, length: int, max_heads: int, max_options: int):
+    """Tokenize ``tasks`` exactly as the native span collator does and pad into the bucket."""
+    if not 1 <= len(tasks) <= max_heads:
+        raise ValueError(f"Expected 1..{max_heads} decision heads, got {len(tasks)}")
+    labels = task_labels(tasks)
+    for name, values in labels.items():
+        if not 1 <= len(values) <= max_options:
+            raise ValueError(f"Head {name!r} has {len(values)} labels; bucket holds 1..{max_options}")
+    collator = ExtractorCollator(processor, is_training=False, max_len=None, architecture="span")
+    batch = collator([(text, classification_schema(tasks).build())])
+    ids = batch.input_ids.numpy()
+    if ids.shape[1] > length:
+        raise ValueError(f"Schema and text require {ids.shape[1]} subwords; bucket holds {length}")
+    groups = batch.schema_special_indices[0]
+    if len(groups) != len(labels):
+        raise ValueError("Schema head count does not match the requested tasks")
+    indices = np.zeros((1, max_heads, max_options), dtype=np.int32)
+    mask = np.zeros((1, max_heads, max_options), dtype=np.float32)
+    for head, (positions, values) in enumerate(zip(groups, labels.values())):
+        # positions[0] is the [P] prompt marker; the rest are one [L] marker per label.
+        markers = list(positions[1:])
+        if len(markers) != len(values):
+            raise ValueError("Label markers were truncated or merged")
+        indices[0, head, : len(markers)] = markers
+        mask[0, head, : len(markers)] = 1.0
+    attention = batch.attention_mask.numpy()
+    pad = processor.tokenizer.pad_token_id
+    return {
+        "input_ids": np.pad(ids, ((0, 0), (0, length - ids.shape[1])), constant_values=pad).astype(np.int32),
+        "attention_mask": np.pad(attention, ((0, 0), (0, length - attention.shape[1]))).astype(np.int32),
+        "marker_indices": indices,
+        "marker_mask": mask,
+    }

scripts/build/runtime.py

@@ -0,0 +1,71 @@
+"""Run the GLiNER2.5-Decide Core ML classifier without loading the original model weights."""
+import argparse
+import json
+from pathlib import Path
+
+import coremltools as ct
+import numpy as np
+
+from preprocessing import load_processor, prepare_decision, task_labels
+
+
+def package_name(precision: str, length: int, max_heads: int, max_options: int) -> str:
+    return f"gliner2_decide_classification_{precision}_L{length}_H{max_heads}_K{max_options}.mlpackage"
+
+
+def decode(tasks: dict, logits: np.ndarray) -> dict:
+    """Apply the native activation and label rules to per-head logits of shape (heads, options)."""
+    results = {}
+    for head, (name, labels) in enumerate(task_labels(tasks).items()):
+        config = tasks[name] if isinstance(tasks[name], dict) else {}
+        values = logits[head, : len(labels)].astype(np.float64)
+        activation = config.get("class_act", "auto")
+        multi = config.get("multi_label", False)
+        if activation == "sigmoid" or (activation == "auto" and multi):
+            probs = 1.0 / (1.0 + np.exp(-values))
+        else:
+            probs = np.exp(values - values.max())
+            probs /= probs.sum()
+        if multi:
+            threshold = config.get("cls_threshold", 0.5)
+            chosen = [{"label": labels[j], "confidence": float(probs[j])} for j in range(len(labels))
+                      if probs[j] >= threshold]
+            best = int(probs.argmax())
+            results[name] = chosen or [{"label": labels[best], "confidence": float(probs[best])}]
+        else:
+            best = int(probs.argmax())
+            results[name] = {"label": labels[best], "confidence": float(probs[best])}
+    return results
+
+
+class CoreMLDecide:
+    def __init__(self, model_dir: str, precision: str = "fp16", length: int = 128, max_heads: int = 4,
+                 max_options: int | None = None, compute_units=ct.ComputeUnit.ALL):
+        model_dir = Path(model_dir)
+        # Published buckets: L128 holds 8 labels per head, L256 and L512 hold 32.
+        max_options = max_options or (8 if length == 128 else 32)
+        self.length, self.max_heads, self.max_options = length, max_heads, max_options
+        self.processor = load_processor(str(model_dir))
+        package = model_dir / package_name(precision, length, max_heads, max_options)
+        self.model = ct.models.MLModel(str(package), compute_units=compute_units)
+
+    def classify(self, text: str, tasks: dict) -> dict:
+        arrays = prepare_decision(self.processor, text, tasks, self.length, self.max_heads, self.max_options)
+        logits = np.asarray(self.model.predict(arrays)["logits"])[0]
+        return decode(tasks, logits)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-dir", required=True)
+    parser.add_argument("--text", required=True)
+    parser.add_argument("--tasks", required=True, help='JSON object, e.g. {"intent": ["a", "b"]}')
+    parser.add_argument("--precision", default="fp16")
+    parser.add_argument("--length", type=int, default=128)
+    args = parser.parse_args()
+    model = CoreMLDecide(args.model_dir, args.precision, args.length)
+    print(json.dumps(model.classify(args.text, json.loads(args.tasks)), indent=2))
+
+
+if __name__ == "__main__":
+    main()

scripts/make_golden.py

@@ -0,0 +1,43 @@
+"""Record native (PyTorch) GLiNER2.5-Decide answers for tests/cases.py into tests/golden.json.
+
+Downloads the pinned source checkpoint (~1.7 GB) on first run. The tests compare the Core ML
+router against this file, so they need neither PyTorch weights nor a network connection.
+
+    uv run python scripts/make_golden.py
+"""
+import json
+import sys
+import warnings
+from pathlib import Path
+
+warnings.filterwarnings("ignore")
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT / "tests"))
+sys.path.insert(0, str(ROOT / "scripts" / "build"))
+
+from gliner2 import AutoExtractor  # noqa: E402
+from huggingface_hub import snapshot_download  # noqa: E402
+
+from cases import CASES  # noqa: E402
+from convert_names import MODEL_ID, MODEL_REVISION  # noqa: E402
+
+
+def main():
+    source = snapshot_download(
+        MODEL_ID, revision=MODEL_REVISION,
+        allow_patterns=["config.json", "encoder_config/*", "model.safetensors", "tokenizer.json",
+                        "tokenizer_config.json", "special_tokens_map.json"],
+    )
+    native = AutoExtractor.from_pretrained(source, map_location="cpu").eval()
+    golden = {}
+    for case in CASES:
+        golden[case["id"]] = native.classify_text(case["text"], case["tasks"], include_confidence=True)
+        print(case["id"], json.dumps(golden[case["id"]]))
+    out = ROOT / "tests" / "golden.json"
+    out.write_text(json.dumps({"source_model": MODEL_ID, "source_revision": MODEL_REVISION, "cases": golden},
+                              indent=2) + "\n")
+    print("wrote", out)
+
+
+if __name__ == "__main__":
+    main()

src/gliner_decide_coreml/__init__.py

@@ -0,0 +1,4 @@
+"""GLiNER2.5-Decide on Core ML: one multifunction package, automatic 64/128/256/512 bucket routing."""
+from .router import BUCKETS, MAX_HEADS, MAX_OPTIONS, DecideRouter, Route
+
+__all__ = ["BUCKETS", "MAX_HEADS", "MAX_OPTIONS", "DecideRouter", "Route"]

src/gliner_decide_coreml/decoding.py

@@ -0,0 +1,63 @@
+"""Per-head logits -> labels, with the native ``classify_text`` activation and threshold rules.
+
+Single-label heads use softmax and return one label; multi-label heads use sigmoid and return
+every label at or above ``cls_threshold`` (or the single best one if none pass), as in
+Fluid Inference's gliner2-5-decide-coreml `runtime.decode` (Apache-2.0).
+"""
+import numpy as np
+
+from .encoding import task_labels
+
+
+def head_config(tasks: dict, name: str) -> dict:
+    return tasks[name] if isinstance(tasks[name], dict) else {}
+
+
+def probabilities(config: dict, logits: np.ndarray) -> np.ndarray:
+    values = logits.astype(np.float64)
+    activation = config.get("class_act", "auto")
+    if activation == "sigmoid" or (activation == "auto" and config.get("multi_label", False)):
+        return 1.0 / (1.0 + np.exp(-values))
+    probs = np.exp(values - values.max())
+    return probs / probs.sum()
+
+
+def format_head(config: dict, labels: list[str], probs: np.ndarray):
+    if config.get("multi_label", False):
+        threshold = config.get("cls_threshold", 0.5)
+        chosen = [{"label": labels[j], "confidence": float(probs[j])}
+                  for j in range(len(labels)) if probs[j] >= threshold]
+        best = int(probs.argmax())
+        return chosen or [{"label": labels[best], "confidence": float(probs[best])}]
+    best = int(probs.argmax())
+    return {"label": labels[best], "confidence": float(probs[best])}
+
+
+def head_probabilities(tasks: dict, logits: np.ndarray) -> dict[str, np.ndarray]:
+    """logits: (heads, options) from the model -> probabilities per head name."""
+    return {
+        name: probabilities(head_config(tasks, name), logits[h, : len(labels)])
+        for h, (name, labels) in enumerate(task_labels(tasks).items())
+    }
+
+
+def decode(tasks: dict, probs: dict[str, np.ndarray]) -> dict:
+    labels = task_labels(tasks)
+    return {name: format_head(head_config(tasks, name), labels[name], p) for name, p in probs.items()}
+
+
+def merge_chunks(tasks: dict, per_chunk: list[dict[str, np.ndarray]]) -> dict[str, np.ndarray]:
+    """Combine per-chunk probabilities for text longer than the largest bucket.
+
+    Heuristic, since no single chunk saw the whole text:
+      single-label -> the distribution from the chunk that is most confident for that head
+      multi-label  -> per-label maximum across chunks (a label is on if any chunk says so)
+    """
+    merged = {}
+    for name in task_labels(tasks):
+        rows = [chunk[name] for chunk in per_chunk]
+        if head_config(tasks, name).get("multi_label", False):
+            merged[name] = np.max(np.stack(rows), axis=0)
+        else:
+            merged[name] = max(rows, key=lambda p: float(p.max()))
+    return merged

src/gliner_decide_coreml/encoding.py

@@ -0,0 +1,95 @@
+"""Text + task schema -> model inputs, using GLiNER2's own processor.
+
+This is the exact preprocessing the Core ML buckets were converted and verified with.
+Adapted from Fluid Inference's gliner2-5-decide-coreml `preprocessing.py` (Apache-2.0).
+"""
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+from gliner2 import Schema
+from gliner2.models.base import load_extractor_tokenizer
+from gliner2.processor import SchemaTransformer
+from gliner2.training.trainer import ExtractorCollator
+
+
+def load_collator(tokenizer_dir: Path) -> ExtractorCollator:
+    """Tokenizer + schema formatter + collator; no PyTorch weights are loaded."""
+    processor = SchemaTransformer(tokenizer=load_extractor_tokenizer(str(tokenizer_dir)), token_pooling="first")
+    return ExtractorCollator(processor, is_training=False, max_len=None, architecture="span")
+
+
+def task_labels(tasks: dict) -> dict[str, list[str]]:
+    """Label names per head, in declaration order."""
+    result = {}
+    for name, config in tasks.items():
+        labels = config["labels"] if isinstance(config, dict) and "labels" in config else config
+        result[name] = list(labels.keys()) if isinstance(labels, dict) else list(labels)
+    return result
+
+
+def classification_schema(tasks: dict) -> Schema:
+    """Same task-dict handling as native ``classify_text``."""
+    schema = Schema()
+    for name, config in tasks.items():
+        if isinstance(config, dict) and "labels" in config:
+            cfg = config.copy()
+            labels = cfg.pop("labels")
+            schema.classification(name, labels, **cfg)
+        else:
+            schema.classification(name, config)
+    return schema
+
+
+@dataclass(frozen=True)
+class Encoded:
+    """One request as the encoder sees it, before padding."""
+
+    input_ids: np.ndarray       # (n,) int32: schema tokens, [SEP_TEXT], text tokens
+    attention_mask: np.ndarray  # (n,) int32
+    markers: list[list[int]]    # per head: token position of each [L] label marker
+
+    @property
+    def length(self) -> int:
+        return len(self.input_ids)
+
+
+def encode(collator: ExtractorCollator, text: str, tasks: dict) -> Encoded:
+    batch = collator([(text, classification_schema(tasks).build())])
+    labels = task_labels(tasks)
+    groups = batch.schema_special_indices[0]
+    if len(groups) != len(labels):
+        raise ValueError("Schema head count does not match the requested tasks")
+    markers = []
+    for positions, values in zip(groups, labels.values()):
+        # positions[0] is the [P] prompt marker; the rest are one [L] marker per label.
+        found = [int(p) for p in positions[1:]]
+        if len(found) != len(values):
+            raise ValueError("Label markers were truncated or merged")
+        markers.append(found)
+    return Encoded(
+        input_ids=batch.input_ids[0].numpy().astype(np.int32),
+        attention_mask=batch.attention_mask[0].numpy().astype(np.int32),
+        markers=markers,
+    )
+
+
+def to_bucket(enc: Encoded, length: int, max_heads: int, max_options: int, pad_id: int) -> dict[str, np.ndarray]:
+    """Pad one request into a fixed (length, heads, options) bucket."""
+    n = enc.length
+    if n > length:
+        raise ValueError(f"Request needs {n} tokens; bucket holds {length}")
+    if len(enc.markers) > max_heads:
+        raise ValueError(f"{len(enc.markers)} heads; bucket holds {max_heads}")
+    ids = np.full((1, length), pad_id, dtype=np.int32)
+    ids[0, :n] = enc.input_ids
+    attention = np.zeros((1, length), dtype=np.int32)
+    attention[0, :n] = enc.attention_mask
+    indices = np.zeros((1, max_heads, max_options), dtype=np.int32)
+    mask = np.zeros((1, max_heads, max_options), dtype=np.float32)
+    for head, markers in enumerate(enc.markers):
+        if len(markers) > max_options:
+            raise ValueError(f"Head has {len(markers)} labels; bucket holds {max_options}")
+        indices[0, head, : len(markers)] = markers
+        mask[0, head, : len(markers)] = 1.0
+    return {"input_ids": ids, "attention_mask": attention, "marker_indices": indices, "marker_mask": mask}

src/gliner_decide_coreml/router.py

@@ -0,0 +1,133 @@
+"""Route each request to the smallest Core ML bucket that fits it.
+
+One multifunction package holds four fixed-shape graphs (functions ``L64``, ``L128``, ``L256``,
+``L512``) that share a single copy of the weights. A request is tokenized once, sent to the
+smallest function whose length fits, and decoded on the host. Text that does not fit even
+``L512`` is split into overlapping word chunks whose results are merged.
+"""
+from dataclasses import dataclass
+from pathlib import Path
+
+import coremltools as ct
+
+from .decoding import decode, head_probabilities, merge_chunks
+from .encoding import Encoded, encode, load_collator, task_labels, to_bucket
+
+BUCKETS = (64, 128, 256, 512)
+MAX_HEADS = 4       # heads scored per model call; more are split across calls
+MAX_OPTIONS = 32    # labels per head
+PACKAGE = "GLiNER2.5-Decide-MultiFn-fp16"
+DEFAULT_MODEL_DIR = Path(__file__).resolve().parents[2] / "models"
+
+
+@dataclass(frozen=True)
+class Route:
+    """How a request was served."""
+
+    tokens: int                # schema + text tokens (largest head group)
+    bucket: int | None         # function used; None when the text had to be chunked
+    chunks: int                # 1 unless the text was longer than the largest bucket
+    calls: int                 # model predictions made
+
+
+def _ensure_compiled(package: Path) -> Path:
+    compiled = package.with_suffix(".mlmodelc")
+    if not compiled.exists() or compiled.stat().st_mtime < package.stat().st_mtime:
+        ct.models.utils.compile_model(str(package), destination_path=str(compiled))
+    return compiled
+
+
+class DecideRouter:
+    """GLiNER2.5-Decide classification on Core ML with automatic bucket selection.
+
+    Loads every bucket at construction, since the first load of a function compiles it for
+    the device, so do this once at startup, not per request.
+    ``compute_units`` defaults to the GPU: on an M1 Pro it beats ``ALL`` at every bucket size.
+    """
+
+    def __init__(self, model_dir: Path | str = DEFAULT_MODEL_DIR,
+                 compute_units: ct.ComputeUnit = ct.ComputeUnit.CPU_AND_GPU,
+                 chunk_overlap_words: int = 32):
+        model_dir = Path(model_dir)
+        self.collator = load_collator(model_dir / "tokenizer")
+        self.pad_id = self.collator.processor.tokenizer.pad_token_id
+        compiled = _ensure_compiled(model_dir / f"{PACKAGE}.mlpackage")
+        self.models = {
+            length: ct.models.CompiledMLModel(str(compiled), compute_units=compute_units, function_name=f"L{length}")
+            for length in BUCKETS
+        }
+        self.chunk_overlap_words = chunk_overlap_words
+
+    def count_tokens(self, text: str, tasks: dict) -> int:
+        return encode(self.collator, text, tasks).length
+
+    def classify(self, text: str, tasks: dict) -> dict:
+        return self.classify_with_route(text, tasks)[0]
+
+    def classify_with_route(self, text: str, tasks: dict) -> tuple[dict, Route]:
+        if not tasks:
+            raise ValueError("At least one classification head is required")
+        for name, labels in task_labels(tasks).items():
+            if not 1 <= len(labels) <= MAX_OPTIONS:
+                raise ValueError(f"Head {name!r} has {len(labels)} labels; supported: 1..{MAX_OPTIONS}")
+        # Heads beyond MAX_HEADS go to separate calls. Heads in one call see each other's labels,
+        # so splitting can shift results slightly compared with a single native pass.
+        names = list(tasks)
+        groups = [{n: tasks[n] for n in names[i:i + MAX_HEADS]} for i in range(0, len(names), MAX_HEADS)]
+        result, routes = {}, []
+        for group in groups:
+            group_result, route = self._classify_group(text, group)
+            result.update(group_result)
+            routes.append(route)
+        largest = max(routes, key=lambda r: r.tokens)
+        return {name: result[name] for name in names}, Route(
+            tokens=largest.tokens,
+            bucket=None if any(r.bucket is None for r in routes) else max(r.bucket for r in routes),
+            chunks=max(r.chunks for r in routes),
+            calls=sum(r.calls for r in routes),
+        )
+
+    def _classify_group(self, text: str, tasks: dict) -> tuple[dict, Route]:
+        enc = encode(self.collator, text, tasks)
+        bucket = self._bucket_for(enc)
+        if bucket is not None:
+            return decode(tasks, self._probabilities(enc, tasks, bucket)), Route(enc.length, bucket, 1, 1)
+        chunks = self._split(text, tasks)
+        per_chunk = []
+        for chunk in chunks:
+            chunk_enc = encode(self.collator, chunk, tasks)
+            per_chunk.append(self._probabilities(chunk_enc, tasks, self._bucket_for(chunk_enc)))
+        return decode(tasks, merge_chunks(tasks, per_chunk)), Route(enc.length, None, len(chunks), len(chunks))
+
+    @staticmethod
+    def _bucket_for(enc: Encoded) -> int | None:
+        return next((length for length in BUCKETS if enc.length <= length), None)
+
+    def _probabilities(self, enc: Encoded, tasks: dict, bucket: int):
+        arrays = to_bucket(enc, bucket, MAX_HEADS, MAX_OPTIONS, self.pad_id)
+        logits = self.models[bucket].predict(arrays)["logits"][0]
+        return head_probabilities(tasks, logits)
+
+    def _split(self, text: str, tasks: dict) -> list[str]:
+        """Greedy word chunks that each fit the largest bucket, overlapping by a few words."""
+        words = text.split()
+        limit = BUCKETS[-1]
+
+        def fits(start: int, end: int) -> bool:
+            return encode(self.collator, " ".join(words[start:end]), tasks).length <= limit
+
+        if not fits(0, 1):
+            raise ValueError(f"The task schema alone does not fit in {limit} tokens")
+        chunks, start = [], 0
+        while True:
+            lo, hi = start + 1, len(words)  # invariant: words[start:lo] fits
+            while lo < hi:
+                mid = (lo + hi + 1) // 2
+                if fits(start, mid):
+                    lo = mid
+                else:
+                    hi = mid - 1
+            chunks.append(" ".join(words[start:lo]))
+            if lo == len(words):
+                return chunks
+            start = max(lo - self.chunk_overlap_words, start + 1)

tests/cases.py

@@ -0,0 +1,104 @@
+"""Test requests, sized to land in each bucket. Shared by scripts/make_golden.py and the tests."""
+
+SUPPORT_TASKS = {
+    "intent": ["order_status", "refund_request", "cancel_subscription", "update_payment", "other"],
+    "urgency": ["low", "normal", "high"],
+}
+
+CASES = [
+    {
+        "id": "short-refund",
+        "bucket": 64,
+        "text": "My subscription renewed on April 15 for 5,400 yen after the service was already down. "
+                "Can I get that charge refunded?",
+        "tasks": SUPPORT_TASKS,
+    },
+    {
+        "id": "laptop-review-aspects",
+        "bucket": 64,
+        "text": "Battery dies before lunch, but the keyboard and the screen are the best I have used on a laptop.",
+        "tasks": {
+            "sentiment": ["positive", "negative", "mixed", "neutral"],
+            "aspects": {"labels": ["battery", "keyboard", "screen", "camera", "price", "support"],
+                        "multi_label": True, "cls_threshold": 0.4},
+        },
+    },
+    {
+        "id": "compliance-email-3-heads",
+        "bucket": 128,
+        "text": "From: compliance@group.example\nSubject: Protocol update, action required today\n\n"
+                "Please confirm the new retention rule is applied before Friday's audit.",
+        "tasks": {
+            "intent": ["fyi", "request", "approval", "complaint", "newsletter", "security_alert"],
+            "urgency": ["low", "normal", "high", "critical"],
+            "route": ["support", "billing", "legal", "security", "finance", "archive"],
+        },
+    },
+    {
+        "id": "descriptions-and-prompt",
+        "bucket": 128,
+        "text": "My subscription renewed on April 15 for 5,400 yen after the service was already down. "
+                "Can I get that charge refunded?",
+        "tasks": {
+            "intent": {"labels": {
+                "order_status": "The customer asks where an order is",
+                "refund_request": "The customer wants money back",
+                "cancel_subscription": "The customer wants to stop a recurring plan",
+                "update_payment": "The customer wants to change a card or billing method",
+                "other": "None of the above",
+            }},
+            "urgency": {"labels": ["low", "normal", "high"], "prompt": "How urgently should support respond?"},
+        },
+    },
+    {
+        "id": "hotel-complaint-4-heads",
+        "bucket": 256,
+        "text": "Hello, I'm writing from room 1408. The air conditioning stopped working yesterday afternoon "
+                "and the room is now above 30 degrees, so neither of us could sleep last night. We called the "
+                "front desk twice and were told a technician would come, but nobody showed up. We have a "
+                "conference presentation tomorrow morning and cannot spend another night like this. Please "
+                "either fix it within the next two hours or move us to another room tonight. If neither is "
+                "possible we will check out early and expect the remaining nights to be refunded. We would "
+                "also like the incidentals hold on our card to be released, since we were told at check-in "
+                "that it would only last one day. Thank you.",
+        "tasks": {
+            "intent": ["maintenance", "room_change", "checkout", "billing", "complaint", "amenity_request"],
+            "priority": ["low", "normal", "high", "urgent"],
+            "needs_human": ["yes", "no"],
+            "topics": {"labels": ["hvac", "billing", "housekeeping", "noise", "safety"],
+                       "multi_label": True, "cls_threshold": 0.4},
+        },
+    },
+    {
+        "id": "incident-report-long",
+        "bucket": 512,
+        "text": "Incident summary for the payments platform, prepared for the weekly operations review. "
+                "On Tuesday at 09:12 UTC the card authorization service began returning timeouts for roughly "
+                "one in five requests. The on-call engineer was paged four minutes later by the latency alert "
+                "and started investigating the database tier, which showed normal load. At 09:31 the team noticed "
+                "that a configuration change deployed the previous evening had lowered the connection pool size "
+                "for the authorization service from two hundred to twenty connections. Under the morning traffic "
+                "peak the pool was exhausted, requests queued, and the upstream gateway gave up after its "
+                "three second deadline. The change had been reviewed, but the reviewer assumed the value was "
+                "per worker rather than per service, and the staging environment does not receive enough "
+                "traffic to reveal the problem. The configuration was rolled back at 09:44 and error rates "
+                "returned to normal within two minutes. In total about eleven thousand authorizations failed; "
+                "most customers retried successfully, but around four hundred orders were abandoned. No card "
+                "data was exposed and there is no indication of any security issue. Follow-up actions: add a "
+                "load test that replays production traffic volumes against staging before configuration "
+                "changes to the authorization service are approved; make pool sizes explicit about their "
+                "scope in the configuration schema; add an alert on connection pool saturation so the "
+                "on-call engineer is pointed at the right component immediately; and review whether the "
+                "gateway deadline should be longer for authorization calls. Finance has asked for an estimate "
+                "of lost revenue from abandoned orders by the end of the week, and customer support has "
+                "prepared a short statement in case merchants ask about the failed transactions. The service "
+                "owner will present these actions at the architecture forum next Thursday and track them to "
+                "completion in the reliability backlog.",
+        "tasks": {
+            "doc_type": ["incident_report", "feature_request", "meeting_notes", "security_advisory", "marketing"],
+            "severity": ["sev1", "sev2", "sev3", "sev4"],
+            "teams": {"labels": ["payments", "security", "finance", "support", "marketing", "legal"],
+                      "multi_label": True, "cls_threshold": 0.4},
+        },
+    },
+]

tests/golden.json

@@ -0,0 +1,108 @@
+{
+  "source_model": "fastino/GLiNER2.5-Decide",
+  "source_revision": "65624f1a0265b3f612bae66a2685a06b94a68a9d",
+  "cases": {
+    "short-refund": {
+      "intent": {
+        "label": "refund_request",
+        "confidence": 0.9987769722938538
+      },
+      "urgency": {
+        "label": "low",
+        "confidence": 0.3725019693374634
+      }
+    },
+    "laptop-review-aspects": {
+      "sentiment": {
+        "label": "positive",
+        "confidence": 0.9994937181472778
+      },
+      "aspects": [
+        {
+          "label": "battery",
+          "confidence": 0.7610353231430054
+        },
+        {
+          "label": "keyboard",
+          "confidence": 0.9939987659454346
+        },
+        {
+          "label": "screen",
+          "confidence": 0.9846210479736328
+        }
+      ]
+    },
+    "compliance-email-3-heads": {
+      "intent": {
+        "label": "request",
+        "confidence": 0.44433778524398804
+      },
+      "urgency": {
+        "label": "critical",
+        "confidence": 0.411828875541687
+      },
+      "route": {
+        "label": "legal",
+        "confidence": 0.5800722241401672
+      }
+    },
+    "descriptions-and-prompt": {
+      "intent": {
+        "label": "refund_request",
+        "confidence": 0.914971113204956
+      },
+      "urgency": {
+        "label": "normal",
+        "confidence": 0.38015663623809814
+      }
+    },
+    "hotel-complaint-4-heads": {
+      "intent": {
+        "label": "checkout",
+        "confidence": 0.9280232191085815
+      },
+      "priority": {
+        "label": "urgent",
+        "confidence": 0.3841314911842346
+      },
+      "needs_human": {
+        "label": "yes",
+        "confidence": 0.6280083060264587
+      },
+      "topics": [
+        {
+          "label": "hvac",
+          "confidence": 0.8771321773529053
+        },
+        {
+          "label": "billing",
+          "confidence": 0.6899283528327942
+        }
+      ]
+    },
+    "incident-report-long": {
+      "doc_type": {
+        "label": "incident_report",
+        "confidence": 0.5693653225898743
+      },
+      "severity": {
+        "label": "sev3",
+        "confidence": 0.33189529180526733
+      },
+      "teams": [
+        {
+          "label": "payments",
+          "confidence": 0.9709568619728088
+        },
+        {
+          "label": "finance",
+          "confidence": 0.8886100053787231
+        },
+        {
+          "label": "support",
+          "confidence": 0.8042415976524353
+        }
+      ]
+    }
+  }
+}

tests/test_router.py

@@ -0,0 +1,57 @@
+"""The Core ML router must reproduce the native PyTorch model (tests/golden.json)."""
+import json
+from pathlib import Path
+
+import pytest
+
+from cases import CASES, SUPPORT_TASKS
+from gliner_decide_coreml import BUCKETS, DecideRouter
+
+GOLDEN = json.loads((Path(__file__).parent / "golden.json").read_text())["cases"]
+CONFIDENCE_TOLERANCE = 0.02  # fp16 Core ML vs fp32 PyTorch
+
+
+@pytest.fixture(scope="module")
+def router():
+    return DecideRouter()
+
+
+def entries(value):
+    return value if isinstance(value, list) else [value]
+
+
+@pytest.mark.parametrize("case", CASES, ids=[c["id"] for c in CASES])
+def test_matches_native(router, case):
+    result, route = router.classify_with_route(case["text"], case["tasks"])
+    assert route.bucket == case["bucket"] and route.calls == 1
+    native = GOLDEN[case["id"]]
+    assert list(result) == list(native)
+    for head, expected in native.items():
+        got = {e["label"]: e["confidence"] for e in entries(result[head])}
+        want = {e["label"]: e["confidence"] for e in entries(expected)}
+        assert set(got) == set(want), f"{head}: {sorted(got)} != {sorted(want)}"
+        for label, confidence in want.items():
+            assert got[label] == pytest.approx(confidence, abs=CONFIDENCE_TOLERANCE), f"{head}/{label}"
+
+
+def test_every_bucket_is_exercised():
+    assert {c["bucket"] for c in CASES} == set(BUCKETS)
+
+
+def test_long_text_is_chunked(router):
+    text = " ".join([CASES[0]["text"]] * 40)  # ~1,000 tokens: larger than the 512 bucket
+    result, route = router.classify_with_route(text, SUPPORT_TASKS)
+    assert route.bucket is None and route.chunks >= 2 and route.calls == route.chunks
+    assert result["intent"]["label"] == "refund_request"
+
+
+def test_more_than_four_heads_are_split_across_calls(router):
+    tasks = {**CASES[4]["tasks"], "language": ["english", "spanish", "german"], "sentiment": ["positive", "negative"]}
+    result, route = router.classify_with_route(CASES[4]["text"], tasks)
+    assert list(result) == list(tasks) and route.calls == 2
+    assert result["language"]["label"] == "english"
+
+
+def test_too_many_labels_is_rejected(router):
+    with pytest.raises(ValueError, match="labels"):
+        router.classify("hello", {"topic": [f"label_{i}" for i in range(33)]})