【RAG召回】BM25算法示例
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 npbm25 = 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 npbm25 = 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 BM25Okapibm25 = 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