word2vec/doc2vec/GloVe/fastTextの学習キット作った

本記事は[自然言語処理 #2 Advent Calendar 2019](https://qiita.com/advent-calendar/2019/nlp2) 21日目の記事です。
ちなみに本日誕生日だったりします。祝ってください。
~~誕生日に締め切り設定するとかなんというMの鑑~~

# はじめに

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

- [word2vec-trainer](https://github.com/stfate/word2vec-trainer)
- [doc2vec-trainer](https://github.com/stfate/doc2vec-trainer)
- [fastText-trainer](https://github.com/stfate/fasttext-trainer)
- [GloVe-trainer](https://github.com/stfate/GloVe-trainer)

word2vec/doc2vec/fastTextは[gensim](https://radimrehurek.com/gensim/)、GloVeは[公式実装](https://nlp.stanford.edu/projects/glove/)のモデルを学習可能です．

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


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

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

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

word2vecの場合：

```python
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`を用意しておく．
ユーザが使いたいデータセットの読み出しクラスをこのクラスを継承する形で実装することで任意のデータセットを扱えるようにします．

```python
class TextDatasetBase(ABC):
    """
    a bass class for text dataset
    
    Attributes
    ----------
    """
    @abstractmethod
    def iter_docs(self):
        """
        iterator of documents
        
        Parameters
        ----------
        """
        yield None
```

[MARD](https://www.upf.edu/web/mtg/mard)用データセットクラスの例

```python
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を例にとって

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

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

## 学習の実行

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

## モデルの使い方

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


# 注意点

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


# おわりに

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