114 lines
5.0 KiB
Swift
114 lines
5.0 KiB
Swift
import Foundation
|
|
|
|
/// MNN(CPU / SME2)推理后端,封装 `MNNLLMBridge` 的文本流式生成。
|
|
/// 与 `LLMSession`/`VLSession` 同款 actor 隔离;跨调用的串行化由上游 `AIRuntime` 闸门保证。
|
|
///
|
|
/// 文本与视觉(图→文)由同一个 Qwen3.5-2B 多模态 MNN 模型承担:`generate` 走文本,
|
|
/// `analyze` 把图片拼成 <img> 标签交给 Omni 内核 imread 解码(需 OMNI 构建,xcframework 已含)。
|
|
/// 已实测可用,真机走此单模型全包路径;模拟器无 MNN,VL 仍回退 MLX(见 `AIRuntime`)。
|
|
actor MNNBackend {
|
|
private var bridge: MNNLLMBridge?
|
|
|
|
/// 末次生成统计(供 AIRuntime 在流结束后取走,性能自检用)。
|
|
private(set) var lastStats: GenerateStats?
|
|
|
|
private func record(_ s: GenerateStats) { lastStats = s }
|
|
|
|
var isLoaded: Bool { bridge?.isLoaded ?? false }
|
|
|
|
/// 从 MNN 模型目录加载(目录含 MNN llm 的 config.json + llm.mnn + 权重 + tokenizer)。
|
|
func load(folderURL: URL) throws {
|
|
let configPath = folderURL.appendingPathComponent("config.json").path
|
|
guard FileManager.default.fileExists(atPath: configPath) else {
|
|
throw AIRuntimeError.notReady
|
|
}
|
|
guard let b = MNNLLMBridge(configPath: configPath) else {
|
|
throw AIRuntimeError.modelLoadFailed("MNN createLLM/load 失败")
|
|
}
|
|
bridge = b
|
|
}
|
|
|
|
func unload() { bridge = nil }
|
|
|
|
/// 文本流式生成。`bridge.generateText` 同步阻塞、逐段回调,放在 detached 线程跑,
|
|
/// 把每段文本 yield 成 `TokenChunk`(含即时 tok/s)。流被取消时调用 `bridge.cancel()`。
|
|
func generate(prompt: String, maxTokens: Int) -> AsyncThrowingStream<TokenChunk, Error> {
|
|
guard let bridge else {
|
|
return AsyncThrowingStream { $0.finish(throwing: AIRuntimeError.notReady) }
|
|
}
|
|
let box = MNNUncheckedBox(bridge)
|
|
return AsyncThrowingStream { continuation in
|
|
let meter = MNNRateMeter()
|
|
let task = Task.detached(priority: .userInitiated) {
|
|
let stats = box.value.generateText(prompt, maxTokens: Int32(maxTokens)) { piece in
|
|
let rate = meter.tick()
|
|
continuation.yield(TokenChunk(text: piece, decodeRate: rate))
|
|
}
|
|
// ObjC 统计对象先抽成 Sendable 的 GenerateStats 再跨 actor 记录。
|
|
await self.record(GenerateStats(
|
|
promptTokens: Int(stats.promptTokens),
|
|
genTokens: Int(stats.genTokens),
|
|
prefillSeconds: stats.prefillMs / 1000.0,
|
|
decodeSeconds: stats.decodeMs / 1000.0
|
|
))
|
|
continuation.finish()
|
|
}
|
|
continuation.onTermination = { _ in
|
|
box.value.cancel()
|
|
task.cancel()
|
|
}
|
|
}
|
|
}
|
|
|
|
/// 图→文(VL)。一次性收集(JSON 抽取不需流式)。桥接里把图片路径拼成 <img> 标签,
|
|
/// MNN Omni 内部 imread 加载(需 OMNI 框架);blocking 调用放 detached 线程。
|
|
func analyze(imageURLs: [URL], prompt: String, maxTokens: Int) async throws -> String {
|
|
guard let bridge else { throw AIRuntimeError.notReady }
|
|
let paths = imageURLs.map(\.path)
|
|
let box = MNNUncheckedBox(bridge)
|
|
return try await withCheckedThrowingContinuation { cont in
|
|
Task.detached(priority: .userInitiated) {
|
|
let sink = MNNTextSink()
|
|
do {
|
|
let stats = try box.value.analyzeImages(paths, prompt: prompt, maxTokens: Int32(maxTokens)) { piece in
|
|
sink.append(piece)
|
|
}
|
|
await self.record(GenerateStats(
|
|
promptTokens: Int(stats.promptTokens),
|
|
genTokens: Int(stats.genTokens),
|
|
prefillSeconds: stats.prefillMs / 1000.0,
|
|
decodeSeconds: stats.decodeMs / 1000.0
|
|
))
|
|
cont.resume(returning: sink.text)
|
|
} catch {
|
|
cont.resume(throwing: AIRuntimeError.inferenceFailed(error.localizedDescription))
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/// 单线程串行回调聚合文本,无竞争。
|
|
private nonisolated final class MNNTextSink: @unchecked Sendable {
|
|
private(set) var text = ""
|
|
func append(_ s: String) { text += s }
|
|
}
|
|
|
|
/// 把非 Sendable 的 ObjC 桥对象安全带过 detached 边界。
|
|
/// 安全性来自 `AIRuntime` 闸门:同一时刻只有一个生成在跑,桥不会被并发访问。
|
|
private nonisolated struct MNNUncheckedBox<T>: @unchecked Sendable {
|
|
let value: T
|
|
init(_ value: T) { self.value = value }
|
|
}
|
|
|
|
/// 即时解码速率计:回调在单线程串行调用,内部计数无竞争。
|
|
private nonisolated final class MNNRateMeter: @unchecked Sendable {
|
|
private let start = Date()
|
|
private var produced = 0
|
|
func tick() -> Double {
|
|
produced += 1
|
|
let elapsed = Date().timeIntervalSince(start)
|
|
return elapsed > 0 ? Double(produced) / elapsed : 0
|
|
}
|
|
}
|