feat: 补充rerank

This commit is contained in:
李如威 2025-07-29 00:38:54 +08:00
parent 34e78d0ce0
commit a0cc4aab3a
4 changed files with 585 additions and 15 deletions

144
RERANK_GUIDE.md Normal file
View File

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

View File

@ -12,34 +12,79 @@ class SimpleRAG(BaseRAG):
self.vector_store.add_texts([doc])
def query(self, question, k=3):
docs = self.vector_store.similarity_search(question, k=k)
# 使用带重排功能的搜索
docs = self.similarity_search_with_rerank(question, k=k)
return docs
def query_without_rerank(self, question, k=3):
# 不使用重排的普通搜索
docs = self.similarity_search(question, k=k)
return docs
def main():
config = {
# 嵌入模型配置
embedding_config = {
"type": "local",
"model_name": "sentence-transformers/all-MiniLM-L6-v2",
"embedding_type": "local",
}
rag = SimpleRAG(embedding_config=config)
print("RAG初始化完成!")
# 重排配置 - 使用基于余弦相似度的重排
rerank_config = {
"enabled": True,
"method": "similarity", # 使用相似度重排(无需额外依赖)
"top_k": 3
}
# 添加一些文档
rag = SimpleRAG(
embedding_config=embedding_config,
rerank_config=rerank_config
)
print("RAG系统含重排功能初始化完成!")
# 添加更多测试文档
documents = [
"苹果是一种水果,味道甜美,营养丰富。",
"苹果公司是一家科技公司生产iPhone和Mac等产品。",
"Python是一种编程语言简单易学功能强大。",
"苹果是一种红色或绿色的水果味道甜美营养丰富含有丰富的维生素C。",
"苹果公司是一家总部位于美国的科技公司以生产iPhone、iPad、Mac等产品而闻名。",
"Python是一种高级编程语言简单易学功能强大广泛用于数据科学和机器学习。",
"苹果树是一种果树,春天开花,秋天结果,需要充足的阳光和水分。",
"苹果派是一种传统的美式甜点,由苹果馅和酥脆的派皮制成。",
"苹果醋是由苹果发酵制成的,具有一定的保健功效,可以帮助消化。",
"iPhone是苹果公司生产的智能手机具有先进的技术和优秀的用户体验。",
"机器学习是人工智能的一个分支Python是机器学习领域最流行的编程语言之一。"
]
print("正在添加文档...")
rag.ingest(documents)
print("文档添加完成!")
print(f"文档添加完成! 共添加了 {len(documents)} 个文档")
# 测试查询
print("\n正在查询: '什么是苹果?'")
result = rag.query("什么是苹果?")
print(f"查询结果: {result}")
# 测试查询并比较重排效果
test_query = "苹果的营养价值和健康效益"
print(f"\n测试查询: '{test_query}'")
print("\n=== 不使用重排的结果 ===")
result_no_rerank = rag.query_without_rerank(test_query, k=3)
for i, doc in enumerate(result_no_rerank, 1):
print(f"{i}. {doc.page_content}")
print("\n=== 使用重排的结果 ===")
result_with_rerank = rag.query(test_query, k=3)
for i, doc in enumerate(result_with_rerank, 1):
print(f"{i}. {doc.page_content}")
# 另一个测试查询
test_query2 = "苹果公司的主要产品"
print(f"\n\n测试查询2: '{test_query2}'")
print("\n=== 不使用重排的结果 ===")
result_no_rerank2 = rag.query_without_rerank(test_query2, k=3)
for i, doc in enumerate(result_no_rerank2, 1):
print(f"{i}. {doc.page_content}")
print("\n=== 使用重排的结果 ===")
result_with_rerank2 = rag.query(test_query2, k=3)
for i, doc in enumerate(result_with_rerank2, 1):
print(f"{i}. {doc.page_content}")
if __name__ == "__main__":

139
examples/rerank_demo.py Normal file
View File

@ -0,0 +1,139 @@
import sys
import os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))
from base_rag import BaseRAG
class RerankRAG(BaseRAG):
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.document_count = 0
def ingest(self, documents):
"""批量添加文档,避免重复"""
if documents:
# 清空现有集合并重新添加所有文档
self.vector_store.delete_collection()
# 重新初始化向量库
from langchain_chroma import Chroma
self.vector_store = Chroma(
collection_name=self.vector_store_name,
embedding_function=self.embedding_model,
persist_directory=self.persist_directory,
)
# 添加所有文档
self.vector_store.add_texts(documents)
self.document_count = len(documents)
print(f"已添加 {self.document_count} 个文档到向量库")
def query_with_scores(self, question, k=5):
"""带分数的查询,用于比较重排效果"""
# 不使用重排的结果
docs_no_rerank = self.similarity_search(question, k=k)
# 使用重排的结果
docs_with_rerank = self.similarity_search_with_rerank(question, k=k)
return docs_no_rerank, docs_with_rerank
def query(self, question, k=3):
return self.similarity_search_with_rerank(question, k=k)
def main():
# 嵌入模型配置
embedding_config = {
"type": "local",
"model_name": "sentence-transformers/all-MiniLM-L6-v2",
}
# 重排配置
rerank_config = {
"enabled": True,
"method": "similarity",
"top_k": 3
}
# 初始化RAG系统
rag = RerankRAG(
vector_store_name="rerank_test",
embedding_config=embedding_config,
rerank_config=rerank_config,
retriever_top_k=5 # 获取更多候选文档
)
print("RAG系统含重排功能初始化完成!\n")
# 测试文档 - 关于不同主题的文档
documents = [
# 水果相关
"苹果是一种非常有营养的水果富含维生素C、纤维和抗氧化剂对心脏健康有益。",
"橙子含有丰富的维生素C是柑橘类水果的代表有助于增强免疫系统。",
"香蕉富含钾元素,能够帮助维持血压稳定,是运动员的理想能量补充。",
# 科技公司相关
"苹果公司Apple Inc.是全球知名的科技公司主要产品包括iPhone、iPad、Mac电脑等。",
"谷歌公司专注于搜索引擎和云计算服务Android操作系统是其重要产品。",
"微软公司开发Windows操作系统和Office办公软件在企业软件领域占据重要地位。",
# 编程语言相关
"Python是一种高级编程语言语法简洁广泛用于数据科学、机器学习和Web开发。",
"Java是面向对象的编程语言具有跨平台特性在企业级开发中应用广泛。",
"JavaScript是Web开发的核心语言可以实现网页的交互功能和动态效果。",
# 健康相关
"规律运动有助于维持身体健康建议每周至少进行150分钟的中等强度有氧运动。",
"均衡饮食是健康的基础,应该多吃蔬菜水果,减少加工食品的摄入。",
"充足的睡眠对身心健康至关重要成年人每天应保证7-9小时的睡眠时间。"
]
print("正在添加测试文档...")
rag.ingest(documents)
print(f"文档添加完成!\n")
# 测试查询
test_queries = [
{
"query": "苹果对健康有什么好处?",
"expected_topic": "应该更偏向水果营养相关的文档"
},
{
"query": "苹果公司的主要业务是什么?",
"expected_topic": "应该更偏向科技公司相关的文档"
},
{
"query": "如何保持身体健康?",
"expected_topic": "应该更偏向健康建议相关的文档"
}
]
for i, test_case in enumerate(test_queries, 1):
query = test_case["query"]
expected = test_case["expected_topic"]
print(f"=== 测试查询 {i}: {query} ===")
print(f"期望结果: {expected}\n")
# 获取两种搜索结果
docs_no_rerank, docs_with_rerank = rag.query_with_scores(query, k=5)
print("📍 不使用重排的结果:")
for j, doc in enumerate(docs_no_rerank[:3], 1):
print(f" {j}. {doc.page_content}")
print("\n🎯 使用重排的结果:")
for j, doc in enumerate(docs_with_rerank[:3], 1):
print(f" {j}. {doc.page_content}")
print("\n" + "="*80 + "\n")
print("重排功能测试完成!")
print("\n说明:")
print("- 重排功能通过计算查询与文档的余弦相似度来重新排序检索结果")
print("- 理论上重排后的结果应该更符合查询意图")
print("- 如果结果相同,可能是因为初始检索结果已经很好,或需要更多样化的测试数据")
if __name__ == "__main__":
main()

View File

@ -1,6 +1,7 @@
from abc import ABC, abstractmethod
from typing import List, Optional, Dict, ClassVar, Union
from typing import List, Optional, Dict, ClassVar, Union, Tuple
import threading
import numpy as np
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_openai import OpenAIEmbeddings
@ -26,6 +27,7 @@ class BaseRAG(ABC):
retriever_top_k: int = 3,
llm: Optional[BaseLLM] = None,
persist_directory: str = "./chroma_db",
rerank_config: Optional[Dict] = None,
):
"""
初始化基础RAG类
@ -34,11 +36,16 @@ class BaseRAG(ABC):
:param retriever_top_k: 检索返回的文档数量
:param llm: 可选的对话模型
:param persist_directory: Chroma持久化目录
:param rerank_config: 重排配置
embedding_config 示例:
本地模型名称: {"type": "local", "model_name": "sentence-transformers/all-MiniLM-L6-v2"}
本地模型路径: {"type": "local", "model_path": "/path/to/your/model"}
OpenAI API: {"type": "openai", "model": "text-embedding-ada-002", "api_key": "sk-..."}
rerank_config 示例:
{"enabled": True, "method": "cross_encoder", "model": "cross-encoder/ms-marco-MiniLM-L-6-v2", "top_k": 3}
{"enabled": True, "method": "cohere", "api_key": "your_cohere_key", "top_k": 3}
"""
self.vector_store_name = vector_store_name
self.embedding_config = embedding_config or {
@ -48,6 +55,7 @@ class BaseRAG(ABC):
self.retriever_top_k = retriever_top_k
self.llm = llm
self.persist_directory = persist_directory
self.rerank_config = rerank_config or {"enabled": False}
# 使用缓存的嵌入模型
config_key = self._get_config_key(self.embedding_config)
@ -55,6 +63,11 @@ class BaseRAG(ABC):
config_key, self.embedding_config
)
# 初始化重排模型
self.reranker = None
if self.rerank_config.get("enabled", False):
self.reranker = self._init_reranker()
# 初始化 Chroma 向量库
self.vector_store = Chroma(
collection_name=vector_store_name,
@ -149,6 +162,215 @@ class BaseRAG(ABC):
f"不支持的嵌入模型类型: {config_type},支持的类型: 'local', 'openai'"
)
def _init_reranker(self):
"""初始化重排模型"""
method = self.rerank_config.get("method", "cross_encoder")
# 相似度重排不需要额外的模型
if method == "similarity":
print("使用基于余弦相似度的重排方法")
return "similarity" # 返回标识符
if method == "cross_encoder":
try:
from sentence_transformers import CrossEncoder
model_name = self.rerank_config.get("model", "cross-encoder/ms-marco-MiniLM-L-6-v2")
print(f"正在加载CrossEncoder重排模型: {model_name}")
return CrossEncoder(model_name)
except ImportError:
print("警告: sentence-transformers未安装无法使用CrossEncoder重排")
return None
elif method == "cohere":
try:
import cohere
api_key = self.rerank_config.get("api_key")
if not api_key:
print("警告: 未提供Cohere API密钥无法使用Cohere重排")
return None
print("正在初始化Cohere重排模型")
return cohere.Client(api_key)
except ImportError:
print("警告: cohere包未安装无法使用Cohere重排")
return None
elif method == "bge":
try:
from FlagEmbedding import FlagReranker
model_name = self.rerank_config.get("model", "BAAI/bge-reranker-base")
print(f"正在加载BGE重排模型: {model_name}")
return FlagReranker(model_name, use_fp16=True)
except ImportError:
print("警告: FlagEmbedding未安装无法使用BGE重排")
return None
else:
print(f"警告: 不支持的重排方法: {method}")
return None
def _rerank_documents(self, query: str, documents: List[Document]) -> List[Document]:
"""对检索到的文档进行重排"""
if not documents:
return documents
method = self.rerank_config.get("method", "cross_encoder")
top_k = self.rerank_config.get("top_k", len(documents))
# 如果是相似度重排,直接调用相似度重排方法
if method == "similarity":
return self._similarity_rerank(query, documents)
# 其他方法需要reranker模型
if not self.reranker:
print(f"重排模型未初始化,使用默认相似度重排")
return self._similarity_rerank(query, documents)
try:
if method == "cross_encoder":
# 准备输入对
query_doc_pairs = [(query, doc.page_content) for doc in documents]
scores = self.reranker.predict(query_doc_pairs)
# 根据分数排序
doc_scores = list(zip(documents, scores))
doc_scores.sort(key=lambda x: x[1], reverse=True)
# 返回top_k个文档
return [doc for doc, score in doc_scores[:top_k]]
elif method == "cohere":
# 使用Cohere API进行重排
texts = [doc.page_content for doc in documents]
response = self.reranker.rerank(
model="rerank-english-v2.0",
query=query,
documents=texts,
top_k=top_k
)
# 根据重排结果重新排序
reranked_docs = []
for result in response.results:
reranked_docs.append(documents[result.index])
return reranked_docs
elif method == "bge":
# 使用BGE重排
query_doc_pairs = [[query, doc.page_content] for doc in documents]
scores = self.reranker.compute_score(query_doc_pairs)
# 处理单个文档的情况
if not isinstance(scores, list):
scores = [scores]
# 根据分数排序
doc_scores = list(zip(documents, scores))
doc_scores.sort(key=lambda x: x[1], reverse=True)
# 返回top_k个文档
return [doc for doc, score in doc_scores[:top_k]]
except Exception as e:
print(f"重排失败: {e}")
return documents
return documents
"""对检索到的文档进行重排"""
if not self.reranker or not documents:
return documents
method = self.rerank_config.get("method", "cross_encoder")
top_k = self.rerank_config.get("top_k", len(documents))
# 如果是相似度重排,直接调用相似度重排方法
if method == "similarity":
return self._similarity_rerank(query, documents)
try:
if method == "cross_encoder":
# 准备输入对
query_doc_pairs = [(query, doc.page_content) for doc in documents]
scores = self.reranker.predict(query_doc_pairs)
# 根据分数排序
doc_scores = list(zip(documents, scores))
doc_scores.sort(key=lambda x: x[1], reverse=True)
# 返回top_k个文档
return [doc for doc, score in doc_scores[:top_k]]
elif method == "cohere":
# 使用Cohere API进行重排
texts = [doc.page_content for doc in documents]
response = self.reranker.rerank(
model="rerank-english-v2.0",
query=query,
documents=texts,
top_k=top_k
)
# 根据重排结果重新排序
reranked_docs = []
for result in response.results:
reranked_docs.append(documents[result.index])
return reranked_docs
elif method == "bge":
# 使用BGE重排
query_doc_pairs = [[query, doc.page_content] for doc in documents]
scores = self.reranker.compute_score(query_doc_pairs)
# 处理单个文档的情况
if not isinstance(scores, list):
scores = [scores]
# 根据分数排序
doc_scores = list(zip(documents, scores))
doc_scores.sort(key=lambda x: x[1], reverse=True)
# 返回top_k个文档
return [doc for doc, score in doc_scores[:top_k]]
except Exception as e:
print(f"重排失败: {e}")
return documents
return documents
def _similarity_rerank(self, query: str, documents: List[Document]) -> List[Document]:
"""基于余弦相似度的简单重排(备选方案)"""
if not documents:
return documents
try:
# 获取查询向量
query_embedding = self.embedding_model.embed_query(query)
# 获取文档向量
doc_texts = [doc.page_content for doc in documents]
doc_embeddings = self.embedding_model.embed_documents(doc_texts)
# 计算余弦相似度
similarities = []
for doc_emb in doc_embeddings:
similarity = np.dot(query_embedding, doc_emb) / (
np.linalg.norm(query_embedding) * np.linalg.norm(doc_emb)
)
similarities.append(similarity)
# 根据相似度排序
doc_similarities = list(zip(documents, similarities))
doc_similarities.sort(key=lambda x: x[1], reverse=True)
top_k = self.rerank_config.get("top_k", len(documents))
return [doc for doc, sim in doc_similarities[:top_k]]
except Exception as e:
print(f"相似度重排失败: {e}")
return documents
def load_and_split_documents(self, file_path: str) -> List[Document]:
"""
加载并切分文档可被子类重写实现不同的切分方式
@ -191,6 +413,26 @@ class BaseRAG(ABC):
k = k or self.retriever_top_k
return self.vector_store.similarity_search(query, k=k)
def similarity_search_with_rerank(self, query: str, k: int = None) -> List[Document]:
"""
带重排功能的相似性搜索
"""
# 首先获取更多的候选文档用于重排
initial_k = k or self.retriever_top_k
if self.rerank_config.get("enabled", False):
# 获取更多候选文档进行重排
initial_k = max(initial_k * 2, 10)
documents = self.vector_store.similarity_search(query, k=initial_k)
# 如果启用了重排,进行重排
if self.rerank_config.get("enabled", False) and documents:
documents = self._rerank_documents(query, documents)
# 返回最终的top_k结果
final_k = k or self.retriever_top_k
return documents[:final_k]
@abstractmethod
def ingest(self, *args, **kwargs):
"""