BERTの日本語事前学習済みモデルでテキスト埋め込みをやってみる
どうも、大阪DI部の大澤です。
汎用言語表現モデルBERTの日本語Wikipediaで事前学習済みのモデルがあったので、BERTモデルを使ったテキストの埋め込みをやってみたいと思います。
以下のエントリではbert-as-serviceを使ったテキストの埋め込みを紹介しました。今回はBERTのリポジトリで公開されているextract_features.pyを使って、テキストの埋め込みを試します。
BERT(Bidirectional Encoder Representations from Transformers)
BERTはGoogleが開発した汎用言語表現モデルです。機械学習で言語表現を学習させることができるため、言語に関わる様々な問題に使う事ができます。
モデル構成の概要
BERTはTransformerと言う機械翻訳モデルのエンコーダー部分を重ねたものになります。ベースモデルであれば12層なので、Transformerが12個重なります。Transformerの前段でトークンと文章とトークンの位置それぞれを埋め込み、その埋め込み表現がTransformerに入力されて、最終的にトークンそれぞれの言語表現(埋め込み表現)が出てくる感じです。
入力されるトークンの先頭は[CLS]
という特殊なトークンで、このトークンに対応する言語表現が文章全体の言語表現に相当します。なので、今回の目的である文章の埋め込み表現を取得するときは、このトークン部分を参照することになります。
学習済みモデルの概要
今回使用するのは京都大学の黒橋・河原研究室で公開されている学習済みBERTモデルです。モデルの紹介ページでは次のように説明されています。
- 入力テキスト: 日本語Wikipedia全て (約1,800万文)
- 入力テキストにJuman++ (v2.0.0-rc2)で形態素解析を行い、さらにBPEを適用しsubwordに分割
- BERT_{BASE}と同じ設定 (12-layer, 768-hidden, 12-heads)
- 30 epoch (1GPUで1epochに約1日かかるのでpretrainingに約30日)
- 語彙数: 32,000 (形態素、subwordを含む)
- max_seq_length: 128
事前学習(pretraining)に約30日...! こういったリソースを公開して頂けるのは非常に有難いですね。
やってみる
環境
- macOS Mojave
- Python 3.7.2
※ TensorFlowなどPythonのライブラリが必要になります。進めていく中でエラーが出る場合は必要に応じて、ライブラリをインストールしてください。
スクリプトの準備
まずは、以下のページから学習済みモデルをダウンロードします。zip形式なので、解凍して適当な場所に配置します。
BERTのリポジトリをクローンしてきます。
git clone https://github.com/google-research/bert
今回はJUMAN++(pyknp)でトークン化するため、bert/tokenization.pyを少し修正します。
JUMAN++を使ってトークン化するためのクラス定義JumanPPTokenizer
をスクリプトの最後に追加します。
class JumanPPTokenizer(BasicTokenizer): def __init__(self): """Constructs a BasicTokenizer. """ from pyknp import Juman self.do_lower_case = False self._jumanpp = Juman() def tokenize(self, text): """Tokenizes a piece of text.""" text = convert_to_unicode(text.replace(' ', '')) text = self._clean_text(text) juman_result = self._jumanpp.analysis(text) split_tokens = [] for mrph in juman_result.mrph_list(): split_tokens.extend(self._run_split_on_punc(mrph.midasi)) output_tokens = whitespace_tokenize(" ".join(split_tokens)) print(split_tokens) return output_tokens
tokenization.py
内のFullTokenizer
を修正します。元々BasicTokenizer
を使っていた箇所をJumanPPTokenizer
に変更します。
def __init__(self, vocab_file, do_lower_case=True): self.vocab = load_vocab(vocab_file) self.inv_vocab = {v: k for k, v in self.vocab.items()} # self.basic_tokenizer = BasicTokenizer(do_lower_case=do_lower_case) self.jumanpp_tokenizer = JumanPPTokenizer() self.wordpiece_tokenizer = WordpieceTokenizer(vocab=self.vocab) def tokenize(self, text): split_tokens = [] # for token in self.basic_tokenizer.tokenize(text): for token in self.jumanpp_tokenizer.tokenize(text): for sub_token in self.wordpiece_tokenizer.tokenize(token): split_tokens.append(sub_token) return split_tokens
extract_features.py
で出力されるテキストはjsonl形式なので、そこから埋め込み表現ベクトルを取り出して、tsvに書き出すスクリプトを用意しておきます。
import json import numpy as np # 参照するレイヤーを指定する TARGET_LAYER = -2 # 参照するトークンを指定する SENTENCE_EMBEDDING_TOKEN = '[CLS]' with open('/tmp/output.jsonl', 'r') as f: output_jsons = f.readlines() embedding_list = [] for output_json in output_jsons: output = json.loads(output_json) for feature in output['features']: if feature['token'] != SENTENCE_EMBEDDING_TOKEN: continue for layer in feature['layers']: if layer['index'] != TARGET_LAYER: continue embedding_list.append(layer['values']) np.savetxt('/tmp/output.tsv', embedding_list, delimiter='\t')
文章を埋め込む
準備ができたので、文章を埋め込んでみます。
extract_features.py
の引数は次の通りです。
- input_file: 入力するテキストを改行区切りで記述したテキストファイル
- output_file: 出力結果を書き出すjsonl形式のファイル
- vocab_file: 学習した単語(トークン)を記述したファイル(トークン化の際に使用)
- 学習済みモデルに含まれる
- bert_config_file: BERTモデルの設定ファイル
- 学習済みモデルに含まれる
- init_checkpoint: BERTモデルに読み込むチェックポイントファイル(学習した重み付けなど)
- 学習済みモデルに含まれる
- do_lower_case: トークン化時に文字を小文字にするかどうか
- layers: 必要なTransformerのエンコーダーは何層目のものか
- 複数のレイヤーの結果が欲しい場合はカンマ区切りで指定すればOK(例:-1,-2,-3)
- 基本的には後の方のレイヤーを使うことが多いと思うのでマイナスの値が便利(-1が最後のレイヤー)
- 最後のレイヤーだと事前学習のタスクの影響が大きすぎるかと思い、今回は最後から2つ目のレイヤーの出力を埋め込みに使用しています(とはいえ、可視化した図を見るとそこまで違いはなかったです)
今回使用するもの以外にもいくつか引数があります。詳細はスクリプトbert/extract_features.pyをご覧ください。
cat << EOS > /tmp/input.txt こんにちは おはよう こんばんは お腹すいた ご飯食べたい 明日晴れてると良いな 明日の天気はどうだろうか 雨降ったら嫌やな EOS export BERT_BASE_DIR="/Users/osawa.yuto/git/ml_pr/sandbox/bert/models/Japanese_L-12_H-768_A-12_E-30_BPE" python ./bert/extract_features.py \ --input_file=/tmp/input.txt \ --output_file=/tmp/output.jsonl \ --vocab_file=$BERT_BASE_DIR/vocab.txt \ --bert_config_file=$BERT_BASE_DIR/bert_config.json \ --init_checkpoint=$BERT_BASE_DIR/bert_model.ckpt \ --do_lower_case False \ --layers -2
数秒から数十秒くらいで処理が完了します。コンソールに出力される情報はかなり多いですが、エラーが吐かれて無ければ問題ないかと思います。
出力されたファイルを確認してみます。 形式は次のようになっています。jsonlなので、文章ごとに以下のjsonが一行ずつあります。
{ "linex_index": 入力された文章が何行目のものか(0ベース), "features": [ { "token": 対応するトークン, "layers": [ { "index": レイヤーの番号, "values": 対応する表現ベクトル } ] } ] }
head -1 /tmp/output.jsonl
jsonlから表現ベクトルだけ取り出してtsv形式で書き出します。
python convert_embedding_jsonl.py head -1 /tmp/output.tsv
可視化
Embedding Projectorで各文章のマッピングを可視化してみたいと思います。先ほど作成した/tmp/input.txt
と/tmp/output.tsv
を入力(Load)することで、簡単に可視化できます。今回はT-SNEで次元圧縮して3次元にマッピングしました。
塊ごとにみていきます。 天気に関連した文章が近くに集まっています。
"お腹すいた"と"ご飯食べたい"も近くにマッピングされていますね。
挨拶系の言葉も近くにまとまっています。
全体的に直感的なイメージに近いマッピングになっていました。
さいごに
今回は日本語Wikipediaで学習済みのBERTモデルを使用して、文章の埋め込みを行いました。もともと公式リポジトリでスクリプトが用意されていたため、トークン化用のスクリプトを修正するだけで出来ました。BERTは色々な自然言語タスクに応用することができるようなので、今後も何か試してみたいと思います。