Skip to main content

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

Project description

BM25-Jieba 中文文本搜索

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

特性

  • 🚀 高性能: Rust 实现,比纯 Python 快数倍
  • 🔤 中文分词: 内置 jieba-rs 分词器
  • 🎯 精确搜索: 经典 BM25 算法
  • 🐍 Python API: 简洁易用的接口

安装

# 开发模式安装
uv run maturin develop

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

快速开始

from bm25_jieba import BM25

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

# 创建并训练模型
bm25 = BM25(k1=1.5, b=0.75)
bm25.fit(documents)

# 搜索
results = bm25.search("机器学习", top_k=3)
for doc_idx, score in results:
    print(f"[{score:.4f}] {documents[doc_idx]}")

API 参考

BM25(k1=1.5, b=0.75)

创建 BM25 实例。

参数 类型 默认值 说明
k1 float 1.5 词频饱和参数
b float 0.75 文档长度归一化参数

fit(documents: list[str])

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

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

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

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 文档,每文档约 50 字):

测试项 结果
索引速度 ~38,000 docs/s
搜索 QPS ~2,000 QPS
搜索延迟 0.5-1.2ms
# 运行性能测试
uv run python tests/benchmark.py

算法验证

测试语料库:19 个文档,6 个查询

验证项 结果 说明
公式正确性 手动计算与实现一致
排序一致性 6/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

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

bm25_jieba-0.1.0.tar.gz (26.8 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

bm25_jieba-0.1.0-cp312-cp312-macosx_11_0_arm64.whl (3.3 MB view details)

Uploaded CPython 3.12macOS 11.0+ ARM64

File details

Details for the file bm25_jieba-0.1.0.tar.gz.

File metadata

  • Download URL: bm25_jieba-0.1.0.tar.gz
  • Upload date:
  • Size: 26.8 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: maturin/1.11.5

File hashes

Hashes for bm25_jieba-0.1.0.tar.gz
Algorithm Hash digest
SHA256 61dfe0e9aefc7ecf79be68e775cc58ce224664a668a4ff59524c5ce187dcb1cc
MD5 c4b573be9a7125189649ccc781aed4f7
BLAKE2b-256 3bce914c22d8808a40053bc5113a97536f336a96c81525c5c4de851c7c451fe6

See more details on using hashes here.

File details

Details for the file bm25_jieba-0.1.0-cp312-cp312-macosx_11_0_arm64.whl.

File metadata

File hashes

Hashes for bm25_jieba-0.1.0-cp312-cp312-macosx_11_0_arm64.whl
Algorithm Hash digest
SHA256 615b027be90a25e27eef4bf46c1e09d24a0dc61ea37ed9729fce0cea3f578fcf
MD5 f013cebcb0bf7e6b2fdef858bc65aae6
BLAKE2b-256 da9ec63890c91304bc184100c9e079bdbb1e81affdd0a8ca44e610914580aae0

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page