BERTを使ったテキスト分類モデルを作る

Googleが開発した自然言語処理であるBERTは、2019年10月25日検索エンジンへの導入を発表して以来、世間一般にも広く知られるようになりました。

GoogleはBERTの論文公開と共に、日本語を含む複数の言語でプレトレーニングのモデルを公開しましたが、日本語のモデルにおいてはサブワード分割処理の関係でそのまま使うのが難しいという問題がありました。

その後有識者の方々により、Wikipedia日本語版の全文に対してjuman++またはMecabで行った分かち書き結果を基にしたサブワードによるプレトレーニングモデルが公開され、日本語におけるBERT利用がますます加速していきました。

ただ、Googleのリリースしたソースコードで独自ネットワークを組むのは初学者には敷居が高く、ディープラーニング経験がない開発者は手が出しにくい状況でした。

今回hugging headsが提供しているtransformersというライブラリから、Mecab版のプレトレーニングモデルを呼び出せるようになり、正式リリースされたtensorflow 2.0のEagerモードの使いやすさも相まって、誰もが容易にBERTを使ったモデルを構築できるようになりました。

どれくらい容易に組めるかを比較するため、本稿では最初にLSTMでモデルを作り、その後BERTを使った自然言語の分類モデル構築を行います。

ディープラーニングにおける自然言語処理

ディープラーニングや機械学習において自然言語処理のタスクを実行する場合、最初にやることはトークナイズとエンコードです。

トークナイズは自然言語をトークンという単位に分割（トークナイズ）することで、その結果のトークンを機械で扱えるよう個別のIDを割り当てることをエンコードといいます。

トークナイズとエンコードは以下のような流れで処理を進めます。
①コーパス内のすべてのテキストを形態素解析器にかけ、トークン単位に分かち書きする。
②機械が処理できるように、トークンごとに固有のIDを付与する。
③最後に、テキストに存在するトークンをすべてIDに置き換える。

処理によっては、この過程の中で頻出する意味のない単語（英語で「a」や「the」のような）を除外（ストップワード）したり、語幹の統一化（ステミング）などの前処理を行います。

LivedoorコーパスをLSTMで分類する

今回の例ではlivedoorコーパスを対象に、コーパス内の文章があるカテゴリに含まれるか否かを分類するモデルをLSTMで構築します。

LSTMは時系列データを扱うRNN(Recurrent Neural Network）の拡張モデルで文章のような時系列データを扱うのに長けており、tensorflow 2.0の場合はkerasのLSTMを使うことで簡単に実装できます。

コーパスの入手

livedoorコーパスはNHN Japan株式会社が運営する「livedoor ニュース」のうち、下記のクリエイティブ・コモンズライセンスが適用されるニュース記事を収集したもので、こちらのリンクからダウンロードができます。
https://www.rondhuit.com/download.html

このページのldcc-20140209.tar.gzをダウンロードして解凍します。

$wget https://www.rondhuit.com/download/ldcc-20140209.tar.gz
$ tar zfx ldcc-20140209.tar.gz

解凍したファイルは以下のようなディレクトリ構造になっています。

textの中身

$tree text -L 1
text
├── CHANGES.txt
├── dokujo-tsushin
├── it-life-hack
├── kaden-channel
├── livedoor-homme
├── movie-enter
├── peachy
├── README.txt
├── smax
├── sports-watch
└── topic-news

text/dokujo-tsushinの中身

# tree text/dokujo-tsushin
text/dokujo-tsushin
├── dokujo-tsushin-4778030.txt
├── dokujo-tsushin-4778031.txt
├── dokujo-tsushin-4782522.txt
├── dokujo-tsushin-4788357.txt
...
├── dokujo-tsushin-6915005.txt
└── LICENSE.txt

データサイズが大きすぎるので、簡略化のためdokujo-tsushinとsports-watch以外のディレクトリをすべて削除しておいてください。

textの中身

$tree text -L 1
text
├── dokujo-tsushin
└── sports-watch

トークナイズを行う

このコーパスに対してTensorflow 2.0の機能を使ってトークナイズ処理を施します。

Mecabのインストール

Mecab本体は以下のリンクから入手できます。
http://taku910.github.io/mecab/

インストール直後のMecabは分かち書きの辞書を持っていないため、辞書を別途入手する必要があります。
多くの場合、IPA辞書またはmecab-ipadic-NEologdを選択します。

IPA辞書
https://drive.google.com/uc?export=download&id=0B4y35FiV1wh7MWVlSDBCSXZMTXM

mecab-ipadic-NEologd
https://github.com/neologd/mecab-ipadic-neologd

pythonからMecabを呼び出す

pythonからMecabを使う場合、mecab-python3またはfugashiのいずれかをインストールする必要があります。
今回はmecab-python3を利用します。

$ pip install mecab-python3

pythonからMecabが呼び出せることを確認します。

import Mecab

tagger = Mecab.Tagger('-Owakati')
print(tagger.parse('すもももももももものうち'))

パーサーに与えられた文字列が分かち書きされていることが確認できます。

すもも も もも も もも の うち

データセット作成

モデル構築に先立ち、LSTMモデルで訓練するためのデータセットを作成します。

データセットを作成するためには対象ディレクトリ配下のすべてのファイルを取得する必要がありますが、Tensorflowではテキストファイルのコーパスを操作するAPIが数多く揃っており、tf.compat.v1.gfile.ListDirectoryを使うと指定したディレクトリ配下のリストを取得できます。

livedoorコーパス内のファイルとディレクトリ一覧を取得

text_dir = os.path.join(os.getcwd(),"text")
tf.compat.v1.gfile.ListDirectory(text_dir)

['dokujo-tsushin',
 'sports-watch']

tf.compat.v1.gfile.ListDirectoryをネストしてファイル一覧を取得し、tf.data.TextLineDataset APIを使って各カテゴリ内のファイルを取得します。

さらにtf.data.TextLineDataset を使い、取得したファイル全体を１行ずつ分割し、データセットオブジェクトとして取り出します。

データセット作成

import os
import tensorflow as tf
import tensorflow_datasets as tfds
import Mecab

text_datasets = []
text_dir = os.path.join(os.getcwd(),"text")
for d in tf.compat.v1.gfile.ListDirectory(text_dir):
    data_dir = os.path.join(text_dir,d)
    if os.path.isdir(data_dir):
        for file_name in tf.compat.v1.gfile.ListDirectory(data_dir):
            text_dataset = tf.data.TextLineDataset(os.path.join(data_dir,file_name))
            text_datasets.append(text_dataset)

parsed_dataset = text_datasets[0]
for text_dataset in text_datasets[1:]:
    parsed_dataset = parsed_dataset.concatenate(text_dataset)

ここで一度、iter関数を使用して作成したデータセットの中身を確認してみます。

plane_dataset = text_datasets[0]
iter_plane_dataset = iter(plane_dataset)
print(next(iter_plane_dataset).numpy().decode('utf-8'))
print(next(iter_plane_dataset).numpy().decode('utf-8'))
print(next(iter_plane_dataset).numpy().decode('utf-8'))

実行結果

http://news.livedoor.com/article/detail/4778030/
2010-05-22T14:30:00+0900
友人代表のスピーチ、独女はどうこなしている？

ディレクトリ内のファイルから、データセットが正しく作成されることが確認できました。

分かち書きデータセットを作成

続いてデータセット内のコーパスにトークナイズ処理を施し、トークンごとのエンコード辞書の作成を行います。

tensorflow_datasets にはmapというパイプライン処理があり、パイプラインを使うことでデータセットを読み込むタイミングで分かち書きを施すことができます。

パイプラインはtf.data.TextLineDatasetのメソッドチェーンにmapを指定し、mapの引数に実際にデータセットに施したい処理を記述します。

ここでは 分かち書きを行うparse関数のラッパーとして parse_text という関数を追加しています。

分かち書きデータセット作成

def parse(text):
    '''
    データセットのテキストを1行ごとに分かち書き
    '''
    return tagger.parse(text.numpy().decode('utf-8')).split("\n")[0]

@tf.function
def parse_text(text,label):
    parsed = tf.py_function(parse,[text],[tf.string])
    return parsed[0],label

text_datasets = []
text_dir = os.path.join(os.getcwd(),"text")
for d in tf.compat.v1.gfile.ListDirectory(text_dir):
    data_dir = os.path.join(text_dir,d)
    label = int(d == "dokujo-tsushin") # ディレクトリがdokujo-tsushinだったらTrue
    if os.path.isdir(data_dir):
        for file_name in tf.compat.v1.gfile.ListDirectory(data_dir):
            text_dataset = tf.data.TextLineDataset(os.path.join(data_dir,file_name)).map(lambda ex: parse_text(ex,label)) # 分かち書きコーパスとラベルのペアをデータセットに指定する
            text_datasets.append(text_dataset)

# テキスト単位の配列に変換する
parsed_dataset = text_datasets[0]
for text_dataset in text_datasets[1:]:
    parsed_dataset = parsed_dataset.concatenate(text_dataset)

中身を取り出すと、テキストが分かち書きされた結果が確認できます。

for txt,label in parsed_dataset.take(3):
    print('{} {}'.format(txt.numpy().decode('utf-8'), label))

結果

http :// news . livedoor . com / article / detail / 6317287 / 1
2012 - 02 - 27 T 14 : 54 : 00 + 0900 1
話題 の 格言 『 ボーダー を 着る 女 は 、 95 ％ モテ ない ！ 』 は 本当 か ？ 1

正しく分かち書きされていることが確認できました。

トークナイザの作成

トークン単位に分かち書きしたデータセットが手に入ったので、このデータセットからエンコーダーを作成します。

Tensorflow Datasetsにはデータセットに対する様々な処理を簡単に施せる便利なAPIが揃っており、 分かち書きしたテキストをトークナイズするためにtfds.features.text.Tokenizer()を使います。

トークナイザを生成

tokenizer = tfds.features.text.Tokenizer()

生成したトークナイザに対して分かち書きしたデータセットを与えると、トークン化した結果を返します。
これをvocabulary_setコレクションに渡し、トークンだけの配列（ボキャブラリ）を作成します。

ボキャブラリ構築

vocabulary_set = set()
for line in parsed_dataset:
    some_tokens = tokenizer.tokenize(line.numpy())
    vocabulary_set.update(some_tokens)

作成したボキャブラリからエンコーダを生成します。

encoder = tfds.features.text.TokenTextEncoder(vocabulary_set, tokenizer=tokenizer)

encoder はボキャブラリ単位にIDを割り当てたトークン辞書として使います。

実際にエンコーダーの中身を確認してみましょう。

sample_text = next(iter_parsed_dataset)
encoded_example = encoder.encode(sample_text.numpy().decode('utf-8'))

for enc in encoded_example:
    print("{} -> {}".format(enc, encoder.decode([enc])))

トークン単位にIDが割り当てられているのが確認できます。

9058 -> もうすぐ
3331 -> ジューン
31397 -> ブライド
27220 -> と
2397 -> 呼ば
22334 -> れる
5673 -> ６月
31920 -> 独
32073 -> 女
20415 -> の
28229 -> 中
18191 -> に
986 -> は
8138 -> 自分
20415 -> の
10202 -> 式
986 -> は
36484 -> まだ
10551 -> な
26617 -> のに
2397 -> 呼ば
34 -> れ
27080 -> て
27726 -> ばかり
27509 -> という
6682 -> お祝い
903 -> 貧乏
5111 -> 状態
20415 -> の
35344 -> 人
4722 -> も
21766 -> 多い
20415 -> の
22018 -> で
986 -> は
31505 -> ない
12417 -> だろ
7199 -> う
23655 -> か
35100 -> さらに
22067 -> 出席
16090 -> 回数
18014 -> を
1399 -> 重ね
27080 -> て
27779 -> いく
27220 -> と
6523 -> こんな
6851 -> お願い
6769 -> ごと
18014 -> を
32709 -> さ
22334 -> れる
30766 -> こと
4722 -> も
3435 -> 少なく
31505 -> ない

補足ですがトークナイザはデフォルトで半角記号などを排除するので、それらが必要な場合は生成時にreserved_tokensで指定しておく必要があります。

reserved_tokensあり

tokenizer = tfds.features.text.Tokenizer(reserved_tokens=["/",":"])

エンコーダを使用して分かち書きテキストをエンコード

LSTMモデルのインプットはテキストそのものではなくエンコードされたIDを渡す必要があります。

トークン化済みのparsed_datasetデータセットに対し、さらにパイプライン処理を施し、先ほど作成したエンコーダーを使ってトークンからIDへの変換を行います。

def dataset_mapper(token,label):
    token = encoder.encode(token.numpy())
    label= np.array([label])
    return token,label

@tf.function
def tf_encode(token,label):
    return tf.py_function(dataset_mapper, [token,label], [tf.int64, tf.int64])

new_dataset = parsed_dataset.map(tf_encode)

トレーニング用にデータセットをバッチ分割します。

BATCH_SIZE = 128
new_dataset = new_dataset.padded_batch(BATCH_SIZE, padded_shapes=([-1],[-1]))
new_dataset = new_dataset.prefetch(tf.data.experimental.AUTOTUNE)

バッチ分割されているか確認します。

text_batch, label_batch = next(iter(new_dataset))
print(text_batch.shape)
print(label_batch.shape)

指定したサイズ(128)でデータセットが分割されているのが確認できました。

(128, 155)
(128, 1)

LSTMモデルの作成

データセットの準備ができたので、モデル作成に入ります。

Tensorflow 2.0はkerasと完全に統合され、Eagerモードがデフォルトになったので直感的にネットワークが組めるようになりました。

X = tf.keras.Input(shape=(None,), batch_size=BATCH_SIZE)
embedded = tf.keras.layers.Embedding(encoder.vocab_size, 128)(X)
lstm = tf.keras.layers.Bidirectional(
      tf.keras.layers.LSTM(128, dropout=0.4, recurrent_dropout=0.4)
  )(embedded)
fully_connected = tf.keras.layers.Dense(units=256, activation='relu')(lstm)
Y = tf.keras.layers.Dense(1, activation='softmax')(fully_connected)

model = tf.keras.Model(inputs=X, outputs=Y)

model.compile(loss='categorical_crossentropy',
              optimizer=tf.keras.optimizers.Adam(1e-7))

model.summary()を呼び出すとモデルの要約を確認できます。

model.summary()

Model: "model_1"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_2 (InputLayer)         [(128, None)]             0         
_________________________________________________________________
embedding_1 (Embedding)      (128, None, 128)          4324352   
_________________________________________________________________
bidirectional_1 (Bidirection (128, 256)                263168    
_________________________________________________________________
dense_1 (Dense)              (128, 256)                65792     
_________________________________________________________________
final_layer (Dense)          (128, 1)                  257       
=================================================================
Total params: 4,653,569
Trainable params: 4,653,569
Non-trainable params: 0
_________________________________________________________________

Trainable params: 4,653,569となっていることから、このモデルでは470万近い変数がトレーニング対象になることが分かります。

モデルのトレーニングにはmodel.fit()を使います。

model.fit(new_dataset,epochs=5)

Epoch 1/5
298/298 [==============================] - 140s 471ms/step - loss: 0.5049 - accuracy: 0.8531
Epoch 2/5
298/298 [==============================] - 137s 459ms/step - loss: 0.7768 - accuracy: 0.5536
Epoch 3/5
298/298 [==============================] - 245s 824ms/step - loss: 0.6405 - accuracy: 0.6531
Epoch 4/5
298/298 [==============================] - 327s 1s/step - loss: 0.5248 - accuracy: 0.8055
Epoch 5/5
298/298 [==============================] - 417s 1s/step - loss: 0.4686 - accuracy: 0.8019

LSTMは時系列にデータを扱うため、シーケンシャルな学習になります。GPUを使ってもさほど効果が無いという意味でお財布には優しいです。ただし学習完了まで時間がかかります。。。

BERTモデルの作成

前項ではlivedoorコーパスに対してLSTMで分類モデルを組んで実際にトレーニングを行いましたが、そのためにデータセットを分かち書きしてトークン化し、エンコーダーを作ってデータセットをID化するという手順を踏む必要がありました。

また、文章の意味合いを理解するために時系列に強いLSTMというモデルを使いましたが、LSTMはシーケンシャルに学習するためGPUを使ってもスケールしないため、GPU資源を有効に使いより良い精度を求めるには、Transformerのような複雑なモデルを自前で構築する必要があります。

BERTはtransformerを発展させたモデルであり、hugging faseのtransformersライブラリを使うことでこれらを簡単に実装することができ、さらにMecabで分かち書きされたサブワード辞書も内包されているため、前項で行った分かち書きやエンコーダー作成のような前処理を書く必要がありません。

早速先ほどのLDSTMの分類モデルをBERTに書き換えてみたいと思います。

BERTモデルの呼び出し

hugging headsのBERTモデルを使うためには transformers というライブラリが必要なので、pip経由でインストールします。

pip install transformers

エンコーダーはAutoTokenizer.from_pretrained()で呼び出せます。bert-base-japaneseはWikipediaの日本語版で事前学習済みのサブワード辞書となります。

from transformers import AutoTokenizer
from transformers import BertJapaneseTokenizer

tokenizer = AutoTokenizer.from_pretrained("bert-base-japanese")

次に事前学習済みのBERTモデルを呼び出します。
モデルの呼び出しはTFBertModel.from_pretrained()を使います。

from transformers import TFBertModel

bert = TFBertModel.from_pretrained('bert-base-japanese')

これでBERTを使う準備が完了しました。

データセット作成

BERTのトークナイザはトークナイズ処理を内包しているので、データセット作成時の分かち書きは不要です。データセット作成処理を一部修正します。

import os
import tensorflow as tf
import tensorflow_datasets as tfds
import numpy as np

@tf.function
def mapping(text,label):
    return text,label

text_datasets = []
text_dir = os.path.join(os.getcwd(),"text")
for d in tf.compat.v1.gfile.ListDirectory(text_dir):
    data_dir = os.path.join(text_dir,d)
    label = int(d == "dokujo-tsushin") # ディレクトリがdokujo-tsushinだったらTrue
    if os.path.isdir(data_dir):
        for file_name in tf.compat.v1.gfile.ListDirectory(data_dir):
            text_dataset = tf.data.TextLineDataset(os.path.join(data_dir,file_name)).map(lambda ex: mapping(ex,label)) 
            text_datasets.append(text_dataset)

# 分かち書きせずテキスト単位の配列に変換する
new_dataset = text_datasets[0]
for text_dataset in text_datasets[1:]:
    new_dataset = new_dataset.concatenate(text_dataset)

結果を確認します。

for txt,label in new_dataset.take(3):
    print('{} {}'.format(txt.numpy().decode('utf-8'), label))

分かち書きされていないことを確認します。

http://news.livedoor.com/article/detail/6634416/ 0
2012-06-07T10:45:00+0900 0
ハーフナー・マイク、番組のはからいに「変な汗出るわ」 0

入力値

BERTはMasked Language ModelとNext Sentence Predictionという2つの事前学習を行うことで汎用的な言語モデルを獲得します。BERTを利用した自然言語処理のモデルは、事前学習の結果に対してファインチューニングを施すことで精度の高いモデル構築が可能になります。
以下に2つのモデルの簡単な説明を示します。

Masked Language Model

Masked Language Modelは与えられた文章の一部の単語をマスクし、その単語を推論するタスクです。これにより文章内の単語間の関係性を学習します。

Next Sentence Prediction

もう一つのタスクであるNext Sentence Predictionは文章の文脈を推論するタスクです。
Next Sentence Predictionでは2つの文章を入力として受け取り、その2文章が文脈的につながりがあるか、まったく関係がないかを学習します。

transformersのBERTを使う場合、データセットを2つのモデルを使う形に合わせる必要があります。
具体的にはパイプライン処理でデータセットをinput_ids、token_type_ids、attention_maskの3つをキーとしたディクショナリに変更する必要があります。

input_idsはテキストをサブワードで分かち書きし、その結果をエンコードしたIDです。事前訓練タスクではinput_idsのうちいくつかの単語がランダムでマスク化され、マスク化された単語を推論します。

token_type_idsはNext Sentence Predictionに使用する入力値で、文章を2つに分割するために使用します。
先行するテキストのトークン位置には0を、後続テキストには1をセットします。ファインチューニングの場合Next Sentence Predictionは行わないのですべて0を指定します。

attentoin_maskは、パディング位置です。
例えばmex_length=128とした場合、入力文字列が128文字に達しない場合は残りのトークンには[PAD]という特別なトークンで字埋めを行います。
これをパディング処理といいますが、attentoin_maskは実際の入力トークンとパディングトークンを区別するために使用します。

それではデータセットをBERTの入力に揃えるためのパイプラインを追加します。

max_length = 128

def tokenize_map_fn(tokenizer, max_length=128):

    """map function for pretrained tokenizer"""
    def _tokenize(text_a, label):
        inputs = tokenizer.encode_plus(
            text_a.numpy().decode('utf-8'),
            add_special_tokens=True,
        )
        input_ids, token_type_ids = inputs["input_ids"], inputs["token_type_ids"]
        attention_mask = [1] * len(input_ids)
        return input_ids, token_type_ids, attention_mask, label

    def _map_fn(text,label):
        out = tf.py_function(_tokenize, inp=[text, label], Tout=(tf.int32, tf.int32, tf.int32, tf.int32))
        return (
            {"input_ids": out[0], "token_type_ids": out[1], "attention_mask": out[2]},
            out[3]
        )

    return _map_fn

def load_dataset(data, tokenizer, max_length=128, train_batch=32):
    train_dataset = data.map(tokenize_map_fn(tokenizer, max_length=max_length))
    train_dataset = train_dataset.shuffle(train_batch).padded_batch(train_batch, padded_shapes=({'input_ids': [-1], 'token_type_ids': [-1], 'attention_mask': [-1]}, []))
    return train_dataset

train_dataset = load_dataset(new_dataset,tokenizer,max_length=max_length,train_batch=BATCH_SIZE)

データセットの中身を確認します。

for data in train_dataset.take(1):
    print('input_ids is')
    tf.print(data[0]['input_ids'])
    print('token_type_ids is')
    tf.print(data[0]['token_type_ids'])
    print('attention_mask is')
    tf.print(data[0]['attention_mask'])

結果を確認します。

input_ids is
[[2 21313 16831 ... 0 0 0]
 [2 908 61 ... 0 0 0]
 [2 9241 590 ... 0 0 0]
 ...
 [2 3 0 ... 0 0 0]
 [2 11695 5 ... 0 0 0]
 [2 3 0 ... 0 0 0]]
token_type_ids is
[[0 0 0 ... 0 0 0]
 [0 0 0 ... 0 0 0]
 [0 0 0 ... 0 0 0]
 ...
 [0 0 0 ... 0 0 0]
 [0 0 0 ... 0 0 0]
 [0 0 0 ... 0 0 0]]
attention_mask is
[[1 1 1 ... 0 0 0]
 [1 1 1 ... 0 0 0]
 [1 1 1 ... 0 0 0]
 ...
 [1 1 0 ... 0 0 0]
 [1 1 1 ... 0 0 0]
 [1 1 0 ... 0 0 0]]

正しくBERTの入力値ができたので、続けてモデル構築を行います。

モデル構築

TFBertModel.from_pretrained('bert-base-japanese') により事前訓練済みのBERTの重みが呼び出せます。
インスタンス化したモデルに、先ほどの入力を与えるとlast_hidden_state、pooler_output、hidden_states、attentionsの4つの値をtupleとして返します。

last_hidden_stateにはモデルの最後の隠れ状態、pooler_outputはCLSと呼ばれる文章の平均値のようなもの、hidden_statesには全隠れ状態が、attentionsにはAttentionレイヤの出力が入っています。

今回はこれらのうちpooler_outputを文章の要約として全層結合に渡し、その結果を学習させます。

input_ids = tf.keras.layers.Input(shape=(max_length, ), dtype='int32', name='input_ids')
attention_mask = tf.keras.layers.Input(shape=(max_length, ), dtype='int32', name='attention_mask')
token_type_ids = tf.keras.layers.Input(shape=(max_length, ), dtype='int32', name='token_type_ids')
inputs = [input_ids, attention_mask, token_type_ids]

bert = TFBertModel.from_pretrained('bert-base-japanese')
bert.trainable = False
x = bert(inputs)

out = x[1]

fully_connected = tf.keras.layers.Dense(256, activation='relu')(out)
Y = tf.keras.layers.Dense(1, activation='sigmoid')(fully_connected)

model = tf.keras.Model(inputs=inputs, outputs=Y)

model.compile(loss=tf.keras.losses.SparseCategoricalCrossentropy(),
              optimizer=tf.keras.optimizers.Adam(1e-7))

モデルができたので、要約を確認します。

model.summary()

Model: "model_1"
__________________________________________________________________________________________________
Layer (type)                    Output Shape         Param #     Connected to                     
==================================================================================================
input_ids (InputLayer)          [(None, 128)]        0                                            
__________________________________________________________________________________________________
attention_mask (InputLayer)     [(None, 128)]        0                                            
__________________________________________________________________________________________________
token_type_ids (InputLayer)     [(None, 128)]        0                                            
__________________________________________________________________________________________________
bert (TFBertMainLayer)          ((None, 128, 768), ( 110617344   input_ids[0][0]                  
                                                                 attention_mask[0][0]             
                                                                 token_type_ids[0][0]             
__________________________________________________________________________________________________
dense_2 (Dense)                 (None, 256)          196864      bert[0][1]                       
__________________________________________________________________________________________________
dense_3 (Dense)                 (None, 1)            257         dense_2[0][0]                    
==================================================================================================
Total params: 110,814,465
Trainable params: 197,121
Non-trainable params: 110,617,344

BERTのモデルは変数が非常に多いためTotal params: 110,814,465となりますが、今回はBERTのプレトレーニングモデルはファインチューニングせず、ネットワークの全結合層だけ訓練するため学習すべき変数の数はTrainable params: 197,121となります。

bert.trainable = Falseを外すとBERTのモデルをファインチューニングできますが、Trainableな変数が一気に増えるのでGPUが必須になります。

LSTM同様model.fit()で訓練を行います。

model.fit(train_dataset,epochs=5)

Epoch 1/5
298/298 [==============================] - 1082s 4s/step - loss: 4.7969
Epoch 2/5
298/298 [==============================] - 1062s 4s/step - loss: 4.8130
Epoch 3/5
298/298 [==============================] - 1071s 4s/step - loss: 4.8130
Epoch 4/5
298/298 [==============================] - 1072s 4s/step - loss: 4.8130
Epoch 5/5
298/298 [==============================] - 1060s 4s/step - loss: 4.8130

Epoch数が5回と少なく、またBERTモジュール自体のトレーニングを行っていないため精度はほぼ向上しませんでした。。。

試しにデータセット内の1文で推論してみます。

sample_pred_text = 'ハーフナー・マイク、番組のはからいに「変な汗出るわ」'
encoded = tokenizer.encode_plus(
            sample_pred_text,
            sample_pred_text,
            add_special_tokens=True,
            max_length=128,
            pad_to_max_length=True,
            return_attention_mask=True
        )

inputs = {"input_ids": tf.expand_dims(encoded["input_ids"],0), 
          "token_type_ids": tf.expand_dims(encoded["token_type_ids"],0), 
          "attention_mask": tf.expand_dims(encoded["attention_mask"],0)
         }

res = model.predict_on_batch(inputs)
res.numpy()

結果

array([[1.]], dtype=float32)

結果は1がdokujo-tsushinですがこの文章はスポーツ記事のようですので、この学習量では正しい結果は返ってきませんでした。

bert.trainableを外し、エポック数を増やすことで改善すると思いますが、それに比例して必要なGPUのスペックも高くなり、訓練時間も長くなりますので気を付けてください。