Python
自然言語処理
scikit-learn
tf-idf

IDFをscikit-learnのライブラリに頼らずに計算する

以前、scikit-learnのライブラリ(TfidfVectorizer)を使って単語のIDF値を計算する場合、TfidfVectorizer側で行っている単語の正規化を意識しないと、自分の意図したIDF値が得られない問題があることについて簡単に触れました。
なので、今回はIDF値をscikit-learnのライブラリを使わずに計算してみようと思います。(TFは求めません。)

IDFの定義

wikipediaを見るとよいです。

ちなみにTfidfVectorizerでは、デフォルトでIDFの定義に以下のような修正が加えられています。

IDF_i = \log \frac{|D| + 1}{{ d : d \ni t_i} + 1} + 1

これは $0$ で割ることを防ぐためと、IDF値が $0$ になってしまったとき、TF-IDFの値が $0$ になってしまうのを防ぐのが目的のようです。今回の実装はこの定義に則って実装してみます。

実装

サンプル文書は以下とします。(テキトーな文章で申し訳ございません...)

D = [
    ["hoge", "fuga", "foo", "bar"],
    ["hoge", "foo", "hogehoge", "bar"],
    ["fuga","foofoo","barhoge", "bar"],
    ["foo", "hoge"],
    ["nyaa", "hogehoge","bar", "foo", "nyaa"],
    ["foo","hoge","barhoge","foo","foo"]
    ]

各文章を形態素解析により、分かち書きの配列としたあとにTF-IDFを求めることを想定しています。

TfidfVectorizerで計算

fit_transformしたあとに_tfidf.idf_にアクセスすればIDF値が取得できます。今回はわかりやすく出力するために辞書型にして出力してみました。

corpus = []
for w in D:
    corpus.append(" ".join(w))
tfv = TfidfVectorizer (dtype=np.float32)
tfv.fit_transform(corpus)
featurenames = tfv.get_feature_names()
idf = tfv._tfidf.idf_
word_idf_dict = {}
for pair in zip(featurenames, idf):
    word_idf_dict[pair[0]] = pair[1]
print(word_idf_dict)
#{'bar': 1.336472236621213,
# 'barhoge': 1.8472978603872037,
# 'foo': 1.1541506798272583,
# 'foofoo': 2.2527629684953681,
# 'fuga': 1.8472978603872037,
# 'hoge': 1.336472236621213,
# 'hogehoge': 1.8472978603872037,
# 'nyaa': 2.2527629684953681}

IDF

愚直ロジック版

愚直に定義に従って計算するなら以下のような感じで求められます。

def idf(word, document):
    n_samples = len(document)
    df = np.sum(np.array([int(word in d) for d in document], dtype="float32"))
    n_samples += 1.0
    df += 1.0
    return np.log(n_samples / df) + 1.0

def getIDFval(in_corpus):
    word_idf_dict = {}
    featurenames = list(set([w for s in in_corpus for w in s]))
    for w in featurenames:
        word_idf_dict[w] = idf(w, in_corpus)
    return word_idf_dict, featurenames

word_idf_dict, featurenames = getIDFval(D)
print(word_idf_dict)
#{'bar': 1.336472236621213,
# 'barhoge': 1.8472978603872037,
# 'foo': 1.1541506798272583,
# 'foofoo': 2.2527629684953681,
# 'fuga': 1.8472978603872037,
# 'hoge': 1.336472236621213,
# 'hogehoge': 1.8472978603872037,
# 'nyaa': 2.2527629684953681}

TfidfVectorizerと同じ値が求まりました。

文書数や単語数が少ない場合はこんな感じでさらっと求めればいいわけですが、如何せん処理が愚直過ぎて、文書数や単語数が数万とかになると、TfidfVectorizerと比べてかなり時間がかかってしまいます。

模索版

現在TfidfVectorizerと同じくらいの速度で計算するロジックを模索中です。以下のような感じはどうかと試してみました。

featurenames = list(set([w for s in D for w in s]))
n_samples = len(D)
word_dict = {}
count = 1
for w in featurenames:
    word_dict[w] = count
    count +=1
id_corpus = []
for d in D:
    id_corpus.append(list(set([word_dict[w] for w in d])))
num_feature = len(featurenames)
X = []
for corpus in id_corpus:
    b = np.zeros(num_feature)
    for i in corpus:
        b[i-1] = 1
    X.append(b)
X = np.array(X)
word_idf_dict = {}
for w in featurenames:
    df = np.sum(X[:,word_dict[w] - 1])
    word_idf_dict[w] = np.log((n_samples + 1.0) / (df + 1.0) ) + 1.0
print(word_idf_dict)
#{'bar': 1.336472236621213,
# 'barhoge': 1.8472978603872037,
# 'foo': 1.1541506798272583,
# 'foofoo': 2.2527629684953681,
# 'fuga': 1.8472978603872037,
# 'hoge': 1.336472236621213,
# 'hogehoge': 1.8472978603872037,
# 'nyaa': 2.2527629684953681}

まず単語に連番を割り振って単語をID化します。
次に各文書に各単語が含まれているかどうかを$0$,$1$で置き換えます。
そうすることで上記の行列 $X$ は 行が文章を表し、列が各単語が含まれているかどうか、を表すことになるので、単語$i$を含む文書数(IDFの分母)を求めるには行列$X$の列ベクトルの要素の合計で求めることができます。
基本的にIDFの分母を求めるのに時間がかかるわけですが、上記のロジックでも愚直版よりかは早いですが、それでもかなり時間がかかってしまう。

比較

  • 文書数約37,000
  • 単語数約66,000 のデータでTfidfVectorizerと上記の模索版で処理速度を算出してみると、
モデル 処理速度(s)
TfidfVectorizer 1.57
愚直版 2020.45
模索版 143.48

愚直版がめちゃくちゃ遅い(当たり前ですが)
模索版は愚直版と比べてかなり高速になってはいるのですが、TfidfVectorizerと比べて模索版もクソロジックであることはよくわかりました...
TfidfVectorizerのソースをよく読んで真似して実装したらそれでいいわけですが、
どなたかすっきりした実装でかつ高速にIDF値を計算できるロジックを知っている方いたらぜひご教授いただけたらと思います...

終わり。