概要
langchainのBM25Retrieverを高速にマージする方法を検討しました。
背景
BM25アルゴリズムはキーワード検索を実施する代表的なアルゴリズムであり、生成AIと検索機能を組み合わせたRAGにおいても使用されることがあります。
langchainのBM25RetrieverはBM25による検索を生成AIと簡易に組み合わせるのに適したライブラリです。今回は、複数のBM25Retrieverを１つにまとめる方法を検討しました。
参考までに、同じくlangchainのTfidfRetrieverのマージ方法については以下の記事で検討したことがあります。
【langchain】TFIDFRetrieverのマージ
環境
環境
- M1 Mac, macOS 14.5
- Python 3.11.2
- rank-bm25 0.2.2
- langchain-community 0.3.1
以下の仮想環境上で実行しています。
uv venv
source .venv/bin/activate
uv pip install rank-bm25
簡単な方法
BM25Retireverをマージするいちばん簡単な方法は各retrieverのコーパスをマージして新たなBM25Retrieverを作ることです。
サンプルのコーパスとクエリを用意します。
corpus_a = [
"The quick brown fox jumps over the lazy dog.",
"A journey of a thousand miles begins with a single step.",
"To be or not to be, that is the question."
]
corpus_b = [
"In the middle of difficulty lies opportunity.",
"What we think, we become.",
"The only limit to our realization of tomorrow is our doubts of today."
]
query = "The only limit to our realization of tomorrow is our doubts of today."
各retrieverからdocsを取り出してマージします。
from langchain_community.retrievers import BM25Retriever
retriever_a = BM25Retriever.from_texts(corpus_a)
retriever_b = BM25Retriever.from_texts(corpus_b)
def merge_bm25retriever_simple(retriever_list: list[BM25Retriever]) -> BM25Retriever:
merged_docs = []
for retriever in retriever_list:
merged_docs += retriever.docs
bm25_params = {}
bm25_params["k1"] = retriever_list[0].vectorizer.k1
bm25_params["b"] = retriever_list[0].vectorizer.b
bm25_params["epsilon"] = retriever_list[0].vectorizer.epsilon
return BM25Retriever.from_documents(merged_docs,
k=retriever_list[0].k,
preprocess_func=retriever_list[0].preprocess_func,
bm25_params = bm25_params)
merged_retriever_simple = merge_bm25retriever_simple([retriever_a, retriever_b])
merged_retriever_simple.invoke(query)
[Document(metadata={}, page_content='The only limit to our realization of tomorrow is our doubts of today.'),
Document(metadata={}, page_content='To be or not to be, that is the question.'),
Document(metadata={}, page_content='The quick brown fox jumps over the lazy dog.'),
Document(metadata={}, page_content='What we think, we become.')]
マージの速度に特に不満を感じなければ、このコードで良いと思います。
高速な方法
コーパスサイズによっては、簡単な方法だと、マージの速度が遅くなるかもしれません。
その場合は、各retrieverで計算済みのIDF辞書をうまく活用することで効率的にマージできるようになります。
BM25Retrieverでは主に以下の２つのプラパティをマージすれば、retriever全体をマージすることができます。
- vectorizer: BM25のスコアリングに使用されるrank_bm25.BM25Okapiのインスタンス
- docs: 検索対象の文書リスト
docsは単なるリストなので、プラスで結合すればよいです。
vectorizerのマージ方法については、別記事で検証しました。
参考: 【python】rank_bm25のインスタンスを高速にマージ
上記記事を参考に、vectorizerをマージする処理を実装していきます。
まずマージに必要な情報をプラパティとして取得できるようにするために、rank_bm25.BM25Okapiのメソッドをオーバーライドします。
from rank_bm25 import BM25Okapi
def _initialize(self, corpus):
nd = {} # word -> number of documents with word
num_doc = 0
for document in corpus:
self.doc_len.append(len(document))
num_doc += len(document)
frequencies = {}
for word in document:
if word not in frequencies:
frequencies[word] = 0
frequencies[word] += 1
self.doc_freqs.append(frequencies)
for word, freq in frequencies.items():
try:
nd[word]+=1
except KeyError:
nd[word] = 1
self.corpus_size += 1
self.avgdl = num_doc / self.corpus_size
# ndを保存
self.nd = nd
return nd
BM25Okapi._initialize = _initialize
次にBM25Okapiのインスタンスを効率的にマージする関数を作成します。
詳細は参考記事に譲りますが、各インスタンスで計算済みの文書頻度辞書を再利用することで、マージが高速化されます。
def merge_bm25(bm25_list: list[BM25Okapi]) -> BM25Okapi:
if not bm25_list:
raise ValueError("The bm25_list is empty")
# Combine all tokenized corpora
merged_corpus_size = 0
merged_doc_freqs = []
merged_doc_len = []
for bm25 in bm25_list:
merged_corpus_size += bm25.corpus_size
merged_doc_freqs += bm25.doc_freqs
merged_doc_len += bm25.doc_len
merged_nd = {}
for bm25 in bm25_list:
for word, nd in bm25.nd.items():
merged_nd[word] = merged_nd.get(word, 0) + nd
merged_bm25 = BM25Okapi(["a"]
, tokenizer = bm25_list[0].tokenizer
, k1 = bm25_list[0].k1
, b = bm25_list[0].b
, epsilon = bm25_list[0].epsilon)
merged_bm25.corpus_size = merged_corpus_size
merged_bm25.doc_freqs = merged_doc_freqs
merged_bm25.doc_len = merged_doc_len
merged_bm25.avgdl = sum(merged_doc_len) / merged_corpus_size
merged_bm25.nd = merged_nd
merged_bm25._calc_idf(merged_nd)
return merged_bm25
最後にBM25Retireverをマージする関数を作ります。この関数は_initializeメソッドがオーバーライドされndがプラパティに保存されるようになったBM25Okapiを内部で使用することを前提としています。
from langchain_community.retrievers import BM25Retriever
def merge_bm25retriever(retriever_list: list[BM25Retriever]) -> BM25Retriever:
if not retriever_list:
raise ValueError("The retriever_list is empty")
merged_docs = []
for retriever in retriever_list:
merged_docs += retriever.docs
vectorizer_list = [retriever.vectorizer for retriever in retriever_list]
merged_vectorizer = merge_bm25(vectorizer_list)
merged_retriever = BM25Retriever(
vectorizer = merged_vectorizer,
docs = merged_docs,
k = retriever_list[0].k,
preprocess_func = retriever_list[0].preprocess_func
)
return merged_retriever
動作確認
簡単な実装と比べて検索結果が同じになるかを確認します。
コーパスとクエリは「簡単な実装」で定義したものと同様です。
retriever_a = BM25Retriever.from_texts(corpus_a)
retriever_b = BM25Retriever.from_texts(corpus_b)
merged_retriever = merge_bm25retriever([retriever_a, retriever_b])
merged_retriever.invoke(query)
[Document(metadata={}, page_content='The only limit to our realization of tomorrow is our doubts of today.'),
Document(metadata={}, page_content='To be or not to be, that is the question.'),
Document(metadata={}, page_content='The quick brown fox jumps over the lazy dog.'),
Document(metadata={}, page_content='What we think, we become.')]
簡単な実装の場合と同じ結果になっているので、大丈夫そうです。
速度
速度差を見やすくするため、大量の文書で実験します。50000件のサブコーパスを２つ用意します。
from datasets import load_dataset
# Load the AG News dataset
dataset = load_dataset('ag_news')
# Define sub-corpus size
sub_corpus_size = 50000
# Extract the text data from the dataset
corpus_a = [item['text'] for item in dataset['train'].select(range(sub_corpus_size))] # First sub_corpus_size items for corpus_a
corpus_b = [item['text'] for item in dataset['train'].select(range(sub_corpus_size, 2 * sub_corpus_size))] # Next sub_corpus_size items for corpus_b
query = "The only limit to our realization of tomorrow is our doubts of today."
2種類のやり方でのマージ速度を比較します。
import time
import numpy as np
iterations = 10
def measure_execution_time(func, *args, iterations=10, **kwargs):
times = []
for _ in range(iterations):
start_time = time.time()
func(*args, **kwargs)
end_time = time.time()
times.append(end_time - start_time)
average_time = np.mean(times)
std_dev_time = np.std(times)
return average_time, std_dev_time
retriever_a = BM25Retriever.from_texts(corpus_a)
retriever_b = BM25Retriever.from_texts(corpus_b)
# Measure time for merging BM25
average_time_merge, std_dev_time_merge = measure_execution_time(merge_bm25retriever_simple, [retriever_a, retriever_b])
print(f"merge_bm25retriever_simple function over {iterations} iterations: {average_time_merge} ± {std_dev_time_merge} seconds")
# Measure time for initializing BM25
average_time_init, std_dev_time_init = measure_execution_time(merge_bm25retriever, [retriever_a, retriever_b])
print(f"merge_bm25retriever function over {iterations} iterations: {average_time_init} ± {std_dev_time_init} seconds")
merge_bm25retriever_simple function over 10 iterations: 2.793686556816101 ± 0.1977367136912639 seconds
merge_bm25retriever function over 10 iterations: 0.1386489152908325 ± 0.0025654566699783777 seconds
1/20くらいに時間が短縮されました。
おわりに
BM25Retrieverの高速なマージ方法を検討し、50000件のコーパスで実験したところ、簡単な方法に比べて、1/20ほどに時間短縮できました。
BM25Okapiの関数をオーバーライドするのはできればやりたくないため、別の方法も引き続き検討したいです。