【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 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