通过结合密集嵌入和BM25构建高级本地LLM RAG管道

基本的检索增强生成（RAG）管道使用编码器模型在给定查询时搜索相似文档。

这也被称为语义搜索，因为编码器将文本转换为高维向量表示（称为嵌入），在该表示中，语义相似的文本彼此接近。

在我们拥有大型语言模型（LLMs）来创建这些向量嵌入之前，BM25算法是一个非常流行的搜索算法。BM25关注重要关键词，并在可用文档中寻找精确匹配。这种方法被称为关键词搜索。

如果您想将RAG管道提升到一个新的水平，您可能想尝试混合搜索。混合搜索结合了关键词搜索和语义搜索的优点，以提高搜索质量。

在本文中，我们将涵盖理论并在Python中实现这三种搜索方法。

RAG 检索

混合搜索是关键词搜索和语义搜索的结合。我们将分别介绍这两种搜索策略，然后再将它们结合起来。

关键词搜索与BM25

BM25是关键词搜索的首选算法。使用BM25，我们可以为语料库中的每个文档计算查询得分。

BM25基于TF-IDF算法，这意味着公式的核心是词频（TF）和逆文档频率（IDF）的乘积。

TF-IDF算法的基本思想是“在较少出现的、更具体的术语上的匹配比在常见术语上的匹配更有价值”[1]。

换句话说，TF-IDF算法寻找包含我们查询中稀有关键词的文档。

BM25算法有许多变体，每种变体都旨在改进原始算法。然而，似乎没有一种算法在系统上优于其他算法[2]。

因此，在实践中，选择一种并坚持使用是可以的。

如果我们查看LangChain的source code ，可以看到它使用了来自rank_bm25 包的BM25Okapi类，这是ATIRE BM25算法的一个稍微修改过的版本[3]。

在ATIRE BM25版本中，计算文档d和给定查询q（由多个术语t组成）得分的公式如下[2]：

N是语料库中的文档数量
df_t是包含术语t的文档数量（也称为文档频率）
tf_td是术语t在文档d中出现的次数（也称为词频）
L_d是文档的长度，L_avg是平均文档长度
有两个经验调优参数：b和k_1

直观上，我们看到公式对所有术语t求和，我们可以将其视为单词。

BM25方程中的左侧因子log(N/df_t)称为逆文档频率。对于像“the”这样的常见词，所有文档可能都包含，因此逆文档频率将为零（因为log(1)为零）。

另一方面，非常稀有的词只会出现在少数文档中，从而增加左侧因子。因此，逆文档频率是术语t所包含信息量的度量。

右侧因子受术语t在文档d中出现次数的影响。

文档d=["I like red cats, black cats, white cats, and brown cats"]对于术语t="cats"具有非常高的词频tf_td，这将导致包含“cats”一词的查询获得高BM25得分。

让我们使用BM25通过Python库rank_bm25来获取一些直观的理解。

pip install rank_bm25

首先，我们加载库并用我们的分词语料库初始化BM25。

from rank_bm25 import BM25Okapi

corpus = [
    "The cat, commonly referred to as the domestic cat or house cat, is a small domesticated carnivorous mammal.",
    "The dog is a domesticated descendant of the wolf.",
    "Humans are the most common and widespread species of primate, and the last surviving species of the genus Homo.",
    "The scientific name Felis catus was proposed by Carl Linnaeus in 1758"
]
tokenized_corpus = [doc.split(" ") for doc in corpus]

bm25 = BM25Okapi(tokenized_corpus)

接下来，我们对查询进行分词。

query = "The cat"
tokenized_query = query.split(" ")

最后，我们使用BM25算法计算得分。高得分表示文档与查询之间的良好匹配。

doc_scores = bm25.get_scores(tokenized_query)

print(doc_scores)
>> [0.92932018 0.21121974 0. 0.1901173]# 文档1、2、3和4的得分

由于BM25寻找精确的术语匹配，因此查询“cats”、“Cat”或“feline”都会导致我们的三个示例文档的得分为doc_scores = [0,0,0]。

使用密集嵌入进行语义搜索

当我们通过密集嵌入执行语义搜索时，我们将单词转换为数值表示。其理念是相似的单词在这种新的数学表示中彼此接近。

文本嵌入是单个单词或整个句子的高维向量。它们被称为密集，因为向量中的每个条目都是一个有意义的数字。相反，当许多向量条目为零时，称为sparse。

在将单词转换为嵌入之前，首先将它们转换为tokens，这是一种从字符串到整数的映射。一个称为编码器的神经网络嵌入模型随后将tokens转换为嵌入。

在将我们文档语料库中的所有文本转换为嵌入后，我们可以执行语义搜索，以查看哪个嵌入文档与我们的嵌入查询最接近。

我们可以通过绘制嵌入维度并找到与我们的查询最接近的文档匹配来可视化此任务。

在数学上，我们使用余弦距离函数找到最接近的匹配。对于两个嵌入向量a和b，我们可以使用点积计算余弦相似度，如下所示：

其中，分子是两个嵌入向量的点积，分母是它们的大小的乘积。

在几何上，余弦相似度是向量之间的角度。余弦相似度得分范围从-1到+1。

余弦相似度得分为-1意味着嵌入a和b完全朝相反方向，0意味着它们的角度为90度（它们无关），而+1意味着它们是相同的。因此，在将搜索查询与文档匹配时，我们寻找接近+1的值。

如果我们事先对嵌入进行归一化，则余弦相似度度量等同于点积相似度度量（分母变为1）。

让我们使用Python包sentence-transformers执行基本的语义搜索。

pip install sentence-transformers

首先，我们加载库并从HuggingFace下载all-MiniLM-L6-v2 编码器模型。该编码器模型经过训练以生成384维的密集嵌入。如果您有OpenAI API密钥，您也可以使用他们的text-embedding模型。

from sentence_transformers import SentenceTransformer

# 1. 加载预训练的句子变换器模型
model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')

然后，我们使用与之前相同的文档语料库。

# 要编码的文档
corpus = [
    "The cat, commonly referred to as the domestic cat or house cat, is a small domesticated carnivorous mammal.",
    "The dog is a domesticated descendant of the wolf.",
    "Humans are the most common and widespread species of primate, and the last surviving species of the genus Homo.",
    "The scientific name Felis catus was proposed by Carl Linnaeus in 1758"
]

# 通过调用 model.encode() 计算嵌入
document_embeddings = model.encode(corpus)

# 完整性检查
print(document_embeddings.shape)
>> (4, 384)

然后，我们嵌入我们的查询：

query = "The cat"
query_embedding = model.encode(query)

最后，我们可以计算余弦相似度得分。我们可以使用sentence_transformers中的实用函数cos_sim，而不是自己编写公式。

from sentence_transformers.util import cos_sim

# 计算文档与查询之间的余弦相似度
scores = cos_sim(document_embeddings, query_embedding)

print(scores)
>> tensor([[0.5716],  # 文档1的得分
>>         [0.2904],  # 文档2的得分
>>         [0.0942],  # 文档3的得分
>>         [0.3157]]) # 文档4的得分

为了展示密集嵌入语义搜索的强大功能，我可以用查询“feline”重新运行代码：

query_embedding = model.encode("feline")

scores = cos_sim(document_embeddings, query_embedding)

print(scores)
>> tensor([[0.4007],
>>         [0.3837],
>>         [0.0966],
>>         [0.3804]])

尽管“feline”这个词在文档语料库中没有出现，但语义搜索仍然将关于猫的文本排名为最高匹配。

语义搜索还是关键词搜索？

哪种搜索方法更好？这要看情况。两者都有优缺点。现在我们知道了它们的工作原理，我们可以看到它们在哪些地方有用，在哪些地方可能会失败。

关键词搜索使用BM25查找查询词的精确匹配。 当我们寻找短语的精确匹配时，这非常有用。

如果我在寻找“猫帽子”，我可能是在寻找这本书/电影。我不想要与帽子或猫语义相似的结果。

关键词搜索的另一个用例是编程。如果我在寻找特定的函数或代码片段，我希望得到一个精确匹配。

另一方面，语义搜索寻找语义相似的内容。 这意味着语义搜索还会找到包含同义词或不同拼写的文档，例如复数形式、大写字母等。

由于这两种算法都有其用例，混合搜索使用两者，然后将它们的结果合并为一个最终排名。

混合搜索的缺点是它比仅运行一种算法需要更多的计算资源。

混合搜索

我们可以使用倒数排名融合 (RRF) 将 BM25 和余弦相似度的结果结合起来。RRF 是一种简单的算法，用于结合不同评分函数的排名 [4]。

首先，我们需要为每个评分算法获取文档排名。在我们的示例中，这将是：

corpus = [
    "The cat, commonly referred to as the domestic cat or house cat, is a small domesticated carnivorous mammal.",
    "The dog is a domesticated descendant of the wolf.",
    "Humans are the most common and widespread species of primate, and the last surviving species of the genus Homo.",
    "The scientific name Felis catus was proposed by Carl Linnaeus in 1758",
]
query = "The cat"

bm25_ranking = [1, 2, 4, 3] # scores = [0.92932018 0.21121974 0. 0.1901173]
cosine_ranking = [1, 3, 4, 2] # scores = [0.5716, 0.2904, 0.0942, 0.3157]

每个文档 d 的组合 RRF 分数的公式如下：

其中 k 是一个参数（原始论文中使用 k=60），r(d) 是来自 BM25 和余弦相似度的排名。

整合所有内容

现在我们可以通过分别执行 BM25 和余弦相似度来实现我们的混合搜索，然后使用 RRF 组合结果。

首先，让我们定义 RRF 的函数和一个辅助函数，将浮动分数转换为整数排名。

import numpy as np

def scores_to_ranking(scores: list[float]) -> list[int]:
    """将浮动分数转换为整数排名（排名 1 为最佳）"""
    return np.argsort(scores)[::-1] + 1


def rrf(keyword_rank: int, semantic_rank: int) -> float:
    """将关键词排名和语义排名组合成混合分数。"""
    k = 60
    rrf_score = 1 / (k + keyword_rank) + 1 / (k + semantic_rank)
    return rrf_score

这是我使用上述概念实现的简单混合搜索。

from rank_bm25 import BM25Okapi
from sentence_transformers import SentenceTransformer
from sentence_transformers.util import cos_sim

model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

def hybrid_search(
    query: str, corpus: list[str], encoder_model: SentenceTransformer
) -> list[int]:
    # bm25
    tokenized_corpus = [doc.split(" ") for doc in corpus]
    tokenized_query = query.split(" ")
    bm25 = BM25Okapi(tokenized_corpus)
    bm25_scores = bm25.get_scores(tokenized_query)
    bm25_ranking = scores_to_ranking(bm25_scores)

    # embeddings
    document_embeddings = model.encode(corpus)
    query_embedding = model.encode(query)
    cos_sim_scores = cos_sim(document_embeddings, query_embedding).flatten().tolist()
    cos_sim_ranking = scores_to_ranking(cos_sim_scores)

    # combine rankings into RRF scores
    hybrid_scores = []
    for i, doc in enumerate(corpus):
        document_ranking = rrf(bm25_ranking[i], cos_sim_ranking[i])
        print(f"Document {i} has the rrf score {document_ranking}")
        hybrid_scores.append(document_ranking)

    # convert RRF scores into final rankings
    hybrid_ranking = scores_to_ranking(hybrid_scores)
    return hybrid_ranking

现在我们可以使用 hybrid_search 进行不同的查询。

hybrid_ranking = hybrid_search(
    query="What is the scientifc name for cats?", corpus=corpus, encoder_model=model
)
print(hybrid_ranking)
>> Document 0 has the rrf score 0.03125
>> Document 1 has the rrf score 0.032266458495966696
>> Document 2 has the rrf score 0.03225806451612903
>> Document 3 has the rrf score 0.032266458495966696
>> [4 2 3 1]

作为下一步，我们可以向我们的文档语料库添加更多知识。在我的文章如何使用重排序提高 LLM RAG 检索中，我将维基百科整合到了知识语料库中。

在混合搜索之上添加一个重排序器将进一步改善整体 RAG 流程。

结论

混合搜索结合了语义搜索和关键词搜索，以产生更好的整体搜索结果。

为了执行关键词搜索，我们实现了BM25算法，该算法寻找重要的关键词匹配。

对于语义搜索，我们使用了余弦相似度和一个预训练的编码器模型，该模型生成稠密的嵌入。

虽然混合搜索可以改善RAG检索，但它也需要比仅运行一种搜索算法更多的计算资源。

混合搜索是改善您的RAG管道的一个有趣的构建模块。试试看！

参考文献

[1] K. Spärck Jones, 术语特异性的统计解释及其在检索中的应用 (1972), 文献学杂志第28卷第1期

[2] A. Trotman, A. Puurula, 和 B. Burgess, BM25和语言模型的改进研究 (2014), ADCS ’14: 第19届澳大利亚文档计算研讨会论文集

[3] A. Trotman, X.-F. Jia, 和 M. Crane, 朝着高效且有效的搜索引擎迈进 (2012), SIGIR 2012 开源信息检索研讨会论文集

[4] G. V. Cormack, C. L. A. Clarke, 和 S. Büttcher, 互惠排名融合优于康多塞法和个体排名学习方法 (2009), 第32届国际ACM SIGIR信息检索研究与开发会议论文集

Barry's Home

如何使用混合搜索来改善 LLM RAG 检索