spb/zyquo-mlx Public MIT

1//2//  InferenceEngine.swift3//  Zyquo MLX4//5//  Author: Simon-Pierre Boucher6//  Mail: contact@spboucher.ai7//89import Foundation10import MLX11import MLXEmbedders12import MLXLMCommon1314/// The in-process inference engine. An actor: one loaded model at a time,15/// streaming generation for LLM/VLM, vector embedding, and verifiable16/// memory release on unload (charter Phase 2).17actor InferenceEngine {1819    static let shared = InferenceEngine()2021    private var adapter: LoadedModelAdapter?22    private(set) var loadedModel: LocalModel?2324    // MARK: - Load / unload2526    /// Load a local model (replacing any previously loaded one).27    func load(model: LocalModel) async throws {28        try await unload()29        adapter = try await ModelTypeAdapters.load(model: model)30        loadedModel = model31    }3233    /// Unload the current model and free GPU buffer cache. Returns the34    /// active-memory delta so callers can verify the release.35    @discardableResult36    func unload() async throws -> Int64 {37        let before = Memory.snapshot().activeMemory38        adapter = nil39        loadedModel = nil40        Memory.clearCache()41        let after = Memory.snapshot().activeMemory42        return Int64(before - after)43    }4445    func memorySnapshot() -> (active: Int, cache: Int, peak: Int) {46        let snap = Memory.snapshot()47        return (snap.activeMemory, snap.cacheMemory, snap.peakMemory)48    }4950    // MARK: - Text / chat generation (LLM & VLM)5152    /// Stream a response for a chat conversation. Images/videos attach to the53    /// user message for VLMs. Cancellation: cancel the surrounding Task.54    func generate(55        messages: [Chat.Message],56        params: GenerationParams57    ) async throws -> AsyncThrowingStream<InferenceEvent, Error> {58        guard let adapter else { throw InferenceEngineError.noModelLoaded }59        guard case .language(let container) = adapter else {60            throw InferenceEngineError.wrongModelKind(expected: "a language model")61        }6263        let userInput = UserInput(chat: messages)64        let mlxParams = params.asMLX6566        return AsyncThrowingStream { continuation in67            let task = Task {68                do {69                    try await container.perform(nonSendable: userInput) { (context: ModelContext, userInput: UserInput) in70                        let input = try await context.processor.prepare(input: userInput)71                        let stream = try MLXLMCommon.generate(72                            input: input, parameters: mlxParams, context: context)73                        for await generation in stream {74                            if Task.isCancelled { break }75                            switch generation {76                            case .chunk(let text):77                                continuation.yield(.chunk(text))78                            case .info(let info):79                                continuation.yield(80                                    .finished(81                                        InferenceStats(82                                            promptTokens: info.promptTokenCount,83                                            generatedTokens: info.generationTokenCount,84                                            ttft: info.promptTime,85                                            generateTime: info.generateTime,86                                            stopReason: String(describing: info.stopReason)87                                        )))88                            case .toolCall:89                                break // tool use surfaces in Phase 690                            }91                        }92                        continuation.finish()93                    }94                } catch {95                    continuation.finish(throwing: error)96                }97            }98            continuation.onTermination = { _ in task.cancel() }99        }100    }101102    // MARK: - Embeddings103104    /// Embed a batch of texts into L2-normalized vectors using the model's105    /// configured pooling strategy (docs/MLX-RESEARCH.md §4.7).106    func embed(texts: [String]) async throws -> [[Float]] {107        guard let adapter else { throw InferenceEngineError.noModelLoaded }108        guard case .embedder(let container) = adapter else {109            throw InferenceEngineError.wrongModelKind(expected: "an embedding model")110        }111112        return await container.perform { (context: EmbedderModelContext) in113            var vectors: [[Float]] = []114            for text in texts {115                let tokens = context.tokenizer.encode(text: text, addSpecialTokens: true)116                let ids = MLXArray(tokens)[.newAxis, 0...]117                let output = context.model(118                    ids, positionIds: nil, tokenTypeIds: nil, attentionMask: nil)119                let pooled = context.pooling(output, normalize: true)120                eval(pooled)121                vectors.append(pooled.squeezed().asArray(Float.self))122            }123            return vectors124        }125    }126127    /// Cosine similarity between two embedding vectors (both L2-normalized →128    /// plain dot product).129    static func cosineSimilarity(_ a: [Float], _ b: [Float]) -> Float {130        guard a.count == b.count, !a.isEmpty else { return 0 }131        return zip(a, b).reduce(0) { $0 + $1.0 * $1.1 }132    }133}134