245 - BM25

#include <vector>
#include <unordered_map>
#include <shared_mutex>
#include <cstdint>
#include <string>
#include <stdexcept>

// 1. 文档元数据：存储计算BM25所需的文档级信息
struct DocMeta {
    uint32_t doc_id;       // 文档唯一ID（用uint32_t而非int，节省4字节/文档）
    uint32_t doc_length;   // 文档长度（术语数量）
    static float avg_len;  // 全局平均文档长度（静态变量，预处理后初始化）
};
float DocMeta::avg_len = 0.0f;

// 2. 倒排索引项：存储单个术语在所有文档中的出现信息
struct InvertedIndexItem {
    uint32_t term_id;                  // 术语唯一ID（映射字符串，减少内存占用）
    std::vector<uint32_t> doc_id_list; // 包含该术语的文档ID列表
    std::vector<float> tf_list;        // 对应文档的TF值（提前计算，避免实时计算）
};

// 3. 全局术语映射表：术语字符串→ID（线程安全读写）
class TermDict {
private:
    std::unordered_map<std::string, uint32_t> term_to_id_;
    std::shared_mutex rw_mutex_;        // 读写锁：读多写少场景优化（C++17特性）
    uint32_t next_term_id_ = 1;         // 术语ID从1开始（0预留为无效值）
public:
    // 线程安全的术语ID获取（不存在则创建）
    uint32_t get_or_create_term_id(const std::string& term) {
        // 读锁：先检查是否存在
        std::shared_lock<std::shared_mutex> read_lock(rw_mutex_);
        if (auto it = term_to_id_.find(term); it != term_to_id_.end()) {
            return it->second;
        }
        read_lock.unlock();

        // 写锁：创建新术语ID
        std::unique_lock<std::shared_mutex> write_lock(rw_mutex_);
        auto [it, success] = term_to_id_.emplace(term, next_term_id_);
        if (success) {
            next_term_id_++;
        }
        return it->second;
    }

    // 仅读接口（用于查询阶段）
    uint32_t get_term_id(const std::string& term) const {
        std::shared_lock<std::shared_mutex> read_lock(rw_mutex_);
        if (auto it = term_to_id_.find(term); it != term_to_id_.end()) {
            return it->second;
        }
        return 0; // 无效ID
    }
};

#include "jieba.h" // 假设集成Jieba C++分词库（https://github.com/yanyiwu/cppjieba）
#include <unordered_set>

// 停用词集合（示例）
const std::unordered_set<std::string> STOP_WORDS = {"的", "了", "是", "在", "C++", "the", "a"};

// 文档预处理函数：输入原始文档，输出（文档元数据，术语ID-词频映射）
std::pair<DocMeta, std::unordered_map<uint32_t, uint32_t>> 
preprocess_document(const std::string& raw_doc, uint32_t doc_id, const TermDict& term_dict, const cppjieba::Jieba& jieba) {
    // 1. 异常处理：空文档直接抛出
    if (raw_doc.empty()) {
        throw std::invalid_argument("Empty document, doc_id: " + std::to_string(doc_id));
    }

    // 2. 分词（Jieba接口调用）
    std::vector<std::string> terms;
    jieba.Cut(raw_doc, terms, true); // 精确分词，去停用词前

    // 3. 过滤停用词+术语ID映射
    std::unordered_map<uint32_t, uint32_t> term_id_to_tf; // 术语ID→词频
    for (const auto& term : terms) {
        // 过滤停用词、空术语
        if (term.empty() || STOP_WORDS.count(term)) {
            continue;
        }
        // 获取术语ID（不存在则返回0，跳过）
        uint32_t term_id = term_dict.get_term_id(term);
        if (term_id == 0) {
            continue;
        }
        term_id_to_tf[term_id]++;
    }

    // 4. 构建文档元数据
    DocMeta doc_meta;
    doc_meta.doc_id = doc_id;
    doc_meta.doc_length = static_cast<uint32_t>(term_id_to_tf.size()); // 文档长度=非停用词术语数

    return {doc_meta, term_id_to_tf};
}

#include <mutex>
#include <vector>
#include <numeric>

class BM25Index {
private:
    std::unordered_map<uint32_t, InvertedIndexItem> index_; // 术语ID→倒排索引项
    std::mutex index_mutex_;                                // 倒排索引写入锁（多线程安全）
    std::vector<DocMeta> all_docs_;                         // 存储所有文档元数据
public:
    // 添加单篇文档到索引（多线程可调用）
    void add_document(const DocMeta& doc_meta, const std::unordered_map<uint32_t, uint32_t>& term_id_to_tf) {
        std::lock_guard<std::mutex> lock(index_mutex_);

        // 1. 记录文档元数据
        all_docs_.push_back(doc_meta);

        // 2. 更新倒排索引：遍历当前文档的所有术语
        for (const auto& [term_id, tf] : term_id_to_tf) {
            auto& index_item = index_[term_id]; // 不存在则自动创建
            index_item.term_id = term_id;
            index_item.doc_id_list.push_back(doc_meta.doc_id);
            // 计算TF值（BM25中TF通常用“术语出现次数”，此处直接存储）
            index_item.tf_list.push_back(static_cast<float>(tf));
        }
    }

    // 索引构建完成后，计算全局平均文档长度
    void calculate_avg_doc_length() {
        if (all_docs_.empty()) {
            DocMeta::avg_len = 0.0f;
            return;
        }
        // 求和所有文档长度（用std::accumulate简化代码）
        uint64_t total_len = std::accumulate(
            all_docs_.begin(), all_docs_.end(), 
            0ULL, [](uint64_t sum, const DocMeta& doc) {
                return sum + doc.doc_length;
            }
        );
        DocMeta::avg_len = static_cast<float>(total_len) / all_docs_.size();
    }

    // 获取倒排索引项（查询阶段用）
    const InvertedIndexItem* get_index_item(uint32_t term_id) const {
        auto it = index_.find(term_id);
        return (it != index_.end()) ? &(it->second) : nullptr;
    }

    // 获取文档元数据（查询阶段用）
    const DocMeta* get_doc_meta(uint32_t doc_id) const {
        // 假设doc_id连续（实际项目可用unordered_map优化查找）
        if (doc_id >= all_docs_.size()) {
            return nullptr;
        }
        return &all_docs_[doc_id];
    }

    // 获取总文档数（计算IDF用）
    uint32_t get_total_docs() const {
        return static_cast<uint32_t>(all_docs_.size());
    }
};

#include <cmath>

// BM25评分配置（可动态调整）
struct BM25Config {
    float k1 = 1.5f; // 词频饱和系数
    float b = 0.75f; // 文档长度归一化系数
};

class BM25Scorer {
private:
    const BM25Index& index_;
    const BM25Config config_;
public:
    BM25Scorer(const BM25Index& index, const BM25Config& config) 
        : index_(index), config_(config) {}

    // 计算单个术语-文档的BM25评分
    float score_term_doc(uint32_t term_id, const DocMeta& doc_meta) const {
        // 1. 获取倒排索引项（不存在则返回0）
        const InvertedIndexItem* index_item = index_.get_index_item(term_id);
        if (index_item == nullptr) {
            return 0.0f;
        }

        // 2. 查找当前文档的TF值（假设doc_id在doc_id_list中有序，用线性查找；实际可用二分优化）
        float tf = 0.0f;
        for (size_t i = 0; i < index_item->doc_id_list.size(); ++i) {
            if (index_item->doc_id_list[i] == doc_meta.doc_id) {
                tf = index_item->tf_list[i];
                break;
            }
        }
        if (tf == 0.0f) {
            return 0.0f;
        }

        // 3. 计算IDF（平滑处理：避免分母为0）
        uint32_t doc_count_with_term = static_cast<uint32_t>(index_item->doc_id_list.size()); // n_q
        uint32_t total_docs = index_.get_total_docs(); // N
        float idf = log( (total_docs - doc_count_with_term + 0.5f) / (doc_count_with_term + 0.5f) ) + 1.0f;
        if (idf < 0) {
            idf = 0.0f; // 过滤负IDF（术语在多数文档中出现，无区分度）
        }

        // 4. 计算文档长度归一化因子
        float len_norm = 1.0f - config_.b + config_.b * (doc_meta.doc_length / DocMeta::avg_len);

        // 5. 计算BM25评分（公式落地）
        float tf_adjusted = (tf * (config_.k1 + 1.0f)) / (tf + config_.k1 * len_norm);
        return idf * tf_adjusted;
    }

    // 计算查询与文档的总评分（多术语求和）
    float score_query_doc(const std::vector<uint32_t>& query_term_ids, uint32_t doc_id) const {
        const DocMeta* doc_meta = index_.get_doc_meta(doc_id);
        if (doc_meta == nullptr) {
            return 0.0f;
        }

        float total_score = 0.0f;
        for (uint32_t term_id : query_term_ids) {
            total_score += score_term_doc(term_id, *doc_meta);
        }
        return total_score;
    }
};