Improve embedding model error handling

Author: DePasqualeOrg

Libraries/Embedders/Configuration.swift

@@ -73,7 +73,7 @@ private class ModelTypeRegistry: @unchecked Sendable {
             creators[rawValue]
         }
         guard let creator else {
-            throw EmbedderError(message: "Unsupported model type.")
+            throw EmbedderError.unsupportedModelType(rawValue)
         }
         return try creator(configuration)
     }

Libraries/Embedders/EmbeddingModel.swift

@@ -49,7 +49,8 @@ public actor ModelContainer {
         async let tokenizerConfigTask = loadTokenizerConfig(
             configuration: configuration, hub: hub)
 
-        self.model = try loadSynchronous(modelDirectory: modelDirectory)
+        self.model = try loadSynchronous(
+            modelDirectory: modelDirectory, modelName: configuration.name)
         self.pooler = loadPooling(modelDirectory: modelDirectory)
 
         let (tokenizerConfig, tokenizerData) = try await tokenizerConfigTask

Libraries/Embedders/Load.swift

@@ -6,8 +6,45 @@ import MLX
 import MLXNN
 import Tokenizers
 
-struct EmbedderError: Error {
-    let message: String
+public enum EmbedderError: LocalizedError {
+    case unsupportedModelType(String)
+    case configurationDecodingError(String, String, DecodingError)
+    case missingTokenizerConfig
+
+    public var errorDescription: String? {
+        switch self {
+        case .unsupportedModelType(let type):
+            return "Unsupported model type: \(type)"
+        case .configurationDecodingError(let file, let modelName, let decodingError):
+            let errorDetail = extractDecodingErrorDetail(decodingError)
+            return "Failed to parse \(file) for model '\(modelName)': \(errorDetail)"
+        case .missingTokenizerConfig:
+            return "Missing tokenizer configuration"
+        }
+    }
+
+    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
+        }
+    }
 }
 
 func prepareModelDirectory(
@@ -51,20 +88,33 @@ public func load(
         hub: hub, configuration: configuration, progressHandler: progressHandler)
 
     async let tokenizerTask = loadTokenizer(configuration: configuration, hub: hub)
-    let model = try loadSynchronous(modelDirectory: modelDirectory)
+    let model = try loadSynchronous(modelDirectory: modelDirectory, modelName: configuration.name)
     let tokenizer = try await tokenizerTask
 
     return (model, tokenizer)
 }
 
-func loadSynchronous(modelDirectory: URL) throws -> EmbeddingModel {
+func loadSynchronous(modelDirectory: URL, modelName: String) throws -> EmbeddingModel {
     // Load config.json once and decode for both base config and model-specific config
     let configurationURL = modelDirectory.appending(component: "config.json")
-    let configData = try Data(contentsOf: configurationURL)
-    let baseConfig = try JSONDecoder().decode(BaseConfiguration.self, from: configData)
+    let configData: Data
+    let baseConfig: BaseConfiguration
+    do {
+        configData = try Data(contentsOf: configurationURL)
+        baseConfig = try JSONDecoder().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 = try modelType.createModel(configuration: configData)
+    let model: EmbeddingModel
+    do {
+        model = try modelType.createModel(configuration: configData)
+    } catch let error as DecodingError {
+        throw EmbedderError.configurationDecodingError(
+            configurationURL.lastPathComponent, modelName, error)
+    }
 
     // load the weights
     var weights = [String: MLXArray]()

Libraries/Embedders/Tokenizer.swift

@@ -45,7 +45,7 @@ func loadTokenizerConfig(configuration: ModelConfiguration, hub: HubApi) async t
     }
 
     guard let tokenizerConfig = try await config.tokenizerConfig else {
-        throw EmbedderError(message: "missing config")
+        throw EmbedderError.missingTokenizerConfig
     }
     let tokenizerData = try await config.tokenizerData
     return (tokenizerConfig, tokenizerData)