Run training/eval/inference.

  Args:
    tpu_job_name: string, name of TPU worker binary
    tpu: string, the Cloud TPU to use for training
    gcp_project: string, project name for the Cloud TPU-enabled project
    tpu_zone: string, GCE zone where the Cloud TPU is located in
    model_dir: string, estimator model_dir
    model_type: a string - either "bitransformer", "lm" or "aligned"
    vocabulary: a vocabulary.Vocabulary or
      (inputs_vocabulary, targets_vocabulary) tuple.
    train_dataset_fn: A function returning a tf.data.Dataset. Must be provided
      for mode=train
    eval_dataset_fn: A function returning a tf.data.Dataset. Must be provided
      for model=eval
    dataset_split: a string
    autostack: boolean, internally combine variables
    checkpoint_path: a string - which checkpoint to load for inference
    mode: string, train/evaluate/infer
    iterations_per_loop: integer, steps per train loop
    save_checkpoints_steps: integer, steps per checkpoint
    eval_steps: integer, number of evaluation steps
    train_steps: Total number of training steps.
    batch_size: An integer or a function with the same signature as
      auto_batch_size().  Mini-batch size for the training. Note that this is
      the global batch size and not the per-shard batch size.
    sequence_length: an integer
    mesh_shape: an input to mtf.convert_to_shape()
    layout_rules: an input to mtf.convert_to_layout_rules()
    get_components_fn: an optional function that gets a list of tuples of
      (metric_names, component) for each component. Required if mode is
      "continuous_eval"

def run(tpu_job_name,
        tpu,
        gcp_project,
        tpu_zone,
        model_dir,
        model_type="bitransformer",
        vocabulary=gin.REQUIRED,
        train_dataset_fn=None,
        eval_dataset_fn=None,
        dataset_split="train",
        autostack=True,
        checkpoint_path="",
        mode="train",
        iterations_per_loop=100,
        save_checkpoints_steps=1000,
        eval_steps=10,
        train_steps=1000000,
        batch_size=auto_batch_size,
        sequence_length=gin.REQUIRED,
        mesh_shape=gin.REQUIRED,
        layout_rules=gin.REQUIRED,
        get_components_fn=None):
  """Run training/eval/inference.

  Args:
    tpu_job_name: string, name of TPU worker binary
    tpu: string, the Cloud TPU to use for training
    gcp_project: string, project name for the Cloud TPU-enabled project
    tpu_zone: string, GCE zone where the Cloud TPU is located in
    model_dir: string, estimator model_dir
    model_type: a string - either "bitransformer", "lm" or "aligned"
    vocabulary: a vocabulary.Vocabulary or
      (inputs_vocabulary, targets_vocabulary) tuple.
    train_dataset_fn: A function returning a tf.data.Dataset. Must be provided
      for mode=train
    eval_dataset_fn: A function returning a tf.data.Dataset. Must be provided
      for model=eval
    dataset_split: a string
    autostack: boolean, internally combine variables
    checkpoint_path: a string - which checkpoint to load for inference
    mode: string, train/evaluate/infer
    iterations_per_loop: integer, steps per train loop
    save_checkpoints_steps: integer, steps per checkpoint
    eval_steps: integer, number of evaluation steps
    train_steps: Total number of training steps.
    batch_size: An integer or a function with the same signature as
      auto_batch_size().  Mini-batch size for the training. Note that this is
      the global batch size and not the per-shard batch size.
    sequence_length: an integer
    mesh_shape: an input to mtf.convert_to_shape()
    layout_rules: an input to mtf.convert_to_layout_rules()
    get_components_fn: an optional function that gets a list of tuples of
      (metric_names, component) for each component. Required if mode is
      "continuous_eval"
  """
  if not isinstance(batch_size, int):
    batch_size = batch_size(sequence_length, mesh_shape, layout_rules)

  tf.logging.info("mode=%s" % mode,)
  tf.logging.info("batch_size=%s" % batch_size,)
  tf.logging.info("sequence_length=%s" % sequence_length,)
  tf.logging.info("mesh_shape=%s" % mesh_shape,)
  tf.logging.info("layout_rules=%s" % layout_rules,)

  if mode == "train" and dataset_split != "train":
    raise ValueError("mode==\"train\" requires dataset_split==\"train\"")

  mesh_shape = mtf.convert_to_shape(mesh_shape)
  layout_rules = mtf.convert_to_layout_rules(layout_rules)

  cluster = tf.contrib.cluster_resolver.TPUClusterResolver(
      tpu if (tpu) else "", zone=tpu_zone, project=gcp_project)

  tf.logging.info(
      "Building TPUConfig with tpu_job_name={}".format(tpu_job_name)
  )
  my_tpu_config = tpu_config.TPUConfig(
      tpu_job_name=tpu_job_name,
      iterations_per_loop=iterations_per_loop,
      num_cores_per_replica=1,
      per_host_input_for_training=tpu_config.InputPipelineConfig.BROADCAST,
  )

  run_config = tpu_config.RunConfig(
      cluster=cluster,
      model_dir=model_dir,
      save_checkpoints_steps=save_checkpoints_steps,
      tpu_config=my_tpu_config)

  transformer_model = build_model(
      model_type=model_type,
      input_vocab_size=inputs_vocabulary(vocabulary).vocab_size,
      output_vocab_size=targets_vocabulary(vocabulary).vocab_size,
      layout_rules=layout_rules,
      mesh_shape=mesh_shape)

  model_fn = tpu_estimator_model_fn(
      model_type=model_type,
      transformer_model=transformer_model,
      model_dir=model_dir,
      use_tpu=tpu,
      mesh_shape=mesh_shape,
      layout_rules=layout_rules,
      batch_size=batch_size,
      sequence_length=sequence_length,
      autostack=autostack,
      metric_names=None)

  estimator = tpu_estimator.TPUEstimator(
      model_fn=model_fn,
      config=run_config,
      train_batch_size=batch_size,
      eval_batch_size=batch_size,
      predict_batch_size=batch_size,
      use_tpu=tpu,
      export_to_tpu=False,
      params={})

  if mode == "train":
    if train_dataset_fn is None:
      raise ValueError("Must provide train_dataset_fn through gin for train.")
    def input_fn(params):
      del params
      dataset = train_dataset_fn(batch_size=batch_size,
                                 sequence_length=sequence_length,
                                 vocabulary=vocabulary,
                                 dataset_split=dataset_split)
      return dataset

    estimator.train(input_fn=input_fn, max_steps=train_steps)
  elif mode == "continuous_eval":
    if get_components_fn is None:
      raise ValueError("Must provide get_components_fn through gin for eval.")
    if eval_dataset_fn is None:
      raise ValueError("Must provide eval_dataset_fn through gin for eval.")
    metrics_inputs = get_components_fn()
    for _ in tf.contrib.training.checkpoints_iterator(estimator.model_dir):
      for metric_names, component in metrics_inputs:
        tf.logging.info("Evaluating {}".format(component.__dict__))
        tf.logging.info("on split {}".format(dataset_split))
        # Prepend eval tag and split name to metric names
        metric_names = [
            "eval/{}/{}".format(dataset_split, n) for n in metric_names
        ]
        # Regenerate the estimator
        model_fn = tpu_estimator_model_fn(
            model_type=model_type,
            transformer_model=transformer_model,
            model_dir=model_dir,
            use_tpu=tpu,
            mesh_shape=mesh_shape,
            layout_rules=layout_rules,
            batch_size=batch_size,
            sequence_length=sequence_length,
            autostack=autostack,
            metric_names=metric_names)
        estimator = tpu_estimator.TPUEstimator(
            model_fn=model_fn,
            config=run_config,
            train_batch_size=batch_size,
            eval_batch_size=batch_size,
            predict_batch_size=batch_size,
            use_tpu=tpu,
            export_to_tpu=False,
            params={})
        def input_fn(params):
          del params
          dataset = eval_dataset_fn(component,  # pylint: disable=cell-var-from-loop
                                    batch_size=batch_size,
                                    sequence_length=sequence_length,
                                    vocabulary=vocabulary,
                                    dataset_split=dataset_split,
                                    pack=False)
          return dataset

        eval_args = {"eval": (input_fn, eval_steps)}
        _ = evaluate(estimator, eval_args)

  elif mode == "infer":
    decode_from_file(
        estimator,
        vocabulary=vocabulary,
        model_type=model_type,
        batch_size=batch_size,
        sequence_length=sequence_length,
        checkpoint_path=checkpoint_path)
  else:
    raise ValueError(
        "unknown mode %s - must be train/evaluate/continuous_eval/infer" % mode)