280 lines
11 KiB
Swift
280 lines
11 KiB
Swift
// Copyright © 2024 Apple Inc.
|
|
|
|
import Foundation
|
|
import MLX
|
|
import MLXLMCommon
|
|
import MLXNN
|
|
|
|
/// Errors encountered during the model loading and initialization process.
|
|
///
|
|
/// This enum provides detailed feedback for failures in model type identification,
|
|
/// file access, JSON decoding, and missing configuration files.
|
|
public enum EmbedderError: LocalizedError {
|
|
|
|
/// The specified `model_type` in `config.json` is not supported by the current implementation.
|
|
case unsupportedModelType(String)
|
|
|
|
/// A required file could not be read from the disk.
|
|
/// - Parameters:
|
|
/// - fileName: The name of the file (e.g., "config.json").
|
|
/// - modelName: The name/ID of the model being loaded.
|
|
/// - error: The underlying system error.
|
|
case configurationFileError(String, String, Error)
|
|
|
|
/// The configuration file exists but contains invalid JSON or missing required fields.
|
|
case configurationDecodingError(String, String, DecodingError)
|
|
|
|
/// A human-readable description of the error.
|
|
public var errorDescription: String? {
|
|
switch self {
|
|
case .unsupportedModelType(let type):
|
|
return "Unsupported model type: \(type)"
|
|
case .configurationFileError(let file, let modelName, let error):
|
|
return "Error reading '\(file)' for model '\(modelName)': \(error.localizedDescription)"
|
|
case .configurationDecodingError(let file, let modelName, let decodingError):
|
|
let errorDetail = extractDecodingErrorDetail(decodingError)
|
|
return "Failed to parse \(file) for model '\(modelName)': \(errorDetail)"
|
|
}
|
|
}
|
|
|
|
/// Internal helper to provide specific details about JSON decoding failures,
|
|
/// such as the exact key path that failed.
|
|
private func extractDecodingErrorDetail(_ error: DecodingError) -> String {
|
|
switch error {
|
|
case .keyNotFound(let key, let context):
|
|
let path = (context.codingPath + [key]).map { $0.stringValue }.joined(separator: ".")
|
|
return "Missing field '\(path)'"
|
|
case .typeMismatch(_, let context):
|
|
let path = context.codingPath.map { $0.stringValue }.joined(separator: ".")
|
|
return "Type mismatch at '\(path)'"
|
|
case .valueNotFound(_, let context):
|
|
let path = context.codingPath.map { $0.stringValue }.joined(separator: ".")
|
|
return "Missing value at '\(path)'"
|
|
case .dataCorrupted(let context):
|
|
if context.codingPath.isEmpty {
|
|
return "Invalid JSON"
|
|
} else {
|
|
let path = context.codingPath.map { $0.stringValue }.joined(separator: ".")
|
|
return "Invalid data at '\(path)'"
|
|
}
|
|
@unknown default:
|
|
return error.localizedDescription
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Resolve model and tokenizer directories from a ``ModelConfiguration``
|
|
/// using a ``Downloader``.
|
|
///
|
|
/// - Parameters:
|
|
/// - downloader: The downloader to use for fetching remote resources.
|
|
/// - configuration: The configuration identifying the model.
|
|
/// - useLatest: When true, always checks the provider for updates.
|
|
/// - progressHandler: A closure to monitor download progress.
|
|
/// - Returns: A tuple of (modelDirectory, tokenizerDirectory).
|
|
func resolveDirectories(
|
|
from downloader: any Downloader,
|
|
configuration: ModelConfiguration,
|
|
useLatest: Bool = false,
|
|
progressHandler: @Sendable @escaping (Progress) -> Void
|
|
) async throws -> (modelDirectory: URL, tokenizerDirectory: URL) {
|
|
let modelDirectory: URL
|
|
switch configuration.id {
|
|
case .id(let id, let revision):
|
|
modelDirectory = try await downloader.download(
|
|
id: id, revision: revision,
|
|
matching: modelDownloadPatterns,
|
|
useLatest: useLatest,
|
|
progressHandler: progressHandler)
|
|
case .directory(let directory):
|
|
modelDirectory = directory
|
|
}
|
|
|
|
let tokenizerDirectory: URL
|
|
switch configuration.tokenizerSource {
|
|
case .id(let id, let revision):
|
|
tokenizerDirectory = try await downloader.download(
|
|
id: id, revision: revision,
|
|
matching: tokenizerDownloadPatterns,
|
|
useLatest: useLatest,
|
|
progressHandler: { _ in })
|
|
case .directory(let directory):
|
|
tokenizerDirectory = directory
|
|
case nil:
|
|
tokenizerDirectory = modelDirectory
|
|
}
|
|
|
|
return (modelDirectory, tokenizerDirectory)
|
|
}
|
|
|
|
/// Asynchronously loads the `EmbeddingModel` and its associated `Tokenizer`.
|
|
///
|
|
/// This is the primary high-level function for initializing an embedding pipeline.
|
|
/// It leverages `async let` to load the tokenizer in parallel while the model
|
|
/// structure is being built synchronously.
|
|
///
|
|
/// - Parameters:
|
|
/// - downloader: The ``Downloader`` to use for fetching remote resources.
|
|
/// - configuration: The model configuration.
|
|
/// - useLatest: When true, always checks the provider for updates.
|
|
/// - progressHandler: A closure for tracking download progress.
|
|
/// - Returns: A tuple containing the initialized `EmbeddingModel` and `Tokenizer`.
|
|
public func load(
|
|
from downloader: any Downloader,
|
|
using tokenizerLoader: any TokenizerLoader,
|
|
configuration: ModelConfiguration,
|
|
useLatest: Bool = false,
|
|
progressHandler: @Sendable @escaping (Progress) -> Void = { _ in }
|
|
) async throws -> (EmbeddingModel, Tokenizer) {
|
|
let (modelDirectory, tokenizerDirectory) = try await resolveDirectories(
|
|
from: downloader, configuration: configuration, useLatest: useLatest,
|
|
progressHandler: progressHandler)
|
|
|
|
async let tokenizerTask = tokenizerLoader.load(from: tokenizerDirectory)
|
|
let model = try loadSynchronous(modelDirectory: modelDirectory, modelName: configuration.name)
|
|
let tokenizer = try await tokenizerTask
|
|
|
|
return (model, tokenizer)
|
|
}
|
|
|
|
/// Synchronously initializes the model architecture, loads weights, and applies quantization.
|
|
///
|
|
/// This function performs the following steps:
|
|
/// 1. Reads and decodes `config.json`.
|
|
/// 2. Instantiates the specific model class based on `model_type`.
|
|
/// 3. Recursively scans the directory for `.safetensors` weight files.
|
|
/// 4. Applies quantization if defined in the configuration.
|
|
/// 5. Updates the model parameters and performs an initial evaluation (`eval`).
|
|
///
|
|
/// - Parameters:
|
|
/// - modelDirectory: The local `URL` containing model files.
|
|
/// - modelName: The display name of the model for error reporting.
|
|
/// - Returns: A fully initialized and weighted `EmbeddingModel`.
|
|
func loadSynchronous(modelDirectory: URL, modelName: String) throws -> EmbeddingModel {
|
|
let configurationURL = modelDirectory.appending(component: "config.json")
|
|
let configData: Data
|
|
do {
|
|
configData = try Data(contentsOf: configurationURL)
|
|
} catch {
|
|
throw EmbedderError.configurationFileError(
|
|
configurationURL.lastPathComponent, modelName, error)
|
|
}
|
|
|
|
let baseConfig: BaseConfiguration
|
|
do {
|
|
baseConfig = try JSONDecoder.json5().decode(BaseConfiguration.self, from: configData)
|
|
} catch let error as DecodingError {
|
|
throw EmbedderError.configurationDecodingError(
|
|
configurationURL.lastPathComponent, modelName, error)
|
|
}
|
|
|
|
let modelType = ModelType(rawValue: baseConfig.modelType)
|
|
let model: EmbeddingModel
|
|
do {
|
|
model = try modelType.createModel(configuration: configData)
|
|
} catch let error as DecodingError {
|
|
throw EmbedderError.configurationDecodingError(
|
|
configurationURL.lastPathComponent, modelName, error)
|
|
}
|
|
|
|
var weights = [String: MLXArray]()
|
|
let enumerator = FileManager.default.enumerator(
|
|
at: modelDirectory, includingPropertiesForKeys: nil)!
|
|
for case let url as URL in enumerator {
|
|
if url.pathExtension == "safetensors" {
|
|
let w = try loadArrays(url: url)
|
|
for (key, value) in w {
|
|
weights[key] = value
|
|
}
|
|
}
|
|
}
|
|
|
|
weights = model.sanitize(weights: weights)
|
|
|
|
if let perLayerQuantization = baseConfig.perLayerQuantization {
|
|
quantize(model: model) { path, module in
|
|
if weights["\(path).scales"] != nil {
|
|
return perLayerQuantization.quantization(layer: path)?.asTuple
|
|
} else {
|
|
return nil
|
|
}
|
|
}
|
|
}
|
|
|
|
let parameters = ModuleParameters.unflattened(weights)
|
|
try model.update(parameters: parameters, verify: [.all])
|
|
|
|
eval(model)
|
|
|
|
return model
|
|
}
|
|
|
|
/// Asynchronously loads a `ModelContainer` for thread-safe model access.
|
|
///
|
|
/// The `ModelContainer` is recommended for applications where multiple threads
|
|
/// or tasks may need to access the embedding model simultaneously.
|
|
///
|
|
/// - Parameters:
|
|
/// - downloader: The ``Downloader`` to use for fetching remote resources.
|
|
/// - configuration: The model configuration.
|
|
/// - useLatest: When true, always checks the provider for updates.
|
|
/// - progressHandler: A closure for tracking download progress.
|
|
/// - Returns: A thread-safe `ModelContainer` instance.
|
|
public func loadModelContainer(
|
|
from downloader: any Downloader,
|
|
using tokenizerLoader: any TokenizerLoader,
|
|
configuration: ModelConfiguration,
|
|
useLatest: Bool = false,
|
|
progressHandler: @Sendable @escaping (Progress) -> Void = { _ in }
|
|
) async throws -> ModelContainer {
|
|
let (modelDirectory, tokenizerDirectory) = try await resolveDirectories(
|
|
from: downloader, configuration: configuration, useLatest: useLatest,
|
|
progressHandler: progressHandler)
|
|
|
|
return try await ModelContainer(
|
|
modelDirectory: modelDirectory,
|
|
tokenizerDirectory: tokenizerDirectory,
|
|
configuration: configuration,
|
|
tokenizerLoader: tokenizerLoader)
|
|
}
|
|
|
|
/// Load an embedding model from a local directory.
|
|
///
|
|
/// No downloader is needed — the model and tokenizer are loaded from
|
|
/// the given directory.
|
|
///
|
|
/// - Parameter directory: The local directory containing model files.
|
|
/// - Returns: A tuple containing the initialized `EmbeddingModel` and `Tokenizer`.
|
|
public func load(
|
|
from directory: URL,
|
|
using tokenizerLoader: any TokenizerLoader
|
|
) async throws -> (EmbeddingModel, Tokenizer) {
|
|
let name =
|
|
directory.deletingLastPathComponent().lastPathComponent + "/"
|
|
+ directory.lastPathComponent
|
|
async let tokenizerTask = tokenizerLoader.load(from: directory)
|
|
let model = try loadSynchronous(modelDirectory: directory, modelName: name)
|
|
let tokenizer = try await tokenizerTask
|
|
return (model, tokenizer)
|
|
}
|
|
|
|
/// Load an embedding model container from a local directory.
|
|
///
|
|
/// No downloader is needed — the model and tokenizer are loaded from
|
|
/// the given directory.
|
|
///
|
|
/// - Parameters:
|
|
/// - directory: The local directory containing model files.
|
|
/// - tokenizerLoader: The ``TokenizerLoader`` to use for loading the tokenizer.
|
|
/// - Returns: A thread-safe `ModelContainer` instance.
|
|
public func loadModelContainer(
|
|
from directory: URL,
|
|
using tokenizerLoader: any TokenizerLoader
|
|
) async throws -> ModelContainer {
|
|
try await ModelContainer(
|
|
modelDirectory: directory,
|
|
tokenizerDirectory: directory,
|
|
configuration: ModelConfiguration(directory: directory),
|
|
tokenizerLoader: tokenizerLoader)
|
|
}
|