Something went wrong. Try again.
Routing embedding stack for Qwen3-Embedding-0.6B on Apple Silicon: orchestrator + MLX + llama.cpp backends
Something went wrong. Try again.
2.9 kB · 74 lines
Swift
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475import Foundationimport Hubimport MLXEmbeddersimport MLXLMCommon
/// Loads Qwen3-Embedding model containers on demand and caches them per size./// Deduplicates concurrent loads of the same size.actor ModelRegistry { private var loaded: [ModelSize: MLXEmbedders.ModelContainer] = [:] private var loading: [ModelSize: Task<MLXEmbedders.ModelContainer, Error>] = [:]
private let hub: HubApi
init(hub: HubApi = HubApi()) { self.hub = hub }
/// Already-loaded sizes (read-only snapshot). func loadedSizes() -> Set<ModelSize> { Set(loaded.keys) }
/// Get or lazily load the container for a given size. func container(for size: ModelSize) async throws -> MLXEmbedders.ModelContainer { if let c = loaded[size] { return c } if let inflight = loading[size] { return try await inflight.value }
let hub = self.hub let task = Task<MLXEmbedders.ModelContainer, Error> { // Check for a local override env var first — lets us bypass the // Hub downloader entirely (useful when the Hub's metadata-file // management fails under sandboxing). let envKey: String switch size { case .s06: envKey = "MLX_MODEL_DIR_0_6B" case .s4: envKey = "MLX_MODEL_DIR_4B" case .s8: envKey = "MLX_MODEL_DIR_8B" } let dir: URL if let localPath = ProcessInfo.processInfo.environment[envKey], !localPath.isEmpty { dir = URL(fileURLWithPath: localPath) } else { let repoId = size.repoId // HubApi.snapshot downloads or returns an existing local snapshot directory. dir = try await hub.snapshot( from: repoId, matching: ["*.json", "*.safetensors", "*.txt", "tokenizer.*", "*.model"] ) } let container = try await MLXEmbedders.loadModelContainer( from: dir, using: HFTokenizerLoader()) return container } loading[size] = task do { let container = try await task.value loaded[size] = container loading[size] = nil return container } catch { loading[size] = nil throw error } }
/// Eagerly load and warm a list of sizes. Failures are thrown after the /// first bad size (warmup is best-effort; callers decide policy). func preload(_ sizes: [ModelSize], maxBatch: Int, maxLen: Int) async throws { for size in sizes { let container = try await container(for: size) await WarmUp.warm(container: container, maxBatch: maxBatch, maxLen: maxLen) fputs("[embed] warmed \(size.rawValue)\n", stderr) } }}