diff --git a/RERANK_GUIDE.md b/RERANK_GUIDE.md new file mode 100644 index 0000000..3a37dc8 --- /dev/null +++ b/RERANK_GUIDE.md @@ -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) +- **描述**: 使用BGE(BAAI 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测试来选择最适合的重排方法 diff --git a/examples/quick_start.py b/examples/quick_start.py index 4a65706..0c4d000 100644 --- a/examples/quick_start.py +++ b/examples/quick_start.py @@ -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", + } + + # 重排配置 - 使用基于余弦相似度的重排 + rerank_config = { + "enabled": True, + "method": "similarity", # 使用相似度重排(无需额外依赖) + "top_k": 3 } - rag = SimpleRAG(embedding_config=config) - print("RAG初始化完成!") + 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__": diff --git a/examples/rerank_demo.py b/examples/rerank_demo.py new file mode 100644 index 0000000..873fde2 --- /dev/null +++ b/examples/rerank_demo.py @@ -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() diff --git a/src/base_rag/core.py b/src/base_rag/core.py index 86de225..d979486 100644 --- a/src/base_rag/core.py +++ b/src/base_rag/core.py @@ -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): """