Brunobkr/llama.cpp_AlgMor24_github

1#include "arg.h"2#include "common.h"3#include "log.h"4#include "llama.h"5#include "sampling.h"6 7#include <algorithm>8#include <clocale>9#include <cstdio>10#include <string>11#include <vector>12 13static void print_usage(int, char ** argv) {14    LOG("\nexample usage:\n");15    LOG("\n    %s -m model.gguf -p \"Hello my name is\" -n 32 -np 4\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.prompt = "Hello my name is";25    params.n_predict = 32;26 27    common_init();28 29    if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_BATCHED, print_usage)) {30        return 1;31    }32 33    // number of parallel batches34    int n_parallel = params.n_parallel;35 36    // total length of the sequences including the prompt37    int n_predict = params.n_predict;38 39    // init LLM40 41    llama_backend_init();42    llama_numa_init(params.numa);43 44    // initialize the model45 46    llama_model_params model_params = common_model_params_to_llama(params);47 48    llama_model * model = llama_model_load_from_file(params.model.path.c_str(), model_params);49 50    if (model == NULL) {51        LOG_ERR("%s: error: unable to load model\n" , __func__);52        return 1;53    }54 55    const llama_vocab * vocab = llama_model_get_vocab(model);56 57    // tokenize the prompt58 59    std::vector<llama_token> tokens_list;60    tokens_list = common_tokenize(vocab, params.prompt, true);61 62    const int n_kv_req = tokens_list.size() + (n_predict - tokens_list.size())*n_parallel;63 64    // initialize the context65 66    llama_context_params ctx_params = common_context_params_to_llama(params);67 68    ctx_params.n_ctx   = n_kv_req;69    ctx_params.n_batch = std::max(n_predict, n_parallel);70 71    auto sparams = llama_sampler_chain_default_params();72    sparams.no_perf = false;73 74    std::vector<llama_sampler_seq_config> sampler_configs;75 76    for (int32_t i = 0; i < n_parallel; ++i) {77        llama_sampler * smpl = llama_sampler_chain_init(sparams);78 79        llama_sampler_chain_add(smpl, llama_sampler_init_top_k(params.sampling.top_k));80        llama_sampler_chain_add(smpl, llama_sampler_init_top_p(params.sampling.top_p, params.sampling.min_keep));81        llama_sampler_chain_add(smpl, llama_sampler_init_temp (params.sampling.temp));82        llama_sampler_chain_add(smpl, llama_sampler_init_dist (params.sampling.seed));83 84        sampler_configs.push_back({ i, smpl });85    }86 87    if (params.sampling.backend_sampling) {88        ctx_params.samplers   = sampler_configs.data();89        ctx_params.n_samplers = sampler_configs.size();90    }91 92    llama_context * ctx = llama_init_from_model(model, ctx_params);93 94    if (ctx == NULL) {95        LOG_ERR("%s: error: failed to create the llama_context\n" , __func__);96        return 1;97    }98 99    const int n_ctx = llama_n_ctx(ctx);100 101    LOG_INF("\n%s: n_predict = %d, n_ctx = %d, n_batch = %u, n_parallel = %d, n_kv_req = %d\n", __func__, n_predict, n_ctx, ctx_params.n_batch, n_parallel, n_kv_req);102 103    // make sure the KV cache is big enough to hold all the prompt and generated tokens104    if (n_kv_req > n_ctx) {105        LOG_ERR("%s: error: n_kv_req (%d) > n_ctx, the required KV cache size is not big enough\n", __func__,  n_kv_req);106        LOG_ERR("%s:        either reduce n_parallel or increase n_ctx\n", __func__);107        return 1;108    }109 110    // print the prompt token-by-token111 112    LOG("\n");113 114    for (auto id : tokens_list) {115        LOG("%s", common_token_to_piece(ctx, id).c_str());116    }117 118    // create a llama_batch119    // we use this object to submit token data for decoding120    llama_batch batch = llama_batch_init(std::max(tokens_list.size(), (size_t) n_parallel), 0, n_parallel);121 122    std::vector<llama_seq_id> seq_ids(n_parallel, 0);123    for (int32_t i = 0; i < n_parallel; ++i) {124        seq_ids[i] = i;125    }126 127    // evaluate the initial prompt128    for (size_t i = 0; i < tokens_list.size(); ++i) {129        common_batch_add(batch, tokens_list[i], i, seq_ids, false);130    }131    GGML_ASSERT(batch.n_tokens == (int) tokens_list.size());132 133    if (llama_model_has_encoder(model)) {134        if (llama_encode(ctx, batch)) {135            LOG_ERR("%s : failed to eval\n", __func__);136            return 1;137        }138 139        llama_token decoder_start_token_id = llama_model_decoder_start_token(model);140        if (decoder_start_token_id == LLAMA_TOKEN_NULL) {141            decoder_start_token_id = llama_vocab_bos(vocab);142        }143 144        common_batch_clear(batch);145        common_batch_add(batch, decoder_start_token_id, 0, seq_ids, false);146    }147 148    // llama_decode will output logits only for the last token of the prompt149    batch.logits[batch.n_tokens - 1] = true;150 151    if (llama_decode(ctx, batch) != 0) {152        LOG_ERR("%s: llama_decode() failed\n", __func__);153        return 1;154    }155 156    //// assign the system KV cache to all parallel sequences157    //// this way, the parallel sequences will "reuse" the prompt tokens without having to copy them158    //for (int32_t i = 1; i < n_parallel; ++i) {159    //    llama_kv_cache_seq_cp(ctx, 0, i, -1, -1);160    //}161 162    if (n_parallel > 1) {163        LOG("\n\n%s: generating %d sequences ...\n", __func__, n_parallel);164    }165 166    // main loop167 168    // we will store the parallel decoded sequences in this vector169    std::vector<std::string> streams(n_parallel);170 171    // remember the batch index of the last token for each parallel sequence172    // we need this to determine which logits to sample from173    std::vector<int32_t> i_batch(n_parallel, batch.n_tokens - 1);174 175    int n_cur    = batch.n_tokens;176    int n_decode = 0;177 178    const auto t_main_start = ggml_time_us();179 180    while (n_cur <= n_predict) {181        // prepare the next batch182        common_batch_clear(batch);183 184        // sample the next token for each parallel sequence / stream185        for (int32_t i = 0; i < n_parallel; ++i) {186            if (i_batch[i] < 0) {187                // the stream has already finished188                continue;189            }190 191            const llama_token new_token_id = llama_sampler_sample(sampler_configs[i].sampler, ctx, i_batch[i]);192 193            // is it an end of generation? -> mark the stream as finished194            if (llama_vocab_is_eog(vocab, new_token_id) || n_cur == n_predict) {195                i_batch[i] = -1;196                LOG("\n");197                if (n_parallel > 1) {198                    LOG_INF("%s: stream %d finished at n_cur = %d", __func__, i, n_cur);199                }200 201                continue;202            }203 204            // if there is only one stream, we print immediately to stdout205            if (n_parallel == 1) {206                LOG("%s", common_token_to_piece(ctx, new_token_id).c_str());207            }208 209            streams[i] += common_token_to_piece(ctx, new_token_id);210 211            i_batch[i] = batch.n_tokens;212 213            // push this new token for next evaluation214            common_batch_add(batch, new_token_id, n_cur, { i }, true);215 216            n_decode += 1;217        }218 219        // all streams are finished220        if (batch.n_tokens == 0) {221            break;222        }223 224        n_cur += 1;225 226        // evaluate the current batch with the transformer model227        if (llama_decode(ctx, batch)) {228            LOG_ERR("%s : failed to eval, return code %d\n", __func__, 1);229            return 1;230        }231    }232 233    if (n_parallel > 1) {234        LOG("\n");235 236        for (int32_t i = 0; i < n_parallel; ++i) {237            LOG("sequence %d:\n\n%s%s\n\n", i, params.prompt.c_str(), streams[i].c_str());238        }239    }240 241    const auto t_main_end = ggml_time_us();242 243    LOG_INF("%s: decoded %d tokens in %.2f s, speed: %.2f t/s\n",244            __func__, n_decode, (t_main_end - t_main_start) / 1000000.0f, n_decode / ((t_main_end - t_main_start) / 1000000.0f));245 246    LOG("\n");247    llama_perf_sampler_print(sampler_configs[0].sampler);248    llama_perf_context_print(ctx);249 250    fprintf(stderr, "\n");251 252    llama_batch_free(batch);253 254    for (auto & sampler_config : sampler_configs) {255        llama_sampler_free(sampler_config.sampler);256    }257 258    llama_free(ctx);259    llama_model_free(model);260 261    llama_backend_free();262 263    return 0;264}265