概要
langchainのキーワード検索用RetrieverであるTFIDFRetrieverとBM25Retrieverで、vectorizerを作るためのコーパスと、Retrierクラス内で保持される検索対象のコーパスを別にする方法を検討しました。
結論としては、TFIDFRetrieverであれば、以下のようにすれば、vectorizerと検索用のコーパスを別々にすることができます。
def create_tfidf_retriever(texts, docids = None, vectorizer = None):
tfidf_array = vectorizer.transform(texts)
docids = docids or list(map(str, range(len(texts))))
documents = [Document(page_content = text, metadata = {"docid": docid}) for text, docid in zip(texts, docids)]
retriever = TFIDFRetriever(vectorizer=vectorizer, docs = documents, tfidf_array = tfidf_array)
return retriever
BM25Retrieverでは、classのoverride等をしない限り、標準的に実現する方法はなさそうでした。
環境
以下をインストールします。
pip install langchain openai tiktoken rank-bm25 scikit-learn sudachipy sudachidict_full
検証に用いたlangchainのバージョンは0.0.324でした。
背景
キーワード検索でよく使われるTF-IDFやBM25などの指標は、コーパス全体における単語の出現頻度をもとに計算されます。したがって、値がコーパスに依存します。
TF-IDFやBM25を算出するライブラリ（TfidfVectorizerやrank-bm25など）では、コーパスごとにvectorizerを作成します。また検索機能も一体となったライブラリでは、検索対象の文章はvectorizerを作るために使用したコーパスと同じになることが多いです。
キーワード検索の利用シーンを想定すると、コーパスが変更になるたびに（例えば、検索対象の文章が１つ追加されるたびに）vectorizerを再計算するのは非効率です。そこで、vectorizerを作るためのコーパスと、検索対象の文章を独立にしたいという要求が生じます。典型的なケースとしては、大規模なコーパスでvectorizerを作り、それを小規模な検索対象の文章に対して使い回したくなることがあります。
そこで、検索拡張生成の実装によく使われるlangchainのRetrieverで、vectorizerと検索対象の文章リストを独立に管理する方法を調べました。具体的には、キーワード検索の指標としてよく使われるTF-IDFとBM25のRetrieverについて調べました。
準備
使うライブラリをimportします。
import json
import numpy as np
import os
from langchain.retrievers import BM25Retriever
from langchain.retrievers import TFIDFRetriever
from sklearn.feature_extraction.text import TfidfVectorizer
from langchain.schema import Document
from sudachipy import tokenizer
from sudachipy import dictionary
import requests
from rank_bm25 import BM25Okapi
import dotenv
import tqdm
from typing import List
TFIDFRetrieverやBM25Retrieverを日本語文章に対して機能させるためには、日本語を分かち書きできるtokenizerが必要ですので用意します。ストップワードは今回は検証なのでなくてもよいです。
参考
- 【langchain】BM25Retriever/TFIDFRetrieverを日本語に対応させる
- RAG処理の改善: langchainでハイブリッド検索を実装してみる(勉強メモ)(ストップワード実装の参考)
# 日本語ストップワード辞書
url = "https://raw.githubusercontent.com/stopwords-iso/stopwords-ja/master/stopwords-ja.txt"
stopwords_jp = requests.get(url).text.split("\n")
def remove_stopwords(tokens: List[str]) -> List[str]:
"""指定した単語リストからストップワードを除去した結果を返す"""
tokens = [token for token in tokens if token not in stopwords_jp]
return tokens
tokenizer_obj = dictionary.Dictionary(dict="full").create()
mode = tokenizer.Tokenizer.SplitMode.A
def tokenize_jp(text: str) -> List[str]:
tokens = tokenizer_obj.tokenize(text ,mode)
words = [token.surface() for token in tokens]
words = remove_stopwords(words)
return words
vectorizer用のコーパス（full_corpus）と検索対象用のコーパス（search_corpus）を定義します。両者が異なってさえいれば良いので、適当です。
full_corpus = [
"こんにちは、今日はいい天気ですね"
, "私たちは卒業します"
, "また会う日まで"
]
search_corpus = [
"卒業します"
]
TFIDFRetriever
普通にfrom_textsなどでRetrieverを作成すると、vectorizer用のコーパスと検索対象の文書が同一になってしまいます。
retriever = TFIDFRetriever.from_texts(full_corpus, tfidf_params={"analyzer":tokenize_jp})
result = retriever.get_relevant_documents(query_text)
print(query_text)
for doc in result:
print(doc)
こんにちは
page_content='こんにちは、今日はいい天気ですね'
page_content='また会う日まで'
page_content='私たちは卒業します'
以下のように
- vectorizerをcikit-learnのTfidfVetorizerで直接定義する
- 検索文書をvectorizerでarrayに変換する
- TFIDFRetrieverのコンストラクタにvectorizerと検索文書のarrayを渡す
という手順を経ることで、vectorizerと検索対象文章を別々にすることができます。
def create_tfidf_vectorizer(texts):
vectorizer = TfidfVectorizer(analyzer=tokenize_jp)
vectorizer.fit(texts)
return vectorizer
def create_tfidf_retriever(texts, docids = None, vectorizer = None):
tfidf_array = vectorizer.transform(texts)
docids = docids or list(map(str, range(len(texts))))
documents = [Document(page_content = text, metadata = {"docid": docid}) for text, docid in zip(texts, docids)]
retriever = TFIDFRetriever(vectorizer=vectorizer, docs = documents, tfidf_array = tfidf_array)
return retriever
vectorizer = create_tfidf_vectorizer(full_corpus)
retriever = create_tfidf_retriever(search_corpus, None, vectorizer=vectorizer)
result = retriever.get_relevant_documents(query_text)
print(query_text)
for doc in result:
print(doc)
こんにちは
page_content='卒業します' metadata={'docid': '0'}
検索結果が1件なので、検索対象をsearch_corpusにすることができています。
vectorrizerがfull_corpusに基づく（search_corpusに基づかない)ことを確かめるために、vectrizerにquery_textをtransformさせてみます。
vec = retriever.vectorizer.transform([query_text])[0]
for word, idx in retriever.vectorizer.vocabulary_.items():
print(word, vec[0, idx])
こんにちは 1.0
、 0.0
今日 0.0
いい 0.0
天気 0.0
ね 0.0
卒業 0.0
会う 0.0
日 0.0
search_corpusに存在しない語彙があるので、vocabularyはvectorizerに基づいていると言えそうです。
ソースコード
TFIDFRetrieverの実装を確認しておきます。
from_textsメソッドを読むと以下の処理が書かれています。
tfidf_params = tfidf_params or {}
vectorizer = TfidfVectorizer(**tfidf_params)
tfidf_array = vectorizer.fit_transform(texts)
metadatas = metadatas or ({} for _ in texts)
docs = [Document(page_content=t, metadata=m) for t, m in zip(texts, metadatas)]
return cls(vectorizer=vectorizer, docs=docs, tfidf_array=tfidf_array, **kwargs)
scikit-learnのTfidfVectorizerを作ってから、tfidf_arrayを作り、textsをDocumentに変換して、クラスのコンストラクタに渡しています。ここではtfidf_arrayはvectorizerを作成したコーパス（texts）を行列に変換したものですが、fitとtransformを分けて書けば、別々にすることもできそうです。
vectorizerとtfidf_arrayが検索時にどのように呼び出されるかを確認するため、_get_relevant_documentsメソッドを確認します。
def _get_relevant_documents(
self, query: str, *, run_manager: CallbackManagerForRetrieverRun
) -> List[Document]:
from sklearn.metrics.pairwise import cosine_similarity
query_vec = self.vectorizer.transform(
[query]
) # Ip -- (n_docs,x), Op -- (n_docs,n_Feats)
results = cosine_similarity(self.tfidf_array, query_vec).reshape(
(-1,)
) # Op -- (n_docs,1) -- Cosine Sim with each doc
return_docs = [self.docs[i] for i in results.argsort()[-self.k :][::-1]]
return return_docs
queryのベクトル化はvectorizerで行われ、類似度の算出はtfidf_arrayに対して行われています。したがって、vectorizerとtfidf_arrayを別々に作って渡しても機能することがわかります。
BM25Retriever
結論から言うと、BM25Retrieverでは、vectorizerと検索対象のコーパスを独立にすることができません。できないということを以下で確認します。
まず普通にfrom_textsなどでRetrieverを作成すると、vectorizer用のコーパスと検索対象の文書が同一になってしまいます。
retriever = BM25Retriever.from_texts(full_corpus, preprocess_func=tokenize_jp)
result = retriever.get_relevant_documents(query_text)
print(query_text)
for doc in result:
print(doc)
こんにちは
page_content='こんにちは、今日はいい天気ですね'
page_content='また会う日まで'
page_content='私たちは卒業します'
次にTFIDFRetrieverと同じように、vectorizerを直接作ってから検索用のコーパスをそれとは別にしてコンストラクタに与えてみます。
def create_bm25_vectorizer(texts):
texts_processed = [tokenize_jp(t) for t in tqdm.tqdm(texts)]
bm25_params = {}
vectorizer = BM25Okapi(texts_processed, **bm25_params)
return vectorizer
def create_bm25_retriever(texts, docids = None, vectorizer = None):
# 無効なコード。このやり方ではvectorizerとcorpus（documents）を独立にはできない
docids = docids or list(map(str, range(len(texts))))
documents = [Document(page_content = text, metadata = {"docid": docid}) for text, docid in zip(texts, docids)]
retriever = BM25Retriever(vectorizer=vectorizer, docs = documents, preprocess_func=tokenize_jp)
return retriever
vectorizer = create_bm25_vectorizer(full_corpus)
retriever = create_bm25_retriever(search_corpus, None, vectorizer=vectorizer)
result = retriever.get_relevant_documents(query_text)
202 _kwargs = kwargs if self._expects_other_args else {}
203 if self._new_arg_supported:
--> 204 result = self._get_relevant_documents(
205 query, run_manager=run_manager, **_kwargs
206 )
207 else:
208 result = self._get_relevant_documents(query, **_kwargs)
...
---> 71 assert self.corpus_size == len(documents), "The documents given don't match the index corpus!"
73 scores = self.get_scores(query)
74 top_n = np.argsort(scores)[::-1][:n]
AssertionError: The documents given don't match the index corpus!
retrieverの作成まではできるのですが、get_relevant_documentsを実行するとエラーになります。
このエラーは、BM25Retrieverが内部で使用しているrank-bm25ライブラリ（BM25Okapiの継承元のBM25のget_top_nメソッド）で発生しています。内容は、corpus_size（今回でいうfull_corpusのサイズ）とdocuments（今回でいうsearch_corpus）のサイズが一致しないというものです。
corpus_sizeはrank-bm25でvectorizerを作成したときに決定される変数で、これをBM25Retriever側から変更する方法はないため、基本的にはBM25Retrieverではvectorizerに使用するコーパスと検索対象のコーパスを別々にすることは難しそうです。
ソースコード
BM25Retrieverのソースコードを一応確認します。
from_textsで以下の処理があります。
texts_processed = [preprocess_func(t) for t in texts]
bm25_params = bm25_params or {}
vectorizer = BM25Okapi(texts_processed, **bm25_params)
metadatas = metadatas or ({} for _ in texts)
docs = [Document(page_content=t, metadata=m) for t, m in zip(texts, metadatas)]
return cls(
vectorizer=vectorizer, docs=docs, preprocess_func=preprocess_func, **kwargs
)
TfidfRetrieverと同様にvectorizerをBM25Okapiで作ってからvectorizerとdocsをそれぞれコンストラクタに渡しています。
次に_get_relevant_documentsを確認します。
def _get_relevant_documents(
self, query: str, *, run_manager: CallbackManagerForRetrieverRun
) -> List[Document]:
processed_query = self.preprocess_func(query)
return_docs = self.vectorizer.get_top_n(processed_query, self.docs, n=self.k)
return return_docs
vectorizer.get_top_nにqueryとdocsを渡しているので、ここだけみると、vectorizerとdocsは別々でもよさそうです。しかし、先に見た通り、get_top_nの内部でエラーになります。
langchainの外に出てrank-bm25のソースコードを確認します。
def get_top_n(self, query, documents, n=5):
assert self.corpus_size == len(documents), "The documents given don't match the index corpus!"
scores = self.get_scores(query)
top_n = np.argsort(scores)[::-1][:n]
return [documents[i] for i in top_n]
まずcorpus_sizeとdocumentsのサイズが違うことでassertionが出てしまうのですが、これを回避したとしても、get_scoresでスコアの行列の取得しています。この行列をargsortして、documentsの要素をiで取得していることから、documentsとget_scoresの対象となる文書群が同一であることを期待されているとわかります。get_scoresが参照する文書はBM25Okapiの初期化時に決定されるため、これを変更するのは面倒です。
ちなみにBM25Okapiにはget_batch_scoresというメソッドがあり、これはqueryとdocidsを入力することで、docidsに対応するサブコーパスを検索対象とできるようです。ただし、検索対象はサブコーパス以外にもできたほうがよいので、完璧ではありません。
vectorizerと検索対象文書を独立にしたいというのは、よくある要望と思われ、rank-bm25にも
issueが立っていますが、今のところ解決されていません。
おわりに
キーワード検索用のRetrieverでvectorizerを作るためのコーパスと、検索対象の文章を独立にする方法を検討しました。
TFIDFRetrieverについては方法が見つかりました。
BM25Retrieverについては、ハードルは高く、クラスをオーバーライドするか、ゼロから実装する必要がありそうです。
キーワード検索としてはBM25のほうが王道のため、悩ましいところですが、現状では、TF-IDFで妥協できるのであれば、妥協するのがよさそうです。