diff --git a/.github/workflows/swift-test.yml b/.github/workflows/swift-test.yml index c29b2d0fd..17e8c8132 100644 --- a/.github/workflows/swift-test.yml +++ b/.github/workflows/swift-test.yml @@ -58,3 +58,87 @@ jobs: -scheme OpenMedKit \ -destination 'generic/platform=iOS Simulator' \ -skipPackagePluginValidation + + build-watchos: + name: Build & Test watchOS Simulator + runs-on: macos-latest + + steps: + - uses: actions/checkout@v7 + + - name: Select Xcode version + run: sudo xcode-select -s /Applications/Xcode.app/Contents/Developer + + - name: Select available watchOS simulator + id: destination + shell: bash + run: | + set -euo pipefail + udid="$( + xcrun simctl list devices available --json | + jq -r '[ + .devices | to_entries[] | + select(.key | contains("watchOS")) | + .value[] | + select(.isAvailable == true) | + .udid + ][0] // empty' + )" + test -n "$udid" + echo "udid=$udid" >> "$GITHUB_OUTPUT" + + - name: Build and run watchOS parity tests + working-directory: swift/OpenMedKit + run: | + xcodebuild test \ + -scheme OpenMedKit \ + -destination 'platform=watchOS Simulator,id=${{ steps.destination.outputs.udid }}' \ + -only-testing:OpenMedKitTests/WatchVisionParityTests \ + -skipPackagePluginValidation + + build-visionos: + name: Build & Test visionOS Simulator + runs-on: macos-latest + + steps: + - uses: actions/checkout@v7 + + - name: Select Xcode version + run: sudo xcode-select -s /Applications/Xcode.app/Contents/Developer + + - name: Build native visionOS simulator target + working-directory: swift/OpenMedKit + shell: bash + run: | + set -euo pipefail + sdk_path="$(xcrun --sdk xrsimulator --show-sdk-path)" + swift build \ + --triple arm64-apple-xros1.0-simulator \ + --sdk "$sdk_path" + + - name: Select available visionOS simulator + id: destination + shell: bash + run: | + set -euo pipefail + udid="$( + xcrun simctl list devices available --json | + jq -r '[ + .devices | to_entries[] | + select(.key | contains("xrOS")) | + .value[] | + select(.isAvailable == true) | + .udid + ][0] // empty' + )" + test -n "$udid" + echo "udid=$udid" >> "$GITHUB_OUTPUT" + + - name: Build and run visionOS parity tests + working-directory: swift/OpenMedKit + run: | + xcodebuild test \ + -scheme OpenMedKit \ + -destination 'platform=visionOS Simulator,id=${{ steps.destination.outputs.udid }}' \ + -only-testing:OpenMedKitTests/WatchVisionParityTests \ + -skipPackagePluginValidation diff --git a/Package.swift b/Package.swift index 9ad9463b5..bf7c74a88 100644 --- a/Package.swift +++ b/Package.swift @@ -7,6 +7,8 @@ let package = Package( platforms: [ .iOS(.v17), .macOS(.v14), + .watchOS(.v10), + .visionOS(.v1), ], products: [ .library( @@ -32,10 +34,26 @@ let package = Package( .target( name: "OpenMedKit", dependencies: [ - .product(name: "Transformers", package: "swift-transformers"), - .product(name: "MLX", package: "mlx-swift"), - .product(name: "MLXNN", package: "mlx-swift"), - .product(name: "ZIPFoundation", package: "ZIPFoundation"), + .product( + name: "Transformers", + package: "swift-transformers", + condition: .when(platforms: [.iOS, .macOS]) + ), + .product( + name: "MLX", + package: "mlx-swift", + condition: .when(platforms: [.iOS, .macOS]) + ), + .product( + name: "MLXNN", + package: "mlx-swift", + condition: .when(platforms: [.iOS, .macOS]) + ), + .product( + name: "ZIPFoundation", + package: "ZIPFoundation", + condition: .when(platforms: [.iOS, .macOS]) + ), ], path: "swift/OpenMedKit/Sources/OpenMedKit", resources: [ diff --git a/docs/runtimes/apple-platforms.md b/docs/runtimes/apple-platforms.md new file mode 100644 index 000000000..05fe363d6 --- /dev/null +++ b/docs/runtimes/apple-platforms.md @@ -0,0 +1,99 @@ +# OpenMedKit Apple Platform Support + +OpenMedKit supports macOS 14+, iOS 17+, watchOS 10+, and visionOS 1+. +The runtime surface depends on the platform so constrained devices do not link +backends that exceed their deployment or memory envelope. + +| Platform | Supported backend | Default model ceiling | Resident RAM ceiling | Maximum sequence | +|---|---|---:|---:|---:| +| macOS 14+ | MLX or CoreML | Base | 900 MB | 512 tokens | +| iOS 17+ / iPadOS 17+ | MLX on a physical device, or CoreML | Tiny | 350 MB | 512 tokens | +| watchOS 10+ | CoreML only | Nano, INT8 | 150 MB | 256 tokens | +| visionOS 1+ | CoreML only | Nano, INT8 | 150 MB | 256 tokens | + +The watchOS and visionOS limits use OpenMed's canonical Nano sub-tier: 10–30M +parameters, at most 150 MB resident memory, and INT8 CoreML artifacts. These +limits are enforced from model metadata before `MLModel` is opened. A Tiny, +Base, over-budget, or non-INT8 model fails closed instead of being loaded. + +## Selecting and loading a constrained CoreML model + +Describe the bundled model candidates and let `PlatformModel` select the +highest-capacity candidate that fits the current target: + +```swift +import OpenMedKit + +let nano = PlatformModelDescriptor( + identifier: "OpenMed-PII-Nano-INT8", + modelURL: Bundle.main.url( + forResource: "OpenMed-PII-Nano-INT8", + withExtension: "mlmodelc" + )!, + id2labelURL: Bundle.main.url( + forResource: "id2label", + withExtension: "json" + )!, + tier: .nano, + parameterCount: 24_000_000, + estimatedResidentMemoryMB: 128, + isINT8: true +) + +let model = try PlatformModel(candidates: [nano]) +``` + +watchOS and visionOS intentionally omit the full MLX and +`swift-transformers` graph. Apps tokenize with the assets bundled beside their +Nano model and pass bounded token IDs, attention masks, and character offsets +to `PlatformModel.predict(...)`. This keeps model and tokenizer access local; +OpenMedKit does not add a cloud fallback. + +## Minimal redaction surface + +`PlatformModel.redact(...)` applies mask or removal redaction to detected +`EntityPrediction` spans without loading another backend: + +```swift +let note = "Patient Ada Lovelace, MRN TEST-123." +let spans = [ + EntityPrediction( + label: "full_name", + text: "Ada Lovelace", + confidence: 0.99, + start: 8, + end: 20 + ), + EntityPrediction( + label: "medical_record_number", + text: "TEST-123", + confidence: 0.99, + start: 26, + end: 34 + ), +] + +let result = PlatformModel.redact(note, entities: spans) +// Patient [FULL_NAME], MRN [MEDICAL_RECORD_NUMBER]. +``` + +Offsets remain character offsets into the original note. The watchOS and +visionOS simulator tests use the same synthetic note and iOS reference spans, +with a one-character tolerance at each boundary. + +## Build and validation + +The Swift workflow performs the normal macOS tests, the iOS simulator build, +and focused parity tests on available watchOS and visionOS simulators. Local +checks use the same package scheme: + +```bash +cd swift/OpenMedKit +xcodebuild build -scheme OpenMedKit \ + -destination 'generic/platform=watchOS Simulator' +xcodebuild build -scheme OpenMedKit \ + -destination 'generic/platform=visionOS Simulator' +``` + +Use only synthetic notes in committed tests and fixtures. OpenMedKit keeps +inference on device and does not log input text or detected span text. diff --git a/docs/swift-openmedkit.md b/docs/swift-openmedkit.md index 0353f684e..ac19cd673 100644 --- a/docs/swift-openmedkit.md +++ b/docs/swift-openmedkit.md @@ -1,11 +1,13 @@ # OpenMedKit (Swift Package) -OpenMedKit is the Swift package for running OpenMed models in **macOS**, **iOS**, and **iPadOS** apps. +OpenMedKit is the Swift package for running OpenMed models in **macOS**, +**iOS**, **iPadOS**, **watchOS**, and **visionOS** apps. OpenMedKit currently supports two Apple backends: - **MLX** for Apple Silicon Macs and real iPhone/iPad devices -- **CoreML** for bundled Apple model packages +- **CoreML** for bundled Apple model packages, including constrained watchOS + and visionOS Nano artifacts Swift MLX supports the first OpenMed artifact families used by the public Apple demos: @@ -22,7 +24,7 @@ ModernBERT, Longformer, EuroBERT, Qwen3, and additional architecture families ar ## Requirements -- iOS 17+ / macOS 14+ +- iOS 17+ / macOS 14+ / watchOS 10+ / visionOS 1+ - Xcode 15+ - For MLX: - Apple Silicon Mac, or @@ -31,6 +33,9 @@ ModernBERT, Longformer, EuroBERT, Qwen3, and additional architecture families ar - a compatible `.mlpackage` or `.mlmodelc` bundle plus `id2label.json` iOS Simulator is **not** a Swift MLX validation target. +watchOS and visionOS use the CoreML-only `PlatformModel` surface and require an +INT8 Nano-tier artifact. See [Apple Platform Support](./runtimes/apple-platforms.md) +for selection limits and simulator validation. ## Apple Platform Matrix @@ -40,6 +45,8 @@ iOS Simulator is **not** a Swift MLX validation target. | Swift app on Apple Silicon macOS | `OpenMedKit` + MLX or CoreML | | Swift app on real iPhone/iPad | `OpenMedKit` + MLX or CoreML | | Swift app on iOS Simulator | CoreML only | +| Swift app on Apple Watch | `PlatformModel` + Nano INT8 CoreML | +| Swift app on Apple Vision Pro | `PlatformModel` + Nano INT8 CoreML | | Older Apple OS support | CoreML | ## Install OpenMedKit diff --git a/mkdocs.yml b/mkdocs.yml index db755e1f7..995253102 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -124,6 +124,7 @@ nav: - On-device Segmenters: on-device-segmenters.md - CoreML Packaging: coreml-export.md - Swift Package (OpenMedKit): swift-openmedkit.md + - OpenMedKit Apple Platforms: runtimes/apple-platforms.md - Swift-Kotlin API Parity: swift-kotlin-parity.md - Project: - FAQ: faq.md diff --git a/swift/OpenMedKit/Package.swift b/swift/OpenMedKit/Package.swift index 41bbd2847..2371b9ca3 100644 --- a/swift/OpenMedKit/Package.swift +++ b/swift/OpenMedKit/Package.swift @@ -8,6 +8,8 @@ let package = Package( platforms: [ .iOS(.v17), .macOS(.v14), + .watchOS(.v10), + .visionOS(.v1), ], products: [ .library( @@ -25,10 +27,26 @@ let package = Package( .target( name: "OpenMedKit", dependencies: [ - .product(name: "Transformers", package: "swift-transformers"), - .product(name: "MLX", package: "mlx-swift"), - .product(name: "MLXNN", package: "mlx-swift"), - .product(name: "ZIPFoundation", package: "ZIPFoundation"), + .product( + name: "Transformers", + package: "swift-transformers", + condition: .when(platforms: [.iOS, .macOS]) + ), + .product( + name: "MLX", + package: "mlx-swift", + condition: .when(platforms: [.iOS, .macOS]) + ), + .product( + name: "MLXNN", + package: "mlx-swift", + condition: .when(platforms: [.iOS, .macOS]) + ), + .product( + name: "ZIPFoundation", + package: "ZIPFoundation", + condition: .when(platforms: [.iOS, .macOS]) + ), ], resources: [ .process("Resources") diff --git a/swift/OpenMedKit/Sources/OpenMedKit/NERPipeline.swift b/swift/OpenMedKit/Sources/OpenMedKit/NERPipeline.swift index 01b01b49e..e82ce4365 100644 --- a/swift/OpenMedKit/Sources/OpenMedKit/NERPipeline.swift +++ b/swift/OpenMedKit/Sources/OpenMedKit/NERPipeline.swift @@ -17,8 +17,40 @@ public class NERPipeline { /// - modelURL: Path to the `.mlmodelc` or `.mlpackage` file. /// - id2labelURL: Path to the `id2label.json` file mapping label IDs to names. /// - maxSeqLength: Maximum input sequence length the model supports. - public init(modelURL: URL, id2labelURL: URL, maxSeqLength: Int = 512) throws { - self.model = try MLModel(contentsOf: try Self.resolveModelURL(modelURL)) + @available(watchOS, unavailable, message: "Use PlatformModel for Nano budget enforcement.") + @available(visionOS, unavailable, message: "Use PlatformModel for Nano budget enforcement.") + public convenience init( + modelURL: URL, + id2labelURL: URL, + maxSeqLength: Int = 512 + ) throws { + try self.init( + resolvedModelURL: Self.resolveModelURL(modelURL), + id2labelURL: id2labelURL, + maxSeqLength: maxSeqLength + ) + } + + convenience init( + validatedDescriptor descriptor: PlatformModelDescriptor, + configuration: PlatformModelConfiguration + ) throws { + guard configuration.allows(descriptor) else { + throw PlatformModelError.noCompatibleModel(configuration.platform) + } + try self.init( + resolvedModelURL: Self.resolveModelURL(descriptor.modelURL), + id2labelURL: descriptor.id2labelURL, + maxSeqLength: configuration.maximumSequenceLength + ) + } + + private init( + resolvedModelURL: URL, + id2labelURL: URL, + maxSeqLength: Int + ) throws { + self.model = try MLModel(contentsOf: resolvedModelURL) self.maxSeqLength = maxSeqLength let data = try Data(contentsOf: id2labelURL) @@ -32,7 +64,11 @@ public class NERPipeline { private static func resolveModelURL(_ modelURL: URL) throws -> URL { switch modelURL.pathExtension.lowercased() { case "mlpackage", "mlmodel": - return try MLModel.compileModel(at: modelURL) + #if os(watchOS) || os(visionOS) + throw NERPipelineError.uncompiledModelUnsupported(modelURL) + #else + return try MLModel.compileModel(at: modelURL) + #endif default: return modelURL } @@ -139,11 +175,14 @@ public class NERPipeline { /// Errors thrown by the NER pipeline. public enum NERPipelineError: Error, LocalizedError { case missingOutput(String) + case uncompiledModelUnsupported(URL) public var errorDescription: String? { switch self { case .missingOutput(let name): return "CoreML model output '\(name)' not found" + case .uncompiledModelUnsupported(let url): + return "\(url.lastPathComponent) must be compiled to .mlmodelc before bundling on watchOS or visionOS" } } } diff --git a/swift/OpenMedKit/Sources/OpenMedKit/OpenMedGLiNERModels.swift b/swift/OpenMedKit/Sources/OpenMedKit/OpenMedGLiNERModels.swift index ea292078a..3284a459a 100644 --- a/swift/OpenMedKit/Sources/OpenMedKit/OpenMedGLiNERModels.swift +++ b/swift/OpenMedKit/Sources/OpenMedKit/OpenMedGLiNERModels.swift @@ -1,769 +1,771 @@ -import Foundation -import MLX -import MLXNN - -private func openMedFloatScalar(_ value: Float, like array: MLXArray) -> MLXArray { - MLXArray(value).asType(array.dtype) -} - -struct OpenMedGLiNERSpanOutput { - let logits: MLXArray - let spanIndex: MLXArray - let spanMask: MLXArray - let promptMask: MLXArray - let wordMask: MLXArray -} - -struct OpenMedGLiClassOutput { - let logits: MLXArray - let classesMask: MLXArray -} - -struct OpenMedGLiNERRelexEntityOutput { - let entityScores: MLXArray - let entityPromptMask: MLXArray - let wordMask: MLXArray - let encoded: OpenMedGLiNERRelexEncoded -} - -struct OpenMedGLiNERRelexRelationOutput { - let pairScores: MLXArray - let pairIndex: MLXArray - let pairMask: MLXArray - let relationPromptMask: MLXArray -} - -struct OpenMedGLiNERRelexEncoded { - let entityPrompts: MLXArray - let entityPromptMask: MLXArray - let relationPrompts: MLXArray - let relationPromptMask: MLXArray - let wordsEmbedding: MLXArray - let wordMask: MLXArray -} - -final class OpenMedProjectionMLP: Module { - @ModuleInfo(key: "linear1") var linear1: Linear - @ModuleInfo(key: "linear2") var linear2: Linear - - init(inputSize: Int, outputSize: Int? = nil) { - let outputSize = outputSize ?? inputSize - _linear1.wrappedValue = Linear(inputSize, outputSize * 4) - _linear2.wrappedValue = Linear(outputSize * 4, outputSize) - } +#if canImport(MLX) && !os(watchOS) && !os(visionOS) + import Foundation + import MLX + import MLXNN - func callAsFunction(_ input: MLXArray) -> MLXArray { - linear2(relu(linear1(input))) + private func openMedFloatScalar(_ value: Float, like array: MLXArray) -> MLXArray { + MLXArray(value).asType(array.dtype) } -} -final class OpenMedBidirectionalLSTM: Module { - @ModuleInfo(key: "forward_lstm") var forwardLSTM: LSTM - @ModuleInfo(key: "backward_lstm") var backwardLSTM: LSTM + struct OpenMedGLiNERSpanOutput { + let logits: MLXArray + let spanIndex: MLXArray + let spanMask: MLXArray + let promptMask: MLXArray + let wordMask: MLXArray + } - init(hiddenSize: Int) { - _forwardLSTM.wrappedValue = LSTM(inputSize: hiddenSize, hiddenSize: hiddenSize / 2) - _backwardLSTM.wrappedValue = LSTM(inputSize: hiddenSize, hiddenSize: hiddenSize / 2) + struct OpenMedGLiClassOutput { + let logits: MLXArray + let classesMask: MLXArray } - private static func reversePadded(_ input: MLXArray, lengths: MLXArray) -> MLXArray { - let steps = MLXArray.arange(input.dim(1), dtype: .int32).expandedDimensions(axis: 0) - let expandedLengths = lengths.expandedDimensions(axis: 1) - let gatherIndex = `where`( - steps .< expandedLengths, - expandedLengths - 1 - steps, - steps - ) - return takeAlong(input, gatherIndex.expandedDimensions(axis: 2), axis: 1) + struct OpenMedGLiNERRelexEntityOutput { + let entityScores: MLXArray + let entityPromptMask: MLXArray + let wordMask: MLXArray + let encoded: OpenMedGLiNERRelexEncoded } - func callAsFunction(_ input: MLXArray, mask: MLXArray) -> MLXArray { - let lengths = sum(mask.asType(.int32), axis: 1) - let (forwardOutput, _) = forwardLSTM(input) - let reversedInput = Self.reversePadded(input, lengths: lengths) - let (backwardReversed, _) = backwardLSTM(reversedInput) - let backwardOutput = Self.reversePadded(backwardReversed, lengths: lengths) - return concatenated([forwardOutput, backwardOutput], axis: -1) - * mask.asType(input.dtype).expandedDimensions(axis: -1) + struct OpenMedGLiNERRelexRelationOutput { + let pairScores: MLXArray + let pairIndex: MLXArray + let pairMask: MLXArray + let relationPromptMask: MLXArray } -} -private func openMedPadEmbeddings(_ rows: [MLXArray], width: Int, embedDim: Int) -> MLXArray { - guard !rows.isEmpty else { - return MLXArray.zeros([0, width, embedDim], type: Float.self) + struct OpenMedGLiNERRelexEncoded { + let entityPrompts: MLXArray + let entityPromptMask: MLXArray + let relationPrompts: MLXArray + let relationPromptMask: MLXArray + let wordsEmbedding: MLXArray + let wordMask: MLXArray } - let paddedRows = rows.map { row -> MLXArray in - let padLength = width - row.dim(0) - guard padLength > 0 else { - return row + final class OpenMedProjectionMLP: Module { + @ModuleInfo(key: "linear1") var linear1: Linear + @ModuleInfo(key: "linear2") var linear2: Linear + + init(inputSize: Int, outputSize: Int? = nil) { + let outputSize = outputSize ?? inputSize + _linear1.wrappedValue = Linear(inputSize, outputSize * 4) + _linear2.wrappedValue = Linear(outputSize * 4, outputSize) } - let padding = MLXArray.zeros([padLength, row.dim(-1)], type: Float.self).asType(row.dtype) - return concatenated([row, padding], axis: 0) - } - return stacked(paddedRows, axis: 0) -} -private func openMedPadMaskRows(_ rows: [[Int]], width: Int) -> MLXArray { - let padded = rows.map { row in - row + Array(repeating: 0, count: max(0, width - row.count)) - } - return MLXArray(padded.flatMap { $0 }, [rows.count, width]).asType(.bool) -} - -private func openMedExtractMarkerEmbeddings( - tokenEmbeddings: MLXArray, - inputIDs: MLXArray, - markerTokenID: Int, - includeMarkerToken: Bool -) -> (MLXArray, MLXArray) { - let batchSize = tokenEmbeddings.dim(0) - let seqLen = tokenEmbeddings.dim(1) - let embedDim = tokenEmbeddings.dim(2) - var rows = [MLXArray]() - var maskRows = [[Int]]() - var maxItems = 0 - - for batchIndex in 0.. MLXArray { + linear2(relu(linear1(input))) + } } - maxItems = max(maxItems, 1) - return ( - openMedPadEmbeddings(rows, width: maxItems, embedDim: embedDim), - openMedPadMaskRows(maskRows, width: maxItems) - ) -} - -private func openMedExtractWordEmbeddings( - tokenEmbeddings: MLXArray, - wordsMask: MLXArray -) -> (MLXArray, MLXArray) { - let batchSize = tokenEmbeddings.dim(0) - let embedDim = tokenEmbeddings.dim(2) - var rows = [MLXArray]() - var maskRows = [[Int]]() - var maxWords = 0 - - for batchIndex in 0.. 0 ? index : nil - } - - guard !tokenPositions.isEmpty else { - rows.append(MLXArray.zeros([0, embedDim], type: Float.self).asType(tokenEmbeddings.dtype)) - maskRows.append([]) - continue - } - - let positionArray = MLXArray(tokenPositions.map(Int32.init), [tokenPositions.count]) - let row = tokenEmbeddings[batchIndex].take(positionArray, axis: 0) - rows.append(row) - maxWords = max(maxWords, tokenPositions.count) - maskRows.append(Array(repeating: 1, count: tokenPositions.count)) - } + final class OpenMedBidirectionalLSTM: Module { + @ModuleInfo(key: "forward_lstm") var forwardLSTM: LSTM + @ModuleInfo(key: "backward_lstm") var backwardLSTM: LSTM - maxWords = max(maxWords, 1) - return ( - openMedPadEmbeddings(rows, width: maxWords, embedDim: embedDim), - openMedPadMaskRows(maskRows, width: maxWords) - ) -} - -private func openMedGatherSpanEndpoints( - startHiddenStates: MLXArray, - endHiddenStates: MLXArray, - spanIndex: MLXArray -) -> (MLXArray, MLXArray) { - let hiddenSize = startHiddenStates.dim(-1) - let startIndex = broadcast( - spanIndex[0..., 0..., 0].expandedDimensions(axis: -1), - to: [spanIndex.dim(0), spanIndex.dim(1), hiddenSize] - ) - let endIndex = broadcast( - spanIndex[0..., 0..., 1].expandedDimensions(axis: -1), - to: [spanIndex.dim(0), spanIndex.dim(1), hiddenSize] - ) - return ( - takeAlong(startHiddenStates, startIndex, axis: 1), - takeAlong(endHiddenStates, endIndex, axis: 1) - ) -} - -private func openMedBuildAllEntityPairs( - spanRep: MLXArray, - spanMask: MLXArray -) -> (MLXArray, MLXArray, MLXArray, MLXArray) { - let batchSize = spanRep.dim(0) - let embedDim = spanRep.dim(2) - var pairRows = [[[Int]]]() - var pairMaskRows = [[Int]]() - var headRows = [MLXArray]() - var tailRows = [MLXArray]() - var maxPairs = 0 - - for batchIndex in 0.. MLXArray { + let steps = MLXArray.arange(input.dim(1), dtype: .int32).expandedDimensions(axis: 0) + let expandedLengths = lengths.expandedDimensions(axis: 1) + let gatherIndex = `where`( + steps .< expandedLengths, + expandedLengths - 1 - steps, + steps + ) + return takeAlong(input, gatherIndex.expandedDimensions(axis: 2), axis: 1) } - let headIndex = MLXArray(pairs.map { Int32($0[0]) }, [pairs.count]) - let tailIndex = MLXArray(pairs.map { Int32($0[1]) }, [pairs.count]) - headRows.append(spanRep[batchIndex].take(headIndex, axis: 0)) - tailRows.append(spanRep[batchIndex].take(tailIndex, axis: 0)) + func callAsFunction(_ input: MLXArray, mask: MLXArray) -> MLXArray { + let lengths = sum(mask.asType(.int32), axis: 1) + let (forwardOutput, _) = forwardLSTM(input) + let reversedInput = Self.reversePadded(input, lengths: lengths) + let (backwardReversed, _) = backwardLSTM(reversedInput) + let backwardOutput = Self.reversePadded(backwardReversed, lengths: lengths) + return concatenated([forwardOutput, backwardOutput], axis: -1) + * mask.asType(input.dtype).expandedDimensions(axis: -1) + } } - maxPairs = max(maxPairs, 1) - let paddedPairs = pairRows.map { row in - row + Array(repeating: [0, 0], count: max(0, maxPairs - row.count)) - } - let flatPairs = paddedPairs.flatMap { $0 }.flatMap { $0 }.map(Int32.init) - - return ( - MLXArray(flatPairs, [batchSize, maxPairs, 2]), - openMedPadMaskRows(pairMaskRows, width: maxPairs), - openMedPadEmbeddings(headRows, width: maxPairs, embedDim: embedDim), - openMedPadEmbeddings(tailRows, width: maxPairs, embedDim: embedDim) - ) -} - -final class OpenMedSpanMarkerV0: Module { - @ModuleInfo(key: "project_start") var projectStart: OpenMedProjectionMLP - @ModuleInfo(key: "project_end") var projectEnd: OpenMedProjectionMLP - @ModuleInfo(key: "out_project") var outputProjection: OpenMedProjectionMLP - - init(hiddenSize: Int) { - _projectStart.wrappedValue = OpenMedProjectionMLP(inputSize: hiddenSize) - _projectEnd.wrappedValue = OpenMedProjectionMLP(inputSize: hiddenSize) - _outputProjection.wrappedValue = OpenMedProjectionMLP( - inputSize: hiddenSize * 2, - outputSize: hiddenSize - ) - } + private func openMedPadEmbeddings(_ rows: [MLXArray], width: Int, embedDim: Int) -> MLXArray { + guard !rows.isEmpty else { + return MLXArray.zeros([0, width, embedDim], type: Float.self) + } - func callAsFunction(_ hiddenStates: MLXArray, spanIndex: MLXArray) -> MLXArray { - let startRep = projectStart(hiddenStates) - let endRep = projectEnd(hiddenStates) - let (startSpanRep, endSpanRep) = openMedGatherSpanEndpoints( - startHiddenStates: startRep, - endHiddenStates: endRep, - spanIndex: spanIndex - ) - return outputProjection(relu(concatenated([startSpanRep, endSpanRep], axis: -1))) - } -} - -final class OpenMedGLiNERSpanModel: Module { - private let configuration: OpenMedMLXBertConfiguration - - @ModuleInfo(key: "deberta") var deberta: OpenMedDebertaV2Model - @ModuleInfo(key: "token_projection") var tokenProjection: Linear - @ModuleInfo(key: "rnn") var rnn: OpenMedBidirectionalLSTM? - @ModuleInfo(key: "span_rep_layer") var spanRepLayer: OpenMedSpanMarkerV0 - @ModuleInfo(key: "prompt_rep_layer") var promptRepLayer: OpenMedProjectionMLP - - init(_ configuration: OpenMedMLXBertConfiguration) { - self.configuration = configuration - _deberta.wrappedValue = OpenMedDebertaV2Model(configuration) - _tokenProjection.wrappedValue = Linear( - configuration.encoderHiddenSize, - configuration.hiddenSize - ) - if configuration.numRNNLayers > 0 { - _rnn.wrappedValue = OpenMedBidirectionalLSTM(hiddenSize: configuration.hiddenSize) + let paddedRows = rows.map { row -> MLXArray in + let padLength = width - row.dim(0) + guard padLength > 0 else { + return row + } + let padding = MLXArray.zeros([padLength, row.dim(-1)], type: Float.self).asType(row.dtype) + return concatenated([row, padding], axis: 0) } - _spanRepLayer.wrappedValue = OpenMedSpanMarkerV0(hiddenSize: configuration.hiddenSize) - _promptRepLayer.wrappedValue = OpenMedProjectionMLP(inputSize: configuration.hiddenSize) + return stacked(paddedRows, axis: 0) } - private func encode( - inputIDs: MLXArray, - attentionMask: MLXArray, - wordsMask: MLXArray - ) -> (MLXArray, MLXArray, MLXArray, MLXArray) { - let hiddenStates = deberta(inputIDs: inputIDs, attentionMask: attentionMask) - let markerTokenID = configuration.classTokenIndex ?? 0 - let (promptEmbeddings, promptMask) = openMedExtractMarkerEmbeddings( - tokenEmbeddings: hiddenStates, - inputIDs: inputIDs, - markerTokenID: markerTokenID, - includeMarkerToken: configuration.embedEntityToken - ) - let (wordEmbeddings, wordMask) = openMedExtractWordEmbeddings( - tokenEmbeddings: hiddenStates, - wordsMask: wordsMask - ) - - var projectedWords = tokenProjection(wordEmbeddings) - let projectedPrompts = tokenProjection(promptEmbeddings) - if let rnn { - projectedWords = rnn(projectedWords, mask: wordMask) + private func openMedPadMaskRows(_ rows: [[Int]], width: Int) -> MLXArray { + let padded = rows.map { row in + row + Array(repeating: 0, count: max(0, width - row.count)) } - return (projectedPrompts, promptMask, projectedWords, wordMask) + return MLXArray(padded.flatMap { $0 }, [rows.count, width]).asType(.bool) } - func callAsFunction( + private func openMedExtractMarkerEmbeddings( + tokenEmbeddings: MLXArray, inputIDs: MLXArray, - attentionMask: MLXArray, - wordsMask: MLXArray, - spanIndex: MLXArray, - spanMask: MLXArray - ) -> OpenMedGLiNERSpanOutput { - let maskedSpanIndex = spanIndex * spanMask.asType(spanIndex.dtype).expandedDimensions(axis: -1) - let (promptEmbeddings, promptMask, wordEmbeddings, wordMask) = encode( - inputIDs: inputIDs, - attentionMask: attentionMask, - wordsMask: wordsMask - ) - let spanRep = spanRepLayer(wordEmbeddings, spanIndex: maskedSpanIndex) - let promptRep = promptRepLayer(promptEmbeddings) - var logits = einsum("bsd,bcd->bsc", spanRep, promptRep) - logits = `where`( - spanMask.expandedDimensions(axis: -1), - logits, - openMedFloatScalar(-1.0e9, like: logits) - ) - return OpenMedGLiNERSpanOutput( - logits: logits, - spanIndex: maskedSpanIndex, - spanMask: spanMask, - promptMask: promptMask, - wordMask: wordMask - ) - } + markerTokenID: Int, + includeMarkerToken: Bool + ) -> (MLXArray, MLXArray) { + let batchSize = tokenEmbeddings.dim(0) + let seqLen = tokenEmbeddings.dim(1) + let embedDim = tokenEmbeddings.dim(2) + var rows = [MLXArray]() + var maskRows = [[Int]]() + var maxItems = 0 - func sanitize(weights: [String: MLXArray]) -> [String: MLXArray] { - weights.filter { key, _ in - key != "deberta.embeddings.position_ids" && !key.hasPrefix("_") - } - } -} + for batchIndex in 0.. MLXArray { - linear2(gelu(linear1(features))) + maxItems = max(maxItems, 1) + return ( + openMedPadEmbeddings(rows, width: maxItems, embedDim: embedDim), + openMedPadMaskRows(maskRows, width: maxItems) + ) } -} -final class OpenMedGLiClassMLPScorer: Module { - @ModuleInfo(key: "linear1") var linear1: Linear - @ModuleInfo(key: "linear2") var linear2: Linear - @ModuleInfo(key: "linear3") var linear3: Linear + private func openMedExtractWordEmbeddings( + tokenEmbeddings: MLXArray, + wordsMask: MLXArray + ) -> (MLXArray, MLXArray) { + let batchSize = tokenEmbeddings.dim(0) + let embedDim = tokenEmbeddings.dim(2) + var rows = [MLXArray]() + var maskRows = [[Int]]() + var maxWords = 0 - init(hiddenSize: Int) { - _linear1.wrappedValue = Linear(hiddenSize * 2, 256) - _linear2.wrappedValue = Linear(256, 128) - _linear3.wrappedValue = Linear(128, 1) - } + for batchIndex in 0.. 0 ? index : nil + } + + guard !tokenPositions.isEmpty else { + rows.append(MLXArray.zeros([0, embedDim], type: Float.self).asType(tokenEmbeddings.dtype)) + maskRows.append([]) + continue + } - func callAsFunction(textRep: MLXArray, labelRep: MLXArray) -> MLXArray { - let batchSize = labelRep.dim(0) - let numLabels = labelRep.dim(1) - let dim = labelRep.dim(2) - let expandedText = broadcast( - textRep.expandedDimensions(axis: 1), - to: [batchSize, numLabels, dim] + let positionArray = MLXArray(tokenPositions.map(Int32.init), [tokenPositions.count]) + let row = tokenEmbeddings[batchIndex].take(positionArray, axis: 0) + rows.append(row) + maxWords = max(maxWords, tokenPositions.count) + maskRows.append(Array(repeating: 1, count: tokenPositions.count)) + } + + maxWords = max(maxWords, 1) + return ( + openMedPadEmbeddings(rows, width: maxWords, embedDim: embedDim), + openMedPadMaskRows(maskRows, width: maxWords) ) - let combined = concatenated([expandedText, labelRep], axis: -1) - return linear3(relu(linear2(relu(linear1(combined))))).squeezed(axis: -1) } -} - -final class OpenMedGLiClassUniEncoderModel: Module { - private let configuration: OpenMedMLXBertConfiguration - - @ModuleInfo(key: "deberta") var deberta: OpenMedDebertaV2Model - @ModuleInfo(key: "classes_projector") var classesProjector: OpenMedGLiClassFeaturesProjector - @ModuleInfo(key: "text_projector") var textProjector: OpenMedGLiClassFeaturesProjector - @ModuleInfo(key: "segment_embeddings") var segmentEmbeddings: Embedding - @ModuleInfo(key: "scorer") var scorer: OpenMedGLiClassMLPScorer - @ParameterInfo(key: "logit_scale") var logitScale: MLXArray - - init(_ configuration: OpenMedMLXBertConfiguration) { - self.configuration = configuration - _deberta.wrappedValue = OpenMedDebertaV2Model(configuration) - _classesProjector.wrappedValue = OpenMedGLiClassFeaturesProjector( - encoderHiddenSize: configuration.encoderHiddenSize, - hiddenSize: configuration.hiddenSize - ) - _textProjector.wrappedValue = OpenMedGLiClassFeaturesProjector( - encoderHiddenSize: configuration.encoderHiddenSize, - hiddenSize: configuration.hiddenSize + + private func openMedGatherSpanEndpoints( + startHiddenStates: MLXArray, + endHiddenStates: MLXArray, + spanIndex: MLXArray + ) -> (MLXArray, MLXArray) { + let hiddenSize = startHiddenStates.dim(-1) + let startIndex = broadcast( + spanIndex[0..., 0..., 0].expandedDimensions(axis: -1), + to: [spanIndex.dim(0), spanIndex.dim(1), hiddenSize] ) - _segmentEmbeddings.wrappedValue = Embedding( - embeddingCount: 3, - dimensions: configuration.encoderHiddenSize + let endIndex = broadcast( + spanIndex[0..., 0..., 1].expandedDimensions(axis: -1), + to: [spanIndex.dim(0), spanIndex.dim(1), hiddenSize] ) - _scorer.wrappedValue = OpenMedGLiClassMLPScorer( - hiddenSize: configuration.encoderHiddenSize + return ( + takeAlong(startHiddenStates, startIndex, axis: 1), + takeAlong(endHiddenStates, endIndex, axis: 1) ) - _logitScale.wrappedValue = MLXArray(Float(configuration.logitScaleInitValue)) } - private func createSegmentIDs(_ inputIDs: MLXArray) -> MLXArray { - let batchSize = inputIDs.dim(0) - let seqLength = inputIDs.dim(1) - let textTokenID = configuration.textTokenIndex ?? -1 - let exampleTokenID = configuration.exampleTokenIndex ?? -1 - var rows = [[Int32]]() + private func openMedBuildAllEntityPairs( + spanRep: MLXArray, + spanMask: MLXArray + ) -> (MLXArray, MLXArray, MLXArray, MLXArray) { + let batchSize = spanRep.dim(0) + let embedDim = spanRep.dim(2) + var pairRows = [[[Int]]]() + var pairMaskRows = [[Int]]() + var headRows = [MLXArray]() + var tailRows = [MLXArray]() + var maxPairs = 0 for batchIndex in 0.. (MLXArray, MLXArray, MLXArray, MLXArray) { - let batchSize = encoderLayer.dim(0) - let seqLength = encoderLayer.dim(1) - let embedDim = encoderLayer.dim(2) - let classTokenID = configuration.classTokenIndex ?? -1 - let textTokenID = configuration.textTokenIndex ?? -1 - var classRows = [MLXArray]() - var maskRows = [[Int]]() - var maxClasses = 0 + final class OpenMedSpanMarkerV0: Module { + @ModuleInfo(key: "project_start") var projectStart: OpenMedProjectionMLP + @ModuleInfo(key: "project_end") var projectEnd: OpenMedProjectionMLP + @ModuleInfo(key: "out_project") var outputProjection: OpenMedProjectionMLP - for batchIndex in 0.. MLXArray { + let startRep = projectStart(hiddenStates) + let endRep = projectEnd(hiddenStates) + let (startSpanRep, endSpanRep) = openMedGatherSpanEndpoints( + startHiddenStates: startRep, + endHiddenStates: endRep, + spanIndex: spanIndex + ) + return outputProjection(relu(concatenated([startSpanRep, endSpanRep], axis: -1))) + } + } + + final class OpenMedGLiNERSpanModel: Module { + private let configuration: OpenMedMLXBertConfiguration + + @ModuleInfo(key: "deberta") var deberta: OpenMedDebertaV2Model + @ModuleInfo(key: "token_projection") var tokenProjection: Linear + @ModuleInfo(key: "rnn") var rnn: OpenMedBidirectionalLSTM? + @ModuleInfo(key: "span_rep_layer") var spanRepLayer: OpenMedSpanMarkerV0 + @ModuleInfo(key: "prompt_rep_layer") var promptRepLayer: OpenMedProjectionMLP + + init(_ configuration: OpenMedMLXBertConfiguration) { + self.configuration = configuration + _deberta.wrappedValue = OpenMedDebertaV2Model(configuration) + _tokenProjection.wrappedValue = Linear( + configuration.encoderHiddenSize, + configuration.hiddenSize + ) + if configuration.numRNNLayers > 0 { + _rnn.wrappedValue = OpenMedBidirectionalLSTM(hiddenSize: configuration.hiddenSize) } - let textStart = tokenIDs.firstIndex(of: textTokenID) ?? seqLength - var rowEmbeddings = [MLXArray]() - - for (classIndex, classPosition) in classPositions.enumerated() { - let startPosition = - configuration.embedClassToken - ? classPosition - : min(classPosition + 1, seqLength - 1) - let endPosition = - classIndex + 1 < classPositions.count - ? classPositions[classIndex + 1] - : textStart - if startPosition >= endPosition { - rowEmbeddings.append(encoderLayer[batchIndex, startPosition]) - continue - } + _spanRepLayer.wrappedValue = OpenMedSpanMarkerV0(hiddenSize: configuration.hiddenSize) + _promptRepLayer.wrappedValue = OpenMedProjectionMLP(inputSize: configuration.hiddenSize) + } - let classTokens = encoderLayer[batchIndex, startPosition.. 0, - sum(classTokens * classAttention.expandedDimensions(axis: -1), axis: 0) / denom, - mean(classTokens, axis: 0) - ) - rowEmbeddings.append(pooled) + private func encode( + inputIDs: MLXArray, + attentionMask: MLXArray, + wordsMask: MLXArray + ) -> (MLXArray, MLXArray, MLXArray, MLXArray) { + let hiddenStates = deberta(inputIDs: inputIDs, attentionMask: attentionMask) + let markerTokenID = configuration.classTokenIndex ?? 0 + let (promptEmbeddings, promptMask) = openMedExtractMarkerEmbeddings( + tokenEmbeddings: hiddenStates, + inputIDs: inputIDs, + markerTokenID: markerTokenID, + includeMarkerToken: configuration.embedEntityToken + ) + let (wordEmbeddings, wordMask) = openMedExtractWordEmbeddings( + tokenEmbeddings: hiddenStates, + wordsMask: wordsMask + ) + + var projectedWords = tokenProjection(wordEmbeddings) + let projectedPrompts = tokenProjection(promptEmbeddings) + if let rnn { + projectedWords = rnn(projectedWords, mask: wordMask) } + return (projectedPrompts, promptMask, projectedWords, wordMask) + } + + func callAsFunction( + inputIDs: MLXArray, + attentionMask: MLXArray, + wordsMask: MLXArray, + spanIndex: MLXArray, + spanMask: MLXArray + ) -> OpenMedGLiNERSpanOutput { + let maskedSpanIndex = spanIndex * spanMask.asType(spanIndex.dtype).expandedDimensions(axis: -1) + let (promptEmbeddings, promptMask, wordEmbeddings, wordMask) = encode( + inputIDs: inputIDs, + attentionMask: attentionMask, + wordsMask: wordsMask + ) + let spanRep = spanRepLayer(wordEmbeddings, spanIndex: maskedSpanIndex) + let promptRep = promptRepLayer(promptEmbeddings) + var logits = einsum("bsd,bcd->bsc", spanRep, promptRep) + logits = `where`( + spanMask.expandedDimensions(axis: -1), + logits, + openMedFloatScalar(-1.0e9, like: logits) + ) + return OpenMedGLiNERSpanOutput( + logits: logits, + spanIndex: maskedSpanIndex, + spanMask: spanMask, + promptMask: promptMask, + wordMask: wordMask + ) + } - let row: MLXArray - if rowEmbeddings.isEmpty { - row = MLXArray.zeros([0, embedDim], type: Float.self).asType(encoderLayer.dtype) - } else { - row = stacked(rowEmbeddings, axis: 0) + func sanitize(weights: [String: MLXArray]) -> [String: MLXArray] { + weights.filter { key, _ in + key != "deberta.embeddings.position_ids" && !key.hasPrefix("_") } - classRows.append(row) - maxClasses = max(maxClasses, row.dim(0)) - maskRows.append(Array(repeating: 1, count: row.dim(0))) } + } - maxClasses = max(maxClasses, 1) - return ( - openMedPadEmbeddings(classRows, width: maxClasses, embedDim: embedDim), - openMedPadMaskRows(maskRows, width: maxClasses), - encoderLayer, - attentionMask - ) + final class OpenMedGLiClassFeaturesProjector: Module { + @ModuleInfo(key: "linear1") var linear1: Linear + @ModuleInfo(key: "linear2") var linear2: Linear + + init(encoderHiddenSize: Int, hiddenSize: Int) { + _linear1.wrappedValue = Linear(encoderHiddenSize, hiddenSize) + _linear2.wrappedValue = Linear(hiddenSize, encoderHiddenSize) + } + + func callAsFunction(_ features: MLXArray) -> MLXArray { + linear2(gelu(linear1(features))) + } } - private func poolText(_ textEmbeddings: MLXArray, textMask: MLXArray) -> MLXArray { - switch configuration.poolingStrategy { - case "mean": - let mask = textMask.asType(textEmbeddings.dtype) - let denominator = maximum( - sum(mask, axis: 1, keepDims: true), - openMedFloatScalar(1.0, like: mask) + final class OpenMedGLiClassMLPScorer: Module { + @ModuleInfo(key: "linear1") var linear1: Linear + @ModuleInfo(key: "linear2") var linear2: Linear + @ModuleInfo(key: "linear3") var linear3: Linear + + init(hiddenSize: Int) { + _linear1.wrappedValue = Linear(hiddenSize * 2, 256) + _linear2.wrappedValue = Linear(256, 128) + _linear3.wrappedValue = Linear(128, 1) + } + + func callAsFunction(textRep: MLXArray, labelRep: MLXArray) -> MLXArray { + let batchSize = labelRep.dim(0) + let numLabels = labelRep.dim(1) + let dim = labelRep.dim(2) + let expandedText = broadcast( + textRep.expandedDimensions(axis: 1), + to: [batchSize, numLabels, dim] ) - return sum(textEmbeddings * mask.expandedDimensions(axis: -1), axis: 1) / denominator - default: - return textEmbeddings[0..., 0, 0...] + let combined = concatenated([expandedText, labelRep], axis: -1) + return linear3(relu(linear2(relu(linear1(combined))))).squeezed(axis: -1) } } - func callAsFunction(inputIDs: MLXArray, attentionMask: MLXArray) -> OpenMedGLiClassOutput { - var embedded = deberta.embeddings(inputIDs: inputIDs, attentionMask: attentionMask) - if configuration.useSegmentEmbeddings { - embedded = embedded + segmentEmbeddings(createSegmentIDs(inputIDs)) - } + final class OpenMedGLiClassUniEncoderModel: Module { + private let configuration: OpenMedMLXBertConfiguration - let hiddenStates = deberta.encoder(embedded, attentionMask: attentionMask) - let (classesEmbedding, classesMask, textEmbeddings, textMask) = extractClassFeatures( - encoderLayer: hiddenStates, - inputIDs: inputIDs, - attentionMask: attentionMask - ) - var pooledOutput = textProjector(poolText(textEmbeddings, textMask: textMask)) - var projectedClasses = classesProjector(classesEmbedding) + @ModuleInfo(key: "deberta") var deberta: OpenMedDebertaV2Model + @ModuleInfo(key: "classes_projector") var classesProjector: OpenMedGLiClassFeaturesProjector + @ModuleInfo(key: "text_projector") var textProjector: OpenMedGLiClassFeaturesProjector + @ModuleInfo(key: "segment_embeddings") var segmentEmbeddings: Embedding + @ModuleInfo(key: "scorer") var scorer: OpenMedGLiClassMLPScorer + @ParameterInfo(key: "logit_scale") var logitScale: MLXArray - if configuration.normalizeFeatures { - let pooledNorm = maximum( - sqrt(sum(pooledOutput * pooledOutput, axis: -1, keepDims: true)), - openMedFloatScalar(1.0e-8, like: pooledOutput) + init(_ configuration: OpenMedMLXBertConfiguration) { + self.configuration = configuration + _deberta.wrappedValue = OpenMedDebertaV2Model(configuration) + _classesProjector.wrappedValue = OpenMedGLiClassFeaturesProjector( + encoderHiddenSize: configuration.encoderHiddenSize, + hiddenSize: configuration.hiddenSize + ) + _textProjector.wrappedValue = OpenMedGLiClassFeaturesProjector( + encoderHiddenSize: configuration.encoderHiddenSize, + hiddenSize: configuration.hiddenSize ) - let classNorm = maximum( - sqrt(sum(projectedClasses * projectedClasses, axis: -1, keepDims: true)), - openMedFloatScalar(1.0e-8, like: projectedClasses) + _segmentEmbeddings.wrappedValue = Embedding( + embeddingCount: 3, + dimensions: configuration.encoderHiddenSize ) - pooledOutput = pooledOutput / pooledNorm - projectedClasses = projectedClasses / classNorm + _scorer.wrappedValue = OpenMedGLiClassMLPScorer( + hiddenSize: configuration.encoderHiddenSize + ) + _logitScale.wrappedValue = MLXArray(Float(configuration.logitScaleInitValue)) } - var logits = scorer(textRep: pooledOutput, labelRep: projectedClasses) - if configuration.normalizeFeatures { - logits = logits * logitScale + private func createSegmentIDs(_ inputIDs: MLXArray) -> MLXArray { + let batchSize = inputIDs.dim(0) + let seqLength = inputIDs.dim(1) + let textTokenID = configuration.textTokenIndex ?? -1 + let exampleTokenID = configuration.exampleTokenIndex ?? -1 + var rows = [[Int32]]() + + for batchIndex in 0.. [String: MLXArray] { - weights.filter { key, _ in - key != "deberta.embeddings.position_ids" && !key.hasPrefix("_") + private func extractClassFeatures( + encoderLayer: MLXArray, + inputIDs: MLXArray, + attentionMask: MLXArray + ) -> (MLXArray, MLXArray, MLXArray, MLXArray) { + let batchSize = encoderLayer.dim(0) + let seqLength = encoderLayer.dim(1) + let embedDim = encoderLayer.dim(2) + let classTokenID = configuration.classTokenIndex ?? -1 + let textTokenID = configuration.textTokenIndex ?? -1 + var classRows = [MLXArray]() + var maskRows = [[Int]]() + var maxClasses = 0 + + for batchIndex in 0..= endPosition { + rowEmbeddings.append(encoderLayer[batchIndex, startPosition]) + continue + } + + let classTokens = encoderLayer[batchIndex, startPosition.. 0, + sum(classTokens * classAttention.expandedDimensions(axis: -1), axis: 0) / denom, + mean(classTokens, axis: 0) + ) + rowEmbeddings.append(pooled) + } + + let row: MLXArray + if rowEmbeddings.isEmpty { + row = MLXArray.zeros([0, embedDim], type: Float.self).asType(encoderLayer.dtype) + } else { + row = stacked(rowEmbeddings, axis: 0) + } + classRows.append(row) + maxClasses = max(maxClasses, row.dim(0)) + maskRows.append(Array(repeating: 1, count: row.dim(0))) + } + + maxClasses = max(maxClasses, 1) + return ( + openMedPadEmbeddings(classRows, width: maxClasses, embedDim: embedDim), + openMedPadMaskRows(maskRows, width: maxClasses), + encoderLayer, + attentionMask + ) } - } -} - -final class OpenMedGLiNERTokenScorer: Module { - @ModuleInfo(key: "proj_token") var tokenProjection: Linear - @ModuleInfo(key: "proj_label") var labelProjection: Linear - @ModuleInfo(key: "out_linear1") var outputLinear1: Linear - @ModuleInfo(key: "out_linear2") var outputLinear2: Linear - - init(hiddenSize: Int) { - _tokenProjection.wrappedValue = Linear(hiddenSize, hiddenSize * 2) - _labelProjection.wrappedValue = Linear(hiddenSize, hiddenSize * 2) - _outputLinear1.wrappedValue = Linear(hiddenSize * 3, hiddenSize * 4) - _outputLinear2.wrappedValue = Linear(hiddenSize * 4, 3) - } - func callAsFunction(tokenRep: MLXArray, labelRep: MLXArray) -> MLXArray { - let batchSize = tokenRep.dim(0) - let seqLen = tokenRep.dim(1) - let hiddenSize = tokenRep.dim(2) - let numClasses = labelRep.dim(1) - let tokenProjected = tokenProjection(tokenRep) - .reshaped(batchSize, seqLen, 1, 2, hiddenSize) - let labelProjected = labelProjection(labelRep) - .reshaped(batchSize, 1, numClasses, 2, hiddenSize) - - let tokenLeft = broadcast( - tokenProjected[0..., 0..., 0..., 0, 0...], - to: [batchSize, seqLen, numClasses, hiddenSize] - ) - let tokenRight = broadcast( - tokenProjected[0..., 0..., 0..., 1, 0...], - to: [batchSize, seqLen, numClasses, hiddenSize] - ) - let labelLeft = broadcast( - labelProjected[0..., 0..., 0..., 0, 0...], - to: [batchSize, seqLen, numClasses, hiddenSize] - ) - let labelRight = broadcast( - labelProjected[0..., 0..., 0..., 1, 0...], - to: [batchSize, seqLen, numClasses, hiddenSize] - ) - let combined = concatenated([tokenLeft, labelLeft, tokenRight * labelRight], axis: -1) - return outputLinear2(relu(outputLinear1(combined))) - } -} - -final class OpenMedGLiNERRelexModel: Module { - private let configuration: OpenMedMLXBertConfiguration - - @ModuleInfo(key: "deberta") var deberta: OpenMedDebertaV2Model - @ModuleInfo(key: "token_projection") var tokenProjection: Linear? - @ModuleInfo(key: "rnn") var rnn: OpenMedBidirectionalLSTM? - @ModuleInfo(key: "scorer") var scorer: OpenMedGLiNERTokenScorer - @ModuleInfo(key: "span_rep_layer") var spanRepLayer: OpenMedSpanMarkerV0 - @ModuleInfo(key: "prompt_rep_layer") var promptRepLayer: OpenMedProjectionMLP - @ModuleInfo(key: "pair_rep_layer") var pairRepLayer: OpenMedProjectionMLP - - init(_ configuration: OpenMedMLXBertConfiguration) { - self.configuration = configuration - _deberta.wrappedValue = OpenMedDebertaV2Model(configuration) - if configuration.encoderHiddenSize != configuration.hiddenSize { - _tokenProjection.wrappedValue = Linear( - configuration.encoderHiddenSize, - configuration.hiddenSize + private func poolText(_ textEmbeddings: MLXArray, textMask: MLXArray) -> MLXArray { + switch configuration.poolingStrategy { + case "mean": + let mask = textMask.asType(textEmbeddings.dtype) + let denominator = maximum( + sum(mask, axis: 1, keepDims: true), + openMedFloatScalar(1.0, like: mask) + ) + return sum(textEmbeddings * mask.expandedDimensions(axis: -1), axis: 1) / denominator + default: + return textEmbeddings[0..., 0, 0...] + } + } + + func callAsFunction(inputIDs: MLXArray, attentionMask: MLXArray) -> OpenMedGLiClassOutput { + var embedded = deberta.embeddings(inputIDs: inputIDs, attentionMask: attentionMask) + if configuration.useSegmentEmbeddings { + embedded = embedded + segmentEmbeddings(createSegmentIDs(inputIDs)) + } + + let hiddenStates = deberta.encoder(embedded, attentionMask: attentionMask) + let (classesEmbedding, classesMask, textEmbeddings, textMask) = extractClassFeatures( + encoderLayer: hiddenStates, + inputIDs: inputIDs, + attentionMask: attentionMask ) + var pooledOutput = textProjector(poolText(textEmbeddings, textMask: textMask)) + var projectedClasses = classesProjector(classesEmbedding) + + if configuration.normalizeFeatures { + let pooledNorm = maximum( + sqrt(sum(pooledOutput * pooledOutput, axis: -1, keepDims: true)), + openMedFloatScalar(1.0e-8, like: pooledOutput) + ) + let classNorm = maximum( + sqrt(sum(projectedClasses * projectedClasses, axis: -1, keepDims: true)), + openMedFloatScalar(1.0e-8, like: projectedClasses) + ) + pooledOutput = pooledOutput / pooledNorm + projectedClasses = projectedClasses / classNorm + } + + var logits = scorer(textRep: pooledOutput, labelRep: projectedClasses) + if configuration.normalizeFeatures { + logits = logits * logitScale + } + return OpenMedGLiClassOutput(logits: logits, classesMask: classesMask) } - if configuration.numRNNLayers > 0 { - _rnn.wrappedValue = OpenMedBidirectionalLSTM(hiddenSize: configuration.hiddenSize) + + func sanitize(weights: [String: MLXArray]) -> [String: MLXArray] { + weights.filter { key, _ in + key != "deberta.embeddings.position_ids" && !key.hasPrefix("_") + } } - _scorer.wrappedValue = OpenMedGLiNERTokenScorer(hiddenSize: configuration.hiddenSize) - _spanRepLayer.wrappedValue = OpenMedSpanMarkerV0(hiddenSize: configuration.hiddenSize) - _promptRepLayer.wrappedValue = OpenMedProjectionMLP(inputSize: configuration.hiddenSize) - _pairRepLayer.wrappedValue = OpenMedProjectionMLP( - inputSize: configuration.hiddenSize * 2, - outputSize: configuration.hiddenSize - ) } - func encode( - inputIDs: MLXArray, - attentionMask: MLXArray, - wordsMask: MLXArray - ) -> OpenMedGLiNERRelexEncoded { - let hiddenStates = deberta(inputIDs: inputIDs, attentionMask: attentionMask) - let (entityPrompts, entityPromptMask) = openMedExtractMarkerEmbeddings( - tokenEmbeddings: hiddenStates, - inputIDs: inputIDs, - markerTokenID: configuration.classTokenIndex ?? 0, - includeMarkerToken: configuration.embedEntityToken - ) - let (relationPrompts, relationPromptMask) = openMedExtractMarkerEmbeddings( - tokenEmbeddings: hiddenStates, - inputIDs: inputIDs, - markerTokenID: configuration.relTokenIndex ?? 0, - includeMarkerToken: configuration.embedRelationToken ?? configuration.embedEntityToken - ) - let (wordEmbeddings, wordMask) = openMedExtractWordEmbeddings( - tokenEmbeddings: hiddenStates, - wordsMask: wordsMask - ) + final class OpenMedGLiNERTokenScorer: Module { + @ModuleInfo(key: "proj_token") var tokenProjection: Linear + @ModuleInfo(key: "proj_label") var labelProjection: Linear + @ModuleInfo(key: "out_linear1") var outputLinear1: Linear + @ModuleInfo(key: "out_linear2") var outputLinear2: Linear - var projectedEntities = entityPrompts - var projectedRelations = relationPrompts - var projectedWords = wordEmbeddings - if let tokenProjection { - projectedEntities = tokenProjection(projectedEntities) - projectedRelations = tokenProjection(projectedRelations) - projectedWords = tokenProjection(projectedWords) - } - if let rnn { - projectedWords = rnn(projectedWords, mask: wordMask) - } - - return OpenMedGLiNERRelexEncoded( - entityPrompts: projectedEntities, - entityPromptMask: entityPromptMask, - relationPrompts: projectedRelations, - relationPromptMask: relationPromptMask, - wordsEmbedding: projectedWords, - wordMask: wordMask - ) - } + init(hiddenSize: Int) { + _tokenProjection.wrappedValue = Linear(hiddenSize, hiddenSize * 2) + _labelProjection.wrappedValue = Linear(hiddenSize, hiddenSize * 2) + _outputLinear1.wrappedValue = Linear(hiddenSize * 3, hiddenSize * 4) + _outputLinear2.wrappedValue = Linear(hiddenSize * 4, 3) + } - func entityScores(encoded: OpenMedGLiNERRelexEncoded) -> MLXArray { - scorer( - tokenRep: encoded.wordsEmbedding, - labelRep: promptRepLayer(encoded.entityPrompts) - ) + func callAsFunction(tokenRep: MLXArray, labelRep: MLXArray) -> MLXArray { + let batchSize = tokenRep.dim(0) + let seqLen = tokenRep.dim(1) + let hiddenSize = tokenRep.dim(2) + let numClasses = labelRep.dim(1) + let tokenProjected = tokenProjection(tokenRep) + .reshaped(batchSize, seqLen, 1, 2, hiddenSize) + let labelProjected = labelProjection(labelRep) + .reshaped(batchSize, 1, numClasses, 2, hiddenSize) + + let tokenLeft = broadcast( + tokenProjected[0..., 0..., 0..., 0, 0...], + to: [batchSize, seqLen, numClasses, hiddenSize] + ) + let tokenRight = broadcast( + tokenProjected[0..., 0..., 0..., 1, 0...], + to: [batchSize, seqLen, numClasses, hiddenSize] + ) + let labelLeft = broadcast( + labelProjected[0..., 0..., 0..., 0, 0...], + to: [batchSize, seqLen, numClasses, hiddenSize] + ) + let labelRight = broadcast( + labelProjected[0..., 0..., 0..., 1, 0...], + to: [batchSize, seqLen, numClasses, hiddenSize] + ) + let combined = concatenated([tokenLeft, labelLeft, tokenRight * labelRight], axis: -1) + return outputLinear2(relu(outputLinear1(combined))) + } } - func relationScores( - encoded: OpenMedGLiNERRelexEncoded, - spanIndex: MLXArray, - spanMask: MLXArray - ) -> OpenMedGLiNERRelexRelationOutput { - let maskedSpanIndex = spanIndex * spanMask.asType(spanIndex.dtype).expandedDimensions(axis: -1) - let spanRep = spanRepLayer(encoded.wordsEmbedding, spanIndex: maskedSpanIndex) - let (pairIndex, pairMask, headRep, tailRep) = openMedBuildAllEntityPairs( - spanRep: spanRep, - spanMask: spanMask - ) - let pairRep = pairRepLayer(concatenated([headRep, tailRep], axis: -1)) - let pairScores = einsum("bnd,bcd->bnc", pairRep, encoded.relationPrompts) - return OpenMedGLiNERRelexRelationOutput( - pairScores: pairScores, - pairIndex: pairIndex, - pairMask: pairMask, - relationPromptMask: encoded.relationPromptMask - ) - } + final class OpenMedGLiNERRelexModel: Module { + private let configuration: OpenMedMLXBertConfiguration + + @ModuleInfo(key: "deberta") var deberta: OpenMedDebertaV2Model + @ModuleInfo(key: "token_projection") var tokenProjection: Linear? + @ModuleInfo(key: "rnn") var rnn: OpenMedBidirectionalLSTM? + @ModuleInfo(key: "scorer") var scorer: OpenMedGLiNERTokenScorer + @ModuleInfo(key: "span_rep_layer") var spanRepLayer: OpenMedSpanMarkerV0 + @ModuleInfo(key: "prompt_rep_layer") var promptRepLayer: OpenMedProjectionMLP + @ModuleInfo(key: "pair_rep_layer") var pairRepLayer: OpenMedProjectionMLP + + init(_ configuration: OpenMedMLXBertConfiguration) { + self.configuration = configuration + _deberta.wrappedValue = OpenMedDebertaV2Model(configuration) + if configuration.encoderHiddenSize != configuration.hiddenSize { + _tokenProjection.wrappedValue = Linear( + configuration.encoderHiddenSize, + configuration.hiddenSize + ) + } + if configuration.numRNNLayers > 0 { + _rnn.wrappedValue = OpenMedBidirectionalLSTM(hiddenSize: configuration.hiddenSize) + } + _scorer.wrappedValue = OpenMedGLiNERTokenScorer(hiddenSize: configuration.hiddenSize) + _spanRepLayer.wrappedValue = OpenMedSpanMarkerV0(hiddenSize: configuration.hiddenSize) + _promptRepLayer.wrappedValue = OpenMedProjectionMLP(inputSize: configuration.hiddenSize) + _pairRepLayer.wrappedValue = OpenMedProjectionMLP( + inputSize: configuration.hiddenSize * 2, + outputSize: configuration.hiddenSize + ) + } - func callAsFunction( - inputIDs: MLXArray, - attentionMask: MLXArray, - wordsMask: MLXArray - ) -> OpenMedGLiNERRelexEntityOutput { - let encoded = encode( - inputIDs: inputIDs, - attentionMask: attentionMask, - wordsMask: wordsMask - ) - return OpenMedGLiNERRelexEntityOutput( - entityScores: entityScores(encoded: encoded), - entityPromptMask: encoded.entityPromptMask, - wordMask: encoded.wordMask, - encoded: encoded - ) - } + func encode( + inputIDs: MLXArray, + attentionMask: MLXArray, + wordsMask: MLXArray + ) -> OpenMedGLiNERRelexEncoded { + let hiddenStates = deberta(inputIDs: inputIDs, attentionMask: attentionMask) + let (entityPrompts, entityPromptMask) = openMedExtractMarkerEmbeddings( + tokenEmbeddings: hiddenStates, + inputIDs: inputIDs, + markerTokenID: configuration.classTokenIndex ?? 0, + includeMarkerToken: configuration.embedEntityToken + ) + let (relationPrompts, relationPromptMask) = openMedExtractMarkerEmbeddings( + tokenEmbeddings: hiddenStates, + inputIDs: inputIDs, + markerTokenID: configuration.relTokenIndex ?? 0, + includeMarkerToken: configuration.embedRelationToken ?? configuration.embedEntityToken + ) + let (wordEmbeddings, wordMask) = openMedExtractWordEmbeddings( + tokenEmbeddings: hiddenStates, + wordsMask: wordsMask + ) + + var projectedEntities = entityPrompts + var projectedRelations = relationPrompts + var projectedWords = wordEmbeddings + if let tokenProjection { + projectedEntities = tokenProjection(projectedEntities) + projectedRelations = tokenProjection(projectedRelations) + projectedWords = tokenProjection(projectedWords) + } + if let rnn { + projectedWords = rnn(projectedWords, mask: wordMask) + } + + return OpenMedGLiNERRelexEncoded( + entityPrompts: projectedEntities, + entityPromptMask: entityPromptMask, + relationPrompts: projectedRelations, + relationPromptMask: relationPromptMask, + wordsEmbedding: projectedWords, + wordMask: wordMask + ) + } - func sanitize(weights: [String: MLXArray]) -> [String: MLXArray] { - weights.filter { key, _ in - key != "deberta.embeddings.position_ids" && !key.hasPrefix("_") + func entityScores(encoded: OpenMedGLiNERRelexEncoded) -> MLXArray { + scorer( + tokenRep: encoded.wordsEmbedding, + labelRep: promptRepLayer(encoded.entityPrompts) + ) + } + + func relationScores( + encoded: OpenMedGLiNERRelexEncoded, + spanIndex: MLXArray, + spanMask: MLXArray + ) -> OpenMedGLiNERRelexRelationOutput { + let maskedSpanIndex = spanIndex * spanMask.asType(spanIndex.dtype).expandedDimensions(axis: -1) + let spanRep = spanRepLayer(encoded.wordsEmbedding, spanIndex: maskedSpanIndex) + let (pairIndex, pairMask, headRep, tailRep) = openMedBuildAllEntityPairs( + spanRep: spanRep, + spanMask: spanMask + ) + let pairRep = pairRepLayer(concatenated([headRep, tailRep], axis: -1)) + let pairScores = einsum("bnd,bcd->bnc", pairRep, encoded.relationPrompts) + return OpenMedGLiNERRelexRelationOutput( + pairScores: pairScores, + pairIndex: pairIndex, + pairMask: pairMask, + relationPromptMask: encoded.relationPromptMask + ) + } + + func callAsFunction( + inputIDs: MLXArray, + attentionMask: MLXArray, + wordsMask: MLXArray + ) -> OpenMedGLiNERRelexEntityOutput { + let encoded = encode( + inputIDs: inputIDs, + attentionMask: attentionMask, + wordsMask: wordsMask + ) + return OpenMedGLiNERRelexEntityOutput( + entityScores: entityScores(encoded: encoded), + entityPromptMask: encoded.entityPromptMask, + wordMask: encoded.wordMask, + encoded: encoded + ) + } + + func sanitize(weights: [String: MLXArray]) -> [String: MLXArray] { + weights.filter { key, _ in + key != "deberta.embeddings.position_ids" && !key.hasPrefix("_") + } } } -} +#endif diff --git a/swift/OpenMedKit/Sources/OpenMedKit/OpenMedKit.swift b/swift/OpenMedKit/Sources/OpenMedKit/OpenMedKit.swift index be8cfce79..7ff70298a 100644 --- a/swift/OpenMedKit/Sources/OpenMedKit/OpenMedKit.swift +++ b/swift/OpenMedKit/Sources/OpenMedKit/OpenMedKit.swift @@ -1,1159 +1,1161 @@ -import CryptoKit -import Dispatch -import Foundation -import MLX -import Tokenizers - -/// OpenMedKit — On-device clinical NLP for iOS and macOS. -/// -/// Provides NER and PII detection using either CoreML or MLX models -/// produced by the OpenMed Python library. -/// -/// ## Quick Start -/// -/// ```swift -/// let modelDirectory = try await OpenMedModelStore.downloadMLXModel( -/// repoID: "OpenMed/OpenMed-PII-ClinicalE5-Small-33M-v1-mlx" -/// ) -/// let openmed = try OpenMed(backend: .mlx(modelDirectoryURL: modelDirectory)) -/// let entities = try openmed.extractPII("Patient John Doe, SSN 123-45-6789") -/// for entity in entities { -/// print(entity) // [first_name] "John Doe" (8:16) conf=0.95 -/// } -/// ``` -public final class OpenMed { - private enum Runtime { - case coreML(NERPipeline) - case mlx(MLXTokenClassificationPipeline) - case privacyFilter(OpenMedPrivacyFilterPipeline) - } - - private let runtime: Runtime - private let tokenizer: (any Tokenizer)? - private let maxSeqLength: Int - - /// Release cached MLX buffers that are no longer referenced by model objects. +#if canImport(MLX) && canImport(Tokenizers) && !os(watchOS) && !os(visionOS) + import CryptoKit + import Dispatch + import Foundation + import MLX + import Tokenizers + + /// OpenMedKit — On-device clinical NLP for iOS and macOS. /// - /// OpenMedKit runtimes own their model weights via ARC. Apps that swap between - /// large on-device models can first drop their runtime references, then call - /// this helper so MLX returns cached Metal buffers instead of carrying them - /// into the next model load. - public static func clearRuntimeMemoryCache() { - Memory.clearCache() - } + /// Provides NER and PII detection using either CoreML or MLX models + /// produced by the OpenMed Python library. + /// + /// ## Quick Start + /// + /// ```swift + /// let modelDirectory = try await OpenMedModelStore.downloadMLXModel( + /// repoID: "OpenMed/OpenMed-PII-ClinicalE5-Small-33M-v1-mlx" + /// ) + /// let openmed = try OpenMed(backend: .mlx(modelDirectoryURL: modelDirectory)) + /// let entities = try openmed.extractPII("Patient John Doe, SSN 123-45-6789") + /// for entity in entities { + /// print(entity) // [first_name] "John Doe" (8:16) conf=0.95 + /// } + /// ``` + public final class OpenMed { + private enum Runtime { + case coreML(NERPipeline) + case mlx(MLXTokenClassificationPipeline) + case privacyFilter(OpenMedPrivacyFilterPipeline) + } - /// Initialize OpenMed with an explicit backend. - public init( - backend: OpenMedBackend, - maxSeqLength: Int = 512 - ) throws { - switch backend { - case .coreML(let modelURL, let id2labelURL, let tokenizerName, let tokenizerFolderURL): - let pipeline = try NERPipeline( - modelURL: modelURL, - id2labelURL: id2labelURL, - maxSeqLength: maxSeqLength - ) - self.runtime = .coreML(pipeline) - self.tokenizer = try Self.loadTokenizer( - tokenizerName: tokenizerName, - tokenizerFolderURL: tokenizerFolderURL - ) - self.maxSeqLength = maxSeqLength + private let runtime: Runtime + private let tokenizer: (any Tokenizer)? + private let maxSeqLength: Int + + /// Release cached MLX buffers that are no longer referenced by model objects. + /// + /// OpenMedKit runtimes own their model weights via ARC. Apps that swap between + /// large on-device models can first drop their runtime references, then call + /// this helper so MLX returns cached Metal buffers instead of carrying them + /// into the next model load. + public static func clearRuntimeMemoryCache() { + Memory.clearCache() + } - case .mlx(let modelDirectoryURL): - let artifact = try OpenMedMLXArtifact(modelDirectoryURL: modelDirectoryURL) - if artifact.family == .openaiPrivacyFilter { - let pipeline = try OpenMedPrivacyFilterPipeline( - artifact: artifact, + /// Initialize OpenMed with an explicit backend. + public init( + backend: OpenMedBackend, + maxSeqLength: Int = 512 + ) throws { + switch backend { + case .coreML(let modelURL, let id2labelURL, let tokenizerName, let tokenizerFolderURL): + let pipeline = try NERPipeline( + modelURL: modelURL, + id2labelURL: id2labelURL, maxSeqLength: maxSeqLength ) - self.runtime = .privacyFilter(pipeline) - self.tokenizer = nil - self.maxSeqLength = pipeline.resolvedMaxSequenceLength - } else { - let pipeline = try MLXTokenClassificationPipeline( - modelDirectoryURL: modelDirectoryURL, - maxSeqLength: maxSeqLength - ) - self.runtime = .mlx(pipeline) + self.runtime = .coreML(pipeline) self.tokenizer = try Self.loadTokenizer( - tokenizerName: pipeline.tokenizerName ?? modelDirectoryURL.path, - tokenizerFolderURL: pipeline.tokenizerDirectoryURL + tokenizerName: tokenizerName, + tokenizerFolderURL: tokenizerFolderURL ) - self.maxSeqLength = pipeline.resolvedMaxSequenceLength + self.maxSeqLength = maxSeqLength + + case .mlx(let modelDirectoryURL): + let artifact = try OpenMedMLXArtifact(modelDirectoryURL: modelDirectoryURL) + if artifact.family == .openaiPrivacyFilter { + let pipeline = try OpenMedPrivacyFilterPipeline( + artifact: artifact, + maxSeqLength: maxSeqLength + ) + self.runtime = .privacyFilter(pipeline) + self.tokenizer = nil + self.maxSeqLength = pipeline.resolvedMaxSequenceLength + } else { + let pipeline = try MLXTokenClassificationPipeline( + modelDirectoryURL: modelDirectoryURL, + maxSeqLength: maxSeqLength + ) + self.runtime = .mlx(pipeline) + self.tokenizer = try Self.loadTokenizer( + tokenizerName: pipeline.tokenizerName ?? modelDirectoryURL.path, + tokenizerFolderURL: pipeline.tokenizerDirectoryURL + ) + self.maxSeqLength = pipeline.resolvedMaxSequenceLength + } } } - } - /// Initialize OpenMed with a CoreML model and tokenizer. - /// - /// - Parameters: - /// - modelURL: URL to the compiled CoreML model (`.mlmodelc` or `.mlpackage`). - /// - id2labelURL: URL to the `id2label.json` label mapping file. - /// - tokenizerName: HuggingFace tokenizer name for text tokenization. - /// - tokenizerFolderURL: Optional local tokenizer asset directory for offline use. - /// - maxSeqLength: Maximum token sequence length (default: 512). - public convenience init( - modelURL: URL, - id2labelURL: URL, - tokenizerName: String = "OpenMed/OpenMed-PII-ClinicalE5-Small-33M-v1", - tokenizerFolderURL: URL? = nil, - maxSeqLength: Int = 512 - ) throws { - try self.init( - backend: .coreML( - modelURL: modelURL, - id2labelURL: id2labelURL, - tokenizerName: tokenizerName, - tokenizerFolderURL: tokenizerFolderURL - ), - maxSeqLength: maxSeqLength - ) - } - - /// Run NER on the given text and return detected entities. - /// - /// - Parameters: - /// - text: Input clinical text. - /// - confidenceThreshold: Minimum confidence to include an entity (default: 0.5). - /// - Returns: Array of detected entities above the confidence threshold. - public func analyzeText( - _ text: String, - confidenceThreshold: Float = 0.5 - ) throws -> [EntityPrediction] { - let entities: [EntityPrediction] - switch runtime { - case .coreML(let pipeline): - let (inputIDs, attentionMask, _, offsets) = try tokenize(text) - entities = try pipeline.predict( - inputIds: inputIDs, - attentionMask: attentionMask, - offsets: offsets, - text: text - ) - case .mlx(let pipeline): - let (inputIDs, attentionMask, tokenTypeIDs, offsets) = try tokenize(text) - entities = try pipeline.predict( - inputIDs: inputIDs, - attentionMask: attentionMask, - tokenTypeIDs: tokenTypeIDs, - offsets: offsets, - text: text + /// Initialize OpenMed with a CoreML model and tokenizer. + /// + /// - Parameters: + /// - modelURL: URL to the compiled CoreML model (`.mlmodelc` or `.mlpackage`). + /// - id2labelURL: URL to the `id2label.json` label mapping file. + /// - tokenizerName: HuggingFace tokenizer name for text tokenization. + /// - tokenizerFolderURL: Optional local tokenizer asset directory for offline use. + /// - maxSeqLength: Maximum token sequence length (default: 512). + public convenience init( + modelURL: URL, + id2labelURL: URL, + tokenizerName: String = "OpenMed/OpenMed-PII-ClinicalE5-Small-33M-v1", + tokenizerFolderURL: URL? = nil, + maxSeqLength: Int = 512 + ) throws { + try self.init( + backend: .coreML( + modelURL: modelURL, + id2labelURL: id2labelURL, + tokenizerName: tokenizerName, + tokenizerFolderURL: tokenizerFolderURL + ), + maxSeqLength: maxSeqLength ) - case .privacyFilter(let pipeline): - entities = try pipeline.predict(text) } - return entities.filter { $0.confidence >= confidenceThreshold } - } - - /// Run PII detection on the given text with OpenMed's smart post-processing. - /// - /// This applies the same high-level PII cleanup used by the Python package: - /// grouped BIO spans, span repair, and semantic-unit merging for items such - /// as dates, SSNs, phone numbers, and emails. - public func extractPII( - _ text: String, - confidenceThreshold: Float = 0.5, - useSmartMerging: Bool = true - ) throws -> [EntityPrediction] { - let entities = try analyzeText(text, confidenceThreshold: confidenceThreshold) - let repairedEntities = PostProcessing.repairEntitySpans(entities, text: text) - - guard useSmartMerging else { - return repairedEntities - } + /// Run NER on the given text and return detected entities. + /// + /// - Parameters: + /// - text: Input clinical text. + /// - confidenceThreshold: Minimum confidence to include an entity (default: 0.5). + /// - Returns: Array of detected entities above the confidence threshold. + public func analyzeText( + _ text: String, + confidenceThreshold: Float = 0.5 + ) throws -> [EntityPrediction] { + let entities: [EntityPrediction] + switch runtime { + case .coreML(let pipeline): + let (inputIDs, attentionMask, _, offsets) = try tokenize(text) + entities = try pipeline.predict( + inputIds: inputIDs, + attentionMask: attentionMask, + offsets: offsets, + text: text + ) + case .mlx(let pipeline): + let (inputIDs, attentionMask, tokenTypeIDs, offsets) = try tokenize(text) + entities = try pipeline.predict( + inputIDs: inputIDs, + attentionMask: attentionMask, + tokenTypeIDs: tokenTypeIDs, + offsets: offsets, + text: text + ) + case .privacyFilter(let pipeline): + entities = try pipeline.predict(text) + } - switch runtime { - case .privacyFilter: - return PostProcessing.mergePIIEntities( - repairedEntities, - text: text, - useSemanticPatterns: true, - preferModelLabels: true, - allowSemanticOnlyMatches: false, - allowSemanticLabelExpansion: false - ) - case .coreML, .mlx: - return PostProcessing.mergePIIEntities( - repairedEntities, - text: text, - useSemanticPatterns: true, - preferModelLabels: true - ) + return entities.filter { $0.confidence >= confidenceThreshold } } - } - - /// De-identify text under a bundled policy profile. - /// - /// Pass `Policy.defaultName` to use the default `hipaa_safe_harbor` - /// posture. The policy argument is explicit so existing method-based - /// `deidentify(_:)` call sites keep resolving to mask redaction. This path - /// does not write the input text or detected span text to stdout, stderr, - /// or logs. - public func deidentify( - _ text: String, - policy: String, - confidenceThreshold: Float = 0.5, - useSmartMerging: Bool = true - ) throws -> PolicyDeidentificationResult { - let resolvedPolicy = try Policy(named: policy) - return try deidentify( - text, - policy: resolvedPolicy, - confidenceThreshold: confidenceThreshold, - useSmartMerging: useSmartMerging - ) - } - /// De-identify text under an already loaded policy profile. - /// - /// Detected span offsets in the returned action records reference the - /// original input text even when replacement lengths differ. - public func deidentify( - _ text: String, - policy: Policy, - confidenceThreshold: Float = 0.5, - useSmartMerging: Bool = true - ) throws -> PolicyDeidentificationResult { - let entities = try extractPII( - text, - confidenceThreshold: confidenceThreshold, - useSmartMerging: useSmartMerging - ) - return Self.deidentify(text, entities: entities, policy: policy) - } + /// Run PII detection on the given text with OpenMed's smart post-processing. + /// + /// This applies the same high-level PII cleanup used by the Python package: + /// grouped BIO spans, span repair, and semantic-unit merging for items such + /// as dates, SSNs, phone numbers, and emails. + public func extractPII( + _ text: String, + confidenceThreshold: Float = 0.5, + useSmartMerging: Bool = true + ) throws -> [EntityPrediction] { + let entities = try analyzeText(text, confidenceThreshold: confidenceThreshold) + let repairedEntities = PostProcessing.repairEntitySpans(entities, text: text) + + guard useSmartMerging else { + return repairedEntities + } - /// Detect and de-identify PII, returning a Python-schema-compatible result. - public func deidentify( - _ text: String, - method: DeidentificationMethod = .mask, - confidenceThreshold: Float = 0.5, - useSmartMerging: Bool = true - ) throws -> DeidentificationResult { - let entities = try extractPII( - text, - confidenceThreshold: confidenceThreshold, - useSmartMerging: useSmartMerging - ) - let deidentifiedText = Self.deidentifiedText( - text, - entities: entities, - method: method - ) - return DeidentificationResult( - originalText: text, - deidentifiedText: deidentifiedText, - entities: entities, - method: method.rawValue - ) - } + switch runtime { + case .privacyFilter: + return PostProcessing.mergePIIEntities( + repairedEntities, + text: text, + useSemanticPatterns: true, + preferModelLabels: true, + allowSemanticOnlyMatches: false, + allowSemanticLabelExpansion: false + ) + case .coreML, .mlx: + return PostProcessing.mergePIIEntities( + repairedEntities, + text: text, + useSemanticPatterns: true, + preferModelLabels: true + ) + } + } - /// Run PII detection over long text using overlapping token windows. - /// - /// The returned entity offsets always reference the original full text. - /// Overlapping duplicate detections are merged before the final smart - /// semantic merge pass. - public func extractPIIChunked( - _ text: String, - confidenceThreshold: Float = 0.5, - chunkTokenLimit: Int = 256, - tokenOverlap: Int = 32, - useSmartMerging: Bool = true - ) throws -> [EntityPrediction] { - let chunks = try makeTokenChunks( - for: text, - chunkTokenLimit: chunkTokenLimit, - tokenOverlap: tokenOverlap - ) - guard chunks.count > 1 else { - return try extractPII( + /// De-identify text under a bundled policy profile. + /// + /// Pass `Policy.defaultName` to use the default `hipaa_safe_harbor` + /// posture. The policy argument is explicit so existing method-based + /// `deidentify(_:)` call sites keep resolving to mask redaction. This path + /// does not write the input text or detected span text to stdout, stderr, + /// or logs. + public func deidentify( + _ text: String, + policy: String, + confidenceThreshold: Float = 0.5, + useSmartMerging: Bool = true + ) throws -> PolicyDeidentificationResult { + let resolvedPolicy = try Policy(named: policy) + return try deidentify( text, + policy: resolvedPolicy, confidenceThreshold: confidenceThreshold, useSmartMerging: useSmartMerging ) } - var chunkEntities: [EntityPrediction] = [] - for chunk in chunks { - let chunkText = Self.substring(text, start: chunk.start, end: chunk.end) + /// De-identify text under an already loaded policy profile. + /// + /// Detected span offsets in the returned action records reference the + /// original input text even when replacement lengths differ. + public func deidentify( + _ text: String, + policy: Policy, + confidenceThreshold: Float = 0.5, + useSmartMerging: Bool = true + ) throws -> PolicyDeidentificationResult { let entities = try extractPII( - chunkText, + text, confidenceThreshold: confidenceThreshold, useSmartMerging: useSmartMerging ) - chunkEntities.append( - contentsOf: entities.compactMap { entity in - Self.offset(entity, by: chunk.start, in: text) - }) + return Self.deidentify(text, entities: entities, policy: policy) } - return mergeChunkedPIIEntities( - chunkEntities, - text: text, - useSmartMerging: useSmartMerging - ) - } + /// Detect and de-identify PII, returning a Python-schema-compatible result. + public func deidentify( + _ text: String, + method: DeidentificationMethod = .mask, + confidenceThreshold: Float = 0.5, + useSmartMerging: Bool = true + ) throws -> DeidentificationResult { + let entities = try extractPII( + text, + confidenceThreshold: confidenceThreshold, + useSmartMerging: useSmartMerging + ) + let deidentifiedText = Self.deidentifiedText( + text, + entities: entities, + method: method + ) + return DeidentificationResult( + originalText: text, + deidentifiedText: deidentifiedText, + entities: entities, + method: method.rawValue + ) + } - // MARK: - Private + /// Run PII detection over long text using overlapping token windows. + /// + /// The returned entity offsets always reference the original full text. + /// Overlapping duplicate detections are merged before the final smart + /// semantic merge pass. + public func extractPIIChunked( + _ text: String, + confidenceThreshold: Float = 0.5, + chunkTokenLimit: Int = 256, + tokenOverlap: Int = 32, + useSmartMerging: Bool = true + ) throws -> [EntityPrediction] { + let chunks = try makeTokenChunks( + for: text, + chunkTokenLimit: chunkTokenLimit, + tokenOverlap: tokenOverlap + ) + guard chunks.count > 1 else { + return try extractPII( + text, + confidenceThreshold: confidenceThreshold, + useSmartMerging: useSmartMerging + ) + } - struct TextChunk: Equatable { - let start: Int - let end: Int - let tokenStart: Int - let tokenEnd: Int - } + var chunkEntities: [EntityPrediction] = [] + for chunk in chunks { + let chunkText = Self.substring(text, start: chunk.start, end: chunk.end) + let entities = try extractPII( + chunkText, + confidenceThreshold: confidenceThreshold, + useSmartMerging: useSmartMerging + ) + chunkEntities.append( + contentsOf: entities.compactMap { entity in + Self.offset(entity, by: chunk.start, in: text) + }) + } - private func tokenize(_ text: String) throws -> ([Int], [Int], [Int], [(Int, Int)]) { - // Use swift-transformers for tokenization - // This ensures token IDs match the Python HuggingFace tokenizer - guard let tokenizer else { - throw TokenizerError.missingConfig + return mergeChunkedPIIEntities( + chunkEntities, + text: text, + useSmartMerging: useSmartMerging + ) } - let inputIds = Array(tokenizer(text, addSpecialTokens: true).prefix(maxSeqLength)) - let tokens = tokenizer.convertIdsToTokens(inputIds).map { $0 ?? "" } - let attentionMask = Array(repeating: 1, count: inputIds.count) - let tokenTypeIDs = Array(repeating: 0, count: inputIds.count) - let offsets = Self.buildOffsets(tokens: tokens, in: text) - return (inputIds, attentionMask, tokenTypeIDs, offsets) - } + // MARK: - Private - func makeTokenChunks( - for text: String, - chunkTokenLimit: Int, - tokenOverlap: Int - ) throws -> [TextChunk] { - guard !text.isEmpty else { - return [] + struct TextChunk: Equatable { + let start: Int + let end: Int + let tokenStart: Int + let tokenEnd: Int } - let tokenOffsets = try tokenOffsets(in: text) - .filter { $0.0 < $0.1 } - - let tokenLimit = max(1, chunkTokenLimit) - guard tokenOffsets.count > tokenLimit else { - return [ - TextChunk( - start: 0, - end: text.count, - tokenStart: 0, - tokenEnd: tokenOffsets.count - ) - ] + private func tokenize(_ text: String) throws -> ([Int], [Int], [Int], [(Int, Int)]) { + // Use swift-transformers for tokenization + // This ensures token IDs match the Python HuggingFace tokenizer + guard let tokenizer else { + throw TokenizerError.missingConfig + } + let inputIds = Array(tokenizer(text, addSpecialTokens: true).prefix(maxSeqLength)) + let tokens = tokenizer.convertIdsToTokens(inputIds).map { $0 ?? "" } + let attentionMask = Array(repeating: 1, count: inputIds.count) + let tokenTypeIDs = Array(repeating: 0, count: inputIds.count) + let offsets = Self.buildOffsets(tokens: tokens, in: text) + + return (inputIds, attentionMask, tokenTypeIDs, offsets) } - let overlap = min(max(0, tokenOverlap), tokenLimit - 1) - var chunks: [TextChunk] = [] - var tokenStart = 0 - - while tokenStart < tokenOffsets.count { - let tokenEnd = min(tokenStart + tokenLimit, tokenOffsets.count) - chunks.append( - TextChunk( - start: tokenOffsets[tokenStart].0, - end: tokenOffsets[tokenEnd - 1].1, - tokenStart: tokenStart, - tokenEnd: tokenEnd - ) - ) + func makeTokenChunks( + for text: String, + chunkTokenLimit: Int, + tokenOverlap: Int + ) throws -> [TextChunk] { + guard !text.isEmpty else { + return [] + } - guard tokenEnd < tokenOffsets.count else { - break + let tokenOffsets = try tokenOffsets(in: text) + .filter { $0.0 < $0.1 } + + let tokenLimit = max(1, chunkTokenLimit) + guard tokenOffsets.count > tokenLimit else { + return [ + TextChunk( + start: 0, + end: text.count, + tokenStart: 0, + tokenEnd: tokenOffsets.count + ) + ] } - tokenStart = max(tokenStart + 1, tokenEnd - overlap) - } - return chunks - } + let overlap = min(max(0, tokenOverlap), tokenLimit - 1) + var chunks: [TextChunk] = [] + var tokenStart = 0 + + while tokenStart < tokenOffsets.count { + let tokenEnd = min(tokenStart + tokenLimit, tokenOffsets.count) + chunks.append( + TextChunk( + start: tokenOffsets[tokenStart].0, + end: tokenOffsets[tokenEnd - 1].1, + tokenStart: tokenStart, + tokenEnd: tokenEnd + ) + ) - private func tokenOffsets(in text: String) throws -> [(Int, Int)] { - switch runtime { - case .coreML, .mlx: - guard let tokenizer else { - throw TokenizerError.missingConfig + guard tokenEnd < tokenOffsets.count else { + break + } + tokenStart = max(tokenStart + 1, tokenEnd - overlap) } - let inputIDs = tokenizer(text, addSpecialTokens: false) - let tokens = tokenizer.convertIdsToTokens(inputIDs).map { $0 ?? "" } - return Self.buildOffsets(tokens: tokens, in: text) - case .privacyFilter(let pipeline): - return try pipeline.tokenOffsets(in: text) + + return chunks } - } - func mergeChunkedPIIEntities( - _ entities: [EntityPrediction], - text: String, - useSmartMerging: Bool - ) -> [EntityPrediction] { - let repaired = PostProcessing.repairEntitySpans( - Self.deduplicateOverlappingEntities(entities), - text: text - ) - - guard useSmartMerging else { - return Self.deduplicateOverlappingEntities(repaired) + private func tokenOffsets(in text: String) throws -> [(Int, Int)] { + switch runtime { + case .coreML, .mlx: + guard let tokenizer else { + throw TokenizerError.missingConfig + } + let inputIDs = tokenizer(text, addSpecialTokens: false) + let tokens = tokenizer.convertIdsToTokens(inputIDs).map { $0 ?? "" } + return Self.buildOffsets(tokens: tokens, in: text) + case .privacyFilter(let pipeline): + return try pipeline.tokenOffsets(in: text) + } } - let merged: [EntityPrediction] - switch runtime { - case .privacyFilter: - merged = PostProcessing.mergePIIEntities( - repaired, - text: text, - useSemanticPatterns: true, - preferModelLabels: true, - allowSemanticOnlyMatches: false, - allowSemanticLabelExpansion: false - ) - case .coreML, .mlx: - merged = PostProcessing.mergePIIEntities( - repaired, - text: text, - useSemanticPatterns: true, - preferModelLabels: true + func mergeChunkedPIIEntities( + _ entities: [EntityPrediction], + text: String, + useSmartMerging: Bool + ) -> [EntityPrediction] { + let repaired = PostProcessing.repairEntitySpans( + Self.deduplicateOverlappingEntities(entities), + text: text ) - } - return Self.deduplicateOverlappingEntities(merged) - } + guard useSmartMerging else { + return Self.deduplicateOverlappingEntities(repaired) + } - static func deidentify( - _ text: String, - entities: [EntityPrediction], - policy: Policy - ) -> PolicyDeidentificationResult { - let candidates = deidentificationCandidates(entities, in: text) - var redactedText = text - var actionRecords: [DeidentifiedSpanAction] = [] - - for entity in candidates { - let canonicalLabel = Policy.canonicalLabel(for: entity.label) - let action = policy.action(for: entity.label) - let original = substring(text, start: entity.start, end: entity.end) - let replacement = replacementText( - for: action, - canonicalLabel: canonicalLabel, - original: original - ) + let merged: [EntityPrediction] + switch runtime { + case .privacyFilter: + merged = PostProcessing.mergePIIEntities( + repaired, + text: text, + useSemanticPatterns: true, + preferModelLabels: true, + allowSemanticOnlyMatches: false, + allowSemanticLabelExpansion: false + ) + case .coreML, .mlx: + merged = PostProcessing.mergePIIEntities( + repaired, + text: text, + useSemanticPatterns: true, + preferModelLabels: true + ) + } - actionRecords.append( - DeidentifiedSpanAction( - label: entity.label, + return Self.deduplicateOverlappingEntities(merged) + } + + static func deidentify( + _ text: String, + entities: [EntityPrediction], + policy: Policy + ) -> PolicyDeidentificationResult { + let candidates = deidentificationCandidates(entities, in: text) + var redactedText = text + var actionRecords: [DeidentifiedSpanAction] = [] + + for entity in candidates { + let canonicalLabel = Policy.canonicalLabel(for: entity.label) + let action = policy.action(for: entity.label) + let original = substring(text, start: entity.start, end: entity.end) + let replacement = replacementText( + for: action, canonicalLabel: canonicalLabel, - action: action, - start: entity.start, - end: entity.end, - confidence: entity.confidence, - replacement: replacement + original: original ) - ) - } - for record in actionRecords.reversed() { - guard let replacement = record.replacement else { - continue + actionRecords.append( + DeidentifiedSpanAction( + label: entity.label, + canonicalLabel: canonicalLabel, + action: action, + start: entity.start, + end: entity.end, + confidence: entity.confidence, + replacement: replacement + ) + ) } - let lowerBound = redactedText.index( - redactedText.startIndex, - offsetBy: record.start - ) - let upperBound = redactedText.index( - lowerBound, - offsetBy: record.end - record.start + + for record in actionRecords.reversed() { + guard let replacement = record.replacement else { + continue + } + let lowerBound = redactedText.index( + redactedText.startIndex, + offsetBy: record.start + ) + let upperBound = redactedText.index( + lowerBound, + offsetBy: record.end - record.start + ) + redactedText.replaceSubrange(lowerBound.. [EntityPrediction] { - let textLength = text.count - let sorted = - entities - .filter { $0.start >= 0 && $0.end > $0.start && $0.end <= textLength } - .sorted { - if $0.start == $1.start { - let lhsLength = entityLength($0) - let rhsLength = entityLength($1) - if lhsLength == rhsLength { - return $0.confidence > $1.confidence + private static func deidentificationCandidates( + _ entities: [EntityPrediction], + in text: String + ) -> [EntityPrediction] { + let textLength = text.count + let sorted = + entities + .filter { $0.start >= 0 && $0.end > $0.start && $0.end <= textLength } + .sorted { + if $0.start == $1.start { + let lhsLength = entityLength($0) + let rhsLength = entityLength($1) + if lhsLength == rhsLength { + return $0.confidence > $1.confidence + } + return lhsLength > rhsLength } - return lhsLength > rhsLength + return $0.start < $1.start } - return $0.start < $1.start - } - var selected: [EntityPrediction] = [] - for entity in sorted { - let overlapsSelected = selected.contains { existing in - entity.start < existing.end && entity.end > existing.start + var selected: [EntityPrediction] = [] + for entity in sorted { + let overlapsSelected = selected.contains { existing in + entity.start < existing.end && entity.end > existing.start + } + if !overlapsSelected { + selected.append(entity) + } } - if !overlapsSelected { - selected.append(entity) + + return selected.sorted { + if $0.start == $1.start { + return $0.end < $1.end + } + return $0.start < $1.start } } - return selected.sorted { - if $0.start == $1.start { - return $0.end < $1.end + private static func replacementText( + for action: PolicyAction, + canonicalLabel: String, + original: String + ) -> String? { + switch action.redactionEquivalent { + case .keep: + return nil + case .mask: + return "[\(canonicalLabel)]" + case .replace: + return "[\(canonicalLabel)_REPLACED]" + case .remove: + return "" + case .hash: + return "\(canonicalLabel)_\(stableHash(original))" + case .redact: + return "[\(canonicalLabel)]" } - return $0.start < $1.start } - } - private static func replacementText( - for action: PolicyAction, - canonicalLabel: String, - original: String - ) -> String? { - switch action.redactionEquivalent { - case .keep: - return nil - case .mask: - return "[\(canonicalLabel)]" - case .replace: - return "[\(canonicalLabel)_REPLACED]" - case .remove: - return "" - case .hash: - return "\(canonicalLabel)_\(stableHash(original))" - case .redact: - return "[\(canonicalLabel)]" + private static func stableHash(_ text: String) -> String { + let digest = SHA256.hash(data: Data(text.utf8)) + return digest.prefix(12).map { String(format: "%02x", $0) }.joined() } - } - private static func stableHash(_ text: String) -> String { - let digest = SHA256.hash(data: Data(text.utf8)) - return digest.prefix(12).map { String(format: "%02x", $0) }.joined() - } - - static func deduplicateOverlappingEntities( - _ entities: [EntityPrediction] - ) -> [EntityPrediction] { - var selected: [EntityPrediction] = [] + static func deduplicateOverlappingEntities( + _ entities: [EntityPrediction] + ) -> [EntityPrediction] { + var selected: [EntityPrediction] = [] + + for entity in entities.sorted(by: entitySort) { + guard + let existingIndex = selected.firstIndex(where: { + areDuplicateCandidates(entity, $0) + }) + else { + selected.append(entity) + continue + } - for entity in entities.sorted(by: entitySort) { - guard - let existingIndex = selected.firstIndex(where: { - areDuplicateCandidates(entity, $0) - }) - else { - selected.append(entity) - continue + if isBetterDuplicate(candidate: entity, existing: selected[existingIndex]) { + selected[existingIndex] = entity + } } - if isBetterDuplicate(candidate: entity, existing: selected[existingIndex]) { - selected[existingIndex] = entity + return selected.sorted { + if $0.start == $1.start { + return $0.end < $1.end + } + return $0.start < $1.start } } - return selected.sorted { - if $0.start == $1.start { - return $0.end < $1.end + private static func offset( + _ entity: EntityPrediction, + by baseOffset: Int, + in text: String + ) -> EntityPrediction? { + let start = entity.start + baseOffset + let end = entity.end + baseOffset + guard start >= 0, end > start, end <= text.count else { + return nil } - return $0.start < $1.start + return EntityPrediction( + label: entity.label, + text: substring(text, start: start, end: end), + confidence: entity.confidence, + start: start, + end: end + ) } - } - private static func offset( - _ entity: EntityPrediction, - by baseOffset: Int, - in text: String - ) -> EntityPrediction? { - let start = entity.start + baseOffset - let end = entity.end + baseOffset - guard start >= 0, end > start, end <= text.count else { - return nil + private static func entitySort( + lhs: EntityPrediction, + rhs: EntityPrediction + ) -> Bool { + if lhs.start == rhs.start { + if lhs.end == rhs.end { + return lhs.confidence > rhs.confidence + } + return entityLength(lhs) > entityLength(rhs) + } + return lhs.start < rhs.start } - return EntityPrediction( - label: entity.label, - text: substring(text, start: start, end: end), - confidence: entity.confidence, - start: start, - end: end - ) - } - private static func entitySort( - lhs: EntityPrediction, - rhs: EntityPrediction - ) -> Bool { - if lhs.start == rhs.start { - if lhs.end == rhs.end { - return lhs.confidence > rhs.confidence + private static func areDuplicateCandidates( + _ lhs: EntityPrediction, + _ rhs: EntityPrediction + ) -> Bool { + guard labelsAreCompatible(lhs.label, rhs.label) else { + return false } - return entityLength(lhs) > entityLength(rhs) + let overlap = min(lhs.end, rhs.end) - max(lhs.start, rhs.start) + guard overlap > 0 else { + return false + } + let shorterLength = max(1, min(entityLength(lhs), entityLength(rhs))) + return Double(overlap) / Double(shorterLength) >= 0.5 } - return lhs.start < rhs.start - } - private static func areDuplicateCandidates( - _ lhs: EntityPrediction, - _ rhs: EntityPrediction - ) -> Bool { - guard labelsAreCompatible(lhs.label, rhs.label) else { - return false - } - let overlap = min(lhs.end, rhs.end) - max(lhs.start, rhs.start) - guard overlap > 0 else { - return false - } - let shorterLength = max(1, min(entityLength(lhs), entityLength(rhs))) - return Double(overlap) / Double(shorterLength) >= 0.5 - } + private static func labelsAreCompatible(_ lhs: String, _ rhs: String) -> Bool { + if lhs == rhs { + return true + } - private static func labelsAreCompatible(_ lhs: String, _ rhs: String) -> Bool { - if lhs == rhs { - return true - } + let normalizedLHS = PostProcessing.normalizeLabel(lhs) + let normalizedRHS = PostProcessing.normalizeLabel(rhs) + if normalizedLHS == normalizedRHS { + return true + } - let normalizedLHS = PostProcessing.normalizeLabel(lhs) - let normalizedRHS = PostProcessing.normalizeLabel(rhs) - if normalizedLHS == normalizedRHS { - return true + let lowerLHS = lhs.lowercased() + let lowerRHS = rhs.lowercased() + let nameTokens = ["name", "person"] + return nameTokens.contains { token in + lowerLHS.contains(token) && lowerRHS.contains(token) + } } - let lowerLHS = lhs.lowercased() - let lowerRHS = rhs.lowercased() - let nameTokens = ["name", "person"] - return nameTokens.contains { token in - lowerLHS.contains(token) && lowerRHS.contains(token) - } - } + private static func isBetterDuplicate( + candidate: EntityPrediction, + existing: EntityPrediction + ) -> Bool { + let candidateLength = entityLength(candidate) + let existingLength = entityLength(existing) - private static func isBetterDuplicate( - candidate: EntityPrediction, - existing: EntityPrediction - ) -> Bool { - let candidateLength = entityLength(candidate) - let existingLength = entityLength(existing) + if candidate.start == existing.start && candidate.end == existing.end { + return candidate.confidence > existing.confidence + } - if candidate.start == existing.start && candidate.end == existing.end { - return candidate.confidence > existing.confidence - } + if candidateLength > existingLength && candidate.confidence >= existing.confidence - 0.10 { + return true + } - if candidateLength > existingLength && candidate.confidence >= existing.confidence - 0.10 { - return true + return candidate.confidence > existing.confidence + 0.05 } - return candidate.confidence > existing.confidence + 0.05 - } - - private static func entityLength(_ entity: EntityPrediction) -> Int { - max(0, entity.end - entity.start) - } - - private static func substring(_ text: String, start: Int, end: Int) -> String { - guard start >= 0, end >= start, end <= text.count else { - return "" + private static func entityLength(_ entity: EntityPrediction) -> Int { + max(0, entity.end - entity.start) } - let lowerBound = text.index(text.startIndex, offsetBy: start) - let upperBound = text.index(lowerBound, offsetBy: end - start) - return String(text[lowerBound.. String { - var redacted = text - for entity in entities.sorted(by: { $0.start > $1.start }) { - guard let range = characterRange(in: redacted, start: entity.start, end: entity.end) else { - continue + private static func substring(_ text: String, start: Int, end: Int) -> String { + guard start >= 0, end >= start, end <= text.count else { + return "" } - redacted.replaceSubrange( - range, - with: replacementText(for: entity, method: method) - ) + let lowerBound = text.index(text.startIndex, offsetBy: start) + let upperBound = text.index(lowerBound, offsetBy: end - start) + return String(text[lowerBound.. String { - switch method { - case .mask: - return "[\(entity.entityType.uppercased())]" - case .remove: - return "" + private static func deidentifiedText( + _ text: String, + entities: [EntityPrediction], + method: DeidentificationMethod + ) -> String { + var redacted = text + for entity in entities.sorted(by: { $0.start > $1.start }) { + guard let range = characterRange(in: redacted, start: entity.start, end: entity.end) else { + continue + } + redacted.replaceSubrange( + range, + with: replacementText(for: entity, method: method) + ) + } + return redacted } - } - private static func characterRange( - in text: String, - start: Int, - end: Int - ) -> Range? { - guard start >= 0, end >= start, end <= text.count else { - return nil + private static func replacementText( + for entity: EntityPrediction, + method: DeidentificationMethod + ) -> String { + switch method { + case .mask: + return "[\(entity.entityType.uppercased())]" + case .remove: + return "" + } } - let lowerBound = text.index(text.startIndex, offsetBy: start) - let upperBound = text.index(lowerBound, offsetBy: end - start) - return lowerBound.. any Tokenizer { - let semaphore = DispatchSemaphore(value: 0) - var result: Result? - Task.detached { - do { - let tokenizer = try await loadTokenizerAsync( - tokenizerName: tokenizerName, - tokenizerFolderURL: tokenizerFolderURL - ) - result = .success(tokenizer) - } catch { - result = .failure(error) + private static func characterRange( + in text: String, + start: Int, + end: Int + ) -> Range? { + guard start >= 0, end >= start, end <= text.count else { + return nil } - semaphore.signal() + let lowerBound = text.index(text.startIndex, offsetBy: start) + let upperBound = text.index(lowerBound, offsetBy: end - start) + return lowerBound.. any Tokenizer { - if let tokenizerFolderURL { - return try loadTokenizerFromDirectory( - tokenizerFolderURL, - fallbackTokenizerName: tokenizerName - ) - } + static func loadTokenizer( + tokenizerName: String, + tokenizerFolderURL: URL? + ) throws -> any Tokenizer { + let semaphore = DispatchSemaphore(value: 0) + var result: Result? + + Task.detached { + do { + let tokenizer = try await loadTokenizerAsync( + tokenizerName: tokenizerName, + tokenizerFolderURL: tokenizerFolderURL + ) + result = .success(tokenizer) + } catch { + result = .failure(error) + } + semaphore.signal() + } - if tokenizerName.contains("/") { - let localDirectory = try await ensureTokenizerAssets(modelID: tokenizerName) - return try loadTokenizerFromDirectory( - localDirectory, - fallbackTokenizerName: tokenizerName - ) + semaphore.wait() + return try result!.get() } - return try await AutoTokenizer.from(pretrained: tokenizerName) - } - - private static func loadTokenizerFromDirectory( - _ directoryURL: URL, - fallbackTokenizerName: String? - ) throws -> any Tokenizer { - let tokenizerDataURL = directoryURL.appending(path: "tokenizer.json") - let tokenizerConfigURL = directoryURL.appending(path: "tokenizer_config.json") - - guard FileManager.default.fileExists(atPath: tokenizerDataURL.path), - FileManager.default.fileExists(atPath: tokenizerConfigURL.path) - else { - if let fallbackTokenizerName { - return try blockingPretrainedTokenizer(named: fallbackTokenizerName) + private static func loadTokenizerAsync( + tokenizerName: String, + tokenizerFolderURL: URL? + ) async throws -> any Tokenizer { + if let tokenizerFolderURL { + return try loadTokenizerFromDirectory( + tokenizerFolderURL, + fallbackTokenizerName: tokenizerName + ) } - throw TokenizerError.missingConfig - } - let preparedDirectory = try prepareTokenizerDirectory(directoryURL) - return try blockingLocalTokenizer(from: preparedDirectory) - } + if tokenizerName.contains("/") { + let localDirectory = try await ensureTokenizerAssets(modelID: tokenizerName) + return try loadTokenizerFromDirectory( + localDirectory, + fallbackTokenizerName: tokenizerName + ) + } - static func patchTokenizerConfigDataIfNeeded( - tokenizerConfigData: Data, - tokenizerData: Data - ) throws -> Data? { - guard - let tokenizerConfig = try JSONSerialization.jsonObject(with: tokenizerConfigData) - as? [String: Any], - let tokenizerDataObject = try JSONSerialization.jsonObject(with: tokenizerData) - as? [String: Any] - else { - return nil + return try await AutoTokenizer.from(pretrained: tokenizerName) } - let modelType = - ((tokenizerDataObject["model"] as? [String: Any])?["type"] as? String)? - .lowercased() - let tokenizerClass = tokenizerConfig["tokenizer_class"] as? String + private static func loadTokenizerFromDirectory( + _ directoryURL: URL, + fallbackTokenizerName: String? + ) throws -> any Tokenizer { + let tokenizerDataURL = directoryURL.appending(path: "tokenizer.json") + let tokenizerConfigURL = directoryURL.appending(path: "tokenizer_config.json") - guard modelType == "unigram" else { - return nil - } + guard FileManager.default.fileExists(atPath: tokenizerDataURL.path), + FileManager.default.fileExists(atPath: tokenizerConfigURL.path) + else { + if let fallbackTokenizerName { + return try blockingPretrainedTokenizer(named: fallbackTokenizerName) + } + throw TokenizerError.missingConfig + } - let shouldForceUnigram = - tokenizerClass == nil - || tokenizerClass == "RobertaTokenizer" - || tokenizerClass == "RobertaTokenizerFast" - || tokenizerClass == "XLMRobertaTokenizer" - || tokenizerClass == "XLMRobertaTokenizerFast" - || tokenizerClass == "DebertaV2Tokenizer" - || tokenizerClass == "DebertaV2TokenizerFast" - || tokenizerClass == "PreTrainedTokenizer" + let preparedDirectory = try prepareTokenizerDirectory(directoryURL) + return try blockingLocalTokenizer(from: preparedDirectory) + } - let hasListShapedExtraSpecialTokens = tokenizerConfig["extra_special_tokens"] is [Any] + static func patchTokenizerConfigDataIfNeeded( + tokenizerConfigData: Data, + tokenizerData: Data + ) throws -> Data? { + guard + let tokenizerConfig = try JSONSerialization.jsonObject(with: tokenizerConfigData) + as? [String: Any], + let tokenizerDataObject = try JSONSerialization.jsonObject(with: tokenizerData) + as? [String: Any] + else { + return nil + } - guard shouldForceUnigram || hasListShapedExtraSpecialTokens else { - return nil - } + let modelType = + ((tokenizerDataObject["model"] as? [String: Any])?["type"] as? String)? + .lowercased() + let tokenizerClass = tokenizerConfig["tokenizer_class"] as? String - var patchedConfig = tokenizerConfig - if shouldForceUnigram { - patchedConfig["tokenizer_class"] = "T5Tokenizer" - } - if let extraSpecialTokens = tokenizerConfig["extra_special_tokens"] as? [Any] { - patchedConfig["extra_special_tokens"] = nil - if patchedConfig["additional_special_tokens"] == nil { - patchedConfig["additional_special_tokens"] = extraSpecialTokens + guard modelType == "unigram" else { + return nil } - } - return try JSONSerialization.data( - withJSONObject: patchedConfig, - options: [.prettyPrinted, .sortedKeys] - ) - } - static func prepareTokenizerDirectory(_ directoryURL: URL) throws -> URL { - let tokenizerDataURL = directoryURL.appending(path: "tokenizer.json") - let tokenizerConfigURL = directoryURL.appending(path: "tokenizer_config.json") - let modelConfigURL = directoryURL.appending(path: "config.json") - - let tokenizerData = try Data(contentsOf: tokenizerDataURL) - let tokenizerConfigData = try Data(contentsOf: tokenizerConfigURL) - let patchedTokenizerConfigData = - try patchTokenizerConfigDataIfNeeded( - tokenizerConfigData: tokenizerConfigData, - tokenizerData: tokenizerData - ) ?? tokenizerConfigData - - if FileManager.default.fileExists(atPath: modelConfigURL.path), - patchedTokenizerConfigData == tokenizerConfigData - { - return directoryURL - } + let shouldForceUnigram = + tokenizerClass == nil + || tokenizerClass == "RobertaTokenizer" + || tokenizerClass == "RobertaTokenizerFast" + || tokenizerClass == "XLMRobertaTokenizer" + || tokenizerClass == "XLMRobertaTokenizerFast" + || tokenizerClass == "DebertaV2Tokenizer" + || tokenizerClass == "DebertaV2TokenizerFast" + || tokenizerClass == "PreTrainedTokenizer" - let preparedDirectory = try preparedTokenizerCacheDirectory(for: directoryURL) - let fileManager = FileManager.default + let hasListShapedExtraSpecialTokens = tokenizerConfig["extra_special_tokens"] is [Any] - if fileManager.fileExists(atPath: preparedDirectory.path) { - try fileManager.removeItem(at: preparedDirectory) - } - try fileManager.createDirectory( - at: preparedDirectory, - withIntermediateDirectories: true - ) - - for fileName in tokenizerAssetFileNames { - let sourceURL = directoryURL.appending(path: fileName) - let destinationURL = preparedDirectory.appending(path: fileName) - guard fileManager.fileExists(atPath: sourceURL.path) else { - continue + guard shouldForceUnigram || hasListShapedExtraSpecialTokens else { + return nil } - if fileName == "tokenizer_config.json" { - continue + + var patchedConfig = tokenizerConfig + if shouldForceUnigram { + patchedConfig["tokenizer_class"] = "T5Tokenizer" } - let fileData = try Data(contentsOf: sourceURL) - try fileData.write(to: destinationURL, options: .atomic) + if let extraSpecialTokens = tokenizerConfig["extra_special_tokens"] as? [Any] { + patchedConfig["extra_special_tokens"] = nil + if patchedConfig["additional_special_tokens"] == nil { + patchedConfig["additional_special_tokens"] = extraSpecialTokens + } + } + return try JSONSerialization.data( + withJSONObject: patchedConfig, + options: [.prettyPrinted, .sortedKeys] + ) } - let preparedModelConfigURL = preparedDirectory.appending(path: "config.json") - if fileManager.fileExists(atPath: modelConfigURL.path) { - let modelConfigData = try Data(contentsOf: modelConfigURL) - try modelConfigData.write(to: preparedModelConfigURL, options: .atomic) - } else { - try Data("{}".utf8).write(to: preparedModelConfigURL, options: .atomic) - } + static func prepareTokenizerDirectory(_ directoryURL: URL) throws -> URL { + let tokenizerDataURL = directoryURL.appending(path: "tokenizer.json") + let tokenizerConfigURL = directoryURL.appending(path: "tokenizer_config.json") + let modelConfigURL = directoryURL.appending(path: "config.json") + + let tokenizerData = try Data(contentsOf: tokenizerDataURL) + let tokenizerConfigData = try Data(contentsOf: tokenizerConfigURL) + let patchedTokenizerConfigData = + try patchTokenizerConfigDataIfNeeded( + tokenizerConfigData: tokenizerConfigData, + tokenizerData: tokenizerData + ) ?? tokenizerConfigData + + if FileManager.default.fileExists(atPath: modelConfigURL.path), + patchedTokenizerConfigData == tokenizerConfigData + { + return directoryURL + } - try patchedTokenizerConfigData.write( - to: preparedDirectory.appending(path: "tokenizer_config.json"), - options: .atomic - ) - return preparedDirectory - } + let preparedDirectory = try preparedTokenizerCacheDirectory(for: directoryURL) + let fileManager = FileManager.default - private static func preparedTokenizerCacheDirectory(for directoryURL: URL) throws -> URL { - let base = - try FileManager.default.url( - for: .cachesDirectory, - in: .userDomainMask, - appropriateFor: nil, - create: true + if fileManager.fileExists(atPath: preparedDirectory.path) { + try fileManager.removeItem(at: preparedDirectory) + } + try fileManager.createDirectory( + at: preparedDirectory, + withIntermediateDirectories: true ) - let leafName = sanitizedCacheComponent(directoryURL.lastPathComponent) - let digest = stableDigest(for: directoryURL.path) - return - base - .appending(path: "OpenMed", directoryHint: .isDirectory) - .appending(path: "PreparedTokenizerAssets", directoryHint: .isDirectory) - .appending(path: "\(leafName)-\(digest)", directoryHint: .isDirectory) - } - private static func blockingLocalTokenizer(from modelFolder: URL) throws -> any Tokenizer { - let semaphore = DispatchSemaphore(value: 0) - var result: Result? + for fileName in tokenizerAssetFileNames { + let sourceURL = directoryURL.appending(path: fileName) + let destinationURL = preparedDirectory.appending(path: fileName) + guard fileManager.fileExists(atPath: sourceURL.path) else { + continue + } + if fileName == "tokenizer_config.json" { + continue + } + let fileData = try Data(contentsOf: sourceURL) + try fileData.write(to: destinationURL, options: .atomic) + } - Task.detached { - do { - result = .success(try await AutoTokenizer.from(modelFolder: modelFolder)) - } catch { - result = .failure(error) + let preparedModelConfigURL = preparedDirectory.appending(path: "config.json") + if fileManager.fileExists(atPath: modelConfigURL.path) { + let modelConfigData = try Data(contentsOf: modelConfigURL) + try modelConfigData.write(to: preparedModelConfigURL, options: .atomic) + } else { + try Data("{}".utf8).write(to: preparedModelConfigURL, options: .atomic) } - semaphore.signal() + + try patchedTokenizerConfigData.write( + to: preparedDirectory.appending(path: "tokenizer_config.json"), + options: .atomic + ) + return preparedDirectory } - semaphore.wait() - return try result!.get() - } + private static func preparedTokenizerCacheDirectory(for directoryURL: URL) throws -> URL { + let base = + try FileManager.default.url( + for: .cachesDirectory, + in: .userDomainMask, + appropriateFor: nil, + create: true + ) + let leafName = sanitizedCacheComponent(directoryURL.lastPathComponent) + let digest = stableDigest(for: directoryURL.path) + return + base + .appending(path: "OpenMed", directoryHint: .isDirectory) + .appending(path: "PreparedTokenizerAssets", directoryHint: .isDirectory) + .appending(path: "\(leafName)-\(digest)", directoryHint: .isDirectory) + } - private static func ensureTokenizerAssets(modelID: String) async throws -> URL { - let directory = try tokenizerCacheDirectory(modelID: modelID) - try FileManager.default.createDirectory( - at: directory, - withIntermediateDirectories: true - ) + private static func blockingLocalTokenizer(from modelFolder: URL) throws -> any Tokenizer { + let semaphore = DispatchSemaphore(value: 0) + var result: Result? - let requiredFiles = [ - "tokenizer.json", - "tokenizer_config.json", - ] - let optionalFiles = tokenizerAssetFileNames.filter { fileName in - !requiredFiles.contains(fileName) - } + Task.detached { + do { + result = .success(try await AutoTokenizer.from(modelFolder: modelFolder)) + } catch { + result = .failure(error) + } + semaphore.signal() + } - for fileName in requiredFiles { - try await downloadTokenizerFile( - modelID: modelID, - relativePath: fileName, - destinationURL: directory.appending(path: fileName), - required: true - ) + semaphore.wait() + return try result!.get() } - for fileName in optionalFiles { - try await downloadTokenizerFile( - modelID: modelID, - relativePath: fileName, - destinationURL: directory.appending(path: fileName), - required: false + private static func ensureTokenizerAssets(modelID: String) async throws -> URL { + let directory = try tokenizerCacheDirectory(modelID: modelID) + try FileManager.default.createDirectory( + at: directory, + withIntermediateDirectories: true ) - } - return directory - } + let requiredFiles = [ + "tokenizer.json", + "tokenizer_config.json", + ] + let optionalFiles = tokenizerAssetFileNames.filter { fileName in + !requiredFiles.contains(fileName) + } - private static func tokenizerCacheDirectory(modelID: String) throws -> URL { - let base = - try FileManager.default.url( - for: .cachesDirectory, - in: .userDomainMask, - appropriateFor: nil, - create: true - ) - let sanitized = modelID.replacingOccurrences(of: "/", with: "__") - return - base - .appending(path: "OpenMed", directoryHint: .isDirectory) - .appending(path: "TokenizerAssets", directoryHint: .isDirectory) - .appending(path: sanitized, directoryHint: .isDirectory) - } + for fileName in requiredFiles { + try await downloadTokenizerFile( + modelID: modelID, + relativePath: fileName, + destinationURL: directory.appending(path: fileName), + required: true + ) + } - private static let tokenizerAssetFileNames = [ - "tokenizer.json", - "tokenizer_config.json", - "special_tokens_map.json", - "vocab.txt", - "vocab.json", - "merges.txt", - "spm.model", - "sentencepiece.bpe.model", - "added_tokens.json", - ] - - private static func downloadTokenizerFile( - modelID: String, - relativePath: String, - destinationURL: URL, - required: Bool - ) async throws { - if FileManager.default.fileExists(atPath: destinationURL.path) { - return - } + for fileName in optionalFiles { + try await downloadTokenizerFile( + modelID: modelID, + relativePath: fileName, + destinationURL: directory.appending(path: fileName), + required: false + ) + } - let encodedModelID = - modelID - .split(separator: "/") - .map { String($0).addingPercentEncoding(withAllowedCharacters: .urlPathAllowed) ?? String($0) } - .joined(separator: "/") - let encodedPath = - relativePath - .split(separator: "/") - .map { String($0).addingPercentEncoding(withAllowedCharacters: .urlPathAllowed) ?? String($0) } - .joined(separator: "/") - - guard - let url = URL( - string: "https://huggingface.co/\(encodedModelID)/resolve/main/\(encodedPath)?download=1" - ) - else { - throw TokenizerError.missingConfig + return directory } - let (data, response) = try await URLSession.shared.data(from: url) - guard let http = response as? HTTPURLResponse else { - throw TokenizerError.missingConfig - } - if http.statusCode == 404 && !required { + private static func tokenizerCacheDirectory(modelID: String) throws -> URL { + let base = + try FileManager.default.url( + for: .cachesDirectory, + in: .userDomainMask, + appropriateFor: nil, + create: true + ) + let sanitized = modelID.replacingOccurrences(of: "/", with: "__") return - } - guard (200..<300).contains(http.statusCode) else { - throw TokenizerError.missingConfig + base + .appending(path: "OpenMed", directoryHint: .isDirectory) + .appending(path: "TokenizerAssets", directoryHint: .isDirectory) + .appending(path: sanitized, directoryHint: .isDirectory) } - try FileManager.default.createDirectory( - at: destinationURL.deletingLastPathComponent(), - withIntermediateDirectories: true - ) - try data.write(to: destinationURL, options: .atomic) - } + private static let tokenizerAssetFileNames = [ + "tokenizer.json", + "tokenizer_config.json", + "special_tokens_map.json", + "vocab.txt", + "vocab.json", + "merges.txt", + "spm.model", + "sentencepiece.bpe.model", + "added_tokens.json", + ] + + private static func downloadTokenizerFile( + modelID: String, + relativePath: String, + destinationURL: URL, + required: Bool + ) async throws { + if FileManager.default.fileExists(atPath: destinationURL.path) { + return + } - private static func blockingPretrainedTokenizer(named name: String) throws -> any Tokenizer { - let semaphore = DispatchSemaphore(value: 0) - var result: Result? + let encodedModelID = + modelID + .split(separator: "/") + .map { String($0).addingPercentEncoding(withAllowedCharacters: .urlPathAllowed) ?? String($0) } + .joined(separator: "/") + let encodedPath = + relativePath + .split(separator: "/") + .map { String($0).addingPercentEncoding(withAllowedCharacters: .urlPathAllowed) ?? String($0) } + .joined(separator: "/") - Task.detached { - do { - result = .success(try await AutoTokenizer.from(pretrained: name)) - } catch { - result = .failure(error) + guard + let url = URL( + string: "https://huggingface.co/\(encodedModelID)/resolve/main/\(encodedPath)?download=1" + ) + else { + throw TokenizerError.missingConfig } - semaphore.signal() + + let (data, response) = try await URLSession.shared.data(from: url) + guard let http = response as? HTTPURLResponse else { + throw TokenizerError.missingConfig + } + if http.statusCode == 404 && !required { + return + } + guard (200..<300).contains(http.statusCode) else { + throw TokenizerError.missingConfig + } + + try FileManager.default.createDirectory( + at: destinationURL.deletingLastPathComponent(), + withIntermediateDirectories: true + ) + try data.write(to: destinationURL, options: .atomic) } - semaphore.wait() - return try result!.get() - } + private static func blockingPretrainedTokenizer(named name: String) throws -> any Tokenizer { + let semaphore = DispatchSemaphore(value: 0) + var result: Result? - private static func sanitizedCacheComponent(_ value: String) -> String { - value - .replacingOccurrences(of: "/", with: "__") - .replacingOccurrences(of: ":", with: "_") - .replacingOccurrences(of: " ", with: "_") - } + Task.detached { + do { + result = .success(try await AutoTokenizer.from(pretrained: name)) + } catch { + result = .failure(error) + } + semaphore.signal() + } - private static func stableDigest(for value: String) -> String { - var hash: UInt64 = 0xcbf2_9ce4_8422_2325 - for byte in value.utf8 { - hash ^= UInt64(byte) - hash &*= 0x100_0000_01b3 + semaphore.wait() + return try result!.get() } - return String(format: "%016llx", hash) - } - static func buildOffsets( - tokens: [String], - in text: String - ) -> [(Int, Int)] { - var offsets: [(Int, Int)] = [] - var cursor = text.startIndex - - for token in tokens { - if isSpecialToken(token) { - offsets.append((0, 0)) - continue + private static func sanitizedCacheComponent(_ value: String) -> String { + value + .replacingOccurrences(of: "/", with: "__") + .replacingOccurrences(of: ":", with: "_") + .replacingOccurrences(of: " ", with: "_") + } + + private static func stableDigest(for value: String) -> String { + var hash: UInt64 = 0xcbf2_9ce4_8422_2325 + for byte in value.utf8 { + hash ^= UInt64(byte) + hash &*= 0x100_0000_01b3 } + return String(format: "%016llx", hash) + } - let normalized = normalize(token: token) - let piece = normalized.piece + static func buildOffsets( + tokens: [String], + in text: String + ) -> [(Int, Int)] { + var offsets: [(Int, Int)] = [] + var cursor = text.startIndex + + for token in tokens { + if isSpecialToken(token) { + offsets.append((0, 0)) + continue + } - if piece.isEmpty { - offsets.append((0, 0)) - continue - } + let normalized = normalize(token: token) + let piece = normalized.piece - var searchStart = cursor - if normalized.skipLeadingWhitespace { - while searchStart < text.endIndex && text[searchStart].isWhitespace { - searchStart = text.index(after: searchStart) + if piece.isEmpty { + offsets.append((0, 0)) + continue } - } - let searchSlice = text[searchStart...] - let exactRange = searchSlice.range(of: piece) - let insensitiveRange = searchSlice.range( - of: piece, - options: [.caseInsensitive, .diacriticInsensitive] - ) + var searchStart = cursor + if normalized.skipLeadingWhitespace { + while searchStart < text.endIndex && text[searchStart].isWhitespace { + searchStart = text.index(after: searchStart) + } + } - let range: Range? - switch (exactRange, insensitiveRange) { - case (let exact?, let insensitive?): - if exact.lowerBound <= insensitive.lowerBound { + let searchSlice = text[searchStart...] + let exactRange = searchSlice.range(of: piece) + let insensitiveRange = searchSlice.range( + of: piece, + options: [.caseInsensitive, .diacriticInsensitive] + ) + + let range: Range? + switch (exactRange, insensitiveRange) { + case (let exact?, let insensitive?): + if exact.lowerBound <= insensitive.lowerBound { + range = exact + } else { + range = insensitive + } + case (let exact?, nil): range = exact - } else { + case (nil, let insensitive?): range = insensitive + case (nil, nil): + range = nil + } + + if let range { + let start = text.distance(from: text.startIndex, to: range.lowerBound) + let end = text.distance(from: text.startIndex, to: range.upperBound) + offsets.append((start, end)) + cursor = range.upperBound + continue } - case (let exact?, nil): - range = exact - case (nil, let insensitive?): - range = insensitive - case (nil, nil): - range = nil - } - if let range { - let start = text.distance(from: text.startIndex, to: range.lowerBound) - let end = text.distance(from: text.startIndex, to: range.upperBound) + let start = text.distance(from: text.startIndex, to: searchStart) + let endIndex = + text.index( + searchStart, + offsetBy: piece.count, + limitedBy: text.endIndex + ) ?? text.endIndex + let end = text.distance(from: text.startIndex, to: endIndex) offsets.append((start, end)) - cursor = range.upperBound - continue + cursor = endIndex } - let start = text.distance(from: text.startIndex, to: searchStart) - let endIndex = - text.index( - searchStart, - offsetBy: piece.count, - limitedBy: text.endIndex - ) ?? text.endIndex - let end = text.distance(from: text.startIndex, to: endIndex) - offsets.append((start, end)) - cursor = endIndex + return offsets } - return offsets - } - - private static func normalize(token: String) -> (piece: String, skipLeadingWhitespace: Bool) { - if token == "Ċ" { - return ("\n", false) - } - if token.hasPrefix("##") { - return (String(token.dropFirst(2)), false) - } - if token.hasPrefix("▁") || token.hasPrefix("Ġ") { - return (String(token.dropFirst()), true) + private static func normalize(token: String) -> (piece: String, skipLeadingWhitespace: Bool) { + if token == "Ċ" { + return ("\n", false) + } + if token.hasPrefix("##") { + return (String(token.dropFirst(2)), false) + } + if token.hasPrefix("▁") || token.hasPrefix("Ġ") { + return (String(token.dropFirst()), true) + } + return (token, false) } - return (token, false) - } - private static func isSpecialToken(_ token: String) -> Bool { - switch token { - case "[CLS]", "[SEP]", "[PAD]", "[MASK]", "", "", "", "": - return true - default: - return false + private static func isSpecialToken(_ token: String) -> Bool { + switch token { + case "[CLS]", "[SEP]", "[PAD]", "[MASK]", "", "", "", "": + return true + default: + return false + } } } -} +#endif diff --git a/swift/OpenMedKit/Sources/OpenMedKit/OpenMedMLXDeberta.swift b/swift/OpenMedKit/Sources/OpenMedKit/OpenMedMLXDeberta.swift index 6dfb9e622..5415aca3a 100644 --- a/swift/OpenMedKit/Sources/OpenMedKit/OpenMedMLXDeberta.swift +++ b/swift/OpenMedKit/Sources/OpenMedKit/OpenMedMLXDeberta.swift @@ -1,630 +1,632 @@ -import Foundation -import MLX -import MLXNN - -private func openMedFloatScalar(_ value: Float, like array: MLXArray) -> MLXArray { - MLXArray(value).asType(array.dtype) -} - -private func openMedLogBucketPosition( - _ relativePosition: MLXArray, - bucketSize: Int, - maxPosition: Int -) -> MLXArray { - let signValue = sign(relativePosition) - let mid = bucketSize / 2 - let absPosition = `where`( - (relativePosition .< mid) .&& (relativePosition .> -mid), - MLXArray(mid - 1).asType(relativePosition.dtype), - abs(relativePosition) - ) - - let absPositionFloat = absPosition.asType(.float32) - let midFloat = MLXArray(Float(mid)) - let logBase = log(MLXArray(Float(maxPosition - 1) / Float(mid))) - let logPosition = ceil(log(absPositionFloat / midFloat) / logBase * Float(mid - 1)) + Float(mid) - let bucketPosition = `where`( - absPosition .<= mid, - relativePosition.asType(logPosition.dtype), - logPosition * signValue.asType(logPosition.dtype) - ) - return bucketPosition.asType(relativePosition.dtype) -} - -private func openMedBuildRelativePosition( - queryLayer: MLXArray, - keyLayer: MLXArray, - bucketSize: Int, - maxPosition: Int -) -> MLXArray { - let querySize = queryLayer.dim(-2) - let keySize = keyLayer.dim(-2) - let queryIDs = MLXArray.arange(querySize, dtype: .int32) - let keyIDs = MLXArray.arange(keySize, dtype: .int32) - var relativePosition = queryIDs.expandedDimensions(axis: 1) - keyIDs.expandedDimensions(axis: 0) - - if bucketSize > 0 && maxPosition > 0 { - relativePosition = openMedLogBucketPosition( - relativePosition, - bucketSize: bucketSize, - maxPosition: maxPosition - ) +#if canImport(MLX) && !os(watchOS) && !os(visionOS) + import Foundation + import MLX + import MLXNN + + private func openMedFloatScalar(_ value: Float, like array: MLXArray) -> MLXArray { + MLXArray(value).asType(array.dtype) } - return relativePosition.asType(.int32).expandedDimensions(axis: 0) -} - -private func openMedBuildRPosition( - queryLayer: MLXArray, - keyLayer: MLXArray, - relativePosition: MLXArray, - positionBuckets: Int, - maxRelativePositions: Int -) -> MLXArray { - if keyLayer.dim(-2) != queryLayer.dim(-2) { - return openMedBuildRelativePosition( - queryLayer: keyLayer, - keyLayer: keyLayer, - bucketSize: positionBuckets, - maxPosition: maxRelativePositions + private func openMedLogBucketPosition( + _ relativePosition: MLXArray, + bucketSize: Int, + maxPosition: Int + ) -> MLXArray { + let signValue = sign(relativePosition) + let mid = bucketSize / 2 + let absPosition = `where`( + (relativePosition .< mid) .&& (relativePosition .> -mid), + MLXArray(mid - 1).asType(relativePosition.dtype), + abs(relativePosition) ) - } - return relativePosition -} - -private func openMedRepeatBatches(_ array: MLXArray, batchSize: Int) -> MLXArray { - let targetShape = [batchSize, array.dim(0), array.dim(1), array.dim(2)] - return broadcast(array.expandedDimensions(axis: 0), to: targetShape) - .reshaped(batchSize * array.dim(0), array.dim(1), array.dim(2)) -} - -final class OpenMedDebertaV2Embeddings: Module { - private let embeddingSize: Int - private let hiddenSize: Int - private let positionBiasedInput: Bool - - @ModuleInfo(key: "word_embeddings") var wordEmbeddings: Embedding - @ModuleInfo(key: "position_embeddings") var positionEmbeddings: Embedding? - @ModuleInfo(key: "token_type_embeddings") var tokenTypeEmbeddings: Embedding? - @ModuleInfo(key: "embed_proj") var embedProjection: Linear? - @ModuleInfo(key: "LayerNorm") var layerNorm: LayerNorm - - init(_ configuration: OpenMedMLXBertConfiguration) { - embeddingSize = configuration.embeddingSize - hiddenSize = configuration.encoderHiddenSize - positionBiasedInput = configuration.positionBiasedInput - - _wordEmbeddings.wrappedValue = Embedding( - embeddingCount: configuration.vocabularySize, - dimensions: embeddingSize + + let absPositionFloat = absPosition.asType(.float32) + let midFloat = MLXArray(Float(mid)) + let logBase = log(MLXArray(Float(maxPosition - 1) / Float(mid))) + let logPosition = ceil(log(absPositionFloat / midFloat) / logBase * Float(mid - 1)) + Float(mid) + let bucketPosition = `where`( + absPosition .<= mid, + relativePosition.asType(logPosition.dtype), + logPosition * signValue.asType(logPosition.dtype) ) - if positionBiasedInput { - _positionEmbeddings.wrappedValue = Embedding( - embeddingCount: configuration.maxPositionEmbeddings, - dimensions: embeddingSize - ) - } - if configuration.typeVocabularySize > 0 { - _tokenTypeEmbeddings.wrappedValue = Embedding( - embeddingCount: configuration.typeVocabularySize, - dimensions: embeddingSize + return bucketPosition.asType(relativePosition.dtype) + } + + private func openMedBuildRelativePosition( + queryLayer: MLXArray, + keyLayer: MLXArray, + bucketSize: Int, + maxPosition: Int + ) -> MLXArray { + let querySize = queryLayer.dim(-2) + let keySize = keyLayer.dim(-2) + let queryIDs = MLXArray.arange(querySize, dtype: .int32) + let keyIDs = MLXArray.arange(keySize, dtype: .int32) + var relativePosition = queryIDs.expandedDimensions(axis: 1) - keyIDs.expandedDimensions(axis: 0) + + if bucketSize > 0 && maxPosition > 0 { + relativePosition = openMedLogBucketPosition( + relativePosition, + bucketSize: bucketSize, + maxPosition: maxPosition ) } - if embeddingSize != hiddenSize { - _embedProjection.wrappedValue = Linear(embeddingSize, hiddenSize, bias: false) - } - _layerNorm.wrappedValue = LayerNorm( - dimensions: hiddenSize, - eps: configuration.layerNormEps - ) + + return relativePosition.asType(.int32).expandedDimensions(axis: 0) } - func callAsFunction( - inputIDs: MLXArray, - tokenTypeIDs: MLXArray? = nil, - attentionMask: MLXArray? = nil + private func openMedBuildRPosition( + queryLayer: MLXArray, + keyLayer: MLXArray, + relativePosition: MLXArray, + positionBuckets: Int, + maxRelativePositions: Int ) -> MLXArray { - let seqLen = inputIDs.dim(1) - var embeddings = wordEmbeddings(inputIDs) - - if let positionEmbeddings { - let positionIDs = MLXArray.arange(seqLen, dtype: inputIDs.dtype) - .expandedDimensions(axis: 0) - embeddings = embeddings + positionEmbeddings(positionIDs) + if keyLayer.dim(-2) != queryLayer.dim(-2) { + return openMedBuildRelativePosition( + queryLayer: keyLayer, + keyLayer: keyLayer, + bucketSize: positionBuckets, + maxPosition: maxRelativePositions + ) } + return relativePosition + } - if let tokenTypeEmbeddings { - let tokenTypeIDs = tokenTypeIDs ?? MLXArray.zeros(like: inputIDs) - embeddings = embeddings + tokenTypeEmbeddings(tokenTypeIDs) - } + private func openMedRepeatBatches(_ array: MLXArray, batchSize: Int) -> MLXArray { + let targetShape = [batchSize, array.dim(0), array.dim(1), array.dim(2)] + return broadcast(array.expandedDimensions(axis: 0), to: targetShape) + .reshaped(batchSize * array.dim(0), array.dim(1), array.dim(2)) + } - if let embedProjection { - embeddings = embedProjection(embeddings) - } + final class OpenMedDebertaV2Embeddings: Module { + private let embeddingSize: Int + private let hiddenSize: Int + private let positionBiasedInput: Bool + + @ModuleInfo(key: "word_embeddings") var wordEmbeddings: Embedding + @ModuleInfo(key: "position_embeddings") var positionEmbeddings: Embedding? + @ModuleInfo(key: "token_type_embeddings") var tokenTypeEmbeddings: Embedding? + @ModuleInfo(key: "embed_proj") var embedProjection: Linear? + @ModuleInfo(key: "LayerNorm") var layerNorm: LayerNorm - embeddings = layerNorm(embeddings) + init(_ configuration: OpenMedMLXBertConfiguration) { + embeddingSize = configuration.embeddingSize + hiddenSize = configuration.encoderHiddenSize + positionBiasedInput = configuration.positionBiasedInput - if let attentionMask { - var mask = attentionMask - if mask.ndim == 4 { - mask = mask[0..., 0, 0, 0...] + _wordEmbeddings.wrappedValue = Embedding( + embeddingCount: configuration.vocabularySize, + dimensions: embeddingSize + ) + if positionBiasedInput { + _positionEmbeddings.wrappedValue = Embedding( + embeddingCount: configuration.maxPositionEmbeddings, + dimensions: embeddingSize + ) } - if mask.ndim == 2 { - mask = mask.expandedDimensions(axis: -1) + if configuration.typeVocabularySize > 0 { + _tokenTypeEmbeddings.wrappedValue = Embedding( + embeddingCount: configuration.typeVocabularySize, + dimensions: embeddingSize + ) } - embeddings = embeddings * mask.asType(embeddings.dtype) + if embeddingSize != hiddenSize { + _embedProjection.wrappedValue = Linear(embeddingSize, hiddenSize, bias: false) + } + _layerNorm.wrappedValue = LayerNorm( + dimensions: hiddenSize, + eps: configuration.layerNormEps + ) } - return embeddings - } -} - -final class OpenMedDisentangledSelfAttention: Module { - private let numAttentionHeads: Int - private let attentionHeadSize: Int - private let allHeadSize: Int - private let shareAttentionKey: Bool - private let positionAttentionTypes: Set - private let relativeAttention: Bool - private let positionBuckets: Int - private let maxRelativePositions: Int - private let positionEmbeddingSize: Int - - @ModuleInfo(key: "query_proj") var queryProjection: Linear - @ModuleInfo(key: "key_proj") var keyProjection: Linear - @ModuleInfo(key: "value_proj") var valueProjection: Linear - @ModuleInfo(key: "pos_key_proj") var positionKeyProjection: Linear? - @ModuleInfo(key: "pos_query_proj") var positionQueryProjection: Linear? - - init(_ configuration: OpenMedMLXBertConfiguration) { - numAttentionHeads = configuration.numAttentionHeads - attentionHeadSize = configuration.encoderHiddenSize / configuration.numAttentionHeads - allHeadSize = numAttentionHeads * attentionHeadSize - shareAttentionKey = configuration.shareAttentionKey - positionAttentionTypes = Set(configuration.positionAttentionTypes) - relativeAttention = configuration.relativeAttention - positionBuckets = configuration.positionBuckets - var resolvedMaxRelativePositions = configuration.maxRelativePositions - if resolvedMaxRelativePositions < 1 { - resolvedMaxRelativePositions = configuration.maxPositionEmbeddings - } - maxRelativePositions = resolvedMaxRelativePositions - positionEmbeddingSize = - positionBuckets > 0 ? positionBuckets : resolvedMaxRelativePositions + func callAsFunction( + inputIDs: MLXArray, + tokenTypeIDs: MLXArray? = nil, + attentionMask: MLXArray? = nil + ) -> MLXArray { + let seqLen = inputIDs.dim(1) + var embeddings = wordEmbeddings(inputIDs) + + if let positionEmbeddings { + let positionIDs = MLXArray.arange(seqLen, dtype: inputIDs.dtype) + .expandedDimensions(axis: 0) + embeddings = embeddings + positionEmbeddings(positionIDs) + } - _queryProjection.wrappedValue = Linear(configuration.encoderHiddenSize, allHeadSize) - _keyProjection.wrappedValue = Linear(configuration.encoderHiddenSize, allHeadSize) - _valueProjection.wrappedValue = Linear(configuration.encoderHiddenSize, allHeadSize) + if let tokenTypeEmbeddings { + let tokenTypeIDs = tokenTypeIDs ?? MLXArray.zeros(like: inputIDs) + embeddings = embeddings + tokenTypeEmbeddings(tokenTypeIDs) + } - if relativeAttention && !shareAttentionKey { - if positionAttentionTypes.contains("c2p") { - _positionKeyProjection.wrappedValue = Linear( - configuration.encoderHiddenSize, - allHeadSize - ) + if let embedProjection { + embeddings = embedProjection(embeddings) } - if positionAttentionTypes.contains("p2c") { - _positionQueryProjection.wrappedValue = Linear( - configuration.encoderHiddenSize, - allHeadSize - ) + + embeddings = layerNorm(embeddings) + + if let attentionMask { + var mask = attentionMask + if mask.ndim == 4 { + mask = mask[0..., 0, 0, 0...] + } + if mask.ndim == 2 { + mask = mask.expandedDimensions(axis: -1) + } + embeddings = embeddings * mask.asType(embeddings.dtype) } - } - } - private func transposeForScores(_ array: MLXArray) -> MLXArray { - let batchSize = array.dim(0) - let seqLen = array.dim(1) - return - array - .reshaped(batchSize, seqLen, numAttentionHeads, attentionHeadSize) - .transposed(0, 2, 1, 3) - .reshaped(batchSize * numAttentionHeads, seqLen, attentionHeadSize) + return embeddings + } } - private func disentangledAttentionBias( - queryLayer: MLXArray, - keyLayer: MLXArray, - relativePosition: MLXArray?, - relativeEmbeddings: MLXArray, - scaleFactor: Int - ) -> MLXArray { - var relativePosition = - relativePosition - ?? openMedBuildRelativePosition( - queryLayer: queryLayer, - keyLayer: keyLayer, - bucketSize: positionBuckets, - maxPosition: maxRelativePositions - ) + final class OpenMedDisentangledSelfAttention: Module { + private let numAttentionHeads: Int + private let attentionHeadSize: Int + private let allHeadSize: Int + private let shareAttentionKey: Bool + private let positionAttentionTypes: Set + private let relativeAttention: Bool + private let positionBuckets: Int + private let maxRelativePositions: Int + private let positionEmbeddingSize: Int + + @ModuleInfo(key: "query_proj") var queryProjection: Linear + @ModuleInfo(key: "key_proj") var keyProjection: Linear + @ModuleInfo(key: "value_proj") var valueProjection: Linear + @ModuleInfo(key: "pos_key_proj") var positionKeyProjection: Linear? + @ModuleInfo(key: "pos_query_proj") var positionQueryProjection: Linear? + + init(_ configuration: OpenMedMLXBertConfiguration) { + numAttentionHeads = configuration.numAttentionHeads + attentionHeadSize = configuration.encoderHiddenSize / configuration.numAttentionHeads + allHeadSize = numAttentionHeads * attentionHeadSize + shareAttentionKey = configuration.shareAttentionKey + positionAttentionTypes = Set(configuration.positionAttentionTypes) + relativeAttention = configuration.relativeAttention + positionBuckets = configuration.positionBuckets + var resolvedMaxRelativePositions = configuration.maxRelativePositions + if resolvedMaxRelativePositions < 1 { + resolvedMaxRelativePositions = configuration.maxPositionEmbeddings + } + maxRelativePositions = resolvedMaxRelativePositions + positionEmbeddingSize = + positionBuckets > 0 ? positionBuckets : resolvedMaxRelativePositions + + _queryProjection.wrappedValue = Linear(configuration.encoderHiddenSize, allHeadSize) + _keyProjection.wrappedValue = Linear(configuration.encoderHiddenSize, allHeadSize) + _valueProjection.wrappedValue = Linear(configuration.encoderHiddenSize, allHeadSize) + + if relativeAttention && !shareAttentionKey { + if positionAttentionTypes.contains("c2p") { + _positionKeyProjection.wrappedValue = Linear( + configuration.encoderHiddenSize, + allHeadSize + ) + } + if positionAttentionTypes.contains("p2c") { + _positionQueryProjection.wrappedValue = Linear( + configuration.encoderHiddenSize, + allHeadSize + ) + } + } + } - if relativePosition.ndim == 2 { - relativePosition = relativePosition.expandedDimensions(axes: [0, 1]) - } else if relativePosition.ndim == 3 { - relativePosition = relativePosition.expandedDimensions(axis: 1) + private func transposeForScores(_ array: MLXArray) -> MLXArray { + let batchSize = array.dim(0) + let seqLen = array.dim(1) + return + array + .reshaped(batchSize, seqLen, numAttentionHeads, attentionHeadSize) + .transposed(0, 2, 1, 3) + .reshaped(batchSize * numAttentionHeads, seqLen, attentionHeadSize) } - relativePosition = relativePosition.asType(.int32) - let attentionSpan = positionEmbeddingSize - let relEmbeddings = relativeEmbeddings[0..<(attentionSpan * 2), 0...] - .expandedDimensions(axis: 0) - - let batchSize = queryLayer.dim(0) / numAttentionHeads - let positionQueryLayer: MLXArray? - let positionKeyLayer: MLXArray? - if shareAttentionKey { - positionQueryLayer = openMedRepeatBatches( - transposeForScores(queryProjection(relEmbeddings)), - batchSize: batchSize - ) - positionKeyLayer = openMedRepeatBatches( - transposeForScores(keyProjection(relEmbeddings)), - batchSize: batchSize - ) - } else { - if positionAttentionTypes.contains("p2c"), let positionQueryProjection { + private func disentangledAttentionBias( + queryLayer: MLXArray, + keyLayer: MLXArray, + relativePosition: MLXArray?, + relativeEmbeddings: MLXArray, + scaleFactor: Int + ) -> MLXArray { + var relativePosition = + relativePosition + ?? openMedBuildRelativePosition( + queryLayer: queryLayer, + keyLayer: keyLayer, + bucketSize: positionBuckets, + maxPosition: maxRelativePositions + ) + + if relativePosition.ndim == 2 { + relativePosition = relativePosition.expandedDimensions(axes: [0, 1]) + } else if relativePosition.ndim == 3 { + relativePosition = relativePosition.expandedDimensions(axis: 1) + } + + relativePosition = relativePosition.asType(.int32) + let attentionSpan = positionEmbeddingSize + let relEmbeddings = relativeEmbeddings[0..<(attentionSpan * 2), 0...] + .expandedDimensions(axis: 0) + + let batchSize = queryLayer.dim(0) / numAttentionHeads + let positionQueryLayer: MLXArray? + let positionKeyLayer: MLXArray? + if shareAttentionKey { positionQueryLayer = openMedRepeatBatches( - transposeForScores(positionQueryProjection(relEmbeddings)), + transposeForScores(queryProjection(relEmbeddings)), batchSize: batchSize ) - } else { - positionQueryLayer = nil - } - if positionAttentionTypes.contains("c2p"), let positionKeyProjection { positionKeyLayer = openMedRepeatBatches( - transposeForScores(positionKeyProjection(relEmbeddings)), + transposeForScores(keyProjection(relEmbeddings)), batchSize: batchSize ) } else { - positionKeyLayer = nil + if positionAttentionTypes.contains("p2c"), let positionQueryProjection { + positionQueryLayer = openMedRepeatBatches( + transposeForScores(positionQueryProjection(relEmbeddings)), + batchSize: batchSize + ) + } else { + positionQueryLayer = nil + } + if positionAttentionTypes.contains("c2p"), let positionKeyProjection { + positionKeyLayer = openMedRepeatBatches( + transposeForScores(positionKeyProjection(relEmbeddings)), + batchSize: batchSize + ) + } else { + positionKeyLayer = nil + } } - } - var score: MLXArray? - - if positionAttentionTypes.contains("c2p"), let positionKeyLayer { - let c2pAttention = queryLayer.matmul(positionKeyLayer.transposed(0, 2, 1)) - let c2pPosition = clip( - relativePosition + attentionSpan, - min: 0, - max: attentionSpan * 2 - 1 - ) - let c2pIndex = broadcast( - c2pPosition.squeezed(axis: 0), - to: [queryLayer.dim(0), queryLayer.dim(1), relativePosition.dim(-1)] - ).asType(.int32) - let c2pScale = openMedFloatScalar( - sqrt(Float(positionKeyLayer.dim(-1) * scaleFactor)), - like: c2pAttention - ) - let c2pScore = takeAlong(c2pAttention, c2pIndex, axis: -1) / c2pScale - score = c2pScore - } - - if positionAttentionTypes.contains("p2c"), let positionQueryLayer { - let rPosition = openMedBuildRPosition( - queryLayer: queryLayer, - keyLayer: keyLayer, - relativePosition: relativePosition, - positionBuckets: positionBuckets, - maxRelativePositions: maxRelativePositions - ) - let p2cPosition = clip( - -rPosition + attentionSpan, - min: 0, - max: attentionSpan * 2 - 1 - ) - let p2cAttention = keyLayer.matmul(positionQueryLayer.transposed(0, 2, 1)) - let p2cIndex = broadcast( - p2cPosition.squeezed(axis: 0), - to: [queryLayer.dim(0), keyLayer.dim(-2), keyLayer.dim(-2)] - ).asType(.int32) - let p2cScale = openMedFloatScalar( - sqrt(Float(positionQueryLayer.dim(-1) * scaleFactor)), - like: p2cAttention - ) - let p2cScore = - takeAlong(p2cAttention, p2cIndex, axis: -1) - .transposed(0, 2, 1) - / p2cScale - score = score.map { $0 + p2cScore } ?? p2cScore - } - - return score - ?? MLXArray.zeros( - [queryLayer.dim(0), queryLayer.dim(1), keyLayer.dim(1)], - type: Float.self - ) - } - - func callAsFunction( - hiddenStates: MLXArray, - attentionMask: MLXArray, - queryStates: MLXArray? = nil, - relativePosition: MLXArray? = nil, - relativeEmbeddings: MLXArray? = nil - ) -> MLXArray { - let queryStates = queryStates ?? hiddenStates - let queryLayer = transposeForScores(queryProjection(queryStates)) - let keyLayer = transposeForScores(keyProjection(hiddenStates)) - let valueLayer = transposeForScores(valueProjection(hiddenStates)) - - var scaleFactor = 1 - if positionAttentionTypes.contains("c2p") { - scaleFactor += 1 - } - if positionAttentionTypes.contains("p2c") { - scaleFactor += 1 - } + var score: MLXArray? - let scale = openMedFloatScalar( - sqrt(Float(queryLayer.dim(-1) * scaleFactor)), - like: keyLayer - ) - var attentionScores = queryLayer.matmul(keyLayer.transposed(0, 2, 1) / scale) + if positionAttentionTypes.contains("c2p"), let positionKeyLayer { + let c2pAttention = queryLayer.matmul(positionKeyLayer.transposed(0, 2, 1)) + let c2pPosition = clip( + relativePosition + attentionSpan, + min: 0, + max: attentionSpan * 2 - 1 + ) + let c2pIndex = broadcast( + c2pPosition.squeezed(axis: 0), + to: [queryLayer.dim(0), queryLayer.dim(1), relativePosition.dim(-1)] + ).asType(.int32) + let c2pScale = openMedFloatScalar( + sqrt(Float(positionKeyLayer.dim(-1) * scaleFactor)), + like: c2pAttention + ) + let c2pScore = takeAlong(c2pAttention, c2pIndex, axis: -1) / c2pScale + score = c2pScore + } - if relativeAttention, let relativeEmbeddings { - attentionScores = - attentionScores - + disentangledAttentionBias( + if positionAttentionTypes.contains("p2c"), let positionQueryLayer { + let rPosition = openMedBuildRPosition( queryLayer: queryLayer, keyLayer: keyLayer, relativePosition: relativePosition, - relativeEmbeddings: relativeEmbeddings, - scaleFactor: scaleFactor + positionBuckets: positionBuckets, + maxRelativePositions: maxRelativePositions + ) + let p2cPosition = clip( + -rPosition + attentionSpan, + min: 0, + max: attentionSpan * 2 - 1 + ) + let p2cAttention = keyLayer.matmul(positionQueryLayer.transposed(0, 2, 1)) + let p2cIndex = broadcast( + p2cPosition.squeezed(axis: 0), + to: [queryLayer.dim(0), keyLayer.dim(-2), keyLayer.dim(-2)] + ).asType(.int32) + let p2cScale = openMedFloatScalar( + sqrt(Float(positionQueryLayer.dim(-1) * scaleFactor)), + like: p2cAttention + ) + let p2cScore = + takeAlong(p2cAttention, p2cIndex, axis: -1) + .transposed(0, 2, 1) + / p2cScale + score = score.map { $0 + p2cScore } ?? p2cScore + } + + return score + ?? MLXArray.zeros( + [queryLayer.dim(0), queryLayer.dim(1), keyLayer.dim(1)], + type: Float.self ) } - let batchSize = hiddenStates.dim(0) - let queryLen = queryLayer.dim(1) - let keyLen = keyLayer.dim(1) - attentionScores = attentionScores.reshaped( - batchSize, - numAttentionHeads, - queryLen, - keyLen - ) + func callAsFunction( + hiddenStates: MLXArray, + attentionMask: MLXArray, + queryStates: MLXArray? = nil, + relativePosition: MLXArray? = nil, + relativeEmbeddings: MLXArray? = nil + ) -> MLXArray { + let queryStates = queryStates ?? hiddenStates + let queryLayer = transposeForScores(queryProjection(queryStates)) + let keyLayer = transposeForScores(keyProjection(hiddenStates)) + let valueLayer = transposeForScores(valueProjection(hiddenStates)) + + var scaleFactor = 1 + if positionAttentionTypes.contains("c2p") { + scaleFactor += 1 + } + if positionAttentionTypes.contains("p2c") { + scaleFactor += 1 + } - attentionScores = `where`( - attentionMask .> 0, - attentionScores, - openMedFloatScalar(-3.4028235e38, like: attentionScores) - ) - let attentionProbs = softmax(attentionScores, axis: -1) - let contextLayer = - attentionProbs - .reshaped(batchSize * numAttentionHeads, queryLen, keyLen) - .matmul(valueLayer) - .reshaped(batchSize, numAttentionHeads, queryLen, attentionHeadSize) - .transposed(0, 2, 1, 3) - return contextLayer.reshaped(batchSize, queryLen, allHeadSize) - } -} + let scale = openMedFloatScalar( + sqrt(Float(queryLayer.dim(-1) * scaleFactor)), + like: keyLayer + ) + var attentionScores = queryLayer.matmul(keyLayer.transposed(0, 2, 1) / scale) + + if relativeAttention, let relativeEmbeddings { + attentionScores = + attentionScores + + disentangledAttentionBias( + queryLayer: queryLayer, + keyLayer: keyLayer, + relativePosition: relativePosition, + relativeEmbeddings: relativeEmbeddings, + scaleFactor: scaleFactor + ) + } -final class OpenMedDebertaV2Attention: Module { - @ModuleInfo(key: "self") var selfAttention: OpenMedDisentangledSelfAttention - @ModuleInfo(key: "out_proj") var outputProjection: Linear + let batchSize = hiddenStates.dim(0) + let queryLen = queryLayer.dim(1) + let keyLen = keyLayer.dim(1) + attentionScores = attentionScores.reshaped( + batchSize, + numAttentionHeads, + queryLen, + keyLen + ) - init(_ configuration: OpenMedMLXBertConfiguration) { - _selfAttention.wrappedValue = OpenMedDisentangledSelfAttention(configuration) - _outputProjection.wrappedValue = Linear( - configuration.encoderHiddenSize, - configuration.encoderHiddenSize - ) + attentionScores = `where`( + attentionMask .> 0, + attentionScores, + openMedFloatScalar(-3.4028235e38, like: attentionScores) + ) + let attentionProbs = softmax(attentionScores, axis: -1) + let contextLayer = + attentionProbs + .reshaped(batchSize * numAttentionHeads, queryLen, keyLen) + .matmul(valueLayer) + .reshaped(batchSize, numAttentionHeads, queryLen, attentionHeadSize) + .transposed(0, 2, 1, 3) + return contextLayer.reshaped(batchSize, queryLen, allHeadSize) + } } - func callAsFunction( - hiddenStates: MLXArray, - attentionMask: MLXArray, - queryStates: MLXArray? = nil, - relativePosition: MLXArray? = nil, - relativeEmbeddings: MLXArray? = nil - ) -> MLXArray { - let context = selfAttention( - hiddenStates: hiddenStates, - attentionMask: attentionMask, - queryStates: queryStates, - relativePosition: relativePosition, - relativeEmbeddings: relativeEmbeddings - ) - return outputProjection(context) - } -} - -final class OpenMedDebertaV2Layer: Module { - private let hiddenActivation: String - - @ModuleInfo(key: "attention") var attention: OpenMedDebertaV2Attention - @ModuleInfo(key: "ln1") var attentionNorm: LayerNorm - @ModuleInfo(key: "linear1") var upProjection: Linear - @ModuleInfo(key: "linear2") var downProjection: Linear - @ModuleInfo(key: "ln2") var outputNorm: LayerNorm - - init(_ configuration: OpenMedMLXBertConfiguration) { - hiddenActivation = configuration.hiddenAct - _attention.wrappedValue = OpenMedDebertaV2Attention(configuration) - _attentionNorm.wrappedValue = LayerNorm( - dimensions: configuration.encoderHiddenSize, - eps: configuration.layerNormEps - ) - _upProjection.wrappedValue = Linear( - configuration.encoderHiddenSize, - configuration.intermediateSize - ) - _downProjection.wrappedValue = Linear( - configuration.intermediateSize, - configuration.encoderHiddenSize - ) - _outputNorm.wrappedValue = LayerNorm( - dimensions: configuration.encoderHiddenSize, - eps: configuration.layerNormEps - ) - } + final class OpenMedDebertaV2Attention: Module { + @ModuleInfo(key: "self") var selfAttention: OpenMedDisentangledSelfAttention + @ModuleInfo(key: "out_proj") var outputProjection: Linear - private func activate(_ array: MLXArray) -> MLXArray { - switch hiddenActivation { - case "gelu", "gelu_fast", "gelu_new": - return gelu(array) - case "relu": - return relu(array) - default: - return gelu(array) + init(_ configuration: OpenMedMLXBertConfiguration) { + _selfAttention.wrappedValue = OpenMedDisentangledSelfAttention(configuration) + _outputProjection.wrappedValue = Linear( + configuration.encoderHiddenSize, + configuration.encoderHiddenSize + ) } - } - func callAsFunction( - hiddenStates: MLXArray, - attentionMask: MLXArray, - queryStates: MLXArray? = nil, - relativePosition: MLXArray? = nil, - relativeEmbeddings: MLXArray? = nil - ) -> MLXArray { - let residual = queryStates ?? hiddenStates - let attentionOutput = attention( - hiddenStates: hiddenStates, - attentionMask: attentionMask, - queryStates: queryStates, - relativePosition: relativePosition, - relativeEmbeddings: relativeEmbeddings - ) - let attentionResidual = attentionNorm(residual + attentionOutput) - let feedForward = downProjection(activate(upProjection(attentionResidual))) - return outputNorm(attentionResidual + feedForward) - } -} - -final class OpenMedDebertaV2Encoder: Module { - private let relativeAttention: Bool - private let positionBuckets: Int - private let maxRelativePositions: Int - private let normalizedRelativeEmbeddings: Set - - @ModuleInfo(key: "layer") var layer: [OpenMedDebertaV2Layer] - @ModuleInfo(key: "rel_embeddings") var relativeEmbeddings: Embedding? - @ModuleInfo(key: "LayerNorm") var layerNorm: LayerNorm? - - init(_ configuration: OpenMedMLXBertConfiguration) { - relativeAttention = configuration.relativeAttention - positionBuckets = configuration.positionBuckets - var resolvedMaxRelativePositions = configuration.maxRelativePositions - if resolvedMaxRelativePositions < 1 { - resolvedMaxRelativePositions = configuration.maxPositionEmbeddings + func callAsFunction( + hiddenStates: MLXArray, + attentionMask: MLXArray, + queryStates: MLXArray? = nil, + relativePosition: MLXArray? = nil, + relativeEmbeddings: MLXArray? = nil + ) -> MLXArray { + let context = selfAttention( + hiddenStates: hiddenStates, + attentionMask: attentionMask, + queryStates: queryStates, + relativePosition: relativePosition, + relativeEmbeddings: relativeEmbeddings + ) + return outputProjection(context) } - maxRelativePositions = resolvedMaxRelativePositions - normalizedRelativeEmbeddings = Set( - configuration.normRelativeEmbedding - .lowercased() - .split(separator: "|") - .map { $0.trimmingCharacters(in: .whitespacesAndNewlines) } - ) + } - _layer.wrappedValue = (0.. 0 ? positionBuckets * 2 : resolvedMaxRelativePositions * 2 - _relativeEmbeddings.wrappedValue = Embedding( - embeddingCount: embeddingCount, - dimensions: configuration.encoderHiddenSize + final class OpenMedDebertaV2Layer: Module { + private let hiddenActivation: String + + @ModuleInfo(key: "attention") var attention: OpenMedDebertaV2Attention + @ModuleInfo(key: "ln1") var attentionNorm: LayerNorm + @ModuleInfo(key: "linear1") var upProjection: Linear + @ModuleInfo(key: "linear2") var downProjection: Linear + @ModuleInfo(key: "ln2") var outputNorm: LayerNorm + + init(_ configuration: OpenMedMLXBertConfiguration) { + hiddenActivation = configuration.hiddenAct + _attention.wrappedValue = OpenMedDebertaV2Attention(configuration) + _attentionNorm.wrappedValue = LayerNorm( + dimensions: configuration.encoderHiddenSize, + eps: configuration.layerNormEps ) - } - if normalizedRelativeEmbeddings.contains("layer_norm") { - _layerNorm.wrappedValue = LayerNorm( + _upProjection.wrappedValue = Linear( + configuration.encoderHiddenSize, + configuration.intermediateSize + ) + _downProjection.wrappedValue = Linear( + configuration.intermediateSize, + configuration.encoderHiddenSize + ) + _outputNorm.wrappedValue = LayerNorm( dimensions: configuration.encoderHiddenSize, eps: configuration.layerNormEps ) } - } - private func getRelativeEmbeddings() -> MLXArray? { - guard relativeAttention, let relativeEmbeddings else { - return nil + private func activate(_ array: MLXArray) -> MLXArray { + switch hiddenActivation { + case "gelu", "gelu_fast", "gelu_new": + return gelu(array) + case "relu": + return relu(array) + default: + return gelu(array) + } } - if normalizedRelativeEmbeddings.contains("layer_norm"), let layerNorm { - return layerNorm(relativeEmbeddings.weight) + + func callAsFunction( + hiddenStates: MLXArray, + attentionMask: MLXArray, + queryStates: MLXArray? = nil, + relativePosition: MLXArray? = nil, + relativeEmbeddings: MLXArray? = nil + ) -> MLXArray { + let residual = queryStates ?? hiddenStates + let attentionOutput = attention( + hiddenStates: hiddenStates, + attentionMask: attentionMask, + queryStates: queryStates, + relativePosition: relativePosition, + relativeEmbeddings: relativeEmbeddings + ) + let attentionResidual = attentionNorm(residual + attentionOutput) + let feedForward = downProjection(activate(upProjection(attentionResidual))) + return outputNorm(attentionResidual + feedForward) } - return relativeEmbeddings.weight } - private func getAttentionMask(_ attentionMask: MLXArray) -> MLXArray { - if attentionMask.ndim <= 2 { - return - attentionMask - .expandedDimensions(axis: 1) - .expandedDimensions(axis: 3) - * attentionMask - .expandedDimensions(axis: 1) - .expandedDimensions(axis: 2) + final class OpenMedDebertaV2Encoder: Module { + private let relativeAttention: Bool + private let positionBuckets: Int + private let maxRelativePositions: Int + private let normalizedRelativeEmbeddings: Set + + @ModuleInfo(key: "layer") var layer: [OpenMedDebertaV2Layer] + @ModuleInfo(key: "rel_embeddings") var relativeEmbeddings: Embedding? + @ModuleInfo(key: "LayerNorm") var layerNorm: LayerNorm? + + init(_ configuration: OpenMedMLXBertConfiguration) { + relativeAttention = configuration.relativeAttention + positionBuckets = configuration.positionBuckets + var resolvedMaxRelativePositions = configuration.maxRelativePositions + if resolvedMaxRelativePositions < 1 { + resolvedMaxRelativePositions = configuration.maxPositionEmbeddings + } + maxRelativePositions = resolvedMaxRelativePositions + normalizedRelativeEmbeddings = Set( + configuration.normRelativeEmbedding + .lowercased() + .split(separator: "|") + .map { $0.trimmingCharacters(in: .whitespacesAndNewlines) } + ) + + _layer.wrappedValue = (0.. 0 ? positionBuckets * 2 : resolvedMaxRelativePositions * 2 + _relativeEmbeddings.wrappedValue = Embedding( + embeddingCount: embeddingCount, + dimensions: configuration.encoderHiddenSize + ) + } + if normalizedRelativeEmbeddings.contains("layer_norm") { + _layerNorm.wrappedValue = LayerNorm( + dimensions: configuration.encoderHiddenSize, + eps: configuration.layerNormEps + ) + } } - if attentionMask.ndim == 3 { - return attentionMask.expandedDimensions(axis: 1) + + private func getRelativeEmbeddings() -> MLXArray? { + guard relativeAttention, let relativeEmbeddings else { + return nil + } + if normalizedRelativeEmbeddings.contains("layer_norm"), let layerNorm { + return layerNorm(relativeEmbeddings.weight) + } + return relativeEmbeddings.weight } - return attentionMask - } - private func getRelativePosition( - hiddenStates: MLXArray, - queryStates: MLXArray? = nil, - relativePosition: MLXArray? = nil - ) -> MLXArray? { - guard relativeAttention, relativePosition == nil else { - return relativePosition + private func getAttentionMask(_ attentionMask: MLXArray) -> MLXArray { + if attentionMask.ndim <= 2 { + return + attentionMask + .expandedDimensions(axis: 1) + .expandedDimensions(axis: 3) + * attentionMask + .expandedDimensions(axis: 1) + .expandedDimensions(axis: 2) + } + if attentionMask.ndim == 3 { + return attentionMask.expandedDimensions(axis: 1) + } + return attentionMask } - if let queryStates { + + private func getRelativePosition( + hiddenStates: MLXArray, + queryStates: MLXArray? = nil, + relativePosition: MLXArray? = nil + ) -> MLXArray? { + guard relativeAttention, relativePosition == nil else { + return relativePosition + } + if let queryStates { + return openMedBuildRelativePosition( + queryLayer: queryStates, + keyLayer: hiddenStates, + bucketSize: positionBuckets, + maxPosition: maxRelativePositions + ) + } return openMedBuildRelativePosition( - queryLayer: queryStates, + queryLayer: hiddenStates, keyLayer: hiddenStates, bucketSize: positionBuckets, maxPosition: maxRelativePositions ) } - return openMedBuildRelativePosition( - queryLayer: hiddenStates, - keyLayer: hiddenStates, - bucketSize: positionBuckets, - maxPosition: maxRelativePositions - ) - } - func callAsFunction(_ hiddenStates: MLXArray, attentionMask: MLXArray) -> MLXArray { - let attentionMask = getAttentionMask(attentionMask) - let relativePosition = getRelativePosition(hiddenStates: hiddenStates) - let relativeEmbeddings = getRelativeEmbeddings() + func callAsFunction(_ hiddenStates: MLXArray, attentionMask: MLXArray) -> MLXArray { + let attentionMask = getAttentionMask(attentionMask) + let relativePosition = getRelativePosition(hiddenStates: hiddenStates) + let relativeEmbeddings = getRelativeEmbeddings() - var states = hiddenStates - for layer in layer { - states = layer( - hiddenStates: states, - attentionMask: attentionMask, - relativePosition: relativePosition, - relativeEmbeddings: relativeEmbeddings - ) + var states = hiddenStates + for layer in layer { + states = layer( + hiddenStates: states, + attentionMask: attentionMask, + relativePosition: relativePosition, + relativeEmbeddings: relativeEmbeddings + ) + } + return states } - return states } -} -final class OpenMedDebertaV2Model: Module { - @ModuleInfo(key: "embeddings") var embeddings: OpenMedDebertaV2Embeddings - @ModuleInfo(key: "encoder") var encoder: OpenMedDebertaV2Encoder + final class OpenMedDebertaV2Model: Module { + @ModuleInfo(key: "embeddings") var embeddings: OpenMedDebertaV2Embeddings + @ModuleInfo(key: "encoder") var encoder: OpenMedDebertaV2Encoder - init(_ configuration: OpenMedMLXBertConfiguration) { - _embeddings.wrappedValue = OpenMedDebertaV2Embeddings(configuration) - _encoder.wrappedValue = OpenMedDebertaV2Encoder(configuration) - } + init(_ configuration: OpenMedMLXBertConfiguration) { + _embeddings.wrappedValue = OpenMedDebertaV2Embeddings(configuration) + _encoder.wrappedValue = OpenMedDebertaV2Encoder(configuration) + } - func callAsFunction( - inputIDs: MLXArray, - attentionMask: MLXArray? = nil, - tokenTypeIDs: MLXArray? = nil - ) -> MLXArray { - let resolvedAttentionMask = attentionMask ?? MLXArray.ones(inputIDs.shape, type: Float.self) - let embeddingOutput = embeddings( - inputIDs: inputIDs, - tokenTypeIDs: tokenTypeIDs, - attentionMask: resolvedAttentionMask - ) - return encoder(embeddingOutput, attentionMask: resolvedAttentionMask) + func callAsFunction( + inputIDs: MLXArray, + attentionMask: MLXArray? = nil, + tokenTypeIDs: MLXArray? = nil + ) -> MLXArray { + let resolvedAttentionMask = attentionMask ?? MLXArray.ones(inputIDs.shape, type: Float.self) + let embeddingOutput = embeddings( + inputIDs: inputIDs, + tokenTypeIDs: tokenTypeIDs, + attentionMask: resolvedAttentionMask + ) + return encoder(embeddingOutput, attentionMask: resolvedAttentionMask) + } } -} +#endif diff --git a/swift/OpenMedKit/Sources/OpenMedKit/OpenMedMLXModel.swift b/swift/OpenMedKit/Sources/OpenMedKit/OpenMedMLXModel.swift index 071411fe9..4484ab174 100644 --- a/swift/OpenMedKit/Sources/OpenMedKit/OpenMedMLXModel.swift +++ b/swift/OpenMedKit/Sources/OpenMedKit/OpenMedMLXModel.swift @@ -1,253 +1,255 @@ -import Foundation -import MLX -import MLXNN - -private func openMedMLXQuantizationMode(_ value: String) -> QuantizationMode { - switch value.lowercased() { - default: - return .affine +#if canImport(MLX) && !os(watchOS) && !os(visionOS) + import Foundation + import MLX + import MLXNN + + private func openMedMLXQuantizationMode(_ value: String) -> QuantizationMode { + switch value.lowercased() { + default: + return .affine + } } -} - -private final class OpenMedBertEmbeddings: Module { - private let typeVocabularySize: Int - private let positionOffset: Int - - @ModuleInfo(key: "word_embeddings") var wordEmbeddings: Embedding - @ModuleInfo(key: "position_embeddings") var positionEmbeddings: Embedding - @ModuleInfo(key: "token_type_embeddings") var tokenTypeEmbeddings: Embedding? - @ModuleInfo(key: "norm") var norm: LayerNorm - - init(_ configuration: OpenMedMLXBertConfiguration) { - typeVocabularySize = configuration.typeVocabularySize - positionOffset = configuration.positionOffset - - _wordEmbeddings.wrappedValue = Embedding( - embeddingCount: configuration.vocabularySize, - dimensions: configuration.hiddenSize - ) - _positionEmbeddings.wrappedValue = Embedding( - embeddingCount: configuration.maxPositionEmbeddings, - dimensions: configuration.hiddenSize - ) - if configuration.typeVocabularySize > 0 { - _tokenTypeEmbeddings.wrappedValue = Embedding( - embeddingCount: configuration.typeVocabularySize, + + private final class OpenMedBertEmbeddings: Module { + private let typeVocabularySize: Int + private let positionOffset: Int + + @ModuleInfo(key: "word_embeddings") var wordEmbeddings: Embedding + @ModuleInfo(key: "position_embeddings") var positionEmbeddings: Embedding + @ModuleInfo(key: "token_type_embeddings") var tokenTypeEmbeddings: Embedding? + @ModuleInfo(key: "norm") var norm: LayerNorm + + init(_ configuration: OpenMedMLXBertConfiguration) { + typeVocabularySize = configuration.typeVocabularySize + positionOffset = configuration.positionOffset + + _wordEmbeddings.wrappedValue = Embedding( + embeddingCount: configuration.vocabularySize, + dimensions: configuration.hiddenSize + ) + _positionEmbeddings.wrappedValue = Embedding( + embeddingCount: configuration.maxPositionEmbeddings, dimensions: configuration.hiddenSize ) + if configuration.typeVocabularySize > 0 { + _tokenTypeEmbeddings.wrappedValue = Embedding( + embeddingCount: configuration.typeVocabularySize, + dimensions: configuration.hiddenSize + ) + } + _norm.wrappedValue = LayerNorm( + dimensions: configuration.hiddenSize, + eps: configuration.layerNormEps + ) } - _norm.wrappedValue = LayerNorm( - dimensions: configuration.hiddenSize, - eps: configuration.layerNormEps - ) - } - func callAsFunction( - _ inputIDs: MLXArray, - tokenTypeIDs: MLXArray? - ) -> MLXArray { - let positionIDs = broadcast( - MLXArray.arange(inputIDs.dim(1)) + positionOffset, - to: inputIDs.shape - ) + func callAsFunction( + _ inputIDs: MLXArray, + tokenTypeIDs: MLXArray? + ) -> MLXArray { + let positionIDs = broadcast( + MLXArray.arange(inputIDs.dim(1)) + positionOffset, + to: inputIDs.shape + ) - var embeddings = wordEmbeddings(inputIDs) + positionEmbeddings(positionIDs) - if typeVocabularySize > 0, let tokenTypeEmbeddings { - let tokenTypeIDs = tokenTypeIDs ?? MLXArray.zeros(like: inputIDs) - embeddings += tokenTypeEmbeddings(tokenTypeIDs) - } + var embeddings = wordEmbeddings(inputIDs) + positionEmbeddings(positionIDs) + if typeVocabularySize > 0, let tokenTypeEmbeddings { + let tokenTypeIDs = tokenTypeIDs ?? MLXArray.zeros(like: inputIDs) + embeddings += tokenTypeEmbeddings(tokenTypeIDs) + } - return norm(embeddings) - } -} - -private final class OpenMedBertEncoderLayer: Module { - let attention: MultiHeadAttention - - @ModuleInfo(key: "ln1") var attentionNorm: LayerNorm - @ModuleInfo(key: "ln2") var outputNorm: LayerNorm - @ModuleInfo(key: "linear1") var upProjection: Linear - @ModuleInfo(key: "linear2") var downProjection: Linear - - init(_ configuration: OpenMedMLXBertConfiguration) { - attention = MultiHeadAttention( - dimensions: configuration.hiddenSize, - numHeads: configuration.numAttentionHeads, - bias: true - ) - _attentionNorm.wrappedValue = LayerNorm( - dimensions: configuration.hiddenSize, - eps: configuration.layerNormEps - ) - _outputNorm.wrappedValue = LayerNorm( - dimensions: configuration.hiddenSize, - eps: configuration.layerNormEps - ) - _upProjection.wrappedValue = Linear( - configuration.hiddenSize, - configuration.intermediateSize - ) - _downProjection.wrappedValue = Linear( - configuration.intermediateSize, - configuration.hiddenSize - ) + return norm(embeddings) + } } - func callAsFunction(_ inputs: MLXArray, mask: MLXArray?) -> MLXArray { - let attentionOutput = attention(inputs, keys: inputs, values: inputs, mask: mask) - let attentionResidual = attentionNorm(inputs + attentionOutput) - let feedForwardOutput = downProjection(gelu(upProjection(attentionResidual))) - return outputNorm(attentionResidual + feedForwardOutput) - } -} + private final class OpenMedBertEncoderLayer: Module { + let attention: MultiHeadAttention -private final class OpenMedBertEncoder: Module { - let layers: [OpenMedBertEncoderLayer] + @ModuleInfo(key: "ln1") var attentionNorm: LayerNorm + @ModuleInfo(key: "ln2") var outputNorm: LayerNorm + @ModuleInfo(key: "linear1") var upProjection: Linear + @ModuleInfo(key: "linear2") var downProjection: Linear - init(_ configuration: OpenMedMLXBertConfiguration) { - layers = (0.. MLXArray { - var hiddenStates = inputs - for layer in layers { - hiddenStates = layer(hiddenStates, mask: mask) + func callAsFunction(_ inputs: MLXArray, mask: MLXArray?) -> MLXArray { + let attentionOutput = attention(inputs, keys: inputs, values: inputs, mask: mask) + let attentionResidual = attentionNorm(inputs + attentionOutput) + let feedForwardOutput = downProjection(gelu(upProjection(attentionResidual))) + return outputNorm(attentionResidual + feedForwardOutput) } - return hiddenStates } -} -final class OpenMedBertForTokenClassification: Module { - @ModuleInfo(key: "embeddings") fileprivate var embeddings: OpenMedBertEmbeddings - @ModuleInfo(key: "classifier") var classifier: Linear + private final class OpenMedBertEncoder: Module { + let layers: [OpenMedBertEncoderLayer] - let configuration: OpenMedMLXBertConfiguration - fileprivate let encoder: OpenMedBertEncoder + init(_ configuration: OpenMedMLXBertConfiguration) { + layers = (0.. MLXArray { + var hiddenStates = inputs + for layer in layers { + hiddenStates = layer(hiddenStates, mask: mask) + } + return hiddenStates + } } - func callAsFunction( - _ inputIDs: MLXArray, - tokenTypeIDs: MLXArray? = nil, - attentionMask: MLXArray? = nil - ) -> MLXArray { - var inputs = inputIDs - if inputs.ndim == 1 { - inputs = inputs.reshaped(1, -1) - } - - let embedded = embeddings(inputs, tokenTypeIDs: tokenTypeIDs) - let mask: MLXArray? - if let attentionMask { - mask = - attentionMask - .asType(embedded.dtype) - .expandedDimensions(axes: [1, 2]) - .log() - } else { - mask = nil - } - - let hiddenStates = encoder(embedded, mask: mask) - return classifier(hiddenStates) - } + final class OpenMedBertForTokenClassification: Module { + @ModuleInfo(key: "embeddings") fileprivate var embeddings: OpenMedBertEmbeddings + @ModuleInfo(key: "classifier") var classifier: Linear - func sanitize(weights: [String: MLXArray]) -> [String: MLXArray] { - weights.filter { key, _ in - key != "embeddings.position_ids" && !key.hasPrefix("_") + let configuration: OpenMedMLXBertConfiguration + fileprivate let encoder: OpenMedBertEncoder + + init(_ configuration: OpenMedMLXBertConfiguration) { + self.configuration = configuration + self.encoder = OpenMedBertEncoder(configuration) + _embeddings.wrappedValue = OpenMedBertEmbeddings(configuration) + _classifier.wrappedValue = Linear(configuration.hiddenSize, configuration.numLabels) } - } -} -enum OpenMedMLXModelLoader { - private static func loadedWeights(for artifact: OpenMedMLXArtifact) throws -> [String: MLXArray] { - try OpenMedMLXWeightArchive.loadWeights(from: artifact.weightCandidateURLs) - } + func callAsFunction( + _ inputIDs: MLXArray, + tokenTypeIDs: MLXArray? = nil, + attentionMask: MLXArray? = nil + ) -> MLXArray { + var inputs = inputIDs + if inputs.ndim == 1 { + inputs = inputs.reshaped(1, -1) + } - static func loadTokenClassifier( - from artifact: OpenMedMLXArtifact - ) throws -> OpenMedBertForTokenClassification { - var weights = try loadedWeights(for: artifact) - let model = OpenMedBertForTokenClassification(artifact.configuration) - weights = model.sanitize(weights: weights) - - if let bits = artifact.configuration.quantizationBits { - let groupSize = artifact.configuration.quantizationGroupSize - let mode = openMedMLXQuantizationMode(artifact.configuration.quantizationMode) - quantize(model: model) { path, _ in - if weights["\(path).scales"] != nil { - return (groupSize, bits, mode) - } else { - return nil - } + let embedded = embeddings(inputs, tokenTypeIDs: tokenTypeIDs) + let mask: MLXArray? + if let attentionMask { + mask = + attentionMask + .asType(embedded.dtype) + .expandedDimensions(axes: [1, 2]) + .log() + } else { + mask = nil } + + let hiddenStates = encoder(embedded, mask: mask) + return classifier(hiddenStates) } - try model.update(parameters: ModuleParameters.unflattened(weights), verify: [.all]) - eval(model) - return model + func sanitize(weights: [String: MLXArray]) -> [String: MLXArray] { + weights.filter { key, _ in + key != "embeddings.position_ids" && !key.hasPrefix("_") + } + } } - static func loadPrivacyFilter( - from artifact: OpenMedMLXArtifact - ) throws -> OpenMedPrivacyFilterForTokenClassification { - let model = OpenMedPrivacyFilterForTokenClassification(artifact.configuration) - var weights = model.sanitize(weights: try loadedWeights(for: artifact)) - - if let bits = artifact.configuration.quantizationBits { - let groupSize = artifact.configuration.quantizationGroupSize - let mode = openMedMLXQuantizationMode(artifact.configuration.quantizationMode) - model.installQuantizedPlaceholders( - where: { weights["\($0).scales"] != nil }, - groupSize: groupSize, - bits: bits, - mode: mode - ) + enum OpenMedMLXModelLoader { + private static func loadedWeights(for artifact: OpenMedMLXArtifact) throws -> [String: MLXArray] { + try OpenMedMLXWeightArchive.loadWeights(from: artifact.weightCandidateURLs) } - try model.update(parameters: ModuleParameters.unflattened(weights), verify: [.all]) - weights.removeAll(keepingCapacity: false) - eval(model.parameters()) - return model - } + static func loadTokenClassifier( + from artifact: OpenMedMLXArtifact + ) throws -> OpenMedBertForTokenClassification { + var weights = try loadedWeights(for: artifact) + let model = OpenMedBertForTokenClassification(artifact.configuration) + weights = model.sanitize(weights: weights) + + if let bits = artifact.configuration.quantizationBits { + let groupSize = artifact.configuration.quantizationGroupSize + let mode = openMedMLXQuantizationMode(artifact.configuration.quantizationMode) + quantize(model: model) { path, _ in + if weights["\(path).scales"] != nil { + return (groupSize, bits, mode) + } else { + return nil + } + } + } - static func loadGLiNERSpanModel( - from artifact: OpenMedMLXArtifact - ) throws -> OpenMedGLiNERSpanModel { - var weights = try loadedWeights(for: artifact) - let model = OpenMedGLiNERSpanModel(artifact.configuration) - weights = model.sanitize(weights: weights) - try model.update(parameters: ModuleParameters.unflattened(weights), verify: [.all]) - eval(model) - return model - } + try model.update(parameters: ModuleParameters.unflattened(weights), verify: [.all]) + eval(model) + return model + } - static func loadGLiClassUniEncoderModel( - from artifact: OpenMedMLXArtifact - ) throws -> OpenMedGLiClassUniEncoderModel { - var weights = try loadedWeights(for: artifact) - let model = OpenMedGLiClassUniEncoderModel(artifact.configuration) - weights = model.sanitize(weights: weights) - try model.update(parameters: ModuleParameters.unflattened(weights), verify: [.all]) - eval(model) - return model - } + static func loadPrivacyFilter( + from artifact: OpenMedMLXArtifact + ) throws -> OpenMedPrivacyFilterForTokenClassification { + let model = OpenMedPrivacyFilterForTokenClassification(artifact.configuration) + var weights = model.sanitize(weights: try loadedWeights(for: artifact)) + + if let bits = artifact.configuration.quantizationBits { + let groupSize = artifact.configuration.quantizationGroupSize + let mode = openMedMLXQuantizationMode(artifact.configuration.quantizationMode) + model.installQuantizedPlaceholders( + where: { weights["\($0).scales"] != nil }, + groupSize: groupSize, + bits: bits, + mode: mode + ) + } + + try model.update(parameters: ModuleParameters.unflattened(weights), verify: [.all]) + weights.removeAll(keepingCapacity: false) + eval(model.parameters()) + return model + } + + static func loadGLiNERSpanModel( + from artifact: OpenMedMLXArtifact + ) throws -> OpenMedGLiNERSpanModel { + var weights = try loadedWeights(for: artifact) + let model = OpenMedGLiNERSpanModel(artifact.configuration) + weights = model.sanitize(weights: weights) + try model.update(parameters: ModuleParameters.unflattened(weights), verify: [.all]) + eval(model) + return model + } - static func loadGLiNERRelexModel( - from artifact: OpenMedMLXArtifact - ) throws -> OpenMedGLiNERRelexModel { - var weights = try loadedWeights(for: artifact) - let model = OpenMedGLiNERRelexModel(artifact.configuration) - weights = model.sanitize(weights: weights) - try model.update(parameters: ModuleParameters.unflattened(weights), verify: [.all]) - eval(model) - return model + static func loadGLiClassUniEncoderModel( + from artifact: OpenMedMLXArtifact + ) throws -> OpenMedGLiClassUniEncoderModel { + var weights = try loadedWeights(for: artifact) + let model = OpenMedGLiClassUniEncoderModel(artifact.configuration) + weights = model.sanitize(weights: weights) + try model.update(parameters: ModuleParameters.unflattened(weights), verify: [.all]) + eval(model) + return model + } + + static func loadGLiNERRelexModel( + from artifact: OpenMedMLXArtifact + ) throws -> OpenMedGLiNERRelexModel { + var weights = try loadedWeights(for: artifact) + let model = OpenMedGLiNERRelexModel(artifact.configuration) + weights = model.sanitize(weights: weights) + try model.update(parameters: ModuleParameters.unflattened(weights), verify: [.all]) + eval(model) + return model + } } -} +#endif diff --git a/swift/OpenMedKit/Sources/OpenMedKit/OpenMedMLXPipeline.swift b/swift/OpenMedKit/Sources/OpenMedKit/OpenMedMLXPipeline.swift index 49a18b591..7e28e3f7d 100644 --- a/swift/OpenMedKit/Sources/OpenMedKit/OpenMedMLXPipeline.swift +++ b/swift/OpenMedKit/Sources/OpenMedKit/OpenMedMLXPipeline.swift @@ -1,115 +1,117 @@ -import Foundation -import MLX +#if canImport(MLX) && !os(watchOS) && !os(visionOS) + import Foundation + import MLX -enum OpenMedMLXRuntimeError: LocalizedError { - case unsupportedPlatform + enum OpenMedMLXRuntimeError: LocalizedError { + case unsupportedPlatform - var errorDescription: String? { - switch self { - case .unsupportedPlatform: - return "Swift MLX inference requires Apple Silicon macOS or a real iPhone/iPad device." + var errorDescription: String? { + switch self { + case .unsupportedPlatform: + return "Swift MLX inference requires Apple Silicon macOS or a real iPhone/iPad device." + } } } -} -final class MLXTokenClassificationPipeline { - private let artifact: OpenMedMLXArtifact - private let model: OpenMedBertForTokenClassification - private let maxSeqLength: Int + final class MLXTokenClassificationPipeline { + private let artifact: OpenMedMLXArtifact + private let model: OpenMedBertForTokenClassification + private let maxSeqLength: Int - init(modelDirectoryURL: URL, maxSeqLength: Int = 512) throws { - guard Self.isRuntimeSupported else { - throw OpenMedMLXRuntimeError.unsupportedPlatform - } + init(modelDirectoryURL: URL, maxSeqLength: Int = 512) throws { + guard Self.isRuntimeSupported else { + throw OpenMedMLXRuntimeError.unsupportedPlatform + } - let artifact = try OpenMedMLXArtifact(modelDirectoryURL: modelDirectoryURL) - self.artifact = artifact - self.model = try OpenMedMLXModelLoader.loadTokenClassifier(from: artifact) - self.maxSeqLength = min(maxSeqLength, artifact.manifest.maxSequenceLength ?? maxSeqLength) - } + let artifact = try OpenMedMLXArtifact(modelDirectoryURL: modelDirectoryURL) + self.artifact = artifact + self.model = try OpenMedMLXModelLoader.loadTokenClassifier(from: artifact) + self.maxSeqLength = min(maxSeqLength, artifact.manifest.maxSequenceLength ?? maxSeqLength) + } - var tokenizerDirectoryURL: URL? { - artifact.tokenizerDirectoryURL - } + var tokenizerDirectoryURL: URL? { + artifact.tokenizerDirectoryURL + } - var tokenizerName: String? { - artifact.tokenizerName - } + var tokenizerName: String? { + artifact.tokenizerName + } - var resolvedMaxSequenceLength: Int { - maxSeqLength - } + var resolvedMaxSequenceLength: Int { + maxSeqLength + } - func predict( - inputIDs: [Int], - attentionMask: [Int], - tokenTypeIDs: [Int], - offsets: [(Int, Int)], - text: String, - strategy: PostProcessing.AggregationStrategy = .average - ) throws -> [EntityPrediction] { - let sequenceLength = inputIDs.count - - let inputArray = MLXArray(inputIDs, [1, sequenceLength]) - let attentionArray = MLXArray(attentionMask, [1, sequenceLength]).asType(.float32) - let tokenTypeArray: MLXArray? = - artifact.configuration.typeVocabularySize > 0 - ? MLXArray(tokenTypeIDs, [1, sequenceLength]) - : nil - - let logits = model( - inputArray, - tokenTypeIDs: tokenTypeArray, - attentionMask: attentionArray - ) - eval(logits) - - let probabilities = softmax(logits[0], axis: -1) - let predictions = probabilities.argMax(axis: -1) - eval(probabilities, predictions) - - let flatProbabilities = probabilities.asArray(Float.self) - let predictedLabelIDs = predictions.asArray(Int32.self).map(Int.init) - let numLabels = artifact.configuration.numLabels - - var tokenPredictions = [PostProcessing.TokenPrediction]() - tokenPredictions.reserveCapacity(sequenceLength) - - for tokenIndex in 0.. [EntityPrediction] { + let sequenceLength = inputIDs.count + + let inputArray = MLXArray(inputIDs, [1, sequenceLength]) + let attentionArray = MLXArray(attentionMask, [1, sequenceLength]).asType(.float32) + let tokenTypeArray: MLXArray? = + artifact.configuration.typeVocabularySize > 0 + ? MLXArray(tokenTypeIDs, [1, sequenceLength]) + : nil + + let logits = model( + inputArray, + tokenTypeIDs: tokenTypeArray, + attentionMask: attentionArray + ) + eval(logits) + + let probabilities = softmax(logits[0], axis: -1) + let predictions = probabilities.argMax(axis: -1) + eval(probabilities, predictions) + + let flatProbabilities = probabilities.asArray(Float.self) + let predictedLabelIDs = predictions.asArray(Int32.self).map(Int.init) + let numLabels = artifact.configuration.numLabels + + var tokenPredictions = [PostProcessing.TokenPrediction]() + tokenPredictions.reserveCapacity(sequenceLength) + + for tokenIndex in 0.. [String: MLXArray] { - for url in candidateURLs where FileManager.default.fileExists(atPath: url.path) { - switch url.pathExtension.lowercased() { - case "safetensors": - return try MLX.loadArrays(url: url) - case "npz": - return try loadNPZ(url: url) - default: - throw OpenMedMLXWeightError.unsupportedWeightFile(url) +#if canImport(MLX) && canImport(ZIPFoundation) && !os(watchOS) && !os(visionOS) + import Foundation + import MLX + import ZIPFoundation + + enum OpenMedMLXWeightError: LocalizedError { + case unsupportedWeightFile(URL) + case missingWeights([URL]) + case invalidNPYHeader(String) + case unsupportedNPYDType(String) + case unsupportedNPYEndianness(String) + case unsupportedFortranOrder(String) + case invalidNPZArchive(URL) + + var errorDescription: String? { + switch self { + case .unsupportedWeightFile(let url): + return "Unsupported MLX weight file: \(url.lastPathComponent)" + case .missingWeights(let urls): + let checked = urls.map(\.lastPathComponent).joined(separator: ", ") + return "No MLX weights found. Checked: \(checked)" + case .invalidNPYHeader(let name): + return "Invalid NumPy array header in \(name)" + case .unsupportedNPYDType(let dtype): + return "Unsupported NumPy dtype for Swift MLX loading: \(dtype)" + case .unsupportedNPYEndianness(let dtype): + return "Unsupported non-little-endian NumPy dtype: \(dtype)" + case .unsupportedFortranOrder(let name): + return "Fortran-ordered NumPy arrays are not supported in \(name)" + case .invalidNPZArchive(let url): + return "Unable to open NPZ archive at \(url.path)" } } - - throw OpenMedMLXWeightError.missingWeights(candidateURLs) } - private static func loadNPZ(url: URL) throws -> [String: MLXArray] { - let archive: Archive - do { - archive = try Archive(url: url, accessMode: .read) - } catch { - throw OpenMedMLXWeightError.invalidNPZArchive(url) + enum OpenMedMLXWeightArchive { + static func loadWeights(from candidateURLs: [URL]) throws -> [String: MLXArray] { + for url in candidateURLs where FileManager.default.fileExists(atPath: url.path) { + switch url.pathExtension.lowercased() { + case "safetensors": + return try MLX.loadArrays(url: url) + case "npz": + return try loadNPZ(url: url) + default: + throw OpenMedMLXWeightError.unsupportedWeightFile(url) + } + } + + throw OpenMedMLXWeightError.missingWeights(candidateURLs) } - var arrays = [String: MLXArray]() - for entry in archive where entry.path.hasSuffix(".npy") { - var data = Data() - _ = try archive.extract(entry) { chunk in - data.append(chunk) + private static func loadNPZ(url: URL) throws -> [String: MLXArray] { + let archive: Archive + do { + archive = try Archive(url: url, accessMode: .read) + } catch { + throw OpenMedMLXWeightError.invalidNPZArchive(url) } - let key = URL(fileURLWithPath: entry.path).deletingPathExtension().lastPathComponent - arrays[key] = try loadNPY(data: data, name: entry.path) - } + var arrays = [String: MLXArray]() + for entry in archive where entry.path.hasSuffix(".npy") { + var data = Data() + _ = try archive.extract(entry) { chunk in + data.append(chunk) + } - return arrays - } + let key = URL(fileURLWithPath: entry.path).deletingPathExtension().lastPathComponent + arrays[key] = try loadNPY(data: data, name: entry.path) + } - private static func loadNPY(data: Data, name: String) throws -> MLXArray { - let header = try parseNPYHeader(data: data, name: name) - guard !header.fortranOrder else { - throw OpenMedMLXWeightError.unsupportedFortranOrder(name) + return arrays } - let payload = data.subdata(in: header.dataOffset.. MLXArray { + let header = try parseNPYHeader(data: data, name: name) + guard !header.fortranOrder else { + throw OpenMedMLXWeightError.unsupportedFortranOrder(name) + } - private static func parseNPYHeader(data: Data, name: String) throws -> ParsedNPYHeader { - let minimumLength = 10 - guard data.count >= minimumLength else { - throw OpenMedMLXWeightError.invalidNPYHeader(name) + let payload = data.subdata(in: header.dataOffset..= 12 else { + private static func parseNPYHeader(data: Data, name: String) throws -> ParsedNPYHeader { + let minimumLength = 10 + guard data.count >= minimumLength else { throw OpenMedMLXWeightError.invalidNPYHeader(name) } - headerLength = - Int(data[8]) | (Int(data[9]) << 8) | (Int(data[10]) << 16) - | (Int(data[11]) << 24) - headerOffset = 12 - default: - throw OpenMedMLXWeightError.invalidNPYHeader(name) - } - let endOffset = headerOffset + headerLength - guard data.count >= endOffset else { - throw OpenMedMLXWeightError.invalidNPYHeader(name) - } + let magic = Data([0x93, 0x4e, 0x55, 0x4d, 0x50, 0x59]) + guard data.prefix(6) == magic else { + throw OpenMedMLXWeightError.invalidNPYHeader(name) + } - let headerData = data.subdata(in: headerOffset..= 12 else { + throw OpenMedMLXWeightError.invalidNPYHeader(name) + } + headerLength = + Int(data[8]) | (Int(data[9]) << 8) | (Int(data[10]) << 16) + | (Int(data[11]) << 24) + headerOffset = 12 + default: + throw OpenMedMLXWeightError.invalidNPYHeader(name) + } - let descr = try match(pattern: "'descr'\\s*:\\s*'([^']+)'", in: header, name: name) - let fortranText = try match( - pattern: "'fortran_order'\\s*:\\s*(True|False)", - in: header, - name: name - ) - let shapeText = try match(pattern: "'shape'\\s*:\\s*\\(([^\\)]*)\\)", in: header, name: name) - - let dtype = try dtype(from: descr) - let shape = parseShape(shapeText) - - return ParsedNPYHeader( - dtype: dtype, - shape: shape, - fortranOrder: fortranText == "True", - dataOffset: endOffset - ) - } + let endOffset = headerOffset + headerLength + guard data.count >= endOffset else { + throw OpenMedMLXWeightError.invalidNPYHeader(name) + } + + let headerData = data.subdata(in: headerOffset.. String { - let regex = try NSRegularExpression(pattern: pattern) - let range = NSRange(header.startIndex.. 1, - let captureRange = Range(result.range(at: 1), in: header) - else { - throw OpenMedMLXWeightError.invalidNPYHeader(name) + let descr = try match(pattern: "'descr'\\s*:\\s*'([^']+)'", in: header, name: name) + let fortranText = try match( + pattern: "'fortran_order'\\s*:\\s*(True|False)", + in: header, + name: name + ) + let shapeText = try match(pattern: "'shape'\\s*:\\s*\\(([^\\)]*)\\)", in: header, name: name) + + let dtype = try dtype(from: descr) + let shape = parseShape(shapeText) + + return ParsedNPYHeader( + dtype: dtype, + shape: shape, + fortranOrder: fortranText == "True", + dataOffset: endOffset + ) } - return String(header[captureRange]) - } - private static func parseShape(_ shapeText: String) -> [Int] { - shapeText - .split(separator: ",") - .map { $0.trimmingCharacters(in: .whitespacesAndNewlines) } - .filter { !$0.isEmpty } - .compactMap(Int.init) - } + private static func match(pattern: String, in header: String, name: String) throws -> String { + let regex = try NSRegularExpression(pattern: pattern) + let range = NSRange(header.startIndex.. 1, + let captureRange = Range(result.range(at: 1), in: header) + else { + throw OpenMedMLXWeightError.invalidNPYHeader(name) + } + return String(header[captureRange]) + } - private static func dtype(from descriptor: String) throws -> DType { - if descriptor.hasPrefix(">") { - throw OpenMedMLXWeightError.unsupportedNPYEndianness(descriptor) + private static func parseShape(_ shapeText: String) -> [Int] { + shapeText + .split(separator: ",") + .map { $0.trimmingCharacters(in: .whitespacesAndNewlines) } + .filter { !$0.isEmpty } + .compactMap(Int.init) } - switch descriptor { - case "|b1", "?": - return .bool - case "|u1", " DType { + if descriptor.hasPrefix(">") { + throw OpenMedMLXWeightError.unsupportedNPYEndianness(descriptor) + } + + switch descriptor { + case "|b1", "?": + return .bool + case "|u1", " DType { - switch configuration.parameterDType.lowercased() { - case "bf16", "bfloat16": - return .bfloat16 - default: - return .float32 +#if canImport(MLX) && !os(watchOS) && !os(visionOS) + import Foundation + import MLX + import MLXNN + + private func privacyFilterParameterDType(_ configuration: OpenMedMLXBertConfiguration) -> DType { + switch configuration.parameterDType.lowercased() { + case "bf16", "bfloat16": + return .bfloat16 + default: + return .float32 + } } -} -private func privacyFilterScalar(_ value: Float, like array: MLXArray) -> MLXArray { - MLXArray(value).asType(array.dtype) -} + private func privacyFilterScalar(_ value: Float, like array: MLXArray) -> MLXArray { + MLXArray(value).asType(array.dtype) + } -private func privacyFilterLinearInput(_ input: MLXArray, for linear: Linear) -> MLXArray { - if linear is Quantized { - return input + private func privacyFilterLinearInput(_ input: MLXArray, for linear: Linear) -> MLXArray { + if linear is Quantized { + return input + } + return input.asType(linear.weight.dtype) } - return input.asType(linear.weight.dtype) -} - -private func privacyFilterExpertInput( - _ input: MLXArray, - for expert: OpenMedPrivacyFilterExpertLinear -) -> MLXArray { - if expert is Quantized { - return input + + private func privacyFilterExpertInput( + _ input: MLXArray, + for expert: OpenMedPrivacyFilterExpertLinear + ) -> MLXArray { + if expert is Quantized { + return input + } + return input.asType(expert.weight.dtype) } - return input.asType(expert.weight.dtype) -} -private final class OpenMedPrivacyFilterRMSNorm: Module { - private let eps: Float + private final class OpenMedPrivacyFilterRMSNorm: Module { + private let eps: Float + + @ParameterInfo(key: "scale") var scale: MLXArray - @ParameterInfo(key: "scale") var scale: MLXArray + init(hiddenSize: Int, eps: Float) { + self.eps = eps + _scale.wrappedValue = MLXArray.ones([hiddenSize], type: Float.self) + super.init() + } - init(hiddenSize: Int, eps: Float) { - self.eps = eps - _scale.wrappedValue = MLXArray.ones([hiddenSize], type: Float.self) - super.init() + func callAsFunction(_ input: MLXArray) -> MLXArray { + let dtype = input.dtype + let normalized = input.asType(.float32) + let variance = mean(normalized * normalized, axis: -1, keepDims: true) + return (normalized * rsqrt(variance + MLXArray(eps)) * scale).asType(dtype) + } } - func callAsFunction(_ input: MLXArray) -> MLXArray { + private func applyPrivacyFilterRotaryEmbedding( + _ input: MLXArray, + cos: MLXArray, + sin: MLXArray + ) -> MLXArray { let dtype = input.dtype - let normalized = input.asType(.float32) - let variance = mean(normalized * normalized, axis: -1, keepDims: true) - return (normalized * rsqrt(variance + MLXArray(eps)) * scale).asType(dtype) - } -} - -private func applyPrivacyFilterRotaryEmbedding( - _ input: MLXArray, - cos: MLXArray, - sin: MLXArray -) -> MLXArray { - let dtype = input.dtype - let shape = input.shape - let headDim = shape[3] - let paired = input.reshaped(shape[0], shape[1], shape[2], headDim / 2, 2) - let x1 = paired[0..., 0..., 0..., 0..., 0] - let x2 = paired[0..., 0..., 0..., 0..., 1] - let cos = cos.expandedDimensions(axes: [0, 2]).asType(dtype) - let sin = sin.expandedDimensions(axes: [0, 2]).asType(dtype) - let rotated = stacked([x1 * cos - x2 * sin, x2 * cos + x1 * sin], axis: -1) - return rotated.reshaped(shape).asType(dtype) -} - -private final class OpenMedPrivacyFilterRotaryEmbedding { - private let headDim: Int - private let base: Float - private let initialContextLength: Int - private let scalingFactor: Float - private let ntkAlpha: Float - private let ntkBeta: Float - - init(_ configuration: OpenMedMLXBertConfiguration) { - self.headDim = configuration.headDim - self.base = configuration.ropeTheta - self.initialContextLength = configuration.initialContextLength - self.scalingFactor = configuration.ropeScalingFactor - self.ntkAlpha = configuration.ropeNTKAlpha - self.ntkBeta = configuration.ropeNTKBeta + let shape = input.shape + let headDim = shape[3] + let paired = input.reshaped(shape[0], shape[1], shape[2], headDim / 2, 2) + let x1 = paired[0..., 0..., 0..., 0..., 0] + let x2 = paired[0..., 0..., 0..., 0..., 1] + let cos = cos.expandedDimensions(axes: [0, 2]).asType(dtype) + let sin = sin.expandedDimensions(axes: [0, 2]).asType(dtype) + let rotated = stacked([x1 * cos - x2 * sin, x2 * cos + x1 * sin], axis: -1) + return rotated.reshaped(shape).asType(dtype) } - private func cosSin(numTokens: Int) -> (MLXArray, MLXArray) { - let halfDim = headDim / 2 - let exponents = MLXArray.arange(0, headDim, step: 2, dtype: .float32) / Float(headDim) - let frequency = pow(base, exponents) - - let inverseFrequency: MLXArray - let concentration: Float - if scalingFactor > 1.0 { - concentration = 0.1 * Foundation.log(scalingFactor) + 1.0 - let halfDimFloat = Float(headDim) / 2.0 - let denominator = Foundation.log(base) - let low = - halfDimFloat - * Foundation.log(Float(initialContextLength) / (ntkBeta * 2.0 * Float.pi)) - / denominator - let high = - halfDimFloat - * Foundation.log(Float(initialContextLength) / (ntkAlpha * 2.0 * Float.pi)) - / denominator - let interpolation = 1.0 / (scalingFactor * frequency) - let extrapolation = 1.0 / frequency - let ramp = (MLXArray.arange(halfDim, dtype: .float32) - low) / (high - low) - let mask = 1.0 - clip(ramp, min: 0.0, max: 1.0) - inverseFrequency = interpolation * (1.0 - mask) + extrapolation * mask - } else { - concentration = 1.0 - inverseFrequency = 1.0 / frequency + private final class OpenMedPrivacyFilterRotaryEmbedding { + private let headDim: Int + private let base: Float + private let initialContextLength: Int + private let scalingFactor: Float + private let ntkAlpha: Float + private let ntkBeta: Float + + init(_ configuration: OpenMedMLXBertConfiguration) { + self.headDim = configuration.headDim + self.base = configuration.ropeTheta + self.initialContextLength = configuration.initialContextLength + self.scalingFactor = configuration.ropeScalingFactor + self.ntkAlpha = configuration.ropeNTKAlpha + self.ntkBeta = configuration.ropeNTKBeta } - let positions = MLXArray.arange(numTokens, dtype: .float32) - let frequencies = positions.expandedDimensions(axis: 1) * inverseFrequency.expandedDimensions(axis: 0) - return (cos(frequencies) * concentration, sin(frequencies) * concentration) - } + private func cosSin(numTokens: Int) -> (MLXArray, MLXArray) { + let halfDim = headDim / 2 + let exponents = MLXArray.arange(0, headDim, step: 2, dtype: .float32) / Float(headDim) + let frequency = pow(base, exponents) + + let inverseFrequency: MLXArray + let concentration: Float + if scalingFactor > 1.0 { + concentration = 0.1 * Foundation.log(scalingFactor) + 1.0 + let halfDimFloat = Float(headDim) / 2.0 + let denominator = Foundation.log(base) + let low = + halfDimFloat + * Foundation.log(Float(initialContextLength) / (ntkBeta * 2.0 * Float.pi)) + / denominator + let high = + halfDimFloat + * Foundation.log(Float(initialContextLength) / (ntkAlpha * 2.0 * Float.pi)) + / denominator + let interpolation = 1.0 / (scalingFactor * frequency) + let extrapolation = 1.0 / frequency + let ramp = (MLXArray.arange(halfDim, dtype: .float32) - low) / (high - low) + let mask = 1.0 - clip(ramp, min: 0.0, max: 1.0) + inverseFrequency = interpolation * (1.0 - mask) + extrapolation * mask + } else { + concentration = 1.0 + inverseFrequency = 1.0 / frequency + } - func callAsFunction(query: MLXArray, key: MLXArray) -> (MLXArray, MLXArray) { - let (cos, sin) = cosSin(numTokens: query.dim(1)) - return ( - applyPrivacyFilterRotaryEmbedding(query, cos: cos, sin: sin), - applyPrivacyFilterRotaryEmbedding(key, cos: cos, sin: sin) - ) - } -} - -private func privacyFilterTopK(_ values: MLXArray, k: Int) -> (MLXArray, MLXArray) { - let indices = argPartition(-values, kth: k - 1, axis: -1)[0..., 0.. MLXArray { - let half = input.dim(-1) / 2 - let glu = minimum(input[0..., 0..., 0.. (MLXArray, MLXArray) { + let (cos, sin) = cosSin(numTokens: query.dim(1)) + return ( + applyPrivacyFilterRotaryEmbedding(query, cos: cos, sin: sin), + applyPrivacyFilterRotaryEmbedding(key, cos: cos, sin: sin) + ) + } } - func callAsFunction(_ input: MLXArray, expertIndices: MLXArray) -> MLXArray { - let inputShape = input.shape - let flatInput = input.reshaped(-1, 1, input.dim(-1)).asType(weight.dtype) - let flatIndices = expertIndices.reshaped(-1).asType(.int32) - var output = gatherMM(flatInput, weight, rhsIndices: flatIndices).squeezed(axis: -2) - output = output + bias.take(flatIndices, axis: 0) - return output.reshaped(Array(inputShape.dropLast()) + [outputSize]) + private func privacyFilterTopK(_ values: MLXArray, k: Int) -> (MLXArray, MLXArray) { + let indices = argPartition(-values, kth: k - 1, axis: -1)[0..., 0.. Module { - OpenMedPrivacyFilterQuantizedExpertLinear( - self, - groupSize: groupSize, - bits: bits, - mode: mode - ) - } -} - -private final class OpenMedPrivacyFilterQuantizedExpertLinear: - OpenMedPrivacyFilterExpertLinear, Quantized -{ - let groupSize: Int - let bits: Int - let mode: QuantizationMode - let scales: MLXArray - let biases: MLXArray? - - init( - _ other: OpenMedPrivacyFilterExpertLinear, - groupSize: Int, - bits: Int, - mode: QuantizationMode - ) { - self.groupSize = groupSize - self.bits = bits - self.mode = mode - let transposedWeight = other.weight.swappedAxes(-1, -2) - let quantizedWeights = MLX.quantized( - transposedWeight, - groupSize: groupSize, - bits: bits, - mode: mode - ) - self.scales = quantizedWeights.scales - self.biases = quantizedWeights.biases - super.init( - numExperts: other.numExperts, - inputSize: other.inputSize, - outputSize: other.outputSize, - weight: quantizedWeights.wq, - bias: other.bias + private func privacyFilterSwiGLU( + _ input: MLXArray, + alpha: Float = 1.702, + limit: Float + ) -> MLXArray { + let half = input.dim(-1) / 2 + let glu = minimum(input[0..., 0..., 0.. MLXArray { - let inputShape = input.shape - let flatInput = input.reshaped(-1, 1, input.dim(-1)) - let flatIndices = expertIndices.reshaped(-1).asType(.int32) - var output = gatherQuantizedMM( - flatInput, - weight, - scales: scales, - biases: biases, - rhsIndices: flatIndices, - transpose: true, - groupSize: groupSize, - bits: bits, - mode: mode - ).squeezed(axis: -2) - output = output + bias.take(flatIndices, axis: 0) - return output.reshaped(Array(inputShape.dropLast()) + [outputSize]) - } -} - -private func privacyFilterLocalAttention( - query: MLXArray, - key: MLXArray, - value: MLXArray, - sinks: MLXArray, - leftContext: Int, - rightContext: Int, - attentionMask: MLXArray? -) -> MLXArray { - let batchSize = query.dim(0) - let numTokens = query.dim(1) - let numKVHeads = query.dim(2) - let queryMultiplier = query.dim(3) - let headDim = query.dim(4) - let window = leftContext + rightContext + 1 - let paddedTokens = numTokens + leftContext + rightContext - - let keyPadded = padded(key, widths: [0, [leftContext, rightContext], 0, 0]) - let valuePadded = padded(value, widths: [0, [leftContext, rightContext], 0, 0]) - let strides = [ - paddedTokens * numKVHeads * headDim, - numKVHeads * headDim, - numKVHeads * headDim, - headDim, - 1, - ] - let keyWindows = asStrided( - keyPadded, - [batchSize, numTokens, window, numKVHeads, headDim], - strides: strides - ) - let valueWindows = asStrided( - valuePadded, - [batchSize, numTokens, window, numKVHeads, headDim], - strides: strides - ) - - var scores = einsum("bthqd,btwhd->bthqw", query, keyWindows).asType(.float32) - let offsets = MLXArray.arange(window, dtype: .int32) - Int32(leftContext) - let positions = - MLXArray.arange(numTokens, dtype: .int32).expandedDimensions(axis: 1) - + offsets.expandedDimensions(axis: 0) - var valid = ((positions .>= 0) & (positions .< Int32(numTokens))) - .expandedDimensions(axes: [0, 2, 3]) - - if let attentionMask { - let maskPadded = padded( - attentionMask.asType(.bool), - widths: [0, [leftContext, rightContext]], - value: MLXArray(false) - ) - let maskWindows = asStrided( - maskPadded, - [batchSize, numTokens, window], - strides: [paddedTokens, 1, 1] - ) - valid = valid & maskWindows.expandedDimensions(axes: [2, 3]) - } + fileprivate init( + numExperts: Int, + inputSize: Int, + outputSize: Int, + weight: MLXArray, + bias: MLXArray + ) { + self.numExperts = numExperts + self.inputSize = inputSize + self.outputSize = outputSize + _weight.wrappedValue = weight + _bias.wrappedValue = bias + super.init() + } - scores = `where`( - valid, - scores, - privacyFilterScalar(-1.0e9, like: scores) - ) - - let sinkScores = (sinks * Foundation.log(2.0)).reshaped(numKVHeads, queryMultiplier) - let broadcastSinkScores = broadcast( - sinkScores.expandedDimensions(axes: [0, 1, 4]), - to: [batchSize, numTokens, numKVHeads, queryMultiplier, 1] - ) - let allWeights = softmax(concatenated([scores, broadcastSinkScores], axis: -1), axis: -1) - let weights = allWeights[0..., 0..., 0..., 0..., 0..bthqd", weights.asType(value.dtype), valueWindows) - return attention.reshaped(batchSize, numTokens, numKVHeads * queryMultiplier * headDim) -} - -private final class OpenMedPrivacyFilterAttentionBlock: Module { - private let headDim: Int - private let numAttentionHeads: Int - private let numKeyValueHeads: Int - private let queryMultiplier: Int - private let leftContext: Int - private let rightContext: Int - private let qkScale: Float - private let rope: OpenMedPrivacyFilterRotaryEmbedding - - @ModuleInfo(key: "norm") var norm: OpenMedPrivacyFilterRMSNorm - @ModuleInfo(key: "qkv") var qkv: Linear - @ModuleInfo(key: "out") var out: Linear - @ParameterInfo(key: "sinks") var sinks: MLXArray - - init(_ configuration: OpenMedMLXBertConfiguration) { - self.headDim = configuration.headDim - self.numAttentionHeads = configuration.numAttentionHeads - self.numKeyValueHeads = configuration.numKeyValueHeads - self.queryMultiplier = configuration.numAttentionHeads / configuration.numKeyValueHeads - self.leftContext = configuration.bidirectionalLeftContext - self.rightContext = configuration.bidirectionalRightContext - self.qkScale = 1.0 / sqrt(sqrt(Float(configuration.headDim))) - self.rope = OpenMedPrivacyFilterRotaryEmbedding(configuration) - let qkvSize = - configuration.headDim - * (configuration.numAttentionHeads + 2 * configuration.numKeyValueHeads) - - _norm.wrappedValue = OpenMedPrivacyFilterRMSNorm( - hiddenSize: configuration.hiddenSize, - eps: configuration.rmsNormEps - ) - _qkv.wrappedValue = Linear(configuration.hiddenSize, qkvSize) - _out.wrappedValue = Linear(configuration.headDim * configuration.numAttentionHeads, configuration.hiddenSize) - _sinks.wrappedValue = MLXArray.zeros([configuration.numAttentionHeads], type: Float.self) - super.init() + func callAsFunction(_ input: MLXArray, expertIndices: MLXArray) -> MLXArray { + let inputShape = input.shape + let flatInput = input.reshaped(-1, 1, input.dim(-1)).asType(weight.dtype) + let flatIndices = expertIndices.reshaped(-1).asType(.int32) + var output = gatherMM(flatInput, weight, rhsIndices: flatIndices).squeezed(axis: -2) + output = output + bias.take(flatIndices, axis: 0) + return output.reshaped(Array(inputShape.dropLast()) + [outputSize]) + } + + func toQuantized(groupSize: Int, bits: Int, mode: QuantizationMode) -> Module { + OpenMedPrivacyFilterQuantizedExpertLinear( + self, + groupSize: groupSize, + bits: bits, + mode: mode + ) + } } - func callAsFunction(_ input: MLXArray, attentionMask: MLXArray?) -> MLXArray { - let batchSize = input.dim(0) - let numTokens = input.dim(1) - let qkvStates = qkv(norm(input)) - let queryEnd = numAttentionHeads * headDim - let keyEnd = queryEnd + numKeyValueHeads * headDim - var query = qkvStates[0..., 0..., 0.. MLXArray { + let inputShape = input.shape + let flatInput = input.reshaped(-1, 1, input.dim(-1)) + let flatIndices = expertIndices.reshaped(-1).asType(.int32) + var output = gatherQuantizedMM( + flatInput, + weight, + scales: scales, + biases: biases, + rhsIndices: flatIndices, + transpose: true, + groupSize: groupSize, + bits: bits, + mode: mode + ).squeezed(axis: -2) + output = output + bias.take(flatIndices, axis: 0) + return output.reshaped(Array(inputShape.dropLast()) + [outputSize]) + } } -} - -private final class OpenMedPrivacyFilterMLPBlock: Module { - private let expertsPerToken: Int - private let swigluLimit: Float - - @ModuleInfo(key: "norm") var norm: OpenMedPrivacyFilterRMSNorm - @ModuleInfo(key: "gate") var gate: Linear - @ModuleInfo(key: "swiglu") var swiglu: OpenMedPrivacyFilterExpertLinear - @ModuleInfo(key: "out") var out: OpenMedPrivacyFilterExpertLinear - - init(_ configuration: OpenMedMLXBertConfiguration) { - self.expertsPerToken = configuration.expertsPerToken - self.swigluLimit = configuration.swigluLimit - let dtype = privacyFilterParameterDType(configuration) - _norm.wrappedValue = OpenMedPrivacyFilterRMSNorm( - hiddenSize: configuration.hiddenSize, - eps: configuration.rmsNormEps + + private func privacyFilterLocalAttention( + query: MLXArray, + key: MLXArray, + value: MLXArray, + sinks: MLXArray, + leftContext: Int, + rightContext: Int, + attentionMask: MLXArray? + ) -> MLXArray { + let batchSize = query.dim(0) + let numTokens = query.dim(1) + let numKVHeads = query.dim(2) + let queryMultiplier = query.dim(3) + let headDim = query.dim(4) + let window = leftContext + rightContext + 1 + let paddedTokens = numTokens + leftContext + rightContext + + let keyPadded = padded(key, widths: [0, [leftContext, rightContext], 0, 0]) + let valuePadded = padded(value, widths: [0, [leftContext, rightContext], 0, 0]) + let strides = [ + paddedTokens * numKVHeads * headDim, + numKVHeads * headDim, + numKVHeads * headDim, + headDim, + 1, + ] + let keyWindows = asStrided( + keyPadded, + [batchSize, numTokens, window, numKVHeads, headDim], + strides: strides ) - _gate.wrappedValue = Linear(configuration.hiddenSize, configuration.numExperts) - _swiglu.wrappedValue = OpenMedPrivacyFilterExpertLinear( - numExperts: configuration.numExperts, - inputSize: configuration.hiddenSize, - outputSize: configuration.intermediateSize * 2, - dtype: dtype + let valueWindows = asStrided( + valuePadded, + [batchSize, numTokens, window, numKVHeads, headDim], + strides: strides ) - _out.wrappedValue = OpenMedPrivacyFilterExpertLinear( - numExperts: configuration.numExperts, - inputSize: configuration.intermediateSize, - outputSize: configuration.hiddenSize, - dtype: dtype + + var scores = einsum("bthqd,btwhd->bthqw", query, keyWindows).asType(.float32) + let offsets = MLXArray.arange(window, dtype: .int32) - Int32(leftContext) + let positions = + MLXArray.arange(numTokens, dtype: .int32).expandedDimensions(axis: 1) + + offsets.expandedDimensions(axis: 0) + var valid = ((positions .>= 0) & (positions .< Int32(numTokens))) + .expandedDimensions(axes: [0, 2, 3]) + + if let attentionMask { + let maskPadded = padded( + attentionMask.asType(.bool), + widths: [0, [leftContext, rightContext]], + value: MLXArray(false) + ) + let maskWindows = asStrided( + maskPadded, + [batchSize, numTokens, window], + strides: [paddedTokens, 1, 1] + ) + valid = valid & maskWindows.expandedDimensions(axes: [2, 3]) + } + + scores = `where`( + valid, + scores, + privacyFilterScalar(-1.0e9, like: scores) ) - super.init() - } - func callAsFunction(_ input: MLXArray) -> MLXArray { - let batchShape = Array(input.shape.dropLast()) - let hiddenSize = input.dim(-1) - let normalized = norm(input).reshaped(-1, hiddenSize) - let gateLogits = gate(privacyFilterLinearInput(normalized, for: gate)).asType(.float32) - let (expertValues, expertIndices) = privacyFilterTopK(gateLogits, k: expertsPerToken) - let expertWeights = softmax(expertValues, axis: -1) / Float(expertsPerToken) - let expandedInput = broadcast( - privacyFilterExpertInput(normalized, for: swiglu).expandedDimensions(axis: 1), - to: [normalized.dim(0), expertsPerToken, hiddenSize] + let sinkScores = (sinks * Foundation.log(2.0)).reshaped(numKVHeads, queryMultiplier) + let broadcastSinkScores = broadcast( + sinkScores.expandedDimensions(axes: [0, 1, 4]), + to: [batchSize, numTokens, numKVHeads, queryMultiplier, 1] ) - var hidden = swiglu(expandedInput, expertIndices: expertIndices).asType(.float32) - hidden = privacyFilterSwiGLU(hidden, limit: swigluLimit) - let output = out( - privacyFilterExpertInput(hidden, for: out), - expertIndices: expertIndices - ).asType(.float32) - let mixed = - sum(output * expertWeights.expandedDimensions(axis: -1), axis: 1) - * Float(expertsPerToken) - return input + mixed.reshaped(batchShape + [hiddenSize]).asType(input.dtype) + let allWeights = softmax(concatenated([scores, broadcastSinkScores], axis: -1), axis: -1) + let weights = allWeights[0..., 0..., 0..., 0..., 0..bthqd", weights.asType(value.dtype), valueWindows) + return attention.reshaped(batchSize, numTokens, numKVHeads * queryMultiplier * headDim) } -} -private final class OpenMedPrivacyFilterTransformerBlock: Module { - @ModuleInfo(key: "attn") var attention: OpenMedPrivacyFilterAttentionBlock - @ModuleInfo(key: "mlp") var mlp: OpenMedPrivacyFilterMLPBlock + private final class OpenMedPrivacyFilterAttentionBlock: Module { + private let headDim: Int + private let numAttentionHeads: Int + private let numKeyValueHeads: Int + private let queryMultiplier: Int + private let leftContext: Int + private let rightContext: Int + private let qkScale: Float + private let rope: OpenMedPrivacyFilterRotaryEmbedding + + @ModuleInfo(key: "norm") var norm: OpenMedPrivacyFilterRMSNorm + @ModuleInfo(key: "qkv") var qkv: Linear + @ModuleInfo(key: "out") var out: Linear + @ParameterInfo(key: "sinks") var sinks: MLXArray + + init(_ configuration: OpenMedMLXBertConfiguration) { + self.headDim = configuration.headDim + self.numAttentionHeads = configuration.numAttentionHeads + self.numKeyValueHeads = configuration.numKeyValueHeads + self.queryMultiplier = configuration.numAttentionHeads / configuration.numKeyValueHeads + self.leftContext = configuration.bidirectionalLeftContext + self.rightContext = configuration.bidirectionalRightContext + self.qkScale = 1.0 / sqrt(sqrt(Float(configuration.headDim))) + self.rope = OpenMedPrivacyFilterRotaryEmbedding(configuration) + let qkvSize = + configuration.headDim + * (configuration.numAttentionHeads + 2 * configuration.numKeyValueHeads) + + _norm.wrappedValue = OpenMedPrivacyFilterRMSNorm( + hiddenSize: configuration.hiddenSize, + eps: configuration.rmsNormEps + ) + _qkv.wrappedValue = Linear(configuration.hiddenSize, qkvSize) + _out.wrappedValue = Linear(configuration.headDim * configuration.numAttentionHeads, configuration.hiddenSize) + _sinks.wrappedValue = MLXArray.zeros([configuration.numAttentionHeads], type: Float.self) + super.init() + } - init(_ configuration: OpenMedMLXBertConfiguration) { - _attention.wrappedValue = OpenMedPrivacyFilterAttentionBlock(configuration) - _mlp.wrappedValue = OpenMedPrivacyFilterMLPBlock(configuration) - super.init() + func callAsFunction(_ input: MLXArray, attentionMask: MLXArray?) -> MLXArray { + let batchSize = input.dim(0) + let numTokens = input.dim(1) + let qkvStates = qkv(norm(input)) + let queryEnd = numAttentionHeads * headDim + let keyEnd = queryEnd + numKeyValueHeads * headDim + var query = qkvStates[0..., 0..., 0.. MLXArray { - mlp(attention(input, attentionMask: attentionMask)) + private final class OpenMedPrivacyFilterMLPBlock: Module { + private let expertsPerToken: Int + private let swigluLimit: Float + + @ModuleInfo(key: "norm") var norm: OpenMedPrivacyFilterRMSNorm + @ModuleInfo(key: "gate") var gate: Linear + @ModuleInfo(key: "swiglu") var swiglu: OpenMedPrivacyFilterExpertLinear + @ModuleInfo(key: "out") var out: OpenMedPrivacyFilterExpertLinear + + init(_ configuration: OpenMedMLXBertConfiguration) { + self.expertsPerToken = configuration.expertsPerToken + self.swigluLimit = configuration.swigluLimit + let dtype = privacyFilterParameterDType(configuration) + _norm.wrappedValue = OpenMedPrivacyFilterRMSNorm( + hiddenSize: configuration.hiddenSize, + eps: configuration.rmsNormEps + ) + _gate.wrappedValue = Linear(configuration.hiddenSize, configuration.numExperts) + _swiglu.wrappedValue = OpenMedPrivacyFilterExpertLinear( + numExperts: configuration.numExperts, + inputSize: configuration.hiddenSize, + outputSize: configuration.intermediateSize * 2, + dtype: dtype + ) + _out.wrappedValue = OpenMedPrivacyFilterExpertLinear( + numExperts: configuration.numExperts, + inputSize: configuration.intermediateSize, + outputSize: configuration.hiddenSize, + dtype: dtype + ) + super.init() + } + + func callAsFunction(_ input: MLXArray) -> MLXArray { + let batchShape = Array(input.shape.dropLast()) + let hiddenSize = input.dim(-1) + let normalized = norm(input).reshaped(-1, hiddenSize) + let gateLogits = gate(privacyFilterLinearInput(normalized, for: gate)).asType(.float32) + let (expertValues, expertIndices) = privacyFilterTopK(gateLogits, k: expertsPerToken) + let expertWeights = softmax(expertValues, axis: -1) / Float(expertsPerToken) + let expandedInput = broadcast( + privacyFilterExpertInput(normalized, for: swiglu).expandedDimensions(axis: 1), + to: [normalized.dim(0), expertsPerToken, hiddenSize] + ) + var hidden = swiglu(expandedInput, expertIndices: expertIndices).asType(.float32) + hidden = privacyFilterSwiGLU(hidden, limit: swigluLimit) + let output = out( + privacyFilterExpertInput(hidden, for: out), + expertIndices: expertIndices + ).asType(.float32) + let mixed = + sum(output * expertWeights.expandedDimensions(axis: -1), axis: 1) + * Float(expertsPerToken) + return input + mixed.reshaped(batchShape + [hiddenSize]).asType(input.dtype) + } } -} -final class OpenMedPrivacyFilterForTokenClassification: Module { - @ModuleInfo(key: "embedding") var embedding: Embedding - @ModuleInfo(key: "block") fileprivate var block: [OpenMedPrivacyFilterTransformerBlock] - @ModuleInfo(key: "norm") fileprivate var norm: OpenMedPrivacyFilterRMSNorm - @ModuleInfo(key: "unembedding") var unembedding: Linear + private final class OpenMedPrivacyFilterTransformerBlock: Module { + @ModuleInfo(key: "attn") var attention: OpenMedPrivacyFilterAttentionBlock + @ModuleInfo(key: "mlp") var mlp: OpenMedPrivacyFilterMLPBlock - let configuration: OpenMedMLXBertConfiguration + init(_ configuration: OpenMedMLXBertConfiguration) { + _attention.wrappedValue = OpenMedPrivacyFilterAttentionBlock(configuration) + _mlp.wrappedValue = OpenMedPrivacyFilterMLPBlock(configuration) + super.init() + } - init(_ configuration: OpenMedMLXBertConfiguration) { - self.configuration = configuration - _embedding.wrappedValue = Embedding( - embeddingCount: configuration.vocabularySize, - dimensions: configuration.hiddenSize - ) - _block.wrappedValue = (0.. MLXArray { + mlp(attention(input, attentionMask: attentionMask)) } - _norm.wrappedValue = OpenMedPrivacyFilterRMSNorm( - hiddenSize: configuration.hiddenSize, - eps: configuration.rmsNormEps - ) - // The original openai/privacy-filter has a bias-less classifier head; - // the Nemotron-PII fine-tunes (`classifier_bias: true`) ship with a - // learned bias. Honor whichever the config requests. - _unembedding.wrappedValue = Linear( - configuration.hiddenSize, - configuration.numLabels, - bias: configuration.classifierBias - ) - super.init() } - func callAsFunction( - _ inputIDs: MLXArray, - attentionMask: MLXArray? = nil - ) -> MLXArray { - precondition(inputIDs.ndim == 2, "Privacy Filter expects input IDs with shape [batch, tokens].") - let resolvedAttentionMask = attentionMask?.asType(.bool) - var hiddenStates = embedding(inputIDs) - for layer in block { - hiddenStates = layer(hiddenStates, attentionMask: resolvedAttentionMask) + final class OpenMedPrivacyFilterForTokenClassification: Module { + @ModuleInfo(key: "embedding") var embedding: Embedding + @ModuleInfo(key: "block") fileprivate var block: [OpenMedPrivacyFilterTransformerBlock] + @ModuleInfo(key: "norm") fileprivate var norm: OpenMedPrivacyFilterRMSNorm + @ModuleInfo(key: "unembedding") var unembedding: Linear + + let configuration: OpenMedMLXBertConfiguration + + init(_ configuration: OpenMedMLXBertConfiguration) { + self.configuration = configuration + _embedding.wrappedValue = Embedding( + embeddingCount: configuration.vocabularySize, + dimensions: configuration.hiddenSize + ) + _block.wrappedValue = (0.. [String: MLXArray] { - weights.filter { key, _ in !key.hasPrefix("_") } - } + func callAsFunction( + _ inputIDs: MLXArray, + attentionMask: MLXArray? = nil + ) -> MLXArray { + precondition(inputIDs.ndim == 2, "Privacy Filter expects input IDs with shape [batch, tokens].") + let resolvedAttentionMask = attentionMask?.asType(.bool) + var hiddenStates = embedding(inputIDs) + for layer in block { + hiddenStates = layer(hiddenStates, attentionMask: resolvedAttentionMask) + } + hiddenStates = norm(hiddenStates) + return unembedding(privacyFilterLinearInput(hiddenStates, for: unembedding)) + } + + func sanitize(weights: [String: MLXArray]) -> [String: MLXArray] { + weights.filter { key, _ in !key.hasPrefix("_") } + } - /// Install quantized-module placeholders for every path flagged by - /// ``hasScales``. Custom MoE expert modules use a shape-only placeholder - /// (no `MLX.quantized` on dummy zero weights); standard `Linear`/`Embedding` - /// layers fall through to MLX's default `quantizeSingle`. Real tensors are - /// loaded afterwards via `update(parameters:)`. - func installQuantizedPlaceholders( - where hasScales: (String) -> Bool, - groupSize: Int, - bits: Int, - mode: QuantizationMode - ) { - quantize( - model: self, - filter: { path, _ in - hasScales(path) ? (groupSize, bits, mode) : nil - }, - apply: { layer, groupSize, bits, mode in - if let expert = layer as? OpenMedPrivacyFilterExpertLinear, - !(expert is OpenMedPrivacyFilterQuantizedExpertLinear) - { - return OpenMedPrivacyFilterQuantizedExpertLinear( - numExperts: expert.numExperts, - inputSize: expert.inputSize, - outputSize: expert.outputSize, + /// Install quantized-module placeholders for every path flagged by + /// ``hasScales``. Custom MoE expert modules use a shape-only placeholder + /// (no `MLX.quantized` on dummy zero weights); standard `Linear`/`Embedding` + /// layers fall through to MLX's default `quantizeSingle`. Real tensors are + /// loaded afterwards via `update(parameters:)`. + func installQuantizedPlaceholders( + where hasScales: (String) -> Bool, + groupSize: Int, + bits: Int, + mode: QuantizationMode + ) { + quantize( + model: self, + filter: { path, _ in + hasScales(path) ? (groupSize, bits, mode) : nil + }, + apply: { layer, groupSize, bits, mode in + if let expert = layer as? OpenMedPrivacyFilterExpertLinear, + !(expert is OpenMedPrivacyFilterQuantizedExpertLinear) + { + return OpenMedPrivacyFilterQuantizedExpertLinear( + numExperts: expert.numExperts, + inputSize: expert.inputSize, + outputSize: expert.outputSize, + groupSize: groupSize, + bits: bits, + mode: mode, + dtype: expert.weight.dtype + ) + } + return quantizeSingle( + layer: layer, groupSize: groupSize, bits: bits, - mode: mode, - dtype: expert.weight.dtype + mode: mode ) } - return quantizeSingle( - layer: layer, - groupSize: groupSize, - bits: bits, - mode: mode - ) - } - ) + ) + } } -} +#endif diff --git a/swift/OpenMedKit/Sources/OpenMedKit/OpenMedPrivacyFilterPipeline.swift b/swift/OpenMedKit/Sources/OpenMedKit/OpenMedPrivacyFilterPipeline.swift index 92bd5ce5e..e1512378b 100644 --- a/swift/OpenMedKit/Sources/OpenMedKit/OpenMedPrivacyFilterPipeline.swift +++ b/swift/OpenMedKit/Sources/OpenMedKit/OpenMedPrivacyFilterPipeline.swift @@ -1,919 +1,921 @@ -import Foundation -import MLX - -enum OpenMedPrivacyFilterError: LocalizedError { - case unsupportedArtifact(String) - case missingTokenizer(URL) - case invalidTokenizer(URL) - case unknownToken(String) - - var errorDescription: String? { - switch self { - case .unsupportedArtifact(let family): - return "Expected OpenAI Privacy Filter MLX artifact, got \(family)." - case .missingTokenizer(let url): - return "Privacy Filter tokenizer.json was not found at \(url.path)." - case .invalidTokenizer(let url): - return "Privacy Filter tokenizer.json is not a supported byte-level BPE tokenizer: \(url.path)." - case .unknownToken(let token): - return "Privacy Filter tokenizer could not map token piece: \(token)." +#if canImport(MLX) && !os(watchOS) && !os(visionOS) + import Foundation + import MLX + + enum OpenMedPrivacyFilterError: LocalizedError { + case unsupportedArtifact(String) + case missingTokenizer(URL) + case invalidTokenizer(URL) + case unknownToken(String) + + var errorDescription: String? { + switch self { + case .unsupportedArtifact(let family): + return "Expected OpenAI Privacy Filter MLX artifact, got \(family)." + case .missingTokenizer(let url): + return "Privacy Filter tokenizer.json was not found at \(url.path)." + case .invalidTokenizer(let url): + return "Privacy Filter tokenizer.json is not a supported byte-level BPE tokenizer: \(url.path)." + case .unknownToken(let token): + return "Privacy Filter tokenizer could not map token piece: \(token)." + } } } -} - -struct OpenMedPrivacyFilterEncodedText { - let tokenIDs: [Int] - let charStarts: [Int] - let charEnds: [Int] - let decodedText: String -} - -final class OpenMedPrivacyFilterTokenizer { - private struct TokenizerJSON: Decodable { - struct Model: Decodable { - let type: String - let vocab: [String: Int] - let merges: [MergeEntry] - } - struct PreTokenizer: Decodable { - struct SplitPattern: Decodable { - let regex: String? + struct OpenMedPrivacyFilterEncodedText { + let tokenIDs: [Int] + let charStarts: [Int] + let charEnds: [Int] + let decodedText: String + } - enum CodingKeys: String, CodingKey { - case regex = "Regex" - } + final class OpenMedPrivacyFilterTokenizer { + private struct TokenizerJSON: Decodable { + struct Model: Decodable { + let type: String + let vocab: [String: Int] + let merges: [MergeEntry] } - struct Entry: Decodable { + struct PreTokenizer: Decodable { + struct SplitPattern: Decodable { + let regex: String? + + enum CodingKeys: String, CodingKey { + case regex = "Regex" + } + } + + struct Entry: Decodable { + let type: String + let pattern: SplitPattern? + let pretokenizers: [Entry]? + } + let type: String let pattern: SplitPattern? let pretokenizers: [Entry]? } - let type: String - let pattern: SplitPattern? - let pretokenizers: [Entry]? - } + let model: Model + let preTokenizer: PreTokenizer? - let model: Model - let preTokenizer: PreTokenizer? - - enum CodingKeys: String, CodingKey { - case model - case preTokenizer = "pre_tokenizer" + enum CodingKeys: String, CodingKey { + case model + case preTokenizer = "pre_tokenizer" + } } - } - private enum MergeEntry: Decodable { - case pair(String, String) + private enum MergeEntry: Decodable { + case pair(String, String) - init(from decoder: Decoder) throws { - let container = try decoder.singleValueContainer() - if let parts = try? container.decode([String].self), parts.count == 2 { + init(from decoder: Decoder) throws { + let container = try decoder.singleValueContainer() + if let parts = try? container.decode([String].self), parts.count == 2 { + self = .pair(parts[0], parts[1]) + return + } + let text = try container.decode(String.self) + let parts = text.split(separator: " ", maxSplits: 1).map(String.init) + guard parts.count == 2 else { + throw DecodingError.dataCorruptedError( + in: container, + debugDescription: "Expected BPE merge pair" + ) + } self = .pair(parts[0], parts[1]) - return } - let text = try container.decode(String.self) - let parts = text.split(separator: " ", maxSplits: 1).map(String.init) - guard parts.count == 2 else { - throw DecodingError.dataCorruptedError( - in: container, - debugDescription: "Expected BPE merge pair" - ) - } - self = .pair(parts[0], parts[1]) } - } - private struct PairKey: Hashable { - let left: String - let right: String - } - - private let vocab: [String: Int] - private let idToToken: [Int: String] - private let mergeRanks: [PairKey: Int] - private let regex: NSRegularExpression? - private let byteEncoder: [UInt8: String] - private let byteDecoder: [UnicodeScalar: UInt8] - private var bpeCache = [String: [String]]() - - init(directoryURL: URL) throws { - let tokenizerURL = directoryURL.appending(path: "tokenizer.json") - guard FileManager.default.fileExists(atPath: tokenizerURL.path) else { - throw OpenMedPrivacyFilterError.missingTokenizer(tokenizerURL) + private struct PairKey: Hashable { + let left: String + let right: String } - let data = try Data(contentsOf: tokenizerURL) - let tokenizer = try JSONDecoder().decode(TokenizerJSON.self, from: data) - guard tokenizer.model.type.lowercased() == "bpe" else { - throw OpenMedPrivacyFilterError.invalidTokenizer(tokenizerURL) - } + private let vocab: [String: Int] + private let idToToken: [Int: String] + private let mergeRanks: [PairKey: Int] + private let regex: NSRegularExpression? + private let byteEncoder: [UInt8: String] + private let byteDecoder: [UnicodeScalar: UInt8] + private var bpeCache = [String: [String]]() - self.vocab = tokenizer.model.vocab - self.idToToken = Dictionary(uniqueKeysWithValues: tokenizer.model.vocab.map { ($0.value, $0.key) }) - self.mergeRanks = Dictionary( - uniqueKeysWithValues: tokenizer.model.merges.enumerated().map { rank, entry in - switch entry { - case .pair(let left, let right): - return (PairKey(left: left, right: right), rank) - } + init(directoryURL: URL) throws { + let tokenizerURL = directoryURL.appending(path: "tokenizer.json") + guard FileManager.default.fileExists(atPath: tokenizerURL.path) else { + throw OpenMedPrivacyFilterError.missingTokenizer(tokenizerURL) } - ) - let byteMaps = Self.makeByteMaps() - self.byteEncoder = byteMaps.encoder - self.byteDecoder = byteMaps.decoder + let data = try Data(contentsOf: tokenizerURL) + let tokenizer = try JSONDecoder().decode(TokenizerJSON.self, from: data) + guard tokenizer.model.type.lowercased() == "bpe" else { + throw OpenMedPrivacyFilterError.invalidTokenizer(tokenizerURL) + } + + self.vocab = tokenizer.model.vocab + self.idToToken = Dictionary(uniqueKeysWithValues: tokenizer.model.vocab.map { ($0.value, $0.key) }) + self.mergeRanks = Dictionary( + uniqueKeysWithValues: tokenizer.model.merges.enumerated().map { rank, entry in + switch entry { + case .pair(let left, let right): + return (PairKey(left: left, right: right), rank) + } + } + ) + + let byteMaps = Self.makeByteMaps() + self.byteEncoder = byteMaps.encoder + self.byteDecoder = byteMaps.decoder - if let pattern = Self.findRegexPattern(in: tokenizer.preTokenizer) { - self.regex = try? NSRegularExpression(pattern: pattern) - } else { - self.regex = nil + if let pattern = Self.findRegexPattern(in: tokenizer.preTokenizer) { + self.regex = try? NSRegularExpression(pattern: pattern) + } else { + self.regex = nil + } } - } - func encode(_ text: String, maxTokens: Int) throws -> OpenMedPrivacyFilterEncodedText { - var tokenIDs = [Int]() - tokenIDs.reserveCapacity(min(maxTokens, max(8, text.count / 3))) + func encode(_ text: String, maxTokens: Int) throws -> OpenMedPrivacyFilterEncodedText { + var tokenIDs = [Int]() + tokenIDs.reserveCapacity(min(maxTokens, max(8, text.count / 3))) - for piece in preTokenize(text) { - let byteLevelPiece = piece.utf8.map { byteEncoder[$0, default: ""] }.joined() - for token in bytePairEncode(byteLevelPiece) { - guard let id = vocab[token] else { - throw OpenMedPrivacyFilterError.unknownToken(token) + for piece in preTokenize(text) { + let byteLevelPiece = piece.utf8.map { byteEncoder[$0, default: ""] }.joined() + for token in bytePairEncode(byteLevelPiece) { + guard let id = vocab[token] else { + throw OpenMedPrivacyFilterError.unknownToken(token) + } + tokenIDs.append(id) + if tokenIDs.count >= maxTokens { + break + } } - tokenIDs.append(id) if tokenIDs.count >= maxTokens { break } } - if tokenIDs.count >= maxTokens { - break - } - } - - let offsets = decodeOffsets(tokenIDs) - return OpenMedPrivacyFilterEncodedText( - tokenIDs: tokenIDs, - charStarts: offsets.charStarts, - charEnds: offsets.charEnds, - decodedText: offsets.text - ) - } - func tokenBytes(tokenID: Int) -> [UInt8] { - guard let token = idToToken[tokenID] else { - return [] + let offsets = decodeOffsets(tokenIDs) + return OpenMedPrivacyFilterEncodedText( + tokenIDs: tokenIDs, + charStarts: offsets.charStarts, + charEnds: offsets.charEnds, + decodedText: offsets.text + ) } - var bytes = [UInt8]() - for scalar in token.unicodeScalars { - if let byte = byteDecoder[scalar] { - bytes.append(byte) - } else { - bytes.append(contentsOf: String(scalar).utf8) + + func tokenBytes(tokenID: Int) -> [UInt8] { + guard let token = idToToken[tokenID] else { + return [] } + var bytes = [UInt8]() + for scalar in token.unicodeScalars { + if let byte = byteDecoder[scalar] { + bytes.append(byte) + } else { + bytes.append(contentsOf: String(scalar).utf8) + } + } + return bytes } - return bytes - } - private static func makeByteMaps() -> ( - encoder: [UInt8: String], - decoder: [UnicodeScalar: UInt8] - ) { - var bytes = Array(33...126) + Array(161...172) + Array(174...255) - var scalars = bytes - var byteSet = Set(bytes) - var next = 0 - for byte in 0..<256 where !byteSet.contains(byte) { - bytes.append(byte) - scalars.append(256 + next) - byteSet.insert(byte) - next += 1 - } + private static func makeByteMaps() -> ( + encoder: [UInt8: String], + decoder: [UnicodeScalar: UInt8] + ) { + var bytes = Array(33...126) + Array(161...172) + Array(174...255) + var scalars = bytes + var byteSet = Set(bytes) + var next = 0 + for byte in 0..<256 where !byteSet.contains(byte) { + bytes.append(byte) + scalars.append(256 + next) + byteSet.insert(byte) + next += 1 + } - var encoder = [UInt8: String]() - var decoder = [UnicodeScalar: UInt8]() - for (byte, scalarValue) in zip(bytes, scalars) { - guard let scalar = UnicodeScalar(scalarValue) else { - continue + var encoder = [UInt8: String]() + var decoder = [UnicodeScalar: UInt8]() + for (byte, scalarValue) in zip(bytes, scalars) { + guard let scalar = UnicodeScalar(scalarValue) else { + continue + } + encoder[UInt8(byte)] = String(scalar) + decoder[scalar] = UInt8(byte) } - encoder[UInt8(byte)] = String(scalar) - decoder[scalar] = UInt8(byte) + return (encoder, decoder) } - return (encoder, decoder) - } - private static func findRegexPattern(in preTokenizer: TokenizerJSON.PreTokenizer?) -> String? { - guard let preTokenizer else { - return nil - } - if preTokenizer.type == "Split", let regex = preTokenizer.pattern?.regex { - return regex - } - for entry in preTokenizer.pretokenizers ?? [] { - if entry.type == "Split", let regex = entry.pattern?.regex { + private static func findRegexPattern(in preTokenizer: TokenizerJSON.PreTokenizer?) -> String? { + guard let preTokenizer else { + return nil + } + if preTokenizer.type == "Split", let regex = preTokenizer.pattern?.regex { return regex } + for entry in preTokenizer.pretokenizers ?? [] { + if entry.type == "Split", let regex = entry.pattern?.regex { + return regex + } + } + return nil } - return nil - } - private func preTokenize(_ text: String) -> [String] { - guard let regex else { - return text.isEmpty ? [] : [text] - } + private func preTokenize(_ text: String) -> [String] { + guard let regex else { + return text.isEmpty ? [] : [text] + } - let fullRange = NSRange(text.startIndex.. [String] { - if let cached = bpeCache[token] { - return cached - } - var word = token.map(String.init) - guard word.count > 1 else { + private func bytePairEncode(_ token: String) -> [String] { + if let cached = bpeCache[token] { + return cached + } + var word = token.map(String.init) + guard word.count > 1 else { + bpeCache[token] = word + return word + } + + while word.count > 1 { + var bestRank = Int.max + var bestPair: PairKey? + for index in 0..<(word.count - 1) { + let pair = PairKey(left: word[index], right: word[index + 1]) + if let rank = mergeRanks[pair], rank < bestRank { + bestRank = rank + bestPair = pair + } + } + guard let bestPair else { + break + } + + var merged = [String]() + var index = 0 + while index < word.count { + if index < word.count - 1, + word[index] == bestPair.left, + word[index + 1] == bestPair.right + { + merged.append(bestPair.left + bestPair.right) + index += 2 + } else { + merged.append(word[index]) + index += 1 + } + } + word = merged + } + bpeCache[token] = word return word } - while word.count > 1 { - var bestRank = Int.max - var bestPair: PairKey? - for index in 0..<(word.count - 1) { - let pair = PairKey(left: word[index], right: word[index + 1]) - if let rank = mergeRanks[pair], rank < bestRank { - bestRank = rank - bestPair = pair - } - } - guard let bestPair else { - break - } - - var merged = [String]() - var index = 0 - while index < word.count { - if index < word.count - 1, - word[index] == bestPair.left, - word[index + 1] == bestPair.right - { - merged.append(bestPair.left + bestPair.right) - index += 2 - } else { - merged.append(word[index]) - index += 1 + private func decodeOffsets(_ tokenIDs: [Int]) -> ( + text: String, + charStarts: [Int], + charEnds: [Int] + ) { + let decodedTokenBytes = tokenIDs.map { self.tokenBytes(tokenID: $0) } + let allBytes = decodedTokenBytes.flatMap { $0 } + let decodedText = String(decoding: allBytes, as: UTF8.self) + + var charByteStarts = [Int]() + var charByteEnds = [Int]() + var byteCursor = 0 + for character in decodedText { + charByteStarts.append(byteCursor) + byteCursor += String(character).utf8.count + charByteEnds.append(byteCursor) + } + + var charStarts = [Int]() + var charEnds = [Int]() + var tokenByteCursor = 0 + for rawBytes in decodedTokenBytes { + let tokenByteStart = tokenByteCursor + let tokenByteEnd = tokenByteStart + rawBytes.count + tokenByteCursor = tokenByteEnd + let start = Self.bisectRight(charByteEnds, value: tokenByteStart) + var end = Self.bisectLeft(charByteStarts, value: tokenByteEnd) + if end < start { + end = start } + charStarts.append(start) + charEnds.append(end) } - word = merged - } - - bpeCache[token] = word - return word - } - private func decodeOffsets(_ tokenIDs: [Int]) -> ( - text: String, - charStarts: [Int], - charEnds: [Int] - ) { - let decodedTokenBytes = tokenIDs.map { self.tokenBytes(tokenID: $0) } - let allBytes = decodedTokenBytes.flatMap { $0 } - let decodedText = String(decoding: allBytes, as: UTF8.self) - - var charByteStarts = [Int]() - var charByteEnds = [Int]() - var byteCursor = 0 - for character in decodedText { - charByteStarts.append(byteCursor) - byteCursor += String(character).utf8.count - charByteEnds.append(byteCursor) + return (decodedText, charStarts, charEnds) } - var charStarts = [Int]() - var charEnds = [Int]() - var tokenByteCursor = 0 - for rawBytes in decodedTokenBytes { - let tokenByteStart = tokenByteCursor - let tokenByteEnd = tokenByteStart + rawBytes.count - tokenByteCursor = tokenByteEnd - let start = Self.bisectRight(charByteEnds, value: tokenByteStart) - var end = Self.bisectLeft(charByteStarts, value: tokenByteEnd) - if end < start { - end = start - } - charStarts.append(start) - charEnds.append(end) - } - - return (decodedText, charStarts, charEnds) - } - - private static func bisectLeft(_ values: [Int], value: Int) -> Int { - var low = 0 - var high = values.count - while low < high { - let mid = (low + high) / 2 - if values[mid] < value { - low = mid + 1 - } else { - high = mid + private static func bisectLeft(_ values: [Int], value: Int) -> Int { + var low = 0 + var high = values.count + while low < high { + let mid = (low + high) / 2 + if values[mid] < value { + low = mid + 1 + } else { + high = mid + } } + return low } - return low - } - private static func bisectRight(_ values: [Int], value: Int) -> Int { - var low = 0 - var high = values.count - while low < high { - let mid = (low + high) / 2 - if values[mid] <= value { - low = mid + 1 - } else { - high = mid + private static func bisectRight(_ values: [Int], value: Int) -> Int { + var low = 0 + var high = values.count + while low < high { + let mid = (low + high) / 2 + if values[mid] <= value { + low = mid + 1 + } else { + high = mid + } } + return low } - return low } -} - -struct OpenMedPrivacyFilterLabelInfo { - let spanClassNames: [String] - let tokenToSpanLabel: [Int: Int] - let tokenBoundaryTags: [Int: String] - let backgroundTokenLabel: Int - let backgroundSpanLabel: Int - - init(id2label: [Int: String]) { - var spanClassNames = ["O"] - var spanLabelLookup = ["O": 0] - var tokenToSpanLabel = [Int: Int]() - var tokenBoundaryTags = [Int: String]() - var backgroundTokenLabel = 0 - - for index in id2label.keys.sorted() { - let label = id2label[index] ?? "O" - if label == "O" { - backgroundTokenLabel = index - tokenToSpanLabel[index] = 0 - continue - } - - let split = Self.splitBoundaryLabel(label) - let spanLabel: Int - if let existing = spanLabelLookup[split.baseLabel] { - spanLabel = existing - } else { - spanLabel = spanClassNames.count - spanClassNames.append(split.baseLabel) - spanLabelLookup[split.baseLabel] = spanLabel + + struct OpenMedPrivacyFilterLabelInfo { + let spanClassNames: [String] + let tokenToSpanLabel: [Int: Int] + let tokenBoundaryTags: [Int: String] + let backgroundTokenLabel: Int + let backgroundSpanLabel: Int + + init(id2label: [Int: String]) { + var spanClassNames = ["O"] + var spanLabelLookup = ["O": 0] + var tokenToSpanLabel = [Int: Int]() + var tokenBoundaryTags = [Int: String]() + var backgroundTokenLabel = 0 + + for index in id2label.keys.sorted() { + let label = id2label[index] ?? "O" + if label == "O" { + backgroundTokenLabel = index + tokenToSpanLabel[index] = 0 + continue + } + + let split = Self.splitBoundaryLabel(label) + let spanLabel: Int + if let existing = spanLabelLookup[split.baseLabel] { + spanLabel = existing + } else { + spanLabel = spanClassNames.count + spanClassNames.append(split.baseLabel) + spanLabelLookup[split.baseLabel] = spanLabel + } + tokenToSpanLabel[index] = spanLabel + tokenBoundaryTags[index] = split.boundary } - tokenToSpanLabel[index] = spanLabel - tokenBoundaryTags[index] = split.boundary - } - self.spanClassNames = spanClassNames - self.tokenToSpanLabel = tokenToSpanLabel - self.tokenBoundaryTags = tokenBoundaryTags - self.backgroundTokenLabel = backgroundTokenLabel - self.backgroundSpanLabel = 0 - } + self.spanClassNames = spanClassNames + self.tokenToSpanLabel = tokenToSpanLabel + self.tokenBoundaryTags = tokenBoundaryTags + self.backgroundTokenLabel = backgroundTokenLabel + self.backgroundSpanLabel = 0 + } - private static func splitBoundaryLabel(_ label: String) -> (boundary: String, baseLabel: String) { - guard label.count > 2 else { + private static func splitBoundaryLabel(_ label: String) -> (boundary: String, baseLabel: String) { + guard label.count > 2 else { + return ("B", label) + } + let boundary = String(label.prefix(1)) + let separatorIndex = label.index(label.startIndex, offsetBy: 1) + let baseIndex = label.index(label.startIndex, offsetBy: 2) + if label[separatorIndex] == "-", ["B", "I", "E", "S"].contains(boundary) { + return (boundary, String(label[baseIndex...])) + } return ("B", label) } - let boundary = String(label.prefix(1)) - let separatorIndex = label.index(label.startIndex, offsetBy: 1) - let baseIndex = label.index(label.startIndex, offsetBy: 2) - if label[separatorIndex] == "-", ["B", "I", "E", "S"].contains(boundary) { - return (boundary, String(label[baseIndex...])) - } - return ("B", label) } -} - -enum OpenMedPrivacyFilterViterbi { - private static let negativeInfinity: Float = -1.0e9 - private static let biasKeys = [ - "transition_bias_background_stay", - "transition_bias_background_to_start", - "transition_bias_inside_to_continue", - "transition_bias_inside_to_end", - "transition_bias_end_to_background", - "transition_bias_end_to_start", - ] - - static func decode( - tokenLogProbabilities: [[Float]], - labelInfo: OpenMedPrivacyFilterLabelInfo, - biases: [String: Float] - ) -> [Int] { - guard !tokenLogProbabilities.isEmpty else { - return [] - } - var resolvedBiases = Dictionary(uniqueKeysWithValues: biasKeys.map { ($0, Float(0.0)) }) - for (key, value) in biases where resolvedBiases[key] != nil { - resolvedBiases[key] = value - } + enum OpenMedPrivacyFilterViterbi { + private static let negativeInfinity: Float = -1.0e9 + private static let biasKeys = [ + "transition_bias_background_stay", + "transition_bias_background_to_start", + "transition_bias_inside_to_continue", + "transition_bias_inside_to_end", + "transition_bias_end_to_background", + "transition_bias_end_to_start", + ] - let scores = buildScores(labelInfo: labelInfo, biases: resolvedBiases) - let numClasses = labelInfo.tokenToSpanLabel.count - var currentScores = (0.. bestScore { - bestScore = score - bestIndex = previousIndex + static func decode( + tokenLogProbabilities: [[Float]], + labelInfo: OpenMedPrivacyFilterLabelInfo, + biases: [String: Float] + ) -> [Int] { + guard !tokenLogProbabilities.isEmpty else { + return [] + } + + var resolvedBiases = Dictionary(uniqueKeysWithValues: biasKeys.map { ($0, Float(0.0)) }) + for (key, value) in biases where resolvedBiases[key] != nil { + resolvedBiases[key] = value + } + + let scores = buildScores(labelInfo: labelInfo, biases: resolvedBiases) + let numClasses = labelInfo.tokenToSpanLabel.count + var currentScores = (0.. bestScore { + bestScore = score + bestIndex = previousIndex + } } + nextScores.append(bestScore + tokenScores[nextIndex]) + paths.append(bestIndex) } - nextScores.append(bestScore + tokenScores[nextIndex]) - paths.append(bestIndex) + currentScores = nextScores + backpointers.append(paths) } - currentScores = nextScores - backpointers.append(paths) - } - let finalScores = currentScores.enumerated().map { index, score in - score + scores.end[index] - } - guard finalScores.contains(where: { $0.isFinite }) else { - return tokenLogProbabilities.map { row in - row.enumerated().max(by: { $0.element < $1.element })?.offset ?? 0 + let finalScores = currentScores.enumerated().map { index, score in + score + scores.end[index] + } + guard finalScores.contains(where: { $0.isFinite }) else { + return tokenLogProbabilities.map { row in + row.enumerated().max(by: { $0.element < $1.element })?.offset ?? 0 + } } - } - var label = finalScores.enumerated().max(by: { $0.element < $1.element })?.offset ?? 0 - var path = [label] - for paths in backpointers.reversed() { - label = paths[label] - path.append(label) - } - return Array(path.reversed()) - } + var label = finalScores.enumerated().max(by: { $0.element < $1.element })?.offset ?? 0 + var path = [label] + for paths in backpointers.reversed() { + label = paths[label] + path.append(label) + } + return Array(path.reversed()) + } + + private static func buildScores( + labelInfo: OpenMedPrivacyFilterLabelInfo, + biases: [String: Float] + ) -> ( + start: [Float], + end: [Float], + transition: [[Float]] + ) { + let numClasses = labelInfo.tokenToSpanLabel.count + var startScores = Array(repeating: negativeInfinity, count: numClasses) + var endScores = Array(repeating: negativeInfinity, count: numClasses) + var transitionScores = Array( + repeating: Array(repeating: negativeInfinity, count: numClasses), + count: numClasses + ) + + for previousIndex in 0.. ( - start: [Float], - end: [Float], - transition: [[Float]] - ) { - let numClasses = labelInfo.tokenToSpanLabel.count - var startScores = Array(repeating: negativeInfinity, count: numClasses) - var endScores = Array(repeating: negativeInfinity, count: numClasses) - var transitionScores = Array( - repeating: Array(repeating: negativeInfinity, count: numClasses), - count: numClasses - ) - - for previousIndex in 0.. Bool { + let nextIsBackground = + nextSpan == labelInfo.backgroundSpanLabel || nextIndex == labelInfo.backgroundTokenLabel + if (nextSpan == nil || nextTag == nil) && !nextIsBackground { + return false + } + + guard let previousSpan, let previousTag else { + return nextIsBackground || nextTag == "B" || nextTag == "S" + } - private static func isValidTransition( - previousTag: String?, - previousSpan: Int?, - nextTag: String?, - nextSpan: Int?, - labelInfo: OpenMedPrivacyFilterLabelInfo, - nextIndex: Int - ) -> Bool { - let nextIsBackground = - nextSpan == labelInfo.backgroundSpanLabel || nextIndex == labelInfo.backgroundTokenLabel - if (nextSpan == nil || nextTag == nil) && !nextIsBackground { + if previousSpan == labelInfo.backgroundSpanLabel { + return nextIsBackground || nextTag == "B" || nextTag == "S" + } + if previousTag == "E" || previousTag == "S" { + return nextIsBackground || nextTag == "B" || nextTag == "S" + } + if previousTag == "B" || previousTag == "I" { + return previousSpan == nextSpan && (nextTag == "I" || nextTag == "E") + } return false } - guard let previousSpan, let previousTag else { - return nextIsBackground || nextTag == "B" || nextTag == "S" - } + private static func transitionBias( + previousTag: String?, + previousSpan: Int?, + nextTag: String?, + nextSpan: Int?, + labelInfo: OpenMedPrivacyFilterLabelInfo, + previousIndex: Int, + nextIndex: Int, + biases: [String: Float] + ) -> Float { + let previousIsBackground = + previousSpan == labelInfo.backgroundSpanLabel || previousIndex == labelInfo.backgroundTokenLabel + let nextIsBackground = + nextSpan == labelInfo.backgroundSpanLabel || nextIndex == labelInfo.backgroundTokenLabel + + if previousIsBackground { + if nextIsBackground { + return biases["transition_bias_background_stay"] ?? 0.0 + } + if nextTag == "B" || nextTag == "S" { + return biases["transition_bias_background_to_start"] ?? 0.0 + } + return 0.0 + } - if previousSpan == labelInfo.backgroundSpanLabel { - return nextIsBackground || nextTag == "B" || nextTag == "S" - } - if previousTag == "E" || previousTag == "S" { - return nextIsBackground || nextTag == "B" || nextTag == "S" - } - if previousTag == "B" || previousTag == "I" { - return previousSpan == nextSpan && (nextTag == "I" || nextTag == "E") - } - return false - } + if previousTag == "B" || previousTag == "I" { + if nextTag == "I" && previousSpan == nextSpan { + return biases["transition_bias_inside_to_continue"] ?? 0.0 + } + if nextTag == "E" && previousSpan == nextSpan { + return biases["transition_bias_inside_to_end"] ?? 0.0 + } + return 0.0 + } - private static func transitionBias( - previousTag: String?, - previousSpan: Int?, - nextTag: String?, - nextSpan: Int?, - labelInfo: OpenMedPrivacyFilterLabelInfo, - previousIndex: Int, - nextIndex: Int, - biases: [String: Float] - ) -> Float { - let previousIsBackground = - previousSpan == labelInfo.backgroundSpanLabel || previousIndex == labelInfo.backgroundTokenLabel - let nextIsBackground = - nextSpan == labelInfo.backgroundSpanLabel || nextIndex == labelInfo.backgroundTokenLabel - - if previousIsBackground { - if nextIsBackground { - return biases["transition_bias_background_stay"] ?? 0.0 - } - if nextTag == "B" || nextTag == "S" { - return biases["transition_bias_background_to_start"] ?? 0.0 + if previousTag == "E" || previousTag == "S" { + if nextIsBackground { + return biases["transition_bias_end_to_background"] ?? 0.0 + } + if nextTag == "B" || nextTag == "S" { + return biases["transition_bias_end_to_start"] ?? 0.0 + } } return 0.0 } + } - if previousTag == "B" || previousTag == "I" { - if nextTag == "I" && previousSpan == nextSpan { - return biases["transition_bias_inside_to_continue"] ?? 0.0 - } - if nextTag == "E" && previousSpan == nextSpan { - return biases["transition_bias_inside_to_end"] ?? 0.0 - } - return 0.0 + final class OpenMedPrivacyFilterPipeline { + private let artifact: OpenMedMLXArtifact + private let model: OpenMedPrivacyFilterForTokenClassification + private let tokenizer: OpenMedPrivacyFilterTokenizer + private let labelInfo: OpenMedPrivacyFilterLabelInfo + private let maxSeqLength: Int + + convenience init(modelDirectoryURL: URL, maxSeqLength: Int = 512) throws { + let artifact = try OpenMedMLXArtifact(modelDirectoryURL: modelDirectoryURL) + try self.init(artifact: artifact, maxSeqLength: maxSeqLength) } - if previousTag == "E" || previousTag == "S" { - if nextIsBackground { - return biases["transition_bias_end_to_background"] ?? 0.0 + init(artifact: OpenMedMLXArtifact, maxSeqLength: Int = 512) throws { + guard MLXTokenClassificationPipeline.isRuntimeSupported else { + throw OpenMedMLXRuntimeError.unsupportedPlatform } - if nextTag == "B" || nextTag == "S" { - return biases["transition_bias_end_to_start"] ?? 0.0 + guard artifact.task == .tokenClassification, + artifact.family == .openaiPrivacyFilter + else { + throw OpenMedPrivacyFilterError.unsupportedArtifact(artifact.manifest.family) } + + self.artifact = artifact + self.model = try OpenMedMLXModelLoader.loadPrivacyFilter(from: artifact) + self.tokenizer = try OpenMedPrivacyFilterTokenizer( + directoryURL: artifact.tokenizerDirectoryURL ?? artifact.directoryURL + ) + self.labelInfo = OpenMedPrivacyFilterLabelInfo(id2label: artifact.id2label) + self.maxSeqLength = min(maxSeqLength, artifact.manifest.maxSequenceLength ?? maxSeqLength) } - return 0.0 - } -} - -final class OpenMedPrivacyFilterPipeline { - private let artifact: OpenMedMLXArtifact - private let model: OpenMedPrivacyFilterForTokenClassification - private let tokenizer: OpenMedPrivacyFilterTokenizer - private let labelInfo: OpenMedPrivacyFilterLabelInfo - private let maxSeqLength: Int - - convenience init(modelDirectoryURL: URL, maxSeqLength: Int = 512) throws { - let artifact = try OpenMedMLXArtifact(modelDirectoryURL: modelDirectoryURL) - try self.init(artifact: artifact, maxSeqLength: maxSeqLength) - } - init(artifact: OpenMedMLXArtifact, maxSeqLength: Int = 512) throws { - guard MLXTokenClassificationPipeline.isRuntimeSupported else { - throw OpenMedMLXRuntimeError.unsupportedPlatform + var resolvedMaxSequenceLength: Int { + maxSeqLength } - guard artifact.task == .tokenClassification, - artifact.family == .openaiPrivacyFilter - else { - throw OpenMedPrivacyFilterError.unsupportedArtifact(artifact.manifest.family) + + func tokenOffsets(in text: String) throws -> [(Int, Int)] { + let encoded = try tokenizer.encode(text, maxTokens: Int.max) + return zip(encoded.charStarts, encoded.charEnds).map { ($0, $1) } } - self.artifact = artifact - self.model = try OpenMedMLXModelLoader.loadPrivacyFilter(from: artifact) - self.tokenizer = try OpenMedPrivacyFilterTokenizer( - directoryURL: artifact.tokenizerDirectoryURL ?? artifact.directoryURL - ) - self.labelInfo = OpenMedPrivacyFilterLabelInfo(id2label: artifact.id2label) - self.maxSeqLength = min(maxSeqLength, artifact.manifest.maxSequenceLength ?? maxSeqLength) - } + func predict( + _ text: String, + confidenceThreshold: Float = 0.0 + ) throws -> [EntityPrediction] { + let encoded = try tokenizer.encode(text, maxTokens: maxSeqLength) + guard !encoded.tokenIDs.isEmpty else { + return [] + } - var resolvedMaxSequenceLength: Int { - maxSeqLength - } + let sequenceLength = encoded.tokenIDs.count + let inputIDs = MLXArray(encoded.tokenIDs.map(Int32.init), [1, sequenceLength]) + let attentionMask = MLXArray.ones([1, sequenceLength], type: Bool.self) + let logits = model(inputIDs, attentionMask: attentionMask).asType(.float32) + let logProbabilities = (logits - logSumExp(logits, axis: -1, keepDims: true))[0] + let probabilities = exp(logProbabilities) + eval(logProbabilities, probabilities) - func tokenOffsets(in text: String) throws -> [(Int, Int)] { - let encoded = try tokenizer.encode(text, maxTokens: Int.max) - return zip(encoded.charStarts, encoded.charEnds).map { ($0, $1) } - } + let flatLogProbabilities = logProbabilities.asArray(Float.self) + let flatProbabilities = probabilities.asArray(Float.self) + let numLabels = artifact.configuration.numLabels + let tokenLogProbabilities = stride(from: 0, to: flatLogProbabilities.count, by: numLabels) + .map { offset in Array(flatLogProbabilities[offset..<(offset + numLabels)]) } + let tokenProbabilities = stride(from: 0, to: flatProbabilities.count, by: numLabels) + .map { offset in Array(flatProbabilities[offset..<(offset + numLabels)]) } + + let predictedIDs = OpenMedPrivacyFilterViterbi.decode( + tokenLogProbabilities: tokenLogProbabilities, + labelInfo: labelInfo, + biases: artifact.configuration.viterbiBiases + ) - func predict( - _ text: String, - confidenceThreshold: Float = 0.0 - ) throws -> [EntityPrediction] { - let encoded = try tokenizer.encode(text, maxTokens: maxSeqLength) - guard !encoded.tokenIDs.isEmpty else { - return [] + let sourceText = encoded.decodedText == text ? text : encoded.decodedText + return decodeGroupedEntities( + predictedIDs: predictedIDs, + probabilities: tokenProbabilities, + charStarts: encoded.charStarts, + charEnds: encoded.charEnds, + text: sourceText, + confidenceThreshold: confidenceThreshold + ) } - let sequenceLength = encoded.tokenIDs.count - let inputIDs = MLXArray(encoded.tokenIDs.map(Int32.init), [1, sequenceLength]) - let attentionMask = MLXArray.ones([1, sequenceLength], type: Bool.self) - let logits = model(inputIDs, attentionMask: attentionMask).asType(.float32) - let logProbabilities = (logits - logSumExp(logits, axis: -1, keepDims: true))[0] - let probabilities = exp(logProbabilities) - eval(logProbabilities, probabilities) - - let flatLogProbabilities = logProbabilities.asArray(Float.self) - let flatProbabilities = probabilities.asArray(Float.self) - let numLabels = artifact.configuration.numLabels - let tokenLogProbabilities = stride(from: 0, to: flatLogProbabilities.count, by: numLabels) - .map { offset in Array(flatLogProbabilities[offset..<(offset + numLabels)]) } - let tokenProbabilities = stride(from: 0, to: flatProbabilities.count, by: numLabels) - .map { offset in Array(flatProbabilities[offset..<(offset + numLabels)]) } - - let predictedIDs = OpenMedPrivacyFilterViterbi.decode( - tokenLogProbabilities: tokenLogProbabilities, - labelInfo: labelInfo, - biases: artifact.configuration.viterbiBiases - ) - - let sourceText = encoded.decodedText == text ? text : encoded.decodedText - return decodeGroupedEntities( - predictedIDs: predictedIDs, - probabilities: tokenProbabilities, - charStarts: encoded.charStarts, - charEnds: encoded.charEnds, - text: sourceText, - confidenceThreshold: confidenceThreshold - ) - } + private func decodeGroupedEntities( + predictedIDs: [Int], + probabilities: [[Float]], + charStarts: [Int], + charEnds: [Int], + text: String, + confidenceThreshold: Float + ) -> [EntityPrediction] { + let spans = labelsToTokenSpans(predictedIDs) + var entities = [EntityPrediction]() + for span in spans { + guard span.tokenStart >= 0, + span.tokenStart < span.tokenEnd, + span.tokenEnd <= charStarts.count + else { + continue + } - private func decodeGroupedEntities( - predictedIDs: [Int], - probabilities: [[Float]], - charStarts: [Int], - charEnds: [Int], - text: String, - confidenceThreshold: Float - ) -> [EntityPrediction] { - let spans = labelsToTokenSpans(predictedIDs) - var entities = [EntityPrediction]() - for span in spans { - guard span.tokenStart >= 0, - span.tokenStart < span.tokenEnd, - span.tokenEnd <= charStarts.count - else { - continue - } + var start = charStarts[span.tokenStart] + var end = charEnds[span.tokenEnd - 1] + trimWhitespace(start: &start, end: &end, text: text) + guard end > start else { + continue + } - var start = charStarts[span.tokenStart] - var end = charEnds[span.tokenEnd - 1] - trimWhitespace(start: &start, end: &end, text: text) - guard end > start else { - continue - } + let scores = (span.tokenStart.. Float? in + guard index < probabilities.count, predictedIDs[index] < probabilities[index].count else { + return nil + } + return probabilities[index][predictedIDs[index]] + } + let confidence = scores.isEmpty ? Float(0.0) : scores.reduce(0.0, +) / Float(scores.count) + guard confidence >= confidenceThreshold else { + continue + } - let scores = (span.tokenStart.. Float? in - guard index < probabilities.count, predictedIDs[index] < probabilities[index].count else { - return nil + let label: String + if span.label >= 0 && span.label < labelInfo.spanClassNames.count { + label = labelInfo.spanClassNames[span.label] + } else { + label = "label_\(span.label)" + } + refineStructuredPIISpan(label: label, start: &start, end: &end, text: text) + guard end > start else { + continue } - return probabilities[index][predictedIDs[index]] + entities.append( + EntityPrediction( + label: label, + text: substring(text, start: start, end: end), + confidence: confidence, + start: start, + end: end + ) + ) } - let confidence = scores.isEmpty ? Float(0.0) : scores.reduce(0.0, +) / Float(scores.count) - guard confidence >= confidenceThreshold else { - continue + return entities + } + + private func refineStructuredPIISpan( + label: String, + start: inout Int, + end: inout Int, + text: String + ) { + let normalizedLabel = label.lowercased() + let span = substring(text, start: start, end: end) + let patterns: [(hint: String, pattern: String)] = [ + ("email", #"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}\b"#), + ("url", #"\b(?:https?://|www\.)[^\s,;)\]]+"#), + ("phone", #"(?:\+?1[\s.-]?)?(?:\(?\d{3}\)?[\s.-]?)\d{3}[\s.-]?\d{4}"#), + ] + + for candidate in patterns where normalizedLabel.contains(candidate.hint) { + guard + let match = span.range( + of: candidate.pattern, + options: [.regularExpression, .caseInsensitive] + ) + else { + continue + } + let lowerOffset = span.distance(from: span.startIndex, to: match.lowerBound) + let upperOffset = span.distance(from: span.startIndex, to: match.upperBound) + start += lowerOffset + end = start + (upperOffset - lowerOffset) + return } - let label: String - if span.label >= 0 && span.label < labelInfo.spanClassNames.count { - label = labelInfo.spanClassNames[span.label] - } else { - label = "label_\(span.label)" - } - refineStructuredPIISpan(label: label, start: &start, end: &end, text: text) - guard end > start else { - continue - } - entities.append( - EntityPrediction( - label: label, - text: substring(text, start: start, end: end), - confidence: confidence, - start: start, - end: end - ) - ) - } - return entities - } - - private func refineStructuredPIISpan( - label: String, - start: inout Int, - end: inout Int, - text: String - ) { - let normalizedLabel = label.lowercased() - let span = substring(text, start: start, end: end) - let patterns: [(hint: String, pattern: String)] = [ - ("email", #"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}\b"#), - ("url", #"\b(?:https?://|www\.)[^\s,;)\]]+"#), - ("phone", #"(?:\+?1[\s.-]?)?(?:\(?\d{3}\)?[\s.-]?)\d{3}[\s.-]?\d{4}"#), - ] - - for candidate in patterns where normalizedLabel.contains(candidate.hint) { - guard - let match = span.range( - of: candidate.pattern, - options: [.regularExpression, .caseInsensitive] - ) - else { - continue + let lowercasedSpan = span.lowercased() + for suffix in [" and", " or"] where lowercasedSpan.hasSuffix(suffix) { + end -= suffix.count + trimWhitespace(start: &start, end: &end, text: text) + return } - let lowerOffset = span.distance(from: span.startIndex, to: match.lowerBound) - let upperOffset = span.distance(from: span.startIndex, to: match.upperBound) - start += lowerOffset - end = start + (upperOffset - lowerOffset) - return - } - - let lowercasedSpan = span.lowercased() - for suffix in [" and", " or"] where lowercasedSpan.hasSuffix(suffix) { - end -= suffix.count - trimWhitespace(start: &start, end: &end, text: text) - return } - } - private func labelsToTokenSpans(_ predictedIDs: [Int]) -> [( - label: Int, - tokenStart: Int, - tokenEnd: Int - )] { - var spans = [(label: Int, tokenStart: Int, tokenEnd: Int)]() - var currentLabel: Int? - var startIndex: Int? - var previousIndex: Int? + private func labelsToTokenSpans(_ predictedIDs: [Int]) -> [( + label: Int, + tokenStart: Int, + tokenEnd: Int + )] { + var spans = [(label: Int, tokenStart: Int, tokenEnd: Int)]() + var currentLabel: Int? + var startIndex: Int? + var previousIndex: Int? - for (tokenIndex, labelID) in predictedIDs.enumerated() { - let spanLabel = labelInfo.tokenToSpanLabel[labelID] - let boundaryTag = labelInfo.tokenBoundaryTags[labelID] + for (tokenIndex, labelID) in predictedIDs.enumerated() { + let spanLabel = labelInfo.tokenToSpanLabel[labelID] + let boundaryTag = labelInfo.tokenBoundaryTags[labelID] - if let previousIndex, tokenIndex != previousIndex + 1 { - if let currentLabel, let startIndex { - spans.append((currentLabel, startIndex, previousIndex + 1)) + if let previousIndex, tokenIndex != previousIndex + 1 { + if let currentLabel, let startIndex { + spans.append((currentLabel, startIndex, previousIndex + 1)) + } + currentLabel = nil + startIndex = nil } - currentLabel = nil - startIndex = nil - } - if spanLabel == nil { - previousIndex = tokenIndex - continue - } - - if spanLabel == labelInfo.backgroundSpanLabel { - if let currentLabel, let startIndex { - spans.append((currentLabel, startIndex, tokenIndex)) + if spanLabel == nil { + previousIndex = tokenIndex + continue } - currentLabel = nil - startIndex = nil - previousIndex = tokenIndex - continue - } - switch boundaryTag { - case "S": - if let currentLabel, let startIndex, let previousIndex { - spans.append((currentLabel, startIndex, previousIndex + 1)) - } - spans.append((spanLabel ?? 0, tokenIndex, tokenIndex + 1)) - currentLabel = nil - startIndex = nil - case "B": - if let currentLabel, let startIndex, let previousIndex { - spans.append((currentLabel, startIndex, previousIndex + 1)) - } - currentLabel = spanLabel - startIndex = tokenIndex - case "I": - if currentLabel == nil || currentLabel != spanLabel { - if let currentLabel, let startIndex, let previousIndex { - spans.append((currentLabel, startIndex, previousIndex + 1)) + if spanLabel == labelInfo.backgroundSpanLabel { + if let currentLabel, let startIndex { + spans.append((currentLabel, startIndex, tokenIndex)) } - currentLabel = spanLabel - startIndex = tokenIndex + currentLabel = nil + startIndex = nil + previousIndex = tokenIndex + continue } - case "E": - if currentLabel == nil || currentLabel != spanLabel || startIndex == nil { + + switch boundaryTag { + case "S": if let currentLabel, let startIndex, let previousIndex { spans.append((currentLabel, startIndex, previousIndex + 1)) } spans.append((spanLabel ?? 0, tokenIndex, tokenIndex + 1)) currentLabel = nil startIndex = nil - } else if let resolvedLabel = currentLabel, let resolvedStart = startIndex { - spans.append((resolvedLabel, resolvedStart, tokenIndex + 1)) - currentLabel = nil - startIndex = nil + case "B": + if let currentLabel, let startIndex, let previousIndex { + spans.append((currentLabel, startIndex, previousIndex + 1)) + } + currentLabel = spanLabel + startIndex = tokenIndex + case "I": + if currentLabel == nil || currentLabel != spanLabel { + if let currentLabel, let startIndex, let previousIndex { + spans.append((currentLabel, startIndex, previousIndex + 1)) + } + currentLabel = spanLabel + startIndex = tokenIndex + } + case "E": + if currentLabel == nil || currentLabel != spanLabel || startIndex == nil { + if let currentLabel, let startIndex, let previousIndex { + spans.append((currentLabel, startIndex, previousIndex + 1)) + } + spans.append((spanLabel ?? 0, tokenIndex, tokenIndex + 1)) + currentLabel = nil + startIndex = nil + } else if let resolvedLabel = currentLabel, let resolvedStart = startIndex { + spans.append((resolvedLabel, resolvedStart, tokenIndex + 1)) + currentLabel = nil + startIndex = nil + } + default: + break } - default: - break - } - previousIndex = tokenIndex - } + previousIndex = tokenIndex + } - if let currentLabel, let startIndex, let previousIndex { - spans.append((currentLabel, startIndex, previousIndex + 1)) + if let currentLabel, let startIndex, let previousIndex { + spans.append((currentLabel, startIndex, previousIndex + 1)) + } + return spans } - return spans - } - private func trimWhitespace(start: inout Int, end: inout Int, text: String) { - while start < end, character(at: start, in: text)?.isWhitespace == true { - start += 1 - } - while end > start, character(at: end - 1, in: text)?.isWhitespace == true { - end -= 1 + private func trimWhitespace(start: inout Int, end: inout Int, text: String) { + while start < end, character(at: start, in: text)?.isWhitespace == true { + start += 1 + } + while end > start, character(at: end - 1, in: text)?.isWhitespace == true { + end -= 1 + } } - } - private func character(at offset: Int, in text: String) -> Character? { - guard offset >= 0, - let index = text.index(text.startIndex, offsetBy: offset, limitedBy: text.endIndex), - index < text.endIndex - else { - return nil + private func character(at offset: Int, in text: String) -> Character? { + guard offset >= 0, + let index = text.index(text.startIndex, offsetBy: offset, limitedBy: text.endIndex), + index < text.endIndex + else { + return nil + } + return text[index] } - return text[index] - } - private func substring(_ text: String, start: Int, end: Int) -> String { - let lower = - text.index(text.startIndex, offsetBy: max(0, start), limitedBy: text.endIndex) - ?? text.endIndex - let upper = - text.index(text.startIndex, offsetBy: max(start, end), limitedBy: text.endIndex) - ?? text.endIndex - guard lower <= upper else { - return "" + private func substring(_ text: String, start: Int, end: Int) -> String { + let lower = + text.index(text.startIndex, offsetBy: max(0, start), limitedBy: text.endIndex) + ?? text.endIndex + let upper = + text.index(text.startIndex, offsetBy: max(start, end), limitedBy: text.endIndex) + ?? text.endIndex + guard lower <= upper else { + return "" + } + return String(text[lower.. [OpenMedZeroShotEntity] { - guard !labels.isEmpty else { - return [] - } + public final class OpenMedZeroShotNER { + private let artifact: OpenMedMLXArtifact + private let model: OpenMedGLiNERSpanModel + private let tokenizer: any Tokenizer + private let promptEncoder: OpenMedGLiNERPromptEncoder + private let maxSeqLength: Int - let split = OpenMedGLiNERPromptEncoder.splitWordsWithOffsets(text) - guard !split.words.isEmpty else { - return [] - } + public init(modelDirectoryURL: URL, maxSeqLength: Int = 512) throws { + guard MLXTokenClassificationPipeline.isRuntimeSupported else { + throw OpenMedMLXRuntimeError.unsupportedPlatform + } - let spec = artifact.manifest.promptSpec - let entityToken = spec?.entityToken ?? "<>" - let separatorToken = spec?.separatorToken ?? "<>" - var promptWords = [String]() - for label in labels { - promptWords.append(entityToken) - promptWords.append(label) - } - promptWords.append(separatorToken) - - let encoded = promptEncoder.encodeWords( - promptWords + split.words, - skipFirstWords: promptWords.count, - maxSeqLength: maxSeqLength, - specialTokenIDs: [ - entityToken: spec?.classTokenIndex ?? artifact.configuration.classTokenIndex, - separatorToken: artifact.configuration.textTokenIndex, - ].compactMapValues { $0 } - ) - let spans = OpenMedGLiNERPromptEncoder.buildCandidateSpanBatch( - wordCount: split.words.count, - maxWidth: artifact.configuration.maxWidth - ) - - let inputIDs = MLXArray(encoded.inputIDs.map(Int32.init), [1, encoded.inputIDs.count]) - let attentionMask = MLXArray(encoded.attentionMask, [1, encoded.attentionMask.count]) - .asType(.float32) - let wordsMask = MLXArray(encoded.wordsMask.map(Int32.init), [1, encoded.wordsMask.count]) - let output = model( - inputIDs: inputIDs, - attentionMask: attentionMask, - wordsMask: wordsMask, - spanIndex: spans.index, - spanMask: spans.mask - ) - let probabilities = sigmoid(output.logits) - eval(probabilities, output.promptMask, output.spanIndex, output.spanMask) - - let promptMaskValues = output.promptMask[0].asArray(Bool.self) - let validPromptCount = min(labels.count, promptMaskValues.filter { $0 }.count) - let scores = probabilities[0].asArray(Float.self) - let scoreLabelWidth = output.logits.dim(2) - let spanIndex = output.spanIndex[0].asArray(Int32.self).map(Int.init) - let spanMask = output.spanMask[0].asArray(Bool.self) - let spanCount = spanMask.count - - var entities = [OpenMedZeroShotEntity]() - for span in 0..= threshold else { + let artifact = try OpenMedMLXArtifact(modelDirectoryURL: modelDirectoryURL) + guard artifact.task == .zeroShotNER, artifact.family == .glinerUniEncoderSpan else { + throw OpenMedZeroShotError.unsupportedArtifact( + expectedTask: OpenMedMLXTask.zeroShotNER.rawValue, + expectedFamily: OpenMedMLXFamily.glinerUniEncoderSpan.rawValue, + actualTask: artifact.manifest.task, + actualFamily: artifact.manifest.family + ) + } + + self.artifact = artifact + self.model = try OpenMedMLXModelLoader.loadGLiNERSpanModel(from: artifact) + self.tokenizer = try OpenMed.loadTokenizer( + tokenizerName: artifact.tokenizerName ?? modelDirectoryURL.path, + tokenizerFolderURL: artifact.tokenizerDirectoryURL + ) + self.promptEncoder = OpenMedGLiNERPromptEncoder(tokenizer: tokenizer) + self.maxSeqLength = min(maxSeqLength, artifact.manifest.maxSequenceLength ?? maxSeqLength) + } + + public func extract( + _ text: String, + labels: [String], + threshold: Float = 0.5, + flatNER: Bool = true + ) throws -> [OpenMedZeroShotEntity] { + guard !labels.isEmpty else { + return [] + } + + let split = OpenMedGLiNERPromptEncoder.splitWordsWithOffsets(text) + guard !split.words.isEmpty else { + return [] + } + + let spec = artifact.manifest.promptSpec + let entityToken = spec?.entityToken ?? "<>" + let separatorToken = spec?.separatorToken ?? "<>" + var promptWords = [String]() + for label in labels { + promptWords.append(entityToken) + promptWords.append(label) + } + promptWords.append(separatorToken) + + let encoded = promptEncoder.encodeWords( + promptWords + split.words, + skipFirstWords: promptWords.count, + maxSeqLength: maxSeqLength, + specialTokenIDs: [ + entityToken: spec?.classTokenIndex ?? artifact.configuration.classTokenIndex, + separatorToken: artifact.configuration.textTokenIndex, + ].compactMapValues { $0 } + ) + let spans = OpenMedGLiNERPromptEncoder.buildCandidateSpanBatch( + wordCount: split.words.count, + maxWidth: artifact.configuration.maxWidth + ) + + let inputIDs = MLXArray(encoded.inputIDs.map(Int32.init), [1, encoded.inputIDs.count]) + let attentionMask = MLXArray(encoded.attentionMask, [1, encoded.attentionMask.count]) + .asType(.float32) + let wordsMask = MLXArray(encoded.wordsMask.map(Int32.init), [1, encoded.wordsMask.count]) + let output = model( + inputIDs: inputIDs, + attentionMask: attentionMask, + wordsMask: wordsMask, + spanIndex: spans.index, + spanMask: spans.mask + ) + let probabilities = sigmoid(output.logits) + eval(probabilities, output.promptMask, output.spanIndex, output.spanMask) + + let promptMaskValues = output.promptMask[0].asArray(Bool.self) + let validPromptCount = min(labels.count, promptMaskValues.filter { $0 }.count) + let scores = probabilities[0].asArray(Float.self) + let scoreLabelWidth = output.logits.dim(2) + let spanIndex = output.spanIndex[0].asArray(Int32.self).map(Int.init) + let spanMask = output.spanMask[0].asArray(Bool.self) + let spanCount = spanMask.count + + var entities = [OpenMedZeroShotEntity]() + for span in 0..= threshold else { + continue + } + entities.append( + OpenMedZeroShotEntity( + text: String(text.characterSlice(start: startChar, end: endChar)), + label: labels[labelIndex], + score: score, + start: startChar, + end: endChar + ) ) - ) + } } - } - return flatNER - ? OpenMedGLiNERPromptEncoder.suppressOverlaps(entities) - : entities.sorted { ($0.start, $0.end, $0.label) < ($1.start, $1.end, $1.label) } + return flatNER + ? OpenMedGLiNERPromptEncoder.suppressOverlaps(entities) + : entities.sorted { ($0.start, $0.end, $0.label) < ($1.start, $1.end, $1.label) } + } } -} -public final class OpenMedZeroShotClassifier { - private let artifact: OpenMedMLXArtifact - private let model: OpenMedGLiClassUniEncoderModel - private let tokenizer: any Tokenizer - private let maxSeqLength: Int + public final class OpenMedZeroShotClassifier { + private let artifact: OpenMedMLXArtifact + private let model: OpenMedGLiClassUniEncoderModel + private let tokenizer: any Tokenizer + private let maxSeqLength: Int - public init(modelDirectoryURL: URL, maxSeqLength: Int = 512) throws { - guard MLXTokenClassificationPipeline.isRuntimeSupported else { - throw OpenMedMLXRuntimeError.unsupportedPlatform - } + public init(modelDirectoryURL: URL, maxSeqLength: Int = 512) throws { + guard MLXTokenClassificationPipeline.isRuntimeSupported else { + throw OpenMedMLXRuntimeError.unsupportedPlatform + } - let artifact = try OpenMedMLXArtifact(modelDirectoryURL: modelDirectoryURL) - guard artifact.task == .zeroShotSequenceClassification, - artifact.family == .gliclassUniEncoder - else { - throw OpenMedZeroShotError.unsupportedArtifact( - expectedTask: OpenMedMLXTask.zeroShotSequenceClassification.rawValue, - expectedFamily: OpenMedMLXFamily.gliclassUniEncoder.rawValue, - actualTask: artifact.manifest.task, - actualFamily: artifact.manifest.family + let artifact = try OpenMedMLXArtifact(modelDirectoryURL: modelDirectoryURL) + guard artifact.task == .zeroShotSequenceClassification, + artifact.family == .gliclassUniEncoder + else { + throw OpenMedZeroShotError.unsupportedArtifact( + expectedTask: OpenMedMLXTask.zeroShotSequenceClassification.rawValue, + expectedFamily: OpenMedMLXFamily.gliclassUniEncoder.rawValue, + actualTask: artifact.manifest.task, + actualFamily: artifact.manifest.family + ) + } + + self.artifact = artifact + self.model = try OpenMedMLXModelLoader.loadGLiClassUniEncoderModel(from: artifact) + self.tokenizer = try OpenMed.loadTokenizer( + tokenizerName: artifact.tokenizerName ?? modelDirectoryURL.path, + tokenizerFolderURL: artifact.tokenizerDirectoryURL ) + self.maxSeqLength = min(maxSeqLength, artifact.manifest.maxSequenceLength ?? maxSeqLength) } - self.artifact = artifact - self.model = try OpenMedMLXModelLoader.loadGLiClassUniEncoderModel(from: artifact) - self.tokenizer = try OpenMed.loadTokenizer( - tokenizerName: artifact.tokenizerName ?? modelDirectoryURL.path, - tokenizerFolderURL: artifact.tokenizerDirectoryURL - ) - self.maxSeqLength = min(maxSeqLength, artifact.manifest.maxSequenceLength ?? maxSeqLength) - } - - public func classify( - _ text: String, - labels: [String], - threshold: Float = 0.5, - prompt: String? = nil - ) throws -> [OpenMedClassification] { - guard !labels.isEmpty else { - return [] - } + public func classify( + _ text: String, + labels: [String], + threshold: Float = 0.5, + prompt: String? = nil + ) throws -> [OpenMedClassification] { + guard !labels.isEmpty else { + return [] + } - let spec = artifact.manifest.promptSpec - let labelToken = spec?.labelToken ?? "<