spb/zyquo-local Public MIT

1//2//  PoCRunner.swift3//  Zyquo Local4//5//  Author: Simon-Pierre Boucher6//  Mail: contact@spboucher.ai7//89import Foundation1011/// CLI proof-of-concept mode — the Phase 2 gate:12/// `ZyquoLocal --poc <model-dir> "<prompt>"` loads an MLX model from a local13/// directory and streams generated tokens to stdout, ending with stats.14enum PoCRunner {15    static func run(arguments: [String]) async {16        guard let flagIndex = arguments.firstIndex(of: "--poc"),17            arguments.count >= flagIndex + 318        else {19            err("usage: ZyquoLocal --poc <model-dir> \"<prompt>\"\n")20            exit(64)21        }22        let directory = URL(fileURLWithPath: arguments[flagIndex + 1])23        let prompt = arguments[flagIndex + 2]2425        let size = (try? FileManager.default.allocatedSizeOfDirectory(at: directory)) ?? 026        let model = LocalModel(27            repoID: directory.pathComponents.suffix(2).joined(separator: "/"),28            directory: directory,29            sizeBytes: size30        )3132        let engine = InferenceEngine()33        do {34            err("Loading \(model.repoID) (\(ByteCountFormatter.string(fromByteCount: size, countStyle: .file)))…\n")35            let loadStart = Date()36            try await engine.load(model: model)37            err(String(format: "Loaded in %.2fs. Generating…\n\n", Date().timeIntervalSince(loadStart)))3839            let conversation = Conversation(modelID: model.repoID)40            try await engine.startSession(conversation: conversation)4142            var stats: GenerationStats?43            let stream = try await engine.generate(prompt: prompt, params: GenerationParams())44            for try await event in stream {45                switch event {46                case .token(let text):47                    out(text)48                case .stats(let s):49                    stats = s50                case .finished(let reason):51                    out("\n")52                    if let s = stats {53                        err(String(54                            format: "\n⚡ %.1f tok/s · %d tokens · %.2fs to first token · prompt %d tokens · peak %@ · stop: %@\n",55                            s.tokensPerSecond,56                            s.generationTokenCount,57                            s.timeToFirstToken,58                            s.promptTokenCount,59                            ByteCountFormatter.string(fromByteCount: Int64(s.peakMemoryBytes), countStyle: .memory),60                            reason.rawValue61                        ))62                    }63                }64            }65            await engine.unload()66            exit(0)67        } catch {68            err("error: \(error.localizedDescription)\n")69            exit(1)70        }71    }7273    private static func out(_ text: String) {74        FileHandle.standardOutput.write(Data(text.utf8))75    }7677    private static func err(_ text: String) {78        FileHandle.standardError.write(Data(text.utf8))79    }80}8182extension FileManager {83    /// Total allocated size of a directory tree in bytes.84    func allocatedSizeOfDirectory(at url: URL) throws -> Int64 {85        var total: Int64 = 086        let keys: Set<URLResourceKey> = [.totalFileAllocatedSizeKey, .isRegularFileKey]87        guard let enumerator = enumerator(at: url, includingPropertiesForKeys: Array(keys)) else {88            return 089        }90        for case let file as URL in enumerator {91            let values = try file.resourceValues(forKeys: keys)92            if values.isRegularFile == true {93                total += Int64(values.totalFileAllocatedSize ?? 0)94            }95        }96        return total97    }98}99