#include "llama-model-saver.h"

#include "ggml.h"
#include "gguf.h"

#include "llama-arch.h"
#include "llama.h"
#include "llama-hparams.h"
#include "llama-model.h"
#include "llama-vocab.h"

#include <cstdint>
#include <string>
#include <unordered_set>

bool llama_model_saver_supports_arch(llm_arch arch) {
    switch (arch) {
        case LLM_ARCH_GEMMA3N:
        case LLM_ARCH_BITNET:
        case LLM_ARCH_T5:
        case LLM_ARCH_APERTUS:
        case LLM_ARCH_STEP35:
        case LLM_ARCH_CLEF: // the head tensors are not saved
            return false;
        default:
            return true;
    }
}

llama_model_saver::llama_model_saver(const struct llama_model * model) :
        gguf_ctx(gguf_init_empty()), gguf_ctx_owned(true), model(model), llm_kv(model->arch) {
    GGML_ASSERT(llama_model_saver_supports_arch(model->arch));
}

llama_model_saver::llama_model_saver(enum llm_arch arch, struct gguf_context * gguf_ctx) :
        gguf_ctx(gguf_ctx == nullptr ? gguf_init_empty() : gguf_ctx), gguf_ctx_owned(gguf_ctx == nullptr), model(nullptr), llm_kv(arch) {}

llama_model_saver::~llama_model_saver() {
    if (gguf_ctx_owned) {
        gguf_free(gguf_ctx);
    }
}

void llama_model_saver::add_kv(const enum llm_kv key, const uint32_t value) {
    gguf_set_val_u32(gguf_ctx, llm_kv(key).c_str(), value);
}

void llama_model_saver::add_kv(const enum llm_kv key, const int32_t value) {
    gguf_set_val_i32(gguf_ctx, llm_kv(key).c_str(), value);
}

void llama_model_saver::add_kv(const enum llm_kv key, const uint64_t value) {
    gguf_set_val_u64(gguf_ctx, llm_kv(key).c_str(), value);
}

void llama_model_saver::add_kv(const enum llm_kv key, const float value) {
    gguf_set_val_f32(gguf_ctx, llm_kv(key).c_str(), value);
}

void llama_model_saver::add_kv(const enum llm_kv key, const bool value) {
    gguf_set_val_bool(gguf_ctx, llm_kv(key).c_str(), value);
}

void llama_model_saver::add_kv(const enum llm_kv key, const char * value) {
    gguf_set_val_str(gguf_ctx, llm_kv(key).c_str(), value);
}

[[noreturn]]
void llama_model_saver::add_kv(const enum llm_kv key, const char value) {
    GGML_UNUSED(key);
    GGML_UNUSED(value);
    GGML_ABORT("fatal error"); // this should never be called, only needed to make the template below compile
}

template <typename Container>
void llama_model_saver::add_kv(const enum llm_kv key, const Container & value, const bool per_layer) {
    GGML_ASSERT(model != nullptr || !per_layer);
    const size_t n_values = per_layer ? size_t(model->hparams.n_layer()) : value.size();
    GGML_ASSERT(n_values <= value.size());

    if (n_values == 0) {
        return;
    }

    if (per_layer) {
        bool all_values_the_same = true;
        for (size_t i = 1; i < n_values; ++i) {
            if (value[i] != value[0]) {
                all_values_the_same = false;
                break;
            }
        }
        if (all_values_the_same) {
            add_kv(key, value[0]);
            return;
        }
    }

    if (std::is_same<typename Container::value_type, uint8_t>::value) {
        gguf_set_arr_data(gguf_ctx, llm_kv(key).c_str(), GGUF_TYPE_UINT8, value.data(), n_values);
    } else if (std::is_same<typename Container::value_type, int8_t>::value) {
        gguf_set_arr_data(gguf_ctx, llm_kv(key).c_str(), GGUF_TYPE_INT8, value.data(), n_values);
    } else if (std::is_same<typename Container::value_type, uint32_t>::value) {
        gguf_set_arr_data(gguf_ctx, llm_kv(key).c_str(), GGUF_TYPE_UINT32, value.data(), n_values);
    } else if (std::is_same<typename Container::value_type, bool>::value) {
        gguf_set_arr_data(gguf_ctx, llm_kv(key).c_str(), GGUF_TYPE_BOOL, value.data(), n_values);
    } else if (std::is_same<typename Container::value_type, int32_t>::value) {
        gguf_set_arr_data(gguf_ctx, llm_kv(key).c_str(), GGUF_TYPE_INT32, value.data(), n_values);
    } else if (std::is_same<typename Container::value_type, uint64_t>::value) {
        gguf_set_arr_data(gguf_ctx, llm_kv(key).c_str(), GGUF_TYPE_UINT64, value.data(), n_values);
    } else if (std::is_same<typename Container::value_type, float>::value) {
        gguf_set_arr_data(gguf_ctx, llm_kv(key).c_str(), GGUF_TYPE_FLOAT32, value.data(), n_values);
    } else if (std::is_same<Container, std::string>::value) {
        gguf_set_val_str(gguf_ctx, llm_kv(key).c_str(), reinterpret_cast<const char *>(value.data()));
    } else {
        GGML_ABORT("fatal error");
    }
}
// instantiate for external usage:
template void llama_model_saver::add_kv<std::vector<uint32_t>>(const enum llm_kv, const std::vector<uint32_t> &, const bool);
template void llama_model_saver::add_kv<std::vector<float>>(const enum llm_kv, const std::vector<float> &, const bool);
template void llama_model_saver::add_kv<std::vector<uint64_t>>(const enum llm_kv, const std::vector<uint64_t> &, const bool);

void llama_model_saver::add_kv(const enum llm_kv key, const std::vector<std::string> & value) {
    std::vector<const char *> tmp(value.size());
    for (size_t i = 0; i < value.size(); ++i) {
        tmp[i] = value[i].c_str();
    }
    gguf_set_arr_str(gguf_ctx, llm_kv(key).c_str(), tmp.data(), tmp.size());
}

void llama_model_saver::add_tensor(const struct ggml_tensor * tensor) {
    if (!tensor) {
        return;
    }
    if (gguf_find_tensor(gguf_ctx, tensor->name) >= 0) {
        const std::string tensor_name = tensor->name;
        GGML_ASSERT(
            tensor_name == "rope_freqs.weight" || tensor_name == "rope_factors_long.weight" ||
            tensor_name == "rope_factors_short.weight"); // FIXME
        return;
    }
    gguf_add_tensor(gguf_ctx, tensor);
}

void llama_model_saver::add_kv_from_model() {
    const llama_hparams & hparams = model->hparams;
    const llama_vocab   & vocab   = model->vocab;

    const int32_t n_vocab = vocab.n_tokens();
    std::vector<std::string> tokens(n_vocab);
    std::vector<float>       scores(n_vocab);
    std::vector<int32_t>     token_types(n_vocab);

    if (vocab.get_type() != LLAMA_VOCAB_TYPE_NONE) {
        for (int32_t id = 0; id < n_vocab; ++id) {
            const llama_vocab::token_data & token_data = vocab.get_token_data(id);

            tokens[id] = token_data.text;
            scores[id] = token_data.score;

            // FIXME should this be treated as flags?
            switch(token_data.attr) {
                case LLAMA_TOKEN_ATTR_UNKNOWN:      token_types[id] = LLAMA_TOKEN_TYPE_UNKNOWN;      break;
                case LLAMA_TOKEN_ATTR_UNUSED:       token_types[id] = LLAMA_TOKEN_TYPE_UNUSED;       break;
                case LLAMA_TOKEN_ATTR_NORMAL:       token_types[id] = LLAMA_TOKEN_TYPE_NORMAL;       break;
                case LLAMA_TOKEN_ATTR_CONTROL:      token_types[id] = LLAMA_TOKEN_TYPE_CONTROL;      break;
                case LLAMA_TOKEN_ATTR_USER_DEFINED: token_types[id] = LLAMA_TOKEN_TYPE_USER_DEFINED; break;
                case LLAMA_TOKEN_ATTR_BYTE:         token_types[id] = LLAMA_TOKEN_TYPE_BYTE;         break;
                // case LLAMA_TOKEN_ATTR_NORMALIZED:   ???
                // case LLAMA_TOKEN_ATTR_LSTRIP:       ???
                // case LLAMA_TOKEN_ATTR_RSTRIP:       ???
                case LLAMA_TOKEN_ATTR_UNDEFINED:
                default:                            token_types[id] = LLAMA_TOKEN_TYPE_UNDEFINED;    break;
            }
        }
    }

    // add_kv(LLM_KV_GENERAL_TYPE,                      ???);
    add_kv(LLM_KV_GENERAL_ARCHITECTURE,              model->arch_name());
    // add_kv(LLM_KV_GENERAL_QUANTIZATION_VERSION,      ???);
    // add_kv(LLM_KV_GENERAL_ALIGNMENT,                 ???);
    // add_kv(LLM_KV_GENERAL_FILE_TYPE,                 ???);
    // add_kv(LLM_KV_GENERAL_SAMPLING_SEQUENCE,         ???);
    // add_kv(LLM_KV_GENERAL_SAMPLING_TOP_K,            ???);
    // add_kv(LLM_KV_GENERAL_SAMPLING_TOP_P,            ???);
    // add_kv(LLM_KV_GENERAL_SAMPLING_MIN_P,            ???);
    // add_kv(LLM_KV_GENERAL_SAMPLING_XTC_PROBABILITY,  ???);
    // add_kv(LLM_KV_GENERAL_SAMPLING_XTC_THRESHOLD,    ???);
    // add_kv(LLM_KV_GENERAL_SAMPLING_TEMP,             ???);
    // add_kv(LLM_KV_GENERAL_SAMPLING_PENALTY_LAST_N,   ???);
    // add_kv(LLM_KV_GENERAL_SAMPLING_PENALTY_REPEAT,   ???);
    // add_kv(LLM_KV_GENERAL_SAMPLING_MIROSTAT,         ???);
    // add_kv(LLM_KV_GENERAL_SAMPLING_MIROSTAT_TAU,     ???);
    // add_kv(LLM_KV_GENERAL_SAMPLING_MIROSTAT_ETA,     ???);
    add_kv(LLM_KV_GENERAL_NAME,                      model->name);

    if (!model->prec_policy.prec_src1.empty()) {
        std::vector<std::string> tensor_names;
        std::vector<int8_t> values;
        tensor_names.reserve(model->prec_policy.prec_src1.size());
        values.reserve(model->prec_policy.prec_src1.size());
        for (const auto & [w, prec] : model->prec_policy.prec_src1) {
            tensor_names.push_back(ggml_get_name(w));
            values.push_back(prec == GGML_PREC_Q8 ? 0 : 1);
        }
        add_kv(LLM_KV_GENERAL_TENSOR_EXTRA_NAME, tensor_names);
        gguf_set_arr_data(gguf_ctx, llm_kv(LLM_KV_GENERAL_TENSOR_EXTRA_PREC_A4).c_str(), GGUF_TYPE_BOOL, values.data(), values.size());
    }
    // add_kv(LLM_KV_GENERAL_AUTHOR,                    ???);
    // add_kv(LLM_KV_GENERAL_VERSION,                   ???);
    // add_kv(LLM_KV_GENERAL_URL,                       ???);
    // add_kv(LLM_KV_GENERAL_DESCRIPTION,               ???);
    // add_kv(LLM_KV_GENERAL_LICENSE,                   ???);
    // add_kv(LLM_KV_GENERAL_SOURCE_URL,                ???);
    // add_kv(LLM_KV_GENERAL_SOURCE_HF_REPO,            ???);

    add_kv(LLM_KV_VOCAB_SIZE,                        vocab.n_tokens());
    add_kv(LLM_KV_CONTEXT_LENGTH,                    hparams.n_ctx_train);
    add_kv(LLM_KV_EMBEDDING_LENGTH,                  hparams.n_embd);
    if (hparams.n_embd_out_impl > 0) {
        add_kv(LLM_KV_EMBEDDING_LENGTH_OUT,          hparams.n_embd_out_impl);
    }
    add_kv(LLM_KV_BLOCK_COUNT,                       hparams.n_layer_all);
    add_kv(LLM_KV_LEADING_DENSE_BLOCK_COUNT,         hparams.n_layer_dense_lead);
    add_kv(LLM_KV_FEED_FORWARD_LENGTH,               hparams.n_ff_arr, true);
    add_kv(LLM_KV_EXPERT_FEED_FORWARD_LENGTH,        hparams.n_ff_exp());
    add_kv(LLM_KV_EXPERT_LATENT_LENGTH,              hparams.n_expert_latent);
    add_kv(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp);
    add_kv(LLM_KV_EXPERT_CHUNK_FEED_FORWARD_LENGTH,  hparams.n_ff_chexp);
    add_kv(LLM_KV_SWIGLU_CLAMP_EXP, std::vector<float>(
            hparams.swiglu_clamp_exp.begin(), hparams.swiglu_clamp_exp.begin() + hparams.n_layer_all));
    add_kv(LLM_KV_SWIGLU_CLAMP_SHEXP, std::vector<float>(
            hparams.swiglu_clamp_shexp.begin(), hparams.swiglu_clamp_shexp.begin() + hparams.n_layer_all));
    add_kv(LLM_KV_USE_PARALLEL_RESIDUAL,             hparams.use_par_res);
    // add_kv(LLM_KV_TENSOR_DATA_LAYOUT,                ???);
    add_kv(LLM_KV_EXPERT_COUNT,                      hparams.n_expert);
    add_kv(LLM_KV_EXPERT_USED_COUNT,                 hparams.n_expert_used());
    add_kv(LLM_KV_EXPERT_SHARED_COUNT,               hparams.n_expert_shared);
    add_kv(LLM_KV_EXPERT_GROUP_COUNT,                hparams.n_expert_groups);
    add_kv(LLM_KV_EXPERT_GROUP_USED_COUNT,           hparams.n_group_used);
    add_kv(LLM_KV_EXPERT_WEIGHTS_SCALE,              hparams.expert_weights_scale);
    add_kv(LLM_KV_EXPERT_WEIGHTS_NORM,               hparams.expert_weights_norm);
    add_kv(LLM_KV_EXPERT_GATING_FUNC,                hparams.expert_gating_func);
    add_kv(LLM_KV_EXPERT_GROUP_SCALE,                hparams.expert_group_scale);
    add_kv(LLM_KV_EXPERTS_PER_GROUP,                 hparams.n_group_experts);
    add_kv(LLM_KV_MOE_EVERY_N_LAYERS,                hparams.moe_every_n_layers);
    add_kv(LLM_KV_NEXTN_PREDICT_LAYERS,              hparams.n_layer_nextn);
    add_kv(LLM_KV_NUM_DEEPSTACK_LAYERS,              hparams.n_deepstack_layers);
    add_kv(LLM_KV_DEEPSTACK_MAPPING,                 hparams.deepstack_mapping_arr);
    add_kv(LLM_KV_POOLING_TYPE,                      uint32_t(hparams.pooling_type));
    add_kv(LLM_KV_LOGIT_SCALE,                       hparams.f_logit_scale);
    add_kv(LLM_KV_DECODER_START_TOKEN_ID,            hparams.dec_start_token_id);
    add_kv(LLM_KV_DECODER_BLOCK_COUNT,               hparams.dec_n_layer);
    add_kv(LLM_KV_ATTN_LOGIT_SOFTCAPPING,            hparams.f_attn_logit_softcapping);
    add_kv(LLM_KV_ROUTER_LOGIT_SOFTCAPPING,          hparams.f_router_logit_softcapping);
    add_kv(LLM_KV_FINAL_LOGIT_SOFTCAPPING,           hparams.f_final_logit_softcapping);
    add_kv(LLM_KV_SWIN_NORM,                         hparams.swin_norm);
    add_kv(LLM_KV_RESCALE_EVERY_N_LAYERS,            hparams.rescale_every_n_layers);
    add_kv(LLM_KV_TIME_MIX_EXTRA_DIM,                hparams.time_mix_extra_dim);
    add_kv(LLM_KV_TIME_DECAY_EXTRA_DIM,              hparams.time_decay_extra_dim);
    add_kv(LLM_KV_RESIDUAL_SCALE,                    hparams.f_residual_scale);
    add_kv(LLM_KV_EMBEDDING_SCALE,                   hparams.f_embedding_scale);
    add_kv(LLM_KV_HRM_LAYERS_PER_STACK,              hparams.n_hrm_layers_per_stack);
    add_kv(LLM_KV_HRM_H_CYCLES,                      hparams.n_hrm_h_cycles);
    add_kv(LLM_KV_HRM_L_CYCLES,                      hparams.n_hrm_l_cycles);
    add_kv(LLM_KV_HRM_PREFIX_LM,                     hparams.hrm_prefix_lm);
    add_kv(LLM_KV_TOKEN_SHIFT_COUNT,                 hparams.token_shift_count);
    add_kv(LLM_KV_INTERLEAVE_MOE_LAYER_STEP,         hparams.n_moe_layer_step);
    // add_kv(LLM_KV_FULL_ATTENTION_INTERVAL,           ???); // saved as LLM_KV_ATTENTION_RECURRENT_LAYERS instead

    add_kv(LLM_KV_ATTENTION_HEAD_COUNT,              hparams.n_head_arr, true);
    add_kv(LLM_KV_ATTENTION_HEAD_COUNT_KV,           hparams.n_head_kv_arr, true);
    add_kv(LLM_KV_ATTENTION_MAX_ALIBI_BIAS,          hparams.f_max_alibi_bias);
    add_kv(LLM_KV_ATTENTION_CLAMP_KQV,               hparams.f_clamp_kqv);
    add_kv(LLM_KV_ATTENTION_KEY_LENGTH,              hparams.n_embd_head_k_full);
    add_kv(LLM_KV_ATTENTION_VALUE_LENGTH,            hparams.n_embd_head_v_full);
    add_kv(LLM_KV_ATTENTION_LAYERNORM_EPS,           hparams.f_norm_eps);
    add_kv(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS,       hparams.f_norm_rms_eps);
    add_kv(LLM_KV_ATTENTION_GROUPNORM_EPS,           hparams.f_norm_group_eps);
    add_kv(LLM_KV_ATTENTION_GROUPNORM_GROUPS,        hparams.n_norm_groups);
    add_kv(LLM_KV_ATTENTION_CAUSAL,                  hparams.causal_attn);
    add_kv(LLM_KV_ATTENTION_Q_LORA_RANK,             hparams.n_lora_q);
    add_kv(LLM_KV_ATTENTION_KV_LORA_RANK,            hparams.n_lora_kv);
    add_kv(LLM_KV_ATTENTION_DECAY_LORA_RANK,         hparams.n_lora_decay);
    add_kv(LLM_KV_ATTENTION_ICLR_LORA_RANK,          hparams.n_lora_iclr);
    add_kv(LLM_KV_ATTENTION_VALUE_RESIDUAL_MIX_LORA_RANK, hparams.n_lora_value_res_mix);
    add_kv(LLM_KV_ATTENTION_GATE_LORA_RANK,          hparams.n_lora_gate);
    add_kv(LLM_KV_ATTENTION_RELATIVE_BUCKETS_COUNT,  hparams.n_rel_attn_bkts);
    add_kv(LLM_KV_ATTENTION_ROPE_PATTERN,            hparams.rope_pattern, true);
    add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW,          hparams.n_swa);
    if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) {
        // never collapsed to a scalar: the loaders read a scalar as a period
        add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, std::vector<uint32_t>(
                hparams.is_swa_impl.begin(), hparams.is_swa_impl.begin() + hparams.n_layer_all));
    }
    add_kv(LLM_KV_ATTENTION_SCALE,                   hparams.f_attention_scale);
    add_kv(LLM_KV_ATTENTION_OUTPUT_SCALE,            hparams.f_attn_out_scale);
    add_kv(LLM_KV_ATTENTION_VALUE_SCALE,             hparams.f_attn_value_scale);
    add_kv(LLM_KV_ATTENTION_TEMPERATURE_LENGTH,      hparams.attn_temp_length);
    add_kv(LLM_KV_ATTENTION_TEMPERATURE_SCALE,       hparams.f_attn_temp_scale);
    add_kv(LLM_KV_ATTENTION_KEY_LENGTH_MLA,          hparams.n_embd_head_k_mla_impl);
    add_kv(LLM_KV_ATTENTION_VALUE_LENGTH_MLA,        hparams.n_embd_head_v_mla_impl);
    add_kv(LLM_KV_ATTENTION_KEY_LENGTH_SWA,          hparams.n_embd_head_k_swa);
    add_kv(LLM_KV_ATTENTION_VALUE_LENGTH_SWA,        hparams.n_embd_head_v_swa);
    add_kv(LLM_KV_ATTENTION_KEY_LENGTH_MLA_SWA,      hparams.n_embd_head_k_mla_swa);
    add_kv(LLM_KV_ATTENTION_VALUE_LENGTH_MLA_SWA,    hparams.n_embd_head_v_mla_swa);
    add_kv(LLM_KV_ATTENTION_KV_LORA_RANK_SWA,        hparams.n_lora_kv_swa);
    add_kv(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT,      hparams.indexer_n_head);
    add_kv(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH,      hparams.indexer_head_size);
    add_kv(LLM_KV_ATTENTION_INDEXER_TOP_K,           hparams.indexer_top_k);
    add_kv(LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE,      hparams.indexer_block_size);
    add_kv(LLM_KV_ATTENTION_INDEXER_KPOOL,           hparams.indexer_kpool);
    add_kv(LLM_KV_ATTENTION_INDEXER_KPOOL_SELECT_TAIL, hparams.indexer_kpool_select_tail);
    add_kv(LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS,    hparams.indexer_local_blocks);
    add_kv(LLM_KV_ATTENTION_INDEXER_TYPES,           hparams.is_indexer_full_impl, true);
    add_kv(LLM_KV_ATTENTION_RECURRENT_LAYERS,        hparams.is_recr_impl, true);
    add_kv(LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT,      hparams.dsv4_o_group_count);
    add_kv(LLM_KV_ATTENTION_OUTPUT_LORA_RANK,        hparams.dsv4_o_lora_rank);
    add_kv(LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE, hparams.dsv4_compress_rope_base);
    if (model->arch == LLM_ARCH_DEEPSEEK4 || hparams.dsv4_hc_mult > 0) {
        // the loader requires one compress ratio per layer, including nextn layers
        const std::vector<uint32_t> compress_ratios(
                hparams.dsv4_compress_ratios.begin(), hparams.dsv4_compress_ratios.begin() + hparams.n_layer_all);
        add_kv(LLM_KV_ATTENTION_COMPRESS_RATIOS,         compress_ratios);
    } else {
        add_kv(LLM_KV_ATTENTION_COMPRESS_RATIOS,         hparams.dsv4_compress_ratios, true);
    }
    add_kv(LLM_KV_HYPER_CONNECTION_COUNT,               hparams.dsv4_hc_mult);
    add_kv(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, hparams.dsv4_hc_sinkhorn_iters);
    add_kv(LLM_KV_HYPER_CONNECTION_EPSILON,             hparams.dsv4_hc_eps);
    add_kv(LLM_KV_HYPER_CONNECTION_MAGNITUDE,           hparams.hc_magnitude);
    add_kv(LLM_KV_HASH_LAYER_COUNT,                     hparams.dsv4_hash_layer_count);
    add_kv(LLM_KV_HYPER_CONNECTION_LOW_RANK,             hparams.hc_low_rank);

    // the PLE group only means anything whole: write all of it or none
    if (hparams.ple_n_heads > 0) {
        std::vector<uint32_t> ple_layers;
        for (uint32_t il = 0; il < hparams.n_layer_all; ++il) {
            if (hparams.is_ple_impl[il]) {
                ple_layers.push_back(il);
            }
        }
        add_kv(LLM_KV_PLE_LAYERS,            ple_layers);
        add_kv(LLM_KV_PLE_NGRAM_SIZE,        hparams.ple_ngram_size);
        add_kv(LLM_KV_PLE_HEADS_PER_NGRAM,   hparams.ple_heads_per_ngram);
        add_kv(LLM_KV_PLE_CONV_KERNEL,       hparams.ple_conv_kernel);
        add_kv(LLM_KV_PLE_EOS_TOKEN_ID,      hparams.ple_eos_token_id);
        add_kv(LLM_KV_EMBEDDING_LENGTH_PER_LAYER, hparams.ple_head_dim);
        add_kv(LLM_KV_PLE_LAYER_MULTIPLIERS, std::vector<uint64_t>(
                hparams.ple_layer_multipliers.begin(),
                hparams.ple_layer_multipliers.begin() + hparams.ple_ngram_size));
        add_kv(LLM_KV_PLE_HEAD_OFFSETS,      std::vector<uint64_t>(
                hparams.ple_head_offsets.begin(),
                hparams.ple_head_offsets.begin() + hparams.ple_n_heads));
        add_kv(LLM_KV_PLE_HEAD_VOCAB_SIZES,  std::vector<uint64_t>(
                hparams.ple_head_vocab_sizes.begin(),
                hparams.ple_head_vocab_sizes.begin() + hparams.ple_n_heads));
    }

    const float rope_scaling_factor = hparams.rope_freq_scale_train == 1.0f ? 0.0f : 1.0f/hparams.rope_freq_scale_train;

    add_kv(LLM_KV_ROPE_DIMENSION_COUNT,              hparams.n_rot_full);
    add_kv(LLM_KV_ROPE_DIMENSION_COUNT_SWA,          hparams.n_rot_swa);
    add_kv(LLM_KV_ROPE_DIMENSION_SECTIONS,           hparams.rope_sections);
    add_kv(LLM_KV_ROPE_FREQ_BASE,                    hparams.rope_freq_base_train);
    add_kv(LLM_KV_ROPE_FREQ_BASE_SWA,                hparams.rope_freq_base_train_swa);
    // add_kv(LLM_KV_ROPE_SCALE_LINEAR,                 rope_scaling_factor); // old name
    add_kv(LLM_KV_ROPE_SCALING_TYPE,                 llama_rope_scaling_type_name(hparams.rope_scaling_type_train));
    add_kv(LLM_KV_ROPE_SCALING_FACTOR,               rope_scaling_factor);
    add_kv(LLM_KV_ROPE_SCALING_ATTN_FACTOR,          hparams.rope_attn_factor);
    add_kv(LLM_KV_ROPE_SCALING_ORIG_CTX_LEN,         hparams.n_ctx_orig_yarn);
    add_kv(LLM_KV_ROPE_SCALING_FINETUNED,            hparams.rope_finetuned);
    add_kv(LLM_KV_ROPE_SCALING_YARN_LOG_MUL,         hparams.rope_yarn_log_mul);
    add_kv(LLM_KV_ROPE_SCALING_YARN_EXT_FACTOR,      hparams.yarn_ext_factor);
    add_kv(LLM_KV_ROPE_SCALING_YARN_ATTN_FACTOR,     hparams.yarn_attn_factor);
    add_kv(LLM_KV_ROPE_SCALING_YARN_BETA_FAST,       hparams.yarn_beta_fast);
    add_kv(LLM_KV_ROPE_SCALING_YARN_BETA_SLOW,       hparams.yarn_beta_slow);

    // TODO: implement split file support
    // add_kv(LLM_KV_SPLIT_NO,                          ???);
    // add_kv(LLM_KV_SPLIT_COUNT,                       ???);
    // add_kv(LLM_KV_SPLIT_TENSORS_COUNT,               ???);

    add_kv(LLM_KV_SSM_INNER_SIZE,                    hparams.ssm_d_inner);
    add_kv(LLM_KV_SSM_CONV_KERNEL,                   hparams.ssm_d_conv);
    add_kv(LLM_KV_SSM_STATE_SIZE,                    hparams.ssm_d_state);
    add_kv(LLM_KV_SSM_TIME_STEP_RANK,                hparams.ssm_dt_rank);
    add_kv(LLM_KV_SSM_GROUP_COUNT,                   hparams.ssm_n_group);
    add_kv(LLM_KV_SSM_DT_B_C_RMS,                    hparams.ssm_dt_b_c_rms);

    add_kv(LLM_KV_KDA_HEAD_DIM,                      hparams.n_embd_head_kda);
    add_kv(LLM_KV_KDA_SAFE_GATE,                     hparams.kda_safe_gate);
    add_kv(LLM_KV_KDA_GATE_LOWER_BOUND,              hparams.kda_gate_lower_bound);

    add_kv(LLM_KV_WKV_HEAD_SIZE,                     hparams.wkv_head_size);

    add_kv(LLM_KV_TOKENIZER_MODEL,                   vocab.get_tokenizer_model());
    add_kv(LLM_KV_TOKENIZER_PRE,                     vocab.get_tokenizer_pre());
    add_kv(LLM_KV_TOKENIZER_LIST,                    tokens);
    add_kv(LLM_KV_TOKENIZER_TOKEN_TYPE,              token_types);
    add_kv(LLM_KV_TOKENIZER_TOKEN_TYPE_COUNT,        vocab.n_token_types());
    add_kv(LLM_KV_TOKENIZER_SCORES,                  scores);
    add_kv(LLM_KV_TOKENIZER_MERGES,                  vocab.get_bpe_merges());
    // FIXME llama_token is type i32 but when reading in a GGUF file u32 is expected, not an issue for writing though
    if (vocab.token_bos()  != LLAMA_TOKEN_NULL) { add_kv(LLM_KV_TOKENIZER_BOS_ID, uint32_t(vocab.token_bos()));  }
    if (vocab.token_eos()  != LLAMA_TOKEN_NULL) { add_kv(LLM_KV_TOKENIZER_EOS_ID, uint32_t(vocab.token_eos()));  }
    if (vocab.token_eot()  != LLAMA_TOKEN_NULL) { add_kv(LLM_KV_TOKENIZER_EOT_ID, uint32_t(vocab.token_eot()));  }
    if (vocab.token_eom()  != LLAMA_TOKEN_NULL) { add_kv(LLM_KV_TOKENIZER_EOM_ID, uint32_t(vocab.token_eom()));  }
    if (vocab.token_unk()  != LLAMA_TOKEN_NULL) { add_kv(LLM_KV_TOKENIZER_UNK_ID, uint32_t(vocab.token_unk()));  }
    if (vocab.token_sep()  != LLAMA_TOKEN_NULL) { add_kv(LLM_KV_TOKENIZER_SEP_ID, uint32_t(vocab.token_sep()));  }
    if (vocab.token_pad()  != LLAMA_TOKEN_NULL) { add_kv(LLM_KV_TOKENIZER_PAD_ID, uint32_t(vocab.token_pad()));  }
    // add_kv(LLM_KV_TOKENIZER_CLS_ID,                  uint32_t(vocab.token_bos())); // deprecated
    // add_kv(LLM_KV_TOKENIZER_MASK_ID,                 ???);
    add_kv(LLM_KV_TOKENIZER_ADD_BOS,                 vocab.get_add_bos());
    add_kv(LLM_KV_TOKENIZER_ADD_EOS,                 vocab.get_add_eos());
    add_kv(LLM_KV_TOKENIZER_ADD_SEP,                 vocab.get_add_sep());
    add_kv(LLM_KV_TOKENIZER_ADD_PREFIX,              vocab.get_add_space_prefix());
    add_kv(LLM_KV_TOKENIZER_REMOVE_EXTRA_WS,         vocab.get_remove_extra_whitespaces());
    add_kv(LLM_KV_TOKENIZER_PRECOMPILED_CHARSMAP,    vocab.get_precompiled_charsmap());
    // add_kv(LLM_KV_TOKENIZER_HF_JSON,                 ???);
    // add_kv(LLM_KV_TOKENIZER_RWKV,                    ???);
    if (vocab.token_fim_pre() != LLAMA_TOKEN_NULL) { add_kv(LLM_KV_TOKENIZER_FIM_PRE_ID, uint32_t(vocab.token_fim_pre())); }
    if (vocab.token_fim_suf() != LLAMA_TOKEN_NULL) { add_kv(LLM_KV_TOKENIZER_FIM_SUF_ID, uint32_t(vocab.token_fim_suf())); }
    if (vocab.token_fim_mid() != LLAMA_TOKEN_NULL) { add_kv(LLM_KV_TOKENIZER_FIM_MID_ID, uint32_t(vocab.token_fim_mid())); }
    if (vocab.token_fim_pad() != LLAMA_TOKEN_NULL) { add_kv(LLM_KV_TOKENIZER_FIM_PAD_ID, uint32_t(vocab.token_fim_pad())); }
    if (vocab.token_fim_rep() != LLAMA_TOKEN_NULL) { add_kv(LLM_KV_TOKENIZER_FIM_REP_ID, uint32_t(vocab.token_fim_rep())); }
    if (vocab.token_fim_sep() != LLAMA_TOKEN_NULL) { add_kv(LLM_KV_TOKENIZER_FIM_SEP_ID, uint32_t(vocab.token_fim_sep())); }

    // TODO: implement LoRA support
    // add_kv(LLM_KV_ADAPTER_TYPE,                      ???);
    // add_kv(LLM_KV_ADAPTER_LORA_ALPHA,                ???);
    // add_kv(LLM_KV_ADAPTER_LORA_TASK_NAME,            ???);
    // add_kv(LLM_KV_ADAPTER_LORA_PROMPT_PREFIX,        ???);
    // add_kv(LLM_KV_ADAPTER_ALORA_INVOCATION_TOKENS,   ???);

    add_kv(LLM_KV_POSNET_EMBEDDING_LENGTH,           hparams.posnet.n_embd);
    add_kv(LLM_KV_POSNET_BLOCK_COUNT,                hparams.posnet.n_layer);

    add_kv(LLM_KV_CONVNEXT_EMBEDDING_LENGTH,         hparams.convnext.n_embd);
    add_kv(LLM_KV_CONVNEXT_BLOCK_COUNT,              hparams.convnext.n_layer);

    add_kv(LLM_KV_CLASSIFIER_OUTPUT_LABELS,          model->classifier_labels);

    add_kv(LLM_KV_SHORTCONV_L_CACHE,                 hparams.n_shortconv_l_cache);

    add_kv(LLM_KV_XIELU_ALPHA_N,                     hparams.xielu_alpha_n);
    add_kv(LLM_KV_XIELU_ALPHA_P,                     hparams.xielu_alpha_p);
    add_kv(LLM_KV_XIELU_BETA,                        hparams.xielu_beta);
    add_kv(LLM_KV_XIELU_EPS,                         hparams.xielu_eps);

    add_kv(LLM_KV_ATTN_RES_BLOCK_SIZE,               hparams.attn_res_block_size);
    add_kv(LLM_KV_ACTIVATION_SITU_BETA,              hparams.situ_beta);
    add_kv(LLM_KV_ACTIVATION_SITU_LINEAR_BETA,       hparams.situ_linear_beta);

    // deprecated
    // add_kv(LLM_KV_TOKENIZER_PREFIX_ID,               ???);
    // add_kv(LLM_KV_TOKENIZER_SUFFIX_ID,               ???);
    // add_kv(LLM_KV_TOKENIZER_MIDDLE_ID,               ???);

    add_kv(LLM_KV_DENSE_2_FEAT_IN,                   hparams.dense_2_feat_in);
    add_kv(LLM_KV_DENSE_2_FEAT_OUT,                  hparams.dense_2_feat_out);
    add_kv(LLM_KV_DENSE_3_FEAT_IN,                   hparams.dense_3_feat_in);
    add_kv(LLM_KV_DENSE_3_FEAT_OUT,                  hparams.dense_3_feat_out);
}

void llama_model_saver::add_tensors_from_model() {
    if (model->output != nullptr &&
            std::string(model->output->name) != std::string(model->tok_embd->name)) {
        add_tensor(model->tok_embd); // some models use the same tensor for tok_embd and output
    }
    add_tensor(model->type_embd);
    add_tensor(model->pos_embd);
    add_tensor(model->tok_norm);
    add_tensor(model->tok_norm_b);
    add_tensor(model->output_norm);
    add_tensor(model->output_norm_b);
    add_tensor(model->output);
    add_tensor(model->output_b);
    add_tensor(model->output_norm_enc);
    add_tensor(model->output_s);
    add_tensor(model->output_in_s);
    add_tensor(model->output_res_score);
    add_tensor(model->nextn_proj_pre);
    add_tensor(model->nextn_proj_post);
    add_tensor(model->cls);
    add_tensor(model->cls_b);
    add_tensor(model->cls_out);
    add_tensor(model->cls_out_b);
    add_tensor(model->cls_norm);
    add_tensor(model->hrm_z_l_init);
    add_tensor(model->hc_head_fn);
    add_tensor(model->hc_head_base);
    add_tensor(model->hc_head_scale);
    add_tensor(model->per_layer_tok_embd);
    add_tensor(model->hc_head_norm);
    add_tensor(model->hc_head_down);
    add_tensor(model->hc_head_up);

    // looped architectures alias physical tensors across cache slots; save each
    // tensor once. a different tensor with an existing name still asserts below
    std::unordered_set<const struct ggml_tensor *> seen;

    for (const struct llama_layer & layer : model->layers) {
        for (size_t i = 0; i < sizeof(layer)/sizeof(struct ggml_tensor *); ++i) {
            const struct ggml_tensor * tensor = reinterpret_cast<const struct ggml_tensor * const *>(&layer)[i];
            if (tensor == nullptr || !seen.insert(tensor).second) {
                continue;
            }
            add_tensor(tensor);
        }
    }
}

void llama_model_saver::save(const std::string & path_model) {
    gguf_write_to_file(gguf_ctx, path_model.c_str(), false);
}

void llama_model_saver::save(FILE * file) {
    gguf_write_to_file_ptr(gguf_ctx, file, false);
}