echodict/llama.cpp

1#include "arg.h"2#include "common.h"3#include "log.h"4#include "ngram-cache.h"5#include "llama.h"6#include "ggml.h"7 8#include <cinttypes>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    const int n_draft = params.speculative.n_max;28 29    // init llama.cpp30    llama_backend_init();31    llama_numa_init(params.numa);32 33    // load the model34    auto llama_init = common_init_from_params(params);35 36    llama_context * ctx = llama_init->context();37 38    // tokenize the prompt39    std::vector<llama_token> inp;40    inp = common_tokenize(ctx, params.prompt, true, true);41 42    common_ngram_cache ngram_cache_context;43    common_ngram_cache ngram_cache_dynamic;44    common_ngram_cache ngram_cache_static;45 46    int64_t t_draft_flat_us = 0;47    int64_t t_draft_us = 0;48 49    {50        const int64_t t_start_draft_us = ggml_time_us();51 52        if (!params.speculative.lookup_cache_static.empty()) {53            try {54                ngram_cache_static = common_ngram_cache_load(params.speculative.lookup_cache_static);55            } catch (std::ifstream::failure const &) {56                LOG_ERR("failed to open static lookup cache: %s", params.speculative.lookup_cache_static.c_str());57                exit(1);58            }59        }60 61        if (!params.speculative.lookup_cache_dynamic.empty()) {62            try {63                ngram_cache_dynamic = common_ngram_cache_load(params.speculative.lookup_cache_dynamic);64            } catch (std::ifstream::failure const &) {} // if the file does not exist it will simply be created at the end of the program65        }66 67        t_draft_flat_us += ggml_time_us() - t_start_draft_us;68    }69 70    const int n_input = inp.size();71    const int n_ctx = llama_n_ctx(ctx);72 73    int n_drafted = 0;74    int n_accept  = 0;75 76    const int64_t t_start_ms = ggml_time_ms();77 78    // Iterate over input tokens in chunks of size n_ctx.79    // Each chunk is treated as if a sequential generation but with pre-determined tokens to ensure reproducibility.80    for (int i_start = 0; i_start + n_ctx < n_input; i_start += n_ctx) {81        const std::vector<llama_token> inp_slice(inp.begin() + i_start, inp.begin() + i_start + n_ctx);82        std::vector<llama_token> pseudo_output;83        pseudo_output.push_back(inp_slice[0]);84 85        while ((int) pseudo_output.size() < n_ctx) {86            // Simulate drafting and decoding from draft:87            std::vector<llama_token> draft;88            draft.push_back(pseudo_output.back());89 90            {91                const int64_t t_start_draft_us = ggml_time_us();92                common_ngram_cache_draft(pseudo_output, draft, n_draft, LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX, ngram_cache_context, ngram_cache_dynamic, ngram_cache_static);93                t_draft_us += ggml_time_us() - t_start_draft_us;94            }95 96            n_drafted += draft.size() - 1;97 98            for (size_t j = 1; j < draft.size() && (int) pseudo_output.size() < n_ctx; ++j) {99                const llama_token ground_truth = inp_slice[pseudo_output.size()];100                const llama_token drafted = draft[j];101 102                if (ground_truth != drafted) {103                    break;104                }105 106                ++n_accept;107                pseudo_output.push_back(ground_truth);108 109                {110                    const int64_t t_start_draft_us = ggml_time_us();111                    common_ngram_cache_update(ngram_cache_context, LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX, pseudo_output, 1, false);112                    t_draft_us += ggml_time_us() - t_start_draft_us;113                }114            }115 116            // After each simulated batch decoding simulate the sampling of a single token:117            if ((int) pseudo_output.size() < n_ctx) {118                pseudo_output.push_back(inp_slice[pseudo_output.size()]);119                {120                    const int64_t t_start_draft_us = ggml_time_us();121                    common_ngram_cache_update(ngram_cache_context, LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX, pseudo_output, 1, false);122                    t_draft_us += ggml_time_us() - t_start_draft_us;123                }124            }125 126            draft.erase(draft.begin());127 128        }129        if (i_start > 0 && i_start / 100000 != (i_start - n_ctx) / 100000) {130            const int64_t t_now_ms = ggml_time_ms();131            const int64_t eta_ms   = (n_input - i_start) * (t_now_ms - t_start_ms) / i_start;132            const int64_t eta_min  = eta_ms / (60*1000);133            const int64_t eta_s    = (eta_ms - 60*1000*eta_min) / 1000;134 135            LOG_INF("lookup-stats: %d/%d done, ETA: %02" PRId64 ":%02" PRId64 "\n", i_start, n_input, eta_min, eta_s);136        }137 138        // After each chunk, update the dynamic ngram cache with the context ngram cache:139        common_ngram_cache_merge(ngram_cache_dynamic, ngram_cache_context);140        ngram_cache_context.clear();141    }142 143    LOG("\n");144 145    LOG_INF("\n");146    LOG_INF("n_draft      = %d\n", n_draft);147    LOG_INF("n_predict    = %d\n", n_input - n_input % n_ctx);148    LOG_INF("n_drafted    = %d\n", n_drafted);149    LOG_INF("t_draft_flat = %.2f ms\n", t_draft_flat_us*1e-3);150    LOG_INF("t_draft      = %.2f ms, %.2f us per token, %.2f tokens per second\n",151            t_draft_us*1e-3, 1.0f*t_draft_us/n_drafted, n_drafted/(1e-6*t_draft_us));152    LOG_INF("n_accept     = %d\n", n_accept);153    LOG_INF("accept       = %.3f%%\n", 100.0f * n_accept / n_drafted);154 155    llama_backend_free();156 157    LOG("\n\n");158 159    return 0;160}161