質問があります? TensorFlowフォーラム訪問フォーラムでコミュニティとつながる

TPUでBERTを使用してGLUEタスクを解決する

TensorFlow.orgで表示 GoogleColabで実行 GitHubで表示ノートブックをダウンロードするTFハブモデルを参照してください

BERTは、自然言語処理における多くの問題を解決するために使用できます。 GLUEベンチマークから、多くのタスクに対してBERTを微調整する方法を学習します

  1. CoLA (Corpus of Linguistic Acceptability):文は文法的に正しいですか?

  2. SST-2 (Stanford Sentiment Treebank):タスクは、特定の文の感情を予測することです。

  3. MRPC (Microsoft Research Paraphrase Corpus):文のペアが意味的に同等であるかどうかを判断します。

  4. QQP (Quora Question Pairs2):質問のペアが意味的に同等であるかどうかを判断します。

  5. MNLI (Multi-Genre Natural Language Inference):前提文と仮説文が与えられた場合、タスクは、前提が仮説を伴うか(含意)、仮説と矛盾するか(矛盾)、どちらでもないか(中立)を予測することです。

  6. QNLI (質問応答自然言語推論):タスクは、コンテキスト文に質問への回答が含まれているかどうかを判断することです。

  7. RTE (テキスト含意の認識):文が特定の仮説を伴うかどうかを判断します。

  8. WNLI (Winograd Natural Language Inference):タスクは、代名詞が置換された文が元の文に含まれているかどうかを予測することです。

このチュートリアルには、TPUでこれらのモデルをトレーニングするための完全なエンドツーエンドのコードが含まれています。 1行を変更することで、このノートブックをGPUで実行することもできます(以下で説明)。

このノートブックでは、次のことを行います。

  • TensorFlowハブからBERTモデルをロードします
  • GLUEタスクの1つを選択し、データセットをダウンロードします
  • テキストを前処理する
  • BERTを微調整します(例は単一文と複数文のデータセットに示されています)
  • トレーニング済みモデルを保存して使用します

セットアップ

別のモデルを使用してテキストを前処理してから、BERTを微調整します。このモデルは、以下にインストールするtensorflow / textに依存します。

pip install -q -U tensorflow-text

Tensorflow /モデルからAdamWオプティマイザーを使用して、BERTを微調整します。BERTもインストールします。

pip install -q -U tf-models-official
pip install -U tfds-nightly
import os
import tensorflow as tf
import tensorflow_hub as hub
import tensorflow_datasets as tfds
import tensorflow_text as text  # A dependency of the preprocessing model
import tensorflow_addons as tfa
from official.nlp import optimization
import numpy as np

tf.get_logger().setLevel('ERROR')

次に、TFHubのクラウドストレージバケットからチェックポイントを直接読み取るようにTFHubを構成します。これは、TPUでTFHubモデルを実行する場合にのみ推奨されます。

この設定がないと、TFHubは圧縮ファイルをダウンロードし、チェックポイントをローカルに抽出します。これらのローカルファイルからロードしようとすると、次のエラーで失敗します。

InvalidArgumentError: Unimplemented: File system scheme '[local]' not implemented

これは、 TPUがCloudStorageバケットから直接読み取ることしかできないためです。

os.environ["TFHUB_MODEL_LOAD_FORMAT"]="UNCOMPRESSED"

TPUワーカーに接続します

次のコードはTPUワーカーに接続し、TensorFlowのデフォルトデバイスをTPUワーカーのCPUデバイスに変更します。また、この1つのTPUワーカーで使用可能な8つの個別のTPUコアにモデルトレーニングを配布するために使用するTPU配布戦略も定義します。詳細については、TensorFlowのTPUガイドを参照してください。

import os

if os.environ['COLAB_TPU_ADDR']:
  cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
  tf.config.experimental_connect_to_cluster(cluster_resolver)
  tf.tpu.experimental.initialize_tpu_system(cluster_resolver)
  strategy = tf.distribute.TPUStrategy(cluster_resolver)
  print('Using TPU')
elif tf.test.is_gpu_available():
  strategy = tf.distribute.MirroredStrategy()
  print('Using GPU')
else:
  raise ValueError('Running on CPU is not recommended.')
Using TPU

TensorFlowハブからモデルを読み込んでいます

ここで、TensorFlowハブからロードするBERTモデルを選択し、微調整できます。選択できるBERTモデルは複数あります。

  • BERT- ベースケースなし、および元のBERT作成者によってリリースされたトレーニング済みウェイトを備えた7つのモデル
  • 小さなBERTの一般的なアーキテクチャは同じですが、Transformerブロックが少ないか小さいため、速度、サイズ、品質の間のトレードオフを調べることができます。
  • ALBERT :レイヤー間でパラメーターを共有することでモデルサイズを(計算時間ではなく)削減する4つの異なるサイズの「ALiteBERT」。
  • BERTエキスパート:すべてBERTベースのアーキテクチャを備えているが、ターゲットタスクとより緊密に連携するために、さまざまな事前トレーニングドメインから選択できる8つのモデル。
  • ElectraはBERTと同じアーキテクチャ(3つの異なるサイズ)を持っていますが、Generative Adversarial Network(GAN)に似たセットアップで弁別子として事前にトレーニングされています。
  • トーキング・ヘッズ・アテンションとゲート付きGELU [ ベースラージ]を備えたBERTには、Transformerアーキテクチャーのコアに2つの改良が加えられています。

詳細については、上記のリンク先のモデルドキュメントを参照してください。

このチュートリアルでは、BERTベースから始めます。より大きく、より新しいモデルを使用して精度を高めたり、より小さなモデルを使用してトレーニング時間を短縮したりできます。モデルを変更するには、1行のコードを切り替えるだけです(以下を参照)。すべての違いは、TensorFlowHubからダウンロードするSavedModelにカプセル化されています。

BERTモデルを選択して微調整します

BERT model selected           : https://tfhub.dev/tensorflow/bert_en_uncased_L-12_H-768_A-12/3
Preprocessing model auto-selected: https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3

テキストを前処理する

BERT colabを使用したテキスト分類では、前処理モデルがBERTエンコーダーに直接埋め込まれて使用されます。

このチュートリアルでは、Dataset.mapを使用して、トレーニング用の入力パイプラインの一部として前処理を実行し、それを推論のためにエクスポートされるモデルにマージする方法を示します。このように、TPU自体は数値入力を必要としますが、トレーニングと推論の両方が生のテキスト入力から機能します。

TPU要件はさておき、入力パイプラインでパフォーマンスの前処理を非同期で実行するのに役立ちます(詳細については、 tf.dataパフォーマンスガイドを参照してください)。

このチュートリアルでは、多入力モデルを構築する方法、およびBERTへの入力のシーケンス長を調整する方法も示します。

前処理モデルを示しましょう。

bert_preprocess = hub.load(tfhub_handle_preprocess)
tok = bert_preprocess.tokenize(tf.constant(['Hello TensorFlow!']))
print(tok)
<tf.RaggedTensor [[[7592], [23435, 12314], [999]]]>

各前処理モデルは、トークンのリスト(上記のtok )とシーケンス長の引数を.bert_pack_inputs(tensors, seq_length)メソッド.bert_pack_inputs(tensors, seq_length)も提供します。これにより、入力がパックされ、BERTモデルで期待される形式でテンソルの辞書が作成されます。

text_preprocessed = bert_preprocess.bert_pack_inputs([tok, tok], tf.constant(20))

print('Shape Word Ids : ', text_preprocessed['input_word_ids'].shape)
print('Word Ids       : ', text_preprocessed['input_word_ids'][0, :16])
print('Shape Mask     : ', text_preprocessed['input_mask'].shape)
print('Input Mask     : ', text_preprocessed['input_mask'][0, :16])
print('Shape Type Ids : ', text_preprocessed['input_type_ids'].shape)
print('Type Ids       : ', text_preprocessed['input_type_ids'][0, :16])
Shape Word Ids :  (1, 20)
Word Ids       :  tf.Tensor(
[  101  7592 23435 12314   999   102  7592 23435 12314   999   102     0
     0     0     0     0], shape=(16,), dtype=int32)
Shape Mask     :  (1, 20)
Input Mask     :  tf.Tensor([1 1 1 1 1 1 1 1 1 1 1 0 0 0 0 0], shape=(16,), dtype=int32)
Shape Type Ids :  (1, 20)
Type Ids       :  tf.Tensor([0 0 0 0 0 0 1 1 1 1 1 0 0 0 0 0], shape=(16,), dtype=int32)

注意すべき詳細は次のとおりです。

  • input_maskマスクを使用すると、モデルでコンテンツとパディングを明確に区別できます。マスクはinput_word_idsと同じ形状であり、 input_word_idsがパディングされていない場所には1が含まれます。
  • input_type_ids同じ形状を有しているinput_maskが、非パディング領域の内側に、トークンがその一部である文を示す0又は1を含んでいます。

次に、このすべてのロジックをカプセル化する前処理モデルを作成します。モデルは文字列を入力として受け取り、BERTに渡すことができる適切にフォーマットされたオブジェクトを返します。

各BERTモデルには特定の前処理モデルがあります。BERTのモデルドキュメントに記載されている適切なモデルを使用してください。

def make_bert_preprocess_model(sentence_features, seq_length=128):
  """Returns Model mapping string features to BERT inputs.

  Args:
    sentence_features: a list with the names of string-valued features.
    seq_length: an integer that defines the sequence length of BERT inputs.

  Returns:
    A Keras Model that can be called on a list or dict of string Tensors
    (with the order or names, resp., given by sentence_features) and
    returns a dict of tensors for input to BERT.
  """

  input_segments = [
      tf.keras.layers.Input(shape=(), dtype=tf.string, name=ft)
      for ft in sentence_features]

  # Tokenize the text to word pieces.
  bert_preprocess = hub.load(tfhub_handle_preprocess)
  tokenizer = hub.KerasLayer(bert_preprocess.tokenize, name='tokenizer')
  segments = [tokenizer(s) for s in input_segments]

  # Optional: Trim segments in a smart way to fit seq_length.
  # Simple cases (like this example) can skip this step and let
  # the next step apply a default truncation to approximately equal lengths.
  truncated_segments = segments

  # Pack inputs. The details (start/end token ids, dict of output tensors)
  # are model-dependent, so this gets loaded from the SavedModel.
  packer = hub.KerasLayer(bert_preprocess.bert_pack_inputs,
                          arguments=dict(seq_length=seq_length),
                          name='packer')
  model_inputs = packer(truncated_segments)
  return tf.keras.Model(input_segments, model_inputs)

前処理モデルを示しましょう。 2つの文の入力(input1とinput2)を使用してテストを作成します。出力は、BERTモデルが入力として期待するものです: input_word_idsinput_masksinput_type_ids

test_preprocess_model = make_bert_preprocess_model(['my_input1', 'my_input2'])
test_text = [np.array(['some random test sentence']),
             np.array(['another sentence'])]
text_preprocessed = test_preprocess_model(test_text)

print('Keys           : ', list(text_preprocessed.keys()))
print('Shape Word Ids : ', text_preprocessed['input_word_ids'].shape)
print('Word Ids       : ', text_preprocessed['input_word_ids'][0, :16])
print('Shape Mask     : ', text_preprocessed['input_mask'].shape)
print('Input Mask     : ', text_preprocessed['input_mask'][0, :16])
print('Shape Type Ids : ', text_preprocessed['input_type_ids'].shape)
print('Type Ids       : ', text_preprocessed['input_type_ids'][0, :16])
Keys           :  ['input_type_ids', 'input_word_ids', 'input_mask']
Shape Word Ids :  (1, 128)
Word Ids       :  tf.Tensor(
[ 101 2070 6721 3231 6251  102 2178 6251  102    0    0    0    0    0
    0    0], shape=(16,), dtype=int32)
Shape Mask     :  (1, 128)
Input Mask     :  tf.Tensor([1 1 1 1 1 1 1 1 1 0 0 0 0 0 0 0], shape=(16,), dtype=int32)
Shape Type Ids :  (1, 128)
Type Ids       :  tf.Tensor([0 0 0 0 0 0 1 1 1 0 0 0 0 0 0 0], shape=(16,), dtype=int32)

定義した2つの入力に注目して、モデルの構造を見てみましょう。

tf.keras.utils.plot_model(test_preprocess_model)
('You must install pydot (`pip install pydot`) and install graphviz (see instructions at https://graphviz.gitlab.io/download/) ', 'for plot_model/model_to_dot to work.')

データセットからのすべての入力に前処理を適用するには、データセットからのmap関数を使用しmap 。結果は、パフォーマンスのためにキャッシュされます

AUTOTUNE = tf.data.AUTOTUNE


def load_dataset_from_tfds(in_memory_ds, info, split, batch_size,
                           bert_preprocess_model):
  is_training = split.startswith('train')
  dataset = tf.data.Dataset.from_tensor_slices(in_memory_ds[split])
  num_examples = info.splits[split].num_examples

  if is_training:
    dataset = dataset.shuffle(num_examples)
    dataset = dataset.repeat()
  dataset = dataset.batch(batch_size)
  dataset = dataset.map(lambda ex: (bert_preprocess_model(ex), ex['label']))
  dataset = dataset.cache().prefetch(buffer_size=AUTOTUNE)
  return dataset, num_examples

モデルを定義する

これで、前処理された入力をBERTエンコーダーに送り、線形分類器を上に配置し(または必要に応じて他のレイヤーの配置)、正則化にドロップアウトを使用することで、文または文のペア分類のモデルを定義する準備が整いました。

def build_classifier_model(num_classes):
  inputs = dict(
      input_word_ids=tf.keras.layers.Input(shape=(None,), dtype=tf.int32),
      input_mask=tf.keras.layers.Input(shape=(None,), dtype=tf.int32),
      input_type_ids=tf.keras.layers.Input(shape=(None,), dtype=tf.int32),
  )

  encoder = hub.KerasLayer(tfhub_handle_encoder, trainable=True, name='encoder')
  net = encoder(inputs)['pooled_output']
  net = tf.keras.layers.Dropout(rate=0.1)(net)
  net = tf.keras.layers.Dense(num_classes, activation=None, name='classifier')(net)
  return tf.keras.Model(inputs, net, name='prediction')

いくつかの前処理された入力でモデルを実行してみましょう。

test_classifier_model = build_classifier_model(2)
bert_raw_result = test_classifier_model(text_preprocessed)
print(tf.sigmoid(bert_raw_result))
tf.Tensor([[0.20310709 0.39733988]], shape=(1, 2), dtype=float32)

モデルの構造を見てみましょう。 3つのBERTの予想される入力を確認できます。

tf.keras.utils.plot_model(test_classifier_model)
('You must install pydot (`pip install pydot`) and install graphviz (see instructions at https://graphviz.gitlab.io/download/) ', 'for plot_model/model_to_dot to work.')

GLUEからタスクを選択してください

GLUEベンチマークスイートのTensorFlowDataSetを使用します。

Colabを使用すると、これらの小さなデータセットをローカルファイルシステムにダウンロードできます。別のTPUワーカーホストはcolabランタイムのローカルファイルシステムにアクセスできないため、以下のコードはそれらを完全にメモリに読み込みます。

より大きなデータセットの場合は、独自のGoogle Cloud Storageバケットを作成し、そこからTPUワーカーにデータを読み取らせる必要があります。詳細については、 TPUガイドをご覧ください

CoLaデータセット(単一文の場合)またはMRPC(複数文の場合)から始めることをお勧めします。これらは小さく、微調整に時間がかからないためです。

Using glue/cola from TFDS
This dataset has 10657 examples
Number of classes: 2
Features ['sentence']
Splits ['test', 'train', 'validation']
Here are some sample rows from glue/cola dataset
['unacceptable', 'acceptable']

sample row 1
b'It is this hat that it is certain that he was wearing.'
label: 1 (acceptable)

sample row 2
b'Her efficient looking up of the answer pleased the boss.'
label: 1 (acceptable)

sample row 3
b'Both the workers will wear carnations.'
label: 1 (acceptable)

sample row 4
b'John enjoyed drawing trees for his syntax homework.'
label: 1 (acceptable)

sample row 5
b'We consider Leslie rather foolish, and Lou a complete idiot.'
label: 1 (acceptable)

データセットは、問題の種類(分類または回帰)とトレーニングに適した損失関数も決定します。

def get_configuration(glue_task):

  loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

  if glue_task == 'glue/cola':
    metrics = tfa.metrics.MatthewsCorrelationCoefficient(num_classes=2)
  else:
    metrics = tf.keras.metrics.SparseCategoricalAccuracy(
        'accuracy', dtype=tf.float32)

  return metrics, loss

モデルをトレーニングする

最後に、選択したデータセットでモデルをエンドツーエンドでトレーニングできます。

分布

上部にあるセットアップコードを思い出してください。これは、colabランタイムを複数のTPUデバイスを持つTPUワーカーに接続しています。それらにトレーニングを配布するには、TPU配布戦略の範囲内でメインのKerasモデルを作成してコンパイルします。 (詳細については、 Kerasを使用した分散トレーニングを参照してください。)

一方、前処理はTPUではなくワーカーホストのCPUで実行されるため、前処理用のKerasモデルと、それにマップされたトレーニングおよび検証データセットは、配布戦略の範囲外で構築されます。 Model.fit()を呼び出すと、渡されたデータセットがモデルのレプリカに配布されます。

オプティマイザ

微調整は、BERT事前トレーニングからのオプティマイザー設定に従います( BERTでテキストを分類する場合のように):概念的な初期学習率の線形減衰を伴うAdamWオプティマイザーを使用し、最初のトレーニングステップの10%( num_warmup_steps )。 BERTの論文と一致して、初期学習率は微調整の方が小さくなっています(5e-5、3e-5、2e-5のベスト)。

epochs = 3
batch_size = 32
init_lr = 2e-5

print(f'Fine tuning {tfhub_handle_encoder} model')
bert_preprocess_model = make_bert_preprocess_model(sentence_features)

with strategy.scope():

  # metric have to be created inside the strategy scope
  metrics, loss = get_configuration(tfds_name)

  train_dataset, train_data_size = load_dataset_from_tfds(
      in_memory_ds, tfds_info, train_split, batch_size, bert_preprocess_model)
  steps_per_epoch = train_data_size // batch_size
  num_train_steps = steps_per_epoch * epochs
  num_warmup_steps = num_train_steps // 10

  validation_dataset, validation_data_size = load_dataset_from_tfds(
      in_memory_ds, tfds_info, validation_split, batch_size,
      bert_preprocess_model)
  validation_steps = validation_data_size // batch_size

  classifier_model = build_classifier_model(num_classes)

  optimizer = optimization.create_optimizer(
      init_lr=init_lr,
      num_train_steps=num_train_steps,
      num_warmup_steps=num_warmup_steps,
      optimizer_type='adamw')

  classifier_model.compile(optimizer=optimizer, loss=loss, metrics=[metrics])

  classifier_model.fit(
      x=train_dataset,
      validation_data=validation_dataset,
      steps_per_epoch=steps_per_epoch,
      epochs=epochs,
      validation_steps=validation_steps)
Fine tuning https://tfhub.dev/tensorflow/bert_en_uncased_L-12_H-768_A-12/3 model
/tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow/python/keras/engine/functional.py:591: UserWarning: Input dict contained keys ['idx', 'label'] which did not match any model input. They will be ignored by the model.
  [n for n in tensors.keys() if n not in ref_input_names])
Epoch 1/3
/tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow/python/framework/indexed_slices.py:449: UserWarning: Converting sparse IndexedSlices(IndexedSlices(indices=Tensor("AdamWeightDecay/gradients/StatefulPartitionedCall:1", shape=(None,), dtype=int32), values=Tensor("clip_by_global_norm/clip_by_global_norm/_0:0", dtype=float32), dense_shape=Tensor("AdamWeightDecay/gradients/StatefulPartitionedCall:2", shape=(None,), dtype=int32))) to a dense Tensor of unknown shape. This may consume a large amount of memory.
  "shape. This may consume a large amount of memory." % value)
266/267 [============================>.] - ETA: 0s - loss: 0.5234 - MatthewsCorrelationCoefficient: 0.0000e+00
/tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow/python/keras/metrics.py:257: UserWarning: Metric MatthewsCorrelationCoefficient implements a `reset_states()` method; rename it to `reset_state()` (without the final "s"). The name `reset_states()` has been deprecated to improve API consistency.
  'consistency.' % (self.__class__.__name__,))
267/267 [==============================] - 81s 73ms/step - loss: 0.5224 - MatthewsCorrelationCoefficient: 0.0000e+00 - val_loss: 0.5404 - val_MatthewsCorrelationCoefficient: 0.0000e+00
Epoch 2/3
267/267 [==============================] - 14s 54ms/step - loss: 0.3468 - MatthewsCorrelationCoefficient: 0.0000e+00 - val_loss: 0.6059 - val_MatthewsCorrelationCoefficient: 0.0000e+00
Epoch 3/3
267/267 [==============================] - 14s 54ms/step - loss: 0.2320 - MatthewsCorrelationCoefficient: 0.0000e+00 - val_loss: 0.6276 - val_MatthewsCorrelationCoefficient: 0.0000e+00

推論のためにエクスポート

前処理部分と、作成したばかりの微調整されたBERTを含む最終モデルを作成します。

推論時には、前処理をモデルの一部にする必要があります(それを行うトレーニングデータに関しては、別個の入力キューがなくなったため)。前処理は単なる計算ではありません。エクスポート用に保存されるKerasモデルにアタッチする必要がある独自のリソース(語彙テーブル)があります。この最終アセンブリが保存されます。

モデルをcolabに保存し、後でダウンロードして将来のために保存することができます( [表示]-> [目次]-> [ファイル] )。

main_save_path = './my_models'
bert_type = tfhub_handle_encoder.split('/')[-2]
saved_model_name = f'{tfds_name.replace("/", "_")}_{bert_type}'

saved_model_path = os.path.join(main_save_path, saved_model_name)

preprocess_inputs = bert_preprocess_model.inputs
bert_encoder_inputs = bert_preprocess_model(preprocess_inputs)
bert_outputs = classifier_model(bert_encoder_inputs)
model_for_export = tf.keras.Model(preprocess_inputs, bert_outputs)

print('Saving', saved_model_path)

# Save everything on the Colab host (even the variables from TPU memory)
save_options = tf.saved_model.SaveOptions(experimental_io_device='/job:localhost')
model_for_export.save(saved_model_path, include_optimizer=False,
                      options=save_options)
Saving ./my_models/glue_cola_bert_en_uncased_L-12_H-768_A-12
WARNING:absl:Found untraced functions such as restored_function_body, restored_function_body, restored_function_body, restored_function_body, restored_function_body while saving (showing 5 of 910). These functions will not be directly callable after loading.

モデルをテストする

最後のステップは、エクスポートされたモデルの結果をテストすることです。

比較のために、モデルをリロードし、データセットから分割されたテストからの入力を使用してテストしてみましょう。

with tf.device('/job:localhost'):
  reloaded_model = tf.saved_model.load(saved_model_path)

ユーティリティメソッド

テスト

with tf.device('/job:localhost'):
  test_dataset = tf.data.Dataset.from_tensor_slices(in_memory_ds[test_split])
  for test_row in test_dataset.shuffle(1000).map(prepare).take(5):
    if len(sentence_features) == 1:
      result = reloaded_model(test_row[0])
    else:
      result = reloaded_model(list(test_row))

    print_bert_results(test_row, result, tfds_name)
sentence: [b'the putter on the table left.']
This sentence is acceptable
BERT raw results: tf.Tensor([-1.9087534  2.614734 ], shape=(2,), dtype=float32)

sentence: [b'Doug cleared the table of dishes.']
This sentence is acceptable
BERT raw results: tf.Tensor([-2.746095   3.2254264], shape=(2,), dtype=float32)

sentence: [b'I ran into the baker from whom I bought these bagels.']
This sentence is acceptable
BERT raw results: tf.Tensor([-2.7611763  2.8297088], shape=(2,), dtype=float32)

sentence: [b'I presented it to John dead.']
This sentence is unacceptable
BERT raw results: tf.Tensor([ 2.3774471 -2.079125 ], shape=(2,), dtype=float32)

sentence: [b'The only thing capable of consuming this food has four legs and flies.']
This sentence is acceptable
BERT raw results: tf.Tensor([-1.4872739  1.9206202], shape=(2,), dtype=float32)

TF Servingでモデルを使用する場合は、名前付きシグニチャの1つを介してSavedModelを呼び出すことに注意してください。入力にいくつかの小さな違いがあることに注意してください。 Pythonでは、次のようにテストできます。

with tf.device('/job:localhost'):
  serving_model = reloaded_model.signatures['serving_default']
  for test_row in test_dataset.shuffle(1000).map(prepare_serving).take(5):
    result = serving_model(**test_row)
    # The 'prediction' key is the classifier's defined model name.
    print_bert_results(list(test_row.values()), result['prediction'], tfds_name)
sentence: b'There was believed to have been a riot in the kitchen.'
This sentence is acceptable
BERT raw results: tf.Tensor([-3.196773   2.9193494], shape=(2,), dtype=float32)

sentence: b'Tom swam the English Channel because he believed that Suzy ex.'
This sentence is acceptable
BERT raw results: tf.Tensor([-1.2277479  1.9688965], shape=(2,), dtype=float32)

sentence: b'I was convinced by John to leave.'
This sentence is acceptable
BERT raw results: tf.Tensor([-1.669661   2.0653794], shape=(2,), dtype=float32)

sentence: b'I can find the baker in whom to place your trust.'
This sentence is acceptable
BERT raw results: tf.Tensor([-2.5297134  2.9251165], shape=(2,), dtype=float32)

sentence: b'John promised Bill to be examined.'
This sentence is unacceptable
BERT raw results: tf.Tensor([ 2.2290623 -1.1185381], shape=(2,), dtype=float32)

できたね!保存したモデルは、プロセスでの提供または単純な推論に使用できます。APIは単純で、コードが少なく、保守が簡単です。

次のステップ

ベースBERTモデルの1つを試したので、他のモデルを試して、より高い精度を達成するか、より小さなモデルバージョンを使用することができます。

他のデータセットで試すこともできます。