code: def main(_):
      if FLAGS.module_import:
          for module in FLAGS.module_import:
              importlib.import_module(module)
  
      if FLAGS.t5_tfds_data_dir:
          t5.data.set_tfds_data_dir_override(FLAGS.t5_tfds_data_dir)
  
      # Add search path for gin files stored in package.
      assert pkg_resources.resource_filename is not None
      gin.add_config_file_search_path(
          pkg_resources.resource_filename(__name__, "gin")
      )
      try:
          suffix = 0
          command_dir = os.path.join(FLAGS.model_dir, "commands")
          tf.io.gfile.makedirs(command_dir)
          command_filename = os.path.join(command_dir, "command")
          while tf.io.gfile.exists(command_filename):
              suffix += 1
              command_filename = os.path.join(
                  command_dir, "command.{}".format(suffix)
              )
          with tf.io.gfile.GFile(command_filename, "w") as f:
              f.write(" ".join(sys.argv))
      except (
          tf.errors.PermissionDeniedError,
          tf.errors.InvalidArgumentError,
      ):
          logging.info(
              "No write access to model directory. Skipping command logging."
          )
  
      utils.parse_gin_defaults_and_flags(
          skip_unknown=(
              FLAGS.skip_all_gin_unknowns
              or (
                  mesh_transformer.DEPRECATED_GIN_REFERENCES
                  + tuple(FLAGS.additional_deprecated_gin_references)
              )
          ),
          finalize_config=False,
      )
      # We must overide this binding explicitly since it is set to a deprecated
      # function or class in many existing configs.
      gin.bind_parameter(
          "run.vocabulary", mesh_transformer.get_vocabulary()
      )
      gin.finalize()
  
      # Set cache dir after loading gin to avoid unintentionally overriding it.
      t5.data.add_global_cache_dirs(FLAGS.additional_task_cache_dirs)
  
      if FLAGS.use_model_api:
          model = mtf_model.MtfModel(
              tpu_job_name=FLAGS.tpu_job_name,
              tpu=FLAGS.tpu,
              gcp_project=FLAGS.gcp_project,
              tpu_zone=FLAGS.tpu_zone,
              tpu_topology=FLAGS.tpu_topology,
              model_parallelism=FLAGS.model_parallelism,
              model_dir=FLAGS.model_dir,
              batch_size=FLAGS.batch_size,
              sequence_length={
                  "inputs": FLAGS.input_sequence_length,
                  "targets": FLAGS.target_sequence_length,
              },
          )
  
          if (
              FLAGS.checkpoint_mode != "specific"
              and FLAGS.checkpoint_steps
          ):
              raise ValueError(
                  "checkpoint_mode is set to %s and checkpoint_steps is "
                  "also set. To use a particular checkpoint, please set "
                  "checkpoint_mode to 'specific'. For other modes, please "
                  "ensure that checkpoint_steps is not set."
                  % FLAGS.checkpoint_mode
              )
  
          if FLAGS.checkpoint_mode == "latest":
              checkpoint_steps = -1
          elif FLAGS.checkpoint_mode == "all":
              checkpoint_steps = "all"
          else:
              checkpoint_steps = [
                  int(c) for c in FLAGS.checkpoint_steps
              ]
  
          if FLAGS.mode == "train":
              model.train(
                  mixture_or_task_name=FLAGS.mixture_or_task,
                  steps=FLAGS.train_steps,
              )
          elif FLAGS.mode == "eval":
              model.eval(
                  mixture_or_task_name=FLAGS.mixture_or_task,
                  checkpoint_steps=checkpoint_steps,
                  summary_dir=FLAGS.eval_summary_dir,
                  split=FLAGS.eval_split,
              )
          elif FLAGS.mode == "finetune":
              if not (
                  FLAGS.checkpoint_mode == "latest"
                  or (
                      FLAGS.checkpoint_mode == "specific"
                      and len(FLAGS.checkpoint_steps) == 1
                  )
              ):
                  raise ValueError(
                      "Must specify a single checkpoint for finetuning a model."
                  )
  
              if isinstance(checkpoint_steps, list):
                  checkpoint_steps = checkpoint_steps[0]
  
              model.finetune(
                  mixture_or_task_name=FLAGS.mixture_or_task,
                  steps=FLAGS.train_steps,
                  pretrained_model_dir=FLAGS.pretrained_model_dir,
                  checkpoint_steps=checkpoint_steps,
              )
          elif FLAGS.mode == "predict":
              model.predict(
                  checkpoint_steps=checkpoint_steps,
                  input_file=FLAGS.input_file,
                  output_file=FLAGS.output_file,
                  beam_size=FLAGS.beam_size,
                  temperature=FLAGS.temperature,
                  keep_top_k=FLAGS.keep_top_k,
              )
          elif FLAGS.mode == "score":
              model.score(
                  FLAGS.input_file,
                  FLAGS.target_file,
                  scores_file=FLAGS.output_file,
                  checkpoint_steps=checkpoint_steps,
              )
          elif FLAGS.mode in ("export_predict", "export_score"):
              if not (
                  FLAGS.checkpoint_mode == "latest"
                  or (
                      FLAGS.checkpoint_mode == "specific"
                      and len(FLAGS.checkpoint_steps) == 1
                  )
              ):
                  raise ValueError(
                      "Must specify a single checkpoint for exporting a model."
                  )
  
              if isinstance(checkpoint_steps, list):
                  checkpoint_steps = checkpoint_steps[0]
  
              model.export(
                  export_dir=FLAGS.export_dir,
                  checkpoint_step=checkpoint_steps,
                  beam_size=FLAGS.beam_size,
                  temperature=FLAGS.temperature,
                  keep_top_k=FLAGS.keep_top_k,
                  eval_with_score=(FLAGS.mode == "export_score"),
              )
          else:
              raise ValueError(
                  "--mode flag must be set when using Model API."
              )
      else:
          if FLAGS.mode:
              raise ValueError(
                  "--mode flag should only be set when using Model API."
              )
          if not FLAGS.tpu:
              with gin.unlock_config():
                  gin.bind_parameter(
                      "utils.get_variable_dtype.slice_dtype", "float32"
                  )
                  gin.bind_parameter(
                      "utils.get_variable_dtype.activation_dtype",
                      "float32",
                  )
          utils.run(
              tpu_job_name=FLAGS.tpu_job_name,
              tpu=FLAGS.tpu,
              gcp_project=FLAGS.gcp_project,
              tpu_zone=FLAGS.tpu_zone,
              model_dir=FLAGS.model_dir,
          )
  
file_path: text-to-text-transfer-transformer/t5/models/mesh_transformer_main.py
imports: import importlib
  import os
  import sys
  from absl import app
  from absl import flags
  from absl import logging
  import gin
  from mesh_tensorflow.transformer import utils
  import pkg_resources
  import t5
  from t5.models import mesh_transformer
  from t5.models import mtf_model
  import tensorflow.compat.v1
label_desc: Use image node_id 1 for calling a global function with example usage: main(_) without return types