欢迎使用 bm25s
,这是一个在 Python 中实现 BM25 的库,允许您根据查询对文档进行排名。BM25 是一种广泛用于文本检索任务的排名函数,也是像 Elasticsearch 这样的搜索服务的核心组件。
它被设计为:
- 快速:
bm25s
以纯 Python 实现,并利用 Scipy 稀疏矩阵存储所有文档标记的预计算分数。这使得在查询时能够极其快速地进行评分,与流行库相比,性能提高了几个数量级(请参见下面的基准测试)。 - 简单:
bm25s
设计为易于使用和理解。您可以通过 pip 安装,并在几分钟内开始使用。它不依赖于 Java 或 Pytorch - 您只需要 Scipy 和 Numpy,以及可选的轻量依赖项用于词干提取。
在下面,我们将 bm25s
与 Elasticsearch 的速度与 rank-bm25
(最流行的 Python 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},
}
安装
您可以通过 pip 安装 bm25s
:
pip install bm25s
如果您想使用词干提取以获得更好的结果,可以安装推荐的(但可选的)依赖项:
# 安装所有额外依赖项
pip install bm25s[full]
# 如果您想使用词干提取以获得更好的结果,可以安装词干提取器
pip install PyStemmer
# 为加速 top-k 选择过程,您可以安装 `jax`
pip install jax[cpu]
快速开始
以下是如何使用 bm25s
的简单示例:
import bm25s
import Stemmer # 可选:用于词干提取
# 在此创建您的语料库
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",
]
# 可选:创建词干提取器
stemmer = Stemmer.Stemmer("english")
# 标记化语料库并仅保留 id(更快且节省内存)
corpus_tokens = bm25s.tokenize(corpus, stopwords="en", stemmer=stemmer)
# 创建 BM25 模型并索引语料库
retriever = bm25s.BM25()
retriever.index(corpus_tokens)
# 查询语料库
query = "does the fish purr like a cat?"
query_tokens = bm25s.tokenize(query, stemmer=stemmer)
# 获取 top-k 结果,结果是一个(文档 id,分数)的元组。两者都是形状为 (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"排名 {i+1} (分数: {score:.2f}): {doc}")
# 您可以将数组保存到目录中...
retriever.save("animal_index_bm25")
# 您可以将语料库与模型一起保存
retriever.save("animal_index_bm25", corpus=corpus)
# ...并在需要时加载它们
import bm25s
reloaded_retriever = bm25s.BM25.load("animal_index_bm25", load_corpus=True)
# 如果不需要语料库,可以设置 load_corpus=False
如需快速索引 200 万文档语料库(自然问题)的示例,请查看 examples/index_nq.py
。
灵活性
bm25s
提供了一个灵活的 API,允许您自定义 BM25 模型和标记化过程。以下是一些您可以使用的选项:
# 您可以提供一个查询列表而不是单个查询
queries = ["What is a cat?", "is the bird a dog?"]
# 如果您不喜欢默认的停用词列表,可以提供您自己的停用词列表
stopwords = ["a", "the"]
# 对于词干提取,使用任何可以在每个单词列表上调用的函数
stemmer_fn = lambda lst: [word for word in lst]
# 标记化查询
query_token_ids = bm25s.tokenize(queries, stopwords=stopwords, stemmer=stemmer_fn)
# 如果您希望标记化器返回字符串而不是标记 id,您可以这样做
query_token_strs = bm25s.tokenize(queries, return_ids=False)
# 您可以使用不同的语料库进行检索,例如,仅使用标题而不是完整文档
titles = ["About Cat", "About Dog", "About Bird", "About Fish"]
# 您还可以选择仅返回文档而省略分数
results = retriever.retrieve(query_token_ids, corpus=titles, k=2, return_as="documents")
# 文档作为形状为 (n_queries, k) 的 numpy 数组返回
for i in range(results.shape[1]):
print(f"排名 {i+1}: {results[0, i]}")
内存高效检索
bm25s
被设计为内存高效。您可以使用 mmap
选项将 BM25 索引作为内存映射文件加载,这使您可以在不将完整索引加载到内存中的情况下加载索引。当您有一个大索引并想节省内存时,这非常有用:
# 创建一个 BM25 索引
# ...
# 假设您有一个大语料库
corpus = [
"a very long document that is very long and has many words",
"another long document that is long and has many words",
# ...
]
# 将 BM25 索引保存到文件
retriever.save("bm25s_very_big_index", corpus=corpus)
# 将 BM25 索引作为内存映射文件加载,这在内存高效的情况下非常有用
# 并减少将完整索引加载到内存中的开销
retriever = bm25s.BM25.load("bm25s_very_big_index", mmap=True)
如需使用 mmap=True
模式进行检索的示例,请查看 examples/retrieve_nq.py
。
变体
您可以在 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
构造函数来更改方法:
# IR 书推荐的默认值为 k1 在 1.2 到 2
# 您可以在需要时指定修订版本并设置 `load_corpus=True`
retriever = BM25HF.load_from_hub(
f"{user}/bm25s-animals", revision="main", load_corpus=True
)
# 如果您想要低内存使用率,可以使用 `mmap=True` 作为内存映射加载
retriever = BM25HF.load_from_hub(
f"{user}/bm25s-animals", load_corpus=True, mmap=True
)
# 查询语料库
query = "鱼会像猫一样呼噜吗?"
# 对查询进行分词
query_tokens = bm25s.tokenize(query)
# 获取前 k 个结果,返回值为 (doc ids, scores) 的元组。两者均为形状为 (n_queries, k) 的数组
results, scores = retriever.retrieve(query_tokens, k=2)
```python
要查看完整示例,请访问:
examples/index_to_hf.py
进行语料库的索引和上传到 Huggingface Hubexamples/retrieve_from_hf.py
从 Huggingface Hub 加载索引和语料库并进行查询。
对比
以下是 bm25s
与其他流行的 BM25 实现的对比基准。我们比较了以下几种实现:
bm25s
:我们基于 Scipy 稀疏矩阵在纯 Python 中实现的 BM25。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 查询/秒的实例以 加粗 显示。
数据集 | BM25S | Elastic | BM25-PT | Rank-BM25 |
---|---|---|---|---|
arguana | 573.91 | 13.67 | 110.51 | 2 |
climate-fever | 13.09 | 4.02 | OOM | 0.03 |
cqadupstack | 170.91 | 13.38 | OOM | 0.77 |
dbpedia-entity | 13.44 | 10.68 | OOM | 0.11 |
fever | 20.19 | 7.45 | OOM | 0.06 |
fiqa | 507.03 | 16.96 | 20.52 | 4.46 |
hotpotqa | 20.88 | 7.11 | OOM | 0.04 |
msmarco | 12.2 | 11.88 | OOM | 0.07 |
nfcorpus | 1196.16 | 45.84 | 256.67 | 224.66 |
nq | 41.85 | 12.16 | OOM | 0.1 |
quora | 183.53 | 21.8 | 6.49 | 1.18 |
scidocs | 767.05 | 17.93 | 41.34 | 9.01 |
scifact | 952.92 | 20.81 | 184.3 | 47.6 |
trec-covid | 85.64 | 7.34 | 3.73 | 1.48 |
webis-touche2020 | 60.59 | 13.53 | OOM | 1.1 |
更多详细的基准测试可以在 bm25-benchmarks 仓库 中找到。
磁盘使用量
bm25s
设计得非常轻量化。这意味着该包的总磁盘使用量非常小,因为它只需要 numpy
(18MB),scipy
(37MB) 的 wheels,而包本身不到 100KB。安装完成后,完整的虚拟环境占用的空间比 rank-bm25
多,但比 pyserini
和 bm25_pt
少:
包 | 磁盘使用量 |
---|---|
venv (无包) | 45MB |
rank-bm25 | 99MB |
bm25s (我们实现的) | 479MB |
bm25_pt | 5346MB |
pyserini | 6976MB |
elastic | 1183MB |
显示详情
虚拟环境的磁盘使用量是通过以下命令计算的:
$ du -s *env-* --block-size=1MB
6976 conda-env-pyserini
5346 venv-bm25-pt
479 venv-bm25s
45 venv-empty
99 venv-rank-bm25
对于 pyserini
,我们使用了 推荐安装 的 conda 环境以处理 Java 依赖。
优化后的 RAM 使用量
bm25s
通过使用内存映射允许显著的内存节省,这使得索引可以存储在磁盘上并按需加载。
在使用 MS MARCO 构建的索引(8.8M 文档,300M+ 令牌)进行 6 个任意查询的测试中,我们得到以下结果:
方法 | 加载索引 (秒) | 检索 (秒) | RAM 使用量 (GB) |
---|---|---|---|
内存映射 | 0.62 | 0.18 | 0.90 |
内存中 | 11.41 | 0.74 | 10.56 |
当您在 Natural Questions 数据集(2M+ 文档)上运行 bm25s
进行 1000 个查询时,内存使用量比内存中版本降低了超过 50%,速度差异可以忽略不计。您可以在 GitHub 仓库 中找到更多信息。
引用
如果您在工作中使用 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},
}