gensimで「文書を全部メモリに載せる」をやめてジェネレータでやるようにした

python 自然言語処理

10GB強のテキストファイルを自然言語処理するにあたり、「とりあえずcsvファイルをpandasで読み込む」というアプローチができないため、きちんとジェネレータを使って「ファイルを１行ずつ読む」という処理を実装した記録です。

なお、「csvファイルをpandasで」が好ましくないことはgensim公式がきちんと名言していました。1

コード

次の処理を行うサンプルコードを用意しました。

  • テキストファイルに対してジェネレータを生やして、
  • gensimの辞書を作成し、
  • 辞書を元に生テキストをBag-of-Wordsへ変換するジェネレータを生やして、
  • トピックモデル訓練する。
  • 最後に、生テキストをトピックモデル のベクトルへ変換するジェネレータを生やす。
def tokenize_text(some_str):
    # 文字列を正規化したり、形態素解析したりする処理

docs = (tokenize_text(line) for line in open('documents.txt'))
dic = corpora.Dictionary(docs)

class Corpus2bow(object): # イテレータとして呼ばれるたびにBag-of-Wordsのジェネレータを渡す
    def __init__(self, file_path):
        self.file_path = file_path
    
    def __iter__(self):
        for line in open(self.file_path):
            yield dic.doc2bow(tokenize_text(line))

corpus = Corpus2bow("documents.txt")
lda = LdaModel(corpus=corpus, num_topics=50, id2word=dic) #トピックモデル を学習

corpus_lda = lda[corpus] # トピックモデルのベクトルを得るジェネレータ

参考情報１：まずgensim公式のお手本をおさらい

公式の introduction と tutorial を読み直しました。
特に「メモリに優しく」という側面で次の点が重要でした。

  1. ジェネレータでテキストファイルを読み込めば、そのままgensimの辞書やモデル訓練に使える。
  2. 学習したモデルを使って文書を潜在意味空間上のベクトルへと写像する部分も、ジェネレータを使って処理できる。

つまり、生テキスト→トークン化→Bag-of-Words→トピックモデルのベクトル という変換を「生テキスト１行ずつイテレーティブに」行うことができます。

radimrehurek.com

参考情報２：Effective Pythonの項目17

ジェネレータが一度末尾に到達すると、以後はStopIteration例外が起きるため、同一のジェネレータを使って再度イテレーションをすることができません。
よって、本書を参考にして「イテレータとして呼ばれるたびにジェネレータを返すクラス」を作りました。ちなみにgensimのチュートリアル中でも似た実装が紹介されていました。

  1. gensim公式のintroductionにMemory independence – there is no need for the whole training corpus to reside fully in RAM at any one timeとある↩