code: def bert_metrics_function(
      outputs: Dict[str, jnp.ndarray], batch: Batch
  ) -> Dict[str, Tuple[float, int]]:
      """Calcualte metrics for the BERT task.
  
      Args:
       outputs: Output of model that has masked LM logits of shape [batch, length,
         vocab_size], and  next sentence prediction logits of shape [batch, 2].
       batch: Batch of data that has 'masked_lm_ids', 'masked_lm_weights' and
         'next_sentence_labels'.
  
      Returns:
        A dict of metrics, in which keys are metrics name and values are tuples of
        (metric, normalizer).
      """
      mlm_logits = outputs["mlm_logits"]
      nsp_logits = outputs["nsp_logits"]
      next_sentence_labels = common_utils.onehot(
          batch["next_sentence_labels"], 2
      )
      batch_weights = batch.get(
          "batch_mask"
      )  # batch_mask might not be defined
      per_ex_nsp_loss = (
          model_utils.weighted_unnormalized_softmax_cross_entropy(
              nsp_logits, next_sentence_labels, batch_weights
          )
      )
      per_ex_nsp_accuracy = model_utils.weighted_correctly_classified(
          nsp_logits, next_sentence_labels, batch_weights
      )
  
      per_ex_mlm_loss = (
          sparse_weighted_unnormalized_softmax_cross_entropy(
              mlm_logits,
              batch["masked_lm_ids"],
              batch["masked_lm_weights"],
              batch_weights,
          )
      )
      per_ex_mlm_accuracy = sparse_weighted_per_example_accuracy(
          mlm_logits,
          batch["masked_lm_ids"],
          batch["masked_lm_weights"],
          batch_weights,
      )
  
      # This psum is required to correctly evaluate with multihost. Only host 0
      # will report the metrics, so we must aggregate across all hosts. The psum
      # will map an array of shape [n_global_devices, batch_size] -> [batch_size]
      # by summing across the devices dimension. The outer sum then sums across the
      # batch dim. The result is then we have summed across all samples in the
      # sharded batch.
      evaluated_metrics = {}
      normalizer = num_examples(mlm_logits, batch_weights)
      for name, value in zip(
          [
              "nsp_loss",
              "nsp_accuracy",
              "mlm_loss",
              "mlm_accuracy",
              "loss",
          ],
          [
              per_ex_nsp_loss,
              per_ex_nsp_accuracy,
              per_ex_mlm_loss,
              per_ex_mlm_accuracy,
              per_ex_nsp_loss + per_ex_mlm_loss,
          ],
      ):
          evaluated_metrics[name] = model_utils.psum_metric_normalizer(
              (value, normalizer)
          )
      return evaluated_metrics
  
file_path: scenic/scenic/projects/baselines/bert/bert_base_model.py
imports: from typing import Callable, Dict, Optional, Tuple, Union
  from flax.training import common_utils
  import jax
  import jax.numpy
  import numpy
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import model_utils
label_desc: Use image node_id 5 for calling a global function with example usage: bert_metrics_function(outputs, batch) and returns: evaluated_metrics