#include "kv-mean-center.h"
#include "log.h"
#include "ggml.h"
#include "gguf.h"
#include <cstring>
bool common_kv_mean_center_write(
const std::string & fname,
const std::vector<common_kv_mean_center_layer> & layers,
bool k_rot) {
size_t n_with_bias = 0;
for (const auto & layer : layers) {
if (!layer.bias.empty()) {
n_with_bias++;
}
}
if (n_with_bias == 0) {
LOG_ERR("%s: no layers with bias data to write\n", __func__);
return false;
}
size_t data_size = 0;
for (const auto & layer : layers) {
if (!layer.bias.empty()) {
data_size += GGML_PAD(ggml_tensor_overhead() + sizeof(float)*layer.bias.size(), GGML_MEM_ALIGN);
}
}
struct ggml_init_params params = {
/*.mem_size =*/ data_size,
/*.mem_buffer =*/ NULL,
/*.no_alloc =*/ false,
};
struct ggml_context * ctx = ggml_init(params);
struct gguf_context * ctx_gguf = gguf_init_empty();
if (!ctx || !ctx_gguf) {
LOG_ERR("%s: failed to allocate ggml/gguf context\n", __func__);
if (ctx) ggml_free(ctx);
if (ctx_gguf) gguf_free(ctx_gguf);
return false;
}
gguf_set_val_str(ctx_gguf, "general.type", "kv-mean-center");
gguf_set_val_bool(ctx_gguf, "kv_mean_center.k_rot", k_rot);
for (const auto & layer : layers) {
if (layer.bias.empty()) {
continue;
}
const std::string name = "kv_bar.blk." + std::to_string(layer.il) + ".k";
ggml_tensor * t = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, (int64_t) layer.bias.size());
ggml_set_name(t, name.c_str());
memcpy(t->data, layer.bias.data(), layer.bias.size()*sizeof(float));
gguf_add_tensor(ctx_gguf, t);
}
const bool ok = gguf_write_to_file(ctx_gguf, fname.c_str(), false);
if (!ok) {
LOG_ERR("%s: failed to write %s\n", __func__, fname.c_str());
}
gguf_free(ctx_gguf);
ggml_free(ctx);
return ok;
}