Brunobkr/llama.cpp_AlgMor24_github

1#include "imatrix-loader.h"2#include "common.h"3#include "log.h"4#include "gguf.h"5 6#include <cmath>7#include <cstring>8#include <fstream>9 10static bool common_imatrix_load_legacy(const std::string & fname, common_imatrix & imatrix) {11    std::ifstream in(fname, std::ios::binary);12    if (!in) {13        LOG_ERR("%s: failed to open %s\n", __func__, fname.c_str());14        return false;15    }16 17    int n_entries;18    in.read((char *) &n_entries, sizeof(n_entries));19    if (in.fail() || n_entries < 1) {20        LOG_ERR("%s: no data in file %s\n", __func__, fname.c_str());21        return false;22    }23 24    for (int i = 0; i < n_entries; ++i) {25        int32_t len = 0;26        in.read((char *) &len, sizeof(len));27        std::vector<char> name_as_vec(len + 1);28        in.read((char *) name_as_vec.data(), len);29        if (in.fail()) {30            LOG_ERR("%s: failed reading name for entry %d from %s\n", __func__, i + 1, fname.c_str());31            return false;32        }33        name_as_vec[len] = 0;34        std::string name{ name_as_vec.data() };35 36        int32_t ncall = 0;37        in.read((char *) &ncall, sizeof(ncall));38        int32_t nval = 0;39        in.read((char *) &nval, sizeof(nval));40        if (in.fail() || nval < 1) {41            LOG_ERR("%s: failed reading number of values for entry %d\n", __func__, i);42            return false;43        }44 45        auto & e = imatrix.entries[std::move(name)];46        e.sums.resize(nval);47        in.read((char *) e.sums.data(), nval * sizeof(float));48        if (in.fail()) {49            LOG_ERR("%s: failed reading data for entry %d\n", __func__, i);50            return false;51        }52 53        e.counts.resize(1);54        e.counts[0] = ncall;55    }56 57    // the trailing data (chunk count + dataset name) is optional58    if (in.peek() != EOF) {59        int32_t n_calls = 0;60        in.read((char *) &n_calls, sizeof(n_calls));61        imatrix.chunk_count = n_calls;62 63        if (!in.fail()) {64            int32_t len = 0;65            in.read((char *) &len, sizeof(len));66            if (!in.fail() && len > 0) {67                std::vector<char> dataset(len + 1, 0);68                in.read(dataset.data(), len);69                if (!in.fail()) {70                    imatrix.datasets.push_back(dataset.data());71                }72            }73        }74    }75 76    imatrix.chunk_size = 0;77    imatrix.is_legacy  = true;78 79    return true;80}81 82bool common_imatrix_load(const std::string & fname, common_imatrix & imatrix) {83    struct ggml_context * ctx = nullptr;84    struct gguf_init_params meta_gguf_params = {85        /* .no_alloc = */ false,86        /* .ctx      = */ &ctx,87    };88    struct gguf_context * ctx_gguf = gguf_init_from_file(fname.c_str(), meta_gguf_params);89    if (!ctx_gguf) {90        return common_imatrix_load_legacy(fname, imatrix);91    }92 93    const int32_t n_entries = gguf_get_n_tensors(ctx_gguf);94    if (n_entries < 1) {95        LOG_ERR("%s: no data in file %s\n", __func__, fname.c_str());96        gguf_free(ctx_gguf);97        ggml_free(ctx);98        return false;99    }100 101    const int64_t datasets_key   = gguf_find_key(ctx_gguf, LLM_KV_IMATRIX_DATASETS);102    const int64_t chunk_count_key = gguf_find_key(ctx_gguf, LLM_KV_IMATRIX_CHUNK_COUNT);103    const int64_t chunk_size_key  = gguf_find_key(ctx_gguf, LLM_KV_IMATRIX_CHUNK_SIZE);104 105    if (datasets_key != -1 && gguf_get_arr_type(ctx_gguf, datasets_key) == GGUF_TYPE_STRING) {106        const int64_t n = gguf_get_arr_n(ctx_gguf, datasets_key);107        imatrix.datasets.reserve(imatrix.datasets.size() + n);108        for (int64_t i = 0; i < n; ++i) {109            imatrix.datasets.push_back(gguf_get_arr_str(ctx_gguf, datasets_key, i));110        }111    }112 113    imatrix.has_metadata = (datasets_key != -1 && chunk_count_key != -1 && chunk_size_key != -1);114    imatrix.chunk_count  = (chunk_count_key != -1) ? gguf_get_val_u32(ctx_gguf, chunk_count_key) : 0;115    imatrix.chunk_size   = (chunk_size_key  != -1) ? gguf_get_val_u32(ctx_gguf, chunk_size_key)  : 0;116 117    const std::string in_sum2_suffix{ ".in_sum2" };118    const std::string counts_suffix{ ".counts" };119 120    std::map<std::string, std::pair<struct ggml_tensor *, struct ggml_tensor *>> sums_counts_for;121 122    for (struct ggml_tensor * cur = ggml_get_first_tensor(ctx); cur; cur = ggml_get_next_tensor(ctx, cur)) {123        std::string name = cur->name;124 125        if (name.empty()) { continue; }126 127        if (string_remove_suffix(name, in_sum2_suffix)) {128            sums_counts_for[std::move(name)].first = cur;129        } else if (string_remove_suffix(name, counts_suffix)) {130            sums_counts_for[std::move(name)].second = cur;131        }132    }133 134    for (const auto & sc : sums_counts_for) {135        const std::string &        name    = sc.first;136        const struct ggml_tensor * in_sum2 = sc.second.first;137        const struct ggml_tensor * counts  = sc.second.second;138 139        if (!in_sum2 || !counts) {140            LOG_ERR("%s: mismatched sums and counts for %s\n", __func__, name.c_str());141            gguf_free(ctx_gguf);142            ggml_free(ctx);143            return false;144        }145 146        auto & e = imatrix.entries[name];147 148        const int64_t nval    = ggml_nelements(in_sum2);149        const int64_t ncounts = ggml_nelements(counts);150 151        e.sums.resize(nval);152        for (int64_t j = 0; j < nval; ++j) {153            e.sums[j] = ((const float *) in_sum2->data)[j];154        }155 156        e.counts.resize(ncounts);157        for (int64_t j = 0; j < ncounts; ++j) {158            e.counts[j] = std::lround(((const float *) counts->data)[j]);159        }160    }161 162    gguf_free(ctx_gguf);163    ggml_free(ctx);164    return true;165}166