diff --git a/README.md b/README.md index ec8e45e..2ae4da3 100644 --- a/README.md +++ b/README.md @@ -1,41 +1,162 @@ -# Base RAG +# BaseRAG 系统功能总结 -简洁的RAG基础库,支持多种embedding模型和Chroma向量数据库。 +## 概述 +BaseRAG 是一个灵活的检索增强生成(RAG)框架,支持多种嵌入模型和重排策略,专注于本地部署和HuggingFace生态系统。 -## 安装 +## 🎯 核心特性 -```bash -pip install base-rag -``` +### 1. 多种嵌入模型支持 +- **本地HuggingFace模型**: 支持模型名称和本地路径两种方式 +- **本地API接口**: 兼容OpenAI API格式的本地嵌入服务 +- **自动回退机制**: API不可用时自动切换到本地模型 -## 使用 +### 2. 文档重排功能 +- **相似度重排**: 基于余弦相似度,无额外依赖 +- **CrossEncoder重排**: 专业重排模型,效果优秀 +- **BGE重排**: 中文支持良好的重排模型 +### 3. 向量存储 +- **Chroma数据库**: 自动持久化,支持多集合管理 +- **线程安全**: 模型缓存和并发访问保护 + +## 🔧 配置示例 + +### 基础配置 ```python from base_rag import BaseRAG class MyRAG(BaseRAG): - def ingest(self, file_path: str): - documents = self.load_and_split_documents(file_path) - self.add_documents_to_vector_store(documents) + def ingest(self, documents): + self.vector_store.add_texts(documents) - def query(self, question: str) -> str: - docs = self.similarity_search(question) - return f"找到 {len(docs)} 个相关文档" + def query(self, question, k=3): + return self.similarity_search_with_rerank(question, k=k) -# OpenAI API -config = { - "type": "openai", - "model": "text-embedding-3-small", - "api_key": "your-api-key" -} - -# 本地模型 -config = { +# 配置 +embedding_config = { "type": "local", "model_name": "sentence-transformers/all-MiniLM-L6-v2" } -rag = MyRAG(embedding_config=config) -rag.ingest("document.txt") -result = rag.query("问题") +rerank_config = { + "enabled": True, + "method": "similarity", + "top_k": 3 +} + +rag = MyRAG( + embedding_config=embedding_config, + rerank_config=rerank_config +) ``` + +### 本地API配置 +```python +embedding_config = { + "type": "api", + "api_url": "http://localhost:8080", + "model": "text-embedding-model" +} +``` + +### 本地模型路径配置 +```python +embedding_config = { + "type": "local", + "model_path": "/path/to/your/model" +} +``` + +## 📁 项目结构 +``` +base_rag/ +├── src/ +│ └── base_rag/ +│ ├── __init__.py # 包入口 +│ └── core.py # 核心BaseRAG类 +├── examples/ +│ ├── quick_start.py # 快速开始示例 +│ ├── rerank_demo.py # 重排功能演示 +│ └── local_api_demo.py # 本地API配置示例 +├── requirements.txt # 依赖列表 +├── pyproject.toml # 包配置 +├── RERANK_GUIDE.md # 重排功能详细说明 +└── README.md # 项目说明 +``` + +## 🚀 快速开始 + +1. **安装依赖** +```bash +pip install -r requirements.txt +``` + +2. **运行示例** +```bash +# 基础功能演示 +python examples/quick_start.py + +# 重排功能演示 +python examples/rerank_demo.py + +# 本地API配置演示 +python examples/local_api_demo.py +``` + +## 📦 可选依赖 + +```bash +# 本地API接口支持 +pip install langchain-openai + +# BGE重排支持 +pip install FlagEmbedding +``` + +## 🔍 API参考 + +### 核心方法 +- `similarity_search(query, k)`: 基础相似性搜索 +- `similarity_search_with_rerank(query, k)`: 带重排的搜索 +- `load_and_split_documents(file_path)`: 文档加载和分割 +- `add_documents_to_vector_store(documents)`: 添加文档到向量库 + +### 抽象方法(需实现) +- `ingest(*args, **kwargs)`: 文档导入逻辑 +- `query(question)`: 问答逻辑 + +## 💡 最佳实践 + +1. **开发阶段**: 使用本地模型,快速原型验证 +2. **生产环境**: + - 小规模: 本地模型 + 相似度重排 + - 大规模: 本地API + CrossEncoder/BGE重排 +3. **模型选择**: 根据语言和领域需求选择合适的嵌入模型 +4. **重排策略**: 在实际数据上测试不同重排方法的效果 + +## 🛠️ 技术特点 + +- **线程安全**: 支持并发访问和模型缓存 +- **错误处理**: 完善的异常处理和回退机制 +- **灵活配置**: 支持多种配置方式和自定义参数 +- **易于扩展**: 抽象设计,便于子类实现特定业务逻辑 + +## 📋 注意事项 + +1. 首次运行会下载模型,需要网络连接 +2. 重排功能会增加查询延迟,但提高结果质量 +3. 不同模型对硬件要求不同,请根据实际情况选择 +4. 建议在生产环境前进行充分测试 + +## 🔄 版本信息 + +- **当前版本**: 0.1.0 +- **Python要求**: >= 3.8 +- **主要依赖**: langchain, chromadb, sentence-transformers + +--- + +更多详细信息请参考: +- [重排功能详细说明](RERANK_GUIDE.md) +- [示例代码](examples/) +- [配置文件](pyproject.toml) diff --git a/RERANK_GUIDE.md b/RERANK_GUIDE.md index 3a37dc8..13f8101 100644 --- a/RERANK_GUIDE.md +++ b/RERANK_GUIDE.md @@ -23,11 +23,16 @@ BaseRAG现在支持对检索到的文档片段进行重排(rerank),这是R - **优点**: 中文支持较好,性能优秀 - **默认模型**: `BAAI/bge-reranker-base` -### 4. Cohere重排 (cohere) -- **描述**: 使用Cohere API进行重排 -- **依赖**: `cohere` -- **优点**: 云端API,无需本地模型 -- **注意**: 需要API密钥 +## 支持的嵌入模型类型 + +### 1. 本地HuggingFace模型 (local) +- **模型名称**: 从HuggingFace Hub下载 +- **本地路径**: 使用已下载的本地模型 + +### 2. 本地部署API接口 (api) +- **描述**: 连接到本地部署的嵌入服务 +- **兼容**: OpenAI API格式的本地服务 +- **回退机制**: API不可用时自动回退到本地模型 ## 使用方法 @@ -36,10 +41,16 @@ BaseRAG现在支持对检索到的文档片段进行重排(rerank),这是R ```python 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", "cohere" + "method": "similarity", # 或 "cross_encoder", "bge" "top_k": 3 # 重排后返回的文档数量 } @@ -50,6 +61,33 @@ rag = MyRAG( ) ``` +### 本地模型配置 + +```python +# 使用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接口配置 + +```python +embedding_config = { + "type": "api", + "api_url": "http://localhost:8000", # 本地API地址 + "model": "your-embedding-model", + "api_key": "optional-key" # 可选 +} +``` + ### CrossEncoder重排配置 ```python @@ -72,17 +110,6 @@ rerank_config = { } ``` -### 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)` @@ -113,27 +140,30 @@ docs = rag.similarity_search("查询问题", k=5) - **相似度重排**: 计算开销小,速度快 - **CrossEncoder重排**: 需要加载额外模型,计算较慢但效果好 - **BGE重排**: 平衡了效果和性能 -- **Cohere重排**: 网络延迟,但无需本地计算资源 +- **本地API**: 网络延迟较小,可扩展性好 ## 最佳实践 1. **开发阶段**: 使用相似度重排进行快速原型验证 2. **生产环境**: 根据具体需求选择CrossEncoder或BGE重排 -3. **云端部署**: 考虑使用Cohere API重排 +3. **分布式部署**: 考虑使用本地API接口 4. **候选文档数量**: 建议设置为最终需要文档数量的2-3倍 5. **模型选择**: 根据语言和领域选择合适的重排模型 ## 依赖安装 ```bash +# 基本依赖 +pip install langchain-huggingface langchain-chroma + # CrossEncoder重排 pip install sentence-transformers # BGE重排 pip install FlagEmbedding -# Cohere重排 -pip install cohere +# 本地API支持(可选) +pip install langchain-openai ``` ## 注意事项 @@ -141,4 +171,5 @@ pip install cohere - 重排会增加查询延迟,但通常能显著提高结果质量 - 如果重排模型加载失败,系统会自动回退到相似度重排 - 不同重排方法的效果可能因数据集和查询类型而异 +- 本地API接口需要兼容OpenAI API格式 - 建议在实际数据上进行A/B测试来选择最适合的重排方法 diff --git a/examples/local_api_demo.py b/examples/local_api_demo.py new file mode 100644 index 0000000..3f49826 --- /dev/null +++ b/examples/local_api_demo.py @@ -0,0 +1,123 @@ +""" +BaseRAG 本地API接口使用示例 + +这个示例展示了如何配置BaseRAG使用本地部署的嵌入API接口, +以及当API不可用时如何自动回退到本地模型。 +""" + +import sys +import os + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) + +from base_rag import BaseRAG + + +class LocalAPIRAG(BaseRAG): + def ingest(self, documents): + """批量添加文档""" + if documents: + self.vector_store.add_texts(documents) + print(f"已向向量库添加 {len(documents)} 个文档") + + def query(self, question, k=3): + """查询文档""" + return self.similarity_search_with_rerank(question, k=k) + + +def demo_local_api(): + """演示本地API配置""" + print("=== 本地API接口配置示例 ===\n") + + # 本地API配置(假设有本地嵌入服务) + api_embedding_config = { + "type": "api", + "api_url": "http://localhost:8080", # 假设的本地API地址 + "model": "text-embedding-model", + "api_key": "optional-key" # 可选 + } + + print("正在尝试连接本地API...") + try: + rag_api = LocalAPIRAG( + vector_store_name="api_test", + embedding_config=api_embedding_config, + rerank_config={"enabled": True, "method": "similarity", "top_k": 3} + ) + print("本地API连接成功!") + except Exception as e: + print(f"本地API连接失败: {e}") + print("系统会自动回退到本地模型") + + +def demo_local_model(): + """演示本地模型配置""" + print("\n=== 本地模型配置示例 ===\n") + + # 本地模型配置 + local_embedding_config = { + "type": "local", + "model_name": "sentence-transformers/all-MiniLM-L6-v2" + } + + rag_local = LocalAPIRAG( + vector_store_name="local_test", + embedding_config=local_embedding_config, + rerank_config={"enabled": True, "method": "similarity", "top_k": 3} + ) + + # 测试文档 + test_documents = [ + "Python是一种高级编程语言,语法简洁明了。", + "机器学习是人工智能的一个重要分支。", + "深度学习使用神经网络来模拟人脑的学习过程。", + "自然语言处理帮助计算机理解和生成人类语言。" + ] + + print("正在添加测试文档...") + rag_local.ingest(test_documents) + + # 测试查询 + query = "什么是机器学习?" + print(f"\n查询: {query}") + + results = rag_local.query(query, k=2) + print("查询结果:") + for i, doc in enumerate(results, 1): + print(f" {i}. {doc.page_content}") + + +def demo_local_path(): + """演示使用本地模型路径的配置""" + print("\n=== 本地模型路径配置示例 ===\n") + + # 假设你有本地下载的模型 + local_path_config = { + "type": "local", + "model_path": "/path/to/your/local/model", # 替换为实际路径 + "model_kwargs": {"device": "cpu"} + } + + print("本地模型路径配置:") + print(f" 路径: {local_path_config['model_path']}") + print(" 注意: 请确保路径存在且包含有效的sentence-transformers模型") + + +def main(): + """主函数""" + print("BaseRAG 本地API和模型配置示例\n") + + # 演示不同的配置方式 + demo_local_api() + demo_local_model() + demo_local_path() + + print("\n=== 配置建议 ===") + print("1. 开发测试: 使用本地模型,快速启动") + print("2. 生产环境: 使用本地API接口,便于扩展和管理") + print("3. 离线部署: 使用本地模型路径,无需网络连接") + print("4. 混合部署: API主用,本地模型备用") + + +if __name__ == "__main__": + main() diff --git a/requirements.txt b/requirements.txt index ab43e84..563529d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,9 +1,11 @@ langchain>=0.3.0 langchain-community>=0.3.0 -langchain-openai>=0.2.0 langchain-chroma>=0.1.0 langchain-huggingface>=0.1.0 chromadb>=0.4.0 -openai>=1.0.0 -tiktoken>=0.5.0 -sentence-transformers>=2.2.0 \ No newline at end of file +sentence-transformers>=2.2.0 +numpy>=1.21.0 + +# 可选依赖 +# langchain-openai>=0.2.0 # 用于本地API接口支持 +# FlagEmbedding>=1.2.0 # 用于BGE重排 \ No newline at end of file diff --git a/src/base_rag/core.py b/src/base_rag/core.py index d979486..a4655dd 100644 --- a/src/base_rag/core.py +++ b/src/base_rag/core.py @@ -4,7 +4,6 @@ import threading import numpy as np from langchain_huggingface import HuggingFaceEmbeddings -from langchain_openai import OpenAIEmbeddings from langchain.embeddings.base import Embeddings from langchain_chroma import Chroma from langchain_community.document_loaders import TextLoader @@ -41,11 +40,12 @@ class BaseRAG(ABC): 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-..."} + 本地部署接口: {"type": "api", "api_url": "http://localhost:8000/embeddings", "model": "your-model"} 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} + {"enabled": True, "method": "bge", "model": "BAAI/bge-reranker-base", "top_k": 3} + {"enabled": True, "method": "similarity", "top_k": 3} """ self.vector_store_name = vector_store_name self.embedding_config = embedding_config or { @@ -87,8 +87,8 @@ class BaseRAG(ABC): return f"local_path_{config['model_path'].replace('/', '_').replace('\\', '_')}" else: return f"local_name_{config.get('model_name', 'default')}" - elif config_type == "openai": - return f"openai_{config.get('model', 'text-embedding-ada-002')}" + elif config_type == "api": + return f"api_{config.get('api_url', 'default').replace('/', '_').replace(':', '_')}" else: return f"{config_type}_{config.get('model', 'default')}" @@ -147,19 +147,36 @@ class BaseRAG(ABC): ), ) - elif config_type == "openai": - from langchain_openai import OpenAIEmbeddings - - return OpenAIEmbeddings( - model=config.get("model", "text-embedding-3-small"), - api_key=config.get("api_key"), - base_url=config.get("api_base"), - max_retries=config.get("max_retries", 3), - ) + elif config_type == "api": + # 本地部署的嵌入API接口 + try: + from langchain_openai import OpenAIEmbeddings + + api_url = config.get("api_url") + if not api_url: + raise ValueError("使用API类型时必须提供api_url") + + print(f"连接到本地嵌入API: {api_url}") + return OpenAIEmbeddings( + model=config.get("model", "text-embedding"), + base_url=api_url, + api_key=config.get("api_key", "dummy"), # 本地API可能不需要密钥 + max_retries=config.get("max_retries", 3), + ) + except ImportError: + print("警告: langchain_openai未安装,无法使用API接口") + # 回退到本地模型 + model_name = config.get("model", "sentence-transformers/all-MiniLM-L6-v2") + print(f"回退到本地模型: {model_name}") + return HuggingFaceEmbeddings( + model_name=model_name, + model_kwargs=config.get("model_kwargs", {"device": "cpu"}), + encode_kwargs=config.get("encode_kwargs", {"normalize_embeddings": True}), + ) else: raise ValueError( - f"不支持的嵌入模型类型: {config_type},支持的类型: 'local', 'openai'" + f"不支持的嵌入模型类型: {config_type},支持的类型: 'local', 'api'" ) def _init_reranker(self): @@ -181,19 +198,6 @@ class BaseRAG(ABC): 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 @@ -205,8 +209,8 @@ class BaseRAG(ABC): return None else: - print(f"警告: 不支持的重排方法: {method}") - return None + print(f"警告: 不支持的重排方法: {method},将使用相似度重排") + return "similarity" def _rerank_documents(self, query: str, documents: List[Document]) -> List[Document]: """对检索到的文档进行重排""" @@ -221,7 +225,7 @@ class BaseRAG(ABC): return self._similarity_rerank(query, documents) # 其他方法需要reranker模型 - if not self.reranker: + if not self.reranker or self.reranker == "similarity": print(f"重排模型未初始化,使用默认相似度重排") return self._similarity_rerank(query, documents) @@ -238,23 +242,6 @@ class BaseRAG(ABC): # 返回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] @@ -272,72 +259,10 @@ class BaseRAG(ABC): 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": + print(f"重排失败: {e},回退到相似度重排") 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 + return self._similarity_rerank(query, documents) def _similarity_rerank(self, query: str, documents: List[Document]) -> List[Document]: """基于余弦相似度的简单重排(备选方案)"""