bert-as-serviceを使って日本語BERTの文エンベディング計算サーバーを作る

はじめに

本記事ではbert-as-serviceを使い、日本語BERTによって文表現ベクトルを計算するサーバーを作ります。

準備

日本語BERTの事前学習済みモデルをディレクトリ bert-jp にダウンロードします。
ダウンロード後のディレクトリの中身は以下のとおりです：

$ ls bert-jp
graph.pbtxt                             model.ckpt-1400000.index                wiki-ja.model
model.ckpt-1400000.data-00000-of-00001  model.ckpt-1400000.meta                 wiki-ja.vocab

次に、bert-as-serviceでモデルをロード出来るようにするために、ファイルを編集します。
ファイル名を変更します:

mv model.ckpt-1400000.index bert_model.ckpt.index
mv model.ckpt-1400000.meta bert_model.ckpt.meta 
mv model.ckpt-1400000.data-00000-of-00001 bert_model.ckpt.data-00000-of-00001

語彙ファイルを作成します（2019/04/20: 先頭行を <unk> から [UNK]に置換する処理を追加）:

cut -f1 wiki-ja.vocab | sed -e "1 s/<unk>/[UNK]/g" > vocab.txt

以下の内容のBERT設定ファイルを作成します:

bert-jp/bert_config.json

{
    "attention_probs_dropout_prob" : 0.1,
    "hidden_act" : "gelu",
    "hidden_dropout_prob" : 0.1,
    "hidden_size" : 768,
    "initializer_range" : 0.02,
    "intermediate_size" : 3072,
    "max_position_embeddings" : 512,
    "num_attention_heads" : 12,
    "num_hidden_layers" : 12,
    "type_vocab_size" : 2,
    "vocab_size" : 32000
}

これで日本語BERTモデルの準備が出来ました。

最後にライブラリをインストールします。

pip install --upgrade
pip install tensorflow-gpu
pip install bert-serving-server  # server
pip install bert-serving-client  # client, independent of `bert-serving-server`

サーバーを立ち上げる

以下のコマンドのように、事前学習済みモデルが含まれるディレクトリと、ワーカーの数を指定して、サーバーを起動します。

bert-serving-start -model_dir ./bert-jp -num_worker=4 

クライアントからサーバーを叩く

サーバーに文の配列を送り、それぞれの文に対応するエンベディングの配列を受け取るクライアントを実装します。クライアントコードは以下の通りです。今回はローカルホストにサーバーを立てたので、ipを0.0.0.0にしています。また、日本語文の分かち書きには事前学習に使われたものと同じsentencepieceのモデルを利用しています。

client.py

from bert_serving.client import BertClient
bc = BertClient(ip='0.0.0.0')

import sentencepiece as spm
s = spm.SentencePieceProcessor()
s.Load('./bert-jp/wiki-ja.model')

def parse(text):
    text = text.lower()
    return s.EncodeAsPieces(text)

texts = ['液体の水は、惑星上における生命の前提条件である。',
        '火星は地球型惑星に分類される。']

parsed_texts = list(map(parse,texts))

print(bc.encode(parsed_texts,is_tokenized=True))

クライアントコードを実行します。

$ python client.py 
[[-0.01896668  0.20814717 -0.229494   ...  0.10796983  0.31211707
  -0.59373856]
 [-0.22285977  0.04470082 -0.48112744 ...  0.11185268  0.02520958
  -0.0751629 ]]

無事、エンベディングを受け取ることが出来ました。

感想

BERTは最近大きな話題を集めていますが、ファインチューニングしたモデルを使う場合には、推論時にも高性能なコンピューティングインスタンスが必要になり、用途によってはコストが高すぎることもあります。

この記事で紹介したようなエンベディング計算サーバーを立てれば、ファインチューニングは行わず、エンベディングを特徴量としてのみ利用することで、GPUが無い安価な推論インスタンスでもBERTを利用することが出来るようになります。