gensimのWord2vecモデルから単語を間引く

しかし、word2vecファイルはサイズがデカいのが難点。

そこで、頻度や重要度の低い単語ベクトルを、word2vecの学習済みファイルから削除する後処理方法を紹介します。

以下のコードで学習とモデル保存が可能です。

モデル保存には2種類あって、「勾配情報も含めたニューラルネットワーク全体の保存」と、「単語ベクトルのみの保存」になります。

後者のほうがファイルサイズが小さくなりますが、再学習が出来なくなります。

from gensim.models import Word2Vec, KeyedVectors

model = Word2Vec(text, size=50, window=10, min_count=3, workers=2,) # 学習

model.save("word2vec.model") # モデル全体の保存
model.wv.save_word2vec_format("word2vec.kv", binary=False) # 単語ベクトルのみの保存

単語ベクトルのみの保存を行うメソッドはいくつかありますが、save_word2vec_formatメソッドで、binary=Falseにすることが重要です。

これにより、wordvec.kvの中身は、以下のような形式のテキストファイルになります。

5 5
の 1.1609643 2.8821776 -0.5543095 -0.2561126 0.5586935
、 -1.7030418 3.704897 2.3130968 -1.7750651 0.9364301
に -1.7022483 5.5853825 2.386544 1.0998139 3.3664994
た -7.971246 5.906819 2.8711908 7.466918 7.0728817
る 2.6325386 1.0028526 -0.81937 12.663451 1.4146792

最初の行の2つの数字が単語数と単語ベクトル次元数になります。
これは例ですので、実際は5単語/5次元なわけなく、30,000単語/300次元とかが一般的だと思います。
2行目以降は、1行あたり1単語の表層形と単語ベクトルが並んでいます。区切り文字は半角スペース。

これを、以下のメソッドを通して、KeyedVectorsオブジェクトとして読み込みます。

wv = KeyedVectors.load_word2vec_format("word2vec.kv", binary=False)

単語を間引く

何かの基準で不要語を削除していくのですが、KeyedVectorsオブジェクトをごにょごにょいじって書き換えようとするのは悪手で、テキストファイルであるword2vec.kv自体を書き換えるほうが早いです。

要はword2vec.kvが、上に載せたフォーマットと同じになっていればよいわけです。
そうなっていれば、gensimがすべてよしなに読み込んでくれます。

一方でword2vec.kvは、ただのテキストファイルなのでプログラミング言語の標準機能で読み込むことができます。
一旦、標準機能でテキストとして読み込み、いらない単語をif文なりで判定して削除した上で、同じフォーマットで再びテキストファイルに書き出すだけでOKです。