code: def main(_):
      tf.logging.set_verbosity(tf.logging.INFO)
  
      bert_config = bert_lib.BertConfig.from_json_file(
          FLAGS.bert_config_file
      )
  
      tf.gfile.MakeDirs(FLAGS.output_dir)
  
      input_train_files = []
      for input_pattern in FLAGS.input_train_files.split(","):
          input_train_files.extend(tf.gfile.Glob(input_pattern))
  
      tf.logging.info("*** Input Training Files ***")
      for input_train_file in input_train_files:
          tf.logging.info("  %s" % input_train_file)
  
      input_eval_files = []
      for input_pattern in FLAGS.input_eval_files.split(","):
          input_eval_files.extend(tf.gfile.Glob(input_pattern))
  
      tf.logging.info("*** Input Evaluation Files ***")
      for input_eval_file in input_eval_files:
          tf.logging.info("  %s" % input_eval_file)
  
      tpu_cluster_resolver = None
      if FLAGS.use_tpu and FLAGS.tpu_name:
          tpu_cluster_resolver = (
              tf.distribute.cluster_resolver.TPUClusterResolver(
                  FLAGS.tpu_name,
                  zone=FLAGS.tpu_zone,
                  project=FLAGS.gcp_project,
              )
          )
  
      run_config = tf_estimator.tpu.RunConfig(
          cluster=tpu_cluster_resolver,
          master=FLAGS.master,
          model_dir=FLAGS.output_dir,
          save_checkpoints_steps=FLAGS.save_checkpoints_steps,
          tpu_config=tf_estimator.tpu.TPUConfig(
              iterations_per_loop=FLAGS.iterations_per_loop,
              num_cores_per_replica=1,
              per_host_input_for_training=tf_estimator.tpu.InputPipelineConfig.BROADCAST,
          ),
      )
  
      model_fn = model_fn_builder(
          bert_config=bert_config,
          init_checkpoint=FLAGS.init_checkpoint,
          learning_rate=FLAGS.learning_rate,
          num_train_steps=FLAGS.num_train_steps,
          num_warmup_steps=FLAGS.num_warmup_steps,
      )
  
      # If TPU is not available, this will fall back to normal Estimator on CPU
      # or GPU.
      estimator = tf_estimator.tpu.TPUEstimator(
          use_tpu=FLAGS.use_tpu,
          model_fn=model_fn,
          config=run_config,
          train_batch_size=FLAGS.train_batch_size,
          eval_batch_size=FLAGS.eval_batch_size,
      )
  
      if FLAGS.mode in ("train_and_eval", "train"):
          tf.logging.info(
              "Set train batch size = %d", FLAGS.train_batch_size
          )
          train_input_fn = input_fn_builder(
              input_files=input_train_files,
              max_seq_length=FLAGS.max_seq_length,
              max_predictions_per_seq=FLAGS.max_predictions_per_seq,
              is_training=True,
          )
  
      if FLAGS.mode in ("train_and_eval", "eval"):
          tf.logging.info(
              "Set eval batch size = %d", FLAGS.eval_batch_size
          )
          eval_input_fn = input_fn_builder(
              input_files=input_eval_files,
              max_seq_length=FLAGS.max_seq_length,
              max_predictions_per_seq=FLAGS.max_predictions_per_seq,
              is_training=False,
          )
  
          try:
              current_step = tf.train.load_variable(
                  FLAGS.output_dir, tf.GraphKeys.GLOBAL_STEP
              )
          except (TypeError, ValueError, tf.errors.NotFoundError):
              current_step = 0
  
          while current_step < FLAGS.num_train_steps:
              if FLAGS.mode == "train_and_eval":
                  # Train for up to steps_per_eval number of steps.
                  # At the end of training, a checkpoint will be written to --model_dir.
                  next_checkpoint = min(
                      current_step + FLAGS.steps_per_eval,
                      FLAGS.num_train_steps,
                  )
              elif FLAGS.mode == "train":
                  next_checkpoint = FLAGS.num_train_steps
  
              if FLAGS.mode in ("train_and_eval", "train"):
                  start_timestamp = (
                      time.time()
                  )  # This time will include compilation time
                  tf.logging.info("Starting to train.")
                  estimator.train(
                      input_fn=train_input_fn, max_steps=next_checkpoint
                  )
                  current_step = next_checkpoint
  
                  tf.logging.info(
                      "Finished training up to step %d. Elapsed seconds %d.",
                      current_step,
                      int(time.time() - start_timestamp),
                  )
  
              if FLAGS.mode in ("train_and_eval", "eval"):
                  tf.logging.info("Starting to evaluate.")
                  result = estimator.evaluate(
                      input_fn=eval_input_fn, steps=FLAGS.max_eval_steps
                  )
                  output_eval_file = os.path.join(
                      FLAGS.output_dir,
                      "eval_results_{}.txt".format(current_step),
                  )
                  with tf.gfile.GFile(output_eval_file, "w") as writer:
                      tf.logging.info("***** Eval results *****")
                      for key in sorted(result.keys()):
                          tf.logging.info(
                              "  %s = %s", key, str(result[key])
                          )
                          writer.write(
                              "%s = %s\n" % (key, str(result[key]))
                          )
                  if FLAGS.mode == "eval":
                      tf.logging.info("Exit eval mode")
                      break
  
file_path: mesh/mesh_tensorflow/bert/run_pretraining.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import os
  import time
  import mesh_tensorflow
  import mesh_tensorflow.bert.bert
  import mesh_tensorflow.bert.optimization
  from six.moves import range
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
label_desc: Use image node_id 4 for calling a global function with example usage: main(_) without return types