# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Train and evaluate the Transformer model.

See README for description of setting the training schedule and evaluating the
BLEU score.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import tempfile

from absl import app
from absl import flags
from absl import logging
import tensorflow as tf

# pylint: disable=g-bad-import-order
from official.transformer import compute_bleu
from official.transformer.utils import tokenizer
from official.transformer.v2 import data_pipeline
from official.transformer.v2 import metrics
from official.transformer.v2 import misc
from official.transformer.v2 import optimizer
from official.transformer.v2 import transformer
from official.transformer.v2 import translate
from official.utils.flags import core as flags_core
from official.utils.logs import logger
from official.utils.misc import keras_utils
from official.utils.misc import distribution_utils

INF = int(1e9)
BLEU_DIR = "bleu"
_SINGLE_SAMPLE = 1


def translate_and_compute_bleu(model,
                               params,
                               subtokenizer,
                               bleu_source,
                               bleu_ref,
                               distribution_strategy=None):
  """Translate file and report the cased and uncased bleu scores.

  Args:
    model: A Keras model, used to generate the translations.
    params: A dictionary, containing the translation related parameters.
    subtokenizer: A subtokenizer object, used for encoding and decoding source
      and translated lines.
    bleu_source: A file containing source sentences for translation.
    bleu_ref: A file containing the reference for the translated sentences.
    distribution_strategy: A platform distribution strategy, used for TPU based
      translation.

  Returns:
    uncased_score: A float, the case insensitive BLEU score.
    cased_score: A float, the case sensitive BLEU score.
  """
  # Create temporary file to store translation.
  tmp = tempfile.NamedTemporaryFile(delete=False)
  tmp_filename = tmp.name

  translate.translate_file(
      model,
      params,
      subtokenizer,
      bleu_source,
      output_file=tmp_filename,
      print_all_translations=False,
      distribution_strategy=distribution_strategy)

  # Compute uncased and cased bleu scores.
  uncased_score = compute_bleu.bleu_wrapper(bleu_ref, tmp_filename, False)
  cased_score = compute_bleu.bleu_wrapper(bleu_ref, tmp_filename, True)
  os.remove(tmp_filename)
  return uncased_score, cased_score


def evaluate_and_log_bleu(model,
                          params,
                          bleu_source,
                          bleu_ref,
                          vocab_file,
                          distribution_strategy=None):
  """Calculate and record the BLEU score.

  Args:
    model: A Keras model, used to generate the translations.
    params: A dictionary, containing the translation related parameters.
    bleu_source: A file containing source sentences for translation.
    bleu_ref: A file containing the reference for the translated sentences.
    vocab_file: A file containing the vocabulary for translation.
    distribution_strategy: A platform distribution strategy, used for TPU based
      translation.

  Returns:
    uncased_score: A float, the case insensitive BLEU score.
    cased_score: A float, the case sensitive BLEU score.
  """
  subtokenizer = tokenizer.Subtokenizer(vocab_file)

  uncased_score, cased_score = translate_and_compute_bleu(
      model, params, subtokenizer, bleu_source, bleu_ref, distribution_strategy)

  logging.info("Bleu score (uncased): %s", uncased_score)
  logging.info("Bleu score (cased): %s", cased_score)
  return uncased_score, cased_score


class TransformerTask(object):
  """Main entry of Transformer model."""

  def __init__(self, flags_obj):
    """Init function of TransformerMain.

    Args:
      flags_obj: Object containing parsed flag values, i.e., FLAGS.

    Raises:
      ValueError: if not using static batch for input data on TPU.
    """
    self.flags_obj = flags_obj
    self.predict_model = None

    # Add flag-defined parameters to params object
    num_gpus = flags_core.get_num_gpus(flags_obj)
    self.params = params = misc.get_model_params(flags_obj.param_set, num_gpus)

    params["num_gpus"] = num_gpus
    params["use_ctl"] = flags_obj.use_ctl
    params["data_dir"] = flags_obj.data_dir
    params["model_dir"] = flags_obj.model_dir
    params["static_batch"] = flags_obj.static_batch
    params["max_length"] = flags_obj.max_length
    params["decode_batch_size"] = flags_obj.decode_batch_size
    params["decode_max_length"] = flags_obj.decode_max_length
    params["padded_decode"] = flags_obj.padded_decode
    params["num_parallel_calls"] = (
        flags_obj.num_parallel_calls or tf.data.experimental.AUTOTUNE)

    params["use_synthetic_data"] = flags_obj.use_synthetic_data
    params["batch_size"] = flags_obj.batch_size or params["default_batch_size"]
    params["repeat_dataset"] = None
    params["dtype"] = flags_core.get_tf_dtype(flags_obj)
    params["enable_tensorboard"] = flags_obj.enable_tensorboard
    params["enable_metrics_in_training"] = flags_obj.enable_metrics_in_training
    params["steps_between_evals"] = flags_obj.steps_between_evals

    self.distribution_strategy = distribution_utils.get_distribution_strategy(
        distribution_strategy=flags_obj.distribution_strategy,
        num_gpus=num_gpus,
        all_reduce_alg=flags_obj.all_reduce_alg,
        num_packs=flags_obj.num_packs,
        tpu_address=flags_obj.tpu or "")
    if self.use_tpu:
      params["num_replicas"] = self.distribution_strategy.num_replicas_in_sync
      if not params["static_batch"]:
        raise ValueError("TPU requires static batch for input data.")
    else:
      logging.info("Running transformer with num_gpus = %d", num_gpus)

    if self.distribution_strategy:
      logging.info("For training, using distribution strategy: %s",
                   self.distribution_strategy)
    else:
      logging.info("Not using any distribution strategy.")

    if params["dtype"] == tf.float16:
      # TODO(reedwm): It's pretty ugly to set the global policy in a constructor
      # like this. What if multiple instances of TransformerTask are created?
      # We should have a better way in the tf.keras.mixed_precision API of doing
      # this.
      loss_scale = flags_core.get_loss_scale(
          flags_obj, default_for_fp16="dynamic")
      policy = tf.compat.v2.keras.mixed_precision.experimental.Policy(
          "mixed_float16", loss_scale=loss_scale)
      tf.compat.v2.keras.mixed_precision.experimental.set_policy(policy)

    elif params["dtype"] == tf.bfloat16:
      policy = tf.compat.v2.keras.mixed_precision.experimental.Policy(
          "mixed_bfloat16")
      tf.compat.v2.keras.mixed_precision.experimental.set_policy(policy)

  @property
  def use_tpu(self):
    if self.distribution_strategy:
      return isinstance(self.distribution_strategy,
                        tf.distribute.experimental.TPUStrategy)
    return False

  def train(self):
    """Trains the model."""
    params = self.params
    flags_obj = self.flags_obj
    # Sets config options.
    keras_utils.set_session_config(enable_xla=flags_obj.enable_xla)

    _ensure_dir(flags_obj.model_dir)
    with distribution_utils.get_strategy_scope(self.distribution_strategy):
      model = transformer.create_model(params, is_train=True)
      opt = self._create_optimizer()

      current_step = 0
      checkpoint = tf.train.Checkpoint(model=model, optimizer=opt)
      latest_checkpoint = tf.train.latest_checkpoint(flags_obj.model_dir)
      if latest_checkpoint:
        checkpoint.restore(latest_checkpoint)
        logging.info("Loaded checkpoint %s", latest_checkpoint)
        current_step = opt.iterations.numpy()

      if params["use_ctl"]:
        train_loss_metric = tf.keras.metrics.Mean(
            "training_loss", dtype=tf.float32)
        if params["enable_tensorboard"]:
          summary_writer = tf.compat.v2.summary.create_file_writer(
              flags_obj.model_dir)
        else:
          summary_writer = tf.compat.v2.summary.create_noop_writer()
        train_metrics = [train_loss_metric]
        if params["enable_metrics_in_training"]:
          train_metrics = train_metrics + model.metrics
      else:
        model.compile(opt)

    model.summary()

    if self.use_tpu:
      # Different from experimental_distribute_dataset,
      # experimental_distribute_datasets_from_function requires
      # per-replica/local batch size.
      params["batch_size"] /= self.distribution_strategy.num_replicas_in_sync
      train_ds = (
          self.distribution_strategy
          .experimental_distribute_datasets_from_function(
              lambda ctx: data_pipeline.train_input_fn(params, ctx)))
    else:
      train_ds = data_pipeline.train_input_fn(params)
      map_data_fn = data_pipeline.map_data_for_transformer_fn
      train_ds = train_ds.map(
          map_data_fn, num_parallel_calls=params["num_parallel_calls"])
    if params["use_ctl"]:
      train_ds_iterator = iter(train_ds)

    callbacks = self._create_callbacks(flags_obj.model_dir, 0, params)

    # TODO(b/139418525): Refactor the custom training loop logic.
    @tf.function
    def train_steps(iterator, steps):
      """Training steps function for TPU runs.

      Args:
        iterator: The input iterator of the training dataset.
        steps: An integer, the number of training steps.

      Returns:
        A float, the loss value.
      """

      def _step_fn(inputs):
        """Per-replica step function."""
        inputs, targets = inputs
        with tf.GradientTape() as tape:
          logits = model([inputs, targets], training=True)
          loss = metrics.transformer_loss(logits, targets,
                                          params["label_smoothing"],
                                          params["vocab_size"])
          # Scales the loss, which results in using the average loss across all
          # of the replicas for backprop.
          scaled_loss = loss / self.distribution_strategy.num_replicas_in_sync

        # De-dupes variables due to keras tracking issues.
        tvars = list({id(v): v for v in model.trainable_variables}.values())
        grads = tape.gradient(scaled_loss, tvars)
        opt.apply_gradients(zip(grads, tvars))
        # For reporting, the metric takes the mean of losses.
        train_loss_metric.update_state(loss)

      for _ in tf.range(steps):
        train_loss_metric.reset_states()
        self.distribution_strategy.experimental_run_v2(
            _step_fn, args=(next(iterator),))

    cased_score, uncased_score = None, None
    cased_score_history, uncased_score_history = [], []
    while current_step < flags_obj.train_steps:
      remaining_steps = flags_obj.train_steps - current_step
      train_steps_per_eval = (
          remaining_steps if remaining_steps < flags_obj.steps_between_evals
          else flags_obj.steps_between_evals)
      current_iteration = current_step // flags_obj.steps_between_evals

      logging.info(
          "Start train iteration at global step:{}".format(current_step))
      history = None
      if params["use_ctl"]:
        if not self.use_tpu:
          raise NotImplementedError(
              "Custom training loop on GPUs is not implemented.")
        # Runs training steps.
        with summary_writer.as_default():
          train_steps(
              train_ds_iterator,
              tf.convert_to_tensor(train_steps_per_eval, dtype=tf.int32))
          current_step += train_steps_per_eval
          train_loss = train_loss_metric.result().numpy().astype(float)
          logging.info("Train Step: %d/%d / loss = %s", current_step,
                       flags_obj.train_steps, train_loss)

          if params["enable_tensorboard"]:
            for metric_obj in train_metrics:
              tf.compat.v2.summary.scalar(metric_obj.name, metric_obj.result(),
                                          current_step)

        checkpoint_name = checkpoint.save(
            os.path.join(flags_obj.model_dir,
                         "ctl_step_{}.ckpt".format(current_step)))
        logging.info("Saved checkpoint to %s", checkpoint_name)
      else:
        if self.use_tpu:
          raise NotImplementedError(
              "Keras model.fit on TPUs is not implemented.")
        history = model.fit(
            train_ds,
            initial_epoch=current_iteration,
            epochs=current_iteration + 1,
            steps_per_epoch=train_steps_per_eval,
            callbacks=callbacks,
            # If TimeHistory is enabled, progress bar would be messy. Increase
            # the verbose level to get rid of it.
            verbose=(2 if flags_obj.enable_time_history else 1))
        current_step += train_steps_per_eval
        logging.info("Train history: {}".format(history.history))

      logging.info("End train iteration at global step:{}".format(current_step))

      if (flags_obj.bleu_source and flags_obj.bleu_ref):
        uncased_score, cased_score = self.eval()
        cased_score_history.append([current_iteration + 1, cased_score])
        uncased_score_history.append([current_iteration + 1, uncased_score])

    stats = ({
        "loss": train_loss
    } if history is None else misc.build_stats(history, callbacks))
    if uncased_score and cased_score:
      stats["bleu_uncased"] = uncased_score
      stats["bleu_cased"] = cased_score
      stats["bleu_uncased_history"] = uncased_score_history
      stats["bleu_cased_history"] = cased_score_history
    return stats

  def eval(self):
    """Evaluates the model."""
    distribution_strategy = self.distribution_strategy if self.use_tpu else None

    # We only want to create the model under DS scope for TPU case.
    # When 'distribution_strategy' is None, a no-op DummyContextManager will
    # be used.
    with distribution_utils.get_strategy_scope(distribution_strategy):
      if not self.predict_model:
        self.predict_model = transformer.create_model(self.params, False)
      self._load_weights_if_possible(
          self.predict_model,
          tf.train.latest_checkpoint(self.flags_obj.model_dir))
      self.predict_model.summary()
    return evaluate_and_log_bleu(
        self.predict_model, self.params, self.flags_obj.bleu_source,
        self.flags_obj.bleu_ref, self.flags_obj.vocab_file,
        distribution_strategy)

  def predict(self):
    """Predicts result from the model."""
    params = self.params
    flags_obj = self.flags_obj

    with tf.name_scope("model"):
      model = transformer.create_model(params, is_train=False)
      self._load_weights_if_possible(
          model, tf.train.latest_checkpoint(self.flags_obj.model_dir))
      model.summary()
    subtokenizer = tokenizer.Subtokenizer(flags_obj.vocab_file)

    ds = data_pipeline.eval_input_fn(params)
    ds = ds.map(lambda x, y: x).take(_SINGLE_SAMPLE)
    ret = model.predict(ds)
    val_outputs, _ = ret
    length = len(val_outputs)
    for i in range(length):
      translate.translate_from_input(val_outputs[i], subtokenizer)

  def _create_callbacks(self, cur_log_dir, init_steps, params):
    """Creates a list of callbacks."""
    sfunc = optimizer.LearningRateFn(params["learning_rate"],
                                     params["hidden_size"],
                                     params["learning_rate_warmup_steps"])
    scheduler_callback = optimizer.LearningRateScheduler(sfunc, init_steps)
    callbacks = misc.get_callbacks(params["steps_between_evals"])
    callbacks.append(scheduler_callback)
    ckpt_full_path = os.path.join(cur_log_dir, "cp-{epoch:04d}.ckpt")
    callbacks.append(
        tf.keras.callbacks.ModelCheckpoint(
            ckpt_full_path, save_weights_only=True))
    return callbacks

  def _load_weights_if_possible(self, model, init_weight_path=None):
    """Loads model weights when it is provided."""
    if init_weight_path:
      logging.info("Load weights: {}".format(init_weight_path))
      # TODO(b/139414977): Having the same variable restoring method for both
      # TPU and GPU.
      if self.use_tpu:
        checkpoint = tf.train.Checkpoint(
            model=model, optimizer=self._create_optimizer())
        checkpoint.restore(init_weight_path)
      else:
        model.load_weights(init_weight_path)
    else:
      logging.info("Weights not loaded from path:{}".format(init_weight_path))

  def _create_optimizer(self):
    """Creates optimizer."""
    params = self.params
    # TODO(b/139414679): Explore the difference between using
    # LearningRateSchedule and callback for GPU runs, and try to merge them.
    lr_schedule = optimizer.LearningRateSchedule(
        params["learning_rate"], params["hidden_size"],
        params["learning_rate_warmup_steps"])
    opt = tf.keras.optimizers.Adam(
        lr_schedule if self.use_tpu else params["learning_rate"],
        params["optimizer_adam_beta1"],
        params["optimizer_adam_beta2"],
        epsilon=params["optimizer_adam_epsilon"])

    if params["dtype"] == tf.float16:
      opt = tf.keras.mixed_precision.experimental.LossScaleOptimizer(
          opt,
          loss_scale=flags_core.get_loss_scale(
              self.flags_obj, default_for_fp16="dynamic"))
    if self.flags_obj.fp16_implementation == "graph_rewrite":
      # Note: when flags_obj.fp16_implementation == "graph_rewrite", dtype as
      # determined by flags_core.get_tf_dtype(flags_obj) would be 'float32'
      # which will ensure tf.compat.v2.keras.mixed_precision and
      # tf.train.experimental.enable_mixed_precision_graph_rewrite do not double
      # up.
      opt = tf.train.experimental.enable_mixed_precision_graph_rewrite(opt)

    return opt


def _ensure_dir(log_dir):
  """Makes log dir if not existed."""
  if not tf.io.gfile.exists(log_dir):
    tf.io.gfile.makedirs(log_dir)


def main(_):
  flags_obj = flags.FLAGS
  with logger.benchmark_context(flags_obj):
    task = TransformerTask(flags_obj)

    # Execute flag override logic for better model performance
    if flags_obj.tf_gpu_thread_mode:
      keras_utils.set_gpu_thread_mode_and_count(
          per_gpu_thread_count=flags_obj.per_gpu_thread_count,
          gpu_thread_mode=flags_obj.tf_gpu_thread_mode,
          num_gpus=flags_obj.num_gpus,
          datasets_num_private_threads=flags_obj.datasets_num_private_threads)

    if flags_obj.mode == "train":
      task.train()
    elif flags_obj.mode == "predict":
      task.predict()
    elif flags_obj.mode == "eval":
      task.eval()
    else:
      raise ValueError("Invalid mode {}".format(flags_obj.mode))


if __name__ == "__main__":
  tf.compat.v1.enable_v2_behavior()
  logging.set_verbosity(logging.INFO)
  misc.define_transformer_flags()
  app.run(main)

tensorflow.convert_to_tensor tensorflow.keras.callbacks.ModelCheckpoint tensorflow.compat.v1.enable_v2_behavior tensorflow.train.latest_checkpoint tensorflow.range tensorflow.io.gfile.exists tensorflow.train.Checkpoint tensorflow.compat.v2.keras.mixed_precision.experimental.Policy tensorflow.compat.v2.summary.create_file_writer tensorflow.io.gfile.makedirs tensorflow.compat.v2.keras.mixed_precision.experimental.set_policy tensorflow.compat.v2.summary.create_noop_writer tensorflow.keras.optimizers.Adam tensorflow.name_scope tensorflow.train.experimental.enable_mixed_precision_graph_rewrite tensorflow.keras.metrics.Mean tensorflow.GradientTape
[(75, 'tempfile.NamedTemporaryFile', 'tempfile.NamedTemporaryFile', ([], {'delete': '(False)'}), False, 'import tempfile\n'), (78, 'official.transformer.v2.translate.translate_file', 'translate.translate_file', (['model', 'params', 'subtokenizer', 'bleu_source'], {'output_file': 'tmp_filename', 'print_all_translations': '(False)', 'distribution_strategy': 'distribution_strategy'}), False, 'from official.transformer.v2 import translate\n'), (88, 'official.transformer.compute_bleu.bleu_wrapper', 'compute_bleu.bleu_wrapper', (['bleu_ref', 'tmp_filename', '(False)'], {}), False, 'from official.transformer import compute_bleu\n'), (89, 'official.transformer.compute_bleu.bleu_wrapper', 'compute_bleu.bleu_wrapper', (['bleu_ref', 'tmp_filename', '(True)'], {}), False, 'from official.transformer import compute_bleu\n'), (90, 'os.remove', 'os.remove', (['tmp_filename'], {}), False, 'import os\n'), (115, 'official.transformer.utils.tokenizer.Subtokenizer', 'tokenizer.Subtokenizer', (['vocab_file'], {}), False, 'from official.transformer.utils import tokenizer\n'), (120, 'absl.logging.info', 'logging.info', (['"""Bleu score (uncased): %s"""', 'uncased_score'], {}), False, 'from absl import logging\n'), (121, 'absl.logging.info', 'logging.info', (['"""Bleu score (cased): %s"""', 'cased_score'], {}), False, 'from absl import logging\n'), (494, 'tensorflow.compat.v1.enable_v2_behavior', 'tf.compat.v1.enable_v2_behavior', ([], {}), True, 'import tensorflow as tf\n'), (495, 'absl.logging.set_verbosity', 'logging.set_verbosity', (['logging.INFO'], {}), False, 'from absl import logging\n'), (496, 'official.transformer.v2.misc.define_transformer_flags', 'misc.define_transformer_flags', ([], {}), False, 'from official.transformer.v2 import misc\n'), (497, 'absl.app.run', 'app.run', (['main'], {}), False, 'from absl import app\n'), (141, 'official.utils.flags.core.get_num_gpus', 'flags_core.get_num_gpus', (['flags_obj'], {}), True, 'from official.utils.flags import core as flags_core\n'), (142, 'official.transformer.v2.misc.get_model_params', 'misc.get_model_params', (['flags_obj.param_set', 'num_gpus'], {}), False, 'from official.transformer.v2 import misc\n'), (159, 'official.utils.flags.core.get_tf_dtype', 'flags_core.get_tf_dtype', (['flags_obj'], {}), True, 'from official.utils.flags import core as flags_core\n'), (164, 'official.utils.misc.distribution_utils.get_distribution_strategy', 'distribution_utils.get_distribution_strategy', ([], {'distribution_strategy': 'flags_obj.distribution_strategy', 'num_gpus': 'num_gpus', 'all_reduce_alg': 'flags_obj.all_reduce_alg', 'num_packs': 'flags_obj.num_packs', 'tpu_address': "(flags_obj.tpu or '')"}), False, 'from official.utils.misc import distribution_utils\n'), (211, 'official.utils.misc.keras_utils.set_session_config', 'keras_utils.set_session_config', ([], {'enable_xla': 'flags_obj.enable_xla'}), False, 'from official.utils.misc import keras_utils\n'), (395, 'official.transformer.utils.tokenizer.Subtokenizer', 'tokenizer.Subtokenizer', (['flags_obj.vocab_file'], {}), False, 'from official.transformer.utils import tokenizer\n'), (397, 'official.transformer.v2.data_pipeline.eval_input_fn', 'data_pipeline.eval_input_fn', (['params'], {}), False, 'from official.transformer.v2 import data_pipeline\n'), (407, 'official.transformer.v2.optimizer.LearningRateFn', 'optimizer.LearningRateFn', (["params['learning_rate']", "params['hidden_size']", "params['learning_rate_warmup_steps']"], {}), False, 'from official.transformer.v2 import optimizer\n'), (410, 'official.transformer.v2.optimizer.LearningRateScheduler', 'optimizer.LearningRateScheduler', (['sfunc', 'init_steps'], {}), False, 'from official.transformer.v2 import optimizer\n'), (411, 'official.transformer.v2.misc.get_callbacks', 'misc.get_callbacks', (["params['steps_between_evals']"], {}), False, 'from official.transformer.v2 import misc\n'), (413, 'os.path.join', 'os.path.join', (['cur_log_dir', '"""cp-{epoch:04d}.ckpt"""'], {}), False, 'import os\n'), (439, 'official.transformer.v2.optimizer.LearningRateSchedule', 'optimizer.LearningRateSchedule', (["params['learning_rate']", "params['hidden_size']", "params['learning_rate_warmup_steps']"], {}), False, 'from official.transformer.v2 import optimizer\n'), (442, 'tensorflow.keras.optimizers.Adam', 'tf.keras.optimizers.Adam', (["(lr_schedule if self.use_tpu else params['learning_rate'])", "params['optimizer_adam_beta1']", "params['optimizer_adam_beta2']"], {'epsilon': "params['optimizer_adam_epsilon']"}), True, 'import tensorflow as tf\n'), (466, 'tensorflow.io.gfile.exists', 'tf.io.gfile.exists', (['log_dir'], {}), True, 'import tensorflow as tf\n'), (467, 'tensorflow.io.gfile.makedirs', 'tf.io.gfile.makedirs', (['log_dir'], {}), True, 'import tensorflow as tf\n'), (472, 'official.utils.logs.logger.benchmark_context', 'logger.benchmark_context', (['flags_obj'], {}), False, 'from official.utils.logs import logger\n'), (175, 'absl.logging.info', 'logging.info', (['"""Running transformer with num_gpus = %d"""', 'num_gpus'], {}), False, 'from absl import logging\n'), (178, 'absl.logging.info', 'logging.info', (['"""For training, using distribution strategy: %s"""', 'self.distribution_strategy'], {}), False, 'from absl import logging\n'), (181, 'absl.logging.info', 'logging.info', (['"""Not using any distribution strategy."""'], {}), False, 'from absl import logging\n'), (188, 'official.utils.flags.core.get_loss_scale', 'flags_core.get_loss_scale', (['flags_obj'], {'default_for_fp16': '"""dynamic"""'}), True, 'from official.utils.flags import core as flags_core\n'), (190, 'tensorflow.compat.v2.keras.mixed_precision.experimental.Policy', 'tf.compat.v2.keras.mixed_precision.experimental.Policy', (['"""mixed_float16"""'], {'loss_scale': 'loss_scale'}), True, 'import tensorflow as tf\n'), (192, 'tensorflow.compat.v2.keras.mixed_precision.experimental.set_policy', 'tf.compat.v2.keras.mixed_precision.experimental.set_policy', (['policy'], {}), True, 'import tensorflow as tf\n'), (214, 'official.utils.misc.distribution_utils.get_strategy_scope', 'distribution_utils.get_strategy_scope', (['self.distribution_strategy'], {}), False, 'from official.utils.misc import distribution_utils\n'), (215, 'official.transformer.v2.transformer.create_model', 'transformer.create_model', (['params'], {'is_train': '(True)'}), False, 'from official.transformer.v2 import transformer\n'), (219, 'tensorflow.train.Checkpoint', 'tf.train.Checkpoint', ([], {'model': 'model', 'optimizer': 'opt'}), True, 'import tensorflow as tf\n'), (220, 'tensorflow.train.latest_checkpoint', 'tf.train.latest_checkpoint', (['flags_obj.model_dir'], {}), True, 'import tensorflow as tf\n'), (252, 'official.transformer.v2.data_pipeline.train_input_fn', 'data_pipeline.train_input_fn', (['params'], {}), False, 'from official.transformer.v2 import data_pipeline\n'), (293, 'tensorflow.range', 'tf.range', (['steps'], {}), True, 'import tensorflow as tf\n'), (358, 'official.transformer.v2.misc.build_stats', 'misc.build_stats', (['history', 'callbacks'], {}), False, 'from official.transformer.v2 import misc\n'), (373, 'official.utils.misc.distribution_utils.get_strategy_scope', 'distribution_utils.get_strategy_scope', (['distribution_strategy'], {}), False, 'from official.utils.misc import distribution_utils\n'), (390, 'tensorflow.name_scope', 'tf.name_scope', (['"""model"""'], {}), True, 'import tensorflow as tf\n'), (391, 'official.transformer.v2.transformer.create_model', 'transformer.create_model', (['params'], {'is_train': '(False)'}), False, 'from official.transformer.v2 import transformer\n'), (403, 'official.transformer.v2.translate.translate_from_input', 'translate.translate_from_input', (['val_outputs[i]', 'subtokenizer'], {}), False, 'from official.transformer.v2 import translate\n'), (415, 'tensorflow.keras.callbacks.ModelCheckpoint', 'tf.keras.callbacks.ModelCheckpoint', (['ckpt_full_path'], {'save_weights_only': '(True)'}), True, 'import tensorflow as tf\n'), (459, 'tensorflow.train.experimental.enable_mixed_precision_graph_rewrite', 'tf.train.experimental.enable_mixed_precision_graph_rewrite', (['opt'], {}), True, 'import tensorflow as tf\n'), (477, 'official.utils.misc.keras_utils.set_gpu_thread_mode_and_count', 'keras_utils.set_gpu_thread_mode_and_count', ([], {'per_gpu_thread_count': 'flags_obj.per_gpu_thread_count', 'gpu_thread_mode': 'flags_obj.tf_gpu_thread_mode', 'num_gpus': 'flags_obj.num_gpus', 'datasets_num_private_threads': 'flags_obj.datasets_num_private_threads'}), False, 'from official.utils.misc import keras_utils\n'), (195, 'tensorflow.compat.v2.keras.mixed_precision.experimental.Policy', 'tf.compat.v2.keras.mixed_precision.experimental.Policy', (['"""mixed_bfloat16"""'], {}), True, 'import tensorflow as tf\n'), (197, 'tensorflow.compat.v2.keras.mixed_precision.experimental.set_policy', 'tf.compat.v2.keras.mixed_precision.experimental.set_policy', (['policy'], {}), True, 'import tensorflow as tf\n'), (223, 'absl.logging.info', 'logging.info', (['"""Loaded checkpoint %s"""', 'latest_checkpoint'], {}), False, 'from absl import logging\n'), (227, 'tensorflow.keras.metrics.Mean', 'tf.keras.metrics.Mean', (['"""training_loss"""'], {'dtype': 'tf.float32'}), True, 'import tensorflow as tf\n'), (332, 'absl.logging.info', 'logging.info', (['"""Saved checkpoint to %s"""', 'checkpoint_name'], {}), False, 'from absl import logging\n'), (375, 'official.transformer.v2.transformer.create_model', 'transformer.create_model', (['self.params', '(False)'], {}), False, 'from official.transformer.v2 import transformer\n'), (378, 'tensorflow.train.latest_checkpoint', 'tf.train.latest_checkpoint', (['self.flags_obj.model_dir'], {}), True, 'import tensorflow as tf\n'), (393, 'tensorflow.train.latest_checkpoint', 'tf.train.latest_checkpoint', (['self.flags_obj.model_dir'], {}), True, 'import tensorflow as tf\n'), (230, 'tensorflow.compat.v2.summary.create_file_writer', 'tf.compat.v2.summary.create_file_writer', (['flags_obj.model_dir'], {}), True, 'import tensorflow as tf\n'), (233, 'tensorflow.compat.v2.summary.create_noop_writer', 'tf.compat.v2.summary.create_noop_writer', ([], {}), True, 'import tensorflow as tf\n'), (250, 'official.transformer.v2.data_pipeline.train_input_fn', 'data_pipeline.train_input_fn', (['params', 'ctx'], {}), False, 'from official.transformer.v2 import data_pipeline\n'), (277, 'tensorflow.GradientTape', 'tf.GradientTape', ([], {}), True, 'import tensorflow as tf\n'), (279, 'official.transformer.v2.metrics.transformer_loss', 'metrics.transformer_loss', (['logits', 'targets', "params['label_smoothing']", "params['vocab_size']"], {}), False, 'from official.transformer.v2 import metrics\n'), (321, 'absl.logging.info', 'logging.info', (['"""Train Step: %d/%d / loss = %s"""', 'current_step', 'flags_obj.train_steps', 'train_loss'], {}), False, 'from absl import logging\n'), (451, 'official.utils.flags.core.get_loss_scale', 'flags_core.get_loss_scale', (['self.flags_obj'], {'default_for_fp16': '"""dynamic"""'}), True, 'from official.utils.flags import core as flags_core\n'), (318, 'tensorflow.convert_to_tensor', 'tf.convert_to_tensor', (['train_steps_per_eval'], {'dtype': 'tf.int32'}), True, 'import tensorflow as tf\n')]