欢迎使用bm25s，这是一个在Python中实现BM25的库，允许您根据查询对文档进行排序。BM25是一种广泛使用的用于文本检索任务的排名函数，是像Elasticsearch这样的搜索服务的关键组件。

它旨在是

• 快速: bm25s是用纯Python实现的，并利用Scipy稀疏矩阵存储所有文档标记的即时计算得分。这允许查询时具有极快的评分，比流行的库快几个数量级（见下文基准测试）。
• 简单: bm25s被设计成易于使用和理解。您可以使用pip安装它，并在几分钟内开始使用。它不依赖于Java或Pytorch - 您需要的只是Scipy和Numpy，以及可选的轻量级依赖项用于词干提取。

以下，我们比较了 bm25s 与 Elasticsearch 在相对于最流行的 BM25 Python 实现 rank-bm25 的加速方面。我们测量了在单个线程环境下，来自 BEIR 的几个流行数据集上的吞吐量，即每秒查询次数（QPS）。

@misc{bm25s,
      title={BM25S: Orders of magnitude faster lexical search via eager sparse scoring}, 
      author={Xing Han Lù},
      year={2024},
      eprint={2407.03618},
      archivePrefix={arXiv},
      primaryClass={cs.IR},
      url={https://arxiv.org/abs/2407.03618}, 
}

[!IMPORTANT] 新增于版本 0.2.0：我们正在推出对 numba 后端的支持，这为更大的数据集提供了大约 2 倍的加速！了解更多信息并在 版本 0.2.0 发布线程 中分享您的想法。

安装

您可以使用 pip 安装 bm25s

pip install bm25s

如果您想使用词干提取以获得更好的结果，您可以安装推荐的（但可选的）依赖项

# Install all extra dependencies
pip install bm25s[full]

# If you want to use stemming for better results, you can install a stemmer
pip install PyStemmer

# To speed up the top-k selection process, you can install `jax`
pip install jax[cpu]

快速入门

以下是如何使用 bm25s 的简单示例

import bm25s
import Stemmer  # optional: for stemming

# 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",
]

# optional: create a stemmer
stemmer = Stemmer.Stemmer("english")

# Tokenize the corpus and only keep the ids (faster and saves memory)
corpus_tokens = bm25s.tokenize(corpus, stopwords="en", stemmer=stemmer)

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

# Query the corpus
query = "does the fish purr like a cat?"
query_tokens = bm25s.tokenize(query, stemmer=stemmer)

# Get top-k results as a tuple of (doc ids, scores). Both are arrays of shape (n_queries, k)
results, scores = retriever.retrieve(query_tokens, corpus=corpus, k=2)

for i in range(results.shape[1]):
    doc, score = results[0, i], scores[0, i]
    print(f"Rank {i+1} (score: {score:.2f}): {doc}")

# You can save the arrays to a directory...
retriever.save("animal_index_bm25")

# You can save the corpus along with the model
retriever.save("animal_index_bm25", corpus=corpus)

# ...and load them when you need them
import bm25s
reloaded_retriever = bm25s.BM25.load("animal_index_bm25", load_corpus=True)
# set load_corpus=False if you don't need the corpus

有关如何快速索引 2M 文档语料库（自然问题）的示例，请参阅 examples/index_nq.py。

灵活性

bm25s 提供了一个灵活的 API，允许您自定义 BM25 模型和分词过程。以下是您可以使用的选项之一

# You can provide a list of queries instead of a single query
queries = ["What is a cat?", "is the bird a dog?"]

# Provide your own stopwords list if you don't like the default one
stopwords = ["a", "the"]

# For stemming, use any function that is callable on each word list
stemmer_fn = lambda lst: [word for word in lst]

# Tokenize the queries
query_token_ids = bm25s.tokenize(queries, stopwords=stopwords, stemmer=stemmer_fn)

# If you want the tokenizer to return strings instead of token ids, you can do this
query_token_strs = bm25s.tokenize(queries, return_ids=False)

# You can use a different corpus for retrieval, e.g., titles instead of full docs
titles = ["About Cat", "About Dog", "About Bird", "About Fish"]

# You can also choose to only return the documents and omit the scores
results = retriever.retrieve(query_token_ids, corpus=titles, k=2, return_as="documents")

# The documents are returned as a numpy array of shape (n_queries, k)
for i in range(results.shape[1]):
    print(f"Rank {i+1}: {results[0, i]}")

内存高效检索

bm25s 被设计为内存高效。您可以使用 mmap 选项将 BM25 索引作为内存映射文件加载，这样您就可以在不将整个索引加载到内存中的情况下加载索引。这在您有一个大型索引并且想要节省内存时很有用

# Create a BM25 index
# ...

# let's say you have a large corpus
corpus = [
    "a very long document that is very long and has many words",
    "another long document that is long and has many words",
    # ...
]
# Save the BM25 index to a file
retriever.save("bm25s_very_big_index", corpus=corpus)

# Load the BM25 index as a memory-mapped file, which is memory efficient
# and reduce overhead of loading the full index into memory
retriever = bm25s.BM25.load("bm25s_very_big_index", mmap=True)

有关如何使用 mmap=True 模式进行检索的示例，请参阅 examples/retrieve_nq.py。

分词

除了使用简单的 bm25s.tokenize 函数外，您还可以使用 Tokenizer 类来自定义分词过程。这在您想使用不同的分词器，或者当您想为查询和文档使用不同的分词过程时很有用

from bm25s.tokenization import Tokenizer

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",
]

# Pick your favorite stemmer, and pass 
stemmer = None
stopwords = ["is"]
splitter = lambda x: x.split() # function or regex pattern
# Create a tokenizer
tokenizer = Tokenizer(
      stemmer=stemmer, stopwords=stopwords, splitter=splitter
)

corpus_tokens = tokenizer.tokenize(corpus)

# let's see what the tokens look like
print("tokens:", corpus_tokens)
print("vocab:", tokenizer.get_vocab_dict())

# note: the vocab dict will either be a dict of `word -> id` if you don't have a stemmer, and a dict of `stemmed word -> stem id` if you do.
# You can save the vocab. it's fine to use the same dir as your index if filename doesn't conflict
tokenizer.save_vocab(save_dir="bm25s_very_big_index")

# loading:
new_tokenizer = Tokenizer(stemmer=stemmer, stopwords=[], splitter=splitter)
new_tokenizer.load_vocab("bm25s_very_big_index")
print("vocab reloaded:", new_tokenizer.get_vocab_dict())

# the same can be done for stopwords
print("stopwords before reload:", new_tokenizer.stopwords)
tokenizer.save_stopwords(save_dir="bm25s_very_big_index")
new_tokenizer.load_stopwords("bm25s_very_big_index")
print("stopwords reloaded:", new_tokenizer.stopwords)

您可以在 examples/tokenizer_class.py 中找到高级示例，包括如何

• 传递词干提取器、停用词和分割函数/正则表达式模式
• 通过 tokenizer.tokenize 调用来控制词汇是否更新（默认情况下，它仅在第一次调用期间更新）
• 使用 tokenizer.reset_vocab() 将分词器重置为其初始状态
• 使用生成器模式使用分词器以节省内存，通过 yield 一次一个文档。
• 将分词器的不同输出传递给 BM25.retrieve 函数。

变体

您可以在 bm25s 中使用以下 BM25 变体（有关更多详细信息，请参阅 Kamphuis et al. 2020）

• 原始实现（method="robertson"）- 我们将 idf>=0 设置为避免负数
• ATIRE（method="atire"）
• BM25L（method="bm25l"）
• BM25+（method="bm25+"）
• Lucene（method="lucene"）

默认情况下，bm25s 使用 method="lucene"，这是 Lucene 的 BM25 实现（确切版本）。您可以通过将 method 参数传递给 BM25 构造函数来更改方法

# The IR book recommends default values of k1 between 1.2 and 2.0, and b=0.75
retriever = bm25s.BM25(method="robertson", k1=1.5, b=0.75)

# For BM25+, BM25L, you need a delta parameter (default is 0.5)
retriever = bm25s.BM25(method="bm25+", delta=1.5)

# You can also choose a different "method" for idf, while keeping the default for the rest
# for example, this is equivalent to rank-bm25 when `epsilon=0`
retriever = bm25s.BM25(method="atire", idf_method="robertson")
# and this is equivalent to bm25-pt
retriever = bm25s.BM25(method="atire", idf_method="lucene")

Hugging Face 集成

bm25 可以与 Hugging Face 的 huggingface_hub 自然协同工作，允许您将模型保存到模型中心。这对于共享 BM25 索引和使用社区模型很有用。

首先，请确保您有一个有效的 Hugging Face 模型中心的访问令牌。这需要将模型保存到中心，或加载私有模型。创建后，您可以将它添加到您的环境变量中（例如，在您的 .bashrc 或 .zshrc 中）

export HUGGING_FACE_HUB_TOKEN="hf_..."

现在，让我们安装 huggingface_hub 库

pip install huggingface_hub

让我们看看如何使用 BM25SHF.save_to_hub 将 BM25 索引保存到 Hugging Face 模型库

import os
import bm25s
from bm25s.hf import BM25HF

# Create a BM25 index
retriever = 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",
]
corpus_tokens = bm25s.tokenize(corpus)
retriever.index(corpus_tokens)

# Set your username and token
user = "your-username"
token = os.environ["HF_TOKEN"]
retriever.save_to_hub(f"{user}/bm25s-animals", token=token, corpus=corpus)
# You can also save it publicly with private=False

然后，您可以使用以下代码从 Hugging Face 模型库加载 BM25 索引

import bm25s
from bm25s.hf import BM25HF

# Load a BM25 index from the Hugging Face model hub
user = "your-username"
retriever = BM25HF.load_from_hub(f"{user}/bm25s-animals")

# you can specify revision and load_corpus=True if needed
retriever = BM25HF.load_from_hub(
    f"{user}/bm25s-animals", revision="main", load_corpus=True
)

# if you want a low-memory usage, you can load as memory map with `mmap=True`
retriever = BM25HF.load_from_hub(
    f"{user}/bm25s-animals", load_corpus=True, mmap=True
)

# Query the corpus
query = "does the fish purr like a cat?"

# Tokenize the query
query_tokens = bm25s.tokenize(query)

# Get top-k results as a tuple of (doc ids, scores). Both are arrays of shape (n_queries, k)
results, scores = retriever.retrieve(query_tokens, k=2)

查看完整示例，请参阅

• examples/index_to_hf.py 用于对语料库进行索引并将其上传到 Huggingface Hub
• examples/retrieve_from_hf.py 用于从 Huggingface Hub 加载索引和语料库并进行查询。

比较

以下是一些基准测试，比较了 bm25s 与其他流行的 BM25 实现。我们比较了以下实现

• bm25s：纯 Python 中 BM25 的实现，由 Scipy 稀疏矩阵驱动。
• rank-bm25（《Rank》）：流行的 Python BM25 实现。
• bm25_pt（《PT》）：Pytorch 的 BM25 实现。
• elasticsearch（《ES》）：具有 BM25 配置的 Elasticsearch。

OOM 表示实现在进行基准测试时耗尽了内存。

吞吐量（每秒查询次数）

我们在各种数据集上比较了 BM25 实现的吞吐量。吞吐量以每秒查询次数（QPS）衡量，在单线程的 Intel Xeon CPU @ 2.70GHz（可在 Kaggle 找到）上。对于 BM25S，我们取 10 次运行的平均值。超过 60 查询/秒的实例以 粗体 显示。

更详细的基准测试可以在 bm25-benchmarks 仓库 中找到。

磁盘使用量

bm25s 设计得非常轻量级。这意味着该软件包的总磁盘使用量最小，因为它仅需要 numpy（18MB）、scipy（37MB）的 wheel，而软件包本身小于 100KB。安装后，完整虚拟环境占用的空间比 rank-bm25 多，但比 pyserini 和 bm25_pt 少。

$ du -s *env-* --block-size=1MB
6976    conda-env-pyserini
5346    venv-bm25-pt
479     venv-bm25s
45      venv-empty
99      venv-rank-bm25

优化内存使用

bm25s 通过使用 内存映射 来实现相当大的内存节省，允许索引存储在磁盘上并在需要时加载。

当使用与 MS MARCO 构建的索引（8.8M 文档，300M+ 令牌）进行 6 个任意查询测试时，我们得到以下结果

当在自然问题数据集（2M+ 文档）上运行 1000 个查询的 bm25s 时，内存使用量比内存版本低 50% 以上，速度差异微乎其微。更多详细信息可以在 GitHub 仓库 中找到。

致谢

• 多语言停用词来自 NLTK 停用词列表。
• numba 实现受到了最初由 baguetter 和 retriv 提出的 numba 实现的启发。
• 函数 bm25s.utils.beir.evaluate 来自 BEIR 库。它遵循与 BEIR 库相同的许可证，即 Apache 2.0。

引用

如果您在工作中使用了 bm25s，请使用以下 BibTeX 格式：

@misc{bm25s,
      title={BM25S: Orders of magnitude faster lexical search via eager sparse scoring}, 
      author={Xing Han Lù},
      year={2024},
      eprint={2407.03618},
      archivePrefix={arXiv},
      primaryClass={cs.IR},
      url={https://arxiv.org/abs/2407.03618}, 
}