echodict/llama.cpp

1#include "arg.h"2#include "ggml.h"3#include "common.h"4#include "ngram-cache.h"5#include "sampling.h"6#include "log.h"7#include "llama.h"8 9#include <clocale>10#include <cstdint>11#include <cstdio>12#include <fstream>13#include <string>14#include <vector>15 16int main(int argc, char ** argv){17    std::setlocale(LC_NUMERIC, "C");18 19    common_params params;20 21    common_init();22 23    if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_LOOKUP)) {24        return 1;25    }26 27    // max. number of additional tokens to draft if match is found28    const int n_draft = params.speculative.n_max;29 30    // init llama.cpp31    llama_backend_init();32    llama_numa_init(params.numa);33 34    // load the model35    auto llama_init = common_init_from_params(params);36 37    auto * model = llama_init->model();38    auto * ctx   = llama_init->context();39 40    const llama_vocab * vocab = llama_model_get_vocab(model);41 42    // tokenize the prompt43    std::vector<llama_token> inp;44    inp = common_tokenize(ctx, params.prompt, true, true);45 46    common_ngram_cache ngram_cache_context;47    common_ngram_cache ngram_cache_dynamic;48    common_ngram_cache ngram_cache_static;49    int64_t t_draft_flat_us = 0;50    int64_t t_draft_us = 0;51 52    {53        // Fill up context ngram cache with tokens from user input:54        const int64_t t_start_draft_us = ggml_time_us();55        common_ngram_cache_update(ngram_cache_context, LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX, inp, inp.size(), false);56 57        if (!params.speculative.lookup_cache_static.empty()) {58            try {59                ngram_cache_static = common_ngram_cache_load(params.speculative.lookup_cache_static);60            } catch (std::ifstream::failure const &) {61                LOG_ERR("failed to open static lookup cache: %s", params.speculative.lookup_cache_static.c_str());62                exit(1);63            }64        }65 66        if (!params.speculative.lookup_cache_dynamic.empty()) {67            try {68                ngram_cache_dynamic = common_ngram_cache_load(params.speculative.lookup_cache_dynamic);69            } catch (std::ifstream::failure const &) {} // if the file does not exist it will simply be created at the end of the program70        }71 72        t_draft_flat_us += ggml_time_us() - t_start_draft_us;73    }74 75    const int max_context_size     = llama_n_ctx(ctx);76    const int max_tokens_list_size = max_context_size - 4;77 78    if ((int) inp.size() > max_tokens_list_size) {79        LOG_ERR("%s: prompt too long (%d tokens, max %d)\n", __func__, (int) inp.size(), max_tokens_list_size);80        return 1;81    }82 83    LOG("\n\n");84 85    for (auto id : inp) {86        LOG("%s", common_token_to_piece(ctx, id).c_str());87    }88 89    fflush(stderr);90 91    const int n_input = inp.size();92 93    const auto t_enc_start = ggml_time_us();94 95    llama_decode(ctx, llama_batch_get_one( inp.data(), n_input - 1));96    llama_decode(ctx, llama_batch_get_one(&inp.back(),           1));97 98    const auto t_enc_end = ggml_time_us();99 100    int n_predict = 0;101    int n_drafted = 0;102    int n_accept  = 0;103 104    int n_past = inp.size();105 106    bool has_eos = false;107 108    struct common_sampler * smpl = common_sampler_init(model, params.sampling);109 110    std::vector<llama_token> draft;111 112    llama_batch batch_tgt = llama_batch_init(llama_n_ctx(ctx), 0, 1);113 114    const auto t_dec_start = ggml_time_us();115 116    while (true) {117        // print current draft sequence118        LOG_DBG("drafted %s\n", string_from(ctx, draft).c_str());119 120        int i_dft = 0;121        while (true) {122            // sample from the target model123            llama_token id = common_sampler_sample(smpl, ctx, i_dft);124 125            common_sampler_accept(smpl, id, true);126 127            const std::string token_str = common_token_to_piece(ctx, id);128 129            if (!params.use_color) {130                LOG("%s", token_str.c_str());131            }132 133            if (llama_vocab_is_eog(vocab, id)) {134                has_eos = true;135            }136 137            ++n_predict;138 139            // check if the target token matches the draft140            if (i_dft < (int) draft.size() && id == draft[i_dft]) {141                LOG_DBG("the sampled target token matches the %dth drafted token (%d, '%s') - accepted\n", i_dft, id, token_str.c_str());142                ++n_accept;143                ++n_past;144                ++i_dft;145                inp.push_back(id);146                {147                    // Update context ngram cache with the newly accepted token:148                    const int64_t t_start_draft_us = ggml_time_us();149                    common_ngram_cache_update(ngram_cache_context, LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX, inp, 1, false);150                    t_draft_us += ggml_time_us() - t_start_draft_us;151                }152 153                if (params.use_color) {154                    // color accepted draft token155                    LOG("\033[34m%s\033[0m", token_str.c_str());156                    fflush(stdout);157                }158                continue;159            }160 161            if (params.use_color) {162                LOG("%s", token_str.c_str());163            }164            fflush(stdout);165 166 167            LOG_DBG("the sampled target token (%d, '%s') did not match, or we ran out of drafted tokens\n", id, token_str.c_str());168 169            draft.clear();170            draft.push_back(id);171            inp.push_back(id);172            {173                // Update context ngram cache with the newly accepted token:174                const int64_t t_start_draft_us = ggml_time_us();175                common_ngram_cache_update(ngram_cache_context, LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX, inp, 1, false);176                t_draft_us += ggml_time_us() - t_start_draft_us;177            }178            break;179        }180 181        if ((params.n_predict > 0 && n_predict > params.n_predict) || has_eos) {182            break;183        }184 185        // KV cache management186        // clean the cache of draft tokens that weren't accepted187        llama_memory_seq_rm(llama_get_memory(ctx), 0, n_past, -1);188 189        common_batch_clear(batch_tgt);190        common_batch_add(batch_tgt, draft[0], n_past, { 0 }, true);191 192        // Draft already contains a single token sampled from the model:193        GGML_ASSERT(draft.size() == 1);194        GGML_ASSERT(draft[0] == inp.back());195        const int64_t t_start_draft_us = ggml_time_us();196 197        common_ngram_cache_draft(inp, draft, n_draft, LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX, ngram_cache_context, ngram_cache_dynamic, ngram_cache_static);198 199        for (size_t i = 1; i < draft.size(); ++i) {200            common_batch_add(batch_tgt, draft[i], n_past + i, { 0 }, true);201        }202 203        t_draft_us += ggml_time_us() - t_start_draft_us;204        n_drafted += draft.size() - 1;205 206        llama_decode(ctx, batch_tgt);207        ++n_past;208 209        draft.erase(draft.begin());210    }211 212    auto t_dec_end = ggml_time_us();213 214    // Update dynamic ngram cache with context ngram cache and save it to disk:215    common_ngram_cache_merge(ngram_cache_dynamic, ngram_cache_context);216    common_ngram_cache_save(ngram_cache_dynamic, params.speculative.lookup_cache_dynamic);217 218    LOG("\n\n");219 220    LOG_INF("encoded %4d tokens in %8.3f seconds, speed: %8.3f t/s\n", n_input,   (t_enc_end - t_enc_start) / 1e6f, inp.size() / ((t_enc_end - t_enc_start) / 1e6f));221    LOG_INF("decoded %4d tokens in %8.3f seconds, speed: %8.3f t/s\n", n_predict, (t_dec_end - t_dec_start) / 1e6f, n_predict  / ((t_dec_end - t_dec_start) / 1e6f));222 223    LOG_INF("\n");224    LOG_INF("n_draft      = %d\n", n_draft);225    LOG_INF("n_predict    = %d\n", n_predict);226    LOG_INF("n_drafted    = %d\n", n_drafted);227    LOG_INF("t_draft_flat = %.2f ms\n", t_draft_flat_us*1e-3);228    LOG_INF("t_draft      = %.2f ms, %.2f us per token, %.2f tokens per second\n",229            t_draft_us*1e-3, 1.0f*t_draft_us/n_drafted, n_drafted/(1e-6*t_draft_us));230    LOG_INF("n_accept     = %d\n", n_accept);231    LOG_INF("accept       = %.3f%%\n", 100.0f * n_accept / n_drafted);232 233    LOG_INF("\ntarget:\n\n");234    common_perf_print(ctx, smpl);235 236    common_sampler_free(smpl);237 238    llama_batch_free(batch_tgt);239 240    llama_backend_free();241 242    LOG("\n\n");243 244    return 0;245}246