```
refactor: 重命名项目名称从"体己"到"康康" 将整个项目的目录结构从"体己"重命名为"康康",包括所有源代码文件、 资源文件、测试文件以及Xcode项目配置文件。此更改涉及项目中所有的 文件路径和应用入口点(App/TijiApp.swift → App/KangkangApp.swift)。 ```
This commit is contained in:
87
康康/AI/LLMSession.swift
Normal file
87
康康/AI/LLMSession.swift
Normal file
@@ -0,0 +1,87 @@
|
||||
import Foundation
|
||||
import MLX
|
||||
import MLXLLM
|
||||
import MLXLMCommon
|
||||
|
||||
/// 封装 MLX 语言模型的流式生成,actor 保证单线程访问。
|
||||
/// 基于 mlx-swift-examples 2.29.1(commit 9bff95ca)的 API。
|
||||
actor LLMSession {
|
||||
let container: ModelContainer
|
||||
|
||||
init(container: ModelContainer) {
|
||||
self.container = container
|
||||
}
|
||||
|
||||
/// 从本地目录加载模型(包含 config.json + weights + tokenizer)。
|
||||
static func load(folderURL: URL) async throws -> LLMSession {
|
||||
#if targetEnvironment(simulator)
|
||||
// MLX 的 iOS Simulator GPU stream 初始化会在部分 Metal backend 路径中 abort。
|
||||
// 模拟器只用于功能调试,强制走 CPU;真机仍保留默认 GPU/ANE 相关路径。
|
||||
Device.setDefault(device: .cpu)
|
||||
#endif
|
||||
|
||||
let configuration = ModelConfiguration(directory: folderURL)
|
||||
let container = try await LLMModelFactory.shared.loadContainer(
|
||||
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 {
|
||||
let parameters = GenerateParameters(
|
||||
maxTokens: maxTokens,
|
||||
temperature: Float(0.6),
|
||||
topP: Float(0.9)
|
||||
)
|
||||
|
||||
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:
|
||||
// 生成完成统计,是流的最后一个事件
|
||||
break
|
||||
|
||||
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() }
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user