TfidfVectorizerのマージの高速化

概要

TfidfVectorizerをマージするときになるべく高速にできる方法を検討しました。
ポイントは以下です。

• 形態素解析時間の短縮: サブセットのvectorizerの作成時には形態素解析をfitの外で実施し、形態素解析の結果を保存しておく。
• IDF計算時間の短縮: 各サブセットでIDFから単語の登場回数を逆算したうえでマージする。

今回の実験では、何も工夫せず全データセットで再fitした場合と比べ、文書10万件の場合に約1/10、100万件の場合に約1/25に処理時間を減らすことができました。

背景

TF-IDFは文書の特徴を捉えるのに便利な指標ですが、IDFがコーパスごとに決定されるため、文書集合が更新されるたびにIDF辞書を作り直す必要があり、処理に時間がかかります。

例えば、以下の検証では100万件の文書をTfidfVectorizer.fitするのに形態素解析（tokenize）とidf計算でそれぞれ48秒、21秒ほどかかっています。

参考:TfidfVectorizer.fitでtokenizeとidf計算にかかる時間の計測

完全に新規なコーパスであれば計算するしかないですが、もしコーパスのサブセットについてIDF計算が完了しているのであれば、それを活用して、計算量を減らせそうな気がします。

そこで、サブセットに対してfitしたvectorizer（より厳密にはIDF辞書）をマージして1つにするときの効率的なやり方を検討しました。

準備

ライブラリをインストールします。

pip install scikit-learn datasets mecab-python3

大きな日本語のデータセットを用意します。
今回はmiracl-corpusを使います。700万件弱の文章です。

import time
import MeCab
from sklearn.feature_extraction.text import TfidfVectorizer
from datasets import load_dataset
import tqdm
dataset = load_dataset("miracl/miracl-corpus", "ja")["train"]
dataset = [d["text"] for d in dataset]
print(len(dataset))

6953614

tokenize用の関数を用意します。今回はmecabを使います。

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

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

コーパスから取得した5万ずつのサブセットでTfidfVectorizerをfitします。
高速化に役立てるため、tokenizeとfitを別々に実施します。

dataset1 = dataset[:50000]
dataset2 = dataset[50000:100000]

# あとで使いまわせるように先にtokenizeしておく
tokenized_dataset1 = [tokenize(text) for text in dataset1]
tokenized_dataset2 = [tokenize(text) for text in dataset2]

# tokenize済みのコーパスを入力するので、入力をそのまま返すだけのanalyzerを定義
tfidf1 = TfidfVectorizer(analyzer=lambda x: x)
tfidf2 = TfidfVectorizer(analyzer=lambda x: x)

# 各データセットに対してベクトル化を実行（例）
tfidf1.fit(tokenized_dataset1)
tfidf2.fit(tokenized_dataset2)

# もしtfidf1やtfidf2をそのまま使う場合は、set_paramでanalyzerを置換可能。今回の実験では不要
# tfidf1.set_params(analyzer=tokenize)
# tfidf2.set_params(analyzer=tokenize)

処理時間を気にしない場合

TfidfVectorizerをマージするにあたって最も簡単な方法は全体のdatasetで再度fitすることです。

merged_tfidf = TfidfVectorizer(analyzer=tokenize)
start_time = time.time()
merged_tfidf.fit(dataset1+dataset2)
simple_fit_time = time.time() - start_time
print(f"simple fit time: {simple_fit_time}")

simple fit time: 5.427338123321533

形態素解析時間の短縮

fitで時間のかかる処理の一つは形態素解析です。
サブセットのvectorizerの作成時に、tokenizeされたサブセットを保存しておけば、マージにて再利用が可能です。vectorizer自体にはtokenizeされた文書の情報は保持されていないため、別途保管する必要があります。

merged_tfidf = TfidfVectorizer(analyzer = lambda x: x)
start_time = time.time()
merged_tfidf.fit(tokenized_dataset1+tokenized_dataset2)
merged_tfidf.set_params(analyzer=tokenize) # analyzerを戻しておく
fit_time = time.time() - start_time
print(f"fit time: {fit_time}")

fit time: 1.606766939163208

形態素解析の時間が短縮され、所要時間が1/3くらいになりました。

idf計算時間の短縮

前節の方法だと、IDFはtokenizeされたすべての単語の登場回数を確認し、再計算されています。したがって単語数をLとしてO(L)の計算量です。
一方で、サブセットごとのIDFは計算済みです。IDFは単語を含む文書の数の逆数に基づく値のため、IDFからサブセットにおける単語の登場文書数を計算できます。したがって単語の種類の数だけ計算を行えば良いことになります。単語種類数をKとするとO(K)の計算量です。

今回用いているmiracl-corpusでは10万件の時点で、KがLの1/5くらいです。KのほうがLより増えるペースが遅いため、この差は件数が増えるほど広がり、O(K)で計算する利点が向上します。

サブセットごとにIDFから単語の登場文書数を逆算し、マージされたIDFを計算するコードは例えば以下のようになります。

import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer

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

上記関数のポイントをいくつか述べます。

• vectorizerのリストとそのvectorizerたちを作るのに用いたdatasetのリストを入力とし、マージされたIDF辞書とvocabularyをもつvectorizerを返す関数です。
• IDFから出現回数を計算するには、コーパスのサイズが必要なため、datasetのリストを合わせて入力しています。結局コーパスサイズ(len)しか用いてないので、最初からlenを入力させるようにするのもありです。
• idf_to_doc_freqが出現回数を計算する関数です。今回はsmooth_idf=Trueの場合の計算式です。vectorizerを作ったときの設定に応じて書き換える必要があります。
• mergeされたvectorizerのパラメータは、入力vectorizerリストの0番目のパラメータを引き継ぎます。入力されたvectorizerのパラメータはすべて同じという前提です。

処理時間を計測してみます。

start_time = time.time()
merged_tfidf = merge_tfidf_vectorizers([tfidf1, tfidf2], [dataset1, dataset2])
fit_time = time.time() - start_time
print(f"fit time: {fit_time}")

fit time: 0.47629809379577637

短くなりました。1/5よりは少し多いですが、サブコーパスごとの単語の種類数の計算が発生しているので、そんなものかもしれません。

精度評価

IDFから文書頻度を計算することによって生じる誤差が検索結果にどの程度影響するか、別記事で軽く検証してみました。

結論としては、全体的に異なる検索結果となっていました。ただし実タスクを想定した評価では結果が変わる可能性もあります。

おまけ

今回の手法は文書数が多いほど恩恵を受けられますので、件数を１０倍にした50万件のサブセットの統合もやってみます。

dataset1 = dataset[:500000]
dataset2 = dataset[500000:1000000]

# あとで使いまわせるように先にtokenizeしておく
tokenized_dataset1 = [tokenize(text) for text in dataset1]
tokenized_dataset2 = [tokenize(text) for text in dataset2]

# tokenize済みのコーパスを入力するので、入力をそのまま返すだけのanalyzerを定義
tfidf1 = TfidfVectorizer(analyzer=lambda x: x)
tfidf2 = TfidfVectorizer(analyzer=lambda x: x)

# 各データセットに対してベクトル化を実行（例）
tfidf1.fit(tokenized_dataset1)
tfidf2.fit(tokenized_dataset2)

merged_tfidf = TfidfVectorizer(analyzer=tokenize)
start_time = time.time()
merged_tfidf.fit(dataset1+dataset2)
simple_fit_time = time.time() - start_time
print(f"simple fit time: {simple_fit_time}")

simple fit time: 58.8294517993927

merged_tfidf = TfidfVectorizer(analyzer = lambda x: x)
start_time = time.time()
merged_tfidf.fit(tokenized_dataset1+tokenized_dataset2)
merged_tfidf.set_params(analyzer=tokenize) # analyzerを戻しておく
fit_time = time.time() - start_time
print(f"fit time: {fit_time}")

fit time: 20.96849012374878

start_time = time.time()
merged_tfidf = merge_tfidf_vectorizers([tfidf1, tfidf2], [dataset1, dataset2], tfidf_params={"analyzer": tokenize})
fit_time = time.time() - start_time
print(f"fit time: {fit_time}")

fit time: 2.217916965484619

何も考えずにfitした場合と比べて1/25ほどの時間短縮になっています。