Felipe97/llama-cpp-compiled

1#include "debug.h"2#include "arg.h"3#include "common.h"4#include "log.h"5#include "llama.h"6 7#include <cstdlib>8#include <string>9#include <vector>10#include <filesystem>11#include <fstream>12#include <optional>13#include <regex>14 15static void print_usage(int /*argc*/, char ** argv) {16    const std::string usage_template = R"(17        example usage:18 19          Print tensors:20 21          {prog} -m model.gguf -p "Hello my name is" --verbose22 23          The tensors to be printed can be filtered with --tensor-filter option.24 25          Save logits/embeddings:26 27          {prog} -m model.gguf -p "Hello my name is" --save-logits28 29          Add --embedding to save embeddings)" "\n";30 31    // Fix the source code indentation above that is introduced by the raw string literal.32    std::string usage = std::regex_replace(usage_template, std::regex("\\n {8}"), "\n");33    usage = std::regex_replace(usage, std::regex("\\{prog\\}"), argv[0]);34    LOG("%s\n", usage.c_str());35}36 37static bool has_pooling(llama_context * ctx) {38    switch (llama_pooling_type(ctx)) {39        case LLAMA_POOLING_TYPE_NONE:40        case LLAMA_POOLING_TYPE_UNSPECIFIED:41            return false;42        default:43            return true;44    }45}46 47struct output_data {48    float *                  data_ptr    = nullptr;49    int                      data_size   = 0;50    std::string              type_suffix;51    std::vector<float>       embd_norm;52    std::string              prompt;53    std::vector<llama_token> tokens;54 55    output_data(llama_context * ctx, const llama_model * model, const common_params & params) {56        const llama_vocab * vocab = llama_model_get_vocab(model);57        const bool add_bos = llama_vocab_get_add_bos(vocab);58 59        tokens = common_tokenize(ctx, params.prompt, add_bos);60        prompt = params.prompt;61 62        if (params.embedding) {63            const int n_embd       = llama_model_n_embd_out(model);64            const bool pooling     = has_pooling(ctx);65            const int n_embd_count = pooling ? 1 : tokens.size();66            const int n_floats     = n_embd * n_embd_count;67 68            float * embd_raw = pooling ? llama_get_embeddings_seq(ctx, 0) : llama_get_embeddings(ctx);69            if (embd_raw == nullptr) {70                throw std::runtime_error("failed to get embeddings from the model");71            }72 73            LOG_DBG("pooling_enabled: %s\n", pooling ? "true" : "false");74            LOG_DBG("n_embd: %d\n", n_embd);75            LOG_DBG("n_floats: %d\n", n_floats);76            LOG_DBG("n_embd_count: %d\n", n_embd_count);77 78            data_ptr    = embd_raw;79            data_size   = n_floats;80            type_suffix = "-embeddings";81 82            if (params.embd_normalize >= 0) {83                embd_norm.resize(n_floats);84                for (int i = 0; i < n_embd_count; i++) {85                    common_embd_normalize(embd_raw+i*n_embd, embd_norm.data()+i*n_embd, n_embd, params.embd_normalize);86                }87                data_ptr = embd_norm.data();88            }89        } else {90            const float * logits = llama_get_logits_ith(ctx, tokens.size() - 1);91            const int n_logits = llama_vocab_n_tokens(vocab);92 93            data_ptr = const_cast<float*>(logits);94            data_size = n_logits;95            type_suffix = "";96        }97    }98};99 100static void save_output_data(const output_data & output, const std::string & model_name, const std::string & output_dir) {101    std::filesystem::create_directory(output_dir);102    auto base_path = std::filesystem::path{output_dir} / ("llamacpp-" + model_name + output.type_suffix);103 104    // Save logits/embeddings to binary file.105    {106        std::filesystem::path filepath{base_path.string() + ".bin"};107        std::ofstream file{filepath, std::ios::binary};108        if (!file) {109            throw std::runtime_error("failed to open binary output file: " + filepath.string());110        }111        file.write(reinterpret_cast<const char*>(output.data_ptr), output.data_size * sizeof(float));112        LOG("Data saved to %s\n", filepath.c_str());113    }114 115    // Save logits/embeddings to text file.116    {117        std::filesystem::path filepath{base_path.string() + ".txt"};118        std::ofstream file{filepath};119        if (!file) {120            throw std::runtime_error("failed to open text output file: " + filepath.string());121        }122        for (int i = 0; i < output.data_size; i++) {123            file << i << ": " << output.data_ptr[i] << '\n';124        }125        LOG("Data saved to %s\n", filepath.c_str());126    }127 128    // Save prompt and tokens to text file.129    {130        std::filesystem::path filepath{base_path.string() + "-prompt.txt"};131        std::ofstream file{filepath};132        if (!file) {133            throw std::runtime_error("failed to open prompt output file: " + filepath.string());134        }135 136        file << "prompt: " << output.prompt << '\n';137        file << "n_tokens: " << output.tokens.size() << '\n';138 139        file << "token ids: ";140        for (size_t i = 0; i < output.tokens.size(); i++) {141            file << output.tokens[i];142            if (i + 1 < output.tokens.size()) {143                file << ", ";144            }145        }146        file << '\n';147        LOG("Prompt saved to %s\n", filepath.c_str());148    }149 150    // Save token ids to binary file.151    {152        std::filesystem::path filepath{base_path.string() + "-tokens.bin"};153        std::ofstream file{filepath, std::ios::binary};154        if (!file) {155            throw std::runtime_error("failed to open tokens binary file: " + filepath.string());156        }157        file.write(reinterpret_cast<const char*>(output.tokens.data()), output.tokens.size() * sizeof(llama_token));158        LOG("Tokens saved to %s\n", filepath.c_str());159    }160 161}162 163static void print_tokenized_prompt(llama_context * ctx, const std::vector<llama_token> & tokens, const std::string & prompt) {164    const llama_model * model = llama_get_model(ctx);165    const llama_vocab * vocab = llama_model_get_vocab(model);166 167    LOG("Model add_bos: %s\n", llama_vocab_get_add_bos(vocab) ? "true" : "false");168    LOG("Input prompt: \"%s\"\n", prompt.c_str());169    LOG("Token ids (%zu):\n", tokens.size());170 171    for (auto id : tokens) {172        std::string piece(128, '\0');173        int n = llama_token_to_piece(vocab, id, piece.data(), piece.size(), 0, true);174        if (n < 0) {175            LOG_ERR("failed to convert token %d to piece\n", id);176            continue;177        }178        piece.resize(n);179        LOG("%s(%d) ", piece.c_str(), id);180    }181    LOG("\n");182}183 184static bool run(llama_context * ctx, const common_params & params) {185    const llama_model * model = llama_get_model(ctx);186    const llama_vocab * vocab = llama_model_get_vocab(model);187 188    const bool add_bos = llama_vocab_get_add_bos(vocab);189 190    std::vector<llama_token> tokens = common_tokenize(ctx, params.prompt, add_bos);191 192    if (tokens.empty()) {193        LOG_ERR("%s : there are not input tokens to process - (try to provide a prompt with '-p')\n", __func__);194        return false;195    }196 197    if (llama_decode(ctx, llama_batch_get_one(tokens.data(), tokens.size()))) {198        LOG_ERR("%s : failed to eval\n", __func__);199        return false;200    }201 202    print_tokenized_prompt(ctx, tokens, params.prompt);203 204    if (params.save_logits) {205        try {206            output_data output {ctx, model, params};207            std::filesystem::path model_path{params.model.path};208            std::string model_name{model_path.stem().string()};209            save_output_data(output, model_name, params.logits_output_dir);210        } catch (const std::exception & e) {211            LOG_ERR("%s : error saving logits: %s\n", __func__, e.what());212        }213    }214 215    return true;216}217 218int main(int argc, char ** argv) {219    common_params params;220 221    common_init();222 223    if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_DEBUG, print_usage)) {224        return 1;225    }226 227    llama_backend_init();228    llama_numa_init(params.numa);229 230    std::optional<common_debug_cb_user_data> cb_data;231    if (!params.save_logits) {232        cb_data.emplace(params, params.tensor_filter);233    }234 235    auto llama_init = common_init_from_params(params);236 237    auto * model = llama_init->model();238    auto * ctx   = llama_init->context();239 240    if (model == nullptr || ctx == nullptr) {241        LOG_ERR("%s : failed to init\n", __func__);242        return 1;243    }244 245    {246        LOG_INF("\n");247        LOG_INF("%s\n", common_params_get_system_info(params).c_str());248        LOG_INF("\n");249    }250 251    if (!run(ctx, params)) {252        return 1;253    }254 255    LOG("\n");256    llama_perf_context_print(ctx);257 258    llama_backend_free();259 260    return 0;261}262