feat: rerank

This commit is contained in:
李如威 2025-07-29 09:43:55 +08:00
parent a0cc4aab3a
commit b2fcd4a69c
5 changed files with 364 additions and 162 deletions

171
README.md
View File

@ -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)

View File

@ -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测试来选择最适合的重排方法

123
examples/local_api_demo.py Normal file
View File

@ -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()

View File

@ -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
sentence-transformers>=2.2.0
numpy>=1.21.0
# 可选依赖
# langchain-openai>=0.2.0 # 用于本地API接口支持
# FlagEmbedding>=1.2.0 # 用于BGE重排

View File

@ -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]:
"""基于余弦相似度的简单重排(备选方案)"""