echodict/llama.cpp

1#include "llama.h"2#include <clocale>3#include <cstdio>4#include <cstring>5#include <string>6#include <vector>7 8static void print_usage(int, char ** argv) {9    printf("\nexample usage:\n");10    printf("\n    %s -m model.gguf [-n n_predict] [-ngl n_gpu_layers] [prompt]\n", argv[0]);11    printf("\n");12}13 14int main(int argc, char ** argv) {15    std::setlocale(LC_NUMERIC, "C");16 17    // path to the model gguf file18    std::string model_path;19    // prompt to generate text from20    std::string prompt = "Hello my name is";21    // number of layers to offload to the GPU22    int ngl = 99;23    // number of tokens to predict24    int n_predict = 32;25 26    // parse command line arguments27 28    {29        int i = 1;30        for (; i < argc; i++) {31            if (strcmp(argv[i], "-m") == 0) {32                if (i + 1 < argc) {33                    model_path = argv[++i];34                } else {35                    print_usage(argc, argv);36                    return 1;37                }38            } else if (strcmp(argv[i], "-n") == 0) {39                if (i + 1 < argc) {40                    try {41                        n_predict = std::stoi(argv[++i]);42                    } catch (...) {43                        print_usage(argc, argv);44                        return 1;45                    }46                } else {47                    print_usage(argc, argv);48                    return 1;49                }50            } else if (strcmp(argv[i], "-ngl") == 0) {51                if (i + 1 < argc) {52                    try {53                        ngl = std::stoi(argv[++i]);54                    } catch (...) {55                        print_usage(argc, argv);56                        return 1;57                    }58                } else {59                    print_usage(argc, argv);60                    return 1;61                }62            } else {63                // prompt starts here64                break;65            }66        }67        if (model_path.empty()) {68            print_usage(argc, argv);69            return 1;70        }71        if (i < argc) {72            prompt = argv[i++];73            for (; i < argc; i++) {74                prompt += " ";75                prompt += argv[i];76            }77        }78    }79 80    // load dynamic backends81 82    ggml_backend_load_all();83 84    // initialize the model85 86    llama_model_params model_params = llama_model_default_params();87    model_params.n_gpu_layers = ngl;88 89    llama_model * model = llama_model_load_from_file(model_path.c_str(), model_params);90 91    if (model == NULL) {92        fprintf(stderr , "%s: error: unable to load model\n" , __func__);93        return 1;94    }95 96    const llama_vocab * vocab = llama_model_get_vocab(model);97    // tokenize the prompt98 99    // find the number of tokens in the prompt100    const int n_prompt = -llama_tokenize(vocab, prompt.c_str(), prompt.size(), NULL, 0, true, true);101 102    // allocate space for the tokens and tokenize the prompt103    std::vector<llama_token> prompt_tokens(n_prompt);104    if (llama_tokenize(vocab, prompt.c_str(), prompt.size(), prompt_tokens.data(), prompt_tokens.size(), true, true) < 0) {105        fprintf(stderr, "%s: error: failed to tokenize the prompt\n", __func__);106        return 1;107    }108 109    // initialize the context110 111    llama_context_params ctx_params = llama_context_default_params();112    // n_ctx is the context size113    ctx_params.n_ctx = n_prompt + n_predict - 1;114    // n_batch is the maximum number of tokens that can be processed in a single call to llama_decode115    ctx_params.n_batch = n_prompt;116    // enable performance counters117    ctx_params.no_perf = false;118 119    llama_context * ctx = llama_init_from_model(model, ctx_params);120 121    if (ctx == NULL) {122        fprintf(stderr , "%s: error: failed to create the llama_context\n" , __func__);123        return 1;124    }125 126    // initialize the sampler127 128    auto sparams = llama_sampler_chain_default_params();129    sparams.no_perf = false;130    llama_sampler * smpl = llama_sampler_chain_init(sparams);131 132    llama_sampler_chain_add(smpl, llama_sampler_init_greedy());133 134    // print the prompt token-by-token135 136    for (auto id : prompt_tokens) {137        char buf[128];138        int n = llama_token_to_piece(vocab, id, buf, sizeof(buf), 0, true);139        if (n < 0) {140            fprintf(stderr, "%s: error: failed to convert token to piece\n", __func__);141            return 1;142        }143        std::string s(buf, n);144        printf("%s", s.c_str());145    }146 147    // prepare a batch for the prompt148 149    llama_batch batch = llama_batch_get_one(prompt_tokens.data(), prompt_tokens.size());150 151    if (llama_model_has_encoder(model)) {152        if (llama_encode(ctx, batch)) {153            fprintf(stderr, "%s : failed to eval\n", __func__);154            return 1;155        }156 157        llama_token decoder_start_token_id = llama_model_decoder_start_token(model);158        if (decoder_start_token_id == LLAMA_TOKEN_NULL) {159            decoder_start_token_id = llama_vocab_bos(vocab);160        }161 162        batch = llama_batch_get_one(&decoder_start_token_id, 1);163    }164 165    // main loop166 167    const auto t_main_start = ggml_time_us();168    int n_decode = 0;169    llama_token new_token_id;170 171    for (int n_pos = 0; n_pos + batch.n_tokens < n_prompt + n_predict; ) {172        // evaluate the current batch with the transformer model173        if (llama_decode(ctx, batch)) {174            fprintf(stderr, "%s : failed to eval, return code %d\n", __func__, 1);175            return 1;176        }177 178        n_pos += batch.n_tokens;179 180        // sample the next token181        {182            new_token_id = llama_sampler_sample(smpl, ctx, -1);183 184            // is it an end of generation?185            if (llama_vocab_is_eog(vocab, new_token_id)) {186                break;187            }188 189            char buf[128];190            int n = llama_token_to_piece(vocab, new_token_id, buf, sizeof(buf), 0, true);191            if (n < 0) {192                fprintf(stderr, "%s: error: failed to convert token to piece\n", __func__);193                return 1;194            }195            std::string s(buf, n);196            printf("%s", s.c_str());197            fflush(stdout);198 199            // prepare the next batch with the sampled token200            batch = llama_batch_get_one(&new_token_id, 1);201 202            n_decode += 1;203        }204    }205 206    printf("\n");207 208    const auto t_main_end = ggml_time_us();209 210    fprintf(stderr, "%s: decoded %d tokens in %.2f s, speed: %.2f t/s\n",211            __func__, n_decode, (t_main_end - t_main_start) / 1000000.0f, n_decode / ((t_main_end - t_main_start) / 1000000.0f));212 213    fprintf(stderr, "\n");214    llama_perf_sampler_print(smpl);215    llama_perf_context_print(ctx);216    fprintf(stderr, "\n");217 218    llama_sampler_free(smpl);219    llama_free(ctx);220    llama_model_free(model);221 222    return 0;223}224