⛓️

【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 に 普通の関数を渡す 方法です。つまり

解決方法1
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 は想定されていないので型エラーになります。

https://github.com/langchain-ai/langchain-community/blob/64a56a408cac8fb07f907c3e5a62ca5c364d14ee/libs/community/langchain_community/retrievers/bm25.py#L24-L25

とりあえず 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 だけなので、それだけ保存してもよいと思います。

https://github.com/langchain-ai/langchain-community/blob/64a56a408cac8fb07f907c3e5a62ca5c364d14ee/libs/community/langchain_community/retrievers/bm25.py#L15-L21

preprocess_func を load 時に渡すべきか

preprocess_func を load 時に渡すのは変だと思うかもしれませんが、VectorStoreRetriever と比較してみると、実は意外とそうでもないことがわかります。VectorStoreRetriever
の機能をまとめると

  1. テキストをベクトル化 (スコアリングのための変換) する機能

  2. DB としての機能

    • 文章、ベクトルの保存 (永続化)
    • 追加、削除
    • 検索 (スコアリング)
    • フィルタリング
    • 出力数制限
  3. 検索結果を他のモジュールへ渡す機能

となります。

テキストのベクトル化は Embedding が行います。通常の DB としての機能は VectorStore が担っており、検索結果を他のモジュールへ渡す機能を VectorStoreRetriever が担っています。VectorStoreRetriever は検索条件の仲介も行います。

BM25Retriever において Embedding に対応するのは preprocess_func だと考えられます。Embedding は VectorStore の保存対象ではなく外から与えられるので、この観点からは preprocess_func も保存対象から外すのは変ではありません。

ついでに BM25Retriever に不足している機能を ↑ に挙げた機能から抜粋すると

  • DB としての機能
    • 文章、単語分割の後のデータの保存 (永続化)
    • 追加、削除
    • フィルタリング

で、結構足りていません。追加、削除は BM25 のアルゴリズム的に実装しづらいのでしょうがないですが、残りの 2 つは LlamaIndex には実装されているので、LangChain で実装されていないのは不思議です。

GitHubで編集を提案

Discussion