echodict/llama.cpp

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