Felipe97/llama-cpp-compiled

1#include "arg.h"2#include "common.h"3#include "log.h"4#include "llama.h"5 6#include <clocale>7#include <cmath>8#include <cstdio>9#include <string>10#include <vector>11#include <algorithm>12 13static void print_usage(int, char ** argv) {14    LOG("\nexample usage:\n");15    LOG("\n    %s -m model.gguf --junk 250 --pos 90 --keep 32 --grp-attn-n 2 [--seed 1234]\n", argv[0]);16    LOG("\n");17}18 19int main(int argc, char ** argv) {20    std::setlocale(LC_NUMERIC, "C");21 22    common_params params;23 24    params.n_junk = 250;25    params.n_keep = 32;26    params.i_pos  = -1;27 28    common_init();29 30    if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_PASSKEY, print_usage)) {31        return 1;32    }33 34    int n_junk = params.n_junk;35    int n_keep = params.n_keep;36    int n_grp  = params.grp_attn_n;37    int i_pos  = params.i_pos;38 39    if (i_pos == -1) {40        i_pos = rand() % n_junk;41    }42 43    const std::string prompt_prefix = "There is an important info hidden inside a lot of irrelevant text. Find it and memorize them. I will quiz you about the important information there.";44    const std::string prompt_suffix = " What is the pass key? The pass key is";45 46    // generate junk text47    params.prompt = prompt_prefix;48 49    const int passkey = rand() % 50000 + 1;50 51    for (int i = 0; i < n_junk; i++) {52        if (i % n_junk == i_pos) {53            params.prompt += " The pass key is " + std::to_string(passkey) + ". Remember it. " + std::to_string(passkey) + " is the pass key.";54        }55 56        params.prompt += " The grass is green. The sky is blue. The sun is yellow. Here we go. There and back again.";57    }58 59    params.prompt += prompt_suffix;60 61    // init LLM62 63    llama_backend_init();64    llama_numa_init(params.numa);65 66    // initialize the model67 68    llama_model_params model_params = common_model_params_to_llama(params);69 70    llama_model * model = llama_model_load_from_file(params.model.path.c_str(), model_params);71 72    if (model == NULL) {73        LOG_ERR("%s: unable to load model\n" , __func__);74        return 1;75    }76 77    const llama_vocab * vocab = llama_model_get_vocab(model);78 79    // initialize the context80 81    llama_context_params ctx_params = common_context_params_to_llama(params);82 83    ctx_params.n_ctx = llama_model_n_ctx_train(model)*n_grp + n_keep;84 85    GGML_ASSERT(ctx_params.n_batch % n_grp == 0 && "n_batch must be divisible by n_grp");86 87    llama_context * ctx = llama_init_from_model(model, ctx_params);88    if (ctx == NULL) {89        LOG_ERR("%s: failed to create the llama_context\n" , __func__);90        return 1;91    }92 93    auto sparams = llama_sampler_chain_default_params();94 95    llama_sampler * smpl = llama_sampler_chain_init(sparams);96 97    llama_sampler_chain_add(smpl, llama_sampler_init_greedy());98 99    // tokenize the prompt100    std::vector<llama_token> tokens_list;101    tokens_list = common_tokenize(ctx, params.prompt, true);102 103    // tokenize the prefix and use it as a sink104    const int n_tokens_prefix = common_tokenize(ctx, prompt_prefix, true).size();105 106    const int n_tokens_all = tokens_list.size();107 108    // we leave a margin of 16 tokens for the generated text - it should contain just the passkey109    const int n_predict = 16;110 111    // total length of the sequences including the prompt112    const int n_len = n_tokens_all + n_predict;113 114    const int n_ctx       = llama_n_ctx(ctx) - n_keep;115    const int n_kv_req    = llama_n_ctx(ctx);116    const int n_batch     = ctx_params.n_batch;117    const int n_batch_grp = ctx_params.n_batch/n_grp;118 119    LOG_INF("\n%s: n_len = %d, n_ctx = %d, n_kv_req = %d, n_grp = %d, n_batch = %d, n_junk = %d, i_pos = %d\n", __func__, n_len, n_ctx, n_kv_req, n_grp, n_batch, n_junk, i_pos);120 121    // print the prompt token-by-token122 123    LOG_INF("\n");124    LOG_INF("prefix tokens: %d\n", n_tokens_prefix);125    LOG_INF("prompt tokens: %d\n", n_tokens_all);126    //LOG_INF("prompt: %s\n", params.prompt.c_str());127 128    llama_batch batch = llama_batch_init(params.n_batch, 0, 1);129 130    int n_past = 0;131 132    auto * mem = llama_get_memory(ctx);133 134    // fill the KV cache135    for (int i = 0; i < n_ctx; i += n_batch) {136        if (i > 0 && n_grp > 1) {137            // if SelfExtend is enabled, we compress the position from the last batch by a factor of n_grp138            const int ib = i/n_batch - 1;139            const int bd = n_batch_grp*(n_grp - 1);140 141            llama_memory_seq_add(mem, 0, n_past - n_batch,         n_past,         ib*bd);142            llama_memory_seq_div(mem, 0, n_past - n_batch + ib*bd, n_past + ib*bd, n_grp);143 144            n_past = llama_memory_seq_pos_max(mem, 0) + 1;145        }146 147        common_batch_clear(batch);148 149        for (int j = 0; j < n_batch && i + j < n_tokens_all; j++) {150            common_batch_add(batch, tokens_list[i + j], n_past++, { 0 }, false);151        }152 153        if (i + n_batch >= n_tokens_all) {154            batch.logits[batch.n_tokens - 1] = true;155        }156 157        if (llama_decode(ctx, batch) != 0) {158            LOG_INF("%s: llama_decode() failed\n", __func__);159            return 1;160        }161 162        LOG_INF("%s: processed: [%6d, %6d)\n", __func__, i, std::min(i + n_batch, n_tokens_all));163 164        if (i + n_batch >= n_tokens_all) {165            break;166        }167    }168 169    for (int i = n_ctx; i < n_tokens_all; i += n_batch) {170        const int n_discard = n_batch;171 172        LOG_INF("%s: shifting KV cache with %d\n", __func__, n_discard);173 174        llama_memory_seq_rm (mem, 0, n_keep            , n_keep + n_discard);175        llama_memory_seq_add(mem, 0, n_keep + n_discard, n_ctx,  -n_discard);176 177        n_past = llama_memory_seq_pos_max(mem, 0) + 1;178 179        common_batch_clear(batch);180 181        for (int j = 0; j < n_batch && i + j < n_tokens_all; j++) {182            common_batch_add(batch, tokens_list[i + j], n_past++, { 0 }, false);183        }184 185        if (i + n_batch >= n_tokens_all) {186            batch.logits[batch.n_tokens - 1] = true;187        }188 189        if (llama_decode(ctx, batch) != 0) {190            LOG_ERR("%s: llama_decode() failed\n", __func__);191            return 1;192        }193 194        LOG_INF("%s: processed: [%6d, %6d)\n", __func__, i, std::min(i + n_batch, n_tokens_all));195    }196 197    {198        const int n_discard = n_past - n_ctx + n_predict;199 200        if (n_discard > 0) {201            LOG_INF("%s: shifting KV cache with %d to free space for the answer\n", __func__, n_discard);202 203            llama_memory_seq_rm (mem, 0, n_keep            , n_keep + n_discard);204            llama_memory_seq_add(mem, 0, n_keep + n_discard, n_ctx,  -n_discard);205 206            n_past = llama_memory_seq_pos_max(mem, 0) + 1;207        }208    }209 210    LOG_INF("\n");211    LOG_INF("%s: passkey = %d, inserted at position %d / %d (token pos: ~%d)\n", __func__, passkey, i_pos, n_junk, (i_pos * n_tokens_all) / n_junk);212    LOG_INF("\n");213 214    // main loop215 216    int n_cur    = n_tokens_all;217    int n_decode = 0;218 219    LOG_INF("%s", prompt_suffix.c_str());220 221    const auto t_main_start = ggml_time_us();222 223    while (n_cur <= n_len) {224        // sample the next token225        {226            const llama_token new_token_id = llama_sampler_sample(smpl, ctx, batch.n_tokens - 1);227 228            // is it an end of generation?229            if (llama_vocab_is_eog(vocab, new_token_id) || n_cur == n_len) {230                LOG("\n");231 232                break;233            }234 235            LOG("%s", common_token_to_piece(ctx, new_token_id).c_str());236 237            n_decode += 1;238 239            // prepare the next batch240            common_batch_clear(batch);241 242            // push this new token for next evaluation243            common_batch_add(batch, new_token_id, n_past++, { 0 }, true);244        }245 246        n_cur += 1;247 248        // evaluate the current batch with the transformer model249        if (llama_decode(ctx, batch)) {250            LOG_ERR("%s : failed to eval, return code %d\n", __func__, 1);251            return 1;252        }253    }254 255    LOG("\n");256 257    const auto t_main_end = ggml_time_us();258 259    LOG_INF("%s: decoded %d tokens in %.2f s, speed: %.2f t/s\n",260            __func__, n_decode, (t_main_end - t_main_start) / 1000000.0f, n_decode / ((t_main_end - t_main_start) / 1000000.0f));261 262    LOG("\n");263    llama_perf_context_print(ctx);264 265    LOG("\n");266 267    llama_sampler_free(smpl);268 269    llama_batch_free(batch);270 271    llama_free(ctx);272    llama_model_free(model);273 274    llama_backend_free();275 276    return 0;277}278