ぼろぼろ平原

困った

gensim.corpora.Dictionaryの単語IDを辞書順にする

gensimのcorpora.Dictionaryは単語IDがランダムに振られてしまい困るのでSortableDictionaryというものを作った。

from gensim import corpora

class SortableDictionary(corpora.Dictionary):
    def __init__(self, *arg, **kwargs):
        super().__init__(*arg, **kwargs)

    def sort(self):
        idmap = dict(zip((v for k, v in sorted(self.token2id.items())), range(len(self.token2id))))
        self.token2id = dict((token, idmap[tokenid]) for token, tokenid in self.token2id.items())
        self.id2token = {}
        self.dfs = dict((idmap[tokenid], freq) for tokenid, freq in self.dfs.items())


### 使い方
# corpora.Dictionary(documents) としていたところをこのように置き換える
dictionary = SortableDictionary(documents)

# 辞書順にソートしたいときにsort()を呼び出す
dictionary.sort()

# 元のcorpora.Dictionaryと同じようにファイルに保存できる
dictionary.save_as_text("foo.txt")

# 保存したファイルは元のcorpora.Dictionaryでも読み込める
dictionary = corpora.Dictionary.load_from_text("foo.txt")