Sudachipyの学習済みword2vecを低メモリ環境で使用する

はじめに

日本語の言語処理を行う際に学習済みword2vecを利用したい場合，
いますぐ使える単語埋め込みベクトルのリスト
のように学習済みデータを入手しgensimを用いるのが簡単だった。
しかし，多くのものはwikipediaなどから学習されており単語数が50000程度だったりして，任意の文章を分かち書きしてベクトル化しようとしても，未知語ばっかりで使い物にならなかったりする。

日本語形態素解析といえばmecabだと思っていたが，最近はワークスアプリケーションが提供しているSudachiというのもある。
さらにSudachiで分かち書き・学習したWord2Vecモデルも提供されており，100 億語規模の国語研日本語ウェブコーパスから単語数360万の学習済みデータとなっている。
Sudachiベースの学習済みWord2Vecモデルを使う
かなり単語数が多いため使い勝手が良さそう。

単語と対応するベクトルが書かれたただのtxtファイルだが，圧縮時で5GB、解凍後で12GBくらいある巨大データである。
学習用のメモリがたくさんあるマシンであれば問題ないが，推論用の低メモリの環境では12Gをメモリに展開するのが難しい場合がある。

でも任意の単語に対してベクトル化するだけであれば，メモリに展開せずにその都度SSD(HDD)から呼び出せば良い。

sudachipyのインストール

https://github.com/WorksApplications/SudachiPy
githubに書いてあるとおりにインストールすればよいが，
環境によってはターミナルでsudachipyコマンドが認識されない場合があった。
sudachidict_fullをセットするためにはsudachipyディレクトリにあるconfig.pyのset_default_dict_packageを動かせば良い。

インストール先のsudachipyディレクトリに行きファイルの最後に下記のように追記し，ターミナルで実行すれば動作した。

config.py

#~~~~~~~~~~~
#~~~~~~~~~~~
import sys
output = sys.stdout
set_default_dict_package('sudachidict_full',output)

メモリに展開せずに学習済みデータを利用

この記事を参考に最初に各行のメモリの場所を覚えておく事により高速に任意の単語ベクトルを呼び出せる。
【Python】行数・列数が大きいCSVファイルから行番号指定で特定の行を効率よく読み出す方法についての検討

学習済みデータをURLから入手し好きな場所へ置いておく。
大規模コーパスと複数粒度分割による日本語単語分散表現

クラスにまとめて単語をベクトルにしたり，文章を分かち書きした後それぞれをベクトル化したものを返せるようにした。

sudachi_w2v.py

import numpy as np
import pickle
import re
import csv
import os
from sudachipy import tokenizer
from sudachipy import dictionary

class Sudachi_w2v:
    def __init__(self, path, sudachiDataPath="sudachiData.pickle"):
        f = open(path, 'r')
        self.file = f
        self.reader = csv.reader(f, delimiter=' ')
        #最初に含有単語リストやメモリアドレスリストを作成する（かなり時間かかる）
        #2回目以降はpickle化したものを読み込む
        if os.path.exists(sudachiDataPath):
            with open(sudachiDataPath, 'rb') as f:
                dataset = pickle.load(f)
            self.offset_list = dataset["offset_list"]
            self.emb_size = dataset["emb_size"]
            self.word2index = dataset["word2index"]
            self.ave_vec = dataset["ave_vec"]
        else:
            txt = f.readline()
            #分散表現の次元数
            self.emb_size = int(txt.split()[1])
            #未知語が来た場合平均ベクトルを返す
            self.ave_vec = np.zeros(self.emb_size, np.float)
            #メモリアドレスリスト
            self.offset_list = []
            word_list = []
            count = 0
            maxCount = int(txt.split()[0])
            while True:
                count+=1
                self.offset_list.append(f.tell())
                if count % 100000 == 0:print(count,"/",maxCount)
                line = f.readline()
                if line == '':break
                line_list = line.split()
                word_list.append(line_list[0])
                self.ave_vec += np.array(line_list[-300:]).astype(np.float)
            self.offset_list.pop()
            self.ave_vec = self.ave_vec/count
            self.word2index = {v:k for k,v in enumerate(word_list)}

            dataset = {}
            dataset["offset_list"] = self.offset_list
            dataset["emb_size"] = self.emb_size
            dataset["word2index"] = self.word2index
            dataset["ave_vec"] = self.ave_vec
            with open(sudachiDataPath, 'wb') as f:
                pickle.dump(dataset, f)

        self.num_rows = len(self.offset_list)
        #sudachiの準備
        self.tokenizer_obj = dictionary.Dictionary().create()
        self.mode = tokenizer.Tokenizer.SplitMode.C

    #単語をベクトル化
    def word2vec(self, word):
        try:
            idx = self.word2index[word]
            result = self.read_row(idx)
            vec = np.array(result[-300:])
            return vec
        except:#単語リストにない場合
            print(word, ": out of wordlist")

    #文章を分かち書きした後，それぞれのベクトルをmatでまとめて返す
    def sentence2mat(self, sentence):
        words = sentence.replace("　"," ").replace("\n"," ")
        words = re.sub(r"\s+", " ", words)
        input_seq = [m.surface().lower() for m in self.tokenizer_obj.tokenize(words, self.mode)]
        input_seq = [s for s in input_seq if s != ' ']

        mat = np.zeros((len(input_seq), self.emb_size))
        input_sentence = []
        for i, word in enumerate(input_seq):
            try:
                idx = self.word2index[word]
                result = self.read_row(idx)
                input_sentence.append(result[0])
                mat[i] = np.array(result[-300:])
            except:#単語リストにない場合平均ベクトルを返す
                input_sentence.append("<UNK>")
                mat[i] = self.ave_vec
        return input_sentence, mat

    def __del__(self):
        self.file.close()

    def read_row(self, idx):
        self.file.seek(self.offset_list[idx])
        return next(self.reader)

使い方は以下のとおり。
一番始めに含有単語リストやメモリアドレスリストを作成する。
これはかなり時間かかる。（数10分くらい）
一回作成した後は作成結果をpickle化しているため，2回目以降はpickleを読み込むことで数秒でインスタンスを作れる。

path = '~学習データの保存場所~/nwjc_sudachi_full_abc_w2v/nwjc.sudachi_full_abc_w2v.txt'
sudachi = Sudachi_w2v(path)

vec = sudachi.word2vec("すだち")
print(vec)
#['0.07975651' '0.08931299' '-0.06070593' '0.46959993' '0.19651023' ~

input_sentence, mat = sudachi.sentence2mat("あきらめたらそこで試合終了だよ")
print(input_sentence, mat)
#(['あきらめ', 'たら', 'そこ', 'で', '試合終了', 'だ', 'よ'], array([[ 1.9788130e-02,  1.1190426e-01, -1.6153505e-01, ...,

sudachiの学習済みデータはかなり単語数が多いため，大抵の単語は変換できるしいろいろ使いやすそう。