Skip to content
JingHao-LeonPublic

About

混合检索 RAG 引擎:BM25×稠密双路、RRF 融合、MMR 重排,附标注数据集与 Recall/MRR/nDCG 评测,零模型下载可运行

Topics

Resources

Stars

0 stars

Watchers

0 watching

Forks

Latest commit

 

History

4 Commits

Folders and files

Repository files navigation

rag-forge

BM25 × 稠密向量双路检索、RRF 融合、MMR 多样性重排、标注数据集评测——纯 Python 零模型下载即可运行,一行命令换装 sentence-transformers 生产级向量。

tests Python License: MIT

为什么是"混合"检索?

BM25 抓得住精确术语(型号、专名、代码),稠密向量懂语义改写("胡编乱造" ≈ "幻觉"),两者错误模式互补。融合用 RRF(Reciprocal Rank Fusion)——只看排名不看原始分,天然免疫两路打分量纲不可比的问题。

📊 内置标注数据集实测(20 篇中文 AI 文档 × 10 个查询)

离线哈希嵌入(Mac/任意机器,零模型下载)与 BGE 真嵌入(RTX 3090) 双组实测:

配置 recall@1 recall@3 recall@5 MRR@10 nDCG@5
BM25 only 0.70 1.00 1.00 0.95 0.955
哈希 dense 0.40 0.80 0.90 0.70 0.736
哈希 hybrid 0.50 0.90 1.00 0.82 0.857
BGE dense(真嵌入) 0.45 0.90 0.90 0.783 0.776
BGE hybrid(真嵌入) 0.80 0.90 0.95 1.000 0.946

三个观察:① 真嵌入 dense 单路 recall@1 只比哈希高 5 分——小语料上 BM25 的精确匹配很难被打败;② 换真嵌入后 hybrid 的 recall@1 从 0.50 跳到 0.80、MRR 满分——两路错误模式互补的价值随嵌入质量同步放大;③ 两种嵌入下 hybrid 都稳定优于 dense 单路,但对 BM25 的优势取决于嵌入质量:哈希嵌入时 hybrid 的 recall@1(0.50)低于 BM25(0.70),换真嵌入后才在 recall@1 与 MRR 上反超——RRF 融合的收益来自两路错误模式互补,并非在所有指标上都保证超过每一路。数据:results/bge_eval_3090.json,复现:examples/eval_bge_gpu.py(3090 实测)。

修过一个真 bug(值得写进测试的那种)

HybridPipeline.build() 曾无条件重建稠密索引,把注入的自定义 embedder 覆盖回哈希——BGE 评测跑出的数字与哈希一字不差才暴露。已改为构造参数 embedder= 并补回归测试 test_pipeline_keeps_custom_embedder_after_build。

✨ 核心设计

  • 中英混合分词:英文小写化 + 中文单字/二元组,无词典依赖(二元组让 BM25 在无分词器时依然能命中「检索增强」)
  • 分隔符对齐分块:按 \n\n → \n → 。 → ; → 空格 逐级切分,随机 fuzz 验证精确分区(默认 200 次试验,"".join(chunks) == 原文,不丢一个字符)
  • BM25Okapi:Robertson 2-epsilon IDF(恒非负),k1/b 可调
  • RRF + 加权融合双方案,单元素/全同分等边界均有测试覆盖
  • MMR 重排:λ 平衡相关性与多样性,治"五条结果一个味"
  • LLM 重排适配器:make_llm_reranker(any_llm) 注入你自己的模型调用,库本身零 LLM 依赖
  • 评测三件套:Recall@k / MRR / nDCG@k,标注格式就是 jsonl,3 分钟接上自己的数据

🚀 快速开始

uv sync   # 仅 numpy + pytest

# 检索
uv run python -m rag_forge.cli search --corpus data/docs.jsonl \
    --query "怎么在一张显卡上微调大模型" --top-k 3

# 评测(bm25 / dense / hybrid 三模式对比)
uv run python -m rag_forge.cli eval --corpus data/docs.jsonl --queries data/queries.jsonl --mode hybrid

作为库使用:

from rag_forge.pipeline import HybridPipeline

pipe = HybridPipeline(chunk_size=300).build({"d1": "文档内容...", "d2": "..."})
results = pipe.search("微调大模型需要多少显存", top_k=3, mode="hybrid")
for r in results:
    print(r.doc_id, r.score, f"bm25#{r.bm25_rank} dense#{r.dense_rank}")

换装生产级向量模型(可选):

from rag_forge.pipeline import HybridPipeline
from rag_forge.dense import SentenceTransformerEmbedder

pipe = HybridPipeline(embedder=SentenceTransformerEmbedder("BAAI/bge-small-zh-v1.5"))
pipe.build(corpus)  # 其余逻辑完全不变

📁 结构

rag_forge/
├── tokenize.py    # 中英混合分词(中文 unigram + bigram)
├── chunking.py    # fixed / recursive 分块,精确分区不变量(fuzz 验证)
├── bm25.py        # BM25Okapi 稀疏索引
├── dense.py       # 哈希嵌入(零依赖)+ sentence-transformers 适配
├── fusion.py      # RRF / 加权 min-max 融合
├── rerank.py      # MMR 多样性重排 + LLM 重排适配器
├── pipeline.py    # HybridPipeline:chunk → 双路索引 → 融合 → 重排
├── eval.py        # Recall@k / MRR / nDCG@k
└── cli.py         # search / eval 命令行
data/              # 20 篇标注文档 + 10 个查询(中文 AI 主题)
tests/             # 16 个测试 + 200 次随机 fuzz(默认强度)

🧪 测试

uv run pytest                              # 16 passed(含默认 200 次随机 fuzz)
uv run pytest tests/test_fuzz_chunking.py -v   # 单独跑分块 fuzz;想加大强度,把
                                           # tests/test_fuzz_chunking.py 里 _run(200) 的参数调大

Roadmap

  • HyDE 查询改写(LLM 生成假设文档再检索)
  • 块级评测(当前为文档级标注)
  • FAISS/HNSW 后端适配(大规模库)

常见问题(FAQ)

Q1:MRR@10 = 1.000 是在什么数据上测的?

自建标注集:20 篇中文 AI 主题文档 × 10 个查询(data/docs.jsonl + data/queries.jsonl),jsonl 格式 {"query": "...", "relevant": ["d01"]} 的文档级标注(部分查询有 2 个相关文档)。BGE 两行的数字来自 RTX 3090 上 examples/eval_bge_gpu.py 的一次实测,原始结果存于 results/bge_eval_3090.json。哈希嵌入三行在任意机器上 uv run python -m rag_forge.cli eval ... 即可复现;换自己的数据只需同格式 jsonl。

Q2:RRF 的 k 怎么选?

代码默认 rrf_k = 60(HybridPipeline 与 reciprocal_rank_fusion 的默认值,也是 RRF 原论文及常见检索系统的经验值)。k 只影响排名贡献 1/(k + rank) 的平滑程度:k 越大融合越"平",单路头部排名的优势被稀释得越多。一般场景直接用默认 60 即可;想更强调"某一路把文档排得很靠前"的信号,可调小 k。

Q3:零模型下载模式和真 BGE 模式有什么区别?

默认的 HashEmbedder 把字符 n-gram 特征哈希到 512 维(sublinear TF + L2 归一化),纯 numpy、零下载、CPU 可跑,但只是浅层语义信号——dense 单路 recall@1 仅 0.40。通过 SentenceTransformerEmbedder("BAAI/bge-small-zh-v1.5") 换装真嵌入(sentence-transformers 是可选依赖 models,需另行安装,首次运行会下载权重)后,同一套混合管线 recall@1 升至 0.80、MRR@10 满分。BM25 一路与分块逻辑在两种模式下完全一致,管线其余部分零改动。

Q4:为什么用 RRF 融合,而不是把两路分数直接加权相加?

BM25 分数与余弦相似度的量纲不可比,直接相加没有意义。RRF 只用排名(score = Σ 1/(k + rank)),天然免疫量纲问题,也不需要调两路的分数缩放。如果确实有可信的分数体系,库另提供 weighted_fusion(逐路 min-max 归一化后加权求和)作为备选。

Q5:评测分数这么高,是玩具数据吧?

对,规模就是玩具级:20 文档 × 10 查询、二值相关、文档级标注。这些数字用于验证工程正确性与配置间的相对对比(BM25/dense/hybrid、哈希/BGE),不能外推为生产规模下的检索质量,详见下方"局限与已知问题"。

局限与已知问题(Limitations)

  • 标注集规模小:20 篇中文 AI 文档 × 10 个查询,MRR@10 = 1.000 等指标均在此小集合上测得,样本量不足以支撑统计显著性,主要价值是管线正确性与相对对比。
  • 指标受标注质量约束:relevant 集合是人工圈定的文档级二值标注(nDCG 亦为二值相关实现),"未标注 ≠ 不相关",指标上限受标注覆盖面限制;块级评测尚未支持(见 Roadmap)。
  • 中文语料为主:内置语料与查询全部为中文 AI 主题;分词是"英文小写化 + 中文单字/二元组"的轻量方案,无真正的分词器依赖,英文及其他语言场景未做评测。
  • 稠密检索是暴力精确搜索:全量矩阵乘算 cosine,无 ANN 索引(FAISS/HNSW 在 Roadmap),语料规模上去后会慢。
  • 哈希嵌入只是浅层语义:默认嵌入的意义在于让库离线可跑、验证管线正确性,生产语义质量需换真嵌入模型,代价是下载权重与推理资源(3090 上 BGE 构建索引 231.5 秒,10 个查询评测 0.55 秒)。
  • hybrid 不是处处更优:BM25 已近满分的小语料上,融合可能不如 BM25 单路(哈希 hybrid recall@1 0.50 < BM25 0.70);hybrid 的收益依赖嵌入质量与两路错误模式的互补程度。

License

MIT

About

混合检索 RAG 引擎:BM25×稠密双路、RRF 融合、MMR 重排,附标注数据集与 Recall/MRR/nDCG 评测,零模型下载可运行

Topics

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages