bm25-jieba 0.2.2

High-performance BM25 Chinese text search with jieba-rs tokenizer (Rust implementation)

BM25-Jieba 中文文本搜索

基于 Rust + PyO3 的高性能 BM25 中文文本搜索库，使用 jieba-rs 进行中文分词。

特性

• 🚀 高性能: Rust 实现，采用 倒排索引 + Block-Max WAND 算法加速，比纯 Python 快数倍
• 💾 持久化: 支持存取索引到磁盘 (MessagePack 格式)，无需重复训练
• 🔤 中文分词: 内置 jieba-rs 分词器
• 🎯 精确搜索: 经典 BM25 算法
• 🆔 自定义 ID: 支持绑定外部文档 ID (如数据库主键 u64)
• 🔠 大小写混合: 支持大小写不敏感搜索
• 🐍 Python 3.11 ~ 3.14: 支持最新 Python 版本

安装

# 开发模式安装
uv run maturin develop

# 或构建 wheel
maturin build --release
pip install target/wheels/*.whl

快速开始

from bm25_jieba import BM25

# 准备文档
documents = [
    "Python是一种广泛使用的高级编程语言",
    "机器学习是人工智能的一个分支",
    "深度学习是机器学习的子领域",
]

# 创建并训练模型 (可选传入自定义 ID)
bm25 = BM25(k1=1.5, b=0.75)
bm25.fit(documents, ids=[101, 102, 103])

# 搜索 (返回自定义 ID 和分数)
results = bm25.search("机器学习", top_k=3)
for doc_id, score in results:
    print(f"ID: {doc_id}, Score: {score:.4f}")

# 保存模型 (无需重复训练)
bm25.save("bm25_model.bin")

# 加载模型
loaded_bm25 = BM25.load("bm25_model.bin")

API 参考

BM25(k1=1.5, b=0.75, lowercase=False)

创建 BM25 实例。

参数	类型	默认值	说明
k1	float	1.5	词频饱和参数
b	float	0.75	文档长度归一化参数
lowercase	bool	False	是否将文本转换为小写（大小写不敏感）

fit(documents: list[str], ids: list[int] = None)

使用文档语料库训练模型。

• ids: 可选，与 documents 长度一致的整数列表 (u64)。
• 如果不提供 ids，默认使用 0..N 作为 ID。

search(query: str, top_k: int = None) -> list[tuple[int, float]]

搜索最相关的文档，返回 (文档 ID, 分数) 列表。

save(path: str)

保存当前索引和配置到文件 (MessagePack 格式)。

load(path: str) -> BM25

从文件加载 BM25 模型。

get_scores(query: str) -> list[float]

获取所有文档的 BM25 分数。

开发

# 安装依赖
uv sync

# 编译并安装
uv run maturin develop

# 运行测试
uv run pytest

# 运行示例
uv run python examples/demo.py

技术栈

组件	版本	用途
PyO3	0.27.2	Rust-Python 绑定
maturin	1.11.5	构建工具
jieba-rs	0.8.1	中文分词

性能测试

在 Apple M1 上测试 (10,000 文档，每文档约 100 字)：

测试项	结果
索引速度	~37,000 docs/s
搜索 QPS	~1,000,000 QPS
搜索延迟	~0.001ms

注：得益于 Block-Max WAND 算法的剪枝优化，搜索性能有数量级提升。

# 运行性能测试
uv run python tests/benchmark.py

算法验证

测试语料库：19 个文档，6 个查询

验证项	结果	说明
公式正确性	✅	手动计算与实现一致
排序一致性	✅	与 rank-bm25 排序完全一致
绝对分数	⚠️	因 IDF +1 修正略有差异（符合预期）

# 运行验证
uv sync --group validation
uv run python tests/validate.py

算法实现说明

本实现采用带 +1 修正的 IDF 公式：

IDF(t) = ln((N - df + 0.5) / (df + 0.5) + 1)

与标准 BM25Okapi 的区别：

公式	特点
标准: ln((N-df+0.5)/(df+0.5))	可能产生负 IDF
本实现: ln(...+1)	保证 IDF ≥ 0

影响：

• ✅ 排序一致 - 与 rank-bm25 等标准实现排序结果相同
• ⚠️ 绝对分数不同 - 因 +1 修正，分数值略有差异
• ✅ 数值稳定 - 无负值，无需额外处理

这种变体在只关心相对排序（而非绝对分数）的场景下完全适用。

License

MIT