Move common Registry types to MLXLMCommon (#229)

* Move VLMTypeRegistry to MLXLMCommon
* Move ModelRegistry to MLXLMCommon
* Move ProcessorTypeRegistry to MLXLMCommon
* Group registries
* Rename ModelRegistry to AbstractModelRegistry
* Add ModelRegistry as a typealias and deprecate it

Author: ibrahimcetin

Libraries/MLXLLM/LLMModelFactory.swift

@@ -20,20 +20,10 @@ private func create<C: Codable, M>(
 /// Registry of model type, e.g 'llama', to functions that can instantiate the model from configuration.
 ///
 /// Typically called via ``LLMModelFactory/load(hub:configuration:progressHandler:)``.
-public class ModelTypeRegistry: @unchecked Sendable {
-
-    /// Creates an empty registry.
-    public init() {
-        self.creators = [:]
-    }
-
-    /// Creates a registry with given creators.
-    public init(creators: [String: @Sendable (URL) throws -> any LanguageModel]) {
-        self.creators = creators
-    }
+public class LLMTypeRegistry: ModelTypeRegistry, @unchecked Sendable {
 
     /// Shared instance with default model types.
-    public static let shared: ModelTypeRegistry = .init(creators: all())
+    public static let shared: LLMTypeRegistry = .init(creators: all())
 
     /// All predefined model types.
     private static func all() -> [String: @Sendable (URL) throws -> any LanguageModel] {
@@ -53,32 +43,6 @@ public class ModelTypeRegistry: @unchecked Sendable {
         ]
     }
 
-    // Note: using NSLock as we have very small (just dictionary get/set)
-    // critical sections and expect no contention.  this allows the methods
-    // to remain synchronous.
-    private let lock = NSLock()
-    private var creators: [String: @Sendable (URL) throws -> any LanguageModel]
-
-    /// Add a new model to the type registry.
-    public func registerModelType(
-        _ type: String, creator: @Sendable @escaping (URL) throws -> any LanguageModel
-    ) {
-        lock.withLock {
-            creators[type] = creator
-        }
-    }
-
-    /// Given a `modelType` and configuration file instantiate a new `LanguageModel`.
-    public func createModel(configuration: URL, modelType: String) throws -> LanguageModel {
-        let creator = lock.withLock {
-            creators[modelType]
-        }
-        guard let creator else {
-            throw ModelFactoryError.unsupportedModelType(modelType)
-        }
-        return try creator(configuration)
-    }
-
 }
 
 /// Registry of models and any overrides that go with them, e.g. prompt augmentation.
@@ -87,23 +51,10 @@ public class ModelTypeRegistry: @unchecked Sendable {
 /// The python tokenizers have a very rich set of implementations and configuration.  The
 /// swift-tokenizers code handles a good chunk of that and this is a place to augment that
 /// implementation, if needed.
-public class ModelRegistry: @unchecked Sendable {
-
-    /// Creates an empty registry.
-    public init() {
-        self.registry = Dictionary()
-    }
-
-    /// Creates a new registry with from given model configurations.
-    public init(modelConfigurations: [ModelConfiguration]) {
-        self.registry = Dictionary(uniqueKeysWithValues: modelConfigurations.map { ($0.name, $0) })
-    }
+public class LLMRegistry: AbstractModelRegistry, @unchecked Sendable {
 
     /// Shared instance with default model configurations.
-    public static let shared = ModelRegistry(modelConfigurations: all())
-
-    private let lock = NSLock()
-    private var registry: [String: ModelConfiguration]
+    public static let shared = LLMRegistry(modelConfigurations: all())
 
     static public let smolLM_135M_4bit = ModelConfiguration(
         id: "mlx-community/SmolLM-135M-Instruct-4bit",
@@ -239,31 +190,11 @@ public class ModelRegistry: @unchecked Sendable {
         ]
     }
 
-    public func register(configurations: [ModelConfiguration]) {
-        lock.withLock {
-            for c in configurations {
-                registry[c.name] = c
-            }
-        }
-    }
-
-    public func configuration(id: String) -> ModelConfiguration {
-        lock.withLock {
-            if let c = registry[id] {
-                return c
-            } else {
-                return ModelConfiguration(id: id)
-            }
-        }
-    }
-
-    public var models: some Collection<ModelConfiguration> & Sendable {
-        lock.withLock {
-            return registry.values
-        }
-    }
 }
 
+@available(*, deprecated, renamed: "LLMRegistry", message: "Please use LLMRegistry directly.")
+public typealias ModelRegistry = LLMRegistry
+
 private struct LLMUserInputProcessor: UserInputProcessor {
 
     let tokenizer: Tokenizer
@@ -304,19 +235,20 @@ private struct LLMUserInputProcessor: UserInputProcessor {
 /// ```
 public class LLMModelFactory: ModelFactory {
 
-    public init(typeRegistry: ModelTypeRegistry, modelRegistry: ModelRegistry) {
+    public init(typeRegistry: ModelTypeRegistry, modelRegistry: AbstractModelRegistry) {
         self.typeRegistry = typeRegistry
         self.modelRegistry = modelRegistry
     }
 
     /// Shared instance with default behavior.
-    public static let shared = LLMModelFactory(typeRegistry: .shared, modelRegistry: .shared)
+    public static let shared = LLMModelFactory(
+        typeRegistry: LLMTypeRegistry.shared, modelRegistry: LLMRegistry.shared)
 
     /// registry of model type, e.g. configuration value `llama` -> configuration and init methods
     public let typeRegistry: ModelTypeRegistry
 
     /// registry of model id to configuration, e.g. `mlx-community/Llama-3.2-3B-Instruct-4bit`
-    public let modelRegistry: ModelRegistry
+    public let modelRegistry: AbstractModelRegistry
 
     public func configuration(id: String) -> ModelConfiguration {
         modelRegistry.configuration(id: id)

Libraries/MLXLMCommon/Registries/AbstractModelRegistry.swift

@@ -0,0 +1,43 @@
+// Copyright © 2024 Apple Inc.
+
+import Foundation
+
+open class AbstractModelRegistry: @unchecked Sendable {
+
+    /// Creates an empty registry.
+    public init() {
+        self.registry = Dictionary()
+    }
+
+    /// Creates a new registry with from given model configurations.
+    public init(modelConfigurations: [ModelConfiguration]) {
+        self.registry = Dictionary(uniqueKeysWithValues: modelConfigurations.map { ($0.name, $0) })
+    }
+
+    private let lock = NSLock()
+    private var registry: [String: ModelConfiguration]
+
+    public func register(configurations: [ModelConfiguration]) {
+        lock.withLock {
+            for c in configurations {
+                registry[c.name] = c
+            }
+        }
+    }
+
+    public func configuration(id: String) -> ModelConfiguration {
+        lock.withLock {
+            if let c = registry[id] {
+                return c
+            } else {
+                return ModelConfiguration(id: id)
+            }
+        }
+    }
+
+    public var models: some Collection<ModelConfiguration> & Sendable {
+        lock.withLock {
+            return registry.values
+        }
+    }
+}

Libraries/MLXLMCommon/Registries/ModelTypeRegistry.swift

@@ -0,0 +1,43 @@
+// Copyright © 2024 Apple Inc.
+
+import Foundation
+
+open class ModelTypeRegistry: @unchecked Sendable {
+
+    /// Creates an empty registry.
+    public init() {
+        self.creators = [:]
+    }
+
+    /// Creates a registry with given creators.
+    public init(creators: [String: @Sendable (URL) throws -> any LanguageModel]) {
+        self.creators = creators
+    }
+
+    // Note: using NSLock as we have very small (just dictionary get/set)
+    // critical sections and expect no contention.  this allows the methods
+    // to remain synchronous.
+    private let lock = NSLock()
+    private var creators: [String: @Sendable (URL) throws -> any LanguageModel]
+
+    /// Add a new model to the type registry.
+    public func registerModelType(
+        _ type: String, creator: @Sendable @escaping (URL) throws -> any LanguageModel
+    ) {
+        lock.withLock {
+            creators[type] = creator
+        }
+    }
+
+    /// Given a `modelType` and configuration file instantiate a new `LanguageModel`.
+    public func createModel(configuration: URL, modelType: String) throws -> LanguageModel {
+        let creator = lock.withLock {
+            creators[modelType]
+        }
+        guard let creator else {
+            throw ModelFactoryError.unsupportedModelType(modelType)
+        }
+        return try creator(configuration)
+    }
+
+}

Libraries/MLXLMCommon/Registries/ProcessorTypeRegistry.swift

@@ -0,0 +1,52 @@
+// Copyright © 2024 Apple Inc.
+
+import Foundation
+import Tokenizers
+
+open class ProcessorTypeRegistry: @unchecked Sendable {
+
+    /// Creates an empty registry.
+    public init() {
+        self.creators = [:]
+    }
+
+    /// Creates a registry with given creators.
+    public init(creators: [String: @Sendable (URL, any Tokenizer) throws -> any UserInputProcessor])
+    {
+        self.creators = creators
+    }
+
+    // Note: using NSLock as we have very small (just dictionary get/set)
+    // critical sections and expect no contention.  this allows the methods
+    // to remain synchronous.
+    private let lock = NSLock()
+
+    private var creators: [String: @Sendable (URL, any Tokenizer) throws -> any UserInputProcessor]
+
+    /// Add a new model to the type registry.
+    public func registerProcessorType(
+        _ type: String,
+        creator: @Sendable @escaping (
+            URL,
+            any Tokenizer
+        ) throws -> any UserInputProcessor
+    ) {
+        lock.withLock {
+            creators[type] = creator
+        }
+    }
+
+    /// Given a `processorType` and configuration file instantiate a new `UserInputProcessor`.
+    public func createModel(configuration: URL, processorType: String, tokenizer: any Tokenizer)
+        throws -> any UserInputProcessor
+    {
+        let creator = lock.withLock {
+            creators[processorType]
+        }
+        guard let creator else {
+            throw ModelFactoryError.unsupportedProcessorType(processorType)
+        }
+        return try creator(configuration, tokenizer)
+    }
+
+}