KBaba7/llama.cpp
0
1#include "arg.h"2#include "base64.hpp"3#include "log.h"4#include "common.h"5#include "sampling.h"6#include "clip.h"7#include "llava.h"8#include "llama.h"9#include "ggml.h"10 11#include <cstdio>12#include <cstdlib>13#include <cstring>14#include <vector>15 16static bool eval_tokens(struct llama_context * ctx_llama, std::vector<llama_token> tokens, int n_batch, int * n_past) {17 int N = (int) tokens.size();18 for (int i = 0; i < N; i += n_batch) {19 int n_eval = (int) tokens.size() - i;20 if (n_eval > n_batch) {21 n_eval = n_batch;22 }23 if (llama_decode(ctx_llama, llama_batch_get_one(&tokens[i], n_eval))) {24 LOG_ERR("%s : failed to eval. token %d/%d (batch size %d, n_past %d)\n", __func__, i, N, n_batch, *n_past);25 return false;26 }27 *n_past += n_eval;28 }29 return true;30}31 32static bool eval_id(struct llama_context * ctx_llama, int id, int * n_past) {33 std::vector<llama_token> tokens;34 tokens.push_back(id);35 return eval_tokens(ctx_llama, tokens, 1, n_past);36}37 38static bool eval_string(struct llama_context * ctx_llama, const char* str, int n_batch, int * n_past, bool add_bos){39 std::string str2 = str;40 std::vector<llama_token> embd_inp = common_tokenize(ctx_llama, str2, add_bos, true);41 eval_tokens(ctx_llama, embd_inp, n_batch, n_past);42 return true;43}44 45static const char * sample(struct common_sampler * smpl,46 struct llama_context * ctx_llama,47 int * n_past) {48 const llama_token id = common_sampler_sample(smpl, ctx_llama, -1);49 common_sampler_accept(smpl, id, true);50 51 const llama_model * model = llama_get_model(ctx_llama);52 const llama_vocab * vocab = llama_model_get_vocab(model);53 54 static std::string ret;55 if (llama_vocab_is_eog(vocab, id)) {56 ret = "</s>";57 } else {58 ret = common_token_to_piece(ctx_llama, id);59 }60 eval_id(ctx_llama, id, n_past);61 return ret.c_str();62}63 64static const char* IMG_BASE64_TAG_BEGIN = "<img src=\"data:image/jpeg;base64,";65static const char* IMG_BASE64_TAG_END = "\">";66 67static void find_image_tag_in_prompt(const std::string& prompt, size_t& begin_out, size_t& end_out) {68 begin_out = prompt.find(IMG_BASE64_TAG_BEGIN);69 end_out = prompt.find(IMG_BASE64_TAG_END, (begin_out == std::string::npos) ? 0UL : begin_out);70}71 72static bool prompt_contains_image(const std::string& prompt) {73 size_t begin, end;74 find_image_tag_in_prompt(prompt, begin, end);75 return (begin != std::string::npos);76}77 78// replaces the base64 image tag in the prompt with `replacement`79static llava_image_embed * llava_image_embed_make_with_prompt_base64(struct clip_ctx * ctx_clip, int n_threads, const std::string& prompt) {80 size_t img_base64_str_start, img_base64_str_end;81 find_image_tag_in_prompt(prompt, img_base64_str_start, img_base64_str_end);82 if (img_base64_str_start == std::string::npos || img_base64_str_end == std::string::npos) {83 LOG_ERR("%s: invalid base64 image tag. must be %s<base64 byte string>%s\n", __func__, IMG_BASE64_TAG_BEGIN, IMG_BASE64_TAG_END);84 return NULL;85 }86 87 auto base64_bytes_start = img_base64_str_start + strlen(IMG_BASE64_TAG_BEGIN);88 auto base64_bytes_count = img_base64_str_end - base64_bytes_start;89 auto base64_str = prompt.substr(base64_bytes_start, base64_bytes_count );90 91 auto required_bytes = base64::required_encode_size(base64_str.size());92 auto img_bytes = std::vector<unsigned char>(required_bytes);93 base64::decode(base64_str.begin(), base64_str.end(), img_bytes.begin());94 95 auto embed = llava_image_embed_make_with_bytes(ctx_clip, n_threads, img_bytes.data(), img_bytes.size());96 if (!embed) {97 LOG_ERR("%s: could not load image from base64 string.\n", __func__);98 return NULL;99 }100 101 return embed;102}103 104static std::string remove_image_from_prompt(const std::string& prompt, const char * replacement = "") {105 size_t begin, end;106 find_image_tag_in_prompt(prompt, begin, end);107 if (begin == std::string::npos || end == std::string::npos) {108 return prompt;109 }110 auto pre = prompt.substr(0, begin);111 auto post = prompt.substr(end + strlen(IMG_BASE64_TAG_END));112 return pre + replacement + post;113}114 115struct llava_context {116 struct clip_ctx * ctx_clip = NULL;117 struct llama_context * ctx_llama = NULL;118 struct llama_model * model = NULL;119};120 121static void print_usage(int, char ** argv) {122 LOG("\n example usage:\n");123 LOG("\n %s -m <llava-v1.5-7b/ggml-model-q5_k.gguf> --mmproj <llava-v1.5-7b/mmproj-model-f16.gguf> --image <path/to/an/image.jpg> --image <path/to/another/image.jpg> [--temp 0.1] [-p \"describe the image in detail.\"]\n", argv[0]);124 LOG("\n note: a lower temperature value like 0.1 is recommended for better quality.\n");125}126 127static struct llava_image_embed * load_image(llava_context * ctx_llava, common_params * params, const std::string & fname) {128 129 // load and preprocess the image130 llava_image_embed * embed = NULL;131 auto prompt = params->prompt;132 if (prompt_contains_image(prompt)) {133 if (!params->image.empty()) {134 LOG_INF("using base64 encoded image instead of command line image path\n");135 }136 embed = llava_image_embed_make_with_prompt_base64(ctx_llava->ctx_clip, params->cpuparams.n_threads, prompt);137 if (!embed) {138 LOG_ERR("%s: can't load image from prompt\n", __func__);139 return NULL;140 }141 params->prompt = remove_image_from_prompt(prompt);142 } else {143 embed = llava_image_embed_make_with_filename(ctx_llava->ctx_clip, params->cpuparams.n_threads, fname.c_str());144 if (!embed) {145 fprintf(stderr, "%s: is %s really an image file?\n", __func__, fname.c_str());146 return NULL;147 }148 }149 150 return embed;151}152 153static void process_prompt(struct llava_context * ctx_llava, struct llava_image_embed * image_embed, common_params * params, const std::string & prompt) {154 int n_past = 0;155 156 const int max_tgt_len = params->n_predict < 0 ? 256 : params->n_predict;157 158 std::string system_prompt, user_prompt;159 size_t image_pos = prompt.find("<image>");160 if (image_pos != std::string::npos) {161 // new templating mode: Provide the full prompt including system message and use <image> as a placeholder for the image162 system_prompt = prompt.substr(0, image_pos);163 user_prompt = prompt.substr(image_pos + std::string("<image>").length());164 LOG_INF("system_prompt: %s\n", system_prompt.c_str());165 if (params->verbose_prompt) {166 auto tmp = common_tokenize(ctx_llava->ctx_llama, system_prompt, true, true);167 for (int i = 0; i < (int) tmp.size(); i++) {168 LOG_INF("%6d -> '%s'\n", tmp[i], common_token_to_piece(ctx_llava->ctx_llama, tmp[i]).c_str());169 }170 }171 LOG_INF("user_prompt: %s\n", user_prompt.c_str());172 if (params->verbose_prompt) {173 auto tmp = common_tokenize(ctx_llava->ctx_llama, user_prompt, true, true);174 for (int i = 0; i < (int) tmp.size(); i++) {175 LOG_INF("%6d -> '%s'\n", tmp[i], common_token_to_piece(ctx_llava->ctx_llama, tmp[i]).c_str());176 }177 }178 } else {179 // llava-1.5 native mode180 system_prompt = "A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions.\nUSER:";181 user_prompt = prompt + "\nASSISTANT:";182 if (params->verbose_prompt) {183 auto tmp = common_tokenize(ctx_llava->ctx_llama, user_prompt, true, true);184 for (int i = 0; i < (int) tmp.size(); i++) {185 LOG_INF("%6d -> '%s'\n", tmp[i], common_token_to_piece(ctx_llava->ctx_llama, tmp[i]).c_str());186 }187 }188 }189 190 eval_string(ctx_llava->ctx_llama, system_prompt.c_str(), params->n_batch, &n_past, true);191 llava_eval_image_embed(ctx_llava->ctx_llama, image_embed, params->n_batch, &n_past);192 eval_string(ctx_llava->ctx_llama, user_prompt.c_str(), params->n_batch, &n_past, false);193 194 // generate the response195 196 LOG("\n");197 198 struct common_sampler * smpl = common_sampler_init(ctx_llava->model, params->sampling);199 if (!smpl) {200 LOG_ERR("%s: failed to initialize sampling subsystem\n", __func__);201 exit(1);202 }203 204 std::string response = "";205 for (int i = 0; i < max_tgt_len; i++) {206 const char * tmp = sample(smpl, ctx_llava->ctx_llama, &n_past);207 response += tmp;208 if (strcmp(tmp, "</s>") == 0) break;209 if (strstr(tmp, "###")) break; // Yi-VL behavior210 LOG("%s", tmp);211 if (strstr(response.c_str(), "<|im_end|>")) break; // Yi-34B llava-1.6 - for some reason those decode not as the correct token (tokenizer works)212 if (strstr(response.c_str(), "<|im_start|>")) break; // Yi-34B llava-1.6213 if (strstr(response.c_str(), "USER:")) break; // mistral llava-1.6214 215 fflush(stdout);216 }217 218 common_sampler_free(smpl);219 LOG("\n");220}221 222static struct llama_model * llava_init(common_params * params) {223 llama_backend_init();224 llama_numa_init(params->numa);225 226 llama_model_params model_params = common_model_params_to_llama(*params);227 228 llama_model * model = llama_model_load_from_file(params->model.c_str(), model_params);229 if (model == NULL) {230 LOG_ERR("%s: unable to load model\n" , __func__);231 return NULL;232 }233 return model;234}235 236static struct llava_context * llava_init_context(common_params * params, llama_model * model) {237 const char * clip_path = params->mmproj.c_str();238 239 auto prompt = params->prompt;240 if (prompt.empty()) {241 prompt = "describe the image in detail.";242 }243 244 auto ctx_clip = clip_model_load(clip_path, /*verbosity=*/ 1);245 246 llama_context_params ctx_params = common_context_params_to_llama(*params);247 ctx_params.n_ctx = params->n_ctx < 2048 ? 2048 : params->n_ctx; // we need a longer context size to process image embeddings248 249 llama_context * ctx_llama = llama_init_from_model(model, ctx_params);250 251 if (ctx_llama == NULL) {252 LOG_ERR("%s: failed to create the llama_context\n" , __func__);253 return NULL;254 }255 256 auto * ctx_llava = (struct llava_context *)malloc(sizeof(llava_context));257 258 ctx_llava->ctx_llama = ctx_llama;259 ctx_llava->ctx_clip = ctx_clip;260 ctx_llava->model = model;261 return ctx_llava;262}263 264static void llava_free(struct llava_context * ctx_llava) {265 if (ctx_llava->ctx_clip) {266 clip_free(ctx_llava->ctx_clip);267 ctx_llava->ctx_clip = NULL;268 }269 270 llama_free(ctx_llava->ctx_llama);271 llama_model_free(ctx_llava->model);272 llama_backend_free();273}274 275int main(int argc, char ** argv) {276 ggml_time_init();277 278 common_params params;279 280 if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_LLAVA, print_usage)) {281 return 1;282 }283 284 common_init();285 286 if (params.mmproj.empty() || (params.image.empty() && !prompt_contains_image(params.prompt))) {287 print_usage(argc, argv);288 return 1;289 }290 291 auto * model = llava_init(¶ms);292 if (model == NULL) {293 fprintf(stderr, "%s: error: failed to init llava model\n", __func__);294 return 1;295 }296 297 if (prompt_contains_image(params.prompt)) {298 auto * ctx_llava = llava_init_context(¶ms, model);299 300 auto * image_embed = load_image(ctx_llava, ¶ms, "");301 302 // process the prompt303 process_prompt(ctx_llava, image_embed, ¶ms, params.prompt);304 305 llama_perf_context_print(ctx_llava->ctx_llama);306 llava_image_embed_free(image_embed);307 ctx_llava->model = NULL;308 llava_free(ctx_llava);309 } else {310 for (auto & image : params.image) {311 auto * ctx_llava = llava_init_context(¶ms, model);312 313 auto * image_embed = load_image(ctx_llava, ¶ms, image);314 if (!image_embed) {315 LOG_ERR("%s: failed to load image %s. Terminating\n\n", __func__, image.c_str());316 return 1;317 }318 319 // process the prompt320 process_prompt(ctx_llava, image_embed, ¶ms, params.prompt);321 322 llama_perf_context_print(ctx_llava->ctx_llama);323 llava_image_embed_free(image_embed);324 ctx_llava->model = NULL;325 llava_free(ctx_llava);326 }327 }328 329 llama_model_free(model);330 331 return 0;332}333