include/neuroflow/alignment_common.hpp

#ifndef NEUROFLOW_ALIGNMENT_COMMON_HPP
#define NEUROFLOW_ALIGNMENT_COMMON_HPP

#include <algorithm>
#include <cmath>
#include <cstring>
#include <fstream>
#include <iostream>
#include <string>
#include <unordered_map>
#include <vector>

#include "causal_lm.hpp"
#include "tensor.hpp"

namespace neuroflow {

struct SFTSample {
    std::string instruction;
    std::string response;
};

struct DPOSample {
    std::string instruction;
    std::string chosen;
    std::string rejected;
};

struct SFTTrainingTensors {
    std::vector<size_t> input_ids;
    std::vector<size_t> target_ids;
    std::vector<float> loss_mask;
    size_t instruction_len;
};

struct DPOTrainingTensors {
    std::vector<size_t> chosen_ids;
    std::vector<size_t> rejected_ids;
    size_t chosen_prompt_len;
    size_t rejected_prompt_len;
};

inline std::string extract_json_string(const std::string& json, const std::string& key) {
    std::string search = "\"" + key + "\"";
    size_t pos = json.find(search);
    if (pos == std::string::npos) return "";
    pos = json.find(':', pos + search.size());
    if (pos == std::string::npos) return "";
    pos++;
    while (pos < json.size() && (json[pos] == ' ' || json[pos] == '\t')) pos++;
    if (pos >= json.size() || json[pos] != '"') return "";
    size_t end = pos + 1;
    while (end < json.size()) {
        if (json[end] == '"' && json[end - 1] != '\\') break;
        if (json[end - 1] == '\\') { end++; continue; }
        end++;
    }
    if (end >= json.size()) return "";
    return json.substr(pos + 1, end - pos - 1);
}

inline size_t extract_json_number(const std::string& json, const std::string& key, size_t default_val = 0) {
    std::string search = "\"" + key + "\"";
    size_t pos = json.find(search);
    if (pos == std::string::npos) return default_val;
    pos = json.find(':', pos + search.size());
    if (pos == std::string::npos) return default_val;
    pos++;
    while (pos < json.size() && (json[pos] == ' ' || json[pos] == '\t')) pos++;
    size_t end = pos;
    while (end < json.size() && (json[end] >= '0' && json[end] <= '9')) end++;
    if (end == pos) return default_val;
    return std::stoul(json.substr(pos, end - pos));
}

inline bool extract_json_bool(const std::string& json, const std::string& key, bool default_val = false) {
    std::string search = "\"" + key + "\"";
    size_t pos = json.find(search);
    if (pos == std::string::npos) return default_val;
    pos = json.find(':', pos + search.size());
    if (pos == std::string::npos) return default_val;
    pos++;
    while (pos < json.size() && (json[pos] == ' ' || json[pos] == '\t')) pos++;
    if (pos + 3 < json.size() && json.substr(pos, 4) == "true") return true;
    if (pos + 4 < json.size() && json.substr(pos, 5) == "false") return false;
    return default_val;
}

inline void unescape_json(std::string& s) {
    size_t pos = 0;
    while ((pos = s.find('\\', pos)) != std::string::npos) {
        if (pos + 1 < s.size()) {
            char c = s[pos + 1];
            if (c == 'n') { s.replace(pos, 2, "\n"); pos++; }
            else if (c == 't') { s.replace(pos, 2, "\t"); pos++; }
            else if (c == 'r') { s.replace(pos, 2, "\r"); pos++; }
            else if (c == '"') { s.replace(pos, 2, "\""); pos++; }
            else if (c == '\\') { s.replace(pos, 2, "\\"); pos++; }
            else if (c == '/') { s.replace(pos, 2, "/"); pos++; }
            else if (c == 'u' && pos + 5 < s.size()) { pos += 6; }
            else { pos += 2; }
        } else { pos++; }
    }
}

inline bool validate_path(const std::string& path) {
    if (path.find("..") != std::string::npos) return false;
    return true;
}

inline void save_lm_checkpoint(const CausalLMHead& model, const std::string& path) {
    CausalLMHead& lm_head = const_cast<CausalLMHead&>(model);
#ifdef USE_CUDA
    auto sync_to_cpu = [](const Tensor& t) {
        if (t.is_on_gpu()) { const_cast<Tensor&>(t).to_cpu(); }
    };
    sync_to_cpu(lm_head.w_embed_);
    sync_to_cpu(lm_head.w_pos_);
    sync_to_cpu(lm_head.dw_kernel_);
    sync_to_cpu(lm_head.pw_conv_->weight);
    sync_to_cpu(lm_head.sae_w_encode_->weight);
    sync_to_cpu(lm_head.sae_w_decode_->weight);
    sync_to_cpu(lm_head.ntm_w_read_->weight);
    sync_to_cpu(lm_head.ntm_w_write_->weight);
    sync_to_cpu(lm_head.ntm_w_erase_->weight);
    sync_to_cpu(lm_head.ntm_memory_);
    sync_to_cpu(lm_head.w_proj_->weight);
    sync_to_cpu(lm_head.w_proj_->bias);
    if (lm_head.bridge_) {
        sync_to_cpu(lm_head.bridge_->weight);
        sync_to_cpu(lm_head.bridge_->bias);
    }
    sync_to_cpu(lm_head.w_out_->weight);
    sync_to_cpu(lm_head.w_out_->bias);
    sync_to_cpu(lm_head.ln_->weight);
    sync_to_cpu(lm_head.ln_->bias);
    for (auto& attn : lm_head.attn_layers_) {
        sync_to_cpu(attn->w_q->weight);
        sync_to_cpu(attn->w_q->bias);
        sync_to_cpu(attn->w_k->weight);
        sync_to_cpu(attn->w_k->bias);
        sync_to_cpu(attn->w_v->weight);
        sync_to_cpu(attn->w_v->bias);
        sync_to_cpu(attn->w_out->weight);
        sync_to_cpu(attn->w_out->bias);
        sync_to_cpu(attn->norm->weight);
        sync_to_cpu(attn->norm->bias);
    }
#endif
    auto sl = [](std::ofstream& o, const std::string& n, const Tensor& t) {
        uint32_t nl = n.size(); o.write((char*)&nl, 4); o.write(n.data(), nl);
        uint32_t nd = t.shape_.size(); o.write((char*)&nd, 4);
        for (auto d : t.shape_) { uint32_t dd = d; o.write((char*)&dd, 4); }
        uint32_t ds = t.data_size_; o.write((char*)&ds, 4);
        o.write((char*)t.data_.get(), ds);
    };
    std::ofstream o(path, std::ios::binary);
    o.write("LMH2", 4);
    sl(o, "w_embed", lm_head.w_embed_);
    sl(o, "w_pos", lm_head.w_pos_);
    sl(o, "dw_kernel", lm_head.dw_kernel_);
    sl(o, "pw_conv.weight", lm_head.pw_conv_->weight);
    sl(o, "sae_encode.weight", lm_head.sae_w_encode_->weight);
    sl(o, "sae_decode.weight", lm_head.sae_w_decode_->weight);
    sl(o, "ntm_read.weight", lm_head.ntm_w_read_->weight);
    sl(o, "ntm_write.weight", lm_head.ntm_w_write_->weight);
    sl(o, "ntm_erase.weight", lm_head.ntm_w_erase_->weight);
    sl(o, "ntm_memory", lm_head.ntm_memory_);
    sl(o, "w_proj.weight", lm_head.w_proj_->weight);
    sl(o, "w_proj.bias", lm_head.w_proj_->bias);
    if (lm_head.bridge_) {
        sl(o, "bridge.weight", lm_head.bridge_->weight);
        sl(o, "bridge.bias", lm_head.bridge_->bias);
    }
    sl(o, "w_out.weight", lm_head.w_out_->weight);
    if (lm_head.w_out_->bias.data_) sl(o, "w_out.bias", lm_head.w_out_->bias);
    sl(o, "ln.weight", lm_head.ln_->weight);
    sl(o, "ln.bias", lm_head.ln_->bias);
    for (size_t i = 0; i < lm_head.attn_layers_.size(); ++i) {
        std::string p = "attn" + std::to_string(i) + ".";
        sl(o, p + "w_q.weight", lm_head.attn_layers_[i]->w_q->weight);
        sl(o, p + "w_q.bias", lm_head.attn_layers_[i]->w_q->bias);
        sl(o, p + "w_k.weight", lm_head.attn_layers_[i]->w_k->weight);
        sl(o, p + "w_k.bias", lm_head.attn_layers_[i]->w_k->bias);
        sl(o, p + "w_v.weight", lm_head.attn_layers_[i]->w_v->weight);
        sl(o, p + "w_v.bias", lm_head.attn_layers_[i]->w_v->bias);
        sl(o, p + "w_out.weight", lm_head.attn_layers_[i]->w_out->weight);
        sl(o, p + "w_out.bias", lm_head.attn_layers_[i]->w_out->bias);
        sl(o, p + "norm.weight", lm_head.attn_layers_[i]->norm->weight);
        sl(o, p + "norm.bias", lm_head.attn_layers_[i]->norm->bias);
    }
    uint32_t z = 0; o.write((char*)&z, 4); o.close();
}

inline bool load_lm_checkpoint(CausalLMHead& lm_head, const std::string& lm_path) {
    std::ifstream ifs(lm_path, std::ios::binary);
    if (!ifs) {
        std::cerr << "LM Head checkpoint未找到: " << lm_path << std::endl;
        return false;
    }

    char magic[5] = {0};
    ifs.read(magic, 4);
    if (std::string(magic) != "LMH2" && std::string(magic) != "LMH1") {
        std::cerr << "LM Head格式不匹配: " << std::string(magic) << " (期望LMH2/LMH1)" << std::endl;
        ifs.close();
        return false;
    }

    auto read_named_tensor = [&]() -> std::pair<std::string, Tensor> {
        uint32_t nl = 0; ifs.read((char*)&nl, 4);
        if (nl == 0 || ifs.eof()) return {"", Tensor()};
        std::string name(nl, '\0'); ifs.read(&name[0], nl);
        uint32_t nd = 0; ifs.read((char*)&nd, 4);
        std::vector<size_t> shape(nd);
        for (size_t i = 0; i < nd; ++i) { uint32_t d = 0; ifs.read((char*)&d, 4); shape[i] = d; }
        uint32_t ds = 0; ifs.read((char*)&ds, 4);
        Tensor t(shape, QuantType::FP32);
        if (t.data_size_ == ds) {
            ifs.read((char*)t.data_.get(), ds);
        } else {
            ifs.seekg(ds, std::ios::cur);
            t = Tensor();
        }
        return {name, std::move(t)};
    };

    std::unordered_map<std::string, Tensor*> tensor_map;
    tensor_map["w_embed"] = &lm_head.w_embed_;
    tensor_map["w_pos"] = &lm_head.w_pos_;
    tensor_map["dw_kernel"] = &lm_head.dw_kernel_;
    tensor_map["pw_conv.weight"] = &lm_head.pw_conv_->weight;
    tensor_map["sae_encode.weight"] = &lm_head.sae_w_encode_->weight;
    tensor_map["sae_decode.weight"] = &lm_head.sae_w_decode_->weight;
    tensor_map["ntm_read.weight"] = &lm_head.ntm_w_read_->weight;
    tensor_map["ntm_write.weight"] = &lm_head.ntm_w_write_->weight;
    tensor_map["ntm_erase.weight"] = &lm_head.ntm_w_erase_->weight;
    tensor_map["ntm_memory"] = &lm_head.ntm_memory_;
    tensor_map["w_proj.weight"] = &lm_head.w_proj_->weight;
    tensor_map["w_proj.bias"] = &lm_head.w_proj_->bias;
    if (lm_head.bridge_) {
        tensor_map["bridge.weight"] = &lm_head.bridge_->weight;
        tensor_map["bridge.bias"] = &lm_head.bridge_->bias;
    }
    tensor_map["w_out.weight"] = &lm_head.w_out_->weight;
    tensor_map["w_out.bias"] = &lm_head.w_out_->bias;
    tensor_map["ln.weight"] = &lm_head.ln_->weight;
    tensor_map["ln.bias"] = &lm_head.ln_->bias;
    for (size_t i = 0; i < lm_head.attn_layers_.size(); ++i) {
        std::string p = "attn" + std::to_string(i) + ".";
        tensor_map[p + "w_q.weight"] = &lm_head.attn_layers_[i]->w_q->weight;
        tensor_map[p + "w_q.bias"] = &lm_head.attn_layers_[i]->w_q->bias;
        tensor_map[p + "w_k.weight"] = &lm_head.attn_layers_[i]->w_k->weight;
        tensor_map[p + "w_k.bias"] = &lm_head.attn_layers_[i]->w_k->bias;
        tensor_map[p + "w_v.weight"] = &lm_head.attn_layers_[i]->w_v->weight;
        tensor_map[p + "w_v.bias"] = &lm_head.attn_layers_[i]->w_v->bias;
        tensor_map[p + "w_out.weight"] = &lm_head.attn_layers_[i]->w_out->weight;
        tensor_map[p + "w_out.bias"] = &lm_head.attn_layers_[i]->w_out->bias;
        tensor_map[p + "norm.weight"] = &lm_head.attn_layers_[i]->norm->weight;
        tensor_map[p + "norm.bias"] = &lm_head.attn_layers_[i]->norm->bias;
    }

    size_t loaded = 0;
    while (ifs) {
        auto [name, tensor] = read_named_tensor();
        if (name.empty()) break;
        auto it = tensor_map.find(name);
        if (it != tensor_map.end() && tensor.numel() > 0) {
            if (it->second->shape_ == tensor.shape_) {
                memcpy(it->second->data_.get(), tensor.data_.get(), tensor.data_size_);
                loaded++;
            } else {
                std::cerr << "  跳过 '" << name << "': 形状不匹配" << std::endl;
            }
        } else if (name.find("w_qkv.") != std::string::npos) {
            std::string base = name.substr(0, name.find("w_qkv."));
            std::string suffix = name.substr(name.find("w_qkv.") + 6);
            for (size_t i = 0; i < lm_head.attn_layers_.size(); ++i) {
                std::string p = "attn" + std::to_string(i) + ".";
                if (base != p) continue;
                size_t d_model = lm_head.config_.d_model;
                size_t head_dim = d_model / lm_head.config_.num_attn_heads;
                size_t n_q = lm_head.attn_layers_[i]->n_q_heads_;
                size_t n_kv = lm_head.attn_layers_[i]->n_kv_heads_;
                if (suffix == "weight" && tensor.shape_.size() == 2 && tensor.shape_[0] == 3 * d_model) {
                    const float* src = tensor.as_fp32();
                    float* dq = lm_head.attn_layers_[i]->w_q->weight.as_fp32();
                    float* dk = lm_head.attn_layers_[i]->w_k->weight.as_fp32();
                    float* dv = lm_head.attn_layers_[i]->w_v->weight.as_fp32();
                    for (size_t r = 0; r < d_model; ++r) {
                        memcpy(dq + r * n_q * head_dim, src + r * 3 * d_model, n_q * head_dim * sizeof(float));
                        memcpy(dk + r * n_kv * head_dim, src + r * 3 * d_model + d_model, n_kv * head_dim * sizeof(float));
                        memcpy(dv + r * n_kv * head_dim, src + r * 3 * d_model + 2 * d_model, n_kv * head_dim * sizeof(float));
                    }
                    loaded++;
                    std::cerr << "  拆分旧格式 '" << name << "' -> w_q/w_k/w_v" << std::endl;
                } else if (suffix == "bias" && tensor.shape_[0] == 3 * d_model) {
                    const float* src = tensor.as_fp32();
                    float* bq = lm_head.attn_layers_[i]->w_q->bias.as_fp32();
                    float* bk = lm_head.attn_layers_[i]->w_k->bias.as_fp32();
                    float* bv = lm_head.attn_layers_[i]->w_v->bias.as_fp32();
                    memcpy(bq, src, n_q * head_dim * sizeof(float));
                    memcpy(bk, src + d_model, n_kv * head_dim * sizeof(float));
                    memcpy(bv, src + 2 * d_model, n_kv * head_dim * sizeof(float));
                    loaded++;
                    std::cerr << "  拆分旧格式 '" << name << "' -> w_q/w_k/w_v bias" << std::endl;
                }
            }
        }
    }

    ifs.close();
    if (lm_head.config_.weight_tying) lm_head.tie_weights();
    std::cerr << "LM Head已恢复: " << loaded << " 个张量从 " << lm_path << std::endl;
    return loaded > 0;
}

} // namespace neuroflow

#endif // NEUROFLOW_ALIGNMENT_COMMON_HPP