Swift と llama.cpp で iPhone 上でローカル LLM を動かす

// swift-tools-version: 6.0
import PackageDescription

let package = Package(
    name: "LocalModel",
    platforms: [.iOS(.v17), .macOS(.v14)],
    products: [.library(name: "LocalModel", targets: ["LocalModel"])],
    targets: [
        // llama.cpp b11887, official release build (MIT). iPhone (arm64) and Mac.
        .binaryTarget(
            name: "llama",
            url: "https://github.com/ggml-org/llama.cpp/releases/download/b11887/llama-b11887-xcframework.zip",
            checksum: "85aa7c9195918470c2de8dadf432bd8d01f5f7de2a4e65c1014491880ce673b6"
        ),
        .target(name: "LocalModel", dependencies: ["llama"]),
    ]
)

import Foundation
import llama

enum LLMError: Error {
    case loadFailed
    case promptTooLong(tokens: Int, limit: Int)
    case decodeFailed
}

/// Owns the llama.cpp pointers and frees them when the last reference goes away.
private final class Handle: @unchecked Sendable {
    let model: OpaquePointer
    let ctx: OpaquePointer
    let vocab: OpaquePointer
    init(model: OpaquePointer, ctx: OpaquePointer) {
        self.model = model
        self.ctx = ctx
        vocab = llama_model_get_vocab(model)
    }
    deinit {
        llama_free(ctx)
        llama_model_free(model)
    }
}

actor LLM {
    private let handle: Handle
    private let template: String?
    let contextLength: Int32 = 4096

    init(modelURL: URL) throws {
        llama_backend_init()

        var mp = llama_model_default_params()
        mp.n_gpu_layers = 99                       // all layers on the GPU (Metal)
        guard let model = llama_model_load_from_file(modelURL.path, mp) else { throw LLMError.loadFailed }

        var cp = llama_context_default_params()
        cp.n_ctx = UInt32(contextLength)           // prompt + answer
        cp.n_batch = cp.n_ctx                      // the whole prompt fits in one llama_decode call
        cp.n_ubatch = 512
        let threads = Int32(max(2, ProcessInfo.processInfo.activeProcessorCount - 2))
        cp.n_threads = threads
        cp.n_threads_batch = threads

        guard let ctx = llama_init_from_model(model, cp) else {
            llama_model_free(model)
            throw LLMError.loadFailed
        }
        handle = Handle(model: model, ctx: ctx)
        template = llama_model_chat_template(model, nil).map { String(cString: $0) }
    }
}

extension LLM {
    func prompt(_ messages: [(role: String, content: String)]) -> String {
        let roles = messages.map { strdup($0.role)! }
        let contents = messages.map { strdup($0.content)! }
        defer { (roles + contents).forEach { free($0) } }
        var chat = zip(roles, contents).map { llama_chat_message(role: $0, content: $1) }

        var buf = [CChar](repeating: 0, count: 16_384)
        var n = llama_chat_apply_template(template, &chat, chat.count, true, &buf, Int32(buf.count))
        if n > buf.count {                          // too small: the return value is the size needed
            buf = [CChar](repeating: 0, count: Int(n) + 1)
            n = llama_chat_apply_template(template, &chat, chat.count, true, &buf, Int32(buf.count))
        }
        if n > 0 { return String(decoding: buf.prefix(Int(n)).map { UInt8(bitPattern: $0) }, as: UTF8.self) }

        // Unknown template: ChatML is the most common format.
        return messages.map { "<|im_start|>\($0.role)\n\($0.content)<|im_end|>\n" }.joined() + "<|im_start|>assistant\n"
    }
}

extension LLM {
    func generate(_ messages: [(role: String, content: String)], maxTokens: Int = 512) throws -> String {
        let ctx = handle.ctx, vocab = handle.vocab
        llama_memory_clear(llama_get_memory(ctx), true)     // start a fresh conversation

        var tokens = tokenize(prompt(messages))
        let limit = Int(contextLength) - maxTokens
        guard tokens.count <= limit else { throw LLMError.promptTooLong(tokens: tokens.count, limit: limit) }

        guard let sampler = llama_sampler_chain_init(llama_sampler_chain_default_params()) else { throw LLMError.decodeFailed }
        defer { llama_sampler_free(sampler) }
        llama_sampler_chain_add(sampler, llama_sampler_init_greedy())

        let ok = tokens.withUnsafeMutableBufferPointer {
            llama_decode(ctx, llama_batch_get_one($0.baseAddress, Int32($0.count))) == 0
        }
        guard ok else { throw LLMError.decodeFailed }

        var bytes: [UInt8] = []
        for _ in 0..<maxTokens {
            var token = llama_sampler_sample(sampler, ctx, -1)
            if llama_vocab_is_eog(vocab, token) { break }
            bytes += piece(token)
            let next = withUnsafeMutablePointer(to: &token) { llama_decode(ctx, llama_batch_get_one($0, 1)) }
            if next != 0 { break }
        }
        return String(decoding: bytes, as: UTF8.self)
    }

    private func tokenize(_ text: String) -> [llama_token] {
        let vocab = handle.vocab
        let len = Int32(text.utf8.count)
        let needed = -llama_tokenize(vocab, text, len, nil, 0, false, true)   // parse special tokens
        guard needed > 0 else { return [] }
        var tokens = [llama_token](repeating: 0, count: Int(needed))
        let n = llama_tokenize(vocab, text, len, &tokens, needed, false, true)
        return n > 0 ? Array(tokens.prefix(Int(n))) : []
    }

    private func piece(_ token: llama_token) -> [UInt8] {
        var buf = [CChar](repeating: 0, count: 64)
        var n = llama_token_to_piece(handle.vocab, token, &buf, Int32(buf.count), 0, false)
        if n < 0 {
            buf = [CChar](repeating: 0, count: Int(-n))
            n = llama_token_to_piece(handle.vocab, token, &buf, Int32(buf.count), 0, false)
        }
        return n > 0 ? buf.prefix(Int(n)).map { UInt8(bitPattern: $0) } : []
    }
}

let llm = try LLM(modelURL: modelFile)
let reply = try await llm.generate([
    (role: "system", content: "Answer in one sentence."),
    (role: "user", content: "Why is the sky blue?"),
])