code: def main(_):
      keras.mixed_precision.set_dtype_policy(
          FLAGS.mixed_precision_policy
      )
  
      logging.info(
          "Benchmarking configs...\n"
          "=========================\n"
          f"MODEL: BERT {FLAGS.model_size}\n"
          f"TASK: glue/mrpc \n"
          f"BATCH_SIZE: {FLAGS.batch_size}\n"
          f"EPOCHS: {FLAGS.epochs}\n"
          "=========================\n"
      )
  
      # Load datasets.
      train_ds, test_ds, validation_ds = load_data()
  
      # Load the model.
      model = load_model()
      # Set loss and metrics.
      loss = keras.losses.SparseCategoricalCrossentropy(
          from_logits=True
      )
      metrics = [keras.metrics.SparseCategoricalAccuracy()]
      # Configure optimizer.
      lr = keras.optimizers.schedules.PolynomialDecay(
          5e-4,
          decay_steps=train_ds.cardinality() * FLAGS.epochs,
          end_learning_rate=0.0,
      )
      optimizer = keras.optimizers.AdamW(lr, weight_decay=0.01)
      optimizer.exclude_from_weight_decay(
          var_names=["LayerNorm", "layer_norm", "bias"]
      )
  
      model.compile(optimizer=optimizer, loss=loss, metrics=metrics)
  
      benchmark_metrics_callback = BenchmarkMetricsCallback(
          start_batch=1,
          stop_batch=train_ds.cardinality().numpy() - 1,
      )
  
      # Start training.
      logging.info("Starting Training...")
  
      st = time.time()
      history = model.fit(
          train_ds,
          validation_data=validation_ds,
          epochs=FLAGS.epochs,
          callbacks=[benchmark_metrics_callback],
      )
  
      wall_time = time.time() - st
      validation_accuracy = history.history[
          "val_sparse_categorical_accuracy"
      ][-1]
      examples_per_second = (
          np.mean(
              np.array(benchmark_metrics_callback.state["throughput"])
          )
          * FLAGS.batch_size
      )
  
      logging.info("Training Finished!")
      logging.info(f"Wall Time: {wall_time:.4f} seconds.")
      logging.info(f"Validation Accuracy: {validation_accuracy:.4f}")
      logging.info(f"examples_per_second: {examples_per_second:.4f}")
  
file_path: keras/benchmarks/model_benchmark/bert_benchmark.py
imports: import time
  import keras_nlp
  import numpy
  import tensorflow
  import tensorflow_datasets
  from absl import app
  from absl import flags
  from absl import logging
  from model_benchmark.benchmark_utils import BenchmarkMetricsCallback
  import keras
label_desc: Use image node_id 3 for calling a global function with example usage: main(_) without return types