.package(
    url: "https://github.com/ml-explore/mlx-swift-examples",
    exact: "2.21.2"
)

.product(name: "MLXLLM", package: "mlx-swift-examples")

// 모델이 사용하는 캐시 제한
MLX.GPU.set(cacheLimit: 20 * 1024 * 1024)

// HF에서 MLX로 구현된 llm 모델을 다운로드 후 사용
let modelConfiguration = LLMRegistry.qwen2_5_1_5b
// local에 저장된 모델도 사용 가능 
//let modelConfiguration = ModelConfiguration(directory: URL(fileURLWithPath: resourcePath))

// 모델 다운로드
let modelContainer = try await LLMModelFactory.shared.loadContainer(
    configuration: modelConfiguration
) {
    [modelConfiguration] progress in
    Task { @MainActor in
        self.modelInfo =
            "Downloading \(modelConfiguration.name): \(Int(progress.fractionCompleted * 100))%"
    }
}

// modelContainer 객체를 통해 사용

try await modelContainer.perform { (context: ModelContext) -> Void in
    let lmInput = try await context.processor.prepare(input: userInput)
    let stream = try MLXLMCommon.generate(input: lmInput, parameters: generateParameters, context: context)
    // generate and output in batches
    for await batch in stream._throttle(for: updateInterval, reducing: Generation.collect) {
        let output = batch.compactMap { $0.chunk }.joined(separator: "")
        if !output.isEmpty {
            Task { @MainActor [output] in
                self.output += output
            }
        }

        if let completion = batch.compactMap({ $0.info }).first {
            Task { @MainActor in
                self.stat = "\(completion.tokensPerSecond) tokens/s"
            }
        }
    }
}

// model unload를 위해 제공되는 API는 없습니다.
modelContainer = nil
// cache를 제거해 메모리 공간 확보
MLX.GPU.set(cacheLimit: 0)
MLX.GPU.clearCache()