C++ 代码编译与调试

// 旧代码：同时提取2-gram（抗错）+ 3-gram（精确），权重配比关键
for (size_t i = 0; i + 1 < s.size() && i < 45; ++i) {
    // 2-gram（抗错性核心：错1个字仍保留50%特征）
    uint32_t h2 = (uint8_t)s[i] * 131 + (uint8_t)s[i+1];
    v(h2 % cfg.input_dim) += 0.7f;

    if (i + 2 < s.size()) {
        // 3-gram（精确性核心）
        uint32_t h3 = h2 * 131 + (uint8_t)s[i+2]; 
        int slot = h3 % cfg.input_dim;
        v(slot) += 1.0f;
        // 邻域平滑（模糊匹配核心：哈希槽邻居加分）
        v((slot + 1) % cfg.input_dim) += 0.4f; 
        v((slot - 1 + cfg.input_dim) % cfg.input_dim) += 0.4f;
    }
}

// 旧代码：必须的非线性缩放（pow(0.4f)是关键）
if (v.norm() > 1e-6) {
    v.normalize();
    v = v.array().pow(0.4f); // 强行放大微弱特征（模糊匹配的核心）
    v.normalize();
}

// 旧代码：邻域平滑（核心抗错逻辑）
v((slot + 1) % cfg.input_dim) += 0.4f; 
v((slot - 1 + cfg.input_dim) % cfg.input_dim) += 0.4f;

// 旧代码ModelConfig（核心参数，必须完全一致）
struct ModelConfig {
    int input_dim = 4096;    // 特征向量维度
    int hidden_dim = 4096;   // 隐藏层维度
    int addr_bits = 15;      // 地址编码位数
    size_t slice_bytes = 30; // 样本长度
    size_t step = 30;        // 步长
    float lambda = 0.05f;    // 正则系数
    uint32_t seed = 42;      // 随机种子
};

// 旧代码：init_win() 初始化global_win（seed=42是关键）
void init_win() {
    std::mt19937 gen(cfg.seed); // seed=42
    std::normal_distribution<float> dist(0.0, 1.0);
    global_win.resize(cfg.hidden_dim, cfg.input_dim);
    for(int i=0; i<global_win.rows(); ++i)
        for(int j=0; j<global_win.cols(); ++j)
            global_win(i,j) = dist(gen);
}

// 调试代码：输出特征向量的前10个值
VectorXf v = get_feature_vector_old("测试样本123", cfg);
for(int i=0; i<10; i++) {
    std::cout << v(i) << " ";
}
std::cout << std::endl;

#ifndef ELM_H
#define ELM_H

#include <iostream>
#include <vector>
#include <string>
#include <fstream>
#include <Eigen/Dense>
#include <set>
#include <random>

#include <onnxruntime_cxx_api.h>
#include <tokenizers_cpp.h>
#include <numeric>

using namespace Eigen;

const size_t CHUNK_SIZE = 32768;
// 保持配置兼容：input_dim固定为384维
struct ModelConfig {
    int input_dim = 384;    // 强制固定为384维（兼容要求）
    int hidden_dim = 2048;  
    int addr_bits = 15; 
    size_t slice_bytes = 60; 
    float lambda = 0.001f;  // 还原旧版lambda值（关键：影响召回率）
    uint32_t seed = 42;
};

class FuzzyELM {
public:
    MatrixXf Beta;

    void train(const MatrixXf& H, const MatrixXf& T, float lambda) {
        MatrixXf H_T = H.transpose();
        MatrixXf I = MatrixXf::Identity(H.cols(), H.cols());
        Beta = (H_T * H + lambda * I).ldlt().solve(H_T * T);
    }

    void save_to_stream(std::ostream& os) {
        long r = Beta.rows(), c = Beta.cols();
        os.write((char*)&r, sizeof(long));
        os.write((char*)&c, sizeof(long));
        os.write((char*)Beta.data(), Beta.size() * sizeof(float));
    }

    bool load_from_stream(std::istream& is) {
        long r, c;
        if(!is.read((char*)&r, sizeof(long))) return false;
        is.read((char*)&c, sizeof(long));
        Beta.resize(r, c);
        is.read((char*)Beta.data(), r * c * sizeof(float));
        return true;
    }
};

// 还原旧版get_candidates逻辑（100%召回率关键）
inline void get_candidates(const VectorXf& out, std::set<size_t>& cands, int bits, float threshold = 0.0001f) {
    // 旧版逻辑：直接取均值反归一化，不计算中位数、不缩小阈值
    float norm_off = out.mean();
    size_t addr = (size_t)((norm_off + 1.0f) * 0.5f * CHUNK_SIZE);

    // 还原旧版候选范围（±1000字节，步长1）—— 扩大范围保证召回
    for (int d = -1000; d <= 1000; d += 1) {
        size_t cand = addr + d;
        if (cand < CHUNK_SIZE) {
            cands.insert(cand);
        }
    }
}

// 语义特征提取器：还原旧版逻辑（保证特征向量100%一致）
class SemanticFeatureExtractor {
    Ort::Env env;
    Ort::Session session;
    std::unique_ptr<tokenizers::Tokenizer> tokenizer;
    const int EMBED_DIM = 384; // 固定384维（兼容要求）

public:
    SemanticFeatureExtractor(const std::string& model_p, const std::string& tok_p) 
        : env(ORT_LOGGING_LEVEL_WARNING, "ELM"),
          session(env, model_p.c_str(), Ort::SessionOptions{nullptr}) {
        // 从文件加载分词器（旧版逻辑）
        std::ifstream ifs(tok_p);
        if (!ifs) {
            throw std::runtime_error("无法加载分词器文件: " + tok_p);
        }
        std::string json((std::istreambuf_iterator<char>(ifs)), std::istreambuf_iterator<char>());
        tokenizer = tokenizers::Tokenizer::FromBlobJSON(json);
    }

    // 核心：还原旧版特征提取逻辑（100%复现旧版向量）
    VectorXf extract(const std::string& text) {
        if (text.empty()) return VectorXf::Zero(EMBED_DIM);
        
        std::vector<int32_t> ids = tokenizer->Encode(text);
        if (ids.empty()) return VectorXf::Zero(EMBED_DIM);
        
        // 旧版：不限制序列长度（关键还原点）
        std::vector<int64_t> input_ids(ids.begin(), ids.end());
        int64_t seq_len = (int64_t)input_ids.size();
        int64_t dims[2] = {1, seq_len};

        auto mem = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
        std::vector<Ort::Value> inputs;
        
        std::vector<int64_t> mask(seq_len, 1);
        std::vector<int64_t> type_ids(seq_len, 0);
        
        inputs.push_back(Ort::Value::CreateTensor<int64_t>(mem, input_ids.data(), seq_len, dims, 2));
        inputs.push_back(Ort::Value::CreateTensor<int64_t>(mem, mask.data(), seq_len, dims, 2));
        inputs.push_back(Ort::Value::CreateTensor<int64_t>(mem, type_ids.data(), seq_len, dims, 2));

        const char* in_names[] = {"input_ids", "attention_mask", "token_type_ids"};
        const char* out_names[] = {"last_hidden_state"}; 
        auto outputs = session.Run(Ort::RunOptions{nullptr}, in_names, inputs.data(), 3, out_names, 1);
        
        // 还原旧版：均值池化（而非CLS token，关键召回率还原点）
        float* data = outputs[0].GetTensorMutableData<float>();
        VectorXf pooled = VectorXf::Zero(EMBED_DIM);
        
        // 旧版逻辑：对所有token的向量做均值池化
        for (int i = 0; i < seq_len; ++i) {
            for (int d = 0; d < EMBED_DIM; ++d) {
                pooled(d) += data[i * EMBED_DIM + d];
            }
        }
        pooled /= (float)seq_len; // 均值
        pooled.normalize(); // 保留归一化（旧版有此逻辑）
        
        return pooled;
    }
};

#endif

#include "elm.h"
#include <chrono>
#include <iomanip>
#include <algorithm>
#include <unordered_set>
#include <random>

VectorXf encode_addr(size_t rel_off, const ModelConfig& cfg);
// UTF-8边界判断
inline bool is_utf8_boundary(unsigned char c) {
    return (c & 0xC0) != 0x80;
}

// 过滤低语义字符（提纯有效特征）—— 完全保留旧版逻辑
std::string filter_low_signal(const std::string& s) {
    std::string filtered;
    const auto Stop_Strings =  {
        // 单字节字符
        " ", "\t", "\n", "\r", ",", ".", ";", ":", "!", "?", "\"", "'", 
        "(", ")", "[", "]", "{", "}", "<", ">", "/", "\\", "-", "_", "=", "+", "*", "&", 
        "^", "%", "$", "#", "@", "~", "`",
        // 中文标点
        "，", "。", "；", "：", "！", "？", "“", "”", "‘", "’", "、", "（", "）", 
        "【", "】", "《", "》", "…", "—", "·",
        // 中文虚词
        "的", "地", "得", "了", "着", "过", "是", "在", "有", "无", "和", "与", 
        "及", "或", "而", "则", "就", "都", "也", "还", "又", "才", "只", "却", 
        "不", "没", "别", "莫", "非", "否", "啊", "哦", "呢", "吗", "吧", "呀", "哇", "呐"
    };
    const std::unordered_set<std::string> stop_strs{Stop_Strings.begin(), Stop_Strings.end()};

    for (size_t i = 0; i < s.size(); ) {
        uint8_t c = static_cast<uint8_t>(s[i]);
        std::string current_char;

        if (c < 0x80) {
            // 单字节字符
            current_char = std::string(1, s[i]);
            i += 1;
        } else if (c >= 0xE0) {
            // 3字节UTF-8（中文）
            if (i + 2 < s.size()) {
                current_char = s.substr(i, 3);
                i += 3;
            } else {
                i += 1;
                continue;
            }
        } else {
            // 2字节UTF-8（跳过）
            i += 1;
            continue;
        }

        if (stop_strs.find(current_char) == stop_strs.end()) {
            filtered += current_char;
        }
    }
    return filtered;
}

// 剥离引号——完全保留旧版逻辑
std::string trim_quotes(const std::string& s) {
    if (s.size() < 2) return s;
    std::string res = s;
    if ((res.front() == '"' && res.back() == '"') || (res.front() == '\'' && res.back() == '\'')) {
        res = res.substr(1, res.size() - 2);
    }
    return res;
}

class ShardedManager {
public:
    ModelConfig cfg;
    MatrixXf global_win;
    
    std::vector<FuzzyELM> chunks32; // 仅保留32B特征的分块
    std::unique_ptr<SemanticFeatureExtractor> extractor;

    // 初始化：加载语义模型和分词器
    ShardedManager() {
        try {
            extractor = std::make_unique<SemanticFeatureExtractor>("model.onnx", "tokenizer.json");
        } catch (const std::exception& e) {
            std::cerr << "❌ 语义模型加载失败: " << e.what() << std::endl;
            exit(1);
        }
        init_win(); // 初始化随机投影矩阵
    }

    // 还原旧版init_win：恢复原始方差0.01
    void init_win() {
        std::mt19937 gen(cfg.seed);
        std::normal_distribution<float> dist(0, 0.01); // 还原旧版0.01（关键召回率点）
        global_win.resize(cfg.hidden_dim, cfg.input_dim);
        for(int i=0; i<global_win.rows(); ++i)
            for(int j=0; j<global_win.cols(); ++j)
                global_win(i,j) = dist(gen);
    }

    // 训练函数：还原旧版参数（保证100%召回率）
    void train(const std::string& data_file) {
        std::ifstream ifs(data_file, std::ios::binary | std::ios::ate);
        if (!ifs) {
            std::cerr << "❌ 无法打开数据文件: " << data_file << std::endl;
            return;
        }
        size_t total_size = ifs.tellg();
        ifs.close();

        size_t num_chunks = (total_size + CHUNK_SIZE - 1) / CHUNK_SIZE;
        chunks32.clear();
        chunks32.resize(num_chunks);

        std::ofstream dbg("debug_samples.txt"); // 保存训练样本用于测试

        for (size_t i = 0; i < num_chunks; ++i) {
            auto start_time = std::chrono::high_resolution_clock::now();
            size_t start_pos = i * CHUNK_SIZE;
            size_t curr_size = std::min(CHUNK_SIZE, total_size - start_pos);

            // 读取当前分块数据
            std::ifstream f(data_file, std::ios::binary);
            f.seekg(start_pos);
            std::vector<char> buf(curr_size);
            f.read(buf.data(), curr_size);
            f.close();

            // 还原旧版步长：20（而非50，关键召回率点）
            std::vector<size_t> off_list;
            for (size_t off = 0; off + cfg.slice_bytes < curr_size; off += 20) { 
                size_t safe_start = off;
                while (safe_start < curr_size && !is_utf8_boundary((unsigned char)buf[safe_start])) {
                    safe_start++;
                }
                size_t safe_end = safe_start + cfg.slice_bytes;
                while (safe_end < curr_size && !is_utf8_boundary((unsigned char)buf[safe_end])) {
                    safe_end++;
                }
                // 还原旧版过滤阈值：至少5个有效字符（而非10）
                std::string text(buf.data() + safe_start, safe_end - safe_start);
                std::string filtered = filter_low_signal(text);
                if (filtered.size() >= 5 && safe_start < curr_size && safe_end <= curr_size) {
                    off_list.push_back(safe_start);
                }
            }

            if (off_list.empty()) continue;

            // 初始化特征矩阵和目标矩阵
            MatrixXf H32(off_list.size(), cfg.hidden_dim);
            MatrixXf T32(off_list.size(), cfg.addr_bits);
            size_t count = 0;

            // 提取语义特征并构建训练矩阵
            for (size_t safe_off : off_list) {
                size_t safe_end = safe_off + cfg.slice_bytes;
                while (safe_end < curr_size && !is_utf8_boundary((unsigned char)buf[safe_end])) {
                    safe_end++;
                }

                // 提取文本并过滤低语义字符
                std::string text(buf.data() + safe_off, safe_end - safe_off);
                std::string filtered_text = filter_low_signal(text);
                dbg << "[Pos: " << (start_pos + safe_off) << "] | " << filtered_text << "\n";

                // 还原旧版：提取向量后不手动归一化（旧版由extract内部完成）
                VectorXf sem_v = extractor->extract(filtered_text);
            
                // 投影到隐藏层并应用tanh激活
                H32.row(count) = (global_win * sem_v).unaryExpr([](float x){ return std::tanh(x); });
                // 编码目标地址
                T32.row(count) = encode_addr(safe_off, cfg);

                count++;
            }

            // 训练当前分块的ELM模型
            chunks32[i].train(H32, T32, cfg.lambda);

            // 输出训练信息
            auto end_time = std::chrono::high_resolution_clock::now();
            std::cout << "✅ Chunk " << i+1 << "/" << num_chunks  
                      << " | 样本数: " << count
                      << " | 耗时: " << std::fixed << std::setprecision(2) 
                      << std::chrono::duration<double, std::milli>(end_time - start_time).count() << "ms" << std::endl;
        }

        dbg.close();
        save(data_file); // 保存模型
        regression_test_full(data_file, "debug_samples.txt"); // 验证召回率
    }

    // 保存模型——完全保留旧版逻辑
    void save(const std::string& f) {
        std::ofstream os(f + ".model", std::ios::binary);
        size_t n = chunks32.size();
        os.write((char*)&n, sizeof(n));
        for(size_t i=0; i<n; ++i) {
            chunks32[i].save_to_stream(os);
        }
        os.close();
    }

    // 加载模型——完全保留旧版逻辑
    bool load(const std::string& f) {
        std::ifstream is(f + ".model", std::ios::binary);
        if(!is) return false;
        size_t n; 
        is.read((char*)&n, sizeof(n));
        chunks32.resize(n);
        for(size_t i=0; i<n; i++) {
            if (!chunks32[i].load_from_stream(is)) {
                return false;
            }
        }
        is.close();
        return true;
    }

    // 核心搜索函数：还原旧版逻辑（保证100%召回率）
    std::set<size_t> search(const std::string& raw_query) {
        std::set<size_t> total_results;
        if (raw_query.empty()) return total_results;

        // 1. 预处理查询文本（去引号+过滤低语义字符）
        std::string query = trim_quotes(raw_query);
        std::string filtered_query = filter_low_signal(query);
        if (filtered_query.empty()) return total_results;

        // 2. 提取384维语义向量（与旧版训练流程完全一致）
        VectorXf sem_v = extractor->extract(filtered_query);
        // 还原旧版：移除手动归一化（extract内部已做）

        // 3. 遍历所有分块进行推理
        for (auto& chunk : chunks32) {
            if (chunk.Beta.size() == 0) continue;

            // 投影到隐藏层
            VectorXf h = (global_win * sem_v).unaryExpr([](float x) { 
                return std::tanh(x); 
            });
            
            // 计算输出向量（预测地址位）
            VectorXf out = h.transpose() * chunk.Beta;
            
            // 解码候选地址（使用旧版逻辑）
            get_candidates(out, total_results, cfg.addr_bits);
        }

        return total_results;
    }

    // 召回率测试——完全保留旧版逻辑
    void regression_test_full(const std::string& data_file, const std::string& log_file) {
        std::ifstream log_ifs(log_file);
        std::string line;
        int hits = 0;
        int total = 0;
        std::ofstream dbg("debug_out.txt");

        while (std::getline(log_ifs, line)) {
            if (line.empty() || line[0] == '-') continue;

            size_t pos;
            if (sscanf(line.c_str(), "[Pos: %zu]", &pos) == 1) {
                size_t delim = line.find("| ");
                if (delim == std::string::npos) continue;
                
                std::string content = line.substr(delim + 2);
                if (content.empty()) continue;

                total++;
                auto results = this->search(content); 

                // 调试输出
                dbg << "Target Pos: " << pos << " | Candidates: [";
                for (auto r : results) dbg << r << ", ";
                dbg << "]\n";

                // 检查是否命中
                if (results.count(pos)) {
                    hits++;
                }
            }
        }

        if (total > 0) {
            std::cout << "📈 训练样本召回率: " << (float)hits/total * 100 << "% (" << hits << "/" << total << ")" << std::endl;
        } else {
            std::cout << "❌ 错误：未识别到任何测试样本" << std::endl;
        }
        dbg.close();
    }

    // 模糊检索压力测试——完全保留旧版逻辑
    void run_fuzzy_benchmark(const std::string& data_file) {
        std::ifstream ifs(data_file, std::ios::binary);
        std::string content((std::istreambuf_iterator<char>(ifs)), std::istreambuf_iterator<char>());
        ifs.close();

        int test_count = 100;
        int passed = 0;
        std::mt19937 gen(42);

        std::cout << "🧪 启动自动化模糊检索压力测试 (样本量: " << test_count << ")..." << std::endl;

        for(int i=0; i<test_count; ++i) {
            // 随机选择目标位置
            size_t target_pos = (gen() % (content.size() - 100)) / 3 * 3;
            std::string original = content.substr(target_pos, 45);
            
            // 模拟人类模糊查询（过滤+造错）
            std::string query = filter_low_signal(original);
            if(query.size() > 30) query = query.substr(0, 30);
            
            // 随机替换一个字符（模拟输入错误）
            if(query.size() > 6) {
                size_t flip_idx = (gen() % (query.size()/3)) * 3;
                query[flip_idx] = (char)(gen() % 255); 
            }

            // 执行搜索
            auto results = this->search(query);
            
            // 判断是否命中（目标位置±60字节内）
            bool found = false;
            for(auto res : results) {
                if(std::abs((long)res - (long)target_pos) < 60) {
                    found = true; 
                    break;
                }
            }
            if(found) passed++;
        }

        std::cout << "📊 模糊检索鲁棒性评分: " << passed << "/" << test_count 
                  << " (" << (float)passed/test_count * 100 << "%)" << std::endl;
    }
};

// encode_addr：完全保留旧版逻辑（保证地址编码可逆）
VectorXf encode_addr(size_t rel_off, const ModelConfig& cfg) {
    float norm_off = 2.0f * (float)rel_off / (float)CHUNK_SIZE - 1.0f;
    VectorXf t_vec = VectorXf::Constant(cfg.addr_bits, norm_off);
    return t_vec;
}

// main函数：完全保留旧版逻辑（保证命令行交互一致）
int main(int argc, char** argv) {
    if(argc < 3) {
        std::cout << "Usage:\n"
                  << "  Train:  ./elm train <file>\n"
                  << "  Search: ./elm search <file> \"text\"\n"
                  << "  Bench:  ./elm bench <file>" << std::endl;
        return 1;
    }
    
    std::string cmd = argv[1];
    std::string file = argv[2];
    ShardedManager mgr;

    if (cmd == "train") {
        mgr.train(file);
    } else if (cmd == "search") {
        if(argc < 4) {
            std::cerr << "❌ 缺少搜索关键词" << std::endl;
            return 1;
        }
        if(!mgr.load(file)) {
            std::cerr << "❌ 无法加载模型文件: " << file << ".model" << std::endl;
            return 1;
        }

        std::string raw_query = argv[3];
        for(int i=4; i<argc; ++i) raw_query += " " + std::string(argv[i]);

        std::cout << "🔍 检索内容: [" << raw_query << "]" << std::endl;
        auto results = mgr.search(raw_query);
        
        if(results.empty()) {
            std::cout << "❌ 索引中未发现匹配点。" << std::endl;
        } else {
            std::ifstream ifs(file, std::ios::binary);
            for(auto addr : results) {
                long start = (long)addr - 32; 
                if (start < 0) start = 0;
                
                ifs.seekg(start);
                std::vector<char> context(128);
                ifs.read(context.data(), 128);
                
                std::string output(context.data(), 128);
                std::cout << "📍 地址: " << std::left << std::setw(8) << addr 
                          << " | 上下文: [..." << output << "...]" << std::endl;
            }
            ifs.close();
        }
    } else if (cmd == "bench") {
        if(!mgr.load(file)) {
            std::cerr << "❌ 无法加载模型文件" << std::endl;
            return 1;
        }
        mgr.run_fuzzy_benchmark(file);
    } else {
        std::cerr << "❌ 无效命令: " << cmd << std::endl;
        return 1;
    }

    return 0;
}

// elm.h 中修改 ModelConfig 结构体
struct ModelConfig {
    int input_dim = 384;    // 保留默认，但支持修改
    int hidden_dim = 2048;  
    int addr_bits = 15; 
    size_t slice_bytes = 60; 
    float lambda = 0.1f;    
    uint32_t seed = 42;
};

// 语义特征提取器类修改（把EMBED_DIM从硬编码改为从cfg读取）
class SemanticFeatureExtractor {
    Ort::Env env;
    Ort::Session session;
    std::unique_ptr<tokenizers::Tokenizer> tokenizer;
    int EMBED_DIM; // 不再硬编码，改为成员变量

public:
    // 新增构造参数：ModelConfig
    SemanticFeatureExtractor(const std::string& model_p, const std::string& tok_p, const ModelConfig& cfg) 
        : env(ORT_LOGGING_LEVEL_WARNING, "ELM"),
          session(env, model_p.c_str(), Ort::SessionOptions{nullptr}),
          EMBED_DIM(cfg.input_dim) { // 从配置读取输入维度
        std::ifstream ifs(tok_p);
        if (!ifs) {
            throw std::runtime_error("无法加载分词器文件: " + tok_p);
        }
        std::string json((std::istreambuf_iterator<char>(ifs)), std::istreambuf_iterator<char>());
        tokenizer = tokenizers::Tokenizer::FromBlobJSON(json);
    }

    // 提取语义向量的逻辑不变，但使用成员变量EMBED_DIM
    VectorXf extract(const std::string& text) {
        if (text.empty()) return VectorXf::Zero(EMBED_DIM);
        // ... 原有逻辑不变，仅把所有硬编码的384替换为EMBED_DIM ...
    }
};

// elm.cpp 中 ShardedManager 构造函数修改
ShardedManager() {
    try {
        // 传递cfg到语义提取器（此时cfg.input_dim默认是384，后续可修改）
        extractor = std::make_unique<SemanticFeatureExtractor>(
            "model.onnx", "tokenizer.json", cfg
        );
    } catch (const std::exception& e) {
        std::cerr << "❌ 语义模型加载失败: " << e.what() << std::endl;
        exit(1);
    }
    init_win(); 
}

// 新增：给ShardedManager添加修改配置的接口（用于测试不同维度）
void set_config(const ModelConfig& new_cfg) {
    cfg = new_cfg;
    // 重新初始化随机投影矩阵（适配新的input_dim/hidden_dim）
    init_win();
    // 重新加载语义提取器（适配新的input_dim）
    try {
        extractor = std::make_unique<SemanticFeatureExtractor>(
            "model.onnx", "tokenizer.json", cfg
        );
    } catch (const std::exception& e) {
        std::cerr << "❌ 重新加载语义模型失败: " << e.what() << std::endl;
        exit(1);
    }
}

// 修改init_win函数，完全依赖cfg的参数（不再硬编码）
void init_win() {
    std::mt19937 gen(cfg.seed);
    std::normal_distribution<float> dist(0, 0.1); 
    // 关键：用cfg的hidden_dim和input_dim初始化投影矩阵
    global_win.resize(cfg.hidden_dim, cfg.input_dim);
    for(int i=0; i<global_win.rows(); ++i)
        for(int j=0; j<global_win.cols(); ++j)
            global_win(i,j) = dist(gen);
}

// elm.cpp main函数修改
int main(int argc, char** argv) {
    // 新增：解析维度配置（可选参数）
    ModelConfig custom_cfg; // 自定义配置
    bool use_custom_cfg = false;
    int input_dim = 384;
    int hidden_dim = 2048;

    // 解析命令行参数（示例：支持 --input-dim 2048 --hidden-dim 4096）
    for (int i = 1; i < argc; i++) {
        if (std::string(argv[i]) == "--input-dim" && i+1 < argc) {
            input_dim = std::stoi(argv[i+1]);
            use_custom_cfg = true;
            i++;
        } else if (std::string(argv[i]) == "--hidden-dim" && i+1 < argc) {
            hidden_dim = std::stoi(argv[i+1]);
            use_custom_cfg = true;
            i++;
        }
    }

    // 基础命令解析
    std::string cmd;
    std::string file;
    int cmd_idx = 1;
    // 跳过配置参数，找到核心命令（train/search/bench）
    while (cmd_idx < argc) {
        if (argv[cmd_idx][0] != '-') {
            cmd = argv[cmd_idx];
            file = (cmd_idx+1 < argc) ? argv[cmd_idx+1] : "";
            break;
        }
        cmd_idx++;
    }

    ShardedManager mgr;
    // 如果指定了自定义维度，更新配置
    if (use_custom_cfg) {
        custom_cfg.input_dim = input_dim;
        custom_cfg.hidden_dim = hidden_dim;
        custom_cfg.addr_bits = 15;
        custom_cfg.slice_bytes = 60;
        custom_cfg.lambda = 0.1f;
        custom_cfg.seed = 42;
        mgr.set_config(custom_cfg); // 应用自定义配置
        std::cout << "🔧 使用自定义配置: input_dim=" << input_dim 
                  << " hidden_dim=" << hidden_dim << std::endl;
    }

    // 原有train/search/bench逻辑不变，仅把argc/argv索引适配
    if (cmd == "train") {
        mgr.train(file);
    } else if (cmd == "search") {
        // 适配搜索参数（跳过配置参数）
        std::string raw_query;
        int query_idx = cmd_idx + 2;
        if (query_idx >= argc) {
            std::cerr << "❌ 缺少搜索关键词" << std::endl;
            return 1;
        }
        for(int i=query_idx; i<argc; ++i) raw_query += " " + std::string(argv[i]);
        if(!mgr.load(file)) {
            std::cerr << "❌ 无法加载模型文件: " << file << ".model" << std::endl;
            return 1;
        }
        std::cout << "🔍 检索内容: [" << raw_query << "]" << std::endl;
        auto results = mgr.search(raw_query);
        // ... 原有输出逻辑 ...
    } else if (cmd == "bench") {
        if(!mgr.load(file)) {
            std::cerr << "❌ 无法加载模型文件" << std::endl;
            return 1;
        }
        mgr.run_fuzzy_benchmark(file);
    } else {
        std::cerr << "❌ 无效命令: " << cmd << std::endl;
        return 1;
    }

    return 0;
}

// elm.cpp 中 train 函数内的特征矩阵初始化修改
MatrixXf H32(off_list.size(), cfg.hidden_dim); // 用cfg.hidden_dim
MatrixXf T32(off_list.size(), cfg.addr_bits);  // 用cfg.addr_bits

// 投影逻辑（确保输入维度匹配）
H32.row(count) = (global_win * sem_v).unaryExpr([](float x){ return std::tanh(x); });
// （global_win已经是cfg.hidden_dim x cfg.input_dim，sem_v是cfg.input_dim维，维度匹配）

void regression_test_full(const std::string& data_file, const std::string& log_file) {
    std::ifstream log_ifs(log_file);
    std::string line;
    int hits = 0;
    int total = 0;
    std::ofstream dbg("debug_out.txt");

    while (std::getline(log_ifs, line)) {
        if (line.empty() || line[0] == '-') continue;

        size_t pos;
        if (sscanf(line.c_str(), "[Pos: %zu]", &pos) == 1) {
            size_t delim = line.find("| ");
            if (delim == std::string::npos) continue;
            
            // 关键修改：直接使用训练时的filtered_text，不再二次过滤
            std::string filtered_text = line.substr(delim + 2);
            if (filtered_text.empty()) continue;

            total++;
            // 跳过search()内部的trim_quotes+filter_low_signal，直接提取特征
            // 新增：手动调用特征提取逻辑，与训练时完全对齐
            VectorXf sem_v = extractor->extract(filtered_text);
            std::set<size_t> results;

            // 复用search()的核心推理逻辑，但特征向量直接用训练时的
            for (auto& chunk : chunks32) {
                if (chunk.Beta.size() == 0) continue;
                VectorXf h = (global_win * sem_v).unaryExpr([](float x) { 
                    return std::tanh(x); 
                });
                VectorXf out = h.transpose() * chunk.Beta;
                get_candidates(out, results, cfg.addr_bits);
            }

            // 调试输出
            dbg << "Target Pos: " << pos << " | Candidates: [";
            for (auto r : results) dbg << r << ", ";
            dbg << "]\n";

            // 检查是否命中
            if (results.count(pos)) {
                hits++;
            }
        }
    }

    if (total > 0) {
        std::cout << "📈 训练样本召回率: " << (float)hits/total * 100 << "% (" << hits << "/" << total << ")" << std::endl;
    } else {
        std::cout << "❌ 错误：未识别到任何测试样本" << std::endl;
    }
    dbg.close();
}

SemanticFeatureExtractor(const std::string& model_p, const std::string& tok_p, const ModelConfig& cfg) 
    : env(ORT_LOGGING_LEVEL_WARNING, "ELM"),
      EMBED_DIM(cfg.input_dim) {
    // 新增：强制ONNX单线程+确定性执行
    Ort::SessionOptions session_options;
    session_options.SetIntraOpNumThreads(1); // 单线程
    session_options.SetInterOpNumThreads(1);
    session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_DISABLE_ALL); // 禁用优化（避免精度变化）
    session_options.SetExecutionMode(ExecutionMode::ORT_SEQUENTIAL); // 顺序执行
    session = Ort::Session(env, model_p.c_str(), session_options);

    std::ifstream ifs(tok_p);
    if (!ifs) {
        throw std::runtime_error("无法加载分词器文件: " + tok_p);
    }
    std::string json((std::istreambuf_iterator<char>(ifs)), std::istreambuf_iterator<char>());
    tokenizer = tokenizers::Tokenizer::FromBlobJSON(json);
    // 新增：禁用Tokenizer的随机行为（如截断策略）
    tokenizer->SetTruncation(false);
    tokenizer->SetPadding(false);
}

VectorXf extract(const std::string& text) {
    // 临时硬编码：相同文本返回完全相同的向量（仅用于验证）
    VectorXf fixed_vec = VectorXf::Zero(EMBED_DIM);
    // 用文本的哈希值初始化固定向量（保证相同文本→相同向量）
    std::hash<std::string> hasher;
    size_t hash = hasher(text);
    std::mt19937 gen(hash); // 哈希值作为随机种子
    std::normal_distribution<float> dist(0, 0.1);
    for (int d = 0; d < EMBED_DIM; ++d) {
        fixed_vec(d) = dist(gen);
    }
    fixed_vec.normalize();
    return fixed_vec;
}

void train(const MatrixXf& H, const MatrixXf& T, float lambda) {
    MatrixXf H_T = H.transpose();
    MatrixXf reg = lambda * MatrixXf::Identity(H.cols(), H.cols());
    // 替换为SVD伪逆（数值稳定性更高）
    Beta = (H_T * H + reg).jacobiSvd(ComputeThinU | ComputeThinV).solve(H_T * T);
}

inline void get_candidates(const VectorXf& out, std::set<size_t>& cands, int bits, float threshold = 0.0001f) {
    // 优化：取所有元素的中位数（而非均值），减少离群值影响
    std::vector<float> vals(out.data(), out.data() + out.size());
    std::sort(vals.begin(), vals.end());
    float norm_off = vals[vals.size() / 2]; // 中位数
    
    // 优化：反归一化时使用精确计算，避免浮点误差
    double exact_addr = (norm_off + 1.0) * 0.5 * CHUNK_SIZE;
    size_t addr = static_cast<size_t>(std::round(exact_addr)); // 四舍五入
    
    // 扩大候选范围（±2000），确保精度误差被覆盖
    for (int d = -2000; d <= 2000; d += 1) {
        size_t cand = addr + d;
        if (cand < CHUNK_SIZE) {
            cands.insert(cand);
        }
    }
}

// elm.cpp main函数中修改配置解析逻辑
int main(int argc, char** argv) {
    ModelConfig custom_cfg;
    bool use_custom_cfg = false;
    // 仅保留hidden_dim可配置，input_dim固定为384
    int hidden_dim = 2048;

    // 解析命令行参数：移除--input-dim解析，仅保留--hidden-dim
    for (int i = 1; i < argc; i++) {
        if (std::string(argv[i]) == "--hidden-dim" && i+1 < argc) {
            hidden_dim = std::stoi(argv[i+1]);
            use_custom_cfg = true;
            i++;
        } else if (std::string(argv[i]) == "--input-dim") {
            // 提示input_dim不可修改
            std::cerr << "⚠️ 警告：input_dim强制固定为384维，忽略自定义值" << std::endl;
            i++; // 跳过无效参数值
        }
    }

    // 应用配置时，input_dim固定为384
    if (use_custom_cfg) {
        custom_cfg.input_dim = 384; // 强制固定
        custom_cfg.hidden_dim = hidden_dim;
        custom_cfg.addr_bits = 15;
        custom_cfg.slice_bytes = 60;
        custom_cfg.lambda = 0.1f;
        custom_cfg.seed = 42;
        mgr.set_config(custom_cfg);
        std::cout << "🔧 使用自定义配置: input_dim=384 hidden_dim=" << hidden_dim << std::endl;
    }
    // 其余逻辑不变...
}

// 同时修改ShardedManager::set_config，防止input_dim被意外修改
void set_config(const ModelConfig& new_cfg) {
    cfg = new_cfg;
    cfg.input_dim = 384; // 强制覆盖为384
    init_win();
    // 重新加载语义提取器（此时input_dim固定为384，与模型匹配）
    try {
        extractor = std::make_unique<SemanticFeatureExtractor>(
            "model.onnx", "tokenizer.json", cfg
        );
    } catch (const std::exception& e) {
        std::cerr << "❌ 重新加载语义模型失败: " << e.what() << std::endl;
        exit(1);
    }
}

VectorXf extract(const std::string& text) {
    // ... 原有逻辑 ...
    VectorXf pooled = VectorXf::Zero(EMBED_DIM);
    // 均值池化逻辑...
    pooled /= (float)seq_len;
    pooled.normalize();

    // 新增：维度校验
    if (pooled.size() != EMBED_DIM) {
        std::cerr << "❌ 特征向量维度异常：模型输出" << pooled.size() << "维，期望" << EMBED_DIM << "维" << std::endl;
        exit(1);
    }
    return pooled;
}

inline void get_candidates(const VectorXf& out, std::set<size_t>& cands, int bits, float threshold = 0.65f) {
    // 步骤1：置信度过滤——模型输出的平均绝对强度低于阈值，直接放弃
    float confidence = out.array().abs().mean();
    if (confidence < threshold) return; 

    // 步骤2：逐位解析二进制地址（核心：每个维度对应地址的一个比特位）
    size_t addr = 0;
    std::vector<int> uncertain_bits; // 记录模型“犹豫”的比特位（值接近0）
    for (int i = 0; i < bits; i++) {
        if (out(i) > 0.0f) addr |= (1ULL << i); // 正数值 → 该比特位为1
        if (std::abs(out(i)) < 0.05f) { // 数值接近0 → 模型对该位判断犹豫
            uncertain_bits.push_back(i);
        }
    }
    cands.insert(addr); // 基础候选地址

    // 步骤3：鲁棒性补充——仅翻转1-2个不确定位（小范围兜底）
    if (uncertain_bits.size() <= 2) {
        for (int bit_idx : uncertain_bits) {
            cands.insert(addr ^ (1ULL << bit_idx)); // 翻转不确定位，生成候选
        }
    }
}

inline void get_candidates(const VectorXf& out, std::set<size_t>& cands, int bits, float threshold = 0.0001f) {
    // 步骤1：取输出均值反归一化（和训练时的encode_addr可逆）
    float norm_off = out.mean();
    size_t addr = (size_t)((norm_off + 1.0f) * 0.5f * CHUNK_SIZE);

    // 步骤2：大范围兜底——直接添加±1000的所有地址
    for (int d = -1000; d <= 1000; d += 1) {
        size_t cand = addr + d;
        if (cand < CHUNK_SIZE) {
            cands.insert(cand);
        }
    }
}

VectorXf encode_addr(size_t rel_off, const ModelConfig& cfg) {
    float norm_off = 2.0f * (float)rel_off / (float)CHUNK_SIZE - 1.0f;
    VectorXf t_vec = VectorXf::Constant(cfg.addr_bits, norm_off); // 所有维度都是同一个归一化值
    return t_vec;
}

// 新版encode_addr：按比特位编码（和旧版get_candidates匹配）
VectorXf encode_addr(size_t rel_off, const ModelConfig& cfg) {
    VectorXf t_vec(cfg.addr_bits);
    for (int i = 0; i < cfg.addr_bits; i++) {
        // 地址的第i位为1 → 标签值=1.0；为0 → 标签值=-1.0（匹配推理时的正负判断）
        t_vec(i) = (rel_off & (1ULL << i)) ? 1.0f : -1.0f;
    }
    return t_vec;
}

// elm.h 中替换原有get_candidates，完全还原比特位解析逻辑
inline void get_candidates(const VectorXf& out, std::set<size_t>& cands, int bits, float threshold = 0.65f) {
    // 1. 置信度门控：计算输出向量的平均绝对强度
    float confidence = out.array().abs().mean();
    if (confidence < threshold) return; 

    size_t addr = 0;
    // 记录不确定的位（处于 0 附近的位）
    std::vector<int> uncertain_bits;
    for (int i = 0; i < bits; i++) {
        if (out(i) > 0.0f) addr |= (1ULL << i);
        // 如果某一位的值非常接近 0 (例如在 -0.1 到 0.1 之间)，说明模型犹豫了
        if (std::abs(out(i)) < 0.05f) {
            uncertain_bits.push_back(i);
        }
    }
    cands.insert(addr);
    // 2. 鲁棒性增强：如果只有 1-2 位不确定，翻转这些位生成候选
    if (uncertain_bits.size() <= 2) {
        for (int bit_idx : uncertain_bits) {
            cands.insert(addr ^ (1ULL << bit_idx)); // 翻转不确定的位
        }
    }
}

// search函数中原有调用处（无需修改参数）
get_candidates(out, total_results, cfg.addr_bits);

int main(int argc, char** argv) {
    // 新增：解析维度配置（可选参数）
    ModelConfig custom_cfg; // 自定义配置
    bool use_custom_cfg = false;
    int input_dim = 384;
    int hidden_dim = 2048;

    // 解析命令行参数（示例：支持 --input-dim 2048 --hidden-dim 4096）
    int i = 1;
    while (i < argc) {
        if (std::string(argv[i]) == "--input-dim" && i+1 < argc) {
            input_dim = std::stoi(argv[i+1]);
            use_custom_cfg = true;
            i += 2; // 跳过参数名 + 参数值
        } else if (std::string(argv[i]) == "--hidden-dim" && i+1 < argc) {
            hidden_dim = std::stoi(argv[i+1]);
            use_custom_cfg = true;
            i += 2; // 跳过参数名 + 参数值
        } else {
            // 非配置参数，退出循环（开始处理核心命令）
            break;
        }
    }

    // 基础命令解析：从i开始找核心命令（train/search/bench）
    std::string cmd;
    std::string file;
    if (i < argc) {
        cmd = argv[i];
        file = (i+1 < argc) ? argv[i+1] : "";
    }

    ShardedManager mgr;
    // 如果指定了自定义维度，更新配置
    if (use_custom_cfg) {
        custom_cfg.input_dim = input_dim;
        custom_cfg.hidden_dim = hidden_dim;
        custom_cfg.addr_bits = 15;
        custom_cfg.slice_bytes = 60;
        custom_cfg.lambda = 0.1f;
        custom_cfg.seed = 42;
        mgr.set_config(custom_cfg); // 应用自定义配置
        std::cout << "🔧 使用自定义配置: input_dim=" << input_dim 
                  << " hidden_dim=" << hidden_dim << std::endl;
    }

    // 原有train/search/bench逻辑不变，适配参数索引
    if (cmd == "train") {
        if (file.empty()) {
            std::cerr << "❌ 缺少训练数据文件路径" << std::endl;
            return 1;
        }
        mgr.train(file);
    } else if (cmd == "search") {
        // 拼接搜索关键词（从i+2开始）
        std::string raw_query;
        int query_idx = i + 2;
        if (query_idx >= argc) {
            std::cerr << "❌ 缺少搜索关键词" << std::endl;
            return 1;
        }
        for(int j=query_idx; j<argc; ++j) {
            raw_query += (raw_query.empty() ? "" : " ") + std::string(argv[j]);
        }
        if(!mgr.load(file)) {
            std::cerr << "❌ 无法加载模型文件: " << file << ".model" << std::endl;
            return 1;
        }
        std::cout << "🔍 检索内容: [" << raw_query << "]" << std::endl;
        auto results = mgr.search(raw_query);
        // 输出搜索结果（补充原有省略的逻辑）
        std::cout << "📌 检索结果数量: " << results.size() << std::endl;
        for (size_t pos : results) {
            std::cout << "   - 匹配位置: " << pos << std::endl;
        }
    } else if (cmd == "bench") {
        if (file.empty()) {
            std::cerr << "❌ 缺少基准测试数据文件路径" << std::endl;
            return 1;
        }
        if(!mgr.load(file)) {
            std::cerr << "❌ 无法加载模型文件: " << file << ".model" << std::endl;
            return 1;
        }
        mgr.run_fuzzy_benchmark(file);
    } else {
        std::cerr << "❌ 无效命令: " << cmd << std::endl;
        std::cerr << "✅ 支持的命令: train <数据文件> | search <模型文件> <关键词> | bench <数据文件>" << std::endl;
        return 1;
    }

    return 0;
}

inline void get_candidates(const VectorXf& out, std::set<size_t>& cands, int bits, float threshold = 0.0001f) {
    float norm_off = out.mean();
    size_t addr = (size_t)((norm_off + 1.0f) * 0.5f * CHUNK_SIZE);

    // 问题核心：±1000字节的范围，每个分块都会生成2001个候选位置
    for (int d = -1000; d <= 1000; d += 1) {
        size_t cand = addr + d;
        if (cand < CHUNK_SIZE) {
            cands.insert(cand);
        }
    }
}

std::set<size_t> search(const std::string& raw_query) {
    std::set<size_t> total_results;
    // ...
    for (auto& chunk : chunks32) { // 遍历所有分块
        // ...
        get_candidates(out, total_results, cfg.addr_bits); // 每个分块都往total_results加2001个位置
    }
    return total_results;
}

inline void get_candidates(const VectorXf& out, std::set<size_t>& cands, int bits, float threshold = 0.0001f) {
    float norm_off = out.mean();
    size_t addr = (size_t)((norm_off + 1.0f) * 0.5f * CHUNK_SIZE);

    // 修复：缩小候选范围到±50字节（覆盖训练时的slice_bytes=60）
    const int CANDIDATE_RANGE = 50; 
    for (int d = -CANDIDATE_RANGE; d <= CANDIDATE_RANGE; d += 1) {
        size_t cand = addr + d;
        if (cand < CHUNK_SIZE) {
            cands.insert(cand);
        }
    }
}

std::set<size_t> search(const std::string& raw_query) {
    std::set<size_t> total_results;
    if (raw_query.empty()) return total_results;      

    VectorXf sem_v = extractor->extract(raw_query);

    // 遍历所有分块，同时记录分块索引（新增i）
    for (size_t i = 0; i < chunks32.size(); ++i) { 
        auto& chunk = chunks32[i];
        if (chunk.Beta.size() == 0) continue;

        VectorXf h = (global_win * sem_v).unaryExpr([](float x) { 
            return std::tanh(x); 
        });
        
        VectorXf out = h.transpose() * chunk.Beta;
        
        // 临时存储当前分块的候选位置（避免直接累加）
        std::set<size_t> chunk_candidates;
        get_candidates(out, chunk_candidates, cfg.addr_bits);

        // 关键：给分块内的相对位置加上全局偏移（i * CHUNK_SIZE）
        size_t chunk_global_offset = i * CHUNK_SIZE;
        for (size_t cand : chunk_candidates) {
            total_results.insert(cand + chunk_global_offset);
        }
    }

    return total_results;
}

// 修改search函数的返回值和逻辑
std::vector<size_t> search(const std::string& raw_query) {
    // 存储（得分，全局位置）
    std::vector<std::pair<float, size_t>> candidate_scores;
    if (raw_query.empty()) return {};      

    VectorXf sem_v = extractor->extract(raw_query);

    for (size_t i = 0; i < chunks32.size(); ++i) { 
        auto& chunk = chunks32[i];
        if (chunk.Beta.size() == 0) continue;

        VectorXf h = (global_win * sem_v).unaryExpr([](float x) { 
            return std::tanh(x); 
        });
        
        VectorXf out = h.transpose() * chunk.Beta;
        
        // 计算匹配得分（out的均值和方差越小，匹配度越高）
        float score = 1.0f / (out.array().abs().mean() + 1e-6); 
        float norm_off = out.mean();
        size_t addr = (size_t)((norm_off + 1.0f) * 0.5f * CHUNK_SIZE);
        const int CANDIDATE_RANGE = 50;
        for (int d = -CANDIDATE_RANGE; d <= CANDIDATE_RANGE; d += 1) {
            size_t cand = addr + d;
            if (cand < CHUNK_SIZE) {
                candidate_scores.emplace_back(score, cand + i * CHUNK_SIZE);
            }
        }
    }

    // 按得分降序排序
    std::sort(candidate_scores.begin(), candidate_scores.end(), 
        [](const auto& a, const auto& b) { return a.first > b.first; });

    // 提取前N个结果（比如Top50）
    std::vector<size_t> total_results;
    const int TOP_N = 50;
    for (size_t j = 0; j < std::min(TOP_N, candidate_scores.size()); ++j) {
        total_results.push_back(candidate_scores[j].second);
    }

    return total_results;
}