本記事は自然言語処理 #2 Advent Calendar 2019 21日目の記事です。
ちなみに本日誕生日だったりします。祝ってください。
~~誕生日に締め切り設定するとかなんというMの鑑~~

はじめに

単語埋め込み界ではこの1年でBERTが猛威を奮ってELMoですらあまり見なくなりつつありますが，
今でもword2vecやGloVeなどのレガシーな分散表現を使いたくなることがあります．
さらに，手持ちのデータで学習したくなることもあります(少なくとも僕は)
ということで，自分用にword2vec/doc2vec/GloVe/fastTextの学習キット作ったので公開します。

word2vec/doc2vec/fastTextはgensim、GloVeは公式実装のモデルを学習可能です．

使い方は各パッケージのREADMEに書いたので，ここでは学習キットの設計思想について書いてみます．

1. モデル学習用関数のAPI共通化

単語分散表現のライブラリ/パッケージはいろいろ出ていますが，
ライブラリごとにデータセットの想定フォーマットが違っていて，
ライブラリに適した形に整形する前処理のスクリプトを都度書いていると
コードがどんどん汚くなっていきます．
なので，テキストデータセット読み出し用のイテレータを共通化して，
関数内でライブラリごとに適したデータフォーマットに整形するようにしています．

def train_*****_model(
    output_model_path,
    iter_docs,
    **kwargs
)

word2vecの場合：

def train_word2vec_model(
    output_model_path,
    iter_docs,
    size=300,
    window=8,
    min_count=5,
    sg=1,
    epoch=5
):
    """
    Parameters
    ----------
    output_model_path : string
        path of Word2Vec model
    iter_docs : iterator
        iterator of documents, which are raw texts
    size : int
        size of word vector
    window : int
        window size of word2vec
    min_count : int
        minimum word count
    sg : int
        word2vec training algorithm (1: skip-gram other:CBOW)
    epoch : int
        number of epochs
    """

iter_docsは文書ごとの単語リストのイテレータです．

2. 任意のテキストデータセットから学習できるようにする

データセットの読み出しAPIを定義した抽象クラスTextDatasetBaseを用意しておく．
ユーザが使いたいデータセットの読み出しクラスをこのクラスを継承する形で実装することで任意のデータセットを扱えるようにします．

class TextDatasetBase(ABC):
    """
    a bass class for text dataset

    Attributes
    ----------
    """
    @abstractmethod
    def iter_docs(self):
        """
        iterator of documents

        Parameters
        ----------
        """
        yield None

MARD用データセットクラスの例

class MARDDataset(TextDatasetBase):
    def __init__(self, word_tokenizer):
        self.root_path = None
        self.word_tokenizer = word_tokenizer

    def iter_docs(self, dataset_path):
        """
        get iterator of texts in one document

        Parameters
        ----------
        dataset_path: string
            path to dataset
        """
        self.root_path = Path(dataset_path)
        reviews_json_fn = self.root_path / "mard_reviews.json"
        with open(reviews_json_fn, "r") as fi:
            for line in fi:
                review_dict = json.loads(line, encoding="utf-8")
                title = review_dict["reviewerID"]
                text = review_dict["reviewText"]
                yield self.word_tokenizer.tokenize(text)

pytorchのDataLoaderとかはこれの2億倍くらい洗練されている気がしますが，自分で思いついたのはこんなものでした．
もっと良い設計あったら教えてください．

使い方

インストール

word2vecを例にとって

git clone git@github.com:stfate/word2vec-trainer.git
cd word2vec-trainer

git submodule init
git submodule update
pip install -r requirements.txt

学習の実行

python train_text_dataset.py -o $OUTPUT_PATH --dictionary-path=$DIC_PATH --corpus-path=$CORPUS_PATH --size=100 --window=8 --min-count=5

モデルの使い方

model_path = "model/word2vec.gensim.model"
model = Word2Vec.load(model_path)

注意点

Wikipediaのような大規模データセットで学習させたときにメモリをバカ食いして落ちることがあるかもしれません．原因調査中．

おわりに

ライブラリのAPI考えるのたのしいよね