【langchain】TFIDFRetrieverのマージ

概要

異なるコーパスから作られた複数のTFIDFRetrieverを1つにする方法を検討しました。

背景

TFIDFRetrieverはlangchainでキーワード検索をするときの有用な選択肢の一つです。
開発途中で検索対象の文書を追加する場合など、TFIDFRetrieverの更新やマージをしたくなることがあります。

FAISSRetrieverなどDenseな埋め込みに基づくRetrieverの場合はdocstoreとembeddingをただ拡張すればよいのですが、TFIDFなどSparseなRetrieverの場合、単語のIDやIDF値がコーパスごとに決定されるため、単純にマージはできません。そこで方法を考えました。

（注：キーワード検索用のSparseなRetrieverとしては、BM25Retrieverのほうがよりメジャーなのですが、BM25Retrieverは内部で使用されているrank_bm25の扱いに習熟できていないため、今回はTFIDFRetrieverを対象とします。）

一番シンプルなのは、複数のデータセットをマージして、TFIDFRetrieverを作り直すことです。
例えばdataset1とdataset2でそれぞれ作ったTFIDFRetrieverをマージしたければ、dataset1+dataset2を作ってTFIDFRetriever.from_textsなどを再度作れば良いです。

ただこの方法は形態素解析やIDFの計算を全てやり直すことになり、datasetの量によっては時間がかかってしまいます。そこでこれを効率化する方法を考えます。

方針

ソースコードによればTFIDFRetrieverのマージすべき属性は以下です。

• docs: List[Document]
• tfidf_array: scipy.sparse._csr.csr_matrix (docsをvectorizerでtransformした行列）
• vectorizer: TfidfVectorizer

それぞれ以下のようにマージできます。

• docsのリストをextendするだけです。
• vectorizer（TfidfVectorizer)のマージ方法は別の記事で検討しました。idfとデータセット長から各単語の出現回数を計算することでマージ後のidfを高速に計算できます。
• tfidf_arrayはいくつかの方向性があります。形態素解析済みのデータセットがあれば、マージ後のvectorizerでtransformすることで作れます。形態素解析済みのデータセットがなければ、もとのtfidf_arrayをinverse_transformして得ることもできます。ただし単語の重複がのぞかれるため誤差があります。それが許容できない場合は、docsを形態素解析する必要があります。

実装

tfidf_arrayのマージはinverse_transformを使うパターンで実装してみます。

ライブラリをimportします。

import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from langchain.retrievers import TFIDFRetriever
from typing import List, Dict, Any, Tuple, Optional
import MeCab

vectorizerをmergeする関数を作ります（参考)

def merge_tfidf_vectorizers(vectorizers, datasets, forced_tfidf_params = None)->TfidfVectorizer:
    """
    複数のTfidfVectorizerインスタンスをマージする関数

    :param vectorizers: TfidfVectorizerインスタンスのリスト
    :param datasets: 各TfidfVectorizerに対応するデータセットのリスト
    :param tfidf_params: TfidfVectorizerのパラメータ
    :return: 統合された新しいTfidfVectorizerインスタンス
    """
    # IDF値から文書頻度を逆算する関数
    def idf_to_doc_freq(idf, n_docs):
        df = (n_docs + 1) / np.exp(idf - 1) - 1
        return np.round(df).astype(int)

    
    # 語彙と文書頻度を統合する
    combined_vocabulary = set()
    combined_df = {}
    n_docs_combined = 0

    for vectorizer, dataset in zip(vectorizers, datasets):
        df = idf_to_doc_freq(vectorizer.idf_, len(dataset))
        vocab_df = dict(zip(vectorizer.get_feature_names_out(), df))

        combined_vocabulary.update(vectorizer.get_feature_names_out())
        for word, freq in vocab_df.items():
            combined_df[word] = combined_df.get(word, 0) + freq

        n_docs_combined += len(dataset)

    combined_vocabulary = sorted(combined_vocabulary)

    # 新しいIDF値を計算
    combined_idf = np.log((n_docs_combined + 1) / (np.array(list(combined_df.values())) + 1)) + 1

    # 0番目の要素のparamをコピー。vectorizers同士のパラメータは全て等しい前提
    tfidf_params = vectorizers[0].get_params()
    tfidf_params["vocabulary"] = combined_vocabulary # vocabularyだけ更新
    forced_tfidf_params = forced_tfidf_params or {}
    tfidf_params.update(forced_tfidf_params) # 強制的に上書きしたいパラメータがあれば上書き

    # 新しいTfidfVectorizerインスタンスを作成してIDF値を設定
    tfidf_combined = TfidfVectorizer(**tfidf_params)
    tfidf_combined.fit(["dummy"])  # ダミーデータで初期化
    tfidf_combined._tfidf.idf_ = combined_idf

    return tfidf_combined

TFIDFRetrieverをマージする関数を作ります。

def merge_tfidf_retrievers(retrievers: List[TFIDFRetriever]):
    """
    複数のTFIDFRetrieverインスタンスをマージする関数

    :param retrievers: TFIDFRetrieverインスタンスのリスト
    :return: 統合された新しいTFIDFRetrieverインスタンス
    """
    # vectorizerのマージ
    vectorizers = [retriever.vectorizer for retriever in retrievers]
    merged_vectorizer = merge_tfidf_vectorizers(vectorizers, [retriever.docs for retriever in retrievers])
    # tfidf_arrayのマージ
    merged_tokenized_docs = []
    for retriever in retrievers:
        tokenized_docs = retriever.vectorizer.inverse_transform(retriever.tfidf_array)
        merged_tokenized_docs.extend(tokenized_docs)
    params = merged_vectorizer.get_params()
    merged_vectorizer.set_params(analyzer=lambda x: x) # transformの時間節約のため一時的にanalyzerを無効化
    merged_tfidf_array = merged_vectorizer.transform(merged_tokenized_docs)
    merged_vectorizer.set_params(analyzer=params["analyzer"]) # analyzerを元に戻す
    # docsのマージ
    merged_docs = [doc for retriever in retrievers for doc in retriever.docs]
    # TFIDFRetrieverインスタンスの作成
    merged_retriever = TFIDFRetriever(docs=merged_docs, vectorizer=merged_vectorizer, tfidf_array=merged_tfidf_array
                                      , k=retrievers[0].k) # kはretrievers間で全て同じと仮定

    return merged_retriever

テストしてみます。

dataset1 = ["文書1からのテキスト", "文書2からのテキスト"]
dataset2 = ["文書3からのテキスト", "文書4からのテキスト"]

# 形態素解析器の初期化（MeCab）
mecab = MeCab.Tagger("-Owakati")

# 形態素解析を行う関数
def tokenize(text):
    return mecab.parse(text).split()

retriever1 = TFIDFRetriever.from_texts(dataset1, tfidf_params={"analyzer": tokenize})
retriever2 = TFIDFRetriever.from_texts(dataset2, tfidf_params={"analyzer": tokenize})

merged_retriever = merge_tfidf_retrievers([retriever1, retriever2])

print("idf")
print("retriever1")
print(retriever1.vectorizer.get_feature_names_out())
print(retriever1.vectorizer.idf_)
print("retriever2")
print(retriever2.vectorizer.get_feature_names_out())
print(retriever2.vectorizer.idf_)
print("merged_retriever")
print(merged_retriever.vectorizer.get_feature_names_out())
print(merged_retriever.vectorizer.idf_)
print("")
print("検索テスト")
print(retriever1.get_relevant_documents("文書1"))
print(retriever2.get_relevant_documents("文書3"))
print(merged_retriever.get_relevant_documents("文書1と文書3"))

idf
retriever1
['1' '2' 'から' 'の' 'テキスト' '文書']
[1.40546511 1.40546511 1.         1.         1.         1.        ]
retriever2
['3' '4' 'から' 'の' 'テキスト' '文書']
[1.40546511 1.40546511 1.         1.         1.         1.        ]
merged_retriever
['1' '2' '3' '4' 'から' 'の' 'テキスト' '文書']
[1.91629073 1.91629073 1.         1.         1.         1.
 1.91629073 1.91629073]

検索テスト
[Document(page_content='文書1からのテキスト'), Document(page_content='文書2からのテキスト')]
[Document(page_content='文書3からのテキスト'), Document(page_content='文書4からのテキスト')]
[Document(page_content='文書1からのテキスト'), Document(page_content='文書3からのテキスト'), Document(page_content='文書4からのテキスト'), Document(page_content='文書2からのテキスト')]

厳密には検証していませんが、なんとなくうまくいっている気がします。

おまけ

tfidf_arrayのマージの別バージョンの方法のコードです。
動くかどうかは試していません。

形態素解析をやり直す

時間がかかってもよければ、これを使うのもありです。

# docsからtfidf_arrayを再計算
def merge_tfidf_retrievers2(retrievers: List[TFIDFRetriever]):
    """
    複数のTFIDFRetrieverインスタンスをマージする関数

    :param retrievers: TFIDFRetrieverインスタンスのリスト
    :return: 統合された新しいTFIDFRetrieverインスタンス
    """
    # vectorizerのマージ
    vectorizers = [retriever.vectorizer for retriever in retrievers]
    merged_vectorizer = merge_tfidf_vectorizers(vectorizers, [retriever.docs for retriever in retrievers])

    # docsのマージ
    merged_docs = [doc for retriever in retrievers for doc in retriever.docs]
    # tfidf_arrayのマージ
    merged_tfidf_array = merged_vectorizer.transform([doc.page_content for doc in merged_docs])
    # インスタンスの初期化
    merged_retriever = TFIDFRetriever(docs=merged_docs, vectorizer=merged_vectorizer, tfidf_array=merged_tfidf_array
                                      , k=retrievers[0].k) # kはretrievers間で全て同じと仮定

    return merged_retriever

形態素解析もIDF計算もやり直す

厳密さは上記と変わらず、単に時間がよりかかるので、あまり使う意味はないです。ただ実装がシンプルなのでメンテナンスはしやすいです。

# docsからtfidf_arrayを再計算
def merge_tfidf_retrievers3(retrievers: List[TFIDFRetriever], tokenized_docs: List[List[str]]):
    """
    複数のTFIDFRetrieverインスタンスをマージする関数

    :param retrievers: TFIDFRetrieverインスタンスのリスト
    :return: 統合された新しいTFIDFRetrieverインスタンス
    """
    # vectorizerのマージ
    vectorizers = [retriever.vectorizer for retriever in retrievers]
    merged_vectorizer = merge_tfidf_vectorizers(vectorizers, [retriever.docs for retriever in retrievers])
    # docsのマージ
    merged_docs = [doc for retriever in retrievers for doc in retriever.docs]
    # tfidf_arrayの計算とインスタンスの初期化
    merged_retriever = TFIDFRetriever.from_documents(merged_docs, tfidf_params=retrievers[0].vectorizer.get_params()
                                                     , k = retrievers[0].k) # kはretrievers間で全て同じと仮定

    return merged_retriever