KBaba7/llama.cpp

1#include "arg.h"2#include "common.h"3#include "log.h"4#include "llama.h"5 6#include <ctime>7 8#if defined(_MSC_VER)9#pragma warning(disable: 4244 4267) // possible loss of data10#endif11 12static std::vector<std::string> split_lines(const std::string & s, const std::string & separator = "\n") {13    std::vector<std::string> lines;14    size_t start = 0;15    size_t end = s.find(separator);16 17    while (end != std::string::npos) {18        lines.push_back(s.substr(start, end - start));19        start = end + separator.length();20        end = s.find(separator, start);21    }22 23    lines.push_back(s.substr(start)); // Add the last part24 25    return lines;26}27 28static void batch_add_seq(llama_batch & batch, const std::vector<int32_t> & tokens, llama_seq_id seq_id) {29    size_t n_tokens = tokens.size();30    for (size_t i = 0; i < n_tokens; i++) {31        common_batch_add(batch, tokens[i], i, { seq_id }, true);32    }33}34 35static void batch_decode(llama_context * ctx, llama_batch & batch, float * output, int n_seq, int n_embd, int embd_norm) {36    const enum llama_pooling_type pooling_type = llama_pooling_type(ctx);37    const struct llama_model * model = llama_get_model(ctx);38 39    // clear previous kv_cache values (irrelevant for embeddings)40    llama_kv_cache_clear(ctx);41 42    // run model43    LOG_INF("%s: n_tokens = %d, n_seq = %d\n", __func__, batch.n_tokens, n_seq);44    if (llama_model_has_encoder(model) && !llama_model_has_decoder(model)) {45        // encoder-only model46        if (llama_encode(ctx, batch) < 0) {47            LOG_ERR("%s : failed to encode\n", __func__);48        }49    } else if (!llama_model_has_encoder(model) && llama_model_has_decoder(model)) {50        // decoder-only model51        if (llama_decode(ctx, batch) < 0) {52            LOG_ERR("%s : failed to decode\n", __func__);53        }54    }55 56    for (int i = 0; i < batch.n_tokens; i++) {57        if (!batch.logits[i]) {58            continue;59        }60 61        const float * embd = nullptr;62        int embd_pos = 0;63 64        if (pooling_type == LLAMA_POOLING_TYPE_NONE) {65            // try to get token embeddings66            embd = llama_get_embeddings_ith(ctx, i);67            embd_pos = i;68            GGML_ASSERT(embd != NULL && "failed to get token embeddings");69        } else {70            // try to get sequence embeddings - supported only when pooling_type is not NONE71            embd = llama_get_embeddings_seq(ctx, batch.seq_id[i][0]);72            embd_pos = batch.seq_id[i][0];73            GGML_ASSERT(embd != NULL && "failed to get sequence embeddings");74        }75 76        float * out = output + embd_pos * n_embd;77        common_embd_normalize(embd, out, n_embd, embd_norm);78    }79}80 81int main(int argc, char ** argv) {82    common_params params;83 84    if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_EMBEDDING)) {85        return 1;86    }87 88    common_init();89 90    params.embedding = true;91    // For non-causal models, batch size must be equal to ubatch size92    params.n_ubatch = params.n_batch;93 94    llama_backend_init();95    llama_numa_init(params.numa);96 97    // load the model98    common_init_result llama_init = common_init_from_params(params);99 100    llama_model * model = llama_init.model.get();101    llama_context * ctx = llama_init.context.get();102 103    if (model == NULL) {104        LOG_ERR("%s: unable to load model\n", __func__);105        return 1;106    }107 108    const llama_vocab * vocab = llama_model_get_vocab(model);109 110    const int n_ctx_train = llama_model_n_ctx_train(model);111    const int n_ctx = llama_n_ctx(ctx);112 113    const enum llama_pooling_type pooling_type = llama_pooling_type(ctx);114 115    if (llama_model_has_encoder(model) && llama_model_has_decoder(model)) {116        LOG_ERR("%s: computing embeddings in encoder-decoder models is not supported\n", __func__);117        return 1;118    }119 120    if (n_ctx > n_ctx_train) {121        LOG_WRN("%s: warning: model was trained on only %d context tokens (%d specified)\n",122                __func__, n_ctx_train, n_ctx);123    }124 125    // print system information126    {127        LOG_INF("\n");128        LOG_INF("%s\n", common_params_get_system_info(params).c_str());129    }130 131    // split the prompt into lines132    std::vector<std::string> prompts = split_lines(params.prompt, params.embd_sep);133 134    // max batch size135    const uint64_t n_batch = params.n_batch;136    GGML_ASSERT(params.n_batch >= params.n_ctx);137 138    // tokenize the prompts and trim139    std::vector<std::vector<int32_t>> inputs;140    for (const auto & prompt : prompts) {141        auto inp = common_tokenize(ctx, prompt, true, true);142        if (inp.size() > n_batch) {143            LOG_ERR("%s: number of tokens in input line (%lld) exceeds batch size (%lld), increase batch size and re-run\n",144                    __func__, (long long int) inp.size(), (long long int) n_batch);145            return 1;146        }147        inputs.push_back(inp);148    }149 150    // check if the last token is SEP151    // it should be automatically added by the tokenizer when 'tokenizer.ggml.add_eos_token' is set to 'true'152    for (auto & inp : inputs) {153        if (inp.empty() || inp.back() != llama_vocab_sep(vocab)) {154            LOG_WRN("%s: last token in the prompt is not SEP\n", __func__);155            LOG_WRN("%s: 'tokenizer.ggml.add_eos_token' should be set to 'true' in the GGUF header\n", __func__);156        }157    }158 159    // tokenization stats160    if (params.verbose_prompt) {161        for (int i = 0; i < (int) inputs.size(); i++) {162            LOG_INF("%s: prompt %d: '%s'\n", __func__, i, prompts[i].c_str());163            LOG_INF("%s: number of tokens in prompt = %zu\n", __func__, inputs[i].size());164            for (int j = 0; j < (int) inputs[i].size(); j++) {165                LOG("%6d -> '%s'\n", inputs[i][j], common_token_to_piece(ctx, inputs[i][j]).c_str());166            }167            LOG("\n\n");168        }169    }170 171    // initialize batch172    const int n_prompts = prompts.size();173    struct llama_batch batch = llama_batch_init(n_batch, 0, 1);174 175    // count number of embeddings176    int n_embd_count = 0;177    if (pooling_type == LLAMA_POOLING_TYPE_NONE) {178        for (int k = 0; k < n_prompts; k++) {179            n_embd_count += inputs[k].size();180        }181    } else {182        n_embd_count = n_prompts;183    }184 185    // allocate output186    const int n_embd = llama_model_n_embd(model);187    std::vector<float> embeddings(n_embd_count * n_embd, 0);188    float * emb = embeddings.data();189 190    // break into batches191    int e = 0; // number of embeddings already stored192    int s = 0; // number of prompts in current batch193    for (int k = 0; k < n_prompts; k++) {194        // clamp to n_batch tokens195        auto & inp = inputs[k];196 197        const uint64_t n_toks = inp.size();198 199        // encode if at capacity200        if (batch.n_tokens + n_toks > n_batch) {201            float * out = emb + e * n_embd;202            batch_decode(ctx, batch, out, s, n_embd, params.embd_normalize);203            e += pooling_type == LLAMA_POOLING_TYPE_NONE ? batch.n_tokens : s;204            s = 0;205            common_batch_clear(batch);206        }207 208        // add to batch209        batch_add_seq(batch, inp, s);210        s += 1;211    }212 213    // final batch214    float * out = emb + e * n_embd;215    batch_decode(ctx, batch, out, s, n_embd, params.embd_normalize);216 217    if (params.embd_out.empty()) {218        LOG("\n");219 220        if (pooling_type == LLAMA_POOLING_TYPE_NONE) {221            for (int j = 0; j < n_embd_count; j++) {222                LOG("embedding %d: ", j);223                for (int i = 0; i < std::min(3, n_embd); i++) {224                    if (params.embd_normalize == 0) {225                        LOG("%6.0f ", emb[j * n_embd + i]);226                    } else {227                        LOG("%9.6f ", emb[j * n_embd + i]);228                    }229                }230                LOG(" ... ");231                for (int i = n_embd - 3; i < n_embd; i++) {232                    if (params.embd_normalize == 0) {233                        LOG("%6.0f ", emb[j * n_embd + i]);234                    } else {235                        LOG("%9.6f ", emb[j * n_embd + i]);236                    }237                }238                LOG("\n");239            }240        } else if (pooling_type == LLAMA_POOLING_TYPE_RANK) {241            for (int j = 0; j < n_embd_count; j++) {242                // NOTE: if you change this log - update the tests in ci/run.sh243                LOG("rerank score %d: %8.3f\n", j, emb[j * n_embd]);244            }245        } else {246            // print the first part of the embeddings or for a single prompt, the full embedding247            for (int j = 0; j < n_prompts; j++) {248                LOG("embedding %d: ", j);249                for (int i = 0; i < (n_prompts > 1 ? std::min(16, n_embd) : n_embd); i++) {250                    if (params.embd_normalize == 0) {251                        LOG("%6.0f ", emb[j * n_embd + i]);252                    } else {253                        LOG("%9.6f ", emb[j * n_embd + i]);254                    }255                }256                LOG("\n");257            }258 259            // print cosine similarity matrix260            if (n_prompts > 1) {261                LOG("\n");262                LOG("cosine similarity matrix:\n\n");263                for (int i = 0; i < n_prompts; i++) {264                    LOG("%6.6s ", prompts[i].c_str());265                }266                LOG("\n");267                for (int i = 0; i < n_prompts; i++) {268                    for (int j = 0; j < n_prompts; j++) {269                        float sim = common_embd_similarity_cos(emb + i * n_embd, emb + j * n_embd, n_embd);270                        LOG("%6.2f ", sim);271                    }272                    LOG("%1.10s", prompts[i].c_str());273                    LOG("\n");274                }275            }276        }277    }278 279    if (params.embd_out == "json" || params.embd_out == "json+" || params.embd_out == "array") {280        const bool notArray = params.embd_out != "array";281 282        LOG(notArray ? "{\n  \"object\": \"list\",\n  \"data\": [\n" : "[");283        for (int j = 0;;) { // at least one iteration (one prompt)284            if (notArray) LOG("    {\n      \"object\": \"embedding\",\n      \"index\": %d,\n      \"embedding\": ",j);285            LOG("[");286            for (int i = 0;;) { // at least one iteration (n_embd > 0)287                LOG(params.embd_normalize == 0 ? "%1.0f" : "%1.7f", emb[j * n_embd + i]);288                i++;289                if (i < n_embd) LOG(","); else break;290            }291            LOG(notArray ? "]\n    }" : "]");292            j++;293            if (j < n_embd_count) LOG(notArray ? ",\n" : ","); else break;294        }295        LOG(notArray ? "\n  ]" : "]\n");296 297        if (params.embd_out == "json+" && n_prompts > 1) {298            LOG(",\n  \"cosineSimilarity\": [\n");299            for (int i = 0;;) { // at least two iteration (n_embd_count > 1)300                LOG("    [");301                for (int j = 0;;) { // at least two iteration (n_embd_count > 1)302                    float sim = common_embd_similarity_cos(emb + i * n_embd, emb + j * n_embd, n_embd);303                    LOG("%6.2f", sim);304                    j++;305                    if (j < n_embd_count) LOG(", "); else break;306                }307                LOG(" ]");308                i++;309                if (i < n_embd_count) LOG(",\n"); else break;310            }311            LOG("\n  ]");312        }313 314        if (notArray) LOG("\n}\n");315    }316 317    LOG("\n");318    llama_perf_context_print(ctx);319 320    // clean up321    llama_batch_free(batch);322    llama_backend_free();323 324    return 0;325}326