kangkang/康康/AI/LLMSession.swift

import Foundation
import MLX
import MLXLLM
import MLXLMCommon
// mlx-swift-lm 3.x 把 HF Hub/分词器栈拆成 opt-in 宏(MLXHuggingFace),
// `#hubDownloader()` / `#huggingFaceTokenizerLoader()` 展开处需要 HuggingFace / Tokenizers 在作用域。
import MLXHuggingFace
import HuggingFace
import Tokenizers

/// 封装 MLX 语言模型的流式生成,actor 保证单线程访问。
/// 基于 mlx-swift-lm 3.31.4 的 API(MLXLLM / MLXLMCommon)。
actor LLMSession {
    let container: ModelContainer

    /// 末次生成统计(取自流末尾的 .info 完成事件,性能自检用)。
    private(set) var lastStats: GenerateStats?

    private func record(_ s: GenerateStats) { lastStats = s }

    init(container: ModelContainer) {
        self.container = container
    }

    /// 在 simulator 把默认设备强切为 CPU(MLX 的 Metal backend 在部分 Sim 路径会 abort)。
    /// 真机走 body 默认设备(GPU/ANE)。
    /// 用 task-scoped `withDefaultDevice`,TaskLocal 会传递到 child Task / actor 调用。
    private static func withDeviceOverride<R>(
        _ body: () async throws -> R
    ) async rethrows -> R {
        #if targetEnvironment(simulator)
        return try await Device.withDefaultDevice(.cpu, body)
        #else
        return try await body()
        #endif
    }

    /// 从本地目录加载模型(包含 config.json + weights + tokenizer)。
    /// Gemma 4 改了回合标记:聊天回合以 `<turn|>`(token 106)结束(3n 时代是 `<end_of_turn>`),
    /// 而分词器自带的 eos 仍是 `<eos>`(1);不显式补 `<turn|>` 会让解码停不下来、把 maxTokens
    /// 跑满还吐回合分隔噪声。与上游 mlx-swift-lm 预置 `gemma4_e2b_it_4bit` 的 extraEOSTokens 一致。
    static func load(folderURL: URL) async throws -> LLMSession {
        let configuration = ModelConfiguration(
            directory: folderURL,
            extraEOSTokens: ["<turn|>"]
        )
        // 3.31.4:loadContainer 需显式传入 Downloader + TokenizerLoader。本地目录配置(.directory)
        // 不会真正下载(resolve 直接用目录),Downloader 仅占位;分词器由 HF AutoTokenizer 从目录加载。
        let container = try await withDeviceOverride {
            try await LLMModelFactory.shared.loadContainer(
                from: #hubDownloader(),
                using: #huggingFaceTokenizerLoader(),
                configuration: configuration
            )
        }
        return LLMSession(container: container)
    }

    /// 流式生成。返回的 AsyncThrowingStream 被取消时,内部 Task 也会取消。
    /// - Parameters:
    ///   - prompt: 原始 prompt 文本(经 processor 转 LMInput)
    ///   - maxTokens: 最大 token 数,由 GenerateParameters 控制
    func generate(prompt: String, maxTokens: Int) -> AsyncThrowingStream<TokenChunk, Error> {
        AsyncThrowingStream { continuation in
            let task = Task {
                do {
                    try await Self.withDeviceOverride {
                        // 低温:本 App 文本任务多为"直答/JSON 抽取",高温随机性会经常吐成非 JSON。
                        // 0.3 + topP 0.85 让输出更确定、JSON 更稳(与 MNN set_config 降温对齐)。
                        // repetitionPenalty:低温 + 无惩罚时,长文本(如「关键指标」列表)会逐行复读
                        // 进入死循环;1.1 的重复惩罚 + 64 token 上下文窗口掐断复读(与 MNN penalty 对齐)。
                        let parameters = GenerateParameters(
                            maxTokens: maxTokens,
                            temperature: Float(0.3),
                            topP: Float(0.85),
                            repetitionPenalty: Float(1.1),
                            repetitionContextSize: 64
                        )

                        try await container.perform { (context: ModelContext) in
                            let userInput = UserInput(prompt: prompt)
                            let lmInput = try await context.processor.prepare(input: userInput)

                            let start = Date()
                            var produced = 0

                            for await event in try MLXLMCommon.generate(
                                input: lmInput,
                                parameters: parameters,
                                context: context
                            ) {
                                if Task.isCancelled { break }

                                switch event {
                                case .chunk(let text):
                                    produced += 1
                                    let elapsed = Date().timeIntervalSince(start)
                                    let rate = elapsed > 0 ? Double(produced) / elapsed : 0
                                    continuation.yield(TokenChunk(text: text, decodeRate: rate))

                                case .info(let info):
                                    // 生成完成统计,是流的最后一个事件
                                    await self.record(GenerateStats(
                                        promptTokens: info.promptTokenCount,
                                        genTokens: info.generationTokenCount,
                                        prefillSeconds: info.promptTime,
                                        decodeSeconds: info.generateTime
                                    ))

                                case .toolCall:
                                    // 纯文本生成不会触发,switch 穷举
                                    break
                                }
                            }
                            // 注:研究笔记里曾建议尾部 MLX.GPU.synchronize() 以确保
                            // GPU 操作全部完成。但 AsyncStream 已经 yield 真实解码后的
                            // 文字,GPU 是否完全空闲不影响数据正确性。去掉此调用同时省
                            // 一份 transitive import MLX 的依赖,简化 SPM 链接。
                        }
                    }
                    continuation.finish()
                } catch {
                    continuation.finish(throwing: error)
                }
            }
            continuation.onTermination = { _ in task.cancel() }
        }
    }
}