KBaba7/llama.cpp

1import Foundation2import llama3 4let arguments = CommandLine.arguments5 6// Check that we have at least one argument (the model path)7guard arguments.count > 1 else {8    print("Usage: swift MODEL_PATH [PROMPT] [PARALLEL]")9    exit(1)10}11 12let modelPath: String = arguments[1]13let prompt: String = arguments.count > 2 ? arguments[2] : "Hello my name is"14let n_parallel: Int = arguments.count > 3 && Int(arguments[3]) != nil ? Int(arguments[3])! : 115 16// total length of the sequences including the prompt17let n_len: Int = 3218 19// init LLM20llama_backend_init()21defer {22    llama_backend_free()23}24 25let model_params = llama_model_default_params()26guard let model = llama_model_load_from_file(modelPath.cString(using: .utf8), model_params) else {27    print("Failed to load model")28    exit(1)29}30defer {31    llama_model_free(model)32}33 34guard let vocab = llama_model_get_vocab(model) else {35    print("Failed to get vocab")36    exit(1)37}38 39var tokens = tokenize(text: prompt, add_bos: true)40 41let n_kv_req = UInt32(tokens.count) + UInt32((n_len - Int(tokens.count)) * n_parallel)42 43var context_params = llama_context_default_params()44context_params.n_ctx = n_kv_req45context_params.n_batch = UInt32(max(n_len, n_parallel))46context_params.n_threads = 847context_params.n_threads_batch = 848 49let context = llama_init_from_model(model, context_params)50guard context != nil else {51    print("Failed to initialize context")52    exit(1)53}54defer {55    llama_free(context)56}57 58var sparams = llama_sampler_chain_default_params()59 60let smpl = llama_sampler_chain_init(sparams)61guard smpl != nil else {62    print("Failed to initialize sampling")63    exit(1)64}65defer {66    llama_sampler_free(smpl)67}68 69llama_sampler_chain_add(smpl, llama_sampler_init_top_k(40));70llama_sampler_chain_add(smpl, llama_sampler_init_top_p(0.9, 1));71llama_sampler_chain_add(smpl, llama_sampler_init_temp (0.4));72llama_sampler_chain_add(smpl, llama_sampler_init_dist (1234));73 74let n_ctx = llama_n_ctx(context)75 76print("\nn_len = \(n_len), n_ctx = \(n_ctx), n_batch = \(context_params.n_batch), n_parallel = \(n_parallel), n_kv_req = \(n_kv_req)\n")77 78if n_kv_req > n_ctx {79    print("error: n_kv_req (%d) > n_ctx, the required KV cache size is not big enough\n", n_kv_req)80    exit(1)81}82 83var buffer: [CChar] = []84for id: llama_token in tokens {85    print(token_to_piece(token: id, buffer: &buffer) ?? "", terminator: "")86}87 88print("\n")89 90var batch = llama_batch_init(max(Int32(tokens.count), Int32(n_parallel)), 0, 1)91defer {92    llama_batch_free(batch)93}94 95// evaluate the initial prompt96batch.n_tokens = Int32(tokens.count)97 98for (i, token) in tokens.enumerated() {99    batch.token[i] = token100    batch.pos[i] = Int32(i)101    batch.n_seq_id[i] = 1102    // batch.seq_id[i][0] = 0103    // TODO: is this the proper way to do this?104    if let seq_id = batch.seq_id[i] {105        seq_id[0] = 0106    }107    batch.logits[i] = 0108}109 110// llama_decode will output logits only for the last token of the prompt111batch.logits[Int(batch.n_tokens) - 1] = 1112 113if llama_decode(context, batch) != 0 {114    print("llama_decode() failed")115    exit(1)116}117 118for i in 1 ..< n_parallel {119    llama_kv_cache_seq_cp(context, 0, Int32(i), 0, batch.n_tokens)120}121 122if n_parallel > 1 {123    print("generating \(n_parallel) sequences ...\n")124}125 126var streams: [String] = .init(repeating: "", count: n_parallel)127var streamBuffers: [[CChar]] = .init(repeating: [], count: n_parallel)128var i_batch = [Int32](repeating: batch.n_tokens - 1, count: n_parallel)129 130var n_cur = batch.n_tokens131var n_decode = 0132 133let t_main_start = ggml_time_us()134 135while n_cur <= n_len {136    // prepare the next batch137    batch.n_tokens = 0138 139    // sample the next token for each parallel sequence / stream140    for i in 0 ..< n_parallel {141        if i_batch[i] < 0 {142            // the stream has already finished143            continue144        }145 146        let new_token_id = llama_sampler_sample(smpl, context, i_batch[i])147 148        // is it an end of stream? -> mark the stream as finished149        if llama_vocab_is_eog(vocab, new_token_id) || n_cur == n_len {150            i_batch[i] = -1151            // print("")152            if n_parallel > 1 {153                print("stream \(i) finished at n_cur = \(n_cur)")154            }155 156            continue157        }158 159        let nextStringPiece = token_to_piece(token: new_token_id, buffer: &streamBuffers[i]) ?? ""160 161        // if there is only one stream, we print immediately to stdout162        if n_parallel == 1 {163            print(nextStringPiece, terminator: "")164        }165        streams[i] += nextStringPiece166 167        // push this new token for next evaluation168        batch.token[Int(batch.n_tokens)] = new_token_id169        batch.pos[Int(batch.n_tokens)] = n_cur170        batch.n_seq_id[Int(batch.n_tokens)] = 1171        if let seq_id = batch.seq_id[Int(batch.n_tokens)] {172            seq_id[0] = Int32(i)173        }174        batch.logits[Int(batch.n_tokens)] = 1175 176        i_batch[i] = batch.n_tokens177 178        batch.n_tokens += 1179 180        n_decode += 1181    }182 183    // all streams are finished184    if batch.n_tokens == 0 {185        break186    }187 188    n_cur += 1189 190    // evaluate the current batch with the transformer model191    if llama_decode(context, batch) != 0 {192        print("llama_decode() failed")193        exit(1)194    }195}196 197if n_parallel > 1 {198    print("\n")199    for (i, stream) in streams.enumerated() {200        print("sequence \(i):\n\n\(prompt)\(stream)\n")201    }202}203 204let t_main_end = ggml_time_us()205 206print("decoded \(n_decode) tokens in \(String(format: "%.2f", Double(t_main_end - t_main_start) / 1_000_000.0)) s, speed: \(String(format: "%.2f", Double(n_decode) / (Double(t_main_end - t_main_start) / 1_000_000.0))) t/s\n\n")207 208llama_perf_sampler_print(smpl)209llama_perf_context_print(context)210 211private func tokenize(text: String, add_bos: Bool) -> [llama_token] {212    let utf8Count = text.utf8.count213    let n_tokens = utf8Count + (add_bos ? 1 : 0)214    let tokens = UnsafeMutablePointer<llama_token>.allocate(capacity: n_tokens)215    let tokenCount = llama_tokenize(vocab, text, Int32(utf8Count), tokens, Int32(n_tokens), add_bos, /*special tokens*/ false)216    var swiftTokens: [llama_token] = []217    for i in 0 ..< tokenCount {218        swiftTokens.append(tokens[Int(i)])219    }220    tokens.deallocate()221    return swiftTokens222}223 224private func token_to_piece(token: llama_token, buffer: inout [CChar]) -> String? {225    var result = [CChar](repeating: 0, count: 8)226    let nTokens = llama_token_to_piece(vocab, token, &result, Int32(result.count), 0, false)227    if nTokens < 0 {228        let actualTokensCount = -Int(nTokens)229        result = .init(repeating: 0, count: actualTokensCount)230        let check = llama_token_to_piece(231            vocab,232            token,233            &result,234            Int32(result.count),235            0,236            false237        )238        assert(check == actualTokensCount)239    } else {240        result.removeLast(result.count - Int(nTokens))241    }242    if buffer.isEmpty, let utfString = String(cString: result + [0], encoding: .utf8) {243        return utfString244    } else {245        buffer.append(contentsOf: result)246        let data = Data(buffer.map { UInt8(bitPattern: $0) })247        if buffer.count >= 4 { // 4 bytes is the max length of a utf8 character so if we're here we need to reset the buffer248            buffer = []249        }250        guard let bufferString = String(data: data, encoding: .utf8) else {251            return nil252        }253        buffer = []254        return bufferString255    }256}257