Felipe97/llama-cpp-compiled

1#include "arg.h"2#include "common.h"3#include "log.h"4#include "llama.h"5 6#include <algorithm>7#include <clocale>8#include <fstream>9#include <iostream> // TODO: remove me10 11static void print_usage(int, char ** argv) {12    LOG("\nexample usage:\n");13    LOG("\n    %s --model ./models/bge-base-en-v1.5-f16.gguf --top-k 3 --context-file README.md --context-file License --chunk-size 100 --chunk-separator .\n", argv[0]);14    LOG("\n");15}16 17struct chunk {18    // filename19    std::string filename;20    // original file position21    size_t filepos;22    // original text data23    std::string textdata;24    // tokenized text data25    std::vector<llama_token> tokens;26    // embedding27    std::vector<float> embedding;28};29 30// chunk file data to chunks of size >= chunk_size31// chunk_separator is the separator between chunks32static std::vector<chunk> chunk_file(const std::string & filename, int chunk_size, const std::string & chunk_separator) {33    std::vector<chunk> chunks;34    std::ifstream f(filename.c_str());35 36    if (!f.is_open()) {37        LOG_ERR("could not open file %s\n", filename.c_str());38        return chunks;39    }40 41    chunk current_chunk;42    char buffer[1024];43    int64_t filepos = 0;44    std::string current;45    while (f.read(buffer, 1024)) {46        current += std::string(buffer, f.gcount());47        size_t pos;48        while ((pos = current.find(chunk_separator)) != std::string::npos) {49            current_chunk.textdata += current.substr(0, pos + chunk_separator.size());50            if ((int) current_chunk.textdata.size() > chunk_size) {51                // save chunk52                current_chunk.filepos = filepos;53                current_chunk.filename = filename;54                chunks.push_back(current_chunk);55                // update filepos56                filepos += (int) current_chunk.textdata.size();57                // reset current_chunk58                current_chunk = chunk();59            }60            current = current.substr(pos + chunk_separator.size());61        }62 63    }64    // add leftover data to last chunk65    if (current_chunk.textdata.size() > 0) {66        if (chunks.empty()) {67            current_chunk.filepos = filepos;68            current_chunk.filename = filename;69            chunks.push_back(current_chunk);70        } else {71            chunks.back().textdata += current_chunk.textdata;72        }73    }74    f.close();75    return chunks;76}77 78static void batch_add_seq(llama_batch & batch, const std::vector<int32_t> & tokens, llama_seq_id seq_id) {79    size_t n_tokens = tokens.size();80    for (size_t i = 0; i < n_tokens; i++) {81        common_batch_add(batch, tokens[i], i, { seq_id }, true);82    }83}84 85static void batch_process(llama_context * ctx, llama_batch & batch, float * output, int n_seq, int n_embd) {86    // clear previous kv_cache values (irrelevant for embeddings)87    llama_memory_clear(llama_get_memory(ctx), false);88 89    // run model90    LOG_INF("%s: n_tokens = %d, n_seq = %d\n", __func__, batch.n_tokens, n_seq);91    if (llama_decode(ctx, batch) < 0) {92        LOG_ERR("%s : failed to process\n", __func__);93    }94 95    for (int i = 0; i < batch.n_tokens; i++) {96        if (!batch.logits[i]) {97            continue;98        }99 100        // try to get sequence embeddings - supported only when pooling_type is not NONE101        const float * embd = llama_get_embeddings_seq(ctx, batch.seq_id[i][0]);102        if (embd == NULL) {103            embd = llama_get_embeddings_ith(ctx, i);104            if (embd == NULL) {105                LOG_ERR("%s: failed to get embeddings for token %d\n", __func__, i);106                continue;107            }108        }109 110        float * out = output + batch.seq_id[i][0] * n_embd;111        common_embd_normalize(embd, out, n_embd, 2);112    }113}114 115int main(int argc, char ** argv) {116    std::setlocale(LC_NUMERIC, "C");117 118    common_params params;119 120    common_init();121 122    if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_RETRIEVAL, print_usage)) {123        return 1;124    }125 126    // For BERT models, batch size must be equal to ubatch size127    params.n_ubatch = params.n_batch;128    params.embedding = true;129 130    if (params.chunk_size <= 0) {131        LOG_ERR("chunk_size must be positive\n");132        return 1;133    }134    if (params.context_files.empty()) {135        LOG_ERR("context_files must be specified\n");136        return 1;137    }138 139    LOG_INF("processing files:\n");140    for (auto & context_file : params.context_files) {141        LOG_INF("%s\n", context_file.c_str());142    }143 144    std::vector<chunk> chunks;145    for (auto & context_file : params.context_files) {146        std::vector<chunk> file_chunk = chunk_file(context_file, params.chunk_size, params.chunk_separator);147        chunks.insert(chunks.end(), file_chunk.begin(), file_chunk.end());148    }149    LOG_INF("Number of chunks: %zu\n", chunks.size());150 151    llama_backend_init();152    llama_numa_init(params.numa);153 154    // load the model155    auto llama_init = common_init_from_params(params);156 157    auto * model = llama_init->model();158    auto * ctx   = llama_init->context();159 160    if (model == NULL) {161        LOG_ERR("%s: unable to load model\n", __func__);162        return 1;163    }164 165    const llama_vocab * vocab = llama_model_get_vocab(model);166 167    const int n_ctx_train = llama_model_n_ctx_train(model);168    const int n_ctx = llama_n_ctx(ctx);169 170    const enum llama_pooling_type pooling_type = llama_pooling_type(ctx);171    if (pooling_type == LLAMA_POOLING_TYPE_NONE) {172        LOG_ERR("%s: pooling type NONE not supported\n", __func__);173        return 1;174    }175 176    if (n_ctx > n_ctx_train) {177        LOG_WRN("%s: warning: model was trained on only %d context tokens (%d specified)\n",178                __func__, n_ctx_train, n_ctx);179    }180 181    // print system information182    {183        LOG_INF("\n");184        LOG_INF("%s\n", common_params_get_system_info(params).c_str());185    }186 187    // max batch size188    const uint64_t n_batch = params.n_batch;189    GGML_ASSERT(params.n_batch >= params.n_ctx);190 191    // tokenize the prompts and trim192    for (auto & chunk : chunks) {193        auto inp = common_tokenize(ctx, chunk.textdata, true, false);194        if (inp.size() > n_batch) {195            LOG_ERR("%s: chunk size (%lld) exceeds batch size (%lld), increase batch size and re-run\n",196                    __func__, (long long int) inp.size(), (long long int) n_batch);197            return 1;198        }199        // add eos if not present200        if (llama_vocab_eos(vocab) >= 0 && (inp.empty() || inp.back() != llama_vocab_eos(vocab))) {201            inp.push_back(llama_vocab_eos(vocab));202        }203        chunk.tokens = inp;204    }205 206    // tokenization stats207    if (params.verbose_prompt) {208        for (int i = 0; i < (int) chunks.size(); i++) {209            LOG_INF("%s: prompt %d: '%s'\n", __func__, i, chunks[i].textdata.c_str());210            LOG_INF("%s: number of tokens in prompt = %zu\n", __func__, chunks[i].tokens.size());211            for (int j = 0; j < (int) chunks[i].tokens.size(); j++) {212                LOG_INF("%6d -> '%s'\n", chunks[i].tokens[j], common_token_to_piece(ctx, chunks[i].tokens[j]).c_str());213            }214            LOG_INF("\n\n");215        }216    }217 218    // initialize batch219    const int n_chunks = chunks.size();220    struct llama_batch batch = llama_batch_init(n_batch, 0, 1);221 222    // allocate output223    const int n_embd_out = llama_model_n_embd_out(model);224    std::vector<float> embeddings(n_chunks * n_embd_out, 0);225    float * emb = embeddings.data();226 227    // break into batches228    unsigned int p = 0; // number of prompts processed already229    unsigned int s = 0; // number of prompts in current batch230    for (int k = 0; k < n_chunks; k++) {231        // clamp to n_batch tokens232        auto & inp = chunks[k].tokens;233 234        const uint64_t n_toks = inp.size();235 236        // encode if at capacity237        if (batch.n_tokens + n_toks > n_batch || s >= llama_n_seq_max(ctx)) {238            float * out = emb + p * n_embd_out;239            batch_process(ctx, batch, out, s, n_embd_out);240            common_batch_clear(batch);241            p += s;242            s = 0;243        }244 245        // add to batch246        batch_add_seq(batch, inp, s);247        s += 1;248    }249 250    // final batch251    float * out = emb + p * n_embd_out;252    batch_process(ctx, batch, out, s, n_embd_out);253 254    // save embeddings to chunks255    for (int i = 0; i < n_chunks; i++) {256        chunks[i].embedding = std::vector<float>(emb + i * n_embd_out, emb + (i + 1) * n_embd_out);257        // clear tokens as they are no longer needed258        chunks[i].tokens.clear();259    }260 261    struct llama_batch query_batch = llama_batch_init(n_batch, 0, 1);262 263    // start loop, receive query and return top k similar chunks based on cosine similarity264    std::string query;265    while (true) {266        LOG("Enter query: ");267        std::getline(std::cin, query);268        std::vector<int32_t> query_tokens = common_tokenize(ctx, query, true);269 270        batch_add_seq(query_batch, query_tokens, 0);271 272        std::vector<float> query_emb(n_embd_out, 0);273        batch_process(ctx, query_batch, query_emb.data(), 1, n_embd_out);274 275        common_batch_clear(query_batch);276 277        // compute cosine similarities278        {279            std::vector<std::pair<int, float>> similarities;280            for (int i = 0; i < n_chunks; i++) {281                float sim = common_embd_similarity_cos(chunks[i].embedding.data(), query_emb.data(), n_embd_out);282                similarities.push_back(std::make_pair(i, sim));283            }284 285            // sort similarities286            std::sort(similarities.begin(), similarities.end(), [](const std::pair<int, float> & a, const std::pair<int, float> & b) {287                return a.second > b.second;288            });289 290            LOG("Top %d similar chunks:\n", params.sampling.top_k);291            for (int i = 0; i < std::min(params.sampling.top_k, (int) chunks.size()); i++) {292                LOG("filename: %s\n", chunks[similarities[i].first].filename.c_str());293                LOG("filepos: %lld\n", (long long int) chunks[similarities[i].first].filepos);294                LOG("similarity: %f\n", similarities[i].second);295                LOG("textdata:\n%s\n", chunks[similarities[i].first].textdata.c_str());296                LOG("--------------------\n");297            }298        }299    }300 301    LOG("\n");302    llama_perf_context_print(ctx);303 304    // clean up305    llama_batch_free(query_batch);306    llama_backend_free();307}308