update for MLXEmbedders

Author: davidkoski

Package.resolved

@@ -26,9 +26,10 @@ extension EmbedderRuntime {
                 embeddings: [], skippedIndices: [], fallbackDescription: nil)
         }
 
-        return try await container.perform { model, tokenizer, pooler in
+        return try await container.perform { context in
             var skippedIndices: [Int] = []
 
+            let tokenizer = context.tokenizer
             let encoded = texts.enumerated().compactMap { index, text -> (Int, [Int])? in
                 let tokens = tokenizer.encode(text: text, addSpecialTokens: true)
                 guard !tokens.isEmpty else {
@@ -58,14 +59,14 @@ extension EmbedderRuntime {
             let mask = (padded .!= padToken)
             let tokenTypes = MLXArray.zeros(like: padded)
 
-            let outputs = model(
+            let outputs = context.model(
                 padded,
                 positionIds: nil,
                 tokenTypeIds: tokenTypes,
                 attentionMask: mask
             )
 
-            let poolingModule = resolvedPooler(for: pooler)
+            let poolingModule = resolvedPooler(for: context.pooling)
             let pooled = poolingModule(
                 outputs,
                 mask: mask,

Tools/embedder-tool/EmbedderRuntime+Embedding.swift

@@ -5,6 +5,7 @@ import ArgumentParser
 import Foundation
 import MLX
 import MLXEmbedders
+import MLXLMCommon
 import Tokenizers
 
 @main
@@ -17,7 +18,7 @@ struct EmbedderTool: AsyncParsableCommand {
         ]
     )
 
-    private static let defaultModelConfiguration = ModelConfiguration.nomic_text_v1_5
+    private static let defaultModelConfiguration = EmbedderRegistry.nomic_text_v1_5
 
     @OptionGroup var model: ModelArguments
     @OptionGroup var corpus: CorpusArguments
@@ -42,9 +43,8 @@ struct EmbedderTool: AsyncParsableCommand {
         -> EmbedderRuntime
     {
         let loadedModel = try await model.load(default: defaultModelConfiguration)
-        let baseStrategy = await loadedModel.container.perform { _, _, pooler in
-            pooler.strategy
-        }
+        let baseStrategy = await loadedModel.container.poolingStrategy
+
         return EmbedderRuntime(
             configuration: loadedModel.configuration,
             container: loadedModel.container,
@@ -58,7 +58,7 @@ struct EmbedderTool: AsyncParsableCommand {
 
 struct EmbedderRuntime {
     let configuration: ModelConfiguration
-    let container: ModelContainer
+    let container: EmbedderModelContainer
     let baseStrategy: Pooling.Strategy
     let strategyOverride: Pooling.Strategy?
     let normalize: Bool

Tools/embedder-tool/EmbedderTool.swift

@@ -1,6 +1,7 @@
 import ArgumentParser
 import Foundation
 import MLXEmbedders
+import MLXLMCommon
 
 struct ListCommand: AsyncParsableCommand {
     static let configuration = CommandConfiguration(
@@ -12,13 +13,13 @@ struct ListCommand: AsyncParsableCommand {
     var includeDirectories = false
 
     func run() async throws {
-        let models = await MainActor.run { Array(ModelConfiguration.models) }
+        let models = await MainActor.run { Array(EmbedderRegistry.shared.models) }
             .sorted { $0.name.localizedCaseInsensitiveCompare($1.name) == .orderedAscending }
 
         for configuration in models {
             switch configuration.id {
-            case .id(let identifier):
-                print(identifier)
+            case .id(let id, let revision):
+                print("\(id)/\(revision)")
             case .directory(let url):
                 if includeDirectories {
                     print(url.path)

Tools/embedder-tool/ListCommand.swift

@@ -28,7 +28,7 @@ struct ModelArguments: ParsableArguments {
             return localConfiguration
         }
 
-        return ModelConfiguration.configuration(id: model)
+        return ModelConfiguration(id: model)
     }
 
     var downloadURL: URL? {
@@ -38,20 +38,22 @@ struct ModelArguments: ParsableArguments {
 
 struct LoadedEmbedderModel {
     let configuration: ModelConfiguration
-    let container: ModelContainer
+    let container: EmbedderModelContainer
 }
 
 extension ModelArguments {
 
     func load(default defaultConfiguration: ModelConfiguration) async throws -> LoadedEmbedderModel
     {
         let configuration = await configuration(default: defaultConfiguration)
-        let hub = makeHub()
+        let hub = #hubDownloader
+        let loader = #huggingFaceTokenizerLoader
 
         print("Loading model \(configuration.name)...")
 
-        let container = try await MLXEmbedders.loadModelContainer(
-            hub: hub,
+        let container = try await EmbedderModelFactory.shared.loadContainer(
+            from: hub,
+            using: loader,
             configuration: configuration,
             progressHandler: { progress in
                 let percentage = Int(progress.fractionCompleted * 100)

Tools/embedder-tool/ModelArguments.swift

@@ -43,7 +43,7 @@ struct ModelArguments: ParsableArguments, Sendable {
     }
 
     @Sendable
-    func load(defaultModel: String, modelFactory: ModelFactory) async throws -> ModelContainer {
+    func load(defaultModel: String, modelFactory: any ModelFactory) async throws -> ModelContainer {
         let modelConfiguration: ModelConfiguration
 
         let modelName = self.model ?? defaultModel
@@ -58,7 +58,7 @@ struct ModelArguments: ParsableArguments, Sendable {
             modelConfiguration = modelFactory.configuration(id: modelName)
         }
 
-        return try await loadModelContainer(
+        return try await modelFactory.loadContainer(
             from: self.downloader,
             using: #huggingFaceTokenizerLoader(),
             configuration: modelConfiguration)
@@ -320,7 +320,7 @@ struct EvaluateCommand: AsyncParsableCommand {
 
     @MainActor
     mutating func run() async throws {
-        let modelFactory: ModelFactory
+        let modelFactory: any ModelFactory
         let defaultModel: ModelConfiguration
 
         // Switch between LLM and VLM based on presence of media