【LangChain】BM25 を永続化する
RAG(Retrieval-Augmented Generation)では、文章をベクトル化し、意味の近いものを検索するのが基本です。しかし専門的な用語を検索したい場合、LLM にその知識がないとその用語の意味がベクトルに反映されず、検索がうまくいかない可能性があります。そのようなケースに対応するため、キーワードベースの検索を組み合わせることが有効であるといわれています。
LangChain には、キーワードベースの検索手法の一つである BM25 を用いた検索機能が実装されています。BM25 は文章を単語の集まりとみなしてスコアリングする、というアルゴリズムなので、日本語のように単語の区切りが明確でない言語の場合、前処理として文章を単語に分割する必要があります。単語分割には Sudachi や Janome, Mecab といったツールを使うのが一般的です。
しかし、単語分割は意外と時間がかかります。なので、あらかじめ単語分割の前処理を行って保存し、検索時に分割結果を参照するという処理にしないと、検索対象の文章を全て毎回単語分割しなくてはいけなくなり、検索にかかる時間が長くなってしまします。
それにも関わらず、LangChain には なぜか単語分割結果を保存する機能が実装されていません。LangChain は設計がおかしい。
保存方法を調べてもなかなか情報がヒットしなかったので、とりあえずの解決方法ではありますが、記事で紹介します。
基本方針
もしまだ LangChain を使っていないなら、代わりに LlamaIndex を使うことをお勧めします。LlamaIndex には BM25 の永続化が実装されています。(ついでにフィルター機能も実装されています。)
どうしても LangChain を使いたい場合は以下のように BM25Retriever 全体を pickle 化するのが手っ取り早いです。
import pickle
from langchain_community.retrievers import BM25Retriever
retriever = BM25Retriever.from_texts(["テキスト", ...])
with open("bm25.pkl", "wb") as f:
pickle.dump(retriever, f)
読み込む時は以下のようにすればよいです。
import pickle
with open("bm25.pkl", "rb") as f:
retriever = pickle.load(f)
これらは BM25Retriever を継承してクラスメソッドとして定義してもよいです。(ちなみに VectorStore の一種である、FAISS も一部 pickle で保存しています。)
この方法でエラーが出なかった場合は、以降の内容は読まなくてよいです。
エラーが出た場合
実装によってはこの方法でエラーが出る場合があります。その場合は、日本語を分割する Sudachi, Janome, Mecab 等のオブジェクトが pickle 化できないことが原因の可能性があります。
import sudachipy
d = sudachipy.Dictionary()
tokenizer = d.create()
pickle.dumps(tokenizer)
TypeError: cannot pickle 'sudachipy.tokenizer.Tokenizer' object
Janome, Mecab でもエラーが出ます (他は試していません)。
とはいっても、エラーが出るのはもう少し特殊な状況で、例えば BM25Retriever に渡す preprocess_func をカスタマイズするために、以下のようにクラスのメンバー変数に tokenizer を持たせる場合に発生します。メンバー変数を pickle 化しようとするため、エラーが起きます。
class Preprocessor:
def __init__(self, args):
d = sudachipy.Dictionary(何か..)
self.tokenizer = d.create(何か...) # <- こいつが原因
# なんかカスタマイズ...
def preprocess(text: str) -> list[str]:
tokens = self.tokenizer.tokenize(text, 何か...)
# 品詞でフィルターしたり...
return [t.surface() for t in tokens]
proc = Preprocessor()
retriever = BM25Retriever.from_texts(
["テキスト", ...],
preprocess_func=proc.preprocess
)
with open("bm25.pkl", "wb") as f:
pickle.dump(retriever, f)
解決方法は 2 つあります。
解決方法 1
解決方法 1 つ目は preprocess_func に 普通の関数を渡す 方法です。つまり
d = sudachipy.Dictionary()
tokenizer = d.create()
def preprocess(test: str) -> list[str]:
tokens = self.tokenizer.tokenize(text)
return [t.surface() for t in tokens]
retriever = BM25Retriever.from_texts(
["テキスト", ...],
preprocess_func=preprocess
)
のようにすることです。これが上手くいく理由は、モジュールのトップレベルで定義された関数は pickle 化されないからです。正確にいうとされるのですが、関数名 (+モジュールおよびクラスの名前) のみが pickle 化され、その中身は保存されないようです (公式サイトより)。
よって関数の処理を変えれば、pickle.load 後のオブジェクトの動作も変わりますし、pickle.dump 後に関数名を変えると pickle.load に失敗します。
解決方法 2
解決方法 2 つ目は preprocess_func を外して pickle 化する 方法です。
retriever.preprocess_func = None
with open("bm25.pkl", "wb") as f:
pickle.dump(retriever, f)
with open("bm25.pkl", "rb") as f:
retriever = pickle.load(f)
retriever.preprocess_func = 省略...
ただし型を考慮すると、preprocess_func = None は想定されていないので型エラーになります。
とりあえず lambda 関数を渡せばいいと思うかもしれませんが、lambda 関数は pickle 化できないのでダメです。選択肢に悩みます。(デフォルト引数に指定されている default_preprocessing_func はありだと思います)
それか、BM25Retriever が pydantic で定義されていることを利用して、 dict にして保存する方法もあります。
_dict = retriever.model_dump()
del _dict["preprocess_func"]
with open("bm25.pkl", "wb") as f:
pickle.dump(retriever, f)
with open("bm25.pkl", "rb") as f:
_dict = pickle.load(f)
_dict["preprocess_func"] = 省略...
BM25Retriever.model_validate(_dict)
いずれにしても、preprocess_func は読み込み時に何らかの方法で渡す必要があります。
ちなみに実際に必要なのはおそらく vectorizer と docs だけなので、それだけ保存してもよいと思います。
preprocess_func を load 時に渡すべきか
preprocess_func を load 時に渡すのは変だと思うかもしれませんが、VectorStoreRetriever と比較してみると、実は意外とそうでもないことがわかります。VectorStoreRetriever
の機能をまとめると
-
テキストをベクトル化 (スコアリングのための変換) する機能
-
DB としての機能
- 文章、ベクトルの保存 (永続化)
- 追加、削除
- 検索 (スコアリング)
- フィルタリング
- 出力数制限
-
検索結果を他のモジュールへ渡す機能
となります。
テキストのベクトル化は Embedding が行います。通常の DB としての機能は VectorStore が担っており、検索結果を他のモジュールへ渡す機能を VectorStoreRetriever が担っています。VectorStoreRetriever は検索条件の仲介も行います。
BM25Retriever において Embedding に対応するのは preprocess_func だと考えられます。Embedding は VectorStore の保存対象ではなく外から与えられるので、この観点からは preprocess_func も保存対象から外すのは変ではありません。
ついでに BM25Retriever に不足している機能を ↑ に挙げた機能から抜粋すると
- DB としての機能
- 文章、単語分割の後のデータの保存 (永続化)
- 追加、削除
- フィルタリング
で、結構足りていません。追加、削除は BM25 のアルゴリズム的に実装しづらいのでしょうがないですが、残りの 2 つは LlamaIndex には実装されているので、LangChain で実装されていないのは不思議です。
Discussion