BM25 for Python：使用 BM25S⚡ 简化依赖并实现高性能

摘要：在这篇博文中，我们探讨了流行词法搜索库的优缺点，并介绍了 BM25S⚡，一个用于在 Python 中实现快速词法搜索的新库，与最流行的基于 Python 的库相比，其速度提高了 500 倍。我们将讨论如何将其与 🤗 Hugging Face Hub 结合使用。

如果您想了解更多技术细节，请查看技术报告！

什么是词法搜索？

词法搜索 是一种搜索类型，通过直接匹配查询中的单词，将查询与文本文档集合（短至一个句子，长至整篇文章）进行比较。例如，给定一个查询，如“什么是猫？”，您将查找包含一个或多个单词的文档：“什么”、“是”、“a”、“猫”。您可以根据每个单词的重要性（通过为每个单词分配分数）以及匹配的单词数量（例如通过将分数相加）为每个文档分配分数。

词法搜索最著名的算法是 BM25，该算法由Stephen Robertson 及其合作者在 1990 年代开发。它被用作高度可扩展的商业搜索引擎（如 Elasticsearch，其 Python 客户端每月下载量达 2400 万）的骨干，并已在 Python 中实现（rank-bm25），并在那里也得到了广泛使用。

Elasticsearch 和 Rank-BM25 的构建目标不同。

Elasticsearch 能够实现高性能和可扩展性，通过利用多节点设置，让您能够实时搜索数十亿文档。但是，要开始使用，您需要设置一个运行 Apache Lucene（用 Java 编写）的 Web 服务器，并通过单独的 Web 客户端（可以使用 pip 安装的 elasticsearch 库）访问它。

另一方面，Rank-BM25 在 Python 中更容易访问，因为它只依赖于 Numpy，并且只需使用 pip 安装，避免了设置服务器的麻烦。此外，其实现相当易读，您可以使用自己的分词器。但是，当您将其用于更大的语料库（例如超过 100 万个文档）时，它无法在查询/秒方面实现相同的性能。

隆重推出 BM25S ⚡：一个快速但低依赖的 BM25 实现

为了在不依赖外部依赖的情况下实现高性能，我最近推出了 BM25S，这是一个新的词法搜索库，它利用 scipy 的稀疏矩阵来实现比 Rank-BM25 快几个数量级的搜索，同时仍保留在 Python 生态系统中（这意味着您可以使用简单的 pip install bm25s 进行安装）。

在单节点、单线程和多核设置（4 个线程）下，它实现了与 Elasticsearch 相当的性能（在来自 BEIR 的 5 个流行数据集上进行基准测试）

您可以在以下链接了解更多信息

• 仓库：github.com/xhluca/bm25s
• 主页：bm25s.github.io
• Huggingface：BEIR 数据集索引

但是，您只需此代码片段即可开始使用

# pip install bm25s
import bm25s

# Create your corpus here
corpus = [
    "a cat is a feline and likes to purr",
    "a dog is the human's best friend and loves to play",
    "a bird is a beautiful animal that can fly",
    "a fish is a creature that lives in water and swims",
]

# Create the BM25 model and index the corpus
retriever = bm25s.BM25(corpus=corpus)
retriever.index(bm25s.tokenize(corpus))

# Query the corpus and get top-k results
query = "does the fish purr like a cat?"
results, scores = retriever.retrieve(bm25s.tokenize(query), k=2)

# Let's see what we got!
doc, score = results[0, 0], scores[0, 0]
print(f"Rank {i+1} (score: {score:.2f}): {doc}")

如果您想要更高级的示例，请查看 README，包括词干提取、停用词和自定义。

🤗 Hugging Face 集成

BM25S 与 Huggingface Hub 紧密集成，让您能够轻松地保存到 Hub 并从 Hub 加载。首先，只需安装库

pip install huggingface_hub

现在，您可以创建索引并使用 BM25HF.save_to_hub 进行保存

import bm25s
from bm25s.hf import BM25HF

# Create your corpus here
corpus = [
    "a cat is a feline and likes to purr",
    "a dog is the human's best friend and loves to play",
    "a bird is a beautiful animal that can fly",
    "a fish is a creature that lives in water and swims",
]

retriever = BM25HF(corpus=corpus)
retriever.index(bm25s.tokenize(corpus))

# Set your username and token
user = "your-username"
retriever.save_to_hub(f"{user}/bm25s-animals")

如果您访问 huggingface.co/{user}/bm25s-animals，您将找到您新上传的 BM25S 索引！

现在，要使用它，只需使用 BM25HF.load_from_hub 加载它

import bm25s
from bm25s.hf import BM25HF

# Load the index
retriever = BM25HF.load_from_hub(f"{user}/bm25s-animals", load_corpus=True)

# You can retrieve now
query = "a cat is a feline"

docs, scores = retriever.retrieve(bm25s.tokenize(query), k=2)

选择 BM25 变体

我非常喜欢一篇关于 BM25 的最新论文，它是 Chris Kamphuis 及合著者撰写的《你指的是哪种 BM25？一项大规模评分变体复现性研究》：它简洁而正式地展示了 BM25 的多种流行变体，并倡导研究人员明确定义他们实现的变体。

根据该论文，BM25S 提供以下变体

• 原始（method="robertson"）
• ATIRE (method="atire")
• BM25L (method="bm25l")
• BM25+ (method="bm25+")
• Lucene (method="lucene") - 默认

您只需指定 BM25(method="") 即可更改此设置。

它是如何实现加速的？

更精确地说，它计算语料库中每个文档所有可能的词级相关性分数，并将它们存储在一个稀疏矩阵中（这个想法受到 Jack Morris 的 bm25-pt 启发）。然后，给定一个查询，您可以将相关词元求和以获得每个文档的相关性分数——这通过简单的 scipy.sparse 操作实现，也可以在 Numpy 中实现（允许我们利用内存映射数组）。

对于没有稀疏矩阵的 BM25 变体，我们将分数偏移一个常数（可以根据单词未出现在文档中时分配的相关性分数轻松计算）；这使我们能够将偏移后的分数存储在稀疏矩阵中，并在检索过程中重新计算精确的相关性分数。

BM25S 会取代其他库吗？

BM25S 并非取代其他实现，而是对其进行补充：Elasticsearch 能够扩展到多个节点，可以用 Java 接口，并允许增量索引更新；Rank-BM25 允许您在搜索过程中更改参数；Pyserini 包含密集和混合搜索算法以及许多可复现性工具；BM25-PT 可以通过 PyTorch 在 GPU 上使用。因此，尝试多种实现并选择最适合您需求的一种非常重要。