<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Copyright 2019 The TensorFlow Authors.
Step1:  Fine-tuning a BERT model
<table class="tfo-notebook-buttons" align="left">
  <td>
    <a target="_blank" href="https
Step2:  Import libraries
Step3:  Resources
The following directory contains the BERT model's configuration, vocabulary, and a pre-trained checkpoint used in this tutorial
Step4:  Load and preprocess the dataset
This example uses the GLUE (General Language Understanding Evaluation) MRPC (Microsoft Research Paraphrase Corpus) dataset from TensorFlow Datasets (TFDS).
This dataset is not set up such that it can be directly fed into the BERT model. The following section handles the necessary preprocessing.
Get the dataset from TensorFlow Datasets
The GLUE MRPC (Dolan and Brockett, 2005) dataset is a corpus of sentence pairs automatically extracted from online news sources, with human annotations for whether the sentences in the pair are semantically equivalent. It has the following attributes
Step5:  The info object describes the dataset and its features
Step6:  The two classes are
Step7:  Here is one example from the training set
Step8:  Preprocess the data
The keys "sentence1" and "sentence2" in the GLUE MRPC dataset contain two input sentences for each example.
Because the BERT model from the Model Garden doesn't take raw text as input, two things need to happen first
Step9:  Let's tokenize a test sentence
Step10:  Learn more about the tokenization process in the Subword tokenization and Tokenizing with TensorFlow Text guides.
Pack the inputs
TensorFlow Model Garden's BERT model doesn't just take the tokenized strings as input. It also expects these to be packed into a particular format. tfm.nlp.layers.BertPackInputs layer can handle the conversion from a list of tokenized sentences to the input format expected by the Model Garden's BERT model.
tfm.nlp.layers.BertPackInputs packs the two input sentences (per example in the MRCP dataset) concatenated together. This input is expected to start with a [CLS] "This is a classification problem" token, and each sentence should end with a [SEP] "Separator" token.
Therefore, the tfm.nlp.layers.BertPackInputs layer's constructor takes the tokenizer's special tokens as an argument. It also needs to know the indices of the tokenizer's special tokens.
Step11:  The packer takes a list of tokenized sentences as input. For example
Step12:  Then, it returns a dictionary containing three outputs
Step13:  Put it all together
Combine these two parts into a keras.layers.Layer that can be attached to your model
Step14:  But for now just apply it to the dataset using Dataset.map, since the dataset you loaded from TFDS is a tf.data.Dataset object
Step15:  Here is an example batch from the processed dataset
Step16:  The input_word_ids contain the token IDs
Step17:  The mask allows the model to cleanly differentiate between the content and the padding. The mask has the same shape as the input_word_ids, and contains a 1 anywhere the input_word_ids is not padding.
Step18:  The "input type" also has the same shape, but inside the non-padded region, contains a 0 or a 1 indicating which sentence the token is a part of.
Step19:  Apply the same preprocessing to the validation and test subsets of the GLUE MRPC dataset
Step20:  Build, train and export the model
Now that you have formatted the data as expected, you can start working on building and training the model.
Build the model
The first step is to download the configuration file—config_dict—for the pre-trained BERT model
Step21:  The configuration file defines the core BERT model from the Model Garden, which is a Keras model that predicts the outputs of num_classes from the inputs with maximum sequence length max_seq_length.
Step22:  The classifier has three inputs and one output
Step23:  Run it on a test batch of data 10 examples from the training set. The output is the logits for the two classes
Step24:  The TransformerEncoder in the center of the classifier above is the bert_encoder.
If you inspect the encoder, notice the stack of Transformer layers connected to those same three inputs
Step25:  Restore the encoder weights
When built, the encoder is randomly initialized. Restore the encoder's weights from the checkpoint
Step26:  Note
Step27:  Linear decay from initial_learning_rate to zero over num_train_steps.
Step28:  Warmup to that value over warmup_steps
Step29:  The overall schedule looks like this
Step30:  Use tf.keras.optimizers.experimental.AdamW to instantiate the optimizer with that schedule
Step31:  Train the model
Set the metric as accuracy and the loss as sparse categorical cross-entropy. Then, compile and train the BERT classifier
Step32:  Now run the fine-tuned model on a custom example to see that it works.
Start by encoding some sentence pairs
Step33:  The model should report class 1 "match" for the first example and class 0 "no-match" for the second
Step34:  Export the model
Often the goal of training a model is to use it for something outside of the Python process that created it. You can do this by exporting the model using tf.saved_model. (Learn more in the Using the SavedModel format guide and the Save and load a model using a distribution strategy tutorial.)
First, build a wrapper class to export the model. This wrapper does two things
Step35:  Create an instance of this export-model and save it
Step36:  Reload the model and compare the results to the original
Step37:  Congratulations! You've used tensorflow_models to build a BERT-classifier, train it, and export for later use.
Optional
Step38:  Select the input preprocessor and the model from TF Hub and wrap them as hub.KerasLayer layers
Step39:  Test run the preprocessor on a batch of data
Step40:  At this point it would be simple to add a classification head yourself.
The Model Garden tfm.nlp.models.BertClassifier class can also build a classifier onto the TF Hub encoder
Step41:  The one downside to loading this model from TF Hub is that the structure of internal Keras layers is not restored. This makes it more difficult to inspect or modify the model.
The BERT encoder model—hub_classifier—is now a single layer
Step42:  For concrete examples of this approach, refer to Solve Glue tasks using BERT.
Optional
Step43:  The advantage to using config objects is that they don't contain any complicated TensorFlow objects, and can be easily serialized to JSON, and rebuilt. Here's the JSON for the above tfm.optimization.OptimizationConfig
Step44:  The tfm.optimization.optimizer_factory.OptimizerFactory can just as easily build the optimizer from the JSON dictionary
<ASSISTANT_TASK:>
Python Code:
#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
Explanation: Copyright 2019 The TensorFlow Authors.
End of explanation
!pip install -q -U "tensorflow-text==2.9.*"
!pip install -q tf-models-official
Explanation: Fine-tuning a BERT model
<table class="tfo-notebook-buttons" align="left">
  <td>
    <a target="_blank" href="https://www.tensorflow.org/text/tutorials/fine_tune_bert"><img src="https://www.tensorflow.org/images/tf_logo_32px.png" />View on TensorFlow.org</a>
  </td>
  <td>
    <a target="_blank" href="https://colab.research.google.com/github/tensorflow/text/blob/master/docs/tutorials/fine_tune_bert.ipynb"><img src="https://www.tensorflow.org/images/colab_logo_32px.png" />Run in Google Colab</a>
  </td>
  <td>
    <a target="_blank" href="https://github.com/tensorflow/text/blob/master/docs/tutorials/fine_tune_bert.ipynb"><img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />View source on GitHub</a>
  </td>
  <td>
    <a href="https://storage.googleapis.com/tensorflow_docs/text/docs/tutorials/fine_tune_bert.ipynb"><img src="https://www.tensorflow.org/images/download_logo_32px.png" />Download notebook</a>
  </td>
  <td>
    <a href="https://tfhub.dev/google/collections/bert"><img src="https://www.tensorflow.org/images/hub_logo_32px.png" />See TF Hub model</a>
  </td>
</table>
This tutorial demonstrates how to fine-tune a Bidirectional Encoder Representations from Transformers (BERT) (Devlin et al., 2018) model using TensorFlow Model Garden.
You can also find the pre-trained BERT model used in this tutorial on TensorFlow Hub (TF Hub). For concrete examples of how to use the models from TF Hub, refer to the Solve Glue tasks using BERT tutorial. If you're just trying to fine-tune a model, the TF Hub tutorial is a good starting point.
On the other hand, if you're interested in doing any deeper customization, follow this tutorial. It shows how to do a lot of things manually, so you can learn how you can customize the workflow from data preprocessing to training, exporting and saving the model.
Setup
Install pip packages
Start by installing the TensorFlow Text and Model Garden pip packages.
tf-models-official is the TensorFlow Model Garden package. Note that it may not include the latest changes in the tensorflow_models GitHub repo. To include the latest changes, you may install tf-models-nightly, which is the nightly Model Garden package created daily automatically.
pip will install all models and dependencies automatically.
End of explanation
import os
import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
import tensorflow_models as tfm
import tensorflow_hub as hub
import tensorflow_datasets as tfds
tfds.disable_progress_bar()
Explanation: Import libraries
End of explanation
gs_folder_bert = "gs://cloud-tpu-checkpoints/bert/v3/uncased_L-12_H-768_A-12"
tf.io.gfile.listdir(gs_folder_bert)
Explanation: Resources
The following directory contains the BERT model's configuration, vocabulary, and a pre-trained checkpoint used in this tutorial:
End of explanation
batch_size=32
glue, info = tfds.load('glue/mrpc',
                       with_info=True,
                       batch_size=32)
glue
Explanation: Load and preprocess the dataset
This example uses the GLUE (General Language Understanding Evaluation) MRPC (Microsoft Research Paraphrase Corpus) dataset from TensorFlow Datasets (TFDS).
This dataset is not set up such that it can be directly fed into the BERT model. The following section handles the necessary preprocessing.
Get the dataset from TensorFlow Datasets
The GLUE MRPC (Dolan and Brockett, 2005) dataset is a corpus of sentence pairs automatically extracted from online news sources, with human annotations for whether the sentences in the pair are semantically equivalent. It has the following attributes:
Number of labels: 2
Size of training dataset: 3668
Size of evaluation dataset: 408
Maximum sequence length of training and evaluation dataset: 128
Begin by loading the MRPC dataset from TFDS:
End of explanation
info.features
Explanation: The info object describes the dataset and its features:
End of explanation
info.features['label'].names
Explanation: The two classes are:
End of explanation
example_batch = next(iter(glue['train']))
for key, value in example_batch.items():
  print(f"{key:9s}: {value[0].numpy()}")
Explanation: Here is one example from the training set:
End of explanation
tokenizer = tfm.nlp.layers.FastWordpieceBertTokenizer(
    vocab_file=os.path.join(gs_folder_bert, "vocab.txt"),
    lower_case=True)
Explanation: Preprocess the data
The keys "sentence1" and "sentence2" in the GLUE MRPC dataset contain two input sentences for each example.
Because the BERT model from the Model Garden doesn't take raw text as input, two things need to happen first:
The text needs to be tokenized (split into word pieces) and converted to indices.
Then, the indices need to be packed into the format that the model expects.
The BERT tokenizer
To fine tune a pre-trained language model from the Model Garden, such as BERT, you need to make sure that you're using exactly the same tokenization, vocabulary, and index mapping as used during training.
The following code rebuilds the tokenizer that was used by the base model using the Model Garden's tfm.nlp.layers.FastWordpieceBertTokenizer layer:
End of explanation
tokens = tokenizer(tf.constant(["Hello TensorFlow!"]))
tokens
Explanation: Let's tokenize a test sentence:
End of explanation
special = tokenizer.get_special_tokens_dict()
special
max_seq_length = 128
packer = tfm.nlp.layers.BertPackInputs(
    seq_length=max_seq_length,
    special_tokens_dict = tokenizer.get_special_tokens_dict())
Explanation: Learn more about the tokenization process in the Subword tokenization and Tokenizing with TensorFlow Text guides.
Pack the inputs
TensorFlow Model Garden's BERT model doesn't just take the tokenized strings as input. It also expects these to be packed into a particular format. tfm.nlp.layers.BertPackInputs layer can handle the conversion from a list of tokenized sentences to the input format expected by the Model Garden's BERT model.
tfm.nlp.layers.BertPackInputs packs the two input sentences (per example in the MRCP dataset) concatenated together. This input is expected to start with a [CLS] "This is a classification problem" token, and each sentence should end with a [SEP] "Separator" token.
Therefore, the tfm.nlp.layers.BertPackInputs layer's constructor takes the tokenizer's special tokens as an argument. It also needs to know the indices of the tokenizer's special tokens.
End of explanation
sentences1 = ["hello tensorflow"]
tok1 = tokenizer(sentences1)
tok1
sentences2 = ["goodbye tensorflow"]
tok2 = tokenizer(sentences2)
tok2
Explanation: The packer takes a list of tokenized sentences as input. For example:
End of explanation
packed = packer([tok1, tok2])
for key, tensor in packed.items():
  print(f"{key:15s}: {tensor[:, :12]}")
Explanation: Then, it returns a dictionary containing three outputs:
input_word_ids: The tokenized sentences packed together.
input_mask: The mask indicating which locations are valid in the other outputs.
input_type_ids: Indicating which sentence each token belongs to.
End of explanation
class BertInputProcessor(tf.keras.layers.Layer):
  def __init__(self, tokenizer, packer):
    super().__init__()
    self.tokenizer = tokenizer
    self.packer = packer
  def call(self, inputs):
    tok1 = self.tokenizer(inputs['sentence1'])
    tok2 = self.tokenizer(inputs['sentence2'])
    packed = self.packer([tok1, tok2])
    if 'label' in inputs:
      return packed, inputs['label']
    else:
      return packed
Explanation: Put it all together
Combine these two parts into a keras.layers.Layer that can be attached to your model:
End of explanation
bert_inputs_processor = BertInputProcessor(tokenizer, packer)
glue_train = glue['train'].map(bert_inputs_processor).prefetch(1)
Explanation: But for now just apply it to the dataset using Dataset.map, since the dataset you loaded from TFDS is a tf.data.Dataset object:
End of explanation
example_inputs, example_labels = next(iter(glue_train))
example_inputs
example_labels
for key, value in example_inputs.items():
  print(f'{key:15s} shape: {value.shape}')
print(f'{"labels":15s} shape: {example_labels.shape}')
Explanation: Here is an example batch from the processed dataset:
End of explanation
plt.pcolormesh(example_inputs['input_word_ids'])
Explanation: The input_word_ids contain the token IDs:
End of explanation
plt.pcolormesh(example_inputs['input_mask'])
Explanation: The mask allows the model to cleanly differentiate between the content and the padding. The mask has the same shape as the input_word_ids, and contains a 1 anywhere the input_word_ids is not padding.
End of explanation
plt.pcolormesh(example_inputs['input_type_ids'])
Explanation: The "input type" also has the same shape, but inside the non-padded region, contains a 0 or a 1 indicating which sentence the token is a part of.
End of explanation
glue_validation = glue['validation'].map(bert_inputs_processor).prefetch(1)
glue_test = glue['test'].map(bert_inputs_processor).prefetch(1)
Explanation: Apply the same preprocessing to the validation and test subsets of the GLUE MRPC dataset:
End of explanation
import json
bert_config_file = os.path.join(gs_folder_bert, "bert_config.json")
config_dict = json.loads(tf.io.gfile.GFile(bert_config_file).read())
config_dict
encoder_config = tfm.nlp.encoders.EncoderConfig({
    'type':'bert',
    'bert': config_dict
})
bert_encoder = tfm.nlp.encoders.build_encoder(encoder_config)
bert_encoder
Explanation: Build, train and export the model
Now that you have formatted the data as expected, you can start working on building and training the model.
Build the model
The first step is to download the configuration file—config_dict—for the pre-trained BERT model:
End of explanation
bert_classifier = tfm.nlp.models.BertClassifier(network=bert_encoder, num_classes=2)
Explanation: The configuration file defines the core BERT model from the Model Garden, which is a Keras model that predicts the outputs of num_classes from the inputs with maximum sequence length max_seq_length.
End of explanation
tf.keras.utils.plot_model(bert_classifier, show_shapes=True, dpi=48)
Explanation: The classifier has three inputs and one output:
End of explanation
bert_classifier(
    example_inputs, training=True).numpy()[:10]
Explanation: Run it on a test batch of data 10 examples from the training set. The output is the logits for the two classes:
End of explanation
tf.keras.utils.plot_model(bert_encoder, show_shapes=True, dpi=48)
Explanation: The TransformerEncoder in the center of the classifier above is the bert_encoder.
If you inspect the encoder, notice the stack of Transformer layers connected to those same three inputs:
End of explanation
checkpoint = tf.train.Checkpoint(encoder=bert_encoder)
checkpoint.read(
    os.path.join(gs_folder_bert, 'bert_model.ckpt')).assert_consumed()
Explanation: Restore the encoder weights
When built, the encoder is randomly initialized. Restore the encoder's weights from the checkpoint:
End of explanation
# Set up epochs and steps
epochs = 3
batch_size = 32
eval_batch_size = 32
train_data_size = info.splits['train'].num_examples
steps_per_epoch = int(train_data_size / batch_size)
num_train_steps = steps_per_epoch * epochs
warmup_steps = int(0.1 * num_train_steps)
initial_learning_rate=2e-5
Explanation: Note: The pretrained TransformerEncoder is also available on TensorFlow Hub. Go to the TF Hub appendix for details.
Set up the optimizer
BERT typically uses the Adam optimizer with weight decay—AdamW (tf.keras.optimizers.experimental.AdamW).
It also employs a learning rate schedule that first warms up from 0 and then decays to 0:
End of explanation
linear_decay = tf.keras.optimizers.schedules.PolynomialDecay(
    initial_learning_rate=initial_learning_rate,
    end_learning_rate=0,
    decay_steps=num_train_steps)
Explanation: Linear decay from initial_learning_rate to zero over num_train_steps.
End of explanation
warmup_schedule = tfm.optimization.lr_schedule.LinearWarmup(
    warmup_learning_rate = 0,
    after_warmup_lr_sched = linear_decay,
    warmup_steps = warmup_steps
)
Explanation: Warmup to that value over warmup_steps:
End of explanation
x = tf.linspace(0, num_train_steps, 1001)
y = [warmup_schedule(xi) for xi in x]
plt.plot(x,y)
plt.xlabel('Train step')
plt.ylabel('Learning rate')
Explanation: The overall schedule looks like this:
End of explanation
optimizer = tf.keras.optimizers.experimental.AdamW(
    lr = warmup_schedule,
    weight_decay=0.01,
)
Explanation: Use tf.keras.optimizers.experimental.AdamW to instantiate the optimizer with that schedule:
End of explanation
metrics = [tf.keras.metrics.SparseCategoricalAccuracy('accuracy', dtype=tf.float32)]
loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
bert_classifier.compile(
    optimizer=optimizer,
    loss=loss,
    metrics=metrics)
bert_classifier.fit(
      glue_train,
      validation_data=(glue_validation),
      batch_size=32,
      epochs=epochs)
Explanation: Train the model
Set the metric as accuracy and the loss as sparse categorical cross-entropy. Then, compile and train the BERT classifier:
End of explanation
my_examples = {
        'sentence1':[
            'The rain in Spain falls mainly on the plain.',
            'Look I fine tuned BERT.'],
        'sentence2':[
            'It mostly rains on the flat lands of Spain.',
            'Is it working? This does not match.']
    }
Explanation: Now run the fine-tuned model on a custom example to see that it works.
Start by encoding some sentence pairs:
End of explanation
ex_packed = bert_inputs_processor(my_examples)
my_logits = bert_classifier(ex_packed, training=False)
result_cls_ids = tf.argmax(my_logits)
result_cls_ids
tf.gather(tf.constant(info.features['label'].names), result_cls_ids)
Explanation: The model should report class 1 "match" for the first example and class 0 "no-match" for the second:
End of explanation
class ExportModel(tf.Module):
  def __init__(self, input_processor, classifier):
    self.input_processor = input_processor
    self.classifier = classifier
  @tf.function(input_signature=[{
      'sentence1': tf.TensorSpec(shape=[None], dtype=tf.string),
      'sentence2': tf.TensorSpec(shape=[None], dtype=tf.string)}])
  def __call__(self, inputs):
    packed = self.input_processor(inputs)
    logits =  self.classifier(packed, training=False)
    result_cls_ids = tf.argmax(logits)
    return {
        'logits': logits,
        'class_id': result_cls_ids,
        'class': tf.gather(
            tf.constant(info.features['label'].names),
            result_cls_ids)
    }
Explanation: Export the model
Often the goal of training a model is to use it for something outside of the Python process that created it. You can do this by exporting the model using tf.saved_model. (Learn more in the Using the SavedModel format guide and the Save and load a model using a distribution strategy tutorial.)
First, build a wrapper class to export the model. This wrapper does two things:
First it packages bert_inputs_processor and bert_classifier together into a single tf.Module, so you can export all the functionalities.
Second it defines a tf.function that implements the end-to-end execution of the model.
Setting the input_signature argument of tf.function lets you define a fixed signature for the tf.function. This can be less surprising than the default automatic retracing behavior.
End of explanation
export_model = ExportModel(bert_inputs_processor, bert_classifier)
import tempfile
export_dir=tempfile.mkdtemp(suffix='_saved_model')
tf.saved_model.save(export_model, export_dir=export_dir,
                    signatures={'serving_default': export_model.__call__})
Explanation: Create an instance of this export-model and save it:
End of explanation
original_logits = export_model(my_examples)['logits']
reloaded = tf.saved_model.load(export_dir)
reloaded_logits = reloaded(my_examples)['logits']
# The results are identical:
print(original_logits.numpy())
print()
print(reloaded_logits.numpy())
print(np.mean(abs(original_logits - reloaded_logits)))
Explanation: Reload the model and compare the results to the original:
End of explanation
import tensorflow_hub as hub
Explanation: Congratulations! You've used tensorflow_models to build a BERT-classifier, train it, and export for later use.
Optional: BERT on TF Hub
<a id="hub_bert"></a>
You can get the BERT model off the shelf from TF Hub. There are many versions available along with their input preprocessors.
This example uses a small version of BERT from TF Hub that was pre-trained using the English Wikipedia and BooksCorpus datasets, similar to the original implementation (Turc et al., 2019).
Start by importing TF Hub:
End of explanation
# Always make sure you use the right preprocessor.
hub_preprocessor = hub.KerasLayer(
    "https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3")
# This is a really small BERT.
hub_encoder = hub.KerasLayer(f"https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-2_H-128_A-2/2",
                             trainable=True)
print(f"The Hub encoder has {len(hub_encoder.trainable_variables)} trainable variables")
Explanation: Select the input preprocessor and the model from TF Hub and wrap them as hub.KerasLayer layers:
End of explanation
hub_inputs = hub_preprocessor(['Hello TensorFlow!'])
hub_inputs 
result = hub_encoder(
    inputs=hub_inputs,
    training=False,
)
print("Pooled output shape:", result['pooled_output'].shape)
print("Sequence output shape:", result['sequence_output'].shape)
Explanation: Test run the preprocessor on a batch of data:
End of explanation
hub_classifier = tfm.nlp.models.BertClassifier(
    bert_encoder,
    num_classes=2,
    dropout_rate=0.1,
    initializer=tf.keras.initializers.TruncatedNormal(
        stddev=0.02))
Explanation: At this point it would be simple to add a classification head yourself.
The Model Garden tfm.nlp.models.BertClassifier class can also build a classifier onto the TF Hub encoder:
End of explanation
tf.keras.utils.plot_model(hub_classifier, show_shapes=True, dpi=64)
Explanation: The one downside to loading this model from TF Hub is that the structure of internal Keras layers is not restored. This makes it more difficult to inspect or modify the model.
The BERT encoder model—hub_classifier—is now a single layer:
End of explanation
optimization_config = tfm.optimization.OptimizationConfig(
    optimizer=tfm.optimization.OptimizerConfig(
        type = "adamw",
        adamw =  tfm.optimization.AdamWeightDecayConfig(
            weight_decay_rate=0.01,
        )),
    learning_rate = tfm.optimization.LrConfig(
        type='polynomial',
        polynomial=tfm.optimization.PolynomialLrConfig(
            initial_learning_rate=2e-5,
            end_learning_rate=0.0,
            decay_steps=num_train_steps)),
    warmup = tfm.optimization.WarmupConfig(
        type='linear',
        linear=tfm.optimization.LinearWarmupConfig(warmup_steps=warmup_steps)
    ))
fac = tfm.optimization.optimizer_factory.OptimizerFactory(optimization_config)
lr = fac.build_learning_rate()
optimizer = fac.build_optimizer(lr=lr)
x = tf.linspace(0, num_train_steps, 1001).numpy()
y = [lr(xi) for xi in x]
plt.plot(x,y)
plt.xlabel('Train step')
plt.ylabel('Learning rate')
Explanation: For concrete examples of this approach, refer to Solve Glue tasks using BERT.
Optional: Optimizer configs
The tensorflow_models package defines serializable config classes that describe how to build the live objects. Earlier in this tutorial, you built the optimizer manually.
For example, the configuration below describes an (almost) identical optimizer built by the optimizer_factory.OptimizerFactory:
End of explanation
optimization_config = optimization_config.as_dict()
optimization_config
Explanation: The advantage to using config objects is that they don't contain any complicated TensorFlow objects, and can be easily serialized to JSON, and rebuilt. Here's the JSON for the above tfm.optimization.OptimizationConfig:
End of explanation
fac = tfm.optimization.optimizer_factory.OptimizerFactory(
    tfm.optimization.OptimizationConfig(optimization_config))
lr = fac.build_learning_rate()
optimizer = fac.build_optimizer(lr=lr)
Explanation: The tfm.optimization.optimizer_factory.OptimizerFactory can just as easily build the optimizer from the JSON dictionary:
End of explanation
<END_TASK>