code: def model_fn_builder(
      bert_config,
      init_checkpoint,
      learning_rate,
      num_train_steps,
      num_warmup_steps,
      use_one_hot_embeddings,
      optimizer,
      poly_power,
      start_warmup_step,
      steps_per_update,
      num_workers,
  ):
      """Returns `model_fn` closure for TPUEstimator."""
  
      def model_fn(
          features, labels, mode, params
      ):  # pylint: disable=unused-argument
          """The `model_fn` for TPUEstimator."""
          tf.logging.info("*** Features ***")
          for name in sorted(features.keys()):
              tf.logging.info(
                  "  name = %s, shape = %s"
                  % (name, features[name].shape)
              )
  
          input_ids = features["input_ids"]
          input_mask = features["input_mask"]
          segment_ids = features["segment_ids"]
          masked_lm_positions = features["masked_lm_positions"]
          masked_lm_ids = features["masked_lm_ids"]
          masked_lm_weights = features["masked_lm_weights"]
          next_sentence_labels = features["next_sentence_labels"]
  
          is_training = mode == tf.estimator.ModeKeys.TRAIN
          tf.logging.info(f"********mode in model_fn*******: {mode}")
  
          model = modeling.BertModel(
              config=bert_config,
              is_training=is_training,
              input_ids=input_ids,
              input_mask=input_mask,
              token_type_ids=segment_ids,
              use_one_hot_embeddings=use_one_hot_embeddings,
          )
  
          global_batch_size = params["batch_size"]
          tf.logging.info(
              f"*** Global batch size in model_fn is {global_batch_size}"
          )
          (
              masked_lm_loss,
              masked_lm_example_loss,
              masked_lm_log_probs,
          ) = get_masked_lm_output(
              bert_config,
              model.get_sequence_output(),
              model.get_embedding_table(),
              masked_lm_positions,
              masked_lm_ids,
              masked_lm_weights,
              num_workers,
          )
  
          (
              next_sentence_loss,
              next_sentence_example_loss,
              next_sentence_log_probs,
          ) = get_next_sentence_output(
              bert_config,
              model.get_pooled_output(),
              next_sentence_labels,
              global_batch_size,
          )
  
          total_loss = masked_lm_loss + next_sentence_loss
  
          tvars = tf.trainable_variables()
  
          initialized_variable_names = {}
          scaffold_fn = None
          if init_checkpoint:
              (
                  assignment_map,
                  initialized_variable_names,
              ) = modeling.get_assignment_map_from_checkpoint(
                  tvars, init_checkpoint
              )
  
              tvar_index = {
                  var.name.replace(":0", ""): var for var in tvars
              }
              assignment_map = collections.OrderedDict(
                  [
                      (name, tvar_index.get(name, value))
                      for name, value in assignment_map.items()
                  ]
              )
  
              tf.train.init_from_checkpoint(
                  init_checkpoint, assignment_map
              )
  
          tf.logging.info("**** Trainable Variables ****")
          for var in tvars:
              init_string = ""
              if var.name in initialized_variable_names:
                  init_string = ", *INIT_FROM_CKPT*"
              tf.logging.info(
                  "  name = %s, shape = %s%s",
                  var.name,
                  var.shape,
                  init_string,
              )
  
          output_spec = None
          if mode == tf.estimator.ModeKeys.TRAIN:
              train_op = optimization.create_optimizer(
                  total_loss,
                  learning_rate,
                  num_train_steps,
                  num_warmup_steps,
                  False,
                  optimizer,
                  poly_power,
                  start_warmup_step,
                  steps_per_update,
              )
  
              output_spec = tf.estimator.EstimatorSpec(
                  mode=mode, loss=total_loss, train_op=train_op
              )
          elif mode == tf.estimator.ModeKeys.EVAL:
  
              def metric_fn(
                  masked_lm_example_loss,
                  masked_lm_log_probs,
                  masked_lm_ids,
                  masked_lm_weights,
                  next_sentence_example_loss,
                  next_sentence_log_probs,
                  next_sentence_labels,
              ):
                  """Computes the loss and accuracy of the model."""
                  masked_lm_log_probs = tf.reshape(
                      masked_lm_log_probs,
                      [-1, masked_lm_log_probs.shape[-1]],
                  )
                  masked_lm_predictions = tf.argmax(
                      masked_lm_log_probs, axis=-1, output_type=tf.int32
                  )
                  masked_lm_example_loss = tf.reshape(
                      masked_lm_example_loss, [-1]
                  )
                  masked_lm_ids = tf.reshape(masked_lm_ids, [-1])
                  masked_lm_weights = tf.reshape(
                      masked_lm_weights, [-1]
                  )
                  masked_lm_accuracy = tf.metrics.accuracy(
                      labels=masked_lm_ids,
                      predictions=masked_lm_predictions,
                      weights=masked_lm_weights,
                  )
                  masked_lm_mean_loss = tf.metrics.mean(
                      values=masked_lm_example_loss,
                      weights=masked_lm_weights,
                  )
  
                  next_sentence_log_probs = tf.reshape(
                      next_sentence_log_probs,
                      [-1, next_sentence_log_probs.shape[-1]],
                  )
                  next_sentence_predictions = tf.argmax(
                      next_sentence_log_probs,
                      axis=-1,
                      output_type=tf.int32,
                  )
                  next_sentence_labels = tf.reshape(
                      next_sentence_labels, [-1]
                  )
                  next_sentence_accuracy = tf.metrics.accuracy(
                      labels=next_sentence_labels,
                      predictions=next_sentence_predictions,
                  )
                  next_sentence_mean_loss = tf.metrics.mean(
                      values=next_sentence_example_loss
                  )
  
                  return {
                      "masked_lm_accuracy": masked_lm_accuracy,
                      "masked_lm_loss": masked_lm_mean_loss,
                      "next_sentence_accuracy": next_sentence_accuracy,
                      "next_sentence_loss": next_sentence_mean_loss,
                  }
  
              eval_metrics = (
                  metric_fn,
                  [
                      masked_lm_example_loss,
                      masked_lm_log_probs,
                      masked_lm_ids,
                      masked_lm_weights,
                      next_sentence_example_loss,
                      next_sentence_log_probs,
                      next_sentence_labels,
                  ],
              )
              output_spec = tf.estimator.EstimatorSpec(
                  mode=mode,
                  loss=total_loss,
                  eval_metric_ops=metric_fn(
                      masked_lm_example_loss,
                      masked_lm_log_probs,
                      masked_lm_ids,
                      masked_lm_weights,
                      next_sentence_example_loss,
                      next_sentence_log_probs,
                      next_sentence_labels,
                  ),
              )
  
          else:
              raise ValueError(
                  "Only TRAIN and EVAL modes are supported: %s" % (mode)
              )
  
          return output_spec
  
      return model_fn
  
file_path: BigDL/python/orca/example/learn/tf/bert/run_pretraining.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  import os
  import absl
  import sys
  import modeling
  import optimization
  import tensorflow.compat.v1
  from bigdl.orca.learn.tf.tf_estimator import TFEstimator
label_desc: Use image node_id 1 for calling a global function with example usage: model_fn_builder(bert_config, init_checkpoint, learning_rate, num_train_steps, num_warmup_steps, use_one_hot_embeddings, optimizer, poly_power, start_warmup_step, steps_per_update, num_workers) and returns: model_fn, output_spec, dict