KBaba7/llama.cpp
0
1#include "llama.h"2#include <cstdio>3#include <cstring>4#include <iostream>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 [-c context_size] [-ngl n_gpu_layers]\n", argv[0]);11 printf("\n");12}13 14int main(int argc, char ** argv) {15 std::string model_path;16 int ngl = 99;17 int n_ctx = 2048;18 19 // parse command line arguments20 for (int i = 1; i < argc; i++) {21 try {22 if (strcmp(argv[i], "-m") == 0) {23 if (i + 1 < argc) {24 model_path = argv[++i];25 } else {26 print_usage(argc, argv);27 return 1;28 }29 } else if (strcmp(argv[i], "-c") == 0) {30 if (i + 1 < argc) {31 n_ctx = std::stoi(argv[++i]);32 } else {33 print_usage(argc, argv);34 return 1;35 }36 } else if (strcmp(argv[i], "-ngl") == 0) {37 if (i + 1 < argc) {38 ngl = std::stoi(argv[++i]);39 } else {40 print_usage(argc, argv);41 return 1;42 }43 } else {44 print_usage(argc, argv);45 return 1;46 }47 } catch (std::exception & e) {48 fprintf(stderr, "error: %s\n", e.what());49 print_usage(argc, argv);50 return 1;51 }52 }53 if (model_path.empty()) {54 print_usage(argc, argv);55 return 1;56 }57 58 // only print errors59 llama_log_set([](enum ggml_log_level level, const char * text, void * /* user_data */) {60 if (level >= GGML_LOG_LEVEL_ERROR) {61 fprintf(stderr, "%s", text);62 }63 }, nullptr);64 65 // load dynamic backends66 ggml_backend_load_all();67 68 // initialize the model69 llama_model_params model_params = llama_model_default_params();70 model_params.n_gpu_layers = ngl;71 72 llama_model * model = llama_model_load_from_file(model_path.c_str(), model_params);73 if (!model) {74 fprintf(stderr , "%s: error: unable to load model\n" , __func__);75 return 1;76 }77 78 const llama_vocab * vocab = llama_model_get_vocab(model);79 80 // initialize the context81 llama_context_params ctx_params = llama_context_default_params();82 ctx_params.n_ctx = n_ctx;83 ctx_params.n_batch = n_ctx;84 85 llama_context * ctx = llama_init_from_model(model, ctx_params);86 if (!ctx) {87 fprintf(stderr , "%s: error: failed to create the llama_context\n" , __func__);88 return 1;89 }90 91 // initialize the sampler92 llama_sampler * smpl = llama_sampler_chain_init(llama_sampler_chain_default_params());93 llama_sampler_chain_add(smpl, llama_sampler_init_min_p(0.05f, 1));94 llama_sampler_chain_add(smpl, llama_sampler_init_temp(0.8f));95 llama_sampler_chain_add(smpl, llama_sampler_init_dist(LLAMA_DEFAULT_SEED));96 97 // helper function to evaluate a prompt and generate a response98 auto generate = [&](const std::string & prompt) {99 std::string response;100 101 const bool is_first = llama_get_kv_cache_used_cells(ctx) == 0;102 103 // tokenize the prompt104 const int n_prompt_tokens = -llama_tokenize(vocab, prompt.c_str(), prompt.size(), NULL, 0, is_first, true);105 std::vector<llama_token> prompt_tokens(n_prompt_tokens);106 if (llama_tokenize(vocab, prompt.c_str(), prompt.size(), prompt_tokens.data(), prompt_tokens.size(), is_first, true) < 0) {107 GGML_ABORT("failed to tokenize the prompt\n");108 }109 110 // prepare a batch for the prompt111 llama_batch batch = llama_batch_get_one(prompt_tokens.data(), prompt_tokens.size());112 llama_token new_token_id;113 while (true) {114 // check if we have enough space in the context to evaluate this batch115 int n_ctx = llama_n_ctx(ctx);116 int n_ctx_used = llama_get_kv_cache_used_cells(ctx);117 if (n_ctx_used + batch.n_tokens > n_ctx) {118 printf("\033[0m\n");119 fprintf(stderr, "context size exceeded\n");120 exit(0);121 }122 123 if (llama_decode(ctx, batch)) {124 GGML_ABORT("failed to decode\n");125 }126 127 // sample the next token128 new_token_id = llama_sampler_sample(smpl, ctx, -1);129 130 // is it an end of generation?131 if (llama_vocab_is_eog(vocab, new_token_id)) {132 break;133 }134 135 // convert the token to a string, print it and add it to the response136 char buf[256];137 int n = llama_token_to_piece(vocab, new_token_id, buf, sizeof(buf), 0, true);138 if (n < 0) {139 GGML_ABORT("failed to convert token to piece\n");140 }141 std::string piece(buf, n);142 printf("%s", piece.c_str());143 fflush(stdout);144 response += piece;145 146 // prepare the next batch with the sampled token147 batch = llama_batch_get_one(&new_token_id, 1);148 }149 150 return response;151 };152 153 std::vector<llama_chat_message> messages;154 std::vector<char> formatted(llama_n_ctx(ctx));155 int prev_len = 0;156 while (true) {157 // get user input158 printf("\033[32m> \033[0m");159 std::string user;160 std::getline(std::cin, user);161 162 if (user.empty()) {163 break;164 }165 166 const char * tmpl = llama_model_chat_template(model, /* name */ nullptr);167 168 // add the user input to the message list and format it169 messages.push_back({"user", strdup(user.c_str())});170 int new_len = llama_chat_apply_template(tmpl, messages.data(), messages.size(), true, formatted.data(), formatted.size());171 if (new_len > (int)formatted.size()) {172 formatted.resize(new_len);173 new_len = llama_chat_apply_template(tmpl, messages.data(), messages.size(), true, formatted.data(), formatted.size());174 }175 if (new_len < 0) {176 fprintf(stderr, "failed to apply the chat template\n");177 return 1;178 }179 180 // remove previous messages to obtain the prompt to generate the response181 std::string prompt(formatted.begin() + prev_len, formatted.begin() + new_len);182 183 // generate a response184 printf("\033[33m");185 std::string response = generate(prompt);186 printf("\n\033[0m");187 188 // add the response to the messages189 messages.push_back({"assistant", strdup(response.c_str())});190 prev_len = llama_chat_apply_template(tmpl, messages.data(), messages.size(), false, nullptr, 0);191 if (prev_len < 0) {192 fprintf(stderr, "failed to apply the chat template\n");193 return 1;194 }195 }196 197 // free resources198 for (auto & msg : messages) {199 free(const_cast<char *>(msg.content));200 }201 llama_sampler_free(smpl);202 llama_free(ctx);203 llama_model_free(model);204 205 return 0;206}207