Add GLiNER2.5-Decide support and Sort Anything / Sort Decisions demos

Package.swift

@@ -45,6 +45,11 @@ let package = Package(
             dependencies: ["FluidUse", "Game2048"],
             exclude: ["README.md"]
         ),
+        .target(name: "SortAnything", dependencies: ["FluidUse"]),
+        .executableTarget(name: "SortAnythingCheck", dependencies: ["SortAnything"]),
+        .executableTarget(name: "SortDecisionsCheck", dependencies: ["SortAnything"]),
+        .executableTarget(name: "SortDecisionsDemo", dependencies: ["SortAnything"], exclude: ["README.md"]),
+        .executableTarget(name: "SortAnythingDemo", dependencies: ["SortAnything"], exclude: ["README.md"]),
         .testTarget(
             name: "FluidUseTests", dependencies: ["FluidUse", "LayaTetris"],
             resources: [.copy("Fixtures")]

Sources/FluidUse/GLiNER2/GLiNER2Manager.swift

@@ -1,27 +1,36 @@
 @preconcurrency import CoreML
 import Foundation
 
-/// On-device, dynamic-label classification with GLiNER 2.5 base or multilingual.
+/// On-device, dynamic-label classification with GLiNER 2.5 small, base, multilingual, or Decide.
 public actor GLiNER2Manager {
-    public static let maximumLength = 128
-    public static let maximumOptions = 8
+    public nonisolated let maximumLength: Int
+    public nonisolated let maximumOptions: Int
 
-    private let model: MLModel
-    private let tokenizer: GLiNER2Tokenizer
+    private nonisolated let markerShape: [Int]
+    private nonisolated let model: MLModel
+    private nonisolated let tokenizer: GLiNER2Tokenizer
     private let inputIds: MLMultiArray
     private let attentionMask: MLMultiArray
     private let markerIndices: MLMultiArray
     private let markerMask: MLMultiArray
     private let features: MLDictionaryFeatureProvider
 
-    public init(model: MLModel, tokenizer: GLiNER2Tokenizer) throws {
-        try Self.validate(model.modelDescription)
+    public init(model: MLModel, tokenizer: GLiNER2Tokenizer, variant: GLiNER2Variant = .base) throws {
+        maximumLength = variant.maximumLength
+        maximumOptions = variant.maximumOptions
+        markerShape = variant.markerShape
+        try Self.validate(model.modelDescription, length: maximumLength, markerShape: markerShape)
         self.model = model
         self.tokenizer = tokenizer
-        inputIds = try MLMultiArray(shape: [1, 128], dataType: .int32)
-        attentionMask = try MLMultiArray(shape: [1, 128], dataType: .int32)
-        markerIndices = try MLMultiArray(shape: [1, 8], dataType: .int32)
-        markerMask = try MLMultiArray(shape: [1, 8], dataType: .float32)
+        let length = NSNumber(value: maximumLength)
+        let markers = markerShape.map { NSNumber(value: $0) }
+        inputIds = try MLMultiArray(shape: [1, length], dataType: .int32)
+        attentionMask = try MLMultiArray(shape: [1, length], dataType: .int32)
+        // Zero-filled, so every head after the first stays masked.
+        markerIndices = try MLMultiArray(shape: markers, dataType: .int32)
+        markerMask = try MLMultiArray(shape: markers, dataType: .float32)
+        markerIndices.dataPointer.initializeMemory(as: Int32.self, repeating: 0, count: markerIndices.count)
+        markerMask.dataPointer.initializeMemory(as: Float.self, repeating: 0, count: markerMask.count)
         features = try MLDictionaryFeatureProvider(dictionary: [
             "input_ids": MLFeatureValue(multiArray: inputIds),
             "attention_mask": MLFeatureValue(multiArray: attentionMask),
@@ -30,7 +39,7 @@ public actor GLiNER2Manager {
         ])
     }
 
-    /// Download and load the published W8 classification package.
+    /// Download and load the variant's published classification package.
     public static func load(
         variant: GLiNER2Variant, cacheDirectory: URL? = nil, computeUnits: MLComputeUnits = .all,
         progress: GLiNER2ModelStore.Progress? = nil
@@ -44,9 +53,9 @@ public actor GLiNER2Manager {
     public static func load(
         from directory: URL, variant: GLiNER2Variant, computeUnits: MLComputeUnits = .all
     ) async throws -> GLiNER2Manager {
-        let tokenizerURL = directory.appendingPathComponent("tokenizer/tokenizer.json")
+        let tokenizerURL = directory.appendingPathComponent(variant.tokenizerPath)
         guard FileManager.default.fileExists(atPath: tokenizerURL.path) else {
-            throw GLiNER2Error.invalidAsset("Missing tokenizer/tokenizer.json")
+            throw GLiNER2Error.invalidAsset("Missing \(variant.tokenizerPath)")
         }
         let tokenizer = try GLiNER2Tokenizer(tokenizerJsonURL: tokenizerURL)
         let package = directory.appendingPathComponent(variant.packageName)
@@ -63,22 +72,23 @@ public actor GLiNER2Manager {
         let configuration = MLModelConfiguration()
         configuration.computeUnits = computeUnits
         let model = try await MLModel.load(contentsOf: modelURL, configuration: configuration)
-        return try GLiNER2Manager(model: model, tokenizer: tokenizer)
+        return try GLiNER2Manager(model: model, tokenizer: tokenizer, variant: variant)
     }
 
-    /// Classify `text` against 1–8 labels using the checkpoint's native schema format.
+    /// Classify `text` against 1…`maximumOptions` labels using the checkpoint's native schema format.
     public func classify(text: String, task: String, labels: [String]) throws -> GLiNER2Answer {
         try Task.checkCancellation()
-        guard (1...Self.maximumOptions).contains(labels.count) else {
-            throw GLiNER2Error.invalidInput("Expected 1–8 labels; received \(labels.count)")
+        guard (1...maximumOptions).contains(labels.count) else {
+            throw GLiNER2Error.invalidInput("Expected 1–\(maximumOptions) labels; received \(labels.count)")
         }
         guard !task.isEmpty else { throw GLiNER2Error.invalidInput("Task must be nonempty") }
         guard !labels.contains(where: \.isEmpty) else {
             throw GLiNER2Error.invalidInput("Labels must be nonempty")
         }
         let sequence = try tokenizer.classificationSequence(text: text, task: task, labels: labels)
-        guard sequence.ids.count <= Self.maximumLength else {
-            throw GLiNER2Error.invalidInput("Schema and text require \(sequence.ids.count) tokens; maximum is 128")
+        guard sequence.ids.count <= maximumLength else {
+            throw GLiNER2Error.invalidInput(
+                "Schema and text require \(sequence.ids.count) tokens; maximum is \(maximumLength)")
         }
         guard sequence.markers.count == labels.count else {
             throw GLiNER2Error.invalidInput("A label marker was lost during tokenization")
@@ -88,6 +98,79 @@ public actor GLiNER2Manager {
             labels: labels, probabilities: values.probabilities, logits: values.logits, tokenCount: sequence.ids.count)
     }
 
+    /// Same result as `classify`, but with per-call inputs and Core ML's async prediction, so several calls can be
+    /// in flight on one model at once (the async API is thread-safe; the synchronous one is not).
+    public nonisolated func classifyConcurrently(
+        text: String, task: String, labels: [String]
+    ) async throws -> GLiNER2Answer {
+        try await classifyConcurrently(text: text, heads: [(task, labels)])[0]
+    }
+
+    /// Answers up to `maximumHeads` classification heads over `text` in one prediction, in the order given.
+    public nonisolated func classifyConcurrently(
+        text: String, heads: [(task: String, labels: [String])]
+    ) async throws -> [GLiNER2Answer] {
+        let headCount = markerShape.count == 3 ? markerShape[1] : 1
+        guard (1...headCount).contains(heads.count) else {
+            throw GLiNER2Error.invalidInput("Expected 1–\(headCount) heads; received \(heads.count)")
+        }
+        for head in heads {
+            guard (1...maximumOptions).contains(head.labels.count), !head.task.isEmpty,
+                !head.labels.contains(where: \.isEmpty)
+            else {
+                throw GLiNER2Error.invalidInput("Each head needs a task and 1–\(maximumOptions) nonempty labels")
+            }
+        }
+        let sequence = try tokenizer.classificationSequence(text: text, heads: heads)
+        guard sequence.ids.count <= maximumLength else {
+            throw GLiNER2Error.invalidInput(
+                "Schema and text require \(sequence.ids.count) tokens; maximum is \(maximumLength)")
+        }
+        guard zip(sequence.markers, heads).allSatisfy({ $0.count == $1.labels.count }) else {
+            throw GLiNER2Error.invalidInput("A label marker was lost during tokenization")
+        }
+        let length = NSNumber(value: maximumLength)
+        let shape = markerShape.map { NSNumber(value: $0) }
+        let ids = try MLMultiArray(shape: [1, length], dataType: .int32)
+        let attention = try MLMultiArray(shape: [1, length], dataType: .int32)
+        let markers = try MLMultiArray(shape: shape, dataType: .int32)
+        let mask = try MLMultiArray(shape: shape, dataType: .float32)
+        let idPointer = ids.dataPointer.assumingMemoryBound(to: Int32.self)
+        let attentionPointer = attention.dataPointer.assumingMemoryBound(to: Int32.self)
+        for index in 0..<maximumLength {
+            idPointer[index] = Int32(index < sequence.ids.count ? sequence.ids[index] : tokenizer.padTokenId)
+            attentionPointer[index] = index < sequence.ids.count ? 1 : 0
+        }
+        markers.dataPointer.initializeMemory(as: Int32.self, repeating: 0, count: markers.count)
+        mask.dataPointer.initializeMemory(as: Float.self, repeating: 0, count: mask.count)
+        let markerPointer = markers.dataPointer.assumingMemoryBound(to: Int32.self)
+        let maskPointer = mask.dataPointer.assumingMemoryBound(to: Float.self)
+        for (head, positions) in sequence.markers.enumerated() {
+            for (slot, position) in positions.enumerated() {
+                markerPointer[head * maximumOptions + slot] = Int32(position)
+                maskPointer[head * maximumOptions + slot] = 1
+            }
+        }
+        let features = try MLDictionaryFeatureProvider(dictionary: [
+            "input_ids": MLFeatureValue(multiArray: ids), "attention_mask": MLFeatureValue(multiArray: attention),
+            "marker_indices": MLFeatureValue(multiArray: markers), "marker_mask": MLFeatureValue(multiArray: mask),
+        ])
+        let output = try await model.prediction(from: features)
+        let allLogits = try read("logits", from: output, count: headCount * maximumOptions)
+        let allProbabilities = try read("probabilities", from: output, count: headCount * maximumOptions)
+        return try heads.enumerated().map { head, request in
+            let range = (head * maximumOptions)..<(head * maximumOptions + request.labels.count)
+            let logits = Array(allLogits[range])
+            let probabilities = Array(allProbabilities[range])
+            guard logits.allSatisfy(\.isFinite), probabilities.allSatisfy(\.isFinite) else {
+                throw GLiNER2Error.invalidOutput("Model returned non-finite scores")
+            }
+            return GLiNER2Answer(
+                labels: request.labels, probabilities: probabilities, logits: logits,
+                tokenCount: sequence.ids.count)
+        }
+    }
+
     /// Exposes the native schema sequence for reference parity checks.
     public nonisolated func tokenSequence(
         text: String, task: String, labels: [String]
@@ -98,13 +181,13 @@ public actor GLiNER2Manager {
     private func predict(ids: [Int], markers: [Int]) throws -> (logits: [Float], probabilities: [Float]) {
         let idPointer = inputIds.dataPointer.assumingMemoryBound(to: Int32.self)
         let attentionPointer = attentionMask.dataPointer.assumingMemoryBound(to: Int32.self)
-        for index in 0..<Self.maximumLength {
+        for index in 0..<maximumLength {
             idPointer[index] = Int32(index < ids.count ? ids[index] : tokenizer.padTokenId)
             attentionPointer[index] = index < ids.count ? 1 : 0
         }
         let markerPointer = markerIndices.dataPointer.assumingMemoryBound(to: Int32.self)
         let maskPointer = markerMask.dataPointer.assumingMemoryBound(to: Float.self)
-        for index in 0..<Self.maximumOptions {
+        for index in 0..<maximumOptions {
             markerPointer[index] = Int32(index < markers.count ? markers[index] : 0)
             maskPointer[index] = index < markers.count ? 1 : 0
         }
@@ -117,20 +200,22 @@ public actor GLiNER2Manager {
         return (Array(logits.prefix(markers.count)), Array(probabilities.prefix(markers.count)))
     }
 
-    private func read(_ name: String, from output: MLFeatureProvider) throws -> [Float] {
+    private nonisolated func read(
+        _ name: String, from output: MLFeatureProvider, count: Int? = nil
+    ) throws -> [Float] {
         guard let array = output.featureValue(for: name)?.multiArrayValue,
-            array.shape.map(\.intValue) == [1, Self.maximumOptions], array.dataType == .float32
-        else { throw GLiNER2Error.invalidOutput("\(name) must be float32 [1, 8]") }
+            array.shape.map(\.intValue) == markerShape, array.dataType == .float32
+        else { throw GLiNER2Error.invalidOutput("\(name) must be float32 \(markerShape)") }
         let pointer = array.dataPointer.assumingMemoryBound(to: Float.self)
-        return (0..<Self.maximumOptions).map { pointer[$0] }
+        return (0..<(count ?? maximumOptions)).map { pointer[$0] }
     }
 
-    private static func validate(_ description: MLModelDescription) throws {
+    private static func validate(_ description: MLModelDescription, length: Int, markerShape: [Int]) throws {
         let expected: [String: ([Int], MLMultiArrayDataType)] = [
-            "input_ids": ([1, 128], .int32),
-            "attention_mask": ([1, 128], .int32),
-            "marker_indices": ([1, 8], .int32),
-            "marker_mask": ([1, 8], .float32),
+            "input_ids": ([1, length], .int32),
+            "attention_mask": ([1, length], .int32),
+            "marker_indices": (markerShape, .int32),
+            "marker_mask": (markerShape, .float32),
         ]
         for (name, requirement) in expected {
             guard let constraint = description.inputDescriptionsByName[name]?.multiArrayConstraint,

Sources/FluidUse/GLiNER2/GLiNER2ModelStore.swift

@@ -1,7 +1,7 @@
 import CryptoKit
 import Foundation
 
-/// Downloads the pinned, eight-bit GLiNER 2.5 classification packages from Hugging Face.
+/// Downloads the pinned GLiNER 2.5 classification packages from Hugging Face.
 public enum GLiNER2ModelStore {
     public typealias Progress = @Sendable (_ file: String, _ bytes: Int64) -> Void
 
@@ -15,6 +15,7 @@ public enum GLiNER2ModelStore {
         case .small: "9dcac8a315ca49e71412cf5dec2ee7b7609b614e"
         case .base: "c1843f2c193b11b05f09ac7f258cb9202d8f5e71"
         case .multilingual: "5dab512eb89b88a3680bd6c86841877c3ea49893"
+        case .decide, .decideLong: "cd0d7b1ef32b10e1e3a5a73c9d9ac8411d819c5a"
         }
     }
 
@@ -88,10 +89,52 @@ public enum GLiNER2ModelStore {
                     path: "\(package)/Data/com.apple.CoreML/weights/weight.bin",
                     sha256: "655460ce8e9420b55130f3b6b976d4797aa197cd14c909aee96967cea24d766e"),
             ]
+        case .decide:
+            return [
+                Asset(path: "config.json", sha256: "e748e5b80575471c91b3f0dd00f513ba58242544fcb7e1236e0021e61abd7673"),
+                Asset(
+                    path: "encoder_config/config.json",
+                    sha256: "bd32f1484ba5a199f7a63df44df3814b839fffcf6e64478323c4689868ef6015"),
+                Asset(
+                    path: "tokenizer.json", sha256: "3ad87d9ffe669147063e70850927dd2da90249e2acc5c8527f1eb65df467bcc8"),
+                Asset(
+                    path: "tokenizer_config.json",
+                    sha256: "323199a4e946039410899f3779f2aa3eaef1500213c512727ad0f623d4f21309"),
+                Asset(
+                    path: "\(package)/Manifest.json",
+                    sha256: "6d62c3f4c7331d836cebf29c541815e0c6d7da9e4612ee45ca726847acbf6ed8"),
+                Asset(
+                    path: "\(package)/Data/com.apple.CoreML/model.mlmodel",
+                    sha256: "b9d4d8a497ea986b5d3163259694e8fc2571c9e7e2cdca2bf1b568bab8111c2a"),
+                Asset(
+                    path: "\(package)/Data/com.apple.CoreML/weights/weight.bin",
+                    sha256: "54501158f56baf0ebe295d99ac251eb0204cea2594881062b731bef9154032e2"),
+            ]
+        case .decideLong:
+            return [
+                Asset(path: "config.json", sha256: "e748e5b80575471c91b3f0dd00f513ba58242544fcb7e1236e0021e61abd7673"),
+                Asset(
+                    path: "encoder_config/config.json",
+                    sha256: "bd32f1484ba5a199f7a63df44df3814b839fffcf6e64478323c4689868ef6015"),
+                Asset(
+                    path: "tokenizer.json", sha256: "3ad87d9ffe669147063e70850927dd2da90249e2acc5c8527f1eb65df467bcc8"),
+                Asset(
+                    path: "tokenizer_config.json",
+                    sha256: "323199a4e946039410899f3779f2aa3eaef1500213c512727ad0f623d4f21309"),
+                Asset(
+                    path: "\(package)/Manifest.json",
+                    sha256: "542249fd40bfd27f7da11b303dda14932e8027be4279aab5e65d1c9abcc730c2"),
+                Asset(
+                    path: "\(package)/Data/com.apple.CoreML/model.mlmodel",
+                    sha256: "626d74aac1448e0666e5a80d79302139dfd84d5ca80c44ec213c1c6a5497d823"),
+                Asset(
+                    path: "\(package)/Data/com.apple.CoreML/weights/weight.bin",
+                    sha256: "54501158f56baf0ebe295d99ac251eb0204cea2594881062b731bef9154032e2"),
+            ]
         }
     }
 
-    /// Ensure one variant's tokenizer and W8 Core ML package exist in the FluidUse cache.
+    /// Ensure one variant's tokenizer and Core ML package exist in the FluidUse cache.
     public static func ensure(
         variant: GLiNER2Variant, cacheDirectory: URL? = nil, progress: Progress? = nil
     ) async throws -> URL {

Sources/FluidUse/GLiNER2/GLiNER2Tokenizer.swift

@@ -152,25 +152,36 @@ public struct GLiNER2Tokenizer: Sendable {
     public func classificationSequence(
         text: String, task: String, labels: [String]
     ) throws -> (ids: [Int], markers: [Int]) {
+        let sequence = try classificationSequence(text: text, heads: [(task, labels)])
+        return (sequence.ids, sequence.markers[0])
+    }
+
+    /// Several classification heads in one schema, joined by `[SEP_STRUCT]` as GLiNER's SchemaTransformer does.
+    /// `markers[h]` holds the token position of each `[L]` marker of head `h`.
+    public func classificationSequence(
+        text: String, heads: [(task: String, labels: [String])]
+    ) throws -> (ids: [Int], markers: [[Int]]) {
         var source = text
         if source.isEmpty || !source.hasSuffix(".") && !source.hasSuffix("!") && !source.hasSuffix("?") {
             source += "."
         }
-        var items = ["(", "[P]", task, "("]
-        for label in labels {
-            items.append("[L]")
-            items.append(label)
-        }
-        items += [")", ")", "[SEP_TEXT]"]
-        items += try Self.splitText(source)
-
         var ids: [Int] = []
-        var markers: [Int] = []
-        for (index, item) in items.enumerated() {
-            if index >= 4 && index < 4 + labels.count * 2 && index.isMultiple(of: 2) {
-                markers.append(ids.count)
+        var markers: [[Int]] = []
+        for (index, head) in heads.enumerated() {
+            if index > 0 { ids += encode("[SEP_STRUCT]") }
+            ids += encode("(") + encode("[P]") + encode(head.task) + encode("(")
+            var positions: [Int] = []
+            for label in head.labels {
+                positions.append(ids.count)
+                ids.append(labelTokenId)
+                ids += encode(label)
             }
-            ids.append(contentsOf: encode(item))
+            ids += encode(")") + encode(")")
+            markers.append(positions)
+        }
+        ids.append(separatorTokenId)
+        for word in try Self.splitText(source) {
+            ids += encode(word)
         }
         return (ids, markers)
     }