Felipe97/llama-cpp-compiled
01.1k
1#include "arg.h"2#include "common.h"3#include "log.h"4#include "llama.h"5 6#include <clocale>7#include <ctime>8#include <algorithm>9 10#if defined(_MSC_VER)11#pragma warning(disable: 4244 4267) // possible loss of data12#endif13 14static std::vector<std::string> split_lines(const std::string & s, const std::string & separator = "\n") {15 std::vector<std::string> lines;16 size_t start = 0;17 size_t end = s.find(separator);18 19 while (end != std::string::npos) {20 lines.push_back(s.substr(start, end - start));21 start = end + separator.length();22 end = s.find(separator, start);23 }24 25 lines.push_back(s.substr(start)); // Add the last part26 27 return lines;28}29 30static void batch_add_seq(llama_batch & batch, const std::vector<int32_t> & tokens, llama_seq_id seq_id) {31 size_t n_tokens = tokens.size();32 for (size_t i = 0; i < n_tokens; i++) {33 common_batch_add(batch, tokens[i], i, { seq_id }, true);34 }35}36 37static void batch_decode(llama_context * ctx, llama_batch & batch, float * output, int n_seq, int n_embd_out, int embd_norm) {38 const enum llama_pooling_type pooling_type = llama_pooling_type(ctx);39 40 // clear previous kv_cache values (irrelevant for embeddings)41 llama_memory_clear(llama_get_memory(ctx), true);42 43 // run model44 LOG_INF("%s: n_tokens = %d, n_seq = %d\n", __func__, batch.n_tokens, n_seq);45 if (llama_decode(ctx, batch) < 0) {46 LOG_ERR("%s : failed to process\n", __func__);47 }48 49 for (int i = 0; i < batch.n_tokens; i++) {50 if (!batch.logits[i]) {51 continue;52 }53 54 const float * embd = nullptr;55 int embd_pos = 0;56 57 if (pooling_type == LLAMA_POOLING_TYPE_NONE) {58 // try to get token embeddings59 embd = llama_get_embeddings_ith(ctx, i);60 embd_pos = i;61 GGML_ASSERT(embd != NULL && "failed to get token embeddings");62 } else {63 // try to get sequence embeddings - supported only when pooling_type is not NONE64 embd = llama_get_embeddings_seq(ctx, batch.seq_id[i][0]);65 embd_pos = batch.seq_id[i][0];66 GGML_ASSERT(embd != NULL && "failed to get sequence embeddings");67 }68 69 float * out = output + embd_pos * n_embd_out;70 common_embd_normalize(embd, out, n_embd_out, embd_norm);71 }72}73 74// plain, pipe-friendly output: one embedding per line75static void print_raw_embeddings(const float * emb,76 int n_embd_count,77 int n_embd,78 const llama_model * model,79 enum llama_pooling_type pooling_type,80 int embd_normalize) {81 const uint32_t n_cls_out = llama_model_n_cls_out(model);82 const bool is_rank = (pooling_type == LLAMA_POOLING_TYPE_RANK);83 const int cols = is_rank ? std::min<int>(n_embd, (int) n_cls_out) : n_embd;84 85 for (int j = 0; j < n_embd_count; ++j) {86 for (int i = 0; i < cols; ++i) {87 if (embd_normalize == 0) {88 LOG("%1.0f%s", emb[j * n_embd + i], (i + 1 < cols ? " " : ""));89 } else {90 LOG("%1.7f%s", emb[j * n_embd + i], (i + 1 < cols ? " " : ""));91 }92 }93 LOG("\n");94 }95}96 97int main(int argc, char ** argv) {98 std::setlocale(LC_NUMERIC, "C");99 100 common_params params;101 102 common_init();103 104 if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_EMBEDDING)) {105 return 1;106 }107 108 params.embedding = true;109 110 // get max number of sequences per batch111 const int n_seq_max = llama_max_parallel_sequences();112 113 // if the number of prompts that would be encoded is known in advance, it's more efficient to specify the114 // --parallel argument accordingly. for convenience, if not specified, we fallback to unified KV cache115 // in order to support any number of prompts116 if (params.n_parallel == 1) {117 LOG_INF("%s: n_parallel == 1 -> unified KV cache is enabled\n", __func__);118 params.kv_unified = true;119 params.n_parallel = n_seq_max;120 }121 122 // utilize the full context123 if (params.n_batch < params.n_ctx) {124 LOG_WRN("%s: setting batch size to %d\n", __func__, params.n_ctx);125 params.n_batch = params.n_ctx;126 }127 128 // for non-causal models, batch size must be equal to ubatch size129 if (params.attention_type != LLAMA_ATTENTION_TYPE_CAUSAL) {130 params.n_ubatch = params.n_batch;131 }132 133 llama_backend_init();134 llama_numa_init(params.numa);135 136 // load the model137 auto llama_init = common_init_from_params(params);138 139 auto * model = llama_init->model();140 auto * ctx = llama_init->context();141 142 if (model == NULL) {143 LOG_ERR("%s: unable to load model\n", __func__);144 return 1;145 }146 147 const llama_vocab * vocab = llama_model_get_vocab(model);148 149 const int n_ctx_train = llama_model_n_ctx_train(model);150 const int n_ctx = llama_n_ctx(ctx);151 152 const enum llama_pooling_type pooling_type = llama_pooling_type(ctx);153 154 if (llama_model_has_encoder(model) && llama_model_has_decoder(model)) {155 LOG_ERR("%s: computing embeddings in encoder-decoder models is not supported\n", __func__);156 return 1;157 }158 159 if (n_ctx > n_ctx_train) {160 LOG_WRN("%s: warning: model was trained on only %d context tokens (%d specified)\n",161 __func__, n_ctx_train, n_ctx);162 }163 164 // print system information165 {166 LOG_INF("\n");167 LOG_INF("%s\n", common_params_get_system_info(params).c_str());168 }169 170 // split the prompt into lines171 std::vector<std::string> prompts = split_lines(params.prompt, params.embd_sep);172 173 // max batch size174 const uint64_t n_batch = params.n_batch;175 176 // get added sep and eos token, if any177 const std::string added_sep_token = llama_vocab_get_add_sep(vocab) ? llama_vocab_get_text(vocab, llama_vocab_sep(vocab)) : "";178 const std::string added_eos_token = llama_vocab_get_add_eos(vocab) ? llama_vocab_get_text(vocab, llama_vocab_eos(vocab)) : "";179 const char * rerank_prompt = llama_model_chat_template(model, "rerank");180 181 // tokenize the prompts and trim182 std::vector<std::vector<int32_t>> inputs;183 for (const auto & prompt : prompts) {184 std::vector<llama_token> inp;185 186 // split classification pairs and insert expected separator tokens187 if (pooling_type == LLAMA_POOLING_TYPE_RANK && prompt.find(params.cls_sep) != std::string::npos) {188 std::vector<std::string> pairs = split_lines(prompt, params.cls_sep);189 if (rerank_prompt != nullptr) {190 const std::string query = pairs[0];191 const std::string doc = pairs[1];192 std::string final_prompt = rerank_prompt;193 string_replace_all(final_prompt, "{query}" , query);194 string_replace_all(final_prompt, "{document}", doc );195 inp = common_tokenize(vocab, final_prompt, true, true);196 } else {197 std::string final_prompt;198 for (size_t i = 0; i < pairs.size(); i++) {199 final_prompt += pairs[i];200 if (i != pairs.size() - 1) {201 if (!added_eos_token.empty()) {202 final_prompt += added_eos_token;203 }204 if (!added_sep_token.empty()) {205 final_prompt += added_sep_token;206 }207 }208 }209 inp = common_tokenize(ctx, final_prompt, true, true);210 }211 } else {212 inp = common_tokenize(ctx, prompt, true, true);213 }214 if (inp.size() > n_batch) {215 LOG_ERR("%s: number of tokens in input line (%lld) exceeds batch size (%lld), increase batch size and re-run\n",216 __func__, (long long int) inp.size(), (long long int) n_batch);217 return 1;218 }219 inputs.push_back(inp);220 }221 222 // check if the last token is SEP/EOS223 // it should be automatically added by the tokenizer when 'tokenizer.ggml.add_eos_token' is set to 'true'224 for (auto & inp : inputs) {225 if (inp.empty() || (inp.back() != llama_vocab_sep(vocab) && inp.back() != llama_vocab_eos(vocab))) {226 LOG_WRN("%s: last token in the prompt is not SEP or EOS\n", __func__);227 LOG_WRN("%s: 'tokenizer.ggml.add_eos_token' should be set to 'true' in the GGUF header\n", __func__);228 }229 }230 231 // tokenization stats232 if (params.verbose_prompt) {233 for (int i = 0; i < (int) inputs.size(); i++) {234 LOG_INF("%s: prompt %d: '%s'\n", __func__, i, prompts[i].c_str());235 LOG_INF("%s: number of tokens in prompt = %zu\n", __func__, inputs[i].size());236 for (int j = 0; j < (int) inputs[i].size(); j++) {237 LOG("%6d -> '%s'\n", inputs[i][j], common_token_to_piece(ctx, inputs[i][j]).c_str());238 }239 LOG("\n\n");240 }241 }242 243 // initialize batch244 const int n_prompts = prompts.size();245 struct llama_batch batch = llama_batch_init(n_batch, 0, 1);246 247 // count number of embeddings248 int n_embd_count = 0;249 if (pooling_type == LLAMA_POOLING_TYPE_NONE) {250 for (int k = 0; k < n_prompts; k++) {251 n_embd_count += inputs[k].size();252 }253 } else {254 n_embd_count = n_prompts;255 }256 257 // allocate output258 const int n_embd_out = llama_model_n_embd_out(model);259 std::vector<float> embeddings(n_embd_count * n_embd_out, 0);260 float * emb = embeddings.data();261 262 // break into batches263 int e = 0; // number of embeddings already stored264 int s = 0; // number of prompts in current batch265 for (int k = 0; k < n_prompts; k++) {266 // clamp to n_batch tokens267 auto & inp = inputs[k];268 269 const uint64_t n_toks = inp.size();270 271 // encode if at capacity272 if (batch.n_tokens + n_toks > n_batch || s >= n_seq_max) {273 float * out = emb + e * n_embd_out;274 batch_decode(ctx, batch, out, s, n_embd_out, params.embd_normalize);275 e += pooling_type == LLAMA_POOLING_TYPE_NONE ? batch.n_tokens : s;276 s = 0;277 common_batch_clear(batch);278 }279 280 // add to batch281 batch_add_seq(batch, inp, s);282 s += 1;283 }284 285 // final batch286 float * out = emb + e * n_embd_out;287 batch_decode(ctx, batch, out, s, n_embd_out, params.embd_normalize);288 289 if (params.embd_out.empty()) {290 LOG("\n");291 292 if (pooling_type == LLAMA_POOLING_TYPE_NONE) {293 for (int j = 0; j < n_embd_count; j++) {294 LOG("embedding %d: ", j);295 for (int i = 0; i < std::min(3, n_embd_out); i++) {296 if (params.embd_normalize == 0) {297 LOG("%6.0f ", emb[j * n_embd_out + i]);298 } else {299 LOG("%9.6f ", emb[j * n_embd_out + i]);300 }301 }302 LOG(" ... ");303 for (int i = n_embd_out - 3; i < n_embd_out; i++) {304 if (params.embd_normalize == 0) {305 LOG("%6.0f ", emb[j * n_embd_out + i]);306 } else {307 LOG("%9.6f ", emb[j * n_embd_out + i]);308 }309 }310 LOG("\n");311 }312 } else if (pooling_type == LLAMA_POOLING_TYPE_RANK) {313 const uint32_t n_cls_out = llama_model_n_cls_out(model);314 std::vector<std::string> cls_out_labels;315 316 for (uint32_t i = 0; i < n_cls_out; i++) {317 const char * label = llama_model_cls_label(model, i);318 const std::string label_i(label == nullptr ? "" : label);319 cls_out_labels.emplace_back(label_i.empty() ? std::to_string(i) : label_i);320 }321 322 for (int j = 0; j < n_embd_count; j++) {323 for (uint32_t i = 0; i < n_cls_out; i++) {324 // NOTE: if you change this log - update the tests in ci/run.sh325 if (n_cls_out == 1) {326 LOG("rerank score %d: %8.3f\n", j, emb[j * n_embd_out]);327 } else {328 LOG("rerank score %d: %8.3f [%s]\n", j, emb[j * n_embd_out + i], cls_out_labels[i].c_str());329 }330 }331 }332 } else {333 // print the first part of the embeddings or for a single prompt, the full embedding334 for (int j = 0; j < n_prompts; j++) {335 LOG("embedding %d: ", j);336 for (int i = 0; i < (n_prompts > 1 ? std::min(16, n_embd_out) : n_embd_out); i++) {337 if (params.embd_normalize == 0) {338 LOG("%6.0f ", emb[j * n_embd_out + i]);339 } else {340 LOG("%9.6f ", emb[j * n_embd_out + i]);341 }342 }343 LOG("\n");344 }345 346 // print cosine similarity matrix347 if (n_prompts > 1) {348 LOG("\n");349 LOG("cosine similarity matrix:\n\n");350 for (int i = 0; i < n_prompts; i++) {351 LOG("%6.6s ", prompts[i].c_str());352 }353 LOG("\n");354 for (int i = 0; i < n_prompts; i++) {355 for (int j = 0; j < n_prompts; j++) {356 float sim = common_embd_similarity_cos(emb + i * n_embd_out, emb + j * n_embd_out, n_embd_out);357 LOG("%6.2f ", sim);358 }359 LOG("%1.10s", prompts[i].c_str());360 LOG("\n");361 }362 }363 }364 }365 366 if (params.embd_out == "json" || params.embd_out == "json+" || params.embd_out == "array") {367 const bool notArray = params.embd_out != "array";368 369 LOG(notArray ? "{\n \"object\": \"list\",\n \"data\": [\n" : "[");370 for (int j = 0;;) { // at least one iteration (one prompt)371 if (notArray) LOG(" {\n \"object\": \"embedding\",\n \"index\": %d,\n \"embedding\": ",j);372 LOG("[");373 for (int i = 0;;) { // at least one iteration (n_embd > 0)374 LOG(params.embd_normalize == 0 ? "%1.0f" : "%1.7f", emb[j * n_embd_out + i]);375 i++;376 if (i < n_embd_out) LOG(","); else break;377 }378 LOG(notArray ? "]\n }" : "]");379 j++;380 if (j < n_embd_count) LOG(notArray ? ",\n" : ","); else break;381 }382 LOG(notArray ? "\n ]" : "]\n");383 384 if (params.embd_out == "json+" && n_prompts > 1) {385 LOG(",\n \"cosineSimilarity\": [\n");386 for (int i = 0;;) { // at least two iteration (n_embd_count > 1)387 LOG(" [");388 for (int j = 0;;) { // at least two iteration (n_embd_count > 1)389 float sim = common_embd_similarity_cos(emb + i * n_embd_out, emb + j * n_embd_out, n_embd_out);390 LOG("%6.2f", sim);391 j++;392 if (j < n_embd_count) LOG(", "); else break;393 }394 LOG(" ]");395 i++;396 if (i < n_embd_count) LOG(",\n"); else break;397 }398 LOG("\n ]");399 }400 401 if (notArray) LOG("\n}\n");402 } else if (params.embd_out == "raw") {403 print_raw_embeddings(emb, n_embd_count, n_embd_out, model, pooling_type, params.embd_normalize);404 }405 406 LOG("\n");407 llama_perf_context_print(ctx);408 409 // clean up410 llama_batch_free(batch);411 llama_backend_free();412 413 return 0;414}415