3.8 KiB
3.8 KiB
BaseRAG 重排功能说明
功能概述
BaseRAG现在支持对检索到的文档片段进行重排(rerank),这是RAG系统中一个重要的优化环节,可以提高检索质量和答案的相关性。
支持的重排方法
1. 相似度重排 (similarity)
- 描述: 基于余弦相似度的重排方法
- 优点: 无需额外依赖,使用现有的嵌入模型
- 适用场景: 轻量级部署,快速原型验证
2. CrossEncoder重排 (cross_encoder)
- 描述: 使用预训练的CrossEncoder模型进行重排
- 依赖:
sentence-transformers
- 优点: 专门训练的重排模型,效果通常更好
- 默认模型:
cross-encoder/ms-marco-MiniLM-L-6-v2
3. BGE重排 (bge)
- 描述: 使用BGE(BAAI General Embedding)重排模型
- 依赖:
FlagEmbedding
- 优点: 中文支持较好,性能优秀
- 默认模型:
BAAI/bge-reranker-base
4. Cohere重排 (cohere)
- 描述: 使用Cohere API进行重排
- 依赖:
cohere
- 优点: 云端API,无需本地模型
- 注意: 需要API密钥
使用方法
基本配置
from base_rag import BaseRAG
# 重排配置
rerank_config = {
"enabled": True,
"method": "similarity", # 或 "cross_encoder", "bge", "cohere"
"top_k": 3 # 重排后返回的文档数量
}
# 初始化RAG系统
rag = MyRAG(
embedding_config=embedding_config,
rerank_config=rerank_config
)
CrossEncoder重排配置
rerank_config = {
"enabled": True,
"method": "cross_encoder",
"model": "cross-encoder/ms-marco-MiniLM-L-6-v2", # 可选,自定义模型
"top_k": 3
}
BGE重排配置
rerank_config = {
"enabled": True,
"method": "bge",
"model": "BAAI/bge-reranker-base", # 可选,自定义模型
"top_k": 3
}
Cohere重排配置
rerank_config = {
"enabled": True,
"method": "cohere",
"api_key": "your_cohere_api_key",
"top_k": 3
}
API方法
similarity_search_with_rerank(query, k=None)
带重排功能的相似性搜索方法,会自动应用配置的重排策略。
# 使用重排搜索
docs = rag.similarity_search_with_rerank("查询问题", k=5)
similarity_search(query, k=None)
原始的相似性搜索方法,不使用重排。
# 不使用重排搜索
docs = rag.similarity_search("查询问题", k=5)
重排工作原理
- 候选检索: 首先从向量数据库检索更多的候选文档(通常是目标数量的2倍)
- 重排计算: 根据配置的重排方法,计算查询与每个候选文档的相关性分数
- 重新排序: 按照相关性分数对文档进行重新排序
- 截取结果: 返回top_k个最相关的文档
性能考虑
- 相似度重排: 计算开销小,速度快
- CrossEncoder重排: 需要加载额外模型,计算较慢但效果好
- BGE重排: 平衡了效果和性能
- Cohere重排: 网络延迟,但无需本地计算资源
最佳实践
- 开发阶段: 使用相似度重排进行快速原型验证
- 生产环境: 根据具体需求选择CrossEncoder或BGE重排
- 云端部署: 考虑使用Cohere API重排
- 候选文档数量: 建议设置为最终需要文档数量的2-3倍
- 模型选择: 根据语言和领域选择合适的重排模型
依赖安装
# CrossEncoder重排
pip install sentence-transformers
# BGE重排
pip install FlagEmbedding
# Cohere重排
pip install cohere
注意事项
- 重排会增加查询延迟,但通常能显著提高结果质量
- 如果重排模型加载失败,系统会自动回退到相似度重排
- 不同重排方法的效果可能因数据集和查询类型而异
- 建议在实际数据上进行A/B测试来选择最适合的重排方法