rank-bm25
功能示例
本篇将通过多个示例,快速展示 rank-bm25
库的核心功能。不使用jieba。
准备工作
首先,确保您已经安装了 rank-bm25
。
pip install rank-bm25
接下来,我们定义一个通用的中文语料库和分词函数。这里我们使用简单的单字切分作为分词方法,以避免引入第三方库。
# 1. 定义一个中文语料库
corpus = [
"北京是中国的首都,也是一座历史悠久的文化名城。",
"上海是中国的经济中心,拥有繁忙的港口和现代化的建筑。",
"深圳是中国科技创新的重要城市,被誉为“中国硅谷”。",
"广州的美食文化闻名全国,是粤菜的发源地。",
"学习人工智能技术需要扎实的数学基础和编程能力。",
"中国的历史源远流长,有许多著名的历史人物和事件。"
]
# 2. 定义分词函数 (不使用jieba,直接按字切分)
def char_tokenizer(text):
"""
一个简单的按字切分的函数。
"""
return list(text)
# 3. 对语料库进行分词
tokenized_corpus = [char_tokenizer(doc) for doc in corpus]
print("分词后的语料库 (部分展示):")
print(tokenized_corpus[0])
# 输出: ['北', '京', '是', '中', '国', '的', '首', '都', ',', '也', '是', '一', '座', '历', '史', '悠', '久', '的', '文', '化', '名', '城', '。']
示例 1:基础用法 - 获取最相关的 N 个文档
这是最常用、最直接的功能:输入一个查询,直接返回最相关的文档列表。
from rank_bm25 import BM25Okapi
# 初始化 BM25 模型
bm25 = BM25Okapi(tokenized_corpus)
查询 = "中国的历史文化"
分词后的查询 = char_tokenizer(查询)
# 使用 get_top_n 获取最相关的 3 个文档
# 参数: (分词后的查询, 原始语料库, n=返回数量)
top_n_docs = bm25.get_top_n(分词后的查询, corpus, n=3)
print(f"查询: '{查询}'")
print("--- 最相关的3个文档 ---")
for doc in top_n_docs:
print(doc)
运行结果:
查询: '中国的历史文化'
--- 最相关的3个文档 ---
北京是中国的首都,也是一座历史悠久的文化名城。
中国的历史源远流长,有许多著名的历史人物和事件。
广州的美食文化闻名全国,是粤菜的发源地。
示例 2:获取所有文档的 BM25 分数
如果你不仅想知道谁最相关,还想知道具体的相关性分数是多少,可以使用 get_scores
。
from rank_bm25 import BM25Okapi
import numpy as np
bm25 = BM25Okapi(tokenized_corpus)
查询 = "中国的经济与科技"
分词后的查询 = char_tokenizer(查询)
# 使用 get_scores 获取每个文档的分数
doc_scores = bm25.get_scores(分词后的查询)
print(f"查询: '{查询}'")
print("--- 所有文档的BM25分数 ---")
print(np.round(doc_scores, 2))
# 你可以手动将分数和文档结合起来排序
带分数的文档 = list(zip(corpus, doc_scores))
排序后的文档 = sorted(带分数的文档, key=lambda item: item[1], reverse=True)
print("\n--- 手动排序后的结果 ---")
for doc, score in 排序后的文档:
print(f"分数: {score:.2f} | 文档: {doc}")
运行结果:
查询: '中国的经济与科技'
--- 所有文档的BM25分数 ---
[2.08 4.29 4.38 0. 2.19 2.08]
--- 手动排序后的结果 ---
分数: 4.38 | 文档: 深圳是中国科技创新的重要城市,被誉为“中国硅谷”。
分数: 4.29 | 文档: 上海是中国的经济中心,拥有繁忙的港口和现代化的建筑。
分数: 2.19 | 文档: 学习人工智能技术需要扎实的数学基础和编程能力。
分数: 2.08 | 文档: 北京是中国的首都,也是一座历史悠久的文化名城。
分数: 2.08 | 文档: 中国的历史源远流长,有许多著名的历史人物和事件。
分数: 0.00 | 文档: 广州的美食文化闻名全国,是粤菜的发源地。
示例 3:批量查询 - 同时处理多个查询
当有多个查询需要执行时,使用 get_batch_scores
会比循环调用 get_scores
更高效。
from rank_bm25 import BM25Okapi
import numpy as np
bm25 = BM25Okapi(tokenized_corpus)
查询列表 = [
"历史名城",
"科技创新",
"美食文化"
]
# 批量分词
分词后的查询列表 = [char_tokenizer(q) for q in 查询列表]
# 使用 get_batch_scores 进行批量计算
批量分数 = bm25.get_batch_scores(分词后的查询列表, list(range(len(corpus))))
print("--- 批量查询分数矩阵 ---")
print("行代表查询 (查询列表),列代表文档 (corpus)")
print(np.round(批量分数, 2))
# 为每个查询找到最相关的文档
print("\n--- 每个查询的最佳匹配 ---")
for i, query in enumerate(查询列表):
best_doc_index = np.argmax(批量分数[i])
print(f"查询 '{query}' 的最佳匹配: {corpus[best_doc_index]}")
运行结果:
--- 批量查询分数矩阵 ---
行代表查询 (查询列表),列代表文档 (corpus)
[[3.11 0. 0. 1.61 0. 3.11]
[0. 0. 3.7 0. 1.61 0. ]
[1.61 0. 0. 3.11 0. 0. ]]
--- 每个查询的最佳匹配 ---
查询 '历史名城' 的最佳匹配: 北京是中国的首都,也是一座历史悠久的文化名城。
查询 '科技创新' 的最佳匹配: 深圳是中国科技创新的重要城市,被誉为“中国硅谷”。
查询 '美食文化' 的最佳匹配: 广州的美食文化闻名全国,是粤菜的发源地。
示例 4:自定义 BM25 参数 (k1
和 b
)
BM25Okapi
模型可以接受两个重要参数 k1
和 b
来微调其行为。
k1
(float, default=1.5): 控制词频(TF)的缩放。值越高,词频对分数的影响就越大。b
(float, default=0.75): 控制文档长度惩罚。值在[0, 1]
区间,b=1
表示完全根据文档长度进行惩罚,b=0
表示不进行惩罚。
from rank_bm25 import BM25Okapi
# 使用默认参数的 BM25 模型
bm25_default = BM25Okapi(tokenized_corpus)
# 创建一个自定义参数的 BM25 模型
# 增强词频影响 (k1=2.0), 减弱文档长度惩罚 (b=0.5)
bm25_custom = BM25Okapi(tokenized_corpus, k1=2.0, b=0.5)
查询 = "中国历史"
分词后的查询 = char_tokenizer(查询)
# 获取两种模型下的分数
scores_default = bm25_default.get_scores(分词后的查询)
scores_custom = bm25_custom.get_scores(分词后的查询)
print(f"查询: '{查询}'")
print(f"默认参数 (k1=1.5, b=0.75) 分数: \n{np.round(scores_default, 2)}")
print(f"自定义参数 (k1=2.0, b=0.5) 分数: \n{np.round(scores_custom, 2)}")
运行结果:
查询: '中国历史'
默认参数 (k1=1.5, b=0.75) 分数:
[2.08 0. 0. 0. 0. 2.08]
自定义参数 (k1=2.0, b=0.5) 分数:
[2.35 0. 0. 0. 0. 2.35]
(注意:调整参数后,分数发生了变化)
示例 5:获取模型内部信息
有时需要查看模型内部的一些统计数据,例如词的文档频率、平均文档长度等。
from rank_bm25 import BM25Okapi
bm25 = BM25Okapi(tokenized_corpus)
# 获取模型计算出的平均文档长度
avg_dl = bm25.avgdl
print(f"平均文档长度: {avg_dl:.2f} 个字")
# 获取语料库的文档总数
doc_count = bm25.doc_count
print(f"文档总数: {doc_count}")
# 获取某个词在多少个文档中出现过 (文档频率)
词 = "中"
doc_freq = bm25.doc_freqs.get(词, 0)
print(f"'{词}' 这个字在 {doc_freq} 个文档中出现过。")
词 = "港"
doc_freq = bm25.doc_freqs.get(词, 0)
print(f"'{港}' 这个字在 {doc_freq} 个文档中出现过。")
# 查看模型为一个词计算的 IDF (逆文档频率) 分数
idf_score = bm25.idf.get(词, 0)
print(f"'{港}' 这个字的 IDF 分数是: {idf_score:.2f}")
运行结果:
平均文档长度: 28.50 个字
文档总数: 6
'中' 这个字在 4 个文档中出现过。
'港' 这个字在 1 个文档中出现过。
'港' 这个字的 IDF 分数是: 1.50