llama.cpp/examples/lookup/lookup.cpp

#include "common.h"
#include "llama.h"

#include <cmath>
#include <cstdio>
#include <string>
#include <vector>

int main(int argc, char ** argv){
    gpt_params params;

    if(gpt_params_parse(argc, argv, params) == false){
        return 1;
    }

    // maximum n-grams to search for in prompt
    const int max_ngram_size = 3;

    // length of the candidate / draft sequence, if match is found
    const int n_draft = 10;

#ifndef LOG_DISABLE_LOGS
    log_set_target(log_filename_generator("lookup", "log"));
    LOG_TEE("Log start\n");
    log_dump_cmdline(argc, argv);
#endif // LOG_DISABLE_LOGS

    // init llama.cpp
    llama_backend_init(params.numa);

    llama_model * model = NULL;
    llama_context * ctx = NULL;

    // load the model
    std::tie(model, ctx) = llama_init_from_gpt_params(params);

    // tokenize the prompt
    const bool add_bos = llama_should_add_bos_token(model);
    LOG("add_bos tgt: %d\n", add_bos);
    
    std::vector<llama_token> inp;
    inp = ::llama_tokenize(ctx, params.prompt, add_bos, true);

    const int max_context_size     = llama_n_ctx(ctx);
    const int max_tokens_list_size = max_context_size - 4;

    if ((int) inp.size() > max_tokens_list_size) {
        fprintf(stderr, "%s: error: prompt too long (%d tokens, max %d)\n", __func__, (int) inp.size(), max_tokens_list_size);
        return 1;
    }

    fprintf(stderr, "\n\n");

    for (auto id : inp) {
        fprintf(stderr, "%s", llama_token_to_piece(ctx, id).c_str());
    }

    fflush(stderr);

    const int n_input = inp.size();

    const auto t_enc_start = ggml_time_us();

    llama_decode(ctx, llama_batch_get_one( inp.data(), n_input - 1, 0,           0));
    llama_decode(ctx, llama_batch_get_one(&inp.back(),           1, n_input - 1, 0));

    const auto t_enc_end = ggml_time_us();

    int n_predict = 0;
    int n_drafted = 0;
    int n_accept  = 0;
    
    int n_past = inp.size();

    bool has_eos = false;

    struct llama_sampling_context * ctx_sampling = llama_sampling_init(params.sparams);

    std::vector<llama_token> draft(n_draft);

    llama_batch batch_tgt = llama_batch_init(params.n_ctx, 0, 1);

    const auto t_dec_start = ggml_time_us();

    while(true){
        // print current draft sequence
        LOG("drafted %s\n", LOG_TOKENS_TOSTR_PRETTY(ctx, draft).c_str());

        int i_dft  = 0;        
        while (true) {
            //LOG("sampling target: s_keep = %3d, i_dft = %3d, i_batch_tgt = %3d\n", s_keep, i_dft, drafts[s_keep].i_batch_tgt[i_dft]);

            // sample from the target model
            llama_token id = llama_sampling_sample(ctx_sampling, ctx, NULL, 0);

            llama_sampling_accept(ctx_sampling, ctx, id, true);

            //LOG("last: %s\n", LOG_TOKENS_TOSTR_PRETTY(ctx_tgt, ctx_sampling->prev).c_str());

            const std::string token_str = llama_token_to_piece(ctx, id);

            printf("%s", token_str.c_str());
            fflush(stdout);

            if (id == llama_token_eos(model)) {
                has_eos = true;
            }

            ++n_predict;

            // check if the target token matches the draft
            if (i_dft < (int) draft.size() && id == draft[i_dft]) {
                LOG("the sampled target token matches the %dth drafted token (%d, '%s') - accepted\n", i_dft, id, token_str.c_str());
                ++n_accept;
                ++n_past;
                ++i_dft;
                
                continue;
            } 
            
            LOG("the sampled target token (%d, '%s') did not match, or we ran out of drafted tokens\n", id, token_str.c_str());

            draft.clear();
            draft.push_back(id);
            // drafts[0].i_batch_tgt.push_back(0);

            // llama_batch_clear(batch_dft);
            // llama_batch_add  (batch_dft, id, n_past_dft, { 0 }, true);

            // llama_kv_cache_seq_rm(ctx_dft, 0, n_past_dft, -1);
            // // LOG("dft batch: %s\n", LOG_BATCH_TOSTR_PRETTY(ctx_dft, batch_dft).c_str());
            // llama_decode         (ctx_dft, batch_dft);

            // ++n_past_dft;
            break;
        }

        if (n_predict > params.n_predict || has_eos) {
            break;
        }

        llama_batch_clear(batch_tgt);
        llama_batch_add(batch_tgt, draft[0], n_past, { 0 }, true);

        bool match = false;
        // generate n_pred tokens through prompt lookup
        for (int ngram_size = max_ngram_size ; ngram_size > 0; --ngram_size){
            if (match){
                break;
            }
            const auto & prev = ctx_sampling->prev;
            int prev_size = prev.size();
            const llama_token * ngram = &prev[prev_size - ngram_size];
            
            for (int i = 0; i <= (int) prev_size - (ngram_size * 2); ++i) {
                if (prev[i] == ngram[0] && prev[i + 1] == ngram[1] && prev[i + 2] == ngram[2]) {
                    const int startIdx = i + ngram_size;
                    const int endIdx = startIdx + n_draft;
                    if (endIdx < prev_size){
                        match = true;
                        for (int j = startIdx; j < endIdx; ++j) {
                            LOG(" - draft candidate %d: %d\n", j, prev[j]);
                            draft.push_back(prev[j]);
                            llama_batch_add(batch_tgt, prev[j], n_past + j + 1, { 1 }, true);
                            ++n_drafted;
                        }
                    }
                }
            }
        }

        llama_decode(ctx, batch_tgt);
        ++n_past;

        draft.erase(draft.begin());
        
        // we have our draft!

    }

    auto t_dec_end = ggml_time_us();

    LOG_TEE("\n\n");

    LOG_TEE("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));
    LOG_TEE("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));

    LOG_TEE("\n");
    LOG_TEE("n_draft   = %d\n", n_draft);
    LOG_TEE("n_predict = %d\n", n_predict);
    LOG_TEE("n_drafted = %d\n", n_drafted);
    LOG_TEE("n_accept  = %d\n", n_accept);
    LOG_TEE("accept    = %.3f%%\n", 100.0f * n_accept / n_drafted);

    LOG_TEE("\ntarget:\n");
    llama_print_timings(ctx);

    llama_sampling_free(ctx_sampling);
    llama_batch_free(batch_tgt);

    llama_free(ctx);
    llama_free_model(model);

    llama_backend_free();

    fprintf(stderr, "\n\n");

    return 0;
}