从单一检索到混合策略
在传统 RAG 系统中,检索环节通常依赖向量相似度搜索。这种方式擅长捕捉语义相似性,但当用户查询包含精确术语、专有名词或特定格式时,纯语义检索往往会失效。
你有没有遇到过这样的场景:查询"IRR 计算公式",系统却返回了大量讨论"投资回报率概念"的文档,而真正包含公式细节的文档反而被排在后面?这是因为向量嵌入在编码时丢失了精确的关键词匹配能力。
解决这个问题需要一种更智能的检索策略——混合检索(Hybrid Retrieval),它同时利用关键词检索和语义检索的优势。
BM25 关键词检索的实现
BM25(Best Matching 25)是一种经典的关键词排序算法,它在信息检索领域被广泛使用。与简单的 TF-IDF 相比,BM25 引入了两个关键改进:
- 词频饱和:通过参数 k1 控制词频的影响上限,避免高频词过度主导
- 文档长度归一化:通过参数 b 平衡长短文档的偏好
class BM25:
"""BM25 关键词检索算法
Okapi BM25 算法实现,用于关键词检索。
"""
def __init__(
self,
k1: float = 1.5,
b: float = 0.75,
epsilon: float = 0.25,
):
self.k1 = k1
self.b = b
self.epsilon = epsilon
self.doc_freqs: dict[str, int] = {}
self.doc_len: list[int] = []
self.avgdl: float = 0
self.doc_count: int = 0
self.doc_term_freqs: list[dict[str, int]] = []
self.idf: dict[str, float] = {}
self.documents: list[Document] = []
def _tokenize(self, text: str) -> list[str]:
"""分词"""
text = text.lower()
tokens = re.findall(r"\w+", text)
return tokens
def fit(self, documents: list[Document]):
"""训练 BM25 模型"""
self.documents = documents
self.doc_count = len(documents)
self.doc_len = []
self.doc_term_freqs = []
self.doc_freqs = {}
for doc in documents:
tokens = self._tokenize(doc.page_content)
self.doc_len.append(len(tokens))
term_freqs = Counter(tokens)
self.doc_term_freqs.append(dict(term_freqs))
for term in term_freqs:
if term not in self.doc_freqs:
self.doc_freqs[term] = 0
self.doc_freqs[term] += 1
self.avgdl = sum(self.doc_len) / self.doc_count if self.doc_count > 0 else 0
self._calc_idf()
这个实现的核心在于 fit 方法:它遍历所有文档,统计词频(term frequency),同时计算逆文档频率(IDF)。IDF 的计算采用了 Okapi 公式的标准形式:
def _calc_idf(self):
"""计算 IDF"""
idf_sum = 0
negative_idfs = []
for term, freq in self.doc_freqs.items():
idf = math.log(self.doc_count - freq + 0.5) - math.log(freq + 0.5)
self.idf[term] = idf
idf_sum += idf
if idf < 0:
negative_idfs.append(term)
avg_idf = idf_sum / len(self.idf) if self.idf else 0
eps = self.epsilon * avg_idf
for term in negative_idfs:
self.idf[term] = eps
注意这里对负 IDF 值的处理:当某个词出现在几乎所有文档中时,其 IDF 可能为负。作者使用了 epsilon 参数作为下限,避免这些通用词对评分产生负面影响。
三种融合策略
混合检索的关键挑战在于如何融合关键词检索和语义检索的结果。HybridRetriever 类提供了三种不同的融合策略:
加权融合(Weighted Fusion)
这是最直观的方法:分别对两种检索结果进行归一化,然后按权重相加。
def _weighted_fusion(
self,
keyword_results: list[tuple[Document, float]],
semantic_results: list[tuple[Document, float]],
k: int,
keyword_weight: float,
semantic_weight: float,
) -> list[tuple[Document, float]]:
"""加权融合"""
doc_scores: dict[str, tuple[Document, float]] = {}
max_kw_score = max((score for _, score in keyword_results), default=1.0) or 1.0
for doc, score in keyword_results:
doc_id = self._get_doc_id(doc)
normalized_score = score / max_kw_score
doc_scores[doc_id] = (doc, normalized_score * keyword_weight)
max_sem_score = max((score for _, score in semantic_results), default=1.0) or 1.0
for doc, score in semantic_results:
doc_id = self._get_doc_id(doc)
normalized_score = 1 - (score / max_sem_score)
if doc_id in doc_scores:
existing_doc, existing_score = doc_scores[doc_id]
doc_scores[doc_id] = (
existing_doc,
existing_score + normalized_score * semantic_weight,
)
else:
doc_scores[doc_id] = (doc, normalized_score * semantic_weight)
results = sorted(doc_scores.values(), key=lambda x: x[1], reverse=True)
return results[:k]
关键点在于语义检索返回的是距离分数(越小越相似),因此需要转换为相似度分数(1 - score/max)。
倒数排名融合(RRF)
RRF 是一种不依赖于具体分数的融合方法,它只关注文档在两个列表中的排名位置。这种方法更加稳健,因为不同检索器的分数尺度可能完全不同。
def _rrf_fusion(
self,
keyword_results: list[tuple[Document, float]],
semantic_results: list[tuple[Document, float]],
k: int,
) -> list[tuple[Document, float]]:
"""Reciprocal Rank Fusion"""
doc_scores: dict[str, tuple[Document, float]] = {}
for rank, (doc, _) in enumerate(keyword_results):
doc_id = self._get_doc_id(doc)
rrf_score = 1 / (self.rrf_k + rank + 1)
doc_scores[doc_id] = (doc, rrf_score)
for rank, (doc, _) in enumerate(semantic_results):
doc_id = self._get_doc_id(doc)
rrf_score = 1 / (self.rrf_k + rank + 1)
if doc_id in doc_scores:
existing_doc, existing_score = doc_scores[doc_id]
doc_scores[doc_id] = (existing_doc, existing_score + rrf_score)
else:
doc_scores[doc_id] = (doc, rrf_score)
results = sorted(doc_scores.values(), key=lambda x: x[1], reverse=True)
return results[:k]
RRF 的核心公式是 1 / (k + rank),其中 k 是一个调节参数(默认为 60)。这个公式的特点是:
- 排名越高(rank 越小),得分越高
- 两个列表中都出现的高排名文档会得到显著更高的综合得分
- 不需要对原始分数做任何归一化处理
分数归一化融合
这是加权融合的变体,但它使用 Min-Max 归一化来处理分数分布不一致的问题。
def _score_fusion(
self,
keyword_results: list[tuple[Document, float]],
semantic_results: list[tuple[Document, float]],
k: int,
keyword_weight: float,
semantic_weight: float,
) -> list[tuple[Document, float]]:
"""分数归一化融合"""
doc_scores: dict[str, tuple[Document, float]] = {}
kw_scores = [score for _, score in keyword_results]
kw_min = min(kw_scores) if kw_scores else 0
kw_max = max(kw_scores) if kw_scores else 1
kw_range = kw_max - kw_min if kw_max != kw_min else 1
for doc, score in keyword_results:
doc_id = self._get_doc_id(doc)
normalized = (score - kw_min) / kw_range if kw_range > 0 else 0.5
doc_scores[doc_id] = (doc, normalized * keyword_weight)
sem_scores = [score for _, score in semantic_results]
sem_min = min(sem_scores) if sem_scores else 0
sem_max = max(sem_scores) if sem_scores else 1
sem_range = sem_max - sem_min if sem_max != sem_min else 1
for doc, score in semantic_results:
doc_id = self._get_doc_id(doc)
normalized = 1 - (score - sem_min) / sem_range if sem_range > 0 else 0.5
if doc_id in doc_scores:
existing_doc, existing_score = doc_scores[doc_id]
doc_scores[doc_id] = (
existing_doc,
existing_score + normalized * semantic_weight,
)
else:
doc_scores[doc_id] = (doc, normalized * semantic_weight)
results = sorted(doc_scores.values(), key=lambda x: x[1], reverse=True)
return results[:k]
这种方法的优点是对异常值不那么敏感,但缺点是如果某个检索结果中没有返回任何文档,归一化会退化为默认值 0.5。
混合 RAG 系统集成
HybridRAGSystem 类将混合检索集成到 RAG 流程中,提供了统一的接口:
class HybridRAGSystem:
"""混合 RAG 系统
结合关键词检索和语义检索的 RAG 系统。
"""
def __init__(
self,
rag_system,
fusion_method: str = "rrf",
keyword_weight: float = 0.3,
):
self.rag = rag_system
self.retriever = HybridRetriever(
fusion_method=fusion_method,
keyword_weight=keyword_weight,
)
self._fitted = False
def fit(self):
"""训练混合检索模型"""
if not self.rag.vector_store:
raise ValueError("RAG 系统的向量存储未初始化")
if hasattr(self.rag.vector_store, "similarity_search"):
docs = self.rag.vector_store.similarity_search("", k=1000)
self.retriever.fit(docs)
self._fitted = True
def generate_answer(
self,
query: str,
k: int = 3,
alpha: float | None = None,
max_context_length: int = 4000,
):
"""生成回答"""
results = self.retrieve(query, k, alpha)
relevant_docs = [doc for doc, _ in results]
if not relevant_docs:
return "抱歉,没有找到相关的文档信息。", []
prompt = self.rag.prompt_builder.build_qa_prompt(
query=query,
documents=relevant_docs,
max_context_length=max_context_length,
)
answer = self.rag.llm_integration.generate(prompt)
return answer, relevant_docs
fit 方法的设计很巧妙:它从向量存储中抽取一批文档(默认 1000 条)来训练 BM25 索引。这避免了需要对整个文档库进行全量索引的性能开销,同时保证了召回率。
源码导航
| 模块 | 文件 | 说明 |
|---|---|---|
| BM25 关键词检索 | hybrid_retriever.py | BM25 算法实现和混合检索器 |
| LLM 集成 | llm_integration.py | LLM 调用、重试、缓存、速率限制 |
| RAG 系统 | rag.py | 完整的 RAG 系统实现 |
快速开始
安装依赖:
pip install langchain-llm-toolkit
使用混合检索 RAG 系统:
from langchain_llm_toolkit.hybrid_retriever import HybridRAGSystem
from langchain_llm_toolkit.rag import RAGSystem
# 初始化基础 RAG 系统
rag = RAGSystem(vector_store_type="qdrant")
# 加载文档并创建向量存储
documents = rag.load_and_process_documents(["document.pdf"])
rag.create_vector_store(documents)
# 初始化混合 RAG 系统
hybrid_rag = HybridRAGSystem(rag_system=rag, fusion_method="rrf")
# 生成回答
answer, docs = hybrid_rag.generate_answer("什么是 IRR?")
print(answer)
你也可以直接在 RAGSystem 中使用混合检索方法:
# 使用内置的混合检索
relevant_docs = rag.retrieve_hybrid(query="IRR 计算公式", k=5, bm25_weight=0.3)
这种方法默认使用 30% 的 BM25 权重和 70% 的语义检索权重,你可以根据实际需求调整 bm25_weight 参数。
发表回复