4.6 KiB
4.6 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
支持的嵌入模型类型
1. 本地HuggingFace模型 (local)
- 模型名称: 从HuggingFace Hub下载
- 本地路径: 使用已下载的本地模型
2. 本地部署API接口 (api)
- 描述: 连接到本地部署的嵌入服务
- 兼容: OpenAI API格式的本地服务
- 回退机制: API不可用时自动回退到本地模型
使用方法
基本配置
from base_rag import BaseRAG
# 嵌入模型配置
embedding_config = {
"type": "local", # 或 "api"
"model_name": "sentence-transformers/all-MiniLM-L6-v2"
}
# 重排配置
rerank_config = {
"enabled": True,
"method": "similarity", # 或 "cross_encoder", "bge"
"top_k": 3 # 重排后返回的文档数量
}
# 初始化RAG系统
rag = MyRAG(
embedding_config=embedding_config,
rerank_config=rerank_config
)
本地模型配置
# 使用HuggingFace模型名称
embedding_config = {
"type": "local",
"model_name": "sentence-transformers/all-MiniLM-L6-v2"
}
# 使用本地模型路径
embedding_config = {
"type": "local",
"model_path": "/path/to/your/local/model"
}
本地API接口配置
embedding_config = {
"type": "api",
"api_url": "http://localhost:8000", # 本地API地址
"model": "your-embedding-model",
"api_key": "optional-key" # 可选
}
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
}
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重排: 平衡了效果和性能
- 本地API: 网络延迟较小,可扩展性好
最佳实践
- 开发阶段: 使用相似度重排进行快速原型验证
- 生产环境: 根据具体需求选择CrossEncoder或BGE重排
- 分布式部署: 考虑使用本地API接口
- 候选文档数量: 建议设置为最终需要文档数量的2-3倍
- 模型选择: 根据语言和领域选择合适的重排模型
依赖安装
# 基本依赖
pip install langchain-huggingface langchain-chroma
# CrossEncoder重排
pip install sentence-transformers
# BGE重排
pip install FlagEmbedding
# 本地API支持(可选)
pip install langchain-openai
注意事项
- 重排会增加查询延迟,但通常能显著提高结果质量
- 如果重排模型加载失败,系统会自动回退到相似度重排
- 不同重排方法的效果可能因数据集和查询类型而异
- 本地API接口需要兼容OpenAI API格式
- 建议在实际数据上进行A/B测试来选择最适合的重排方法