本記事は自然言語処理 #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考えるのたのしいよね