code: def main(_):
      tf.logging.set_verbosity(tf.logging.INFO)
  
      processors = {
          "cola": ColaProcessor,
          "mnli": MnliProcessor,
          "mrpc": MrpcProcessor,
          "xnli": XnliProcessor,
      }
  
      tokenization.validate_case_matches_checkpoint(
          FLAGS.do_lower_case, FLAGS.init_checkpoint
      )
  
      if (
          not FLAGS.do_train
          and not FLAGS.do_eval
          and not FLAGS.do_predict
      ):
          raise ValueError(
              "At least one of `do_train`, `do_eval` or `do_predict' must be True."
          )
  
      bert_config = bert_lib.BertConfig.from_json_file(
          FLAGS.bert_config_file
      )
  
      if FLAGS.max_seq_length > bert_config.max_position_embeddings:
          raise ValueError(
              "Cannot use sequence length %d because the BERT model "
              "was only trained up to sequence length %d"
              % (
                  FLAGS.max_seq_length,
                  bert_config.max_position_embeddings,
              )
          )
  
      tf.gfile.MakeDirs(FLAGS.output_dir)
  
      task_name = FLAGS.task_name.lower()
  
      if task_name not in processors:
          raise ValueError("Task not found: %s" % (task_name))
  
      processor = processors[task_name]()
  
      label_list = processor.get_labels()
  
      tokenizer = tokenization.FullTokenizer(
          vocab_file=FLAGS.vocab_file, do_lower_case=FLAGS.do_lower_case
      )
  
      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,
          ),
      )
  
      train_examples = None
      num_train_steps = None
      num_warmup_steps = None
      if FLAGS.do_train:
          train_examples = processor.get_train_examples(FLAGS.data_dir)
          num_train_steps = int(
              len(train_examples)
              / FLAGS.train_batch_size
              * FLAGS.num_train_epochs
          )
          num_warmup_steps = int(
              num_train_steps * FLAGS.warmup_proportion
          )
  
      model_fn = model_fn_builder(
          bert_config=bert_config,
          num_labels=len(label_list),
          init_checkpoint=FLAGS.init_checkpoint,
          learning_rate=FLAGS.learning_rate,
          num_train_steps=num_train_steps,
          num_warmup_steps=num_warmup_steps,
          use_tpu=FLAGS.use_tpu,
      )
  
      # 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,
          predict_batch_size=FLAGS.predict_batch_size,
      )
  
      if FLAGS.do_train:
          if FLAGS.cached_train_file:
              train_file = FLAGS.cached_train_file
          else:
              train_file = os.path.join(
                  FLAGS.output_dir, "train.tf_record"
              )
              file_based_convert_examples_to_features(
                  train_examples,
                  label_list,
                  FLAGS.max_seq_length,
                  tokenizer,
                  train_file,
              )
              tf.logging.info("***** Running training *****")
              tf.logging.info(
                  "  Num examples = %d", len(train_examples)
              )
              tf.logging.info(
                  "  Batch size = %d", FLAGS.train_batch_size
              )
              tf.logging.info("  Num steps = %d", num_train_steps)
          train_input_fn = file_based_input_fn_builder(
              input_file=train_file,
              seq_length=FLAGS.max_seq_length,
              is_training=True,
              drop_remainder=True,
          )
          estimator.train(
              input_fn=train_input_fn, max_steps=num_train_steps
          )
  
      if FLAGS.do_eval:
          eval_examples = processor.get_dev_examples(FLAGS.data_dir)
          num_actual_eval_examples = len(eval_examples)
          if FLAGS.use_tpu:
              # TPU requires a fixed batch size for all batches, therefore the number
              # of examples must be a multiple of the batch size, or else examples
              # will get dropped. So we pad with fake examples which are ignored
              # later on. These do NOT count towards the metric (all tf.metrics
              # support a per-instance weight, and these get a weight of 0.0).
              while len(eval_examples) % FLAGS.eval_batch_size != 0:
                  eval_examples.append(PaddingInputExample())
  
          eval_file = os.path.join(FLAGS.output_dir, "eval.tf_record")
          file_based_convert_examples_to_features(
              eval_examples,
              label_list,
              FLAGS.max_seq_length,
              tokenizer,
              eval_file,
          )
  
          tf.logging.info("***** Running evaluation *****")
          tf.logging.info(
              "  Num examples = %d (%d actual, %d padding)",
              len(eval_examples),
              num_actual_eval_examples,
              len(eval_examples) - num_actual_eval_examples,
          )
          tf.logging.info("  Batch size = %d", FLAGS.eval_batch_size)
  
          # This tells the estimator to run through the entire set.
          eval_steps = None
          # However, if running eval on the TPU, you will need to specify the
          # number of steps.
          if FLAGS.use_tpu:
              assert len(eval_examples) % FLAGS.eval_batch_size == 0
              eval_steps = int(
                  len(eval_examples) // FLAGS.eval_batch_size
              )
  
          eval_drop_remainder = True if FLAGS.use_tpu else False
          eval_input_fn = file_based_input_fn_builder(
              input_file=eval_file,
              seq_length=FLAGS.max_seq_length,
              is_training=False,
              drop_remainder=eval_drop_remainder,
          )
  
          result = estimator.evaluate(
              input_fn=eval_input_fn, steps=eval_steps
          )
  
          output_eval_file = os.path.join(
              FLAGS.output_dir, "eval_results.txt"
          )
          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.do_predict:
          predict_examples = processor.get_test_examples(FLAGS.data_dir)
          num_actual_predict_examples = len(predict_examples)
          if FLAGS.use_tpu:
              # TPU requires a fixed batch size for all batches, therefore the number
              # of examples must be a multiple of the batch size, or else examples
              # will get dropped. So we pad with fake examples which are ignored
              # later on.
              while (
                  len(predict_examples) % FLAGS.predict_batch_size != 0
              ):
                  predict_examples.append(PaddingInputExample())
  
          predict_file = os.path.join(
              FLAGS.output_dir, "predict.tf_record"
          )
          file_based_convert_examples_to_features(
              predict_examples,
              label_list,
              FLAGS.max_seq_length,
              tokenizer,
              predict_file,
          )
  
          tf.logging.info("***** Running prediction*****")
          tf.logging.info(
              "  Num examples = %d (%d actual, %d padding)",
              len(predict_examples),
              num_actual_predict_examples,
              len(predict_examples) - num_actual_predict_examples,
          )
          tf.logging.info("  Batch size = %d", FLAGS.predict_batch_size)
  
          predict_drop_remainder = True if FLAGS.use_tpu else False
          predict_input_fn = file_based_input_fn_builder(
              input_file=predict_file,
              seq_length=FLAGS.max_seq_length,
              is_training=False,
              drop_remainder=predict_drop_remainder,
          )
  
          result = estimator.predict(input_fn=predict_input_fn)
  
          output_predict_file = os.path.join(
              FLAGS.output_dir, "test_results.tsv"
          )
          with tf.gfile.GFile(output_predict_file, "w") as writer:
              num_written_lines = 0
              tf.logging.info("***** Predict results *****")
              for i, prediction in enumerate(result):
                  probabilities = prediction["probabilities"]
                  if i >= num_actual_predict_examples:
                      break
                  output_line = (
                      "\t".join(
                          str(class_probability)
                          for class_probability in probabilities
                      )
                      + "\n"
                  )
                  writer.write(output_line)
                  num_written_lines += 1
          assert num_written_lines == num_actual_predict_examples
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  import csv
  import os
  import mesh_tensorflow
  import mesh_tensorflow.bert.bert
  import mesh_tensorflow.bert.optimization
  import mesh_tensorflow.bert.tokenization
  from six.moves import range
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
label_desc: Use image node_id 9 for calling a global function with example usage: main(_) without return types