class_docstr: APIs to help create and train a BERT-based text classifier.
class_docstr_tok: ["APIs","to","help","create","and","train","a","BERT-based","text","classifier","."]
docstr: Loads an optimizer with a learning rate schedule.
  
  The decay steps in the learning rate schedule depend on the
  `steps_per_epoch` which may depend on the size of the training data.
  
  Args:
    train_data: Training data.
code: def _create_optimizer(self, train_data: text_ds.Dataset):
      """Loads an optimizer with a learning rate schedule.
  
      The decay steps in the learning rate schedule depend on the
      `steps_per_epoch` which may depend on the size of the training data.
  
      Args:
        train_data: Training data.
      """
      self._hparams.steps_per_epoch = model_util.get_steps_per_epoch(
          steps_per_epoch=self._hparams.steps_per_epoch,
          batch_size=self._hparams.batch_size,
          train_data=train_data,
      )
      total_steps = self._hparams.steps_per_epoch * self._hparams.epochs
      warmup_steps = int(total_steps * 0.1)
      initial_lr = self._hparams.learning_rate
      # Implements linear decay of the learning rate.
      lr_schedule = tf.keras.optimizers.schedules.PolynomialDecay(
          initial_learning_rate=initial_lr,
          decay_steps=total_steps,
          end_learning_rate=self._hparams.end_learning_rate,
          power=1.0,
      )
      if warmup_steps:
          lr_schedule = model_util.WarmUp(
              initial_learning_rate=initial_lr,
              decay_schedule_fn=lr_schedule,
              warmup_steps=warmup_steps,
          )
      if self._hparams.optimizer == hp.BertOptimizer.ADAMW:
          self._optimizer = tf.keras.optimizers.experimental.AdamW(
              lr_schedule,
              weight_decay=self._hparams.weight_decay,
              epsilon=1e-6,
              global_clipnorm=1.0,
          )
          self._optimizer.exclude_from_weight_decay(
              var_names=["LayerNorm", "layer_norm", "bias"]
          )
      elif self._hparams.optimizer == hp.BertOptimizer.LAMB:
          self._optimizer = tfa_optimizers.LAMB(
              lr_schedule,
              weight_decay_rate=self._hparams.weight_decay,
              epsilon=1e-6,
              exclude_from_weight_decay=[
                  "LayerNorm",
                  "layer_norm",
                  "bias",
              ],
              global_clipnorm=1.0,
          )
      else:
          raise ValueError(
              "BertHParams.optimizer must be set to ADAM or "
              f"LAMB. Got {self._hparams.optimizer}."
          )
  
imports: import abc
  import os
  import tempfile
  from typing import Any, Optional, Sequence, Tuple
  import tensorflow
  from tensorflow_addons import optimizers
  import tensorflow_hub
  from mediapipe.model_maker.python.core.data import dataset
  from mediapipe.model_maker.python.core.tasks import classifier
  from mediapipe.model_maker.python.core.utils import hub_loader
  from mediapipe.model_maker.python.core.utils import loss_functions
  from mediapipe.model_maker.python.core.utils import metrics
  from mediapipe.model_maker.python.core.utils import model_util
  from mediapipe.model_maker.python.core.utils import quantization
  from mediapipe.model_maker.python.text.text_classifier import bert_tokenizer
  from mediapipe.model_maker.python.text.text_classifier import dataset
  from mediapipe.model_maker.python.text.text_classifier import hyperparameters
  from mediapipe.model_maker.python.text.text_classifier import model_options
  from mediapipe.model_maker.python.text.text_classifier import model_spec
  from mediapipe.model_maker.python.text.text_classifier import model_with_tokenizer
  from mediapipe.model_maker.python.text.text_classifier import preprocessor
  from mediapipe.model_maker.python.text.text_classifier import text_classifier_options
  from mediapipe.tasks.python.metadata.metadata_writers import metadata_writer
  from mediapipe.tasks.python.metadata.metadata_writers import text_classifier
label_desc: Use image node_id 7 for calling the _BertClassifier obj's underlying member method code with example usage: obj._create_optimizer(train_data) without return types