晴耕雨読

working in the fields on fine days and reading books on rainy days

scikit-learn で tf-idf を求める

Python で scikit-learn を使った tf-idf の求め方について説明します。

定義

TF とは Term Frequency の略で、単語の出現頻度を表します。

IDF とは Inverse Document Frequency の略で、逆文書頻度を表します。 この指標は、ある単語が多くの文書で出現するほど値は下がります。 つまり、特定の文書にしか出現しない単語の重要度を上げる役割を果たします。

TF と IDF を掛け合わせたものが TF-IDF です。 TF に対して IDF を掛けることで、IDF が一般語フィルタとして働きます。

ただし、実際にはここで定義した以外の計算式もあります。 色々な TF, IDF の定義は tf–idf - Wikipedia に書いてあります。 計算式の違いは、方言みたいなものです。

プログラム

実際に TF-IDF をプログラムで計算してみます。 たとえば、3つの文書(コーパス)が次のような内容だとします。

  • 文書1 :「ドキュメント集合においてドキュメントの単語に付けられる」
  • 文書2 :「情報検索において単語への重み付けに使える」
  • 文書3 :「ドキュメントで出現したすべての単語の総数」

これらを分かち書きしたデータを与えることで、TF-IDF を求めることができます。

from sklearn.feature_extraction.text import TfidfVectorizer

docs = [
    'ドキュメント 集合 において ドキュメント の 単語 に 付けられる',
    '情報検索 において 単語 へ の 重み付け に 使える',
    'ドキュメント で 出現した すべて の 単語 の 総数',
]
vectorizer = TfidfVectorizer(max_df=0.9) # tf-idfの計算
#                            ^ 文書全体の90%以上で出現する単語は無視する
X = vectorizer.fit_transform(docs)
print('feature_names:', vectorizer.get_feature_names())

words = vectorizer.get_feature_names()
for doc_id, vec in zip(range(len(docs)), X.toarray()):
    print('doc_id:', doc_id)
    for w_id, tfidf in sorted(enumerate(vec), key=lambda x: x[1], reverse=True):
        lemma = words[w_id]
        print('\t{0:s}: {1:f}'.format(lemma, tfidf))

出力結果:

feature_names: ['すべて', 'において', 'ドキュメント', '付けられる', '使える', '出現した', '情報検索', '総数', '重み付け', '集合']
doc_id: 0
	ドキュメント: 0.687703
	付けられる: 0.452123
	集合: 0.452123
	において: 0.343851
	すべて: 0.000000
	使える: 0.000000
	出現した: 0.000000
	情報検索: 0.000000
	総数: 0.000000
	重み付け: 0.000000
doc_id: 1
	使える: 0.528635
	情報検索: 0.528635
	重み付け: 0.528635
	において: 0.402040
	すべて: 0.000000
	ドキュメント: 0.000000
	付けられる: 0.000000
	出現した: 0.000000
	総数: 0.000000
	集合: 0.000000
doc_id: 2
	すべて: 0.528635
	出現した: 0.528635
	総数: 0.528635
	ドキュメント: 0.402040
	において: 0.000000
	付けられる: 0.000000
	使える: 0.000000
	情報検索: 0.000000
	重み付け: 0.000000
	集合: 0.000000

手計算で確認

たとえば文書1の単語「ドキュメント」は、文書1で2回出現し、単語が現れる文書数は2なので、 tf-idf を計算すると、次のようになります。

しかし、scikit-learn で出した値と一致しません。 調べてみると、どうやら TF-IDF には様々な定義があるようで、さらに scikit-learn の TF-IDF は標準的な式ではないらしいです。 詳しく書かれているブログとして、Python: scikit-learn と色々な TF-IDF の定義について - CUBE SUGAR CONTAINER があるので、こちらを参照してください。