はじめての自然言語処理

今回は BERT における事前学習の改良手法である ELECTRA の検証です。ELECTRA はモデルサイズ、データ、計算量が同一条件であればオリジナルの BERT を凌ぐ性能とのことなので結果が楽しみなところです。事前学習をした後のファインチューニングは、いつも livedoor News Corpus の文書分類ばかりだったので、今回は固有表現抽出を試すことにしました。

1. はじめに

今回は BERT における事前学習の改良手法である ELECTRA ¹ の検証です。 BERT に関しては 第3回 で取り上げていますが、トークン化が Sentencepiece である為、トークン単位での処理に難がありました²。今回は ELECTRA を試すにあたり、そのあたりの対応も入れ、 Megagon Labs さんから公開されている UD_Japanese-GSD v2.6-NE ³ を使って固有表現抽出の実験をしたいと思います。

BERT が初めての方は、 第3回 の2章, 3章あたりに目を通して戻ってきて頂けると、より理解がしやすいと思います。

日本語で BERT を使った固有表現が使いたいだけであれば、 HuggingFace の Transformers を使ったほうが手軽なのですが、今回 ELECTRA の検証を行ったのは 「 より品質の良い事前学習済みモデルが手に入るかも」と欲がでたこと、Google から公開される BERT がらみのコードを動かそうとすると、Transformers では都合が悪いこと⁴ などが理由でしょうか。

それでは ELECTRA の紹介からはじめていきましょう。

2. ELECTRA

ELECTRA は 2020年に Google が提案した BERT の事前学習手法の改良です。 BERT の事前学習の改良手法としては RoBERTa ⁵ がありますが、 ELECTRA もその類です。 ですので、事前学習後に出来上がるモデルの構造は基本的に BERT と同一になります。

オリジナルの BERT は “[MASK]” トークンで虫食いにした入力文の穴埋めを行う Masked LM と “[SEP]” で結合した前後二つの文が連続しているか否かを判定する Next Sentence Prediction を目的関数としていました。 ELECTRA では以下のように大小二つの BERT を使い、見た目がGANっぽい学習をします。

動きとしては、以下のようになります。

1. 入力シーケンスの 15%(SMALL, BASE サイズの場合) もしくは 25%(LARGE サイズの場合) を “[MASK]” に置き換え。
2. Generator が “[MASK]” トークンの穴埋めをする。
3. Discriminator は Generator が穴埋めしたシーケンスを入力とし、各トークンが本物（元の文章のママ）か偽物（Generator が穴埋めした）かを判断する。

GAN「っぽい」とは見ため的に Generator と Discriminator がいて、 Discriminator が真偽判定するのが「っぽい」という話であり、敵対的学習をしている訳ではありません。 Generator は Masked LM で学習しており、 Disicriminator を欺こうとはしていません。

Generator はそれっぽいトークンを予測しつつも、いい塩梅に間違ってもらわないといけないので、Discriminator の 1/4 (SMALL, LARGE サイズの場合) もしくは 1/3 (BASE サイズの場合) の規模になっています。

細かい数式の内容は省きますが、事前学習の目的関数は以下のようになります。

L_MLM が Masked LM による Generator 側の学習、 L_Disc がトークンの真偽判定による Discriminator 側の学習、λ はハイパーパラメータです。また、Generator と Discriminator の間にはサンプリングのステップ（上図で言うと真ん中の “the” と “ate” ですね）が入っているので、Discriminator の損失は Generator へは伝播されません。 最後に、Next Sentence Prediction はオミットされています。これは RoBERTa もそうでしたね。

事前学習が終わったら、Generator のパラメータは捨ててしまい、Discriminator のみを用いることになります。

ELECTRA の利点

オリジナルの BERT の Masked LM では “[MASK]” に置換された全体の15%のトークンからしか学習できていなかったのですが、トークンの真偽判定の形とすることで全ての入力トークンを学習に利用できるようになり、計算効率が向上しています。また Discriminator へ “[MASK]” が入力されることがなくなり、BERT に存在した事前学習時とファインチューニング（及び後続タスク）時のミスマッチが解消されています。

ELECTRA の論文では ¼ 程度の計算量で RoBERTa, XLNet に近い性能を出し、同じ計算量であれば、それを超える性能が確認されたとのことです。

それでは、日本語データを使って実際に ELECTRA の学習をしてみましょう。

3. 事前学習

それでは ELECTRA の核心部分である事前学習からはじめていきましょう。いつものように、記事内のコードスニペットは、特に断りがない場合は Google Colaboratory (以下、Colab)で動かす想定にしています。

セットアップ

まずは、"ランタイム“ -> “ランタイムタイプの変更” で “ハードウェアアクセラレータ” を “TPU” にしてください。 ELECTRA のコードは Tensorflow 1.15 が前提になっているので、以下のマジックコマンドを動かしておきます。

%tensorflow_version 1.x 

加工済みの学習データやチェックポイントは GCS に保存するので認証しておきます。

from google.colab import auth
auth.authenticate_user()

さて、今回は分かち書きが単語境界を跨がないようにしたいです。 どうしようかと思案した結果、HuggingFace Transformers の Tokenizer の API が ELECTRA のそれとよく似ていたので、Transformers の日本語 BERT のトークナイザを流用することにしました。そんな訳で Transformers と MeCab 関連をインストールします。

!pip install transformers
!apt-get install mecab mecab-ipadic-utf8
!pip install mecab-python3
!pip install fugashi ipadic

最後に ELECTRA のコードは Github で公開されているのでクローンします。

!git clone https://github.com/google-research/electra

これでセットアップが完了です。次は事前学習データです。

事前学習データの準備

事前学習の元ネタには Tensorflow Datasets(以後、TFDS) の日本語 Wikipedia を使用しました。 そのままでは見出し的な文章の体をなしていないテキストが含まれているので、シンプルに「。」で終了しない行を除外、記事と記事は空行で区切って一旦テキストに落とします。

import tensorflow_datasets as tfds
import tensorflow as tf

ds = tfds.load(name='wikipedia/20190301.ja', 
          shuffle_files=False,
          download=True,
          try_gcs=True)

train_ds = ds["train"].batch(128).prefetch(10)

all_titles = []
all_texts = []

for example in tfds.as_numpy(train_ds):
  titles, texts = example["title"], example["text"]
  for title, text in zip(titles, texts):
    all_titles.append(title.decode('utf-8'))
    all_texts.append(text.decode('utf-8'))

with open("input.txt", "w") as f:
  for text in all_texts:
    lines = [line.strip() for line in text.split("\n")]
    for line in lines:
      if len(line) == 0 or not line.endswith("。"):
        continue
      f.write(line + "\n")
    f.write("\n")

del all_titles
del all_texts

次に以下のような加工をしました。

• 行末の空白は除去、空白のみの行は削除
• "。” の後が"」"、")“、"）”,“]"だった場合、"。"の後で改行
• "。"で始まる行は削除
• "Category:” で始まる行、"thumb|“を含む行は除去

%%bash
cat << EOF > preprocess.sh
#!/bin/bash
FILE=\$1
if [ \$# -ne 1 ]; then
  echo "Usage: ./preprocess.sh INPUT_TEXT"
  exit 1
fi
echo "Processing \${FILE}"
sed -i -e '/<doc id/,+1d; s/<\/doc>//g' \${FILE}
sed -i -e 's/ *$//g; s/。\([^」|)|）|"]\)/。\n\1/g; s/^[ ]*//g' \${FILE}
sed -i -e '/^。/d' \${FILE}
sed -i -e '/^Category/d; /thumb|/d' \${FILE}

EOF
chmod 744 preprocess.sh
./preprocess.sh input.txt

データが出来たらディレクトリを作って移動しておきます。

!mkdir data
!mv input.txt data

次に vocab.txt を取得します。これは Google が英語で学習したときの語彙ファイルだと思うのですが、なんでこれが必要なんだか思い出せません。 トークナイザを差し替えるので要らないんだけれど、ないと入力チェックでエラーになるとか、そんな話だったような気が。。。。

!curl -o vocab.txt https://storage.googleapis.com/electra-data/vocab.txt 

ELECTRA のコードをコードの検索パスに追加して、前処理のパラメータとなる名前付きタプルを定義します。

import sys
sys.path.append("./electra")
from collections import namedtuple

from electra.build_pretraining_dataset import *

Args = namedtuple('Args', (
    'corpus_dir',
    'vocab_file',
    'output_dir',
    'max_seq_length',
    'num_processes',
    'blanks_separate_docs',
    'do_lower_case'
  ))

パラメータの定義はこんな感じです。

args = Args(
    corpus_dir = './data',
    vocab_file = './vocab.txt',
    output_dir = 'gs://somewhere/electra/max_seq_length_512/pretrain_tfrecords',
    max_seq_length = 512,
    num_processes = 1,
    blanks_separate_docs = True,
    do_lower_case = True
    )

トークナイザとして Transformers の日本語 BERT のトークナイザを借りてしまいます。

from transformers import BertJapaneseTokenizer
bert_japanese_tokenizer = BertJapaneseTokenizer.from_pretrained('cl-tohoku/bert-base-japanese-whole-word-masking')

次に ExampleBuilder のコンストラクタを書き換えて、トークナイザを差し替えます。

init_org = ExampleBuilder.__init__

def __init__(self, tokenizer, max_length):
    self._tokenizer = bert_japanese_tokenizer
    self._current_sentences = []
    self._current_length = 0
    self._max_length = max_length
    self._target_length = max_length

ExampleBuilder.__init__ = __init__

後は動かすだけです（けっこう時間がかかります）。

utils.rmkdir(args.output_dir)
write_examples(0, args)
# Job 0: Creating example writer
# Job 0: Writing tf examples
# Job 0: Done!

学習データの準備ができたので事前学習を回しましょう。

事前学習の実行

まずは ELECTRA をロードしてトークナイザを取得します。

import sys
sys.path.append("./electra")
from collections import namedtuple
from transformers import BertJapaneseTokenizer
bert_japanese_tokenizer = BertJapaneseTokenizer.from_pretrained('cl-tohoku/bert-base-japanese-whole-word-masking')

次に語彙ファイルを作成して保存しておきます。

with open("vocab.txt", "w") as f:
  f.write("\n".join(bert_japanese_tokenizer.vocab))
!gsutil cp vocab.txt gs://somewhere/electra/max_seq_length_512/

設定用のパラメータクラスと引数用の名前付きタプルを準備します。

from electra.run_pretraining import *
from electra.configure_pretraining import PretrainingConfig

Args = namedtuple('Args', (
    'data_dir',
    'model_name',
    'hparams'
  ))

Colab の Free TPU で回すので、設定したパラメータはこんな感じです。 今回は BASE サイズでシーケンス長 512 のモデルとしました。train_batch_size は論文に記載のある 256 では OOM で落ちたので、128 としました。 num_train_steps は train_batch_size が半分で学習データ量も少なく悩ましいですが、"とりあえず” 論文どおりの 766000 にしています。

args = Args(
    data_dir = 'gs://somewhere/electra/max_seq_length_512/',
    model_name = 'electra_base_wiki_ja',
    hparams = '''
      {
        "use_tpu": true, 
        "num_tpu_cores": 8,  
        "model_size": "base", 
        "max_seq_length": 512,
        "generator_hidden_size": 0.33333,
        "learning_rate": 2e-4,
        "train_batch_size": 128,
        "embedding_size": 768,
        "num_train_steps": 766000,
        "vocab_file": "gs://rddl-nlp/electra/max_seq_length_512/vocab.txt",
        "vocab_size": %d,
        "save_checkpoints_steps": 1000
      }
    ''' % bert_japanese_tokenizer.vocab_size
)

model_size パラメータを "base" と指定すれば関連パラメータが適切に設定される風なのですが、じつはソースコードは以下のようになっています。 最初、論文記載のパラメータ値がわざわざコメントにしてあるので、「論文の値よりおいしい設定を見つけたのかな？」と "model_size"="base" だけ設定して実験してみたところ、結果はボロボロでした。。。そんな訳で上記のパラメータは model_size の他に embedding_size、generator_hidden_size を明示的に上書き設定しています。

# https://github.com/google-research/electra/blob/master/configure_pretraining.py#L114-L125 から引用
114    # defaults for different-sized model
115    if self.model_size == "small":
116      self.embedding_size = 128
117    # Here are the hyperparameters we used for larger models; see Table 6 in the
118    # paper for the full hyperparameters
119    # else:
120    #   self.max_seq_length = 512
121    #   self.learning_rate = 2e-4
122    #   if self.model_size == "base":
123    #     self.embedding_size = 768
124    #     self.generator_hidden_size = 0.33333
124    #     self.train_batch_size = 256

定義した設定値で config.json を作って保存します。

config = PretrainingConfig(args.model_name, args.data_dir, **json.loads(args.hparams))

from util import training_utils
bert_config = training_utils.get_bert_config(config)
with open("config.json", "w") as f:
  f.write(bert_config.to_json_string())
!gsutil cp config.json gs://somewhere/electra/max_seq_length_512/

ちなみに中身はこんな感じです。

!cat 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
# }

その他、必要なものをインポートして、

import os
import pprint
import json
import tensorflow as tf

TPUアドレスを確認します。

assert 'COLAB_TPU_ADDR' in os.environ, 'ERROR: Not connected to a TPU runtime; please see the first cell in this notebook for instructions!'
TPU_ADDRESS = 'grpc://' + os.environ['COLAB_TPU_ADDR']
print('TPU address is', TPU_ADDRESS)

# TPU address is grpc://10.121.78.10:8470

Colab の Free TPU は 8コアです。

with tf.Session(TPU_ADDRESS) as session:
  print('TPU devices:')
  pprint.pprint(session.list_devices())
  with open('/content/adc.json', 'r') as f:
    auth_info = json.load(f)
  tf.contrib.cloud.configure_gcs(session, credentials=auth_info)

# TPU devices:
# [_DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:CPU:0, CPU, -1, 15478590230924548042),
#  _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 17179869184, 10470576999448356837),
#  _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:0, TPU, 17179869184, 10174756583359959919),
#  _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:1, TPU, 17179869184, 12581305773076099137),
#  _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:2, TPU, 17179869184, 5239439326212537618),
#  _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:3, TPU, 17179869184, 15506925854819515197),
#  _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:4, TPU, 17179869184, 5515832412837802144),
#  _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:5, TPU, 17179869184, 10014908513685979210),
#  _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:6, TPU, 17179869184, 2176223573355481611),
#  _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU:7, TPU, 17179869184, 16624288726840276573),
#  _DeviceAttributes(/job:tpu_worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 8589934592, 9419692871657185934)]

以下は run_pretraining.py をコピーして TPUClusterResolver の取得部分を書き換えています。

# copied from https://github.com/google-research/electra/blob/1ed085288ed265898cdb5deaf55bd56da252af80/run_pretraining.py
def tpu_train_or_eval(config: configure_pretraining.PretrainingConfig):
  """Run pre-training or evaluate the pre-trained model."""
  if config.do_train == config.do_eval:
    raise ValueError("Exactly one of `do_train` or `do_eval` must be True.")
  if config.debug:
    utils.rmkdir(config.model_dir)
  utils.heading("Config:")
  utils.log_config(config)

  is_per_host = tf.estimator.tpu.InputPipelineConfig.PER_HOST_V2

  tpu_cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(TPU_ADDRESS, zone=None, project=None)

  tpu_config = tf.estimator.tpu.TPUConfig(
      iterations_per_loop=config.iterations_per_loop,
      num_shards=(config.num_tpu_cores if config.do_train else
                  config.num_tpu_cores),
      tpu_job_name=config.tpu_job_name,
      per_host_input_for_training=is_per_host)
  run_config = tf.estimator.tpu.RunConfig(
      cluster=tpu_cluster_resolver,
      model_dir=config.model_dir,
      save_checkpoints_steps=config.save_checkpoints_steps,
      tpu_config=tpu_config)
  model_fn = model_fn_builder(config=config)
  estimator = tf.estimator.tpu.TPUEstimator(
      use_tpu=config.use_tpu,
      model_fn=model_fn,
      config=run_config,
      train_batch_size=config.train_batch_size,
      eval_batch_size=config.eval_batch_size)

  if config.do_train:
    utils.heading("Running training")
    estimator.train(input_fn=pretrain_data.get_input_fn(config, True),
                    max_steps=config.num_train_steps)
  if config.do_eval:
    utils.heading("Running evaluation")
    result = estimator.evaluate(
        input_fn=pretrain_data.get_input_fn(config, False),
        steps=config.num_eval_steps)
    for key in sorted(result.keys()):
      utils.log("  {:} = {:}".format(key, str(result[key])))
    return result

ここまで来たら、以下のようにして学習を開始します。チェックポイントは GCS に書き出されているので時間切れになったら、再実行で停止したところから再開できます。

tf.logging.set_verbosity(tf.logging.WARN)
tpu_train_or_eval(config)

# ================================================================================
# Config:
# ================================================================================
# debug False
# disallow_correct False
# disc_weight 50.0
# do_eval False
# ...
# ================================================================================
# Running training
# ================================================================================
# WARNING:tensorflow:From /tensorflow-1.15.2/python3.6/tensorflow_core/python/ops/resource_variable_ops.py:1630: calling BaseResourceVariable.__init__ (from tensorflow.python.ops.resource_variable_ops) with constraint is deprecated and will be removed in a future version.
# Instructions for updating:
# If using Keras pass *_constraint arguments to layers.
# ...

学習曲線はこんな感じです。

さて、今回は事前学習モデルの質を livedoor News Corpus の文書分類の試行5回の平均で評価することにしました。 とりあえず 766000 ステップ学習して見たところ、F1 スコアで 95.760 でした。なんというか"さすが ELECTRA 感(?)“がない結果です。

学習曲線を見ていると、回せばもう少し伸びそうな気がしたので、766000 ステップ終了時点を起点に総ステップ数を 1532000 に伸ばして学習を継続（学習レートが線形下降じゃなくなってしまいますが）し、ところどころチェックポイントを抜き取り検査して精度をみてみました。

ガタガタで精度のピークが出てるんだかないんだかよくわかりませんね。ですが、図に示していない部分で試した分もあわせて 1000000 ステップ辺りがピークっぽいということで、今回は、1001000 ステップ時点(F1スコア=97.03)のチェックポイントを使ってファインチューニングすることにしました。

本当は最初から綺麗に 1000000 ステップ学習しなおして出す予定だったのですが、livedoor News Corpus 文書分類の F1 スコアで 97 が出てたのを見て「もうこれでいいや」な気分になってしまって。。。

それではファインチューニングをしていきましょう。

4. 固有表現抽出データセットでのファインチューニング

それではファインチューニングして固有表現抽出をしてみましょう。

Colab で新しくノートブックを開いて、"ランタイム” -> “ランタイムタイプの変更” で “ハードウェアアクセラレータ” を “GPU” にしてください。

まずはデータセットを ELECTRA のコードに合わせて加工していきます。

UD_Japanese-GSD v2.6-NE の取得と加工

今回はデータセットに UD_Japanese-GSD v2.6-NE を使います。 UD_Japanese-GSD v2.6-NE は約8,000 文に Universal Dependencies(UD) の依存関係情報が付与された v2.6 に Megagon Labs さんが固有表現ラベルを追加して公開して下さったものになります。

SpaCy 2.3.0, GiNZA 4.0.0, Stanza 1.0.1 でベンチマークしたスコアがこちらの資料⁶に掲載されているので引用しておきます。この資料は SpaCy と GiNZA の関係性や学習に使用されたデータなど分かりやすくまとめられているのでお勧めです。

Ent 列が固有表現抽出のスコアです。今回はこの辺りが目標になりますね。

まずは BIO 形式のデータを取得します。

!wget https://github.com/megagonlabs/UD_Japanese-GSD/releases/download/v2.6-NE/train.bio
!wget https://github.com/megagonlabs/UD_Japanese-GSD/releases/download/v2.6-NE/dev.bio
!wget https://github.com/megagonlabs/UD_Japanese-GSD/releases/download/v2.6-NE/test.bio

BIO 形式とは、それぞれ “Begin”, “Inside”, “Outside” の略で複数トークンにまたがる固有表現が以下のようにラベル付けされているものです。

!head -39 test.bio | tail -8

# です    O
# 。   O
#
# 幸福    B-NORP
# の   I-NORP
# 科学    I-NORP
# 側   O
# から    O

これを一旦、 第2回 で紹介した Rasa の Markdown形式に変換します。 Rasa の形式にするのは「精度が似たようなものなら軽い方がいいよね」ということで比較しやすくするためです。やっつけ仕事のコードはこんな感じです。

def bio2rasa(file):
  with open(file, "r") as f:
    lines = f.readlines()
    lines = [line.strip() for line in lines]

  records = []
  record = "- "
  current_ent = ""
  for i, line in enumerate(lines):
    token_bioent = line.split("\t")
    if len(token_bioent) < 2:
      records.append(record)
      record="- "
      continue

    token, bioent = token_bioent
    bioent = bioent.split("-")

    if len(bioent) < 2: # "O" record.
      if current_ent != "":
        record += "](" + current_ent + ")" +  token
      else:
        record += token
      current_ent=""
      continue

    bio, ent = bioent
    if bio == "B": # "B" record.
      if current_ent != "":
        current_ent=ent
        record += "](" + current_ent + ")[" +  token
      else:
        current_ent=ent
        record += "[" + token
    else:  # "I" record
      record += token

  records.append(record)
  with open(file + ".md", "w") as f:
    f.write("\n".join(records))  

変換してみます。

bio2rasa("train.bio")
bio2rasa("dev.bio")
bio2rasa("test.bio")

出来上がったファイルはこんな感じです。

!head -6 train.bio.md

# - ホッケーにはデンジャラスプレーの反則があるので、膝より上にボールを浮かすことは基本的に反則になるが、その例外の一つがこのスクープである。
# - また行きたい、そんな気持ちにさせてくれるお店です。
# - 手に持った特殊な刃物を使ったアクロバティックな体術や、[揚羽](PERSON)と[薄羽](PERSON)同様にクナイや忍具を使って攻撃してくる。
# - [3年次](DATE)にはトータルオフェンスで[2,892ヤード](QUANTITY)を獲得し、これは大学記録となった。
# - 葬儀の最中ですよ!
# - [1998年度](DATE)に着手し、[道の駅遠山郷](FAC)北側から[かぐら大橋](FAC)南詰現道交点までの[1.060km](QUANTITY)のみ開通済み。

ちょうど1行=1文の状態になったので文数をカウントしておきましょう。確かに約8,000件ですね。

wc -l *.bio.md

#    500 dev.bio.md
#    543 test.bio.md
#   7032 train.bio.md
#   8075 total

ELECTRA のコードは BIO ではなく BIOES 形式(“Begin”, “Inside”, “Outside”, “Ending”, “Single”)が前提のコードなので、ここからさらに変換します。

事前学習で使った Transformers のトークナイザをロードして、

from transformers import BertJapaneseTokenizer
bert_japanese_tokenizer = BertJapaneseTokenizer.from_pretrained('cl-tohoku/bert-base-japanese-whole-word-masking')

Rasa の Markdown形式のファイルをメモリにロードする関数です。

def load_nlu_bench_md(file):
  with open(file, "r") as f:
    lines = f.readlines()
    lines = [line.strip()[2:] for line in lines if line[0]=="-"]
  return lines

Rasa NLU形式のサンプルをパースしてプレーン文字列とエンティティ集合に変換する関数。

def parse_example(example):
  mode = ""
  def head_guard(i):
    return i > 0
  def tail_guard(i):
    return i < len(example)-1

  pos=-1
  head=0
  tail=0
  entities=[]
  types=[]
  entity=""
  type=""
  mode="O"
  plain_text=""

  for i, c in enumerate(example):
    # count up normal characters and build plain text vertion.
    if (c not in [ "[", "]", "(", ")"] and mode != "C") or (c in ["(", ")"] and mode in ["B", "I", "E"]):
      pos+=1
      plain_text+=c

    # prev and next character  
    prev =  example[i-1] if head_guard(i) else ""
    next =  example[i+1] if tail_guard(i) else ""

    # handle entity mode
    if prev == "[":
      mode="B"
    if next == "]":
      mode="E"
    if mode=="B" and c != "]":
      if next == "]":
        mode="S"
      if prev != "[":
        mode="I"
    if c=="]":
      mode="O"
    # build entity 
    if mode in ["B", "I", "S", "E"]:
      if len(entity)==0:
        head=pos
      entity+=c
    if mode in ["S", "E"]:
      entities.append((entity, head, pos))
      entity=""

    # handle entity class mode and build entity type.
    if c == ")" and mode in ["C"]:
      mode="O"
      types.append(type)
      type=""
    if mode=="C":
      type+=c
    if c == "(" and mode not in ["B", "I"] :
      mode="C"
  return plain_text, [{"entity":e[0], "start": e[1], "end": e[2], "type":t} for e, t in zip(entities, types)]  

ちなみに、この関数ですが出力される entity の end は INCLUSIVE です。"私はエンティティです。" の場合は {"entity": "エンティティ", "start":2, "end":7 } になります。普通に entity[start:end] な感じで切り出せるほうが気持ちいいのですが、なぜこうしたのか思い出せません。何か理由があったような、なかったような。。。

つぎは、トークン列、プレーン文字列、エンティティ集合からトークン単位のBIOESラベルに変換する関数です。

def create_token_label(token_span, entities):
  start = token_span[1]
  end = token_span[2]
  #print("token_span : %s" % (token_span,))
  for entity in entities:
    #print("   entity : %s" % (entity,))
    if start == entity["start"] and end == entity["end"]:
      return "S-%s" % entity["type"]
    if start == entity["start"] and end < entity["end"]:
      return "B-%s" % entity["type"]      
    if start > entity["start"] and end < entity["end"]:
      return "I-%s" % entity["type"]  
    if start > entity["start"] and end == entity["end"]:
      return "E-%s" % entity["type"]
  return "O"

def create_token_level_labels(tokens, plain_text, entities):
  pos = 0
  i = 0
  token_spans = []
  for token in tokens:
    span=[]
    for c in token:
      #print(" %d %s" % (pos, c))
      span.append(pos)
      pos+=1
    token_spans.append((token, span[0], span[-1]))
  token_level_labels = []
  for token_span in token_spans:
    token_level_labels.append(create_token_label(token_span, entities))

  prev = "0"
  for token_level_label in token_level_labels:
    head = token_level_label[0]
    assert not(prev == "O" and head in ["I", "E"])
    assert not(prev in ["B", "I"] and head =="O")

  return token_level_labels

こんどはサブワードに分かれたトークンをまとめる関数です。

import re
def joint_sub_word_tokens(tokens):
  word_level_tokens = []
  for token in tokens:
    if token.startswith('##'):
      token = re.sub("^##", "", token)
      word_level_tokens[-1]+=token
    else:
      word_level_tokens.append(token)
  return word_level_tokens

ようやく変換処理の本体です。

def convert_to_token_level_labels(filename, lines):
  examples = []
  for line in lines:
    plain_text, entities = parse_example(line)
    tokens = joint_sub_word_tokens(bert_japanese_tokenizer.tokenize(plain_text))
    token_level_labels = create_token_level_labels(tokens, plain_text, entities)
    assert len(tokens) == len(token_level_labels)
    examples.append((tokens, token_level_labels))
  with open(filename, "w") as f:
    records = []
    for example in examples:
      for token, label in zip(example[0], example[1]):
        record = "%s %s" % (token, label)
        f.write(record+"\n")
      f.write("\n")

各ファイルを変換します。

convert_to_token_level_labels("train.txt", load_nlu_bench_md("train.bio.md"))
convert_to_token_level_labels("dev.txt", load_nlu_bench_md("dev.bio.md"))
convert_to_token_level_labels("test.txt", load_nlu_bench_md("test.bio.md"))

変換後の行数はこんな感じで、

!wc -l *.txt

#  12217 dev.txt
#   13177 test.txt
#  169092 train.txt
#  194486 total

変換後のファイルはこんな感じです。一応大丈夫そうな感じですね。 以下の例には BIOES の “I” がないですが 3 トークン以上の固有表現で “B” と “E” に挟まれた部分が “I” でラベリングされます。

!head -85 train.txt | tail -20

# な O
# 体 O
# 術 O
# や O
# 、 O
# 揚羽 S-PERSON
# と O
# 薄 B-PERSON
# 羽 E-PERSON
# 同様 O
# に O
# クナイ O
# や O
# 忍 O
# 具 O
# を O
# 使っ O
# て O
# 攻撃 O
# し O

ファインチューニングでは各トークン単位に BIOES 形式のラベルを分類問題で解くように学習する訳ですが、 BIOES 形式のラベルを推定した後の固有表現の切り出しはざくっと以下のような処理になります。この場合は “ベルゲン港” が “FAC"の固有表現です。

ソースコードで言うとこの辺りですね。

# https://github.com/google-research/electra/blob/79111328070e491b287c307906701ebc61091eb2/finetune/tagging/tagging_utils.py#L23-L40 から引用
23 def get_span_labels(sentence_tags, inv_label_mapping=None):
24   """Go from token-level labels to list of entities (start, end, class)."""
25   if inv_label_mapping:
26     sentence_tags = [inv_label_mapping[i] for i in sentence_tags]
27   span_labels = []
28   last = 'O'
29   start = -1
30   for i, tag in enumerate(sentence_tags):
31     pos, _ = (None, 'O') if tag == 'O' else tag.split('-')
32     if (pos == 'S' or pos == 'B' or tag == 'O') and last != 'O':
33       span_labels.append((start, i - 1, last.split('-')[-1]))
34     if pos == 'B' or pos == 'S' or last == 'O':
35       start = i
36     last = tag
37   if sentence_tags[-1] != 'O':
38     span_labels.append((start, len(sentence_tags) - 1,
39                         sentence_tags[-1].split('-')[-1]))
40   return span_labels

「“B-*” と “E-*” で固有表現のクラスが違っていたら？」とか、「“B-*” の後に”E-*”がなかったら？」とか考えてしまいますが、あまり細かいことは気にしない感じですかね。ただ、「だったら BIO でよくない？ 予測するラベルの種類少ない方が精度あがりそうじゃない？」という気分になってきました。

そんな訳で BIO 形式に変換してファインチューニング用のディレクトリに入れておきます⁷。

!mkdir -p data/finetuning_data/chunk
!cat train.txt | sed -e 's/ S-/ B-/' -e  's/ E-/ I-/' > data/finetuning_data/chunk/train.txt
!cat dev.txt | sed -e 's/ S-/ B-/' -e  's/ E-/ I-/' > data/finetuning_data/chunk/dev.txt
!cat test.txt | sed -e 's/ S-/ B-/' -e  's/ E-/ I-/' > data/finetuning_data/chunk/test.txt

ファインチューニングの実行

ここからファインチューニングを実行していきます。

まずトークナイザをインポートしてデータ格納ディレクトリを定義します。

from transformers import BertJapaneseTokenizer
bert_japanese_tokenizer = BertJapaneseTokenizer.from_pretrained('cl-tohoku/bert-base-japanese-whole-word-masking')
DATA_DIR = "./data"

必要なパッケージや関数をインポートします。

import sys
sys.path.append("./electra")
import configure_finetuning
from run_finetuning import run_finetuning
import json
import tensorflow as tf

ELECTRA のコードでは固有表現抽出に finetune.tagging.tagging_tasks.Chunking を利用します。ただトークナイザの処理が一部互換性がない（というかなくて当たり前なのですが）ので、finetune.tagging.tagging_tasks.tokenize_and_align() を差し替えておきます。

import unicodedata
from model import tokenization

def is_control(char):
  if char == "\t" or char == "\n" or char == "\r":
    return False
  cat = unicodedata.category(char)
  if cat.startswith("C"):
    return True
  return False

def is_whitespace(char):
  if char == " " or char == "\t" or char == "\n" or char == "\r":
    return True
  cat = unicodedata.category(char)
  if cat == "Zs":
    return True
  return False

def clean_text(text):
    output = []
    for char in text:
      cp = ord(char)
      if cp == 0 or cp == 0xfffd or is_control(char):
        continue
      if is_whitespace(char):
        output.append(" ")
      else:
        output.append(char)
    return "".join(output)

def tokenize_and_align(tokenizer, words, cased=False):
  """Splits up words into subword-level tokens."""
  words = ["[CLS]"] + list(words) + ["[SEP]"]
  tokenized_words = []
  for word in words:
    word = tokenization.convert_to_unicode(word)
    word = clean_text(word)
    if word == "[CLS]" or word == "[SEP]":
      word_toks = [word]
    else:
      if not cased:
        word = word.lower()
      word_toks = [word]
    tokenized_word = []
    for word_tok in word_toks:
      tokenized_word += tokenizer.tokenize(word_tok)
    tokenized_words.append(tokenized_word)
  assert len(tokenized_words) == len(words)
  return tokenized_words

from finetune.tagging import tagging_tasks
tagging_tasks.tokenize_and_align = tokenize_and_align

次に Chuking タスクに渡されるトークナイザを変更するために、get_tasks() と get_task() を入れ替えます。

from finetune.tagging.tagging_tasks import Chunking
from finetune import task_builder 

def get_tasks(config: configure_finetuning.FinetuningConfig):
  return [get_task(config, task_name, bert_japanese_tokenizer) for task_name in config.task_names]

def get_task(config: configure_finetuning.FinetuningConfig, task_name, tokenizer):
  if task_name == "chunk":
    return Chunking(config, tokenizer)
  else:
    raise ValueError("Unknown task " + task_name)

task_builder.get_tasks = get_tasks
task_builder.get_task = get_task

次に事前学習済みのモデルを所定のディレクトリに用意します。

!mkdir -p data/models
!mkdir -p data/models/electra_base_wiki_ja
!gsutil cp gs://somewhere/electra/max_seq_length_512/models/electra_base_wiki_ja/* ./data/models/electra_base_wiki_ja
!gsutil cp gs://somewhere/electra/max_seq_length_512/vocab.txt ./data/models/electra_base_wiki_ja
!ls data/models/electra_base_wiki_ja

# checkpoint   model.ckpt-1001000.data-00000-of-00001  model.ckpt-1001000.meta
# graph.pbtxt  model.ckpt-1001000.index          vocab.txt

ELECTRA のコードは {DATA_DIR}/fintuning_data/{task_name} ディレクトリに訓練/検証/テストセットがそれぞれ train.txt, dev.txt, test.txt として配置されていることを前提とします。 eval の対象とするセットをパラメータで指定できると良いのですが、 ELECTRA のコードを見ると以下のようになっていてテストセットでの精度出力に対応しているのは GLUE と SQuARD 2.0 だけのようです。

# https://github.com/google-research/electra/blob/79111328070e491b287c307906701ebc61091eb2/run_finetuning.py#L273-L297 から引用
273    if config.do_eval:
274      heading("Run dev set evaluation")
275      results.append(model_runner.evaluate())
276      write_results(config, results)
277      if config.write_test_outputs and trial <= config.n_writes_test:
278        heading("Running on the test set and writing the predictions")
279        for task in tasks:
280          # Currently only writing preds for GLUE and SQuAD 2.0 is supported
281          if task.name in ["cola", "mrpc", "mnli", "sst", "rte", "qnli", "qqp",
282                           "sts"]:
...          省略                 
285          elif task.name == "squad":
...          省略
295          else:
296            utils.log("Skipping task", task.name,
297                      "- writing predictions is not supported for this task")

これを修正するのも面倒ですし、Early Stopping するわけでもないので、eval でテストセットでの精度が確認できるようにファイルを差し替えておきます。

!mv data/finetuning_data/chunk/dev.txt data/finetuning_data/chunk/dev.txt.org
!mv data/finetuning_data/chunk/test.txt data/finetuning_data/chunk/dev.txt

それでは、ようやくファインチューニングの実行です。

num_trials を指定することで指定回数の試験を行ってくれます。1回の試験に3時間ほどかかったので、"num_trials": 5 とかすると終了前に Colab のインスタンスが消えてしまいます。./data の内容はそのまま GCS において gs://somewhere/data とか指定しても動きますので複数回実行するときは GCS に置いておいたほうが無難かもしれませんね。

MODEL_NAME = 'electra_base_wiki_ja'
HPARAMS =  '''
      {
        "use_tpu": false,   
        "model_size": "base", 
        "max_seq_length": 512,
        "embedding_size": 768,
        "vocab_file": "./data/models/electra_base_wiki_ja/vocab.txt",
        "vocab_size": %d,
        "save_checkpoints_steps": 10000,
        "num_train_epochs": 20,
        "num_trials": 5,
        "task_names": ["chunk"],
        "train_batch_size": 8,
        "eval_batch_size": 4
      }
    ''' % bert_japanese_tokenizer.vocab_size
HPARAMS = " ".join(HPARAMS.splitlines())

hparams = json.loads(HPARAMS)
config = configure_finetuning.FinetuningConfig(MODEL_NAME, DATA_DIR, **hparams)

tf.logging.set_verbosity(tf.logging.ERROR)
run_finetuning(config)

# ================================================================================
# Config: model=electra_base_wiki_ja, trial 1/5
# ================================================================================
# answerable_classifier True
# answerable_uses_start_logits True
# answerable_weight 0.5
# beam_size 20
...

今回は 5 回実験して以下のような結果です。⁸。

!cat data/models/electra_base_wiki_ja/results/chunk_results.txt

# chunk: precision: 85.02 - recall: 84.32 - f1: 84.67 - loss: 0.58
# chunk: precision: 82.87 - recall: 83.83 - f1: 83.35 - loss: 0.60
# chunk: precision: 83.97 - recall: 83.83 - f1: 83.90 - loss: 0.56
# chunk: precision: 84.44 - recall: 84.16 - f1: 84.30 - loss: 0.57
# chunk: precision: 84.81 - recall: 83.83 - f1: 84.32 - loss: 0.57

平均すると、F1 = 84.1 になりました。

import pickle
import numpy as np
with open("data/models/electra_base_wiki_ja/results/chunk_results.pkl", "rb") as f:
  results = pickle.load(f)
print("Mean F1 score = %5.3f" % np.array([result["chunk"]["f1"] for result in results]).mean())

# Mean F1 score = 84.108

他のモデルとの比較は以下のとおりです。

• BERT は Hugging Face Transformers の "cl-tohoku/bert-base-japanese-whole-word-masking" を用い、こちら⁹で公開されているコードで試験した 5 回平均の結果です。最大シーケンス長、バッチサイズとエポック数は今回の ELECTRA の実験に合わせました。
• ”*“ が付いたモノは ⁶ からの転記です。
• Rasa NLU は本連載の 第2回 で紹介した修正を適用したものです。MeCab で形態素解析して CRF に掛けました。

上図の ELECTRA と BERT は DLフレームワークが Tensorflow と PyTorch なので純粋な事前学習モデルの比較にはなっていません。ですが、一応は ELECTRA がギリ上回って面目を保った（？）感じでしょうか。 ELECTRA に関しては事前学習時のバッチサイズを 256 に戻したり、Generator と Discriminator のサイズ比(0.3333)を調整すれば、もう少し伸びるんではないかと思います。

Rasa には厳しい結果になりました。今回試したのはバージョン 1.10.9 です。最新は 2.0.3 なので少し古いですが、最新版でも独自の固有表現抽出モデルの構築には CRF を使うようなので、精度的にはさほど変わらないのではないかと思います。ただ、独自の固有表現モデルの必要がない場合は、SpaCy を 2.3.0 に入れ替え、内部的に CRF ではなく SpaCy の事前学習済み固有表現抽出モデルを用いるように設定してサーバ部分だけ使う手はあるかもしれません（試してはないですが）。

せっかく固有表現抽出モデルができたので SavedModel にエクスポートして推論してみましょう。

5. SavedModel へのエクスポートと推論

エクスポートの処理はこんな感じになります。 ELECTRA のコードを流用しながら作成しており、以下を前提としたコードになっています。

• オリジナルの finetune.Task クラスにない get_export_module() でタスクの推論結果を取得できる。
• finetune に SavedModel の入り口になる関数を生成する build_serving_input_fn() が存在する。

前者はタスクによって推論結果の形はいろいろなので、その部分を finetune.Task の具象クラスに任せたい、後者はシーケンスの最大長を固定で書き込みたくないということで、こんな形になりました¹⁰。

%%bash
cat << EOF > run_export.py
# coding=utf-8
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import argparse
import collections
import json

import tensorflow.compat.v1 as tf

import configure_finetuning
from finetune import preprocessing
from finetune import task_builder
from model import modeling
from model import optimization
from util import training_utils
from util import utils

import importlib
import os
import finetune


class ExportModel(object):
  def __init__(self, config: configure_finetuning.FinetuningConfig, tasks,
               is_training, features, num_train_steps):
    bert_config = training_utils.get_bert_config(config)
    self.bert_config = bert_config
    if config.debug:
      bert_config.num_hidden_layers = 3
      bert_config.hidden_size = 144
      bert_config.intermediate_size = 144 * 4
      bert_config.num_attention_heads = 4
    assert config.max_seq_length <= bert_config.max_position_embeddings
    bert_model = modeling.BertModel(
        bert_config=bert_config,
        is_training=is_training,
        input_ids=features["input_ids"],
        input_mask=features["input_mask"],
        token_type_ids=features["segment_ids"],
        use_one_hot_embeddings=config.use_tpu,
        embedding_size=config.embedding_size)
    percent_done = (tf.cast(tf.train.get_or_create_global_step(), tf.float32) /
                    tf.cast(num_train_steps, tf.float32))

    self.outputs = {}
    for task in tasks:
      with tf.variable_scope("task_specific/" + task.name):
        task_outputs = task.get_export_module(
            bert_model, features, is_training, percent_done)
        self.outputs[task.name] = task_outputs

def model_fn_builder(config: configure_finetuning.FinetuningConfig, tasks,
                     num_train_steps, pretraining_config=None):

  def model_fn(features, labels, mode, params):
    utils.log("Building model...")
    is_training = (mode == tf.estimator.ModeKeys.TRAIN)
    model = ExportModel(
        config, tasks, is_training, features, num_train_steps)

    init_checkpoint = config.init_checkpoint
    if pretraining_config is not None:
      init_checkpoint = tf.train.latest_checkpoint(pretraining_config.model_dir)
      utils.log("Using checkpoint", init_checkpoint)
    tvars = tf.trainable_variables()
    scaffold_fn = None
    if init_checkpoint:
      assignment_map, _ = modeling.get_assignment_map_from_checkpoint(
          tvars, init_checkpoint)
      if config.use_tpu:
        def tpu_scaffold():
          tf.train.init_from_checkpoint(init_checkpoint, assignment_map)
          return tf.train.Scaffold()
        scaffold_fn = tpu_scaffold
      else:
        tf.train.init_from_checkpoint(init_checkpoint, assignment_map)

    assert mode == tf.estimator.ModeKeys.PREDICT
    output_spec = tf.estimator.tpu.TPUEstimatorSpec(
        mode=mode,
        predictions=utils.flatten_dict(model.outputs),
        scaffold_fn=scaffold_fn)

    utils.log("Building complete")
    return output_spec

  return model_fn


class ModelRunner(object):

  def __init__(self, config: configure_finetuning.FinetuningConfig, tasks,
               pretraining_config=None):
    self._config = config
    self._tasks = tasks
    self._preprocessor = preprocessing.Preprocessor(config, self._tasks)

    is_per_host = tf.estimator.tpu.InputPipelineConfig.PER_HOST_V2
    tpu_cluster_resolver = None
    if config.use_tpu and config.tpu_name:
      tpu_cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(
          config.tpu_name, zone=config.tpu_zone, project=config.gcp_project)
    tpu_config = tf.estimator.tpu.TPUConfig(
        iterations_per_loop=config.iterations_per_loop,
        num_shards=config.num_tpu_cores,
        per_host_input_for_training=is_per_host,
        tpu_job_name=config.tpu_job_name)
    run_config = tf.estimator.tpu.RunConfig(
        cluster=tpu_cluster_resolver,
        model_dir=config.model_dir,
        save_checkpoints_steps=config.save_checkpoints_steps,
        save_checkpoints_secs=None,
        tpu_config=tpu_config)

    self._train_input_fn, self.train_steps = None, 0
    model_fn = model_fn_builder(
        config=config,
        tasks=self._tasks,
        num_train_steps=self.train_steps,
        pretraining_config=pretraining_config)
    self._estimator = tf.estimator.tpu.TPUEstimator(
        use_tpu=config.use_tpu,
        model_fn=model_fn,
        config=run_config,
        train_batch_size=config.train_batch_size,
        eval_batch_size=config.eval_batch_size,
        predict_batch_size=config.predict_batch_size)

  def export(self, export_dir):
    return {task.name: self.export_task(task, export_dir) for task in self._tasks}

  def export_task(self, task, export_dir):
    utils.log("Exporting", task.name)
    self._estimator._export_to_tpu = False
    self._estimator.export_savedmodel(os.path.join(export_dir, task.name), 
            finetune.build_serving_input_fn(task))

def run_export(config: configure_finetuning.FinetuningConfig, export_dir: str):

  results = []
  trial = 1
  heading_info = "model={:}, trial {:}/{:}".format(
      config.model_name, trial, config.num_trials)
  heading = lambda msg: utils.heading(msg + ": " + heading_info)
  heading("Config")

  utils.log_config(config)
  generic_model_dir = config.model_dir
  tasks = task_builder.get_tasks(config)

  config.model_dir = generic_model_dir + "_1"
  model_runner = ModelRunner(config, tasks)
  model_runner.export(export_dir)

def main():
  parser = argparse.ArgumentParser(description=__doc__)
  parser.add_argument("--data-dir", required=True,
                      help="Location of data files (model weights, etc).")
  parser.add_argument("--model-name", required=True,
                      help="The name of the model being fine-tuned.")
  parser.add_argument("--hparams", default="{}",
                      help="JSON dict of model hyperparameters.")
  parser.add_argument('--module_import', action='append', help='module to import')
  parser.add_argument("--export-dir", required=True,
                      help="Location of saved_model.")
  args = parser.parse_args()

  for module in args.module_import:
    print("importing module %s" % module)
    importlib.import_module(module)

  if args.hparams.endswith(".json"):
    hparams = utils.load_json(args.hparams)
  else:
    hparams = json.loads(args.hparams)
  tf.logging.set_verbosity(tf.logging.ERROR)
  run_export(configure_finetuning.FinetuningConfig(
      args.model_name, args.data_dir, **hparams), args.export_dir)


if __name__ == "__main__":
  main()
EOF

勝手な前提を立てたので、その埋め合わせが以下のコードになります。 ファインチューニング時にも行ったトークナイザの差し替えをしつつ、 get_export_module() の実装を Chunking クラスに、 SavedModel の入り口になる関数を生成する build_serving_input_fn() を finetuning にねじ込みます。

%%bash
cat << EOF > fix_tagging.py
# Swap tagging_tasks.tokenize_and_align
import unicodedata
from model import tokenization
def is_control(char):
  if char == "\t" or char == "\n" or char == "\r":
    return False
  cat = unicodedata.category(char)
  if cat.startswith("C"):
    return True
  return False

def is_whitespace(char):
  if char == " " or char == "\t" or char == "\n" or char == "\r":
    return True
  cat = unicodedata.category(char)
  if cat == "Zs":
    return True
  return False

def clean_text(text):
    output = []
    for char in text:
      cp = ord(char)
      if cp == 0 or cp == 0xfffd or is_control(char):
        continue
      if is_whitespace(char):
        output.append(" ")
      else:
        output.append(char)
    return "".join(output)

def tokenize_and_align(tokenizer, words, cased=False):
  """Splits up words into subword-level tokens."""
  words = ["[CLS]"] + list(words) + ["[SEP]"]
  #basic_tokenizer = tokenizer.basic_tokenizer
  tokenized_words = []
  for word in words:
    word = tokenization.convert_to_unicode(word)
    word = clean_text(word)
    if word == "[CLS]" or word == "[SEP]":
      word_toks = [word]
    else:
      if not cased:
        word = word.lower()
      word_toks = [word]
    tokenized_word = []
    for word_tok in word_toks:
      tokenized_word += tokenizer.tokenize(word_tok)
    tokenized_words.append(tokenized_word)
  assert len(tokenized_words) == len(words)
  return tokenized_words

from finetune.tagging import tagging_tasks
tagging_tasks.tokenize_and_align = tokenize_and_align

# Swap get_task and get_tasks
import configure_finetuning
from finetune import task_builder 
from finetune.tagging.tagging_tasks import Chunking
from transformers import BertJapaneseTokenizer
bert_japanese_tokenizer = BertJapaneseTokenizer.from_pretrained('cl-tohoku/bert-base-japanese-whole-word-masking')

def get_tasks(config: configure_finetuning.FinetuningConfig):
  return [get_task(config, task_name, bert_japanese_tokenizer) for task_name in config.task_names]

def get_task(config: configure_finetuning.FinetuningConfig, task_name, tokenizer):
  if task_name == "chunk":
    return Chunking(config, tokenizer)
  else:
    raise ValueError("Unknown task " + task_name)

task_builder.get_tasks = get_tasks
task_builder.get_task = get_task

# Define input func
import tensorflow.compat.v1 as tf

def build_serving_input_fn(task):
  def serving_input_fn():
    input_ids = tf.placeholder(
      tf.int32, [None, task.config.max_seq_length], name='input_ids')
    input_mask = tf.placeholder(
      tf.int32, [None, task.config.max_seq_length], name='input_mask')
    segment_ids = tf.placeholder(
      tf.int32, [None, task.config.max_seq_length], name='segment_ids')
    labeled_positions = tf.placeholder(
      tf.int32, [None, task.config.max_seq_length], name='labeled_positions')

    input_fn = tf.estimator.export.build_raw_serving_input_receiver_fn({
      'input_ids': input_ids,
      'input_mask': input_mask,
      'segment_ids': segment_ids,
      'labeled_positions': labeled_positions,
    })()
    return input_fn
  return serving_input_fn

import finetune
finetune.build_serving_input_fn = build_serving_input_fn

# Swap prediction module.
from pretrain import pretrain_helpers

def get_export_module(self, bert_model, features, is_training, percent_done):
  n_classes = len(self._get_label_mapping())
  reprs = bert_model.get_sequence_output()
  reprs = pretrain_helpers.gather_positions(
      reprs, features["labeled_positions"])
  logits = tf.layers.dense(reprs, n_classes)
  probs = tf.nn.softmax(logits)
  return dict(
      probs=probs,
      predictions=tf.argmax(probs, axis=-1)
  )
Chunking.get_export_module = get_export_module
EOF

ここまで来たらエクスポートを実行します。ファインチューニング時に GPU のメモリを掴んだままになっているかもしれないので、ランタイムを再起動しておきましょう。

以下のようにして ./run_export.py を実行します。さっきの帳尻合わせのコードは --module_import で流し込んでいます。

!export PYTHONPATH=${PYTHONPATH}:./electra && \
  \
  python3 ./run_export.py \
    --module_import fix_tagging \
    --export-dir ./export \
    --data-dir ./data \
    --model-name electra_base_wiki_ja    \
    --hparams '{"model_size": "base",    \
                "task_names": ["chunk"], \
                "use_tpu": false,        \
                "max_seq_length": 512,   \
                "embedding_size": 768,   \
                "vocab_size": 32000,     \
                "eval_batch_size": 8}'

無事にエクスポートが出来ました。

!ls export/chunk

# 1604018488

エクスポートされたモデルのインタフェースを確認してみます。

!saved_model_cli show --dir ./export/chunk/1604018488 --tag_set serve --signature_def serving_default

# The given SavedModel SignatureDef contains the following input(s):
#   inputs['input_ids'] tensor_info:
#       dtype: DT_INT32
#       shape: (-1, 512)
#       name: input_ids_1:0
#   inputs['input_mask'] tensor_info:
#       dtype: DT_INT32
#       shape: (-1, 512)
#       name: input_mask_1:0
#   inputs['labeled_positions'] tensor_info:
#       dtype: DT_INT32
#       shape: (-1, 512)
#       name: labeled_positions_1:0
#   inputs['segment_ids'] tensor_info:
#       dtype: DT_INT32
#       shape: (-1, 512)
#       name: segment_ids_1:0
# The given SavedModel SignatureDef contains the following output(s):
#   outputs['chunk_predictions'] tensor_info:
#       dtype: DT_INT64
#       shape: (-1, -1)
#       name: task_specific/chunk/ArgMax:0
#   outputs['chunk_probs'] tensor_info:
#       dtype: DT_FLOAT
#       shape: (-1, -1, 93)
#       name: task_specific/chunk/Softmax:0
# Method name is: tensorflow/serving/predict

入力側

• inputs['input_ids']
• inputs['input_mask']
• inputs['segment_ids']

以上の3つは通常の BERT の入力そのものです（忘れた人は 第3回 を読んでみて下さい）。

• inputs['labeled_positions']

labeled_positions は「本日は晴天なり」であれば、 [1, 3, 4, 6] になります。ちょっとわかりにくいですが、トークン化すると["[CLS]", "本", "##日", "は", "晴", "##天", "なり", "[SEP]"] となり、特殊トークン("[*]")やサブワード("##")を除いたトークンのインデックス系列になります。

ようするに「本日」が固有表現かどうかは「本」トークンに対応する出力だけ見て、「##日」は見ていない訳です。 そのあたりのロジックが以下の 205 行目のところですね。

https://github.com/google-research/electra/blob/79111328070e491b287c307906701ebc61091eb2/finetune/tagging/tagging_tasks.py#L201-L207 から引用
201  def get_prediction_module(
202      self, bert_model, features, is_training, percent_done):
203    n_classes = len(self._get_label_mapping())
204    reprs = bert_model.get_sequence_output()
205    reprs = pretrain_helpers.gather_positions(
206        reprs, features[self.name + "_labeled_positions"])
207    logits = tf.layers.dense(reprs, n_classes)

出力側

出力側の第1軸はバッチサイズ、第2軸は、サブワード化しないトークン（「本日は晴天なり」であれば["本日”, “は”, “晴天”, “なり”]）のシーケンス長になります。chunk_probs の第3軸はラベルの総数（BIOES の “B-”, “I-”… と PERSON, ORG, LOC … の掛け合わせ）です。

• outputs['chunk_predictions']
• outputs['chunk_probs']

chunk_predictions は整数型でラベルのインデックスが返ってきます。このインデックスと元のラベル文字列との対応関係は以下の pickle に格納されているので、忘れずに回収しておきましょう。

!find . -name *label_mapping.pkl 

# ./data/models/electra_base_wiki_ja/finetuning_tfrecords/chunk_tfrecords/chunk_label_mapping.pkl

中身はこんな感じですね。

import pickle
with open("./data/models/electra_base_wiki_ja/finetuning_tfrecords/chunk_tfrecords/chunk_label_mapping.pkl", "rb") as f:
  chunk_label_mapping = pickle.load(f)
for label, index in chunk_label_mapping.items():
  print("%s = %d" % (label, index))

# B-CARDINAL = 0
# B-DATE = 1
# B-EVENT = 2
# B-FAC = 3
...

プレディクションの実行

プレディクションもしてみましょう。過去の連載では Tensorflow Serving にデプロイして REST と gRPC での呼び出し方を紹介したので、今回は tf.Session.run() で呼び出してみます¹¹。

それではもう一度、ランタイムを再起動して Tensorflow 1.x を選択します。

%tensorflow_version 1.x

各種のインポートです。

import sys
sys.path.append("./electra")
import os
import numpy as np
import pickle
import tensorflow as tf
import re
from finetune.tagging.tagging_utils import get_span_labels

先ほどのラベル文字列とインデックスのマッピングをロードする関数です。

def load_label_mapping(label_mapping):
  with open(label_mapping, "rb") as f:
    label_mapping = pickle.load(f)
  inv_label_mapping = {}
  for k, v in label_mapping.items():
    inv_label_mapping[v] = k
  return label_mapping, inv_label_mapping

これはサブワードに分割されたシーケンスを単語単位に結合する関数です。

def joint_sub_word_tokens(tokens):
  word_level_tokens = []
  for token in tokens:
    if token.startswith('##'):
      token = re.sub("^##", "", token)
      word_level_tokens[-1]+=token
    else:
      word_level_tokens.append(token)
  return word_level_tokens  

# こんな感じです。
# ['ノルウェー',
#  'の',
#  '劇',
#  '作家',
#  'ヘンリック・イプセン',
#  'の'
# ...

これは単語単位の系列とサブワード単位の系列のマッピングを作る関数です。

def build_words_to_tokens(sub_tokens):
  words_to_tokens = [] 
  for sub_token in sub_tokens: 
    if sub_token.startswith('##'):
      words_to_tokens[-1].append(sub_token)
    else:
      words_to_tokens.append([sub_token])
  return words_to_tokens

# こんな感じです。
# [['ノルウェー'],
#  ['の'],
#  ['劇'],
#  ['作家'],
#  ['ヘン', '##リック', '##・', '##イ', '##プ', '##セン'],
#  ['の'],
# ...

次に labeled_positions の生成関数です。

def build_labeled_positions(words_to_tokens):
  tagged_positions = []
  pos_sub_words = 0
  for word_tokens in words_to_tokens:
    if "[CLS]" not in word_tokens and "[SEP]" not in word_tokens and "[PAD]" not in word_tokens:
      tagged_positions.append(pos_sub_words)
    pos_sub_words += len(word_tokens)
  return tagged_positions

# こんな感じです。
# [0, 1, 2, 3, 4, 10, 11, ...]

今までの関数をまとめて BERT への入力を作る関数です。

def build_features(text, max_length):
  input_ids = tokenizer.encode(text, add_special_tokens=True, 
                pad_to_max_length=True, max_length=max_length) 
  input_mask = [1]* len(tokenizer.encode(text, add_special_tokens=True, 
                pad_to_max_length=False, max_length=max_length))
  input_mask = input_mask + [0]*(max_length-len(input_mask))
  segment_ids = [0] * max_length
  sub_tokens = tokenizer.tokenize(text)
  tokens = joint_sub_word_tokens(sub_tokens)
  sub_tokens = tokenizer.convert_ids_to_tokens(input_ids)
  words_to_tokens = build_words_to_tokens(sub_tokens)
  labeled_positions =  build_labeled_positions(words_to_tokens)
  labeled_positions = labeled_positions + [0] * (max_length - len(labeled_positions))

  return input_ids, input_mask, segment_ids, labeled_positions, sub_tokens, tokens

これが SavedModel を呼び出す処理です。

def predict(input_ids, input_mask, segment_ids, labeled_positions):
  with tf.Session(graph=tf.Graph()) as sess:
    meta_graph = tf.saved_model.loader.load(sess, [tf.saved_model.tag_constants.SERVING], export_dir)
    model_signature = meta_graph.signature_def['serving_default']
    input_signature = model_signature.inputs
    output_signature = model_signature.outputs

    # input
    key_input_ids = sess.graph.get_tensor_by_name("input_ids_1:0")
    key_input_mask = sess.graph.get_tensor_by_name("input_mask_1:0")
    key_segment_ids = sess.graph.get_tensor_by_name("segment_ids_1:0")
    key_labeled_positions = sess.graph.get_tensor_by_name("labeled_positions_1:0")

    # output
    key_probs = sess.graph.get_tensor_by_name("task_specific/chunk/Softmax:0")
    key_preds = sess.graph.get_tensor_by_name("task_specific/chunk/ArgMax:0")

    preds, probs = sess.run([key_preds, key_probs], feed_dict={
                            key_input_ids: input_ids,
                            key_input_mask: input_mask,
                            key_segment_ids: segment_ids,
                            key_labeled_positions: labeled_positions
                          })
    return preds, probs

ちなみに finetune.tagging.tagging_utils からインポートした span_labels は予測したラベルインデックス系列を以下のように加工します。

span_labels = get_span_labels(preds[0], inv_label_mapping)
# [(0, 0, 'GPE'), (4, 4, 'PERSON'), (8, 8, 'WORK_OF_ART')]

こちらは上記の span_labels とトークン系列からエンティティを切り出す関数です。

def get_entities(span_labels, tokens, probs):
  entities = []
  for span_label in span_labels:
    start = span_label[0]
    end = span_label[1]
    entity_type = span_label[2]
    span_tokens = tokens[start:end+1]
    entity_text = "".join(span_tokens)
    start_pos = 0
    end_pos = 0
    for i, token in enumerate(tokens):
      if i < start:
        start_pos += len(token)
      if i < end +1 :
        end_pos += len(token)
    entity = {
      "text": entity_text,
      "type": entity_type,
      "start": start_pos,
      "end": end_pos
    }
    entities.append(entity)
  return entities

全部をひとまとめにして実行してみます。

text = "ノルウェーの劇作家ヘンリック・イプセンの劇詩『ペール・ギュント』の登場人物に因んで命名された。"
max_length = 512
export_dir = "./export/chunk/1604018488"
label_mapping, inv_label_mapping = load_label_mapping("./data/models/electra_base_wiki_ja/finetuning_tfrecords/chunk_tfrecords/chunk_label_mapping.pkl")

input_ids, input_mask, segment_ids, labeled_positions, sub_tokens, tokens = build_features(text, max_length)
predicts, probs = predict([input_ids], [input_mask], [segment_ids], [labeled_positions])
span_labels = get_span_labels(predicts[0], inv_label_mapping)
get_entities(span_labels, tokens)

# [{'end': 5, 'start': 0, 'text': 'ノルウェー', 'type': 'GPE'},
#  {'end': 19, 'start': 9, 'text': 'ヘンリック・イプセン', 'type': 'PERSON'},
#  {'end': 31, 'start': 23, 'text': 'ペール・ギュント', 'type': 'WORK_OF_ART'}]

いい感じに固有表現を抽出できています。ここまでくれば（呼び出す前と後が面倒ですが）、 Tensorflow Serving にのせて gRPC や REST で利用するのも難しくないですね。

最後に言い訳です。この処理、BertJapaneseTokenizer でばらしたトークンを "".join() で結合した文字列に対する固有表現の文字列と開始・終了位置の計算になっています。つまり、"三代目 J SOUL BROTHERS" から空白が落ちて "三代目JSOULBROTHERS" になってしまいます。対応するなら元の文字列とのマッピングを計算して、開始・終了位置を再計算、エンティティ文字列を再切り出しするのが一番ラクですかねー。

6. おわりに

今回は ELECTRA の事前学習から固有表現抽出の推論までやってみました。 BERT を単体で利用するパターンはもう何回か紹介しているので、次回は今回作成した ELECTRA の事前学習モデルを使って複数の BERT を組み合わたモデルを試してみたいと思います。まずは ORQA でクイズに答えるモデルを作ってみたいです。公開されているコードはゴリゴリの Tensorflow なのでがんばったら組み込めるんじゃないかと。