# Copyright 2017 Google Inc. 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.
# ==============================================================================
#
# Copyright (c) 2019, NVIDIA CORPORATION. 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.

"""TensorFlow NMT model implementation."""
from __future__ import print_function

import argparse
import os
import random
import sys
import subprocess

# import matplotlib.image as mpimg
import numpy as np
import time
import tensorflow as tf

import dllogger

import estimator
from utils import evaluation_utils
from utils import iterator_utils
from utils import misc_utils as utils
from utils import vocab_utils
from variable_mgr import constants

utils.check_tensorflow_version()

FLAGS = None


# LINT.IfChange
def add_arguments(parser):
  """Build ArgumentParser."""
  parser.register("type", "bool", lambda v: v.lower() == "true")

  # network
  parser.add_argument(
      "--num_units", type=int, default=1024, help="Network size.")
  parser.add_argument(
      "--num_layers", type=int, default=4, help="Network depth.")
  parser.add_argument("--num_encoder_layers", type=int, default=None,
                      help="Encoder depth, equal to num_layers if None.")
  parser.add_argument("--num_decoder_layers", type=int, default=None,
                      help="Decoder depth, equal to num_layers if None.")
  parser.add_argument(
      "--encoder_type",
      type=str,
      default="gnmt",
      help="""\
      uni | bi | gnmt.
      For bi, we build num_encoder_layers/2 bi-directional layers.
      For gnmt, we build 1 bi-directional layer, and (num_encoder_layers - 1)
        uni-directional layers.\
      """)
  parser.add_argument(
      "--residual",
      type="bool",
      nargs="?",
      const=True,
      default=True,
      help="Whether to add residual connections.")
  parser.add_argument("--time_major", type="bool", nargs="?", const=True,
                      default=True,
                      help="Whether to use time-major mode for dynamic RNN.")
  parser.add_argument("--num_embeddings_partitions", type=int, default=0,
                      help="Number of partitions for embedding vars.")

  # attention mechanisms
  parser.add_argument(
      "--attention",
      type=str,
      default="normed_bahdanau",
      help="""\
      luong | scaled_luong | bahdanau | normed_bahdanau or set to "" for no
      attention\
      """)
  parser.add_argument(
      "--attention_architecture",
      type=str,
      default="gnmt_v2",
      help="""\
      standard | gnmt | gnmt_v2.
      standard: use top layer to compute attention.
      gnmt: GNMT style of computing attention, use previous bottom layer to
          compute attention.
      gnmt_v2: similar to gnmt, but use current bottom layer to compute
          attention.\
      """)
  parser.add_argument(
      "--output_attention", type="bool", nargs="?", const=True,
      default=True,
      help="""\
      Only used in standard attention_architecture. Whether use attention as
      the cell output at each timestep.
      .\
      """)
  parser.add_argument(
      "--pass_hidden_state", type="bool", nargs="?", const=True,
      default=True,
      help="""\
      Whether to pass encoder's hidden state to decoder when using an attention
      based model.\
      """)

  # optimizer
  parser.add_argument(
      "--optimizer", type=str, default="adam", help="sgd | adam")
  parser.add_argument(
      "--learning_rate",
      type=float,
      default=5e-4,
      help="Learning rate. Adam: 0.001 | 0.0001")
  parser.add_argument("--warmup_steps", type=int, default=200,
                      help="How many steps we inverse-decay learning.")
  parser.add_argument("--warmup_scheme", type=str, default="t2t", help="""\
      How to warmup learning rates. Options include:
        t2t: Tensor2Tensor's way, start with lr 100 times smaller, then
             exponentiate until the specified lr.\
      """)
  parser.add_argument(
      "--decay_scheme", type=str, default="luong234", help="""\
      How we decay learning rate. Options include:
        luong234: after 2/3 num train steps, we start halving the learning rate
          for 4 times before finishing.
        luong5: after 1/2 num train steps, we start halving the learning rate
          for 5 times before finishing.\
        luong10: after 1/2 num train steps, we start halving the learning rate
          for 10 times before finishing.\
      """)

  parser.add_argument(
      "--max_train_epochs", type=int, default=6, help="Max number of epochs.")
  parser.add_argument(
      "--target_bleu", type=float, default=None, help="Target bleu.")
  parser.add_argument("--colocate_gradients_with_ops", type="bool", nargs="?",
                      const=True,
                      default=True,
                      help=("Whether try colocating gradients with "
                            "corresponding op"))
  parser.add_argument("--label_smoothing", type=float, default=0.1,
                      help=("If nonzero, smooth the labels towards "
                            "1/num_classes."))

  # initializer
  parser.add_argument("--init_op", type=str, default="uniform",
                      help="uniform | glorot_normal | glorot_uniform")
  parser.add_argument("--init_weight", type=float, default=0.1,
                      help=("for uniform init_op, initialize weights "
                            "between [-this, this]."))

  # data
  parser.add_argument(
      "--src", type=str, default="en", help="Source suffix, e.g., en.")
  parser.add_argument(
      "--tgt", type=str, default="de", help="Target suffix, e.g., de.")
  parser.add_argument(
      "--data_dir", type=str, default="data/wmt16_de_en",
      help="Training/eval data directory.")

  parser.add_argument(
      "--train_prefix",
      type=str,
      default="train.tok.clean.bpe.32000",
      help="Train prefix, expect files with src/tgt suffixes.")
  parser.add_argument(
      "--test_prefix",
      type=str,
      default="newstest2014.tok.bpe.32000",
      help="Test prefix, expect files with src/tgt suffixes.")

  parser.add_argument(
      "--translate_file",
      type=str,
      help="File to translate, works only with translate mode")

  parser.add_argument(
      "--output_dir", type=str, default="results",
      help="Store log/model files.")

  # Vocab
  parser.add_argument(
      "--vocab_prefix",
      type=str,
      default="vocab.bpe.32000",
      help="""\
      Vocab prefix, expect files with src/tgt suffixes.\
      """)

  parser.add_argument(
      "--embed_prefix",
      type=str,
      default=None,
      help="""\
      Pretrained embedding prefix, expect files with src/tgt suffixes.
      The embedding files should be Glove formatted txt files.\
      """)
  parser.add_argument("--sos", type=str, default="<s>",
                      help="Start-of-sentence symbol.")
  parser.add_argument("--eos", type=str, default="</s>",
                      help="End-of-sentence symbol.")
  parser.add_argument(
      "--share_vocab",
      type="bool",
      nargs="?",
      const=True,
      default=True,
      help="""\
      Whether to use the source vocab and embeddings for both source and
      target.\
      """)
  parser.add_argument("--check_special_token", type="bool", default=True,
                      help="""\
                      Whether check special sos, eos, unk tokens exist in the
                      vocab files.\
                      """)

  # Sequence lengths
  parser.add_argument(
      "--src_max_len",
      type=int,
      default=50,
      help="Max length of src sequences during training (including EOS).")
  parser.add_argument(
      "--tgt_max_len",
      type=int,
      default=50,
      help="Max length of tgt sequences during training (including BOS).")
  parser.add_argument("--src_max_len_infer", type=int, default=None,
                      help="Max length of src sequences during inference (including EOS).")
  parser.add_argument("--tgt_max_len_infer", type=int, default=80,
                      help="""\
      Max length of tgt sequences during inference (including BOS).  Also use to restrict the
      maximum decoding length.\
      """)

  # Default settings works well (rarely need to change)
  parser.add_argument("--unit_type", type=str, default="lstm",
                      help="lstm | gru | layer_norm_lstm | nas")
  parser.add_argument("--forget_bias", type=float, default=0.0,
                      help="Forget bias for BasicLSTMCell.")
  parser.add_argument("--dropout", type=float, default=0.2,
                      help="Dropout rate (not keep_prob)")
  parser.add_argument("--max_gradient_norm", type=float, default=5.0,
                      help="Clip gradients to this norm.")
  parser.add_argument("--batch_size", type=int, default=128, help="Total batch size.")

  parser.add_argument(
      "--num_buckets",
      type=int,
      default=5,
      help="Put data into similar-length buckets (only for training).")

  # SPM
  parser.add_argument("--subword_option", type=str, default="bpe",
                      choices=["", "bpe", "spm"],
                      help="""\
                      Set to bpe or spm to activate subword desegmentation.\
                      """)

  # Experimental encoding feature.
  parser.add_argument("--use_char_encode", type="bool", default=False,
                      help="""\
                      Whether to split each word or bpe into character, and then
                      generate the word-level representation from the character
                      reprentation.
                      """)

  # Misc
  parser.add_argument(
      "--save_checkpoints_steps", type=int, default=2000,
      help="save_checkpoints_steps")
  parser.add_argument(
      "--log_step_count_steps", type=int, default=10,
      help=("The frequency, in number of global steps, that the global step "
            "and the loss will be logged during training"))
  parser.add_argument(
      "--num_gpus", type=int, default=1, help="Number of gpus in each worker.")
  parser.add_argument("--hparams_path", type=str, default=None,
      help=("Path to standard hparams json file that overrides"
            "hparams values from FLAGS."))
  parser.add_argument(
      "--random_seed",
      type=int,
      default=1,
      help="Random seed (>0, set a specific seed).")
  parser.add_argument("--language_model", type="bool", nargs="?",
                      const=True, default=False,
                      help="True to train a language model, ignoring encoder")

  # Inference
  parser.add_argument("--ckpt", type=str, default=None,
                      help="Checkpoint file to load a model for inference. (defaults to newest checkpoint)")
  parser.add_argument(
      "--infer_batch_size",
      type=int,
      default=128,
      help="Batch size for inference mode.")
  parser.add_argument("--detokenizer_file", type=str,
                      default=None,
                      help=("""Detokenizer script file. Default: DATA_DIR/mosesdecoder/scripts/tokenizer/detokenizer.perl"""))
  parser.add_argument("--tokenizer_file", type=str,
                      default=None,
                      help=("""Tokenizer script file. Default: DATA_DIR/mosesdecoder/scripts/tokenizer/tokenizer.perl"""))

  # Advanced inference arguments
  parser.add_argument("--infer_mode", type=str, default="beam_search",
                      choices=["greedy", "beam_search"],
                      help="Which type of decoder to use during inference.")
  parser.add_argument("--beam_width", type=int, default=5,
                      help=("""\
      beam width when using beam search decoder. If 0, use standard
      decoder with greedy helper.\
      """))
  parser.add_argument(
      "--length_penalty_weight",
      type=float,
      default=0.6,
      help="Length penalty for beam search.")
  parser.add_argument(
      "--coverage_penalty_weight",
      type=float,
      default=0.1,
      help="Coverage penalty for beam search.")

  # Job info
  parser.add_argument("--num_workers", type=int, default=1,
                      help="Number of workers (inference only).")

  parser.add_argument("--amp", action='store_true',
                      help="use amp for training and inference")
  parser.add_argument("--use_fastmath", type="bool", default=False,
                      help="use_fastmath for training and inference")
  parser.add_argument("--use_fp16", type="bool", default=False,
                      help="use_fp16 for training and inference")
  parser.add_argument(
      "--fp16_loss_scale",
      type=float,
      default=128,
      help="If fp16 is enabled, the loss is multiplied by this amount "
      "right before gradients are computed, then each gradient "
      "is divided by this amount. Mathematically, this has no "
      "effect, but it helps avoid fp16 underflow. Set to 1 to "
      "effectively disable.")
  parser.add_argument(
      "--enable_auto_loss_scale",
      type="bool",
      default=True,
      help="If True and use_fp16 is True, automatically adjust the "
      "loss scale during training.")
  parser.add_argument(
      "--fp16_inc_loss_scale_every_n",
      type=int,
      default=128,
      help="If fp16 is enabled and enable_auto_loss_scale is "
      "True, increase the loss scale every n steps.")
  parser.add_argument(
      "--check_tower_loss_numerics",
      type="bool",
      default=False,  # Set to false for xla.compile()
      help="whether to check tower loss numerics")
  parser.add_argument(
      "--use_fp32_batch_matmul",
      type="bool",
      default=False,
      help="Whether to use fp32 batch matmul")

  # Performance
  # XLA
  parser.add_argument(
      "--force_inputs_padding",
      type="bool",
      default=False,
      help="Force padding input batch to src_max_len and tgt_max_len")
  parser.add_argument(
      "--use_xla",
      type="bool",
      default=False,
      help="Use xla to compile a few selected locations, mostly Defuns.")
  parser.add_argument(
      "--xla_compile",
      type="bool",
      default=False,
      help="Use xla.compile() for each tower's fwd and bak pass.")
  parser.add_argument(
      "--use_autojit_xla",
      type="bool",
      default=False,
      help="Use auto jit xla.")
  # GPU knobs
  parser.add_argument(
      "--use_pintohost_optimizer",
      type="bool",
      default=False,
      help="whether to use PinToHost optimizer")
  parser.add_argument(
      "--use_cudnn_lstm",
      type="bool",
      default=False,
      help="whether to use cudnn_lstm for encoder, non residual layers")
  parser.add_argument(
      "--use_loose_bidi_cudnn_lstm",
      type="bool",
      default=False,
      help="whether to use loose bidi cudnn_lstm")
  parser.add_argument(
      "--use_fused_lstm",
      type="bool",
      default=True,
      help="whether to use fused lstm and variant. If enabled, training will "
      "use LSTMBlockFusedCell, infer will use LSTMBlockCell when appropriate.")
  parser.add_argument(
      "--use_fused_lstm_dec",
      type="bool",
      default=False,
      help="whether to use fused lstm for decoder (training only).")
  parser.add_argument(
      "--gpu_indices",
      type=str,
      default="",
      help="Indices of worker GPUs in ring order")
  # Graph knobs
  parser.add_argument("--parallel_iterations", type=int, default=10,
                      help="number of parallel iterations in dynamic_rnn")
  parser.add_argument("--use_dist_strategy", type="bool", default=False,
                      help="whether to use distribution strategy")
  parser.add_argument(
      "--hierarchical_copy",
      type="bool",
      default=False,
      help="Use hierarchical copies. Currently only optimized for "
      "use on a DGX-1 with 8 GPUs and may perform poorly on "
      "other hardware. Requires --num_gpus > 1, and only "
      "recommended when --num_gpus=8")
  parser.add_argument(
      "--network_topology",
      type=constants.NetworkTopology,
      default=constants.NetworkTopology.DGX1,
      choices=list(constants.NetworkTopology))
  parser.add_argument(
      "--use_block_lstm",
      type="bool",
      default=False,
      help="whether to use block lstm")
  parser.add_argument(
      "--use_defun",
      type="bool",
      default=False,
      help="whether to use Defun")

  # Gradient tricks
  parser.add_argument(
      "--gradient_repacking",
      type=int,
      default=0,
      help="Use gradient repacking. It"
      "currently only works with replicated mode. At the end of"
      "of each step, it repacks the gradients for more efficient"
      "cross-device transportation. A non-zero value specifies"
      "the number of split packs that will be formed.")
  parser.add_argument(
      "--compact_gradient_transfer",
      type="bool",
      default=True,
      help="Compact gradient as much as possible for cross-device transfer and "
      "aggregation.")
  parser.add_argument(
      "--all_reduce_spec",
      type=str,
      default="nccl",
      help="A specification of the all_reduce algorithm to be used "
      "for reducing gradients.  For more details, see "
      "parse_all_reduce_spec in variable_mgr.py.  An "
      "all_reduce_spec has BNF form:\n"
      "int ::= positive whole number\n"
      "g_int ::= int[KkMGT]?\n"
      "alg_spec ::= alg | alg#int\n"
      "range_spec ::= alg_spec | alg_spec/alg_spec\n"
      "spec ::= range_spec | range_spec:g_int:range_spec\n"
      "NOTE: not all syntactically correct constructs are "
      "supported.\n\n"
      "Examples:\n "
      "\"xring\" == use one global ring reduction for all "
      "tensors\n"
      "\"pscpu\" == use CPU at worker 0 to reduce all tensors\n"
      "\"nccl\" == use NCCL to locally reduce all tensors.  "
      "Limited to 1 worker.\n"
      "\"nccl/xring\" == locally (to one worker) reduce values "
      "using NCCL then ring reduce across workers.\n"
      "\"pscpu:32k:xring\" == use pscpu algorithm for tensors of "
      "size up to 32kB, then xring for larger tensors.")
  parser.add_argument(
      "--agg_small_grads_max_bytes",
      type=int,
      default=0,
      help="If > 0, try to aggregate tensors of less than this "
      "number of bytes prior to all-reduce.")
  parser.add_argument(
      "--agg_small_grads_max_group",
      type=int,
      default=10,
      help="When aggregating small tensors for all-reduce do not "
      "aggregate more than this many into one new tensor.")
  parser.add_argument(
      "--allreduce_merge_scope",
      type=int,
      default=1,
      help="Establish a name scope around this many "
      "gradients prior to creating the all-reduce operations. "
      "It may affect the ability of the backend to merge "
      "parallel ops.")
  # Other knobs
  parser.add_argument(
      "--local_parameter_device",
      type=str,
      default="gpu",
      help="Device to use as parameter server: cpu or gpu. For "
      "distributed training, it can affect where caching of "
      "variables happens.")
  parser.add_argument(
      "--use_resource_vars",
      type="bool",
      default=False,
      help="Use resource variables instead of normal variables. "
      "Resource variables are slower, but this option is useful "
      "for debugging their performance.")


  parser.add_argument("--debug", type="bool", default=False,
                      help="Debug train and eval")
  parser.add_argument(
      "--debug_num_train_steps", type=int, default=None, help="Num steps to train.")
  parser.add_argument("--show_metrics", type="bool", default=True,
                      help="whether to show detailed metrics")
  parser.add_argument("--clip_grads", type="bool", default=True,
                      help="whether to clip gradients")
  parser.add_argument("--profile", type="bool", default=False,
                      help="If generate profile")
  parser.add_argument("--profile_save_steps", type=int, default=10,
                      help="Save timeline every N steps.")

  parser.add_argument("--use_dynamic_rnn", type="bool", default=True)
  parser.add_argument("--use_synthetic_data", type="bool", default=False)
  parser.add_argument(
      "--mode", type=str, default="train_and_eval",
      choices=("train_and_eval", "infer", "translate"))


def create_hparams(flags):
  """Create training hparams."""
  return tf.contrib.training.HParams(
      # Data
      src=flags.src,
      tgt=flags.tgt,
      train_prefix=os.path.join(flags.data_dir, flags.train_prefix),
      test_prefix=os.path.join(flags.data_dir, flags.test_prefix),
      translate_file=flags.translate_file,
      vocab_prefix=os.path.join(flags.data_dir, flags.vocab_prefix),
      embed_prefix=flags.embed_prefix,
      output_dir=flags.output_dir,

      # Networks
      num_units=flags.num_units,
      num_encoder_layers=(flags.num_encoder_layers or flags.num_layers),
      num_decoder_layers=(flags.num_decoder_layers or flags.num_layers),
      dropout=flags.dropout,
      unit_type=flags.unit_type,
      encoder_type=flags.encoder_type,
      residual=flags.residual,
      time_major=flags.time_major,
      num_embeddings_partitions=flags.num_embeddings_partitions,

      # Attention mechanisms
      attention=flags.attention,
      attention_architecture=flags.attention_architecture,
      output_attention=flags.output_attention,
      pass_hidden_state=flags.pass_hidden_state,

      # Train
      optimizer=flags.optimizer,
      max_train_epochs=flags.max_train_epochs,
      target_bleu=flags.target_bleu,
      label_smoothing=flags.label_smoothing,
      batch_size=flags.batch_size,
      init_op=flags.init_op,
      init_weight=flags.init_weight,
      max_gradient_norm=flags.max_gradient_norm,
      learning_rate=flags.learning_rate,
      warmup_steps=flags.warmup_steps,
      warmup_scheme=flags.warmup_scheme,
      decay_scheme=flags.decay_scheme,
      colocate_gradients_with_ops=flags.colocate_gradients_with_ops,

      # Data constraints
      num_buckets=flags.num_buckets,
      src_max_len=flags.src_max_len,
      tgt_max_len=flags.tgt_max_len,

      # Inference
      src_max_len_infer=flags.src_max_len_infer,
      tgt_max_len_infer=flags.tgt_max_len_infer,
      ckpt=flags.ckpt,
      infer_batch_size=flags.infer_batch_size,
      detokenizer_file=flags.detokenizer_file if flags.detokenizer_file is not None \
                       else os.path.join(flags.data_dir, 'mosesdecoder/scripts/tokenizer/detokenizer.perl'),
      tokenizer_file=flags.tokenizer_file if flags.tokenizer_file is not None \
                       else os.path.join(flags.data_dir, 'mosesdecoder/scripts/tokenizer/tokenizer.perl'),

      # Advanced inference arguments
      infer_mode=flags.infer_mode,
      beam_width=flags.beam_width,
      length_penalty_weight=flags.length_penalty_weight,
      coverage_penalty_weight=flags.coverage_penalty_weight,

      # Vocab
      sos=flags.sos if flags.sos else vocab_utils.SOS,
      eos=flags.eos if flags.eos else vocab_utils.EOS,
      subword_option=flags.subword_option,
      check_special_token=flags.check_special_token,
      use_char_encode=flags.use_char_encode,

      # Misc
      forget_bias=flags.forget_bias,
      num_gpus=flags.num_gpus,
      save_checkpoints_steps=flags.save_checkpoints_steps,
      log_step_count_steps=flags.log_step_count_steps,
      epoch_step=0,  # record where we were within an epoch.
      share_vocab=flags.share_vocab,
      random_seed=flags.random_seed,
      language_model=flags.language_model,

      amp=flags.amp,
      use_fastmath=flags.use_fastmath,
      use_fp16=flags.use_fp16,
      fp16_loss_scale=flags.fp16_loss_scale,
      enable_auto_loss_scale=flags.enable_auto_loss_scale,
      fp16_inc_loss_scale_every_n=flags.fp16_inc_loss_scale_every_n,
      check_tower_loss_numerics=flags.check_tower_loss_numerics,
      use_fp32_batch_matmul=flags.use_fp32_batch_matmul,

      # Performance
      # GPU knbs
      force_inputs_padding=flags.force_inputs_padding,
      use_xla=flags.use_xla,
      xla_compile=flags.xla_compile,
      use_autojit_xla=flags.use_autojit_xla,
      use_pintohost_optimizer=flags.use_pintohost_optimizer,
      use_cudnn_lstm=flags.use_cudnn_lstm,
      use_loose_bidi_cudnn_lstm=flags.use_loose_bidi_cudnn_lstm,
      use_fused_lstm=flags.use_fused_lstm,
      use_fused_lstm_dec=flags.use_fused_lstm_dec,
      gpu_indices=flags.gpu_indices,
      # Graph knobs
      parallel_iterations=flags.parallel_iterations,
      use_dynamic_rnn=flags.use_dynamic_rnn,
      use_dist_strategy=flags.use_dist_strategy,
      hierarchical_copy=flags.hierarchical_copy,
      network_topology=flags.network_topology,
      use_block_lstm=flags.use_block_lstm,
      # Grad tricks
      gradient_repacking=flags.gradient_repacking,
      compact_gradient_transfer=flags.compact_gradient_transfer,
      all_reduce_spec=flags.all_reduce_spec,
      agg_small_grads_max_bytes=flags.agg_small_grads_max_bytes,
      agg_small_grads_max_group=flags.agg_small_grads_max_group,
      allreduce_merge_scope=flags.allreduce_merge_scope,
      # Other knobs
      local_parameter_device=("cpu" if flags.num_gpus ==0
                              else flags.local_parameter_device),
      use_resource_vars=flags.use_resource_vars,

      debug=flags.debug,
      debug_num_train_steps=flags.debug_num_train_steps,
      clip_grads=flags.clip_grads,
      profile=flags.profile,
      profile_save_steps=flags.profile_save_steps,
      show_metrics=flags.show_metrics,

      use_synthetic_data=flags.use_synthetic_data,
      mode=flags.mode,
  )


def _add_argument(hparams, key, value, update=True):
  """Add an argument to hparams; if exists, change the value if update==True."""
  if hasattr(hparams, key):
    if update:
      setattr(hparams, key, value)
  else:
    hparams.add_hparam(key, value)


def extend_hparams(hparams):
  """Add new arguments to hparams."""
  # Sanity checks
  if hparams.encoder_type == "bi" and hparams.num_encoder_layers % 2 != 0:
    raise ValueError("For bi, num_encoder_layers %d should be even" %
                     hparams.num_encoder_layers)
  if (hparams.attention_architecture in ["gnmt"] and
      hparams.num_encoder_layers < 2):
    raise ValueError("For gnmt attention architecture, "
                     "num_encoder_layers %d should be >= 2" %
                     hparams.num_encoder_layers)
  if hparams.subword_option and hparams.subword_option not in ["spm", "bpe"]:
    raise ValueError("subword option must be either spm, or bpe")
  if hparams.infer_mode == "beam_search" and hparams.beam_width <= 0:
    raise ValueError("beam_width must greater than 0 when using beam_search"
                     "decoder.")
  if hparams.mode == "translate" and not hparams.translate_file:
    raise ValueError("--translate_file flag must be specified in translate mode")

  # Different number of encoder / decoder layers
  assert hparams.num_encoder_layers and hparams.num_decoder_layers
  if hparams.num_encoder_layers != hparams.num_decoder_layers:
    hparams.pass_hidden_state = False
    utils.print_out("Num encoder layer %d is different from num decoder layer"
                    " %d, so set pass_hidden_state to False" % (
                        hparams.num_encoder_layers,
                        hparams.num_decoder_layers))

  # Set residual layers
  num_encoder_residual_layers = 0
  num_decoder_residual_layers = 0
  if hparams.residual:
    if hparams.num_encoder_layers > 1:
      num_encoder_residual_layers = hparams.num_encoder_layers - 1
    if hparams.num_decoder_layers > 1:
      num_decoder_residual_layers = hparams.num_decoder_layers - 1

    if hparams.encoder_type == "gnmt":
      # The first unidirectional layer (after the bi-directional layer) in
      # the GNMT encoder can't have residual connection due to the input is
      # the concatenation of fw_cell and bw_cell's outputs.
      num_encoder_residual_layers = hparams.num_encoder_layers - 2

      # Compatible for GNMT models
      if hparams.num_encoder_layers == hparams.num_decoder_layers:
        num_decoder_residual_layers = num_encoder_residual_layers
  _add_argument(hparams, "num_encoder_residual_layers",
                num_encoder_residual_layers)
  _add_argument(hparams, "num_decoder_residual_layers",
                num_decoder_residual_layers)

  # Language modeling
  if hparams.language_model:
    hparams.attention = ""
    hparams.attention_architecture = ""
    hparams.pass_hidden_state = False
    hparams.share_vocab = True
    hparams.src = hparams.tgt
    utils.print_out("For language modeling, we turn off attention and "
                    "pass_hidden_state; turn on share_vocab; set src to tgt.")

  ## Vocab
  # Get vocab file names first
  if hparams.vocab_prefix:
    src_vocab_file = hparams.vocab_prefix + "." + hparams.src
    tgt_vocab_file = hparams.vocab_prefix + "." + hparams.tgt
  else:
    raise ValueError("hparams.vocab_prefix must be provided.")

  # Source vocab
  src_vocab_size, src_vocab_file = vocab_utils.check_vocab(
      src_vocab_file,
      hparams.output_dir,
      check_special_token=hparams.check_special_token,
      sos=hparams.sos,
      eos=hparams.eos,
      unk=vocab_utils.UNK,
      pad_vocab=True)

  # Target vocab
  if hparams.share_vocab:
    utils.print_out("  using source vocab for target")
    tgt_vocab_file = src_vocab_file
    tgt_vocab_size = src_vocab_size
  else:
    tgt_vocab_size, tgt_vocab_file = vocab_utils.check_vocab(
        tgt_vocab_file,
        hparams.output_dir,
        check_special_token=hparams.check_special_token,
        sos=hparams.sos,
        eos=hparams.eos,
        unk=vocab_utils.UNK)
  _add_argument(hparams, "src_vocab_size", src_vocab_size)
  _add_argument(hparams, "tgt_vocab_size", tgt_vocab_size)
  _add_argument(hparams, "src_vocab_file", src_vocab_file)
  _add_argument(hparams, "tgt_vocab_file", tgt_vocab_file)

  # Num embedding partitions
  _add_argument(
      hparams, "num_enc_emb_partitions", hparams.num_embeddings_partitions)
  _add_argument(
      hparams, "num_dec_emb_partitions", hparams.num_embeddings_partitions)

  # Pretrained Embeddings
  _add_argument(hparams, "src_embed_file", "")
  _add_argument(hparams, "tgt_embed_file", "")
  if hparams.embed_prefix:
    src_embed_file = hparams.embed_prefix + "." + hparams.src
    tgt_embed_file = hparams.embed_prefix + "." + hparams.tgt

    if tf.gfile.Exists(src_embed_file):
      utils.print_out("  src_embed_file %s exist" % src_embed_file)
      hparams.src_embed_file = src_embed_file

      utils.print_out(
          "For pretrained embeddings, set num_enc_emb_partitions to 1")
      hparams.num_enc_emb_partitions = 1
    else:
      utils.print_out("  src_embed_file %s doesn't exist" % src_embed_file)

    if tf.gfile.Exists(tgt_embed_file):
      utils.print_out("  tgt_embed_file %s exist" % tgt_embed_file)
      hparams.tgt_embed_file = tgt_embed_file

      utils.print_out(
          "For pretrained embeddings, set num_dec_emb_partitions to 1")
      hparams.num_dec_emb_partitions = 1
    else:
      utils.print_out("  tgt_embed_file %s doesn't exist" % tgt_embed_file)

  # Evaluation
  metric = "bleu"
  best_metric_dir = os.path.join(hparams.output_dir, "best_" + metric)
  tf.gfile.MakeDirs(best_metric_dir)
  _add_argument(hparams, "best_" + metric, 0, update=False)
  _add_argument(hparams, "best_" + metric + "_dir", best_metric_dir)

  return hparams


def create_or_load_hparams(default_hparams, hparams_path):
  """Create hparams or load hparams from output_dir."""
  hparams = utils.maybe_parse_standard_hparams(default_hparams, hparams_path)
  hparams = extend_hparams(hparams)
  # Print HParams
  utils.print_hparams(hparams)
  return hparams


def run_main(flags, default_hparams, estimator_fn):
  """Run main."""
  # Random
  random_seed = flags.random_seed
  if random_seed is not None and random_seed > 0:
    utils.print_out("# Set random seed to %d" % random_seed)
    random.seed(random_seed)
    np.random.seed(random_seed)
    tf.set_random_seed(random_seed)

  # Model output directory
  output_dir = flags.output_dir
  if output_dir and not tf.gfile.Exists(output_dir):
    utils.print_out("# Creating output directory %s ..." % output_dir)
    tf.gfile.MakeDirs(output_dir)

  # Load hparams.
  hparams = create_or_load_hparams(default_hparams, flags.hparams_path)

  # Train or Evaluation
  estimator_fn(hparams)
  return hparams

def tokenize(hparams, file, tokenized_file):
  utils.print_out("tokenizing {} -> {}".format(file, tokenized_file))
  with open(file, 'rb') as input_file:
    with open(tokenized_file, 'wb') as output_file:
      subprocess.run([hparams.tokenizer_file, '-l', hparams.src], stdin=input_file, stdout=output_file)

def detokenize(hparams, file, detokenized_file):
  utils.print_out("detokenizing {} -> {}".format(file, detokenized_file))
  with open(file, 'rb') as input_file:
    with open(detokenized_file, 'wb') as output_file:
      subprocess.run([hparams.detokenizer_file, '-l', hparams.tgt], stdin=input_file, stdout=output_file)

def main(unused_argv):
  experiment_start = time.time()

  tf.logging.set_verbosity(tf.logging.INFO)

  if FLAGS.use_fp16 and FLAGS.use_dist_strategy:
    raise ValueError("use_fp16 and use_dist_strategy aren't compatible")

  if FLAGS.use_fp16 + FLAGS.amp + FLAGS.use_fastmath > 1:
    raise ValueError("Only one of use_fp16, amp, use_fastmath can be set")

  if FLAGS.amp:
    utils.print_out('Enabling TF-AMP')

    os.environ['TF_ENABLE_AUTO_MIXED_PRECISION'] = '1'

  if FLAGS.use_fastmath:
    utils.print_out('Enabling FastMath')

    os.environ["TF_ENABLE_CUBLAS_TENSOR_OP_MATH_FP32"] = '1'
    os.environ["TF_ENABLE_CUDNN_TENSOR_OP_MATH_FP32"] = '1'
    os.environ["TF_ENABLE_CUDNN_RNN_TENSOR_OP_MATH_FP32"] = '1'

  # Set up hacky envvars.
  # Hack that affects Defun in attention_wrapper.py
  active_xla_option_nums = np.sum([FLAGS.use_xla, FLAGS.use_autojit_xla,
                                   FLAGS.xla_compile])
  if active_xla_option_nums > 1:
    raise ValueError(
        "Only one of use_xla, xla_compile, use_autojit_xla can be set")

  os.environ["use_xla"] = str(FLAGS.use_xla).lower()
  if FLAGS.use_xla:
    os.environ["use_defun"] = str(True).lower()
  else:
    os.environ["use_defun"] = str(FLAGS.use_defun).lower()
  utils.print_out("use_defun is %s for attention" % os.environ["use_defun"])

  # TODO(jamesqin): retire this config after Cuda9.1
  os.environ["use_fp32_batch_matmul"] = ("true" if FLAGS.use_fp32_batch_matmul
                                         else "false")
  os.environ["xla_compile"] = "true" if FLAGS.xla_compile else "false"
  os.environ["force_inputs_padding"] = (
      "true" if FLAGS.force_inputs_padding else "false")

  if FLAGS.mode == "train":
    utils.print_out("Running training mode.")
    default_hparams = create_hparams(FLAGS)
    run_main(FLAGS, default_hparams, estimator.train_fn)
  elif FLAGS.mode == "infer" or FLAGS.mode == "translate":
    if FLAGS.mode == "infer":
        utils.print_out("Running inference mode.")
        translate_mode = False
    else:
        utils.print_out("Running translate mode on file {}.".format(FLAGS.translate_file))
        translate_mode = True

    # Random
    random_seed = FLAGS.random_seed
    if random_seed is not None and random_seed > 0:
      utils.print_out("# Set random seed to %d" % random_seed)
      random.seed(random_seed)
      np.random.seed(random_seed)
      tf.set_random_seed(random_seed)

    # Model output directory
    output_dir = FLAGS.output_dir
    if output_dir and not tf.gfile.Exists(output_dir):
      utils.print_out("# Creating output directory %s ..." % output_dir)
      tf.gfile.MakeDirs(output_dir)

    dllogger.init(backends=[
        dllogger.StdOutBackend(dllogger.Verbosity.DEFAULT),
        dllogger.JSONStreamBackend(dllogger.Verbosity.VERBOSE, os.path.join(FLAGS.output_dir, FLAGS.mode + '-report.json')),
    ])
    dllogger.log('PARAMETER', vars(FLAGS))

    # Load hparams.
    default_hparams = create_hparams(FLAGS)
    default_hparams.num_buckets = 1
    # The estimator model_fn is written in a way allowing train hparams to be
    # passed in infer mode.
    hparams = create_or_load_hparams(default_hparams, FLAGS.hparams_path)
    utils.print_out("infer_hparams:")
    utils.print_hparams(hparams)

    if translate_mode:
      tokenize(hparams, hparams.translate_file, hparams.translate_file + ".tok")

    eval_sentences, eval_src_tokens, _ = iterator_utils.get_effective_epoch_size(hparams, train=False)

    # Run evaluation when there's a new checkpoint
    tf.logging.info("Starting to evaluate...")
    eval_start = time.time()
    _, (eval_speed, eval_latencies), eval_output_tokens = estimator.eval_fn(hparams, hparams.ckpt, only_translate=translate_mode)
    eval_end = time.time()
    eval_delta = eval_end - eval_start
    utils.print_out("eval time for ckpt: %.2f mins (%.2f sent/sec, %.2f tokens/sec)" %
                    (eval_delta / 60., eval_speed, eval_speed * (eval_src_tokens + eval_output_tokens) / eval_sentences), f=sys.stderr)
    logging_data = {
      'infer_speed_sent': eval_speed,
      'infer_speed_toks': eval_speed * (eval_src_tokens + eval_output_tokens) / eval_sentences,
    }
    for lat in sorted(eval_latencies):
      utils.print_out("eval latency_%s for ckpt: %.2f ms" % (lat, eval_latencies[lat] * 1000))
      logging_data['infer_latency_{}'.format(lat)] = eval_latencies[lat] * 1000

    dllogger.log((), logging_data)
    dllogger.flush()


    if translate_mode:
      detokenize(hparams, hparams.translate_file + ".trans.tok", hparams.translate_file + ".trans")

  else:
    assert FLAGS.mode == "train_and_eval"
    utils.print_out("Running train and eval mode.")

    # Random
    random_seed = FLAGS.random_seed
    if random_seed is not None and random_seed > 0:
      utils.print_out("# Set random seed to %d" % random_seed)
      random.seed(random_seed)
      np.random.seed(random_seed)
      tf.set_random_seed(random_seed)

    # Model output directory
    output_dir = FLAGS.output_dir
    if output_dir and not tf.gfile.Exists(output_dir):
      utils.print_out("# Creating output directory %s ..." % output_dir)
      tf.gfile.MakeDirs(output_dir)

    dllogger.init(backends=[
        dllogger.StdOutBackend(dllogger.Verbosity.DEFAULT),
        dllogger.JSONStreamBackend(dllogger.Verbosity.VERBOSE, os.path.join(FLAGS.output_dir, FLAGS.mode + '-report.json')),
    ])
    dllogger.log('PARAMETER', vars(FLAGS))
    dllogger.metadata("bleu", {"unit": None})
    dllogger.metadata("train_speed_sent", {"unit": "sequences/s"})
    dllogger.metadata("train_speed_toks", {"unit": "tokens/s"})

    # Load hparams.
    default_hparams = create_hparams(FLAGS)

    hparams = create_or_load_hparams(default_hparams, FLAGS.hparams_path)
    utils.print_out("training hparams:")
    utils.print_hparams(hparams)
    with tf.gfile.GFile(os.path.join(output_dir, "train_hparams.txt"), "w") as f:
      f.write(utils.serialize_hparams(hparams) + "\n")

    # The estimator model_fn is written in a way allowing train hparams to be
    # passed in infer mode.
    infer_hparams = tf.contrib.training.HParams(**hparams.values())
    infer_hparams.num_buckets = 1
    utils.print_out("infer_hparams:")
    utils.print_hparams(infer_hparams)
    with tf.gfile.GFile(os.path.join(output_dir, "infer_hparams.txt"), "w") as f:
      f.write(utils.serialize_hparams(infer_hparams) + "\n")

    epochs = 0
    should_stop = epochs >= FLAGS.max_train_epochs

    train_sentences, train_src_tokens, train_tgt_tokens = iterator_utils.get_effective_epoch_size(hparams)
    eval_sentences, eval_src_tokens, _ = iterator_utils.get_effective_epoch_size(hparams, train=False)

    while not should_stop:
      utils.print_out("Starting epoch %d" % epochs)
      try:
        train_start = time.time()
        train_speed, _ = estimator.train_fn(hparams)
      except tf.errors.OutOfRangeError:
        utils.print_out("training hits OutOfRangeError", f=sys.stderr)

      train_end = time.time()
      train_delta = train_end - train_start
      utils.print_out("training time for epoch %d: %.2f mins (%.2f sent/sec, %.2f tokens/sec)" %
                      (epochs + 1, train_delta / 60., train_speed, train_speed * (train_src_tokens + train_tgt_tokens) / train_sentences), f=sys.stderr)
      logging_data = {
        'train_speed_sent': train_speed,
        'train_speed_toks': train_speed * (train_src_tokens + train_tgt_tokens) / train_sentences,
      }

      # This is probably sub-optimal, doing eval per-epoch
      eval_start = time.time()
      bleu_score, (eval_speed, eval_latencies), eval_output_tokens = estimator.eval_fn(infer_hparams)
      eval_end = time.time()
      eval_delta = eval_end - eval_start
      utils.print_out("eval time for epoch %d: %.2f mins (%.2f sent/sec, %.2f tokens/sec)" %
                      (epochs + 1, eval_delta / 60., eval_speed, eval_speed * (eval_src_tokens + eval_output_tokens) / eval_sentences), f=sys.stderr)
      logging_data.update({
        'bleu': bleu_score,
        'infer_speed_sent': eval_speed,
        'infer_speed_toks': eval_speed * (eval_src_tokens + eval_output_tokens) / eval_sentences,
      })
      for lat in sorted(eval_latencies):
        utils.print_out("eval latency_%s for epoch %d: %.2f ms" % (lat, epochs + 1, eval_latencies[lat] * 1000))
        logging_data['eval_latency_{}'.format(lat)] = eval_latencies[lat] * 1000

      dllogger.log((epochs,), logging_data)
      dllogger.flush()


      if FLAGS.debug or (FLAGS.target_bleu is not None and bleu_score > FLAGS.target_bleu):
        should_stop = True
        utils.print_out(
            "Stop job since target bleu is reached at epoch %d ." % epochs,
            f=sys.stderr)

      epochs += 1
      if epochs >= FLAGS.max_train_epochs:
        should_stop = True
        utils.print_out("Stop job since max_train_epochs is reached.",
                        f=sys.stderr)

    dllogger.log((), logging_data)
    dllogger.flush()

  experiment_end = time.time()
  utils.print_out('Experiment took {} min'.format((experiment_end - experiment_start) / 60))


if __name__ == "__main__":
  nmt_parser = argparse.ArgumentParser()
  add_arguments(nmt_parser)
  FLAGS, unparsed = nmt_parser.parse_known_args()
  tf.app.run(main=main, argv=[sys.argv[0]] + unparsed)

e26faece7f568406d4f9c5d6435a31723eff3f81
<mask token>


def add_arguments(parser):
    """Build ArgumentParser."""
    parser.register('type', 'bool', lambda v: v.lower() == 'true')
    parser.add_argument('--num_units', type=int, default=1024, help=
        'Network size.')
    parser.add_argument('--num_layers', type=int, default=4, help=
        'Network depth.')
    parser.add_argument('--num_encoder_layers', type=int, default=None,
        help='Encoder depth, equal to num_layers if None.')
    parser.add_argument('--num_decoder_layers', type=int, default=None,
        help='Decoder depth, equal to num_layers if None.')
    parser.add_argument('--encoder_type', type=str, default='gnmt', help=
        """      uni | bi | gnmt.
      For bi, we build num_encoder_layers/2 bi-directional layers.
      For gnmt, we build 1 bi-directional layer, and (num_encoder_layers - 1)
        uni-directional layers.      """
        )
    parser.add_argument('--residual', type='bool', nargs='?', const=True,
        default=True, help='Whether to add residual connections.')
    parser.add_argument('--time_major', type='bool', nargs='?', const=True,
        default=True, help='Whether to use time-major mode for dynamic RNN.')
    parser.add_argument('--num_embeddings_partitions', type=int, default=0,
        help='Number of partitions for embedding vars.')
    parser.add_argument('--attention', type=str, default='normed_bahdanau',
        help=
        """      luong | scaled_luong | bahdanau | normed_bahdanau or set to "" for no
      attention      """
        )
    parser.add_argument('--attention_architecture', type=str, default=
        'gnmt_v2', help=
        """      standard | gnmt | gnmt_v2.
      standard: use top layer to compute attention.
      gnmt: GNMT style of computing attention, use previous bottom layer to
          compute attention.
      gnmt_v2: similar to gnmt, but use current bottom layer to compute
          attention.      """
        )
    parser.add_argument('--output_attention', type='bool', nargs='?', const
        =True, default=True, help=
        """      Only used in standard attention_architecture. Whether use attention as
      the cell output at each timestep.
      .      """
        )
    parser.add_argument('--pass_hidden_state', type='bool', nargs='?',
        const=True, default=True, help=
        """      Whether to pass encoder's hidden state to decoder when using an attention
      based model.      """
        )
    parser.add_argument('--optimizer', type=str, default='adam', help=
        'sgd | adam')
    parser.add_argument('--learning_rate', type=float, default=0.0005, help
        ='Learning rate. Adam: 0.001 | 0.0001')
    parser.add_argument('--warmup_steps', type=int, default=200, help=
        'How many steps we inverse-decay learning.')
    parser.add_argument('--warmup_scheme', type=str, default='t2t', help=
        """      How to warmup learning rates. Options include:
        t2t: Tensor2Tensor's way, start with lr 100 times smaller, then
             exponentiate until the specified lr.      """
        )
    parser.add_argument('--decay_scheme', type=str, default='luong234',
        help=
        """      How we decay learning rate. Options include:
        luong234: after 2/3 num train steps, we start halving the learning rate
          for 4 times before finishing.
        luong5: after 1/2 num train steps, we start halving the learning rate
          for 5 times before finishing.        luong10: after 1/2 num train steps, we start halving the learning rate
          for 10 times before finishing.      """
        )
    parser.add_argument('--max_train_epochs', type=int, default=6, help=
        'Max number of epochs.')
    parser.add_argument('--target_bleu', type=float, default=None, help=
        'Target bleu.')
    parser.add_argument('--colocate_gradients_with_ops', type='bool', nargs
        ='?', const=True, default=True, help=
        'Whether try colocating gradients with corresponding op')
    parser.add_argument('--label_smoothing', type=float, default=0.1, help=
        'If nonzero, smooth the labels towards 1/num_classes.')
    parser.add_argument('--init_op', type=str, default='uniform', help=
        'uniform | glorot_normal | glorot_uniform')
    parser.add_argument('--init_weight', type=float, default=0.1, help=
        'for uniform init_op, initialize weights between [-this, this].')
    parser.add_argument('--src', type=str, default='en', help=
        'Source suffix, e.g., en.')
    parser.add_argument('--tgt', type=str, default='de', help=
        'Target suffix, e.g., de.')
    parser.add_argument('--data_dir', type=str, default='data/wmt16_de_en',
        help='Training/eval data directory.')
    parser.add_argument('--train_prefix', type=str, default=
        'train.tok.clean.bpe.32000', help=
        'Train prefix, expect files with src/tgt suffixes.')
    parser.add_argument('--test_prefix', type=str, default=
        'newstest2014.tok.bpe.32000', help=
        'Test prefix, expect files with src/tgt suffixes.')
    parser.add_argument('--translate_file', type=str, help=
        'File to translate, works only with translate mode')
    parser.add_argument('--output_dir', type=str, default='results', help=
        'Store log/model files.')
    parser.add_argument('--vocab_prefix', type=str, default=
        'vocab.bpe.32000', help=
        '      Vocab prefix, expect files with src/tgt suffixes.      ')
    parser.add_argument('--embed_prefix', type=str, default=None, help=
        """      Pretrained embedding prefix, expect files with src/tgt suffixes.
      The embedding files should be Glove formatted txt files.      """
        )
    parser.add_argument('--sos', type=str, default='<s>', help=
        'Start-of-sentence symbol.')
    parser.add_argument('--eos', type=str, default='</s>', help=
        'End-of-sentence symbol.')
    parser.add_argument('--share_vocab', type='bool', nargs='?', const=True,
        default=True, help=
        """      Whether to use the source vocab and embeddings for both source and
      target.      """
        )
    parser.add_argument('--check_special_token', type='bool', default=True,
        help=
        """                      Whether check special sos, eos, unk tokens exist in the
                      vocab files.                      """
        )
    parser.add_argument('--src_max_len', type=int, default=50, help=
        'Max length of src sequences during training (including EOS).')
    parser.add_argument('--tgt_max_len', type=int, default=50, help=
        'Max length of tgt sequences during training (including BOS).')
    parser.add_argument('--src_max_len_infer', type=int, default=None, help
        ='Max length of src sequences during inference (including EOS).')
    parser.add_argument('--tgt_max_len_infer', type=int, default=80, help=
        """      Max length of tgt sequences during inference (including BOS).  Also use to restrict the
      maximum decoding length.      """
        )
    parser.add_argument('--unit_type', type=str, default='lstm', help=
        'lstm | gru | layer_norm_lstm | nas')
    parser.add_argument('--forget_bias', type=float, default=0.0, help=
        'Forget bias for BasicLSTMCell.')
    parser.add_argument('--dropout', type=float, default=0.2, help=
        'Dropout rate (not keep_prob)')
    parser.add_argument('--max_gradient_norm', type=float, default=5.0,
        help='Clip gradients to this norm.')
    parser.add_argument('--batch_size', type=int, default=128, help=
        'Total batch size.')
    parser.add_argument('--num_buckets', type=int, default=5, help=
        'Put data into similar-length buckets (only for training).')
    parser.add_argument('--subword_option', type=str, default='bpe',
        choices=['', 'bpe', 'spm'], help=
        '                      Set to bpe or spm to activate subword desegmentation.                      '
        )
    parser.add_argument('--use_char_encode', type='bool', default=False,
        help=
        """                      Whether to split each word or bpe into character, and then
                      generate the word-level representation from the character
                      reprentation.
                      """
        )
    parser.add_argument('--save_checkpoints_steps', type=int, default=2000,
        help='save_checkpoints_steps')
    parser.add_argument('--log_step_count_steps', type=int, default=10,
        help=
        'The frequency, in number of global steps, that the global step and the loss will be logged during training'
        )
    parser.add_argument('--num_gpus', type=int, default=1, help=
        'Number of gpus in each worker.')
    parser.add_argument('--hparams_path', type=str, default=None, help=
        'Path to standard hparams json file that overrideshparams values from FLAGS.'
        )
    parser.add_argument('--random_seed', type=int, default=1, help=
        'Random seed (>0, set a specific seed).')
    parser.add_argument('--language_model', type='bool', nargs='?', const=
        True, default=False, help=
        'True to train a language model, ignoring encoder')
    parser.add_argument('--ckpt', type=str, default=None, help=
        'Checkpoint file to load a model for inference. (defaults to newest checkpoint)'
        )
    parser.add_argument('--infer_batch_size', type=int, default=128, help=
        'Batch size for inference mode.')
    parser.add_argument('--detokenizer_file', type=str, default=None, help=
        'Detokenizer script file. Default: DATA_DIR/mosesdecoder/scripts/tokenizer/detokenizer.perl'
        )
    parser.add_argument('--tokenizer_file', type=str, default=None, help=
        'Tokenizer script file. Default: DATA_DIR/mosesdecoder/scripts/tokenizer/tokenizer.perl'
        )
    parser.add_argument('--infer_mode', type=str, default='beam_search',
        choices=['greedy', 'beam_search'], help=
        'Which type of decoder to use during inference.')
    parser.add_argument('--beam_width', type=int, default=5, help=
        """      beam width when using beam search decoder. If 0, use standard
      decoder with greedy helper.      """
        )
    parser.add_argument('--length_penalty_weight', type=float, default=0.6,
        help='Length penalty for beam search.')
    parser.add_argument('--coverage_penalty_weight', type=float, default=
        0.1, help='Coverage penalty for beam search.')
    parser.add_argument('--num_workers', type=int, default=1, help=
        'Number of workers (inference only).')
    parser.add_argument('--amp', action='store_true', help=
        'use amp for training and inference')
    parser.add_argument('--use_fastmath', type='bool', default=False, help=
        'use_fastmath for training and inference')
    parser.add_argument('--use_fp16', type='bool', default=False, help=
        'use_fp16 for training and inference')
    parser.add_argument('--fp16_loss_scale', type=float, default=128, help=
        'If fp16 is enabled, the loss is multiplied by this amount right before gradients are computed, then each gradient is divided by this amount. Mathematically, this has no effect, but it helps avoid fp16 underflow. Set to 1 to effectively disable.'
        )
    parser.add_argument('--enable_auto_loss_scale', type='bool', default=
        True, help=
        'If True and use_fp16 is True, automatically adjust the loss scale during training.'
        )
    parser.add_argument('--fp16_inc_loss_scale_every_n', type=int, default=
        128, help=
        'If fp16 is enabled and enable_auto_loss_scale is True, increase the loss scale every n steps.'
        )
    parser.add_argument('--check_tower_loss_numerics', type='bool', default
        =False, help='whether to check tower loss numerics')
    parser.add_argument('--use_fp32_batch_matmul', type='bool', default=
        False, help='Whether to use fp32 batch matmul')
    parser.add_argument('--force_inputs_padding', type='bool', default=
        False, help='Force padding input batch to src_max_len and tgt_max_len')
    parser.add_argument('--use_xla', type='bool', default=False, help=
        'Use xla to compile a few selected locations, mostly Defuns.')
    parser.add_argument('--xla_compile', type='bool', default=False, help=
        "Use xla.compile() for each tower's fwd and bak pass.")
    parser.add_argument('--use_autojit_xla', type='bool', default=False,
        help='Use auto jit xla.')
    parser.add_argument('--use_pintohost_optimizer', type='bool', default=
        False, help='whether to use PinToHost optimizer')
    parser.add_argument('--use_cudnn_lstm', type='bool', default=False,
        help='whether to use cudnn_lstm for encoder, non residual layers')
    parser.add_argument('--use_loose_bidi_cudnn_lstm', type='bool', default
        =False, help='whether to use loose bidi cudnn_lstm')
    parser.add_argument('--use_fused_lstm', type='bool', default=True, help
        =
        'whether to use fused lstm and variant. If enabled, training will use LSTMBlockFusedCell, infer will use LSTMBlockCell when appropriate.'
        )
    parser.add_argument('--use_fused_lstm_dec', type='bool', default=False,
        help='whether to use fused lstm for decoder (training only).')
    parser.add_argument('--gpu_indices', type=str, default='', help=
        'Indices of worker GPUs in ring order')
    parser.add_argument('--parallel_iterations', type=int, default=10, help
        ='number of parallel iterations in dynamic_rnn')
    parser.add_argument('--use_dist_strategy', type='bool', default=False,
        help='whether to use distribution strategy')
    parser.add_argument('--hierarchical_copy', type='bool', default=False,
        help=
        'Use hierarchical copies. Currently only optimized for use on a DGX-1 with 8 GPUs and may perform poorly on other hardware. Requires --num_gpus > 1, and only recommended when --num_gpus=8'
        )
    parser.add_argument('--network_topology', type=constants.
        NetworkTopology, default=constants.NetworkTopology.DGX1, choices=
        list(constants.NetworkTopology))
    parser.add_argument('--use_block_lstm', type='bool', default=False,
        help='whether to use block lstm')
    parser.add_argument('--use_defun', type='bool', default=False, help=
        'whether to use Defun')
    parser.add_argument('--gradient_repacking', type=int, default=0, help=
        'Use gradient repacking. Itcurrently only works with replicated mode. At the end ofof each step, it repacks the gradients for more efficientcross-device transportation. A non-zero value specifiesthe number of split packs that will be formed.'
        )
    parser.add_argument('--compact_gradient_transfer', type='bool', default
        =True, help=
        'Compact gradient as much as possible for cross-device transfer and aggregation.'
        )
    parser.add_argument('--all_reduce_spec', type=str, default='nccl', help
        =
        """A specification of the all_reduce algorithm to be used for reducing gradients.  For more details, see parse_all_reduce_spec in variable_mgr.py.  An all_reduce_spec has BNF form:
int ::= positive whole number
g_int ::= int[KkMGT]?
alg_spec ::= alg | alg#int
range_spec ::= alg_spec | alg_spec/alg_spec
spec ::= range_spec | range_spec:g_int:range_spec
NOTE: not all syntactically correct constructs are supported.

Examples:
 "xring" == use one global ring reduction for all tensors
"pscpu" == use CPU at worker 0 to reduce all tensors
"nccl" == use NCCL to locally reduce all tensors.  Limited to 1 worker.
"nccl/xring" == locally (to one worker) reduce values using NCCL then ring reduce across workers.
"pscpu:32k:xring" == use pscpu algorithm for tensors of size up to 32kB, then xring for larger tensors."""
        )
    parser.add_argument('--agg_small_grads_max_bytes', type=int, default=0,
        help=
        'If > 0, try to aggregate tensors of less than this number of bytes prior to all-reduce.'
        )
    parser.add_argument('--agg_small_grads_max_group', type=int, default=10,
        help=
        'When aggregating small tensors for all-reduce do not aggregate more than this many into one new tensor.'
        )
    parser.add_argument('--allreduce_merge_scope', type=int, default=1,
        help=
        'Establish a name scope around this many gradients prior to creating the all-reduce operations. It may affect the ability of the backend to merge parallel ops.'
        )
    parser.add_argument('--local_parameter_device', type=str, default='gpu',
        help=
        'Device to use as parameter server: cpu or gpu. For distributed training, it can affect where caching of variables happens.'
        )
    parser.add_argument('--use_resource_vars', type='bool', default=False,
        help=
        'Use resource variables instead of normal variables. Resource variables are slower, but this option is useful for debugging their performance.'
        )
    parser.add_argument('--debug', type='bool', default=False, help=
        'Debug train and eval')
    parser.add_argument('--debug_num_train_steps', type=int, default=None,
        help='Num steps to train.')
    parser.add_argument('--show_metrics', type='bool', default=True, help=
        'whether to show detailed metrics')
    parser.add_argument('--clip_grads', type='bool', default=True, help=
        'whether to clip gradients')
    parser.add_argument('--profile', type='bool', default=False, help=
        'If generate profile')
    parser.add_argument('--profile_save_steps', type=int, default=10, help=
        'Save timeline every N steps.')
    parser.add_argument('--use_dynamic_rnn', type='bool', default=True)
    parser.add_argument('--use_synthetic_data', type='bool', default=False)
    parser.add_argument('--mode', type=str, default='train_and_eval',
        choices=('train_and_eval', 'infer', 'translate'))


def create_hparams(flags):
    """Create training hparams."""
    return tf.contrib.training.HParams(src=flags.src, tgt=flags.tgt,
        train_prefix=os.path.join(flags.data_dir, flags.train_prefix),
        test_prefix=os.path.join(flags.data_dir, flags.test_prefix),
        translate_file=flags.translate_file, vocab_prefix=os.path.join(
        flags.data_dir, flags.vocab_prefix), embed_prefix=flags.
        embed_prefix, output_dir=flags.output_dir, num_units=flags.
        num_units, num_encoder_layers=flags.num_encoder_layers or flags.
        num_layers, num_decoder_layers=flags.num_decoder_layers or flags.
        num_layers, dropout=flags.dropout, unit_type=flags.unit_type,
        encoder_type=flags.encoder_type, residual=flags.residual,
        time_major=flags.time_major, num_embeddings_partitions=flags.
        num_embeddings_partitions, attention=flags.attention,
        attention_architecture=flags.attention_architecture,
        output_attention=flags.output_attention, pass_hidden_state=flags.
        pass_hidden_state, optimizer=flags.optimizer, max_train_epochs=
        flags.max_train_epochs, target_bleu=flags.target_bleu,
        label_smoothing=flags.label_smoothing, batch_size=flags.batch_size,
        init_op=flags.init_op, init_weight=flags.init_weight,
        max_gradient_norm=flags.max_gradient_norm, learning_rate=flags.
        learning_rate, warmup_steps=flags.warmup_steps, warmup_scheme=flags
        .warmup_scheme, decay_scheme=flags.decay_scheme,
        colocate_gradients_with_ops=flags.colocate_gradients_with_ops,
        num_buckets=flags.num_buckets, src_max_len=flags.src_max_len,
        tgt_max_len=flags.tgt_max_len, src_max_len_infer=flags.
        src_max_len_infer, tgt_max_len_infer=flags.tgt_max_len_infer, ckpt=
        flags.ckpt, infer_batch_size=flags.infer_batch_size,
        detokenizer_file=flags.detokenizer_file if flags.detokenizer_file
         is not None else os.path.join(flags.data_dir,
        'mosesdecoder/scripts/tokenizer/detokenizer.perl'), tokenizer_file=
        flags.tokenizer_file if flags.tokenizer_file is not None else os.
        path.join(flags.data_dir,
        'mosesdecoder/scripts/tokenizer/tokenizer.perl'), infer_mode=flags.
        infer_mode, beam_width=flags.beam_width, length_penalty_weight=
        flags.length_penalty_weight, coverage_penalty_weight=flags.
        coverage_penalty_weight, sos=flags.sos if flags.sos else
        vocab_utils.SOS, eos=flags.eos if flags.eos else vocab_utils.EOS,
        subword_option=flags.subword_option, check_special_token=flags.
        check_special_token, use_char_encode=flags.use_char_encode,
        forget_bias=flags.forget_bias, num_gpus=flags.num_gpus,
        save_checkpoints_steps=flags.save_checkpoints_steps,
        log_step_count_steps=flags.log_step_count_steps, epoch_step=0,
        share_vocab=flags.share_vocab, random_seed=flags.random_seed,
        language_model=flags.language_model, amp=flags.amp, use_fastmath=
        flags.use_fastmath, use_fp16=flags.use_fp16, fp16_loss_scale=flags.
        fp16_loss_scale, enable_auto_loss_scale=flags.
        enable_auto_loss_scale, fp16_inc_loss_scale_every_n=flags.
        fp16_inc_loss_scale_every_n, check_tower_loss_numerics=flags.
        check_tower_loss_numerics, use_fp32_batch_matmul=flags.
        use_fp32_batch_matmul, force_inputs_padding=flags.
        force_inputs_padding, use_xla=flags.use_xla, xla_compile=flags.
        xla_compile, use_autojit_xla=flags.use_autojit_xla,
        use_pintohost_optimizer=flags.use_pintohost_optimizer,
        use_cudnn_lstm=flags.use_cudnn_lstm, use_loose_bidi_cudnn_lstm=
        flags.use_loose_bidi_cudnn_lstm, use_fused_lstm=flags.
        use_fused_lstm, use_fused_lstm_dec=flags.use_fused_lstm_dec,
        gpu_indices=flags.gpu_indices, parallel_iterations=flags.
        parallel_iterations, use_dynamic_rnn=flags.use_dynamic_rnn,
        use_dist_strategy=flags.use_dist_strategy, hierarchical_copy=flags.
        hierarchical_copy, network_topology=flags.network_topology,
        use_block_lstm=flags.use_block_lstm, gradient_repacking=flags.
        gradient_repacking, compact_gradient_transfer=flags.
        compact_gradient_transfer, all_reduce_spec=flags.all_reduce_spec,
        agg_small_grads_max_bytes=flags.agg_small_grads_max_bytes,
        agg_small_grads_max_group=flags.agg_small_grads_max_group,
        allreduce_merge_scope=flags.allreduce_merge_scope,
        local_parameter_device='cpu' if flags.num_gpus == 0 else flags.
        local_parameter_device, use_resource_vars=flags.use_resource_vars,
        debug=flags.debug, debug_num_train_steps=flags.
        debug_num_train_steps, clip_grads=flags.clip_grads, profile=flags.
        profile, profile_save_steps=flags.profile_save_steps, show_metrics=
        flags.show_metrics, use_synthetic_data=flags.use_synthetic_data,
        mode=flags.mode)


def _add_argument(hparams, key, value, update=True):
    """Add an argument to hparams; if exists, change the value if update==True."""
    if hasattr(hparams, key):
        if update:
            setattr(hparams, key, value)
    else:
        hparams.add_hparam(key, value)


<mask token>


def create_or_load_hparams(default_hparams, hparams_path):
    """Create hparams or load hparams from output_dir."""
    hparams = utils.maybe_parse_standard_hparams(default_hparams, hparams_path)
    hparams = extend_hparams(hparams)
    utils.print_hparams(hparams)
    return hparams


def run_main(flags, default_hparams, estimator_fn):
    """Run main."""
    random_seed = flags.random_seed
    if random_seed is not None and random_seed > 0:
        utils.print_out('# Set random seed to %d' % random_seed)
        random.seed(random_seed)
        np.random.seed(random_seed)
        tf.set_random_seed(random_seed)
    output_dir = flags.output_dir
    if output_dir and not tf.gfile.Exists(output_dir):
        utils.print_out('# Creating output directory %s ...' % output_dir)
        tf.gfile.MakeDirs(output_dir)
    hparams = create_or_load_hparams(default_hparams, flags.hparams_path)
    estimator_fn(hparams)
    return hparams


def tokenize(hparams, file, tokenized_file):
    utils.print_out('tokenizing {} -> {}'.format(file, tokenized_file))
    with open(file, 'rb') as input_file:
        with open(tokenized_file, 'wb') as output_file:
            subprocess.run([hparams.tokenizer_file, '-l', hparams.src],
                stdin=input_file, stdout=output_file)


def detokenize(hparams, file, detokenized_file):
    utils.print_out('detokenizing {} -> {}'.format(file, detokenized_file))
    with open(file, 'rb') as input_file:
        with open(detokenized_file, 'wb') as output_file:
            subprocess.run([hparams.detokenizer_file, '-l', hparams.tgt],
                stdin=input_file, stdout=output_file)


def main(unused_argv):
    experiment_start = time.time()
    tf.logging.set_verbosity(tf.logging.INFO)
    if FLAGS.use_fp16 and FLAGS.use_dist_strategy:
        raise ValueError("use_fp16 and use_dist_strategy aren't compatible")
    if FLAGS.use_fp16 + FLAGS.amp + FLAGS.use_fastmath > 1:
        raise ValueError('Only one of use_fp16, amp, use_fastmath can be set')
    if FLAGS.amp:
        utils.print_out('Enabling TF-AMP')
        os.environ['TF_ENABLE_AUTO_MIXED_PRECISION'] = '1'
    if FLAGS.use_fastmath:
        utils.print_out('Enabling FastMath')
        os.environ['TF_ENABLE_CUBLAS_TENSOR_OP_MATH_FP32'] = '1'
        os.environ['TF_ENABLE_CUDNN_TENSOR_OP_MATH_FP32'] = '1'
        os.environ['TF_ENABLE_CUDNN_RNN_TENSOR_OP_MATH_FP32'] = '1'
    active_xla_option_nums = np.sum([FLAGS.use_xla, FLAGS.use_autojit_xla,
        FLAGS.xla_compile])
    if active_xla_option_nums > 1:
        raise ValueError(
            'Only one of use_xla, xla_compile, use_autojit_xla can be set')
    os.environ['use_xla'] = str(FLAGS.use_xla).lower()
    if FLAGS.use_xla:
        os.environ['use_defun'] = str(True).lower()
    else:
        os.environ['use_defun'] = str(FLAGS.use_defun).lower()
    utils.print_out('use_defun is %s for attention' % os.environ['use_defun'])
    os.environ['use_fp32_batch_matmul'
        ] = 'true' if FLAGS.use_fp32_batch_matmul else 'false'
    os.environ['xla_compile'] = 'true' if FLAGS.xla_compile else 'false'
    os.environ['force_inputs_padding'
        ] = 'true' if FLAGS.force_inputs_padding else 'false'
    if FLAGS.mode == 'train':
        utils.print_out('Running training mode.')
        default_hparams = create_hparams(FLAGS)
        run_main(FLAGS, default_hparams, estimator.train_fn)
    elif FLAGS.mode == 'infer' or FLAGS.mode == 'translate':
        if FLAGS.mode == 'infer':
            utils.print_out('Running inference mode.')
            translate_mode = False
        else:
            utils.print_out('Running translate mode on file {}.'.format(
                FLAGS.translate_file))
            translate_mode = True
        random_seed = FLAGS.random_seed
        if random_seed is not None and random_seed > 0:
            utils.print_out('# Set random seed to %d' % random_seed)
            random.seed(random_seed)
            np.random.seed(random_seed)
            tf.set_random_seed(random_seed)
        output_dir = FLAGS.output_dir
        if output_dir and not tf.gfile.Exists(output_dir):
            utils.print_out('# Creating output directory %s ...' % output_dir)
            tf.gfile.MakeDirs(output_dir)
        dllogger.init(backends=[dllogger.StdOutBackend(dllogger.Verbosity.
            DEFAULT), dllogger.JSONStreamBackend(dllogger.Verbosity.VERBOSE,
            os.path.join(FLAGS.output_dir, FLAGS.mode + '-report.json'))])
        dllogger.log('PARAMETER', vars(FLAGS))
        default_hparams = create_hparams(FLAGS)
        default_hparams.num_buckets = 1
        hparams = create_or_load_hparams(default_hparams, FLAGS.hparams_path)
        utils.print_out('infer_hparams:')
        utils.print_hparams(hparams)
        if translate_mode:
            tokenize(hparams, hparams.translate_file, hparams.
                translate_file + '.tok')
        eval_sentences, eval_src_tokens, _ = (iterator_utils.
            get_effective_epoch_size(hparams, train=False))
        tf.logging.info('Starting to evaluate...')
        eval_start = time.time()
        _, (eval_speed, eval_latencies
            ), eval_output_tokens = estimator.eval_fn(hparams, hparams.ckpt,
            only_translate=translate_mode)
        eval_end = time.time()
        eval_delta = eval_end - eval_start
        utils.print_out(
            'eval time for ckpt: %.2f mins (%.2f sent/sec, %.2f tokens/sec)' %
            (eval_delta / 60.0, eval_speed, eval_speed * (eval_src_tokens +
            eval_output_tokens) / eval_sentences), f=sys.stderr)
        logging_data = {'infer_speed_sent': eval_speed, 'infer_speed_toks':
            eval_speed * (eval_src_tokens + eval_output_tokens) /
            eval_sentences}
        for lat in sorted(eval_latencies):
            utils.print_out('eval latency_%s for ckpt: %.2f ms' % (lat, 
                eval_latencies[lat] * 1000))
            logging_data['infer_latency_{}'.format(lat)] = eval_latencies[lat
                ] * 1000
        dllogger.log((), logging_data)
        dllogger.flush()
        if translate_mode:
            detokenize(hparams, hparams.translate_file + '.trans.tok', 
                hparams.translate_file + '.trans')
    else:
        assert FLAGS.mode == 'train_and_eval'
        utils.print_out('Running train and eval mode.')
        random_seed = FLAGS.random_seed
        if random_seed is not None and random_seed > 0:
            utils.print_out('# Set random seed to %d' % random_seed)
            random.seed(random_seed)
            np.random.seed(random_seed)
            tf.set_random_seed(random_seed)
        output_dir = FLAGS.output_dir
        if output_dir and not tf.gfile.Exists(output_dir):
            utils.print_out('# Creating output directory %s ...' % output_dir)
            tf.gfile.MakeDirs(output_dir)
        dllogger.init(backends=[dllogger.StdOutBackend(dllogger.Verbosity.
            DEFAULT), dllogger.JSONStreamBackend(dllogger.Verbosity.VERBOSE,
            os.path.join(FLAGS.output_dir, FLAGS.mode + '-report.json'))])
        dllogger.log('PARAMETER', vars(FLAGS))
        dllogger.metadata('bleu', {'unit': None})
        dllogger.metadata('train_speed_sent', {'unit': 'sequences/s'})
        dllogger.metadata('train_speed_toks', {'unit': 'tokens/s'})
        default_hparams = create_hparams(FLAGS)
        hparams = create_or_load_hparams(default_hparams, FLAGS.hparams_path)
        utils.print_out('training hparams:')
        utils.print_hparams(hparams)
        with tf.gfile.GFile(os.path.join(output_dir, 'train_hparams.txt'), 'w'
            ) as f:
            f.write(utils.serialize_hparams(hparams) + '\n')
        infer_hparams = tf.contrib.training.HParams(**hparams.values())
        infer_hparams.num_buckets = 1
        utils.print_out('infer_hparams:')
        utils.print_hparams(infer_hparams)
        with tf.gfile.GFile(os.path.join(output_dir, 'infer_hparams.txt'), 'w'
            ) as f:
            f.write(utils.serialize_hparams(infer_hparams) + '\n')
        epochs = 0
        should_stop = epochs >= FLAGS.max_train_epochs
        train_sentences, train_src_tokens, train_tgt_tokens = (iterator_utils
            .get_effective_epoch_size(hparams))
        eval_sentences, eval_src_tokens, _ = (iterator_utils.
            get_effective_epoch_size(hparams, train=False))
        while not should_stop:
            utils.print_out('Starting epoch %d' % epochs)
            try:
                train_start = time.time()
                train_speed, _ = estimator.train_fn(hparams)
            except tf.errors.OutOfRangeError:
                utils.print_out('training hits OutOfRangeError', f=sys.stderr)
            train_end = time.time()
            train_delta = train_end - train_start
            utils.print_out(
                'training time for epoch %d: %.2f mins (%.2f sent/sec, %.2f tokens/sec)'
                 % (epochs + 1, train_delta / 60.0, train_speed, 
                train_speed * (train_src_tokens + train_tgt_tokens) /
                train_sentences), f=sys.stderr)
            logging_data = {'train_speed_sent': train_speed,
                'train_speed_toks': train_speed * (train_src_tokens +
                train_tgt_tokens) / train_sentences}
            eval_start = time.time()
            bleu_score, (eval_speed, eval_latencies
                ), eval_output_tokens = estimator.eval_fn(infer_hparams)
            eval_end = time.time()
            eval_delta = eval_end - eval_start
            utils.print_out(
                'eval time for epoch %d: %.2f mins (%.2f sent/sec, %.2f tokens/sec)'
                 % (epochs + 1, eval_delta / 60.0, eval_speed, eval_speed *
                (eval_src_tokens + eval_output_tokens) / eval_sentences), f
                =sys.stderr)
            logging_data.update({'bleu': bleu_score, 'infer_speed_sent':
                eval_speed, 'infer_speed_toks': eval_speed * (
                eval_src_tokens + eval_output_tokens) / eval_sentences})
            for lat in sorted(eval_latencies):
                utils.print_out('eval latency_%s for epoch %d: %.2f ms' % (
                    lat, epochs + 1, eval_latencies[lat] * 1000))
                logging_data['eval_latency_{}'.format(lat)] = eval_latencies[
                    lat] * 1000
            dllogger.log((epochs,), logging_data)
            dllogger.flush()
            if (FLAGS.debug or FLAGS.target_bleu is not None and bleu_score >
                FLAGS.target_bleu):
                should_stop = True
                utils.print_out(
                    'Stop job since target bleu is reached at epoch %d .' %
                    epochs, f=sys.stderr)
            epochs += 1
            if epochs >= FLAGS.max_train_epochs:
                should_stop = True
                utils.print_out('Stop job since max_train_epochs is reached.',
                    f=sys.stderr)
        dllogger.log((), logging_data)
        dllogger.flush()
    experiment_end = time.time()
    utils.print_out('Experiment took {} min'.format((experiment_end -
        experiment_start) / 60))


<mask token>

<mask token>


def add_arguments(parser):
    """Build ArgumentParser."""
    parser.register('type', 'bool', lambda v: v.lower() == 'true')
    parser.add_argument('--num_units', type=int, default=1024, help=
        'Network size.')
    parser.add_argument('--num_layers', type=int, default=4, help=
        'Network depth.')
    parser.add_argument('--num_encoder_layers', type=int, default=None,
        help='Encoder depth, equal to num_layers if None.')
    parser.add_argument('--num_decoder_layers', type=int, default=None,
        help='Decoder depth, equal to num_layers if None.')
    parser.add_argument('--encoder_type', type=str, default='gnmt', help=
        """      uni | bi | gnmt.
      For bi, we build num_encoder_layers/2 bi-directional layers.
      For gnmt, we build 1 bi-directional layer, and (num_encoder_layers - 1)
        uni-directional layers.      """
        )
    parser.add_argument('--residual', type='bool', nargs='?', const=True,
        default=True, help='Whether to add residual connections.')
    parser.add_argument('--time_major', type='bool', nargs='?', const=True,
        default=True, help='Whether to use time-major mode for dynamic RNN.')
    parser.add_argument('--num_embeddings_partitions', type=int, default=0,
        help='Number of partitions for embedding vars.')
    parser.add_argument('--attention', type=str, default='normed_bahdanau',
        help=
        """      luong | scaled_luong | bahdanau | normed_bahdanau or set to "" for no
      attention      """
        )
    parser.add_argument('--attention_architecture', type=str, default=
        'gnmt_v2', help=
        """      standard | gnmt | gnmt_v2.
      standard: use top layer to compute attention.
      gnmt: GNMT style of computing attention, use previous bottom layer to
          compute attention.
      gnmt_v2: similar to gnmt, but use current bottom layer to compute
          attention.      """
        )
    parser.add_argument('--output_attention', type='bool', nargs='?', const
        =True, default=True, help=
        """      Only used in standard attention_architecture. Whether use attention as
      the cell output at each timestep.
      .      """
        )
    parser.add_argument('--pass_hidden_state', type='bool', nargs='?',
        const=True, default=True, help=
        """      Whether to pass encoder's hidden state to decoder when using an attention
      based model.      """
        )
    parser.add_argument('--optimizer', type=str, default='adam', help=
        'sgd | adam')
    parser.add_argument('--learning_rate', type=float, default=0.0005, help
        ='Learning rate. Adam: 0.001 | 0.0001')
    parser.add_argument('--warmup_steps', type=int, default=200, help=
        'How many steps we inverse-decay learning.')
    parser.add_argument('--warmup_scheme', type=str, default='t2t', help=
        """      How to warmup learning rates. Options include:
        t2t: Tensor2Tensor's way, start with lr 100 times smaller, then
             exponentiate until the specified lr.      """
        )
    parser.add_argument('--decay_scheme', type=str, default='luong234',
        help=
        """      How we decay learning rate. Options include:
        luong234: after 2/3 num train steps, we start halving the learning rate
          for 4 times before finishing.
        luong5: after 1/2 num train steps, we start halving the learning rate
          for 5 times before finishing.        luong10: after 1/2 num train steps, we start halving the learning rate
          for 10 times before finishing.      """
        )
    parser.add_argument('--max_train_epochs', type=int, default=6, help=
        'Max number of epochs.')
    parser.add_argument('--target_bleu', type=float, default=None, help=
        'Target bleu.')
    parser.add_argument('--colocate_gradients_with_ops', type='bool', nargs
        ='?', const=True, default=True, help=
        'Whether try colocating gradients with corresponding op')
    parser.add_argument('--label_smoothing', type=float, default=0.1, help=
        'If nonzero, smooth the labels towards 1/num_classes.')
    parser.add_argument('--init_op', type=str, default='uniform', help=
        'uniform | glorot_normal | glorot_uniform')
    parser.add_argument('--init_weight', type=float, default=0.1, help=
        'for uniform init_op, initialize weights between [-this, this].')
    parser.add_argument('--src', type=str, default='en', help=
        'Source suffix, e.g., en.')
    parser.add_argument('--tgt', type=str, default='de', help=
        'Target suffix, e.g., de.')
    parser.add_argument('--data_dir', type=str, default='data/wmt16_de_en',
        help='Training/eval data directory.')
    parser.add_argument('--train_prefix', type=str, default=
        'train.tok.clean.bpe.32000', help=
        'Train prefix, expect files with src/tgt suffixes.')
    parser.add_argument('--test_prefix', type=str, default=
        'newstest2014.tok.bpe.32000', help=
        'Test prefix, expect files with src/tgt suffixes.')
    parser.add_argument('--translate_file', type=str, help=
        'File to translate, works only with translate mode')
    parser.add_argument('--output_dir', type=str, default='results', help=
        'Store log/model files.')
    parser.add_argument('--vocab_prefix', type=str, default=
        'vocab.bpe.32000', help=
        '      Vocab prefix, expect files with src/tgt suffixes.      ')
    parser.add_argument('--embed_prefix', type=str, default=None, help=
        """      Pretrained embedding prefix, expect files with src/tgt suffixes.
      The embedding files should be Glove formatted txt files.      """
        )
    parser.add_argument('--sos', type=str, default='<s>', help=
        'Start-of-sentence symbol.')
    parser.add_argument('--eos', type=str, default='</s>', help=
        'End-of-sentence symbol.')
    parser.add_argument('--share_vocab', type='bool', nargs='?', const=True,
        default=True, help=
        """      Whether to use the source vocab and embeddings for both source and
      target.      """
        )
    parser.add_argument('--check_special_token', type='bool', default=True,
        help=
        """                      Whether check special sos, eos, unk tokens exist in the
                      vocab files.                      """
        )
    parser.add_argument('--src_max_len', type=int, default=50, help=
        'Max length of src sequences during training (including EOS).')
    parser.add_argument('--tgt_max_len', type=int, default=50, help=
        'Max length of tgt sequences during training (including BOS).')
    parser.add_argument('--src_max_len_infer', type=int, default=None, help
        ='Max length of src sequences during inference (including EOS).')
    parser.add_argument('--tgt_max_len_infer', type=int, default=80, help=
        """      Max length of tgt sequences during inference (including BOS).  Also use to restrict the
      maximum decoding length.      """
        )
    parser.add_argument('--unit_type', type=str, default='lstm', help=
        'lstm | gru | layer_norm_lstm | nas')
    parser.add_argument('--forget_bias', type=float, default=0.0, help=
        'Forget bias for BasicLSTMCell.')
    parser.add_argument('--dropout', type=float, default=0.2, help=
        'Dropout rate (not keep_prob)')
    parser.add_argument('--max_gradient_norm', type=float, default=5.0,
        help='Clip gradients to this norm.')
    parser.add_argument('--batch_size', type=int, default=128, help=
        'Total batch size.')
    parser.add_argument('--num_buckets', type=int, default=5, help=
        'Put data into similar-length buckets (only for training).')
    parser.add_argument('--subword_option', type=str, default='bpe',
        choices=['', 'bpe', 'spm'], help=
        '                      Set to bpe or spm to activate subword desegmentation.                      '
        )
    parser.add_argument('--use_char_encode', type='bool', default=False,
        help=
        """                      Whether to split each word or bpe into character, and then
                      generate the word-level representation from the character
                      reprentation.
                      """
        )
    parser.add_argument('--save_checkpoints_steps', type=int, default=2000,
        help='save_checkpoints_steps')
    parser.add_argument('--log_step_count_steps', type=int, default=10,
        help=
        'The frequency, in number of global steps, that the global step and the loss will be logged during training'
        )
    parser.add_argument('--num_gpus', type=int, default=1, help=
        'Number of gpus in each worker.')
    parser.add_argument('--hparams_path', type=str, default=None, help=
        'Path to standard hparams json file that overrideshparams values from FLAGS.'
        )
    parser.add_argument('--random_seed', type=int, default=1, help=
        'Random seed (>0, set a specific seed).')
    parser.add_argument('--language_model', type='bool', nargs='?', const=
        True, default=False, help=
        'True to train a language model, ignoring encoder')
    parser.add_argument('--ckpt', type=str, default=None, help=
        'Checkpoint file to load a model for inference. (defaults to newest checkpoint)'
        )
    parser.add_argument('--infer_batch_size', type=int, default=128, help=
        'Batch size for inference mode.')
    parser.add_argument('--detokenizer_file', type=str, default=None, help=
        'Detokenizer script file. Default: DATA_DIR/mosesdecoder/scripts/tokenizer/detokenizer.perl'
        )
    parser.add_argument('--tokenizer_file', type=str, default=None, help=
        'Tokenizer script file. Default: DATA_DIR/mosesdecoder/scripts/tokenizer/tokenizer.perl'
        )
    parser.add_argument('--infer_mode', type=str, default='beam_search',
        choices=['greedy', 'beam_search'], help=
        'Which type of decoder to use during inference.')
    parser.add_argument('--beam_width', type=int, default=5, help=
        """      beam width when using beam search decoder. If 0, use standard
      decoder with greedy helper.      """
        )
    parser.add_argument('--length_penalty_weight', type=float, default=0.6,
        help='Length penalty for beam search.')
    parser.add_argument('--coverage_penalty_weight', type=float, default=
        0.1, help='Coverage penalty for beam search.')
    parser.add_argument('--num_workers', type=int, default=1, help=
        'Number of workers (inference only).')
    parser.add_argument('--amp', action='store_true', help=
        'use amp for training and inference')
    parser.add_argument('--use_fastmath', type='bool', default=False, help=
        'use_fastmath for training and inference')
    parser.add_argument('--use_fp16', type='bool', default=False, help=
        'use_fp16 for training and inference')
    parser.add_argument('--fp16_loss_scale', type=float, default=128, help=
        'If fp16 is enabled, the loss is multiplied by this amount right before gradients are computed, then each gradient is divided by this amount. Mathematically, this has no effect, but it helps avoid fp16 underflow. Set to 1 to effectively disable.'
        )
    parser.add_argument('--enable_auto_loss_scale', type='bool', default=
        True, help=
        'If True and use_fp16 is True, automatically adjust the loss scale during training.'
        )
    parser.add_argument('--fp16_inc_loss_scale_every_n', type=int, default=
        128, help=
        'If fp16 is enabled and enable_auto_loss_scale is True, increase the loss scale every n steps.'
        )
    parser.add_argument('--check_tower_loss_numerics', type='bool', default
        =False, help='whether to check tower loss numerics')
    parser.add_argument('--use_fp32_batch_matmul', type='bool', default=
        False, help='Whether to use fp32 batch matmul')
    parser.add_argument('--force_inputs_padding', type='bool', default=
        False, help='Force padding input batch to src_max_len and tgt_max_len')
    parser.add_argument('--use_xla', type='bool', default=False, help=
        'Use xla to compile a few selected locations, mostly Defuns.')
    parser.add_argument('--xla_compile', type='bool', default=False, help=
        "Use xla.compile() for each tower's fwd and bak pass.")
    parser.add_argument('--use_autojit_xla', type='bool', default=False,
        help='Use auto jit xla.')
    parser.add_argument('--use_pintohost_optimizer', type='bool', default=
        False, help='whether to use PinToHost optimizer')
    parser.add_argument('--use_cudnn_lstm', type='bool', default=False,
        help='whether to use cudnn_lstm for encoder, non residual layers')
    parser.add_argument('--use_loose_bidi_cudnn_lstm', type='bool', default
        =False, help='whether to use loose bidi cudnn_lstm')
    parser.add_argument('--use_fused_lstm', type='bool', default=True, help
        =
        'whether to use fused lstm and variant. If enabled, training will use LSTMBlockFusedCell, infer will use LSTMBlockCell when appropriate.'
        )
    parser.add_argument('--use_fused_lstm_dec', type='bool', default=False,
        help='whether to use fused lstm for decoder (training only).')
    parser.add_argument('--gpu_indices', type=str, default='', help=
        'Indices of worker GPUs in ring order')
    parser.add_argument('--parallel_iterations', type=int, default=10, help
        ='number of parallel iterations in dynamic_rnn')
    parser.add_argument('--use_dist_strategy', type='bool', default=False,
        help='whether to use distribution strategy')
    parser.add_argument('--hierarchical_copy', type='bool', default=False,
        help=
        'Use hierarchical copies. Currently only optimized for use on a DGX-1 with 8 GPUs and may perform poorly on other hardware. Requires --num_gpus > 1, and only recommended when --num_gpus=8'
        )
    parser.add_argument('--network_topology', type=constants.
        NetworkTopology, default=constants.NetworkTopology.DGX1, choices=
        list(constants.NetworkTopology))
    parser.add_argument('--use_block_lstm', type='bool', default=False,
        help='whether to use block lstm')
    parser.add_argument('--use_defun', type='bool', default=False, help=
        'whether to use Defun')
    parser.add_argument('--gradient_repacking', type=int, default=0, help=
        'Use gradient repacking. Itcurrently only works with replicated mode. At the end ofof each step, it repacks the gradients for more efficientcross-device transportation. A non-zero value specifiesthe number of split packs that will be formed.'
        )
    parser.add_argument('--compact_gradient_transfer', type='bool', default
        =True, help=
        'Compact gradient as much as possible for cross-device transfer and aggregation.'
        )
    parser.add_argument('--all_reduce_spec', type=str, default='nccl', help
        =
        """A specification of the all_reduce algorithm to be used for reducing gradients.  For more details, see parse_all_reduce_spec in variable_mgr.py.  An all_reduce_spec has BNF form:
int ::= positive whole number
g_int ::= int[KkMGT]?
alg_spec ::= alg | alg#int
range_spec ::= alg_spec | alg_spec/alg_spec
spec ::= range_spec | range_spec:g_int:range_spec
NOTE: not all syntactically correct constructs are supported.

Examples:
 "xring" == use one global ring reduction for all tensors
"pscpu" == use CPU at worker 0 to reduce all tensors
"nccl" == use NCCL to locally reduce all tensors.  Limited to 1 worker.
"nccl/xring" == locally (to one worker) reduce values using NCCL then ring reduce across workers.
"pscpu:32k:xring" == use pscpu algorithm for tensors of size up to 32kB, then xring for larger tensors."""
        )
    parser.add_argument('--agg_small_grads_max_bytes', type=int, default=0,
        help=
        'If > 0, try to aggregate tensors of less than this number of bytes prior to all-reduce.'
        )
    parser.add_argument('--agg_small_grads_max_group', type=int, default=10,
        help=
        'When aggregating small tensors for all-reduce do not aggregate more than this many into one new tensor.'
        )
    parser.add_argument('--allreduce_merge_scope', type=int, default=1,
        help=
        'Establish a name scope around this many gradients prior to creating the all-reduce operations. It may affect the ability of the backend to merge parallel ops.'
        )
    parser.add_argument('--local_parameter_device', type=str, default='gpu',
        help=
        'Device to use as parameter server: cpu or gpu. For distributed training, it can affect where caching of variables happens.'
        )
    parser.add_argument('--use_resource_vars', type='bool', default=False,
        help=
        'Use resource variables instead of normal variables. Resource variables are slower, but this option is useful for debugging their performance.'
        )
    parser.add_argument('--debug', type='bool', default=False, help=
        'Debug train and eval')
    parser.add_argument('--debug_num_train_steps', type=int, default=None,
        help='Num steps to train.')
    parser.add_argument('--show_metrics', type='bool', default=True, help=
        'whether to show detailed metrics')
    parser.add_argument('--clip_grads', type='bool', default=True, help=
        'whether to clip gradients')
    parser.add_argument('--profile', type='bool', default=False, help=
        'If generate profile')
    parser.add_argument('--profile_save_steps', type=int, default=10, help=
        'Save timeline every N steps.')
    parser.add_argument('--use_dynamic_rnn', type='bool', default=True)
    parser.add_argument('--use_synthetic_data', type='bool', default=False)
    parser.add_argument('--mode', type=str, default='train_and_eval',
        choices=('train_and_eval', 'infer', 'translate'))


def create_hparams(flags):
    """Create training hparams."""
    return tf.contrib.training.HParams(src=flags.src, tgt=flags.tgt,
        train_prefix=os.path.join(flags.data_dir, flags.train_prefix),
        test_prefix=os.path.join(flags.data_dir, flags.test_prefix),
        translate_file=flags.translate_file, vocab_prefix=os.path.join(
        flags.data_dir, flags.vocab_prefix), embed_prefix=flags.
        embed_prefix, output_dir=flags.output_dir, num_units=flags.
        num_units, num_encoder_layers=flags.num_encoder_layers or flags.
        num_layers, num_decoder_layers=flags.num_decoder_layers or flags.
        num_layers, dropout=flags.dropout, unit_type=flags.unit_type,
        encoder_type=flags.encoder_type, residual=flags.residual,
        time_major=flags.time_major, num_embeddings_partitions=flags.
        num_embeddings_partitions, attention=flags.attention,
        attention_architecture=flags.attention_architecture,
        output_attention=flags.output_attention, pass_hidden_state=flags.
        pass_hidden_state, optimizer=flags.optimizer, max_train_epochs=
        flags.max_train_epochs, target_bleu=flags.target_bleu,
        label_smoothing=flags.label_smoothing, batch_size=flags.batch_size,
        init_op=flags.init_op, init_weight=flags.init_weight,
        max_gradient_norm=flags.max_gradient_norm, learning_rate=flags.
        learning_rate, warmup_steps=flags.warmup_steps, warmup_scheme=flags
        .warmup_scheme, decay_scheme=flags.decay_scheme,
        colocate_gradients_with_ops=flags.colocate_gradients_with_ops,
        num_buckets=flags.num_buckets, src_max_len=flags.src_max_len,
        tgt_max_len=flags.tgt_max_len, src_max_len_infer=flags.
        src_max_len_infer, tgt_max_len_infer=flags.tgt_max_len_infer, ckpt=
        flags.ckpt, infer_batch_size=flags.infer_batch_size,
        detokenizer_file=flags.detokenizer_file if flags.detokenizer_file
         is not None else os.path.join(flags.data_dir,
        'mosesdecoder/scripts/tokenizer/detokenizer.perl'), tokenizer_file=
        flags.tokenizer_file if flags.tokenizer_file is not None else os.
        path.join(flags.data_dir,
        'mosesdecoder/scripts/tokenizer/tokenizer.perl'), infer_mode=flags.
        infer_mode, beam_width=flags.beam_width, length_penalty_weight=
        flags.length_penalty_weight, coverage_penalty_weight=flags.
        coverage_penalty_weight, sos=flags.sos if flags.sos else
        vocab_utils.SOS, eos=flags.eos if flags.eos else vocab_utils.EOS,
        subword_option=flags.subword_option, check_special_token=flags.
        check_special_token, use_char_encode=flags.use_char_encode,
        forget_bias=flags.forget_bias, num_gpus=flags.num_gpus,
        save_checkpoints_steps=flags.save_checkpoints_steps,
        log_step_count_steps=flags.log_step_count_steps, epoch_step=0,
        share_vocab=flags.share_vocab, random_seed=flags.random_seed,
        language_model=flags.language_model, amp=flags.amp, use_fastmath=
        flags.use_fastmath, use_fp16=flags.use_fp16, fp16_loss_scale=flags.
        fp16_loss_scale, enable_auto_loss_scale=flags.
        enable_auto_loss_scale, fp16_inc_loss_scale_every_n=flags.
        fp16_inc_loss_scale_every_n, check_tower_loss_numerics=flags.
        check_tower_loss_numerics, use_fp32_batch_matmul=flags.
        use_fp32_batch_matmul, force_inputs_padding=flags.
        force_inputs_padding, use_xla=flags.use_xla, xla_compile=flags.
        xla_compile, use_autojit_xla=flags.use_autojit_xla,
        use_pintohost_optimizer=flags.use_pintohost_optimizer,
        use_cudnn_lstm=flags.use_cudnn_lstm, use_loose_bidi_cudnn_lstm=
        flags.use_loose_bidi_cudnn_lstm, use_fused_lstm=flags.
        use_fused_lstm, use_fused_lstm_dec=flags.use_fused_lstm_dec,
        gpu_indices=flags.gpu_indices, parallel_iterations=flags.
        parallel_iterations, use_dynamic_rnn=flags.use_dynamic_rnn,
        use_dist_strategy=flags.use_dist_strategy, hierarchical_copy=flags.
        hierarchical_copy, network_topology=flags.network_topology,
        use_block_lstm=flags.use_block_lstm, gradient_repacking=flags.
        gradient_repacking, compact_gradient_transfer=flags.
        compact_gradient_transfer, all_reduce_spec=flags.all_reduce_spec,
        agg_small_grads_max_bytes=flags.agg_small_grads_max_bytes,
        agg_small_grads_max_group=flags.agg_small_grads_max_group,
        allreduce_merge_scope=flags.allreduce_merge_scope,
        local_parameter_device='cpu' if flags.num_gpus == 0 else flags.
        local_parameter_device, use_resource_vars=flags.use_resource_vars,
        debug=flags.debug, debug_num_train_steps=flags.
        debug_num_train_steps, clip_grads=flags.clip_grads, profile=flags.
        profile, profile_save_steps=flags.profile_save_steps, show_metrics=
        flags.show_metrics, use_synthetic_data=flags.use_synthetic_data,
        mode=flags.mode)


def _add_argument(hparams, key, value, update=True):
    """Add an argument to hparams; if exists, change the value if update==True."""
    if hasattr(hparams, key):
        if update:
            setattr(hparams, key, value)
    else:
        hparams.add_hparam(key, value)


def extend_hparams(hparams):
    """Add new arguments to hparams."""
    if hparams.encoder_type == 'bi' and hparams.num_encoder_layers % 2 != 0:
        raise ValueError('For bi, num_encoder_layers %d should be even' %
            hparams.num_encoder_layers)
    if hparams.attention_architecture in ['gnmt'
        ] and hparams.num_encoder_layers < 2:
        raise ValueError(
            'For gnmt attention architecture, num_encoder_layers %d should be >= 2'
             % hparams.num_encoder_layers)
    if hparams.subword_option and hparams.subword_option not in ['spm', 'bpe']:
        raise ValueError('subword option must be either spm, or bpe')
    if hparams.infer_mode == 'beam_search' and hparams.beam_width <= 0:
        raise ValueError(
            'beam_width must greater than 0 when using beam_searchdecoder.')
    if hparams.mode == 'translate' and not hparams.translate_file:
        raise ValueError(
            '--translate_file flag must be specified in translate mode')
    assert hparams.num_encoder_layers and hparams.num_decoder_layers
    if hparams.num_encoder_layers != hparams.num_decoder_layers:
        hparams.pass_hidden_state = False
        utils.print_out(
            'Num encoder layer %d is different from num decoder layer %d, so set pass_hidden_state to False'
             % (hparams.num_encoder_layers, hparams.num_decoder_layers))
    num_encoder_residual_layers = 0
    num_decoder_residual_layers = 0
    if hparams.residual:
        if hparams.num_encoder_layers > 1:
            num_encoder_residual_layers = hparams.num_encoder_layers - 1
        if hparams.num_decoder_layers > 1:
            num_decoder_residual_layers = hparams.num_decoder_layers - 1
        if hparams.encoder_type == 'gnmt':
            num_encoder_residual_layers = hparams.num_encoder_layers - 2
            if hparams.num_encoder_layers == hparams.num_decoder_layers:
                num_decoder_residual_layers = num_encoder_residual_layers
    _add_argument(hparams, 'num_encoder_residual_layers',
        num_encoder_residual_layers)
    _add_argument(hparams, 'num_decoder_residual_layers',
        num_decoder_residual_layers)
    if hparams.language_model:
        hparams.attention = ''
        hparams.attention_architecture = ''
        hparams.pass_hidden_state = False
        hparams.share_vocab = True
        hparams.src = hparams.tgt
        utils.print_out(
            'For language modeling, we turn off attention and pass_hidden_state; turn on share_vocab; set src to tgt.'
            )
    if hparams.vocab_prefix:
        src_vocab_file = hparams.vocab_prefix + '.' + hparams.src
        tgt_vocab_file = hparams.vocab_prefix + '.' + hparams.tgt
    else:
        raise ValueError('hparams.vocab_prefix must be provided.')
    src_vocab_size, src_vocab_file = vocab_utils.check_vocab(src_vocab_file,
        hparams.output_dir, check_special_token=hparams.check_special_token,
        sos=hparams.sos, eos=hparams.eos, unk=vocab_utils.UNK, pad_vocab=True)
    if hparams.share_vocab:
        utils.print_out('  using source vocab for target')
        tgt_vocab_file = src_vocab_file
        tgt_vocab_size = src_vocab_size
    else:
        tgt_vocab_size, tgt_vocab_file = vocab_utils.check_vocab(tgt_vocab_file
            , hparams.output_dir, check_special_token=hparams.
            check_special_token, sos=hparams.sos, eos=hparams.eos, unk=
            vocab_utils.UNK)
    _add_argument(hparams, 'src_vocab_size', src_vocab_size)
    _add_argument(hparams, 'tgt_vocab_size', tgt_vocab_size)
    _add_argument(hparams, 'src_vocab_file', src_vocab_file)
    _add_argument(hparams, 'tgt_vocab_file', tgt_vocab_file)
    _add_argument(hparams, 'num_enc_emb_partitions', hparams.
        num_embeddings_partitions)
    _add_argument(hparams, 'num_dec_emb_partitions', hparams.
        num_embeddings_partitions)
    _add_argument(hparams, 'src_embed_file', '')
    _add_argument(hparams, 'tgt_embed_file', '')
    if hparams.embed_prefix:
        src_embed_file = hparams.embed_prefix + '.' + hparams.src
        tgt_embed_file = hparams.embed_prefix + '.' + hparams.tgt
        if tf.gfile.Exists(src_embed_file):
            utils.print_out('  src_embed_file %s exist' % src_embed_file)
            hparams.src_embed_file = src_embed_file
            utils.print_out(
                'For pretrained embeddings, set num_enc_emb_partitions to 1')
            hparams.num_enc_emb_partitions = 1
        else:
            utils.print_out("  src_embed_file %s doesn't exist" %
                src_embed_file)
        if tf.gfile.Exists(tgt_embed_file):
            utils.print_out('  tgt_embed_file %s exist' % tgt_embed_file)
            hparams.tgt_embed_file = tgt_embed_file
            utils.print_out(
                'For pretrained embeddings, set num_dec_emb_partitions to 1')
            hparams.num_dec_emb_partitions = 1
        else:
            utils.print_out("  tgt_embed_file %s doesn't exist" %
                tgt_embed_file)
    metric = 'bleu'
    best_metric_dir = os.path.join(hparams.output_dir, 'best_' + metric)
    tf.gfile.MakeDirs(best_metric_dir)
    _add_argument(hparams, 'best_' + metric, 0, update=False)
    _add_argument(hparams, 'best_' + metric + '_dir', best_metric_dir)
    return hparams


def create_or_load_hparams(default_hparams, hparams_path):
    """Create hparams or load hparams from output_dir."""
    hparams = utils.maybe_parse_standard_hparams(default_hparams, hparams_path)
    hparams = extend_hparams(hparams)
    utils.print_hparams(hparams)
    return hparams


def run_main(flags, default_hparams, estimator_fn):
    """Run main."""
    random_seed = flags.random_seed
    if random_seed is not None and random_seed > 0:
        utils.print_out('# Set random seed to %d' % random_seed)
        random.seed(random_seed)
        np.random.seed(random_seed)
        tf.set_random_seed(random_seed)
    output_dir = flags.output_dir
    if output_dir and not tf.gfile.Exists(output_dir):
        utils.print_out('# Creating output directory %s ...' % output_dir)
        tf.gfile.MakeDirs(output_dir)
    hparams = create_or_load_hparams(default_hparams, flags.hparams_path)
    estimator_fn(hparams)
    return hparams


def tokenize(hparams, file, tokenized_file):
    utils.print_out('tokenizing {} -> {}'.format(file, tokenized_file))
    with open(file, 'rb') as input_file:
        with open(tokenized_file, 'wb') as output_file:
            subprocess.run([hparams.tokenizer_file, '-l', hparams.src],
                stdin=input_file, stdout=output_file)


def detokenize(hparams, file, detokenized_file):
    utils.print_out('detokenizing {} -> {}'.format(file, detokenized_file))
    with open(file, 'rb') as input_file:
        with open(detokenized_file, 'wb') as output_file:
            subprocess.run([hparams.detokenizer_file, '-l', hparams.tgt],
                stdin=input_file, stdout=output_file)


def main(unused_argv):
    experiment_start = time.time()
    tf.logging.set_verbosity(tf.logging.INFO)
    if FLAGS.use_fp16 and FLAGS.use_dist_strategy:
        raise ValueError("use_fp16 and use_dist_strategy aren't compatible")
    if FLAGS.use_fp16 + FLAGS.amp + FLAGS.use_fastmath > 1:
        raise ValueError('Only one of use_fp16, amp, use_fastmath can be set')
    if FLAGS.amp:
        utils.print_out('Enabling TF-AMP')
        os.environ['TF_ENABLE_AUTO_MIXED_PRECISION'] = '1'
    if FLAGS.use_fastmath:
        utils.print_out('Enabling FastMath')
        os.environ['TF_ENABLE_CUBLAS_TENSOR_OP_MATH_FP32'] = '1'
        os.environ['TF_ENABLE_CUDNN_TENSOR_OP_MATH_FP32'] = '1'
        os.environ['TF_ENABLE_CUDNN_RNN_TENSOR_OP_MATH_FP32'] = '1'
    active_xla_option_nums = np.sum([FLAGS.use_xla, FLAGS.use_autojit_xla,
        FLAGS.xla_compile])
    if active_xla_option_nums > 1:
        raise ValueError(
            'Only one of use_xla, xla_compile, use_autojit_xla can be set')
    os.environ['use_xla'] = str(FLAGS.use_xla).lower()
    if FLAGS.use_xla:
        os.environ['use_defun'] = str(True).lower()
    else:
        os.environ['use_defun'] = str(FLAGS.use_defun).lower()
    utils.print_out('use_defun is %s for attention' % os.environ['use_defun'])
    os.environ['use_fp32_batch_matmul'
        ] = 'true' if FLAGS.use_fp32_batch_matmul else 'false'
    os.environ['xla_compile'] = 'true' if FLAGS.xla_compile else 'false'
    os.environ['force_inputs_padding'
        ] = 'true' if FLAGS.force_inputs_padding else 'false'
    if FLAGS.mode == 'train':
        utils.print_out('Running training mode.')
        default_hparams = create_hparams(FLAGS)
        run_main(FLAGS, default_hparams, estimator.train_fn)
    elif FLAGS.mode == 'infer' or FLAGS.mode == 'translate':
        if FLAGS.mode == 'infer':
            utils.print_out('Running inference mode.')
            translate_mode = False
        else:
            utils.print_out('Running translate mode on file {}.'.format(
                FLAGS.translate_file))
            translate_mode = True
        random_seed = FLAGS.random_seed
        if random_seed is not None and random_seed > 0:
            utils.print_out('# Set random seed to %d' % random_seed)
            random.seed(random_seed)
            np.random.seed(random_seed)
            tf.set_random_seed(random_seed)
        output_dir = FLAGS.output_dir
        if output_dir and not tf.gfile.Exists(output_dir):
            utils.print_out('# Creating output directory %s ...' % output_dir)
            tf.gfile.MakeDirs(output_dir)
        dllogger.init(backends=[dllogger.StdOutBackend(dllogger.Verbosity.
            DEFAULT), dllogger.JSONStreamBackend(dllogger.Verbosity.VERBOSE,
            os.path.join(FLAGS.output_dir, FLAGS.mode + '-report.json'))])
        dllogger.log('PARAMETER', vars(FLAGS))
        default_hparams = create_hparams(FLAGS)
        default_hparams.num_buckets = 1
        hparams = create_or_load_hparams(default_hparams, FLAGS.hparams_path)
        utils.print_out('infer_hparams:')
        utils.print_hparams(hparams)
        if translate_mode:
            tokenize(hparams, hparams.translate_file, hparams.
                translate_file + '.tok')
        eval_sentences, eval_src_tokens, _ = (iterator_utils.
            get_effective_epoch_size(hparams, train=False))
        tf.logging.info('Starting to evaluate...')
        eval_start = time.time()
        _, (eval_speed, eval_latencies
            ), eval_output_tokens = estimator.eval_fn(hparams, hparams.ckpt,
            only_translate=translate_mode)
        eval_end = time.time()
        eval_delta = eval_end - eval_start
        utils.print_out(
            'eval time for ckpt: %.2f mins (%.2f sent/sec, %.2f tokens/sec)' %
            (eval_delta / 60.0, eval_speed, eval_speed * (eval_src_tokens +
            eval_output_tokens) / eval_sentences), f=sys.stderr)
        logging_data = {'infer_speed_sent': eval_speed, 'infer_speed_toks':
            eval_speed * (eval_src_tokens + eval_output_tokens) /
            eval_sentences}
        for lat in sorted(eval_latencies):
            utils.print_out('eval latency_%s for ckpt: %.2f ms' % (lat, 
                eval_latencies[lat] * 1000))
            logging_data['infer_latency_{}'.format(lat)] = eval_latencies[lat
                ] * 1000
        dllogger.log((), logging_data)
        dllogger.flush()
        if translate_mode:
            detokenize(hparams, hparams.translate_file + '.trans.tok', 
                hparams.translate_file + '.trans')
    else:
        assert FLAGS.mode == 'train_and_eval'
        utils.print_out('Running train and eval mode.')
        random_seed = FLAGS.random_seed
        if random_seed is not None and random_seed > 0:
            utils.print_out('# Set random seed to %d' % random_seed)
            random.seed(random_seed)
            np.random.seed(random_seed)
            tf.set_random_seed(random_seed)
        output_dir = FLAGS.output_dir
        if output_dir and not tf.gfile.Exists(output_dir):
            utils.print_out('# Creating output directory %s ...' % output_dir)
            tf.gfile.MakeDirs(output_dir)
        dllogger.init(backends=[dllogger.StdOutBackend(dllogger.Verbosity.
            DEFAULT), dllogger.JSONStreamBackend(dllogger.Verbosity.VERBOSE,
            os.path.join(FLAGS.output_dir, FLAGS.mode + '-report.json'))])
        dllogger.log('PARAMETER', vars(FLAGS))
        dllogger.metadata('bleu', {'unit': None})
        dllogger.metadata('train_speed_sent', {'unit': 'sequences/s'})
        dllogger.metadata('train_speed_toks', {'unit': 'tokens/s'})
        default_hparams = create_hparams(FLAGS)
        hparams = create_or_load_hparams(default_hparams, FLAGS.hparams_path)
        utils.print_out('training hparams:')
        utils.print_hparams(hparams)
        with tf.gfile.GFile(os.path.join(output_dir, 'train_hparams.txt'), 'w'
            ) as f:
            f.write(utils.serialize_hparams(hparams) + '\n')
        infer_hparams = tf.contrib.training.HParams(**hparams.values())
        infer_hparams.num_buckets = 1
        utils.print_out('infer_hparams:')
        utils.print_hparams(infer_hparams)
        with tf.gfile.GFile(os.path.join(output_dir, 'infer_hparams.txt'), 'w'
            ) as f:
            f.write(utils.serialize_hparams(infer_hparams) + '\n')
        epochs = 0
        should_stop = epochs >= FLAGS.max_train_epochs
        train_sentences, train_src_tokens, train_tgt_tokens = (iterator_utils
            .get_effective_epoch_size(hparams))
        eval_sentences, eval_src_tokens, _ = (iterator_utils.
            get_effective_epoch_size(hparams, train=False))
        while not should_stop:
            utils.print_out('Starting epoch %d' % epochs)
            try:
                train_start = time.time()
                train_speed, _ = estimator.train_fn(hparams)
            except tf.errors.OutOfRangeError:
                utils.print_out('training hits OutOfRangeError', f=sys.stderr)
            train_end = time.time()
            train_delta = train_end - train_start
            utils.print_out(
                'training time for epoch %d: %.2f mins (%.2f sent/sec, %.2f tokens/sec)'
                 % (epochs + 1, train_delta / 60.0, train_speed, 
                train_speed * (train_src_tokens + train_tgt_tokens) /
                train_sentences), f=sys.stderr)
            logging_data = {'train_speed_sent': train_speed,
                'train_speed_toks': train_speed * (train_src_tokens +
                train_tgt_tokens) / train_sentences}
            eval_start = time.time()
            bleu_score, (eval_speed, eval_latencies
                ), eval_output_tokens = estimator.eval_fn(infer_hparams)
            eval_end = time.time()
            eval_delta = eval_end - eval_start
            utils.print_out(
                'eval time for epoch %d: %.2f mins (%.2f sent/sec, %.2f tokens/sec)'
                 % (epochs + 1, eval_delta / 60.0, eval_speed, eval_speed *
                (eval_src_tokens + eval_output_tokens) / eval_sentences), f
                =sys.stderr)
            logging_data.update({'bleu': bleu_score, 'infer_speed_sent':
                eval_speed, 'infer_speed_toks': eval_speed * (
                eval_src_tokens + eval_output_tokens) / eval_sentences})
            for lat in sorted(eval_latencies):
                utils.print_out('eval latency_%s for epoch %d: %.2f ms' % (
                    lat, epochs + 1, eval_latencies[lat] * 1000))
                logging_data['eval_latency_{}'.format(lat)] = eval_latencies[
                    lat] * 1000
            dllogger.log((epochs,), logging_data)
            dllogger.flush()
            if (FLAGS.debug or FLAGS.target_bleu is not None and bleu_score >
                FLAGS.target_bleu):
                should_stop = True
                utils.print_out(
                    'Stop job since target bleu is reached at epoch %d .' %
                    epochs, f=sys.stderr)
            epochs += 1
            if epochs >= FLAGS.max_train_epochs:
                should_stop = True
                utils.print_out('Stop job since max_train_epochs is reached.',
                    f=sys.stderr)
        dllogger.log((), logging_data)
        dllogger.flush()
    experiment_end = time.time()
    utils.print_out('Experiment took {} min'.format((experiment_end -
        experiment_start) / 60))


<mask token>

<mask token>
utils.check_tensorflow_version()
<mask token>


def add_arguments(parser):
    """Build ArgumentParser."""
    parser.register('type', 'bool', lambda v: v.lower() == 'true')
    parser.add_argument('--num_units', type=int, default=1024, help=
        'Network size.')
    parser.add_argument('--num_layers', type=int, default=4, help=
        'Network depth.')
    parser.add_argument('--num_encoder_layers', type=int, default=None,
        help='Encoder depth, equal to num_layers if None.')
    parser.add_argument('--num_decoder_layers', type=int, default=None,
        help='Decoder depth, equal to num_layers if None.')
    parser.add_argument('--encoder_type', type=str, default='gnmt', help=
        """      uni | bi | gnmt.
      For bi, we build num_encoder_layers/2 bi-directional layers.
      For gnmt, we build 1 bi-directional layer, and (num_encoder_layers - 1)
        uni-directional layers.      """
        )
    parser.add_argument('--residual', type='bool', nargs='?', const=True,
        default=True, help='Whether to add residual connections.')
    parser.add_argument('--time_major', type='bool', nargs='?', const=True,
        default=True, help='Whether to use time-major mode for dynamic RNN.')
    parser.add_argument('--num_embeddings_partitions', type=int, default=0,
        help='Number of partitions for embedding vars.')
    parser.add_argument('--attention', type=str, default='normed_bahdanau',
        help=
        """      luong | scaled_luong | bahdanau | normed_bahdanau or set to "" for no
      attention      """
        )
    parser.add_argument('--attention_architecture', type=str, default=
        'gnmt_v2', help=
        """      standard | gnmt | gnmt_v2.
      standard: use top layer to compute attention.
      gnmt: GNMT style of computing attention, use previous bottom layer to
          compute attention.
      gnmt_v2: similar to gnmt, but use current bottom layer to compute
          attention.      """
        )
    parser.add_argument('--output_attention', type='bool', nargs='?', const
        =True, default=True, help=
        """      Only used in standard attention_architecture. Whether use attention as
      the cell output at each timestep.
      .      """
        )
    parser.add_argument('--pass_hidden_state', type='bool', nargs='?',
        const=True, default=True, help=
        """      Whether to pass encoder's hidden state to decoder when using an attention
      based model.      """
        )
    parser.add_argument('--optimizer', type=str, default='adam', help=
        'sgd | adam')
    parser.add_argument('--learning_rate', type=float, default=0.0005, help
        ='Learning rate. Adam: 0.001 | 0.0001')
    parser.add_argument('--warmup_steps', type=int, default=200, help=
        'How many steps we inverse-decay learning.')
    parser.add_argument('--warmup_scheme', type=str, default='t2t', help=
        """      How to warmup learning rates. Options include:
        t2t: Tensor2Tensor's way, start with lr 100 times smaller, then
             exponentiate until the specified lr.      """
        )
    parser.add_argument('--decay_scheme', type=str, default='luong234',
        help=
        """      How we decay learning rate. Options include:
        luong234: after 2/3 num train steps, we start halving the learning rate
          for 4 times before finishing.
        luong5: after 1/2 num train steps, we start halving the learning rate
          for 5 times before finishing.        luong10: after 1/2 num train steps, we start halving the learning rate
          for 10 times before finishing.      """
        )
    parser.add_argument('--max_train_epochs', type=int, default=6, help=
        'Max number of epochs.')
    parser.add_argument('--target_bleu', type=float, default=None, help=
        'Target bleu.')
    parser.add_argument('--colocate_gradients_with_ops', type='bool', nargs
        ='?', const=True, default=True, help=
        'Whether try colocating gradients with corresponding op')
    parser.add_argument('--label_smoothing', type=float, default=0.1, help=
        'If nonzero, smooth the labels towards 1/num_classes.')
    parser.add_argument('--init_op', type=str, default='uniform', help=
        'uniform | glorot_normal | glorot_uniform')
    parser.add_argument('--init_weight', type=float, default=0.1, help=
        'for uniform init_op, initialize weights between [-this, this].')
    parser.add_argument('--src', type=str, default='en', help=
        'Source suffix, e.g., en.')
    parser.add_argument('--tgt', type=str, default='de', help=
        'Target suffix, e.g., de.')
    parser.add_argument('--data_dir', type=str, default='data/wmt16_de_en',
        help='Training/eval data directory.')
    parser.add_argument('--train_prefix', type=str, default=
        'train.tok.clean.bpe.32000', help=
        'Train prefix, expect files with src/tgt suffixes.')
    parser.add_argument('--test_prefix', type=str, default=
        'newstest2014.tok.bpe.32000', help=
        'Test prefix, expect files with src/tgt suffixes.')
    parser.add_argument('--translate_file', type=str, help=
        'File to translate, works only with translate mode')
    parser.add_argument('--output_dir', type=str, default='results', help=
        'Store log/model files.')
    parser.add_argument('--vocab_prefix', type=str, default=
        'vocab.bpe.32000', help=
        '      Vocab prefix, expect files with src/tgt suffixes.      ')
    parser.add_argument('--embed_prefix', type=str, default=None, help=
        """      Pretrained embedding prefix, expect files with src/tgt suffixes.
      The embedding files should be Glove formatted txt files.      """
        )
    parser.add_argument('--sos', type=str, default='<s>', help=
        'Start-of-sentence symbol.')
    parser.add_argument('--eos', type=str, default='</s>', help=
        'End-of-sentence symbol.')
    parser.add_argument('--share_vocab', type='bool', nargs='?', const=True,
        default=True, help=
        """      Whether to use the source vocab and embeddings for both source and
      target.      """
        )
    parser.add_argument('--check_special_token', type='bool', default=True,
        help=
        """                      Whether check special sos, eos, unk tokens exist in the
                      vocab files.                      """
        )
    parser.add_argument('--src_max_len', type=int, default=50, help=
        'Max length of src sequences during training (including EOS).')
    parser.add_argument('--tgt_max_len', type=int, default=50, help=
        'Max length of tgt sequences during training (including BOS).')
    parser.add_argument('--src_max_len_infer', type=int, default=None, help
        ='Max length of src sequences during inference (including EOS).')
    parser.add_argument('--tgt_max_len_infer', type=int, default=80, help=
        """      Max length of tgt sequences during inference (including BOS).  Also use to restrict the
      maximum decoding length.      """
        )
    parser.add_argument('--unit_type', type=str, default='lstm', help=
        'lstm | gru | layer_norm_lstm | nas')
    parser.add_argument('--forget_bias', type=float, default=0.0, help=
        'Forget bias for BasicLSTMCell.')
    parser.add_argument('--dropout', type=float, default=0.2, help=
        'Dropout rate (not keep_prob)')
    parser.add_argument('--max_gradient_norm', type=float, default=5.0,
        help='Clip gradients to this norm.')
    parser.add_argument('--batch_size', type=int, default=128, help=
        'Total batch size.')
    parser.add_argument('--num_buckets', type=int, default=5, help=
        'Put data into similar-length buckets (only for training).')
    parser.add_argument('--subword_option', type=str, default='bpe',
        choices=['', 'bpe', 'spm'], help=
        '                      Set to bpe or spm to activate subword desegmentation.                      '
        )
    parser.add_argument('--use_char_encode', type='bool', default=False,
        help=
        """                      Whether to split each word or bpe into character, and then
                      generate the word-level representation from the character
                      reprentation.
                      """
        )
    parser.add_argument('--save_checkpoints_steps', type=int, default=2000,
        help='save_checkpoints_steps')
    parser.add_argument('--log_step_count_steps', type=int, default=10,
        help=
        'The frequency, in number of global steps, that the global step and the loss will be logged during training'
        )
    parser.add_argument('--num_gpus', type=int, default=1, help=
        'Number of gpus in each worker.')
    parser.add_argument('--hparams_path', type=str, default=None, help=
        'Path to standard hparams json file that overrideshparams values from FLAGS.'
        )
    parser.add_argument('--random_seed', type=int, default=1, help=
        'Random seed (>0, set a specific seed).')
    parser.add_argument('--language_model', type='bool', nargs='?', const=
        True, default=False, help=
        'True to train a language model, ignoring encoder')
    parser.add_argument('--ckpt', type=str, default=None, help=
        'Checkpoint file to load a model for inference. (defaults to newest checkpoint)'
        )
    parser.add_argument('--infer_batch_size', type=int, default=128, help=
        'Batch size for inference mode.')
    parser.add_argument('--detokenizer_file', type=str, default=None, help=
        'Detokenizer script file. Default: DATA_DIR/mosesdecoder/scripts/tokenizer/detokenizer.perl'
        )
    parser.add_argument('--tokenizer_file', type=str, default=None, help=
        'Tokenizer script file. Default: DATA_DIR/mosesdecoder/scripts/tokenizer/tokenizer.perl'
        )
    parser.add_argument('--infer_mode', type=str, default='beam_search',
        choices=['greedy', 'beam_search'], help=
        'Which type of decoder to use during inference.')
    parser.add_argument('--beam_width', type=int, default=5, help=
        """      beam width when using beam search decoder. If 0, use standard
      decoder with greedy helper.      """
        )
    parser.add_argument('--length_penalty_weight', type=float, default=0.6,
        help='Length penalty for beam search.')
    parser.add_argument('--coverage_penalty_weight', type=float, default=
        0.1, help='Coverage penalty for beam search.')
    parser.add_argument('--num_workers', type=int, default=1, help=
        'Number of workers (inference only).')
    parser.add_argument('--amp', action='store_true', help=
        'use amp for training and inference')
    parser.add_argument('--use_fastmath', type='bool', default=False, help=
        'use_fastmath for training and inference')
    parser.add_argument('--use_fp16', type='bool', default=False, help=
        'use_fp16 for training and inference')
    parser.add_argument('--fp16_loss_scale', type=float, default=128, help=
        'If fp16 is enabled, the loss is multiplied by this amount right before gradients are computed, then each gradient is divided by this amount. Mathematically, this has no effect, but it helps avoid fp16 underflow. Set to 1 to effectively disable.'
        )
    parser.add_argument('--enable_auto_loss_scale', type='bool', default=
        True, help=
        'If True and use_fp16 is True, automatically adjust the loss scale during training.'
        )
    parser.add_argument('--fp16_inc_loss_scale_every_n', type=int, default=
        128, help=
        'If fp16 is enabled and enable_auto_loss_scale is True, increase the loss scale every n steps.'
        )
    parser.add_argument('--check_tower_loss_numerics', type='bool', default
        =False, help='whether to check tower loss numerics')
    parser.add_argument('--use_fp32_batch_matmul', type='bool', default=
        False, help='Whether to use fp32 batch matmul')
    parser.add_argument('--force_inputs_padding', type='bool', default=
        False, help='Force padding input batch to src_max_len and tgt_max_len')
    parser.add_argument('--use_xla', type='bool', default=False, help=
        'Use xla to compile a few selected locations, mostly Defuns.')
    parser.add_argument('--xla_compile', type='bool', default=False, help=
        "Use xla.compile() for each tower's fwd and bak pass.")
    parser.add_argument('--use_autojit_xla', type='bool', default=False,
        help='Use auto jit xla.')
    parser.add_argument('--use_pintohost_optimizer', type='bool', default=
        False, help='whether to use PinToHost optimizer')
    parser.add_argument('--use_cudnn_lstm', type='bool', default=False,
        help='whether to use cudnn_lstm for encoder, non residual layers')
    parser.add_argument('--use_loose_bidi_cudnn_lstm', type='bool', default
        =False, help='whether to use loose bidi cudnn_lstm')
    parser.add_argument('--use_fused_lstm', type='bool', default=True, help
        =
        'whether to use fused lstm and variant. If enabled, training will use LSTMBlockFusedCell, infer will use LSTMBlockCell when appropriate.'
        )
    parser.add_argument('--use_fused_lstm_dec', type='bool', default=False,
        help='whether to use fused lstm for decoder (training only).')
    parser.add_argument('--gpu_indices', type=str, default='', help=
        'Indices of worker GPUs in ring order')
    parser.add_argument('--parallel_iterations', type=int, default=10, help
        ='number of parallel iterations in dynamic_rnn')
    parser.add_argument('--use_dist_strategy', type='bool', default=False,
        help='whether to use distribution strategy')
    parser.add_argument('--hierarchical_copy', type='bool', default=False,
        help=
        'Use hierarchical copies. Currently only optimized for use on a DGX-1 with 8 GPUs and may perform poorly on other hardware. Requires --num_gpus > 1, and only recommended when --num_gpus=8'
        )
    parser.add_argument('--network_topology', type=constants.
        NetworkTopology, default=constants.NetworkTopology.DGX1, choices=
        list(constants.NetworkTopology))
    parser.add_argument('--use_block_lstm', type='bool', default=False,
        help='whether to use block lstm')
    parser.add_argument('--use_defun', type='bool', default=False, help=
        'whether to use Defun')
    parser.add_argument('--gradient_repacking', type=int, default=0, help=
        'Use gradient repacking. Itcurrently only works with replicated mode. At the end ofof each step, it repacks the gradients for more efficientcross-device transportation. A non-zero value specifiesthe number of split packs that will be formed.'
        )
    parser.add_argument('--compact_gradient_transfer', type='bool', default
        =True, help=
        'Compact gradient as much as possible for cross-device transfer and aggregation.'
        )
    parser.add_argument('--all_reduce_spec', type=str, default='nccl', help
        =
        """A specification of the all_reduce algorithm to be used for reducing gradients.  For more details, see parse_all_reduce_spec in variable_mgr.py.  An all_reduce_spec has BNF form:
int ::= positive whole number
g_int ::= int[KkMGT]?
alg_spec ::= alg | alg#int
range_spec ::= alg_spec | alg_spec/alg_spec
spec ::= range_spec | range_spec:g_int:range_spec
NOTE: not all syntactically correct constructs are supported.

Examples:
 "xring" == use one global ring reduction for all tensors
"pscpu" == use CPU at worker 0 to reduce all tensors
"nccl" == use NCCL to locally reduce all tensors.  Limited to 1 worker.
"nccl/xring" == locally (to one worker) reduce values using NCCL then ring reduce across workers.
"pscpu:32k:xring" == use pscpu algorithm for tensors of size up to 32kB, then xring for larger tensors."""
        )
    parser.add_argument('--agg_small_grads_max_bytes', type=int, default=0,
        help=
        'If > 0, try to aggregate tensors of less than this number of bytes prior to all-reduce.'
        )
    parser.add_argument('--agg_small_grads_max_group', type=int, default=10,
        help=
        'When aggregating small tensors for all-reduce do not aggregate more than this many into one new tensor.'
        )
    parser.add_argument('--allreduce_merge_scope', type=int, default=1,
        help=
        'Establish a name scope around this many gradients prior to creating the all-reduce operations. It may affect the ability of the backend to merge parallel ops.'
        )
    parser.add_argument('--local_parameter_device', type=str, default='gpu',
        help=
        'Device to use as parameter server: cpu or gpu. For distributed training, it can affect where caching of variables happens.'
        )
    parser.add_argument('--use_resource_vars', type='bool', default=False,
        help=
        'Use resource variables instead of normal variables. Resource variables are slower, but this option is useful for debugging their performance.'
        )
    parser.add_argument('--debug', type='bool', default=False, help=
        'Debug train and eval')
    parser.add_argument('--debug_num_train_steps', type=int, default=None,
        help='Num steps to train.')
    parser.add_argument('--show_metrics', type='bool', default=True, help=
        'whether to show detailed metrics')
    parser.add_argument('--clip_grads', type='bool', default=True, help=
        'whether to clip gradients')
    parser.add_argument('--profile', type='bool', default=False, help=
        'If generate profile')
    parser.add_argument('--profile_save_steps', type=int, default=10, help=
        'Save timeline every N steps.')
    parser.add_argument('--use_dynamic_rnn', type='bool', default=True)
    parser.add_argument('--use_synthetic_data', type='bool', default=False)
    parser.add_argument('--mode', type=str, default='train_and_eval',
        choices=('train_and_eval', 'infer', 'translate'))


def create_hparams(flags):
    """Create training hparams."""
    return tf.contrib.training.HParams(src=flags.src, tgt=flags.tgt,
        train_prefix=os.path.join(flags.data_dir, flags.train_prefix),
        test_prefix=os.path.join(flags.data_dir, flags.test_prefix),
        translate_file=flags.translate_file, vocab_prefix=os.path.join(
        flags.data_dir, flags.vocab_prefix), embed_prefix=flags.
        embed_prefix, output_dir=flags.output_dir, num_units=flags.
        num_units, num_encoder_layers=flags.num_encoder_layers or flags.
        num_layers, num_decoder_layers=flags.num_decoder_layers or flags.
        num_layers, dropout=flags.dropout, unit_type=flags.unit_type,
        encoder_type=flags.encoder_type, residual=flags.residual,
        time_major=flags.time_major, num_embeddings_partitions=flags.
        num_embeddings_partitions, attention=flags.attention,
        attention_architecture=flags.attention_architecture,
        output_attention=flags.output_attention, pass_hidden_state=flags.
        pass_hidden_state, optimizer=flags.optimizer, max_train_epochs=
        flags.max_train_epochs, target_bleu=flags.target_bleu,
        label_smoothing=flags.label_smoothing, batch_size=flags.batch_size,
        init_op=flags.init_op, init_weight=flags.init_weight,
        max_gradient_norm=flags.max_gradient_norm, learning_rate=flags.
        learning_rate, warmup_steps=flags.warmup_steps, warmup_scheme=flags
        .warmup_scheme, decay_scheme=flags.decay_scheme,
        colocate_gradients_with_ops=flags.colocate_gradients_with_ops,
        num_buckets=flags.num_buckets, src_max_len=flags.src_max_len,
        tgt_max_len=flags.tgt_max_len, src_max_len_infer=flags.
        src_max_len_infer, tgt_max_len_infer=flags.tgt_max_len_infer, ckpt=
        flags.ckpt, infer_batch_size=flags.infer_batch_size,
        detokenizer_file=flags.detokenizer_file if flags.detokenizer_file
         is not None else os.path.join(flags.data_dir,
        'mosesdecoder/scripts/tokenizer/detokenizer.perl'), tokenizer_file=
        flags.tokenizer_file if flags.tokenizer_file is not None else os.
        path.join(flags.data_dir,
        'mosesdecoder/scripts/tokenizer/tokenizer.perl'), infer_mode=flags.
        infer_mode, beam_width=flags.beam_width, length_penalty_weight=
        flags.length_penalty_weight, coverage_penalty_weight=flags.
        coverage_penalty_weight, sos=flags.sos if flags.sos else
        vocab_utils.SOS, eos=flags.eos if flags.eos else vocab_utils.EOS,
        subword_option=flags.subword_option, check_special_token=flags.
        check_special_token, use_char_encode=flags.use_char_encode,
        forget_bias=flags.forget_bias, num_gpus=flags.num_gpus,
        save_checkpoints_steps=flags.save_checkpoints_steps,
        log_step_count_steps=flags.log_step_count_steps, epoch_step=0,
        share_vocab=flags.share_vocab, random_seed=flags.random_seed,
        language_model=flags.language_model, amp=flags.amp, use_fastmath=
        flags.use_fastmath, use_fp16=flags.use_fp16, fp16_loss_scale=flags.
        fp16_loss_scale, enable_auto_loss_scale=flags.
        enable_auto_loss_scale, fp16_inc_loss_scale_every_n=flags.
        fp16_inc_loss_scale_every_n, check_tower_loss_numerics=flags.
        check_tower_loss_numerics, use_fp32_batch_matmul=flags.
        use_fp32_batch_matmul, force_inputs_padding=flags.
        force_inputs_padding, use_xla=flags.use_xla, xla_compile=flags.
        xla_compile, use_autojit_xla=flags.use_autojit_xla,
        use_pintohost_optimizer=flags.use_pintohost_optimizer,
        use_cudnn_lstm=flags.use_cudnn_lstm, use_loose_bidi_cudnn_lstm=
        flags.use_loose_bidi_cudnn_lstm, use_fused_lstm=flags.
        use_fused_lstm, use_fused_lstm_dec=flags.use_fused_lstm_dec,
        gpu_indices=flags.gpu_indices, parallel_iterations=flags.
        parallel_iterations, use_dynamic_rnn=flags.use_dynamic_rnn,
        use_dist_strategy=flags.use_dist_strategy, hierarchical_copy=flags.
        hierarchical_copy, network_topology=flags.network_topology,
        use_block_lstm=flags.use_block_lstm, gradient_repacking=flags.
        gradient_repacking, compact_gradient_transfer=flags.
        compact_gradient_transfer, all_reduce_spec=flags.all_reduce_spec,
        agg_small_grads_max_bytes=flags.agg_small_grads_max_bytes,
        agg_small_grads_max_group=flags.agg_small_grads_max_group,
        allreduce_merge_scope=flags.allreduce_merge_scope,
        local_parameter_device='cpu' if flags.num_gpus == 0 else flags.
        local_parameter_device, use_resource_vars=flags.use_resource_vars,
        debug=flags.debug, debug_num_train_steps=flags.
        debug_num_train_steps, clip_grads=flags.clip_grads, profile=flags.
        profile, profile_save_steps=flags.profile_save_steps, show_metrics=
        flags.show_metrics, use_synthetic_data=flags.use_synthetic_data,
        mode=flags.mode)


def _add_argument(hparams, key, value, update=True):
    """Add an argument to hparams; if exists, change the value if update==True."""
    if hasattr(hparams, key):
        if update:
            setattr(hparams, key, value)
    else:
        hparams.add_hparam(key, value)


def extend_hparams(hparams):
    """Add new arguments to hparams."""
    if hparams.encoder_type == 'bi' and hparams.num_encoder_layers % 2 != 0:
        raise ValueError('For bi, num_encoder_layers %d should be even' %
            hparams.num_encoder_layers)
    if hparams.attention_architecture in ['gnmt'
        ] and hparams.num_encoder_layers < 2:
        raise ValueError(
            'For gnmt attention architecture, num_encoder_layers %d should be >= 2'
             % hparams.num_encoder_layers)
    if hparams.subword_option and hparams.subword_option not in ['spm', 'bpe']:
        raise ValueError('subword option must be either spm, or bpe')
    if hparams.infer_mode == 'beam_search' and hparams.beam_width <= 0:
        raise ValueError(
            'beam_width must greater than 0 when using beam_searchdecoder.')
    if hparams.mode == 'translate' and not hparams.translate_file:
        raise ValueError(
            '--translate_file flag must be specified in translate mode')
    assert hparams.num_encoder_layers and hparams.num_decoder_layers
    if hparams.num_encoder_layers != hparams.num_decoder_layers:
        hparams.pass_hidden_state = False
        utils.print_out(
            'Num encoder layer %d is different from num decoder layer %d, so set pass_hidden_state to False'
             % (hparams.num_encoder_layers, hparams.num_decoder_layers))
    num_encoder_residual_layers = 0
    num_decoder_residual_layers = 0
    if hparams.residual:
        if hparams.num_encoder_layers > 1:
            num_encoder_residual_layers = hparams.num_encoder_layers - 1
        if hparams.num_decoder_layers > 1:
            num_decoder_residual_layers = hparams.num_decoder_layers - 1
        if hparams.encoder_type == 'gnmt':
            num_encoder_residual_layers = hparams.num_encoder_layers - 2
            if hparams.num_encoder_layers == hparams.num_decoder_layers:
                num_decoder_residual_layers = num_encoder_residual_layers
    _add_argument(hparams, 'num_encoder_residual_layers',
        num_encoder_residual_layers)
    _add_argument(hparams, 'num_decoder_residual_layers',
        num_decoder_residual_layers)
    if hparams.language_model:
        hparams.attention = ''
        hparams.attention_architecture = ''
        hparams.pass_hidden_state = False
        hparams.share_vocab = True
        hparams.src = hparams.tgt
        utils.print_out(
            'For language modeling, we turn off attention and pass_hidden_state; turn on share_vocab; set src to tgt.'
            )
    if hparams.vocab_prefix:
        src_vocab_file = hparams.vocab_prefix + '.' + hparams.src
        tgt_vocab_file = hparams.vocab_prefix + '.' + hparams.tgt
    else:
        raise ValueError('hparams.vocab_prefix must be provided.')
    src_vocab_size, src_vocab_file = vocab_utils.check_vocab(src_vocab_file,
        hparams.output_dir, check_special_token=hparams.check_special_token,
        sos=hparams.sos, eos=hparams.eos, unk=vocab_utils.UNK, pad_vocab=True)
    if hparams.share_vocab:
        utils.print_out('  using source vocab for target')
        tgt_vocab_file = src_vocab_file
        tgt_vocab_size = src_vocab_size
    else:
        tgt_vocab_size, tgt_vocab_file = vocab_utils.check_vocab(tgt_vocab_file
            , hparams.output_dir, check_special_token=hparams.
            check_special_token, sos=hparams.sos, eos=hparams.eos, unk=
            vocab_utils.UNK)
    _add_argument(hparams, 'src_vocab_size', src_vocab_size)
    _add_argument(hparams, 'tgt_vocab_size', tgt_vocab_size)
    _add_argument(hparams, 'src_vocab_file', src_vocab_file)
    _add_argument(hparams, 'tgt_vocab_file', tgt_vocab_file)
    _add_argument(hparams, 'num_enc_emb_partitions', hparams.
        num_embeddings_partitions)
    _add_argument(hparams, 'num_dec_emb_partitions', hparams.
        num_embeddings_partitions)
    _add_argument(hparams, 'src_embed_file', '')
    _add_argument(hparams, 'tgt_embed_file', '')
    if hparams.embed_prefix:
        src_embed_file = hparams.embed_prefix + '.' + hparams.src
        tgt_embed_file = hparams.embed_prefix + '.' + hparams.tgt
        if tf.gfile.Exists(src_embed_file):
            utils.print_out('  src_embed_file %s exist' % src_embed_file)
            hparams.src_embed_file = src_embed_file
            utils.print_out(
                'For pretrained embeddings, set num_enc_emb_partitions to 1')
            hparams.num_enc_emb_partitions = 1
        else:
            utils.print_out("  src_embed_file %s doesn't exist" %
                src_embed_file)
        if tf.gfile.Exists(tgt_embed_file):
            utils.print_out('  tgt_embed_file %s exist' % tgt_embed_file)
            hparams.tgt_embed_file = tgt_embed_file
            utils.print_out(
                'For pretrained embeddings, set num_dec_emb_partitions to 1')
            hparams.num_dec_emb_partitions = 1
        else:
            utils.print_out("  tgt_embed_file %s doesn't exist" %
                tgt_embed_file)
    metric = 'bleu'
    best_metric_dir = os.path.join(hparams.output_dir, 'best_' + metric)
    tf.gfile.MakeDirs(best_metric_dir)
    _add_argument(hparams, 'best_' + metric, 0, update=False)
    _add_argument(hparams, 'best_' + metric + '_dir', best_metric_dir)
    return hparams


def create_or_load_hparams(default_hparams, hparams_path):
    """Create hparams or load hparams from output_dir."""
    hparams = utils.maybe_parse_standard_hparams(default_hparams, hparams_path)
    hparams = extend_hparams(hparams)
    utils.print_hparams(hparams)
    return hparams


def run_main(flags, default_hparams, estimator_fn):
    """Run main."""
    random_seed = flags.random_seed
    if random_seed is not None and random_seed > 0:
        utils.print_out('# Set random seed to %d' % random_seed)
        random.seed(random_seed)
        np.random.seed(random_seed)
        tf.set_random_seed(random_seed)
    output_dir = flags.output_dir
    if output_dir and not tf.gfile.Exists(output_dir):
        utils.print_out('# Creating output directory %s ...' % output_dir)
        tf.gfile.MakeDirs(output_dir)
    hparams = create_or_load_hparams(default_hparams, flags.hparams_path)
    estimator_fn(hparams)
    return hparams


def tokenize(hparams, file, tokenized_file):
    utils.print_out('tokenizing {} -> {}'.format(file, tokenized_file))
    with open(file, 'rb') as input_file:
        with open(tokenized_file, 'wb') as output_file:
            subprocess.run([hparams.tokenizer_file, '-l', hparams.src],
                stdin=input_file, stdout=output_file)


def detokenize(hparams, file, detokenized_file):
    utils.print_out('detokenizing {} -> {}'.format(file, detokenized_file))
    with open(file, 'rb') as input_file:
        with open(detokenized_file, 'wb') as output_file:
            subprocess.run([hparams.detokenizer_file, '-l', hparams.tgt],
                stdin=input_file, stdout=output_file)


def main(unused_argv):
    experiment_start = time.time()
    tf.logging.set_verbosity(tf.logging.INFO)
    if FLAGS.use_fp16 and FLAGS.use_dist_strategy:
        raise ValueError("use_fp16 and use_dist_strategy aren't compatible")
    if FLAGS.use_fp16 + FLAGS.amp + FLAGS.use_fastmath > 1:
        raise ValueError('Only one of use_fp16, amp, use_fastmath can be set')
    if FLAGS.amp:
        utils.print_out('Enabling TF-AMP')
        os.environ['TF_ENABLE_AUTO_MIXED_PRECISION'] = '1'
    if FLAGS.use_fastmath:
        utils.print_out('Enabling FastMath')
        os.environ['TF_ENABLE_CUBLAS_TENSOR_OP_MATH_FP32'] = '1'
        os.environ['TF_ENABLE_CUDNN_TENSOR_OP_MATH_FP32'] = '1'
        os.environ['TF_ENABLE_CUDNN_RNN_TENSOR_OP_MATH_FP32'] = '1'
    active_xla_option_nums = np.sum([FLAGS.use_xla, FLAGS.use_autojit_xla,
        FLAGS.xla_compile])
    if active_xla_option_nums > 1:
        raise ValueError(
            'Only one of use_xla, xla_compile, use_autojit_xla can be set')
    os.environ['use_xla'] = str(FLAGS.use_xla).lower()
    if FLAGS.use_xla:
        os.environ['use_defun'] = str(True).lower()
    else:
        os.environ['use_defun'] = str(FLAGS.use_defun).lower()
    utils.print_out('use_defun is %s for attention' % os.environ['use_defun'])
    os.environ['use_fp32_batch_matmul'
        ] = 'true' if FLAGS.use_fp32_batch_matmul else 'false'
    os.environ['xla_compile'] = 'true' if FLAGS.xla_compile else 'false'
    os.environ['force_inputs_padding'
        ] = 'true' if FLAGS.force_inputs_padding else 'false'
    if FLAGS.mode == 'train':
        utils.print_out('Running training mode.')
        default_hparams = create_hparams(FLAGS)
        run_main(FLAGS, default_hparams, estimator.train_fn)
    elif FLAGS.mode == 'infer' or FLAGS.mode == 'translate':
        if FLAGS.mode == 'infer':
            utils.print_out('Running inference mode.')
            translate_mode = False
        else:
            utils.print_out('Running translate mode on file {}.'.format(
                FLAGS.translate_file))
            translate_mode = True
        random_seed = FLAGS.random_seed
        if random_seed is not None and random_seed > 0:
            utils.print_out('# Set random seed to %d' % random_seed)
            random.seed(random_seed)
            np.random.seed(random_seed)
            tf.set_random_seed(random_seed)
        output_dir = FLAGS.output_dir
        if output_dir and not tf.gfile.Exists(output_dir):
            utils.print_out('# Creating output directory %s ...' % output_dir)
            tf.gfile.MakeDirs(output_dir)
        dllogger.init(backends=[dllogger.StdOutBackend(dllogger.Verbosity.
            DEFAULT), dllogger.JSONStreamBackend(dllogger.Verbosity.VERBOSE,
            os.path.join(FLAGS.output_dir, FLAGS.mode + '-report.json'))])
        dllogger.log('PARAMETER', vars(FLAGS))
        default_hparams = create_hparams(FLAGS)
        default_hparams.num_buckets = 1
        hparams = create_or_load_hparams(default_hparams, FLAGS.hparams_path)
        utils.print_out('infer_hparams:')
        utils.print_hparams(hparams)
        if translate_mode:
            tokenize(hparams, hparams.translate_file, hparams.
                translate_file + '.tok')
        eval_sentences, eval_src_tokens, _ = (iterator_utils.
            get_effective_epoch_size(hparams, train=False))
        tf.logging.info('Starting to evaluate...')
        eval_start = time.time()
        _, (eval_speed, eval_latencies
            ), eval_output_tokens = estimator.eval_fn(hparams, hparams.ckpt,
            only_translate=translate_mode)
        eval_end = time.time()
        eval_delta = eval_end - eval_start
        utils.print_out(
            'eval time for ckpt: %.2f mins (%.2f sent/sec, %.2f tokens/sec)' %
            (eval_delta / 60.0, eval_speed, eval_speed * (eval_src_tokens +
            eval_output_tokens) / eval_sentences), f=sys.stderr)
        logging_data = {'infer_speed_sent': eval_speed, 'infer_speed_toks':
            eval_speed * (eval_src_tokens + eval_output_tokens) /
            eval_sentences}
        for lat in sorted(eval_latencies):
            utils.print_out('eval latency_%s for ckpt: %.2f ms' % (lat, 
                eval_latencies[lat] * 1000))
            logging_data['infer_latency_{}'.format(lat)] = eval_latencies[lat
                ] * 1000
        dllogger.log((), logging_data)
        dllogger.flush()
        if translate_mode:
            detokenize(hparams, hparams.translate_file + '.trans.tok', 
                hparams.translate_file + '.trans')
    else:
        assert FLAGS.mode == 'train_and_eval'
        utils.print_out('Running train and eval mode.')
        random_seed = FLAGS.random_seed
        if random_seed is not None and random_seed > 0:
            utils.print_out('# Set random seed to %d' % random_seed)
            random.seed(random_seed)
            np.random.seed(random_seed)
            tf.set_random_seed(random_seed)
        output_dir = FLAGS.output_dir
        if output_dir and not tf.gfile.Exists(output_dir):
            utils.print_out('# Creating output directory %s ...' % output_dir)
            tf.gfile.MakeDirs(output_dir)
        dllogger.init(backends=[dllogger.StdOutBackend(dllogger.Verbosity.
            DEFAULT), dllogger.JSONStreamBackend(dllogger.Verbosity.VERBOSE,
            os.path.join(FLAGS.output_dir, FLAGS.mode + '-report.json'))])
        dllogger.log('PARAMETER', vars(FLAGS))
        dllogger.metadata('bleu', {'unit': None})
        dllogger.metadata('train_speed_sent', {'unit': 'sequences/s'})
        dllogger.metadata('train_speed_toks', {'unit': 'tokens/s'})
        default_hparams = create_hparams(FLAGS)
        hparams = create_or_load_hparams(default_hparams, FLAGS.hparams_path)
        utils.print_out('training hparams:')
        utils.print_hparams(hparams)
        with tf.gfile.GFile(os.path.join(output_dir, 'train_hparams.txt'), 'w'
            ) as f:
            f.write(utils.serialize_hparams(hparams) + '\n')
        infer_hparams = tf.contrib.training.HParams(**hparams.values())
        infer_hparams.num_buckets = 1
        utils.print_out('infer_hparams:')
        utils.print_hparams(infer_hparams)
        with tf.gfile.GFile(os.path.join(output_dir, 'infer_hparams.txt'), 'w'
            ) as f:
            f.write(utils.serialize_hparams(infer_hparams) + '\n')
        epochs = 0
        should_stop = epochs >= FLAGS.max_train_epochs
        train_sentences, train_src_tokens, train_tgt_tokens = (iterator_utils
            .get_effective_epoch_size(hparams))
        eval_sentences, eval_src_tokens, _ = (iterator_utils.
            get_effective_epoch_size(hparams, train=False))
        while not should_stop:
            utils.print_out('Starting epoch %d' % epochs)
            try:
                train_start = time.time()
                train_speed, _ = estimator.train_fn(hparams)
            except tf.errors.OutOfRangeError:
                utils.print_out('training hits OutOfRangeError', f=sys.stderr)
            train_end = time.time()
            train_delta = train_end - train_start
            utils.print_out(
                'training time for epoch %d: %.2f mins (%.2f sent/sec, %.2f tokens/sec)'
                 % (epochs + 1, train_delta / 60.0, train_speed, 
                train_speed * (train_src_tokens + train_tgt_tokens) /
                train_sentences), f=sys.stderr)
            logging_data = {'train_speed_sent': train_speed,
                'train_speed_toks': train_speed * (train_src_tokens +
                train_tgt_tokens) / train_sentences}
            eval_start = time.time()
            bleu_score, (eval_speed, eval_latencies
                ), eval_output_tokens = estimator.eval_fn(infer_hparams)
            eval_end = time.time()
            eval_delta = eval_end - eval_start
            utils.print_out(
                'eval time for epoch %d: %.2f mins (%.2f sent/sec, %.2f tokens/sec)'
                 % (epochs + 1, eval_delta / 60.0, eval_speed, eval_speed *
                (eval_src_tokens + eval_output_tokens) / eval_sentences), f
                =sys.stderr)
            logging_data.update({'bleu': bleu_score, 'infer_speed_sent':
                eval_speed, 'infer_speed_toks': eval_speed * (
                eval_src_tokens + eval_output_tokens) / eval_sentences})
            for lat in sorted(eval_latencies):
                utils.print_out('eval latency_%s for epoch %d: %.2f ms' % (
                    lat, epochs + 1, eval_latencies[lat] * 1000))
                logging_data['eval_latency_{}'.format(lat)] = eval_latencies[
                    lat] * 1000
            dllogger.log((epochs,), logging_data)
            dllogger.flush()
            if (FLAGS.debug or FLAGS.target_bleu is not None and bleu_score >
                FLAGS.target_bleu):
                should_stop = True
                utils.print_out(
                    'Stop job since target bleu is reached at epoch %d .' %
                    epochs, f=sys.stderr)
            epochs += 1
            if epochs >= FLAGS.max_train_epochs:
                should_stop = True
                utils.print_out('Stop job since max_train_epochs is reached.',
                    f=sys.stderr)
        dllogger.log((), logging_data)
        dllogger.flush()
    experiment_end = time.time()
    utils.print_out('Experiment took {} min'.format((experiment_end -
        experiment_start) / 60))


if __name__ == '__main__':
    nmt_parser = argparse.ArgumentParser()
    add_arguments(nmt_parser)
    FLAGS, unparsed = nmt_parser.parse_known_args()
    tf.app.run(main=main, argv=[sys.argv[0]] + unparsed)

<mask token>
utils.check_tensorflow_version()
FLAGS = None


def add_arguments(parser):
    """Build ArgumentParser."""
    parser.register('type', 'bool', lambda v: v.lower() == 'true')
    parser.add_argument('--num_units', type=int, default=1024, help=
        'Network size.')
    parser.add_argument('--num_layers', type=int, default=4, help=
        'Network depth.')
    parser.add_argument('--num_encoder_layers', type=int, default=None,
        help='Encoder depth, equal to num_layers if None.')
    parser.add_argument('--num_decoder_layers', type=int, default=None,
        help='Decoder depth, equal to num_layers if None.')
    parser.add_argument('--encoder_type', type=str, default='gnmt', help=
        """      uni | bi | gnmt.
      For bi, we build num_encoder_layers/2 bi-directional layers.
      For gnmt, we build 1 bi-directional layer, and (num_encoder_layers - 1)
        uni-directional layers.      """
        )
    parser.add_argument('--residual', type='bool', nargs='?', const=True,
        default=True, help='Whether to add residual connections.')
    parser.add_argument('--time_major', type='bool', nargs='?', const=True,
        default=True, help='Whether to use time-major mode for dynamic RNN.')
    parser.add_argument('--num_embeddings_partitions', type=int, default=0,
        help='Number of partitions for embedding vars.')
    parser.add_argument('--attention', type=str, default='normed_bahdanau',
        help=
        """      luong | scaled_luong | bahdanau | normed_bahdanau or set to "" for no
      attention      """
        )
    parser.add_argument('--attention_architecture', type=str, default=
        'gnmt_v2', help=
        """      standard | gnmt | gnmt_v2.
      standard: use top layer to compute attention.
      gnmt: GNMT style of computing attention, use previous bottom layer to
          compute attention.
      gnmt_v2: similar to gnmt, but use current bottom layer to compute
          attention.      """
        )
    parser.add_argument('--output_attention', type='bool', nargs='?', const
        =True, default=True, help=
        """      Only used in standard attention_architecture. Whether use attention as
      the cell output at each timestep.
      .      """
        )
    parser.add_argument('--pass_hidden_state', type='bool', nargs='?',
        const=True, default=True, help=
        """      Whether to pass encoder's hidden state to decoder when using an attention
      based model.      """
        )
    parser.add_argument('--optimizer', type=str, default='adam', help=
        'sgd | adam')
    parser.add_argument('--learning_rate', type=float, default=0.0005, help
        ='Learning rate. Adam: 0.001 | 0.0001')
    parser.add_argument('--warmup_steps', type=int, default=200, help=
        'How many steps we inverse-decay learning.')
    parser.add_argument('--warmup_scheme', type=str, default='t2t', help=
        """      How to warmup learning rates. Options include:
        t2t: Tensor2Tensor's way, start with lr 100 times smaller, then
             exponentiate until the specified lr.      """
        )
    parser.add_argument('--decay_scheme', type=str, default='luong234',
        help=
        """      How we decay learning rate. Options include:
        luong234: after 2/3 num train steps, we start halving the learning rate
          for 4 times before finishing.
        luong5: after 1/2 num train steps, we start halving the learning rate
          for 5 times before finishing.        luong10: after 1/2 num train steps, we start halving the learning rate
          for 10 times before finishing.      """
        )
    parser.add_argument('--max_train_epochs', type=int, default=6, help=
        'Max number of epochs.')
    parser.add_argument('--target_bleu', type=float, default=None, help=
        'Target bleu.')
    parser.add_argument('--colocate_gradients_with_ops', type='bool', nargs
        ='?', const=True, default=True, help=
        'Whether try colocating gradients with corresponding op')
    parser.add_argument('--label_smoothing', type=float, default=0.1, help=
        'If nonzero, smooth the labels towards 1/num_classes.')
    parser.add_argument('--init_op', type=str, default='uniform', help=
        'uniform | glorot_normal | glorot_uniform')
    parser.add_argument('--init_weight', type=float, default=0.1, help=
        'for uniform init_op, initialize weights between [-this, this].')
    parser.add_argument('--src', type=str, default='en', help=
        'Source suffix, e.g., en.')
    parser.add_argument('--tgt', type=str, default='de', help=
        'Target suffix, e.g., de.')
    parser.add_argument('--data_dir', type=str, default='data/wmt16_de_en',
        help='Training/eval data directory.')
    parser.add_argument('--train_prefix', type=str, default=
        'train.tok.clean.bpe.32000', help=
        'Train prefix, expect files with src/tgt suffixes.')
    parser.add_argument('--test_prefix', type=str, default=
        'newstest2014.tok.bpe.32000', help=
        'Test prefix, expect files with src/tgt suffixes.')
    parser.add_argument('--translate_file', type=str, help=
        'File to translate, works only with translate mode')
    parser.add_argument('--output_dir', type=str, default='results', help=
        'Store log/model files.')
    parser.add_argument('--vocab_prefix', type=str, default=
        'vocab.bpe.32000', help=
        '      Vocab prefix, expect files with src/tgt suffixes.      ')
    parser.add_argument('--embed_prefix', type=str, default=None, help=
        """      Pretrained embedding prefix, expect files with src/tgt suffixes.
      The embedding files should be Glove formatted txt files.      """
        )
    parser.add_argument('--sos', type=str, default='<s>', help=
        'Start-of-sentence symbol.')
    parser.add_argument('--eos', type=str, default='</s>', help=
        'End-of-sentence symbol.')
    parser.add_argument('--share_vocab', type='bool', nargs='?', const=True,
        default=True, help=
        """      Whether to use the source vocab and embeddings for both source and
      target.      """
        )
    parser.add_argument('--check_special_token', type='bool', default=True,
        help=
        """                      Whether check special sos, eos, unk tokens exist in the
                      vocab files.                      """
        )
    parser.add_argument('--src_max_len', type=int, default=50, help=
        'Max length of src sequences during training (including EOS).')
    parser.add_argument('--tgt_max_len', type=int, default=50, help=
        'Max length of tgt sequences during training (including BOS).')
    parser.add_argument('--src_max_len_infer', type=int, default=None, help
        ='Max length of src sequences during inference (including EOS).')
    parser.add_argument('--tgt_max_len_infer', type=int, default=80, help=
        """      Max length of tgt sequences during inference (including BOS).  Also use to restrict the
      maximum decoding length.      """
        )
    parser.add_argument('--unit_type', type=str, default='lstm', help=
        'lstm | gru | layer_norm_lstm | nas')
    parser.add_argument('--forget_bias', type=float, default=0.0, help=
        'Forget bias for BasicLSTMCell.')
    parser.add_argument('--dropout', type=float, default=0.2, help=
        'Dropout rate (not keep_prob)')
    parser.add_argument('--max_gradient_norm', type=float, default=5.0,
        help='Clip gradients to this norm.')
    parser.add_argument('--batch_size', type=int, default=128, help=
        'Total batch size.')
    parser.add_argument('--num_buckets', type=int, default=5, help=
        'Put data into similar-length buckets (only for training).')
    parser.add_argument('--subword_option', type=str, default='bpe',
        choices=['', 'bpe', 'spm'], help=
        '                      Set to bpe or spm to activate subword desegmentation.                      '
        )
    parser.add_argument('--use_char_encode', type='bool', default=False,
        help=
        """                      Whether to split each word or bpe into character, and then
                      generate the word-level representation from the character
                      reprentation.
                      """
        )
    parser.add_argument('--save_checkpoints_steps', type=int, default=2000,
        help='save_checkpoints_steps')
    parser.add_argument('--log_step_count_steps', type=int, default=10,
        help=
        'The frequency, in number of global steps, that the global step and the loss will be logged during training'
        )
    parser.add_argument('--num_gpus', type=int, default=1, help=
        'Number of gpus in each worker.')
    parser.add_argument('--hparams_path', type=str, default=None, help=
        'Path to standard hparams json file that overrideshparams values from FLAGS.'
        )
    parser.add_argument('--random_seed', type=int, default=1, help=
        'Random seed (>0, set a specific seed).')
    parser.add_argument('--language_model', type='bool', nargs='?', const=
        True, default=False, help=
        'True to train a language model, ignoring encoder')
    parser.add_argument('--ckpt', type=str, default=None, help=
        'Checkpoint file to load a model for inference. (defaults to newest checkpoint)'
        )
    parser.add_argument('--infer_batch_size', type=int, default=128, help=
        'Batch size for inference mode.')
    parser.add_argument('--detokenizer_file', type=str, default=None, help=
        'Detokenizer script file. Default: DATA_DIR/mosesdecoder/scripts/tokenizer/detokenizer.perl'
        )
    parser.add_argument('--tokenizer_file', type=str, default=None, help=
        'Tokenizer script file. Default: DATA_DIR/mosesdecoder/scripts/tokenizer/tokenizer.perl'
        )
    parser.add_argument('--infer_mode', type=str, default='beam_search',
        choices=['greedy', 'beam_search'], help=
        'Which type of decoder to use during inference.')
    parser.add_argument('--beam_width', type=int, default=5, help=
        """      beam width when using beam search decoder. If 0, use standard
      decoder with greedy helper.      """
        )
    parser.add_argument('--length_penalty_weight', type=float, default=0.6,
        help='Length penalty for beam search.')
    parser.add_argument('--coverage_penalty_weight', type=float, default=
        0.1, help='Coverage penalty for beam search.')
    parser.add_argument('--num_workers', type=int, default=1, help=
        'Number of workers (inference only).')
    parser.add_argument('--amp', action='store_true', help=
        'use amp for training and inference')
    parser.add_argument('--use_fastmath', type='bool', default=False, help=
        'use_fastmath for training and inference')
    parser.add_argument('--use_fp16', type='bool', default=False, help=
        'use_fp16 for training and inference')
    parser.add_argument('--fp16_loss_scale', type=float, default=128, help=
        'If fp16 is enabled, the loss is multiplied by this amount right before gradients are computed, then each gradient is divided by this amount. Mathematically, this has no effect, but it helps avoid fp16 underflow. Set to 1 to effectively disable.'
        )
    parser.add_argument('--enable_auto_loss_scale', type='bool', default=
        True, help=
        'If True and use_fp16 is True, automatically adjust the loss scale during training.'
        )
    parser.add_argument('--fp16_inc_loss_scale_every_n', type=int, default=
        128, help=
        'If fp16 is enabled and enable_auto_loss_scale is True, increase the loss scale every n steps.'
        )
    parser.add_argument('--check_tower_loss_numerics', type='bool', default
        =False, help='whether to check tower loss numerics')
    parser.add_argument('--use_fp32_batch_matmul', type='bool', default=
        False, help='Whether to use fp32 batch matmul')
    parser.add_argument('--force_inputs_padding', type='bool', default=
        False, help='Force padding input batch to src_max_len and tgt_max_len')
    parser.add_argument('--use_xla', type='bool', default=False, help=
        'Use xla to compile a few selected locations, mostly Defuns.')
    parser.add_argument('--xla_compile', type='bool', default=False, help=
        "Use xla.compile() for each tower's fwd and bak pass.")
    parser.add_argument('--use_autojit_xla', type='bool', default=False,
        help='Use auto jit xla.')
    parser.add_argument('--use_pintohost_optimizer', type='bool', default=
        False, help='whether to use PinToHost optimizer')
    parser.add_argument('--use_cudnn_lstm', type='bool', default=False,
        help='whether to use cudnn_lstm for encoder, non residual layers')
    parser.add_argument('--use_loose_bidi_cudnn_lstm', type='bool', default
        =False, help='whether to use loose bidi cudnn_lstm')
    parser.add_argument('--use_fused_lstm', type='bool', default=True, help
        =
        'whether to use fused lstm and variant. If enabled, training will use LSTMBlockFusedCell, infer will use LSTMBlockCell when appropriate.'
        )
    parser.add_argument('--use_fused_lstm_dec', type='bool', default=False,
        help='whether to use fused lstm for decoder (training only).')
    parser.add_argument('--gpu_indices', type=str, default='', help=
        'Indices of worker GPUs in ring order')
    parser.add_argument('--parallel_iterations', type=int, default=10, help
        ='number of parallel iterations in dynamic_rnn')
    parser.add_argument('--use_dist_strategy', type='bool', default=False,
        help='whether to use distribution strategy')
    parser.add_argument('--hierarchical_copy', type='bool', default=False,
        help=
        'Use hierarchical copies. Currently only optimized for use on a DGX-1 with 8 GPUs and may perform poorly on other hardware. Requires --num_gpus > 1, and only recommended when --num_gpus=8'
        )
    parser.add_argument('--network_topology', type=constants.
        NetworkTopology, default=constants.NetworkTopology.DGX1, choices=
        list(constants.NetworkTopology))
    parser.add_argument('--use_block_lstm', type='bool', default=False,
        help='whether to use block lstm')
    parser.add_argument('--use_defun', type='bool', default=False, help=
        'whether to use Defun')
    parser.add_argument('--gradient_repacking', type=int, default=0, help=
        'Use gradient repacking. Itcurrently only works with replicated mode. At the end ofof each step, it repacks the gradients for more efficientcross-device transportation. A non-zero value specifiesthe number of split packs that will be formed.'
        )
    parser.add_argument('--compact_gradient_transfer', type='bool', default
        =True, help=
        'Compact gradient as much as possible for cross-device transfer and aggregation.'
        )
    parser.add_argument('--all_reduce_spec', type=str, default='nccl', help
        =
        """A specification of the all_reduce algorithm to be used for reducing gradients.  For more details, see parse_all_reduce_spec in variable_mgr.py.  An all_reduce_spec has BNF form:
int ::= positive whole number
g_int ::= int[KkMGT]?
alg_spec ::= alg | alg#int
range_spec ::= alg_spec | alg_spec/alg_spec
spec ::= range_spec | range_spec:g_int:range_spec
NOTE: not all syntactically correct constructs are supported.

Examples:
 "xring" == use one global ring reduction for all tensors
"pscpu" == use CPU at worker 0 to reduce all tensors
"nccl" == use NCCL to locally reduce all tensors.  Limited to 1 worker.
"nccl/xring" == locally (to one worker) reduce values using NCCL then ring reduce across workers.
"pscpu:32k:xring" == use pscpu algorithm for tensors of size up to 32kB, then xring for larger tensors."""
        )
    parser.add_argument('--agg_small_grads_max_bytes', type=int, default=0,
        help=
        'If > 0, try to aggregate tensors of less than this number of bytes prior to all-reduce.'
        )
    parser.add_argument('--agg_small_grads_max_group', type=int, default=10,
        help=
        'When aggregating small tensors for all-reduce do not aggregate more than this many into one new tensor.'
        )
    parser.add_argument('--allreduce_merge_scope', type=int, default=1,
        help=
        'Establish a name scope around this many gradients prior to creating the all-reduce operations. It may affect the ability of the backend to merge parallel ops.'
        )
    parser.add_argument('--local_parameter_device', type=str, default='gpu',
        help=
        'Device to use as parameter server: cpu or gpu. For distributed training, it can affect where caching of variables happens.'
        )
    parser.add_argument('--use_resource_vars', type='bool', default=False,
        help=
        'Use resource variables instead of normal variables. Resource variables are slower, but this option is useful for debugging their performance.'
        )
    parser.add_argument('--debug', type='bool', default=False, help=
        'Debug train and eval')
    parser.add_argument('--debug_num_train_steps', type=int, default=None,
        help='Num steps to train.')
    parser.add_argument('--show_metrics', type='bool', default=True, help=
        'whether to show detailed metrics')
    parser.add_argument('--clip_grads', type='bool', default=True, help=
        'whether to clip gradients')
    parser.add_argument('--profile', type='bool', default=False, help=
        'If generate profile')
    parser.add_argument('--profile_save_steps', type=int, default=10, help=
        'Save timeline every N steps.')
    parser.add_argument('--use_dynamic_rnn', type='bool', default=True)
    parser.add_argument('--use_synthetic_data', type='bool', default=False)
    parser.add_argument('--mode', type=str, default='train_and_eval',
        choices=('train_and_eval', 'infer', 'translate'))


def create_hparams(flags):
    """Create training hparams."""
    return tf.contrib.training.HParams(src=flags.src, tgt=flags.tgt,
        train_prefix=os.path.join(flags.data_dir, flags.train_prefix),
        test_prefix=os.path.join(flags.data_dir, flags.test_prefix),
        translate_file=flags.translate_file, vocab_prefix=os.path.join(
        flags.data_dir, flags.vocab_prefix), embed_prefix=flags.
        embed_prefix, output_dir=flags.output_dir, num_units=flags.
        num_units, num_encoder_layers=flags.num_encoder_layers or flags.
        num_layers, num_decoder_layers=flags.num_decoder_layers or flags.
        num_layers, dropout=flags.dropout, unit_type=flags.unit_type,
        encoder_type=flags.encoder_type, residual=flags.residual,
        time_major=flags.time_major, num_embeddings_partitions=flags.
        num_embeddings_partitions, attention=flags.attention,
        attention_architecture=flags.attention_architecture,
        output_attention=flags.output_attention, pass_hidden_state=flags.
        pass_hidden_state, optimizer=flags.optimizer, max_train_epochs=
        flags.max_train_epochs, target_bleu=flags.target_bleu,
        label_smoothing=flags.label_smoothing, batch_size=flags.batch_size,
        init_op=flags.init_op, init_weight=flags.init_weight,
        max_gradient_norm=flags.max_gradient_norm, learning_rate=flags.
        learning_rate, warmup_steps=flags.warmup_steps, warmup_scheme=flags
        .warmup_scheme, decay_scheme=flags.decay_scheme,
        colocate_gradients_with_ops=flags.colocate_gradients_with_ops,
        num_buckets=flags.num_buckets, src_max_len=flags.src_max_len,
        tgt_max_len=flags.tgt_max_len, src_max_len_infer=flags.
        src_max_len_infer, tgt_max_len_infer=flags.tgt_max_len_infer, ckpt=
        flags.ckpt, infer_batch_size=flags.infer_batch_size,
        detokenizer_file=flags.detokenizer_file if flags.detokenizer_file
         is not None else os.path.join(flags.data_dir,
        'mosesdecoder/scripts/tokenizer/detokenizer.perl'), tokenizer_file=
        flags.tokenizer_file if flags.tokenizer_file is not None else os.
        path.join(flags.data_dir,
        'mosesdecoder/scripts/tokenizer/tokenizer.perl'), infer_mode=flags.
        infer_mode, beam_width=flags.beam_width, length_penalty_weight=
        flags.length_penalty_weight, coverage_penalty_weight=flags.
        coverage_penalty_weight, sos=flags.sos if flags.sos else
        vocab_utils.SOS, eos=flags.eos if flags.eos else vocab_utils.EOS,
        subword_option=flags.subword_option, check_special_token=flags.
        check_special_token, use_char_encode=flags.use_char_encode,
        forget_bias=flags.forget_bias, num_gpus=flags.num_gpus,
        save_checkpoints_steps=flags.save_checkpoints_steps,
        log_step_count_steps=flags.log_step_count_steps, epoch_step=0,
        share_vocab=flags.share_vocab, random_seed=flags.random_seed,
        language_model=flags.language_model, amp=flags.amp, use_fastmath=
        flags.use_fastmath, use_fp16=flags.use_fp16, fp16_loss_scale=flags.
        fp16_loss_scale, enable_auto_loss_scale=flags.
        enable_auto_loss_scale, fp16_inc_loss_scale_every_n=flags.
        fp16_inc_loss_scale_every_n, check_tower_loss_numerics=flags.
        check_tower_loss_numerics, use_fp32_batch_matmul=flags.
        use_fp32_batch_matmul, force_inputs_padding=flags.
        force_inputs_padding, use_xla=flags.use_xla, xla_compile=flags.
        xla_compile, use_autojit_xla=flags.use_autojit_xla,
        use_pintohost_optimizer=flags.use_pintohost_optimizer,
        use_cudnn_lstm=flags.use_cudnn_lstm, use_loose_bidi_cudnn_lstm=
        flags.use_loose_bidi_cudnn_lstm, use_fused_lstm=flags.
        use_fused_lstm, use_fused_lstm_dec=flags.use_fused_lstm_dec,
        gpu_indices=flags.gpu_indices, parallel_iterations=flags.
        parallel_iterations, use_dynamic_rnn=flags.use_dynamic_rnn,
        use_dist_strategy=flags.use_dist_strategy, hierarchical_copy=flags.
        hierarchical_copy, network_topology=flags.network_topology,
        use_block_lstm=flags.use_block_lstm, gradient_repacking=flags.
        gradient_repacking, compact_gradient_transfer=flags.
        compact_gradient_transfer, all_reduce_spec=flags.all_reduce_spec,
        agg_small_grads_max_bytes=flags.agg_small_grads_max_bytes,
        agg_small_grads_max_group=flags.agg_small_grads_max_group,
        allreduce_merge_scope=flags.allreduce_merge_scope,
        local_parameter_device='cpu' if flags.num_gpus == 0 else flags.
        local_parameter_device, use_resource_vars=flags.use_resource_vars,
        debug=flags.debug, debug_num_train_steps=flags.
        debug_num_train_steps, clip_grads=flags.clip_grads, profile=flags.
        profile, profile_save_steps=flags.profile_save_steps, show_metrics=
        flags.show_metrics, use_synthetic_data=flags.use_synthetic_data,
        mode=flags.mode)


def _add_argument(hparams, key, value, update=True):
    """Add an argument to hparams; if exists, change the value if update==True."""
    if hasattr(hparams, key):
        if update:
            setattr(hparams, key, value)
    else:
        hparams.add_hparam(key, value)


def extend_hparams(hparams):
    """Add new arguments to hparams."""
    if hparams.encoder_type == 'bi' and hparams.num_encoder_layers % 2 != 0:
        raise ValueError('For bi, num_encoder_layers %d should be even' %
            hparams.num_encoder_layers)
    if hparams.attention_architecture in ['gnmt'
        ] and hparams.num_encoder_layers < 2:
        raise ValueError(
            'For gnmt attention architecture, num_encoder_layers %d should be >= 2'
             % hparams.num_encoder_layers)
    if hparams.subword_option and hparams.subword_option not in ['spm', 'bpe']:
        raise ValueError('subword option must be either spm, or bpe')
    if hparams.infer_mode == 'beam_search' and hparams.beam_width <= 0:
        raise ValueError(
            'beam_width must greater than 0 when using beam_searchdecoder.')
    if hparams.mode == 'translate' and not hparams.translate_file:
        raise ValueError(
            '--translate_file flag must be specified in translate mode')
    assert hparams.num_encoder_layers and hparams.num_decoder_layers
    if hparams.num_encoder_layers != hparams.num_decoder_layers:
        hparams.pass_hidden_state = False
        utils.print_out(
            'Num encoder layer %d is different from num decoder layer %d, so set pass_hidden_state to False'
             % (hparams.num_encoder_layers, hparams.num_decoder_layers))
    num_encoder_residual_layers = 0
    num_decoder_residual_layers = 0
    if hparams.residual:
        if hparams.num_encoder_layers > 1:
            num_encoder_residual_layers = hparams.num_encoder_layers - 1
        if hparams.num_decoder_layers > 1:
            num_decoder_residual_layers = hparams.num_decoder_layers - 1
        if hparams.encoder_type == 'gnmt':
            num_encoder_residual_layers = hparams.num_encoder_layers - 2
            if hparams.num_encoder_layers == hparams.num_decoder_layers:
                num_decoder_residual_layers = num_encoder_residual_layers
    _add_argument(hparams, 'num_encoder_residual_layers',
        num_encoder_residual_layers)
    _add_argument(hparams, 'num_decoder_residual_layers',
        num_decoder_residual_layers)
    if hparams.language_model:
        hparams.attention = ''
        hparams.attention_architecture = ''
        hparams.pass_hidden_state = False
        hparams.share_vocab = True
        hparams.src = hparams.tgt
        utils.print_out(
            'For language modeling, we turn off attention and pass_hidden_state; turn on share_vocab; set src to tgt.'
            )
    if hparams.vocab_prefix:
        src_vocab_file = hparams.vocab_prefix + '.' + hparams.src
        tgt_vocab_file = hparams.vocab_prefix + '.' + hparams.tgt
    else:
        raise ValueError('hparams.vocab_prefix must be provided.')
    src_vocab_size, src_vocab_file = vocab_utils.check_vocab(src_vocab_file,
        hparams.output_dir, check_special_token=hparams.check_special_token,
        sos=hparams.sos, eos=hparams.eos, unk=vocab_utils.UNK, pad_vocab=True)
    if hparams.share_vocab:
        utils.print_out('  using source vocab for target')
        tgt_vocab_file = src_vocab_file
        tgt_vocab_size = src_vocab_size
    else:
        tgt_vocab_size, tgt_vocab_file = vocab_utils.check_vocab(tgt_vocab_file
            , hparams.output_dir, check_special_token=hparams.
            check_special_token, sos=hparams.sos, eos=hparams.eos, unk=
            vocab_utils.UNK)
    _add_argument(hparams, 'src_vocab_size', src_vocab_size)
    _add_argument(hparams, 'tgt_vocab_size', tgt_vocab_size)
    _add_argument(hparams, 'src_vocab_file', src_vocab_file)
    _add_argument(hparams, 'tgt_vocab_file', tgt_vocab_file)
    _add_argument(hparams, 'num_enc_emb_partitions', hparams.
        num_embeddings_partitions)
    _add_argument(hparams, 'num_dec_emb_partitions', hparams.
        num_embeddings_partitions)
    _add_argument(hparams, 'src_embed_file', '')
    _add_argument(hparams, 'tgt_embed_file', '')
    if hparams.embed_prefix:
        src_embed_file = hparams.embed_prefix + '.' + hparams.src
        tgt_embed_file = hparams.embed_prefix + '.' + hparams.tgt
        if tf.gfile.Exists(src_embed_file):
            utils.print_out('  src_embed_file %s exist' % src_embed_file)
            hparams.src_embed_file = src_embed_file
            utils.print_out(
                'For pretrained embeddings, set num_enc_emb_partitions to 1')
            hparams.num_enc_emb_partitions = 1
        else:
            utils.print_out("  src_embed_file %s doesn't exist" %
                src_embed_file)
        if tf.gfile.Exists(tgt_embed_file):
            utils.print_out('  tgt_embed_file %s exist' % tgt_embed_file)
            hparams.tgt_embed_file = tgt_embed_file
            utils.print_out(
                'For pretrained embeddings, set num_dec_emb_partitions to 1')
            hparams.num_dec_emb_partitions = 1
        else:
            utils.print_out("  tgt_embed_file %s doesn't exist" %
                tgt_embed_file)
    metric = 'bleu'
    best_metric_dir = os.path.join(hparams.output_dir, 'best_' + metric)
    tf.gfile.MakeDirs(best_metric_dir)
    _add_argument(hparams, 'best_' + metric, 0, update=False)
    _add_argument(hparams, 'best_' + metric + '_dir', best_metric_dir)
    return hparams


def create_or_load_hparams(default_hparams, hparams_path):
    """Create hparams or load hparams from output_dir."""
    hparams = utils.maybe_parse_standard_hparams(default_hparams, hparams_path)
    hparams = extend_hparams(hparams)
    utils.print_hparams(hparams)
    return hparams


def run_main(flags, default_hparams, estimator_fn):
    """Run main."""
    random_seed = flags.random_seed
    if random_seed is not None and random_seed > 0:
        utils.print_out('# Set random seed to %d' % random_seed)
        random.seed(random_seed)
        np.random.seed(random_seed)
        tf.set_random_seed(random_seed)
    output_dir = flags.output_dir
    if output_dir and not tf.gfile.Exists(output_dir):
        utils.print_out('# Creating output directory %s ...' % output_dir)
        tf.gfile.MakeDirs(output_dir)
    hparams = create_or_load_hparams(default_hparams, flags.hparams_path)
    estimator_fn(hparams)
    return hparams


def tokenize(hparams, file, tokenized_file):
    utils.print_out('tokenizing {} -> {}'.format(file, tokenized_file))
    with open(file, 'rb') as input_file:
        with open(tokenized_file, 'wb') as output_file:
            subprocess.run([hparams.tokenizer_file, '-l', hparams.src],
                stdin=input_file, stdout=output_file)


def detokenize(hparams, file, detokenized_file):
    utils.print_out('detokenizing {} -> {}'.format(file, detokenized_file))
    with open(file, 'rb') as input_file:
        with open(detokenized_file, 'wb') as output_file:
            subprocess.run([hparams.detokenizer_file, '-l', hparams.tgt],
                stdin=input_file, stdout=output_file)


def main(unused_argv):
    experiment_start = time.time()
    tf.logging.set_verbosity(tf.logging.INFO)
    if FLAGS.use_fp16 and FLAGS.use_dist_strategy:
        raise ValueError("use_fp16 and use_dist_strategy aren't compatible")
    if FLAGS.use_fp16 + FLAGS.amp + FLAGS.use_fastmath > 1:
        raise ValueError('Only one of use_fp16, amp, use_fastmath can be set')
    if FLAGS.amp:
        utils.print_out('Enabling TF-AMP')
        os.environ['TF_ENABLE_AUTO_MIXED_PRECISION'] = '1'
    if FLAGS.use_fastmath:
        utils.print_out('Enabling FastMath')
        os.environ['TF_ENABLE_CUBLAS_TENSOR_OP_MATH_FP32'] = '1'
        os.environ['TF_ENABLE_CUDNN_TENSOR_OP_MATH_FP32'] = '1'
        os.environ['TF_ENABLE_CUDNN_RNN_TENSOR_OP_MATH_FP32'] = '1'
    active_xla_option_nums = np.sum([FLAGS.use_xla, FLAGS.use_autojit_xla,
        FLAGS.xla_compile])
    if active_xla_option_nums > 1:
        raise ValueError(
            'Only one of use_xla, xla_compile, use_autojit_xla can be set')
    os.environ['use_xla'] = str(FLAGS.use_xla).lower()
    if FLAGS.use_xla:
        os.environ['use_defun'] = str(True).lower()
    else:
        os.environ['use_defun'] = str(FLAGS.use_defun).lower()
    utils.print_out('use_defun is %s for attention' % os.environ['use_defun'])
    os.environ['use_fp32_batch_matmul'
        ] = 'true' if FLAGS.use_fp32_batch_matmul else 'false'
    os.environ['xla_compile'] = 'true' if FLAGS.xla_compile else 'false'
    os.environ['force_inputs_padding'
        ] = 'true' if FLAGS.force_inputs_padding else 'false'
    if FLAGS.mode == 'train':
        utils.print_out('Running training mode.')
        default_hparams = create_hparams(FLAGS)
        run_main(FLAGS, default_hparams, estimator.train_fn)
    elif FLAGS.mode == 'infer' or FLAGS.mode == 'translate':
        if FLAGS.mode == 'infer':
            utils.print_out('Running inference mode.')
            translate_mode = False
        else:
            utils.print_out('Running translate mode on file {}.'.format(
                FLAGS.translate_file))
            translate_mode = True
        random_seed = FLAGS.random_seed
        if random_seed is not None and random_seed > 0:
            utils.print_out('# Set random seed to %d' % random_seed)
            random.seed(random_seed)
            np.random.seed(random_seed)
            tf.set_random_seed(random_seed)
        output_dir = FLAGS.output_dir
        if output_dir and not tf.gfile.Exists(output_dir):
            utils.print_out('# Creating output directory %s ...' % output_dir)
            tf.gfile.MakeDirs(output_dir)
        dllogger.init(backends=[dllogger.StdOutBackend(dllogger.Verbosity.
            DEFAULT), dllogger.JSONStreamBackend(dllogger.Verbosity.VERBOSE,
            os.path.join(FLAGS.output_dir, FLAGS.mode + '-report.json'))])
        dllogger.log('PARAMETER', vars(FLAGS))
        default_hparams = create_hparams(FLAGS)
        default_hparams.num_buckets = 1
        hparams = create_or_load_hparams(default_hparams, FLAGS.hparams_path)
        utils.print_out('infer_hparams:')
        utils.print_hparams(hparams)
        if translate_mode:
            tokenize(hparams, hparams.translate_file, hparams.
                translate_file + '.tok')
        eval_sentences, eval_src_tokens, _ = (iterator_utils.
            get_effective_epoch_size(hparams, train=False))
        tf.logging.info('Starting to evaluate...')
        eval_start = time.time()
        _, (eval_speed, eval_latencies
            ), eval_output_tokens = estimator.eval_fn(hparams, hparams.ckpt,
            only_translate=translate_mode)
        eval_end = time.time()
        eval_delta = eval_end - eval_start
        utils.print_out(
            'eval time for ckpt: %.2f mins (%.2f sent/sec, %.2f tokens/sec)' %
            (eval_delta / 60.0, eval_speed, eval_speed * (eval_src_tokens +
            eval_output_tokens) / eval_sentences), f=sys.stderr)
        logging_data = {'infer_speed_sent': eval_speed, 'infer_speed_toks':
            eval_speed * (eval_src_tokens + eval_output_tokens) /
            eval_sentences}
        for lat in sorted(eval_latencies):
            utils.print_out('eval latency_%s for ckpt: %.2f ms' % (lat, 
                eval_latencies[lat] * 1000))
            logging_data['infer_latency_{}'.format(lat)] = eval_latencies[lat
                ] * 1000
        dllogger.log((), logging_data)
        dllogger.flush()
        if translate_mode:
            detokenize(hparams, hparams.translate_file + '.trans.tok', 
                hparams.translate_file + '.trans')
    else:
        assert FLAGS.mode == 'train_and_eval'
        utils.print_out('Running train and eval mode.')
        random_seed = FLAGS.random_seed
        if random_seed is not None and random_seed > 0:
            utils.print_out('# Set random seed to %d' % random_seed)
            random.seed(random_seed)
            np.random.seed(random_seed)
            tf.set_random_seed(random_seed)
        output_dir = FLAGS.output_dir
        if output_dir and not tf.gfile.Exists(output_dir):
            utils.print_out('# Creating output directory %s ...' % output_dir)
            tf.gfile.MakeDirs(output_dir)
        dllogger.init(backends=[dllogger.StdOutBackend(dllogger.Verbosity.
            DEFAULT), dllogger.JSONStreamBackend(dllogger.Verbosity.VERBOSE,
            os.path.join(FLAGS.output_dir, FLAGS.mode + '-report.json'))])
        dllogger.log('PARAMETER', vars(FLAGS))
        dllogger.metadata('bleu', {'unit': None})
        dllogger.metadata('train_speed_sent', {'unit': 'sequences/s'})
        dllogger.metadata('train_speed_toks', {'unit': 'tokens/s'})
        default_hparams = create_hparams(FLAGS)
        hparams = create_or_load_hparams(default_hparams, FLAGS.hparams_path)
        utils.print_out('training hparams:')
        utils.print_hparams(hparams)
        with tf.gfile.GFile(os.path.join(output_dir, 'train_hparams.txt'), 'w'
            ) as f:
            f.write(utils.serialize_hparams(hparams) + '\n')
        infer_hparams = tf.contrib.training.HParams(**hparams.values())
        infer_hparams.num_buckets = 1
        utils.print_out('infer_hparams:')
        utils.print_hparams(infer_hparams)
        with tf.gfile.GFile(os.path.join(output_dir, 'infer_hparams.txt'), 'w'
            ) as f:
            f.write(utils.serialize_hparams(infer_hparams) + '\n')
        epochs = 0
        should_stop = epochs >= FLAGS.max_train_epochs
        train_sentences, train_src_tokens, train_tgt_tokens = (iterator_utils
            .get_effective_epoch_size(hparams))
        eval_sentences, eval_src_tokens, _ = (iterator_utils.
            get_effective_epoch_size(hparams, train=False))
        while not should_stop:
            utils.print_out('Starting epoch %d' % epochs)
            try:
                train_start = time.time()
                train_speed, _ = estimator.train_fn(hparams)
            except tf.errors.OutOfRangeError:
                utils.print_out('training hits OutOfRangeError', f=sys.stderr)
            train_end = time.time()
            train_delta = train_end - train_start
            utils.print_out(
                'training time for epoch %d: %.2f mins (%.2f sent/sec, %.2f tokens/sec)'
                 % (epochs + 1, train_delta / 60.0, train_speed, 
                train_speed * (train_src_tokens + train_tgt_tokens) /
                train_sentences), f=sys.stderr)
            logging_data = {'train_speed_sent': train_speed,
                'train_speed_toks': train_speed * (train_src_tokens +
                train_tgt_tokens) / train_sentences}
            eval_start = time.time()
            bleu_score, (eval_speed, eval_latencies
                ), eval_output_tokens = estimator.eval_fn(infer_hparams)
            eval_end = time.time()
            eval_delta = eval_end - eval_start
            utils.print_out(
                'eval time for epoch %d: %.2f mins (%.2f sent/sec, %.2f tokens/sec)'
                 % (epochs + 1, eval_delta / 60.0, eval_speed, eval_speed *
                (eval_src_tokens + eval_output_tokens) / eval_sentences), f
                =sys.stderr)
            logging_data.update({'bleu': bleu_score, 'infer_speed_sent':
                eval_speed, 'infer_speed_toks': eval_speed * (
                eval_src_tokens + eval_output_tokens) / eval_sentences})
            for lat in sorted(eval_latencies):
                utils.print_out('eval latency_%s for epoch %d: %.2f ms' % (
                    lat, epochs + 1, eval_latencies[lat] * 1000))
                logging_data['eval_latency_{}'.format(lat)] = eval_latencies[
                    lat] * 1000
            dllogger.log((epochs,), logging_data)
            dllogger.flush()
            if (FLAGS.debug or FLAGS.target_bleu is not None and bleu_score >
                FLAGS.target_bleu):
                should_stop = True
                utils.print_out(
                    'Stop job since target bleu is reached at epoch %d .' %
                    epochs, f=sys.stderr)
            epochs += 1
            if epochs >= FLAGS.max_train_epochs:
                should_stop = True
                utils.print_out('Stop job since max_train_epochs is reached.',
                    f=sys.stderr)
        dllogger.log((), logging_data)
        dllogger.flush()
    experiment_end = time.time()
    utils.print_out('Experiment took {} min'.format((experiment_end -
        experiment_start) / 60))


if __name__ == '__main__':
    nmt_parser = argparse.ArgumentParser()
    add_arguments(nmt_parser)
    FLAGS, unparsed = nmt_parser.parse_known_args()
    tf.app.run(main=main, argv=[sys.argv[0]] + unparsed)

# Copyright 2017 Google Inc. 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.
# ==============================================================================
#
# Copyright (c) 2019, NVIDIA CORPORATION. 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.

"""TensorFlow NMT model implementation."""
from __future__ import print_function

import argparse
import os
import random
import sys
import subprocess

# import matplotlib.image as mpimg
import numpy as np
import time
import tensorflow as tf

import dllogger

import estimator
from utils import evaluation_utils
from utils import iterator_utils
from utils import misc_utils as utils
from utils import vocab_utils
from variable_mgr import constants

utils.check_tensorflow_version()

FLAGS = None


# LINT.IfChange
def add_arguments(parser):
  """Build ArgumentParser."""
  parser.register("type", "bool", lambda v: v.lower() == "true")

  # network
  parser.add_argument(
      "--num_units", type=int, default=1024, help="Network size.")
  parser.add_argument(
      "--num_layers", type=int, default=4, help="Network depth.")
  parser.add_argument("--num_encoder_layers", type=int, default=None,
                      help="Encoder depth, equal to num_layers if None.")
  parser.add_argument("--num_decoder_layers", type=int, default=None,
                      help="Decoder depth, equal to num_layers if None.")
  parser.add_argument(
      "--encoder_type",
      type=str,
      default="gnmt",
      help="""\
      uni | bi | gnmt.
      For bi, we build num_encoder_layers/2 bi-directional layers.
      For gnmt, we build 1 bi-directional layer, and (num_encoder_layers - 1)
        uni-directional layers.\
      """)
  parser.add_argument(
      "--residual",
      type="bool",
      nargs="?",
      const=True,
      default=True,
      help="Whether to add residual connections.")
  parser.add_argument("--time_major", type="bool", nargs="?", const=True,
                      default=True,
                      help="Whether to use time-major mode for dynamic RNN.")
  parser.add_argument("--num_embeddings_partitions", type=int, default=0,
                      help="Number of partitions for embedding vars.")

  # attention mechanisms
  parser.add_argument(
      "--attention",
      type=str,
      default="normed_bahdanau",
      help="""\
      luong | scaled_luong | bahdanau | normed_bahdanau or set to "" for no
      attention\
      """)
  parser.add_argument(
      "--attention_architecture",
      type=str,
      default="gnmt_v2",
      help="""\
      standard | gnmt | gnmt_v2.
      standard: use top layer to compute attention.
      gnmt: GNMT style of computing attention, use previous bottom layer to
          compute attention.
      gnmt_v2: similar to gnmt, but use current bottom layer to compute
          attention.\
      """)
  parser.add_argument(
      "--output_attention", type="bool", nargs="?", const=True,
      default=True,
      help="""\
      Only used in standard attention_architecture. Whether use attention as
      the cell output at each timestep.
      .\
      """)
  parser.add_argument(
      "--pass_hidden_state", type="bool", nargs="?", const=True,
      default=True,
      help="""\
      Whether to pass encoder's hidden state to decoder when using an attention
      based model.\
      """)

  # optimizer
  parser.add_argument(
      "--optimizer", type=str, default="adam", help="sgd | adam")
  parser.add_argument(
      "--learning_rate",
      type=float,
      default=5e-4,
      help="Learning rate. Adam: 0.001 | 0.0001")
  parser.add_argument("--warmup_steps", type=int, default=200,
                      help="How many steps we inverse-decay learning.")
  parser.add_argument("--warmup_scheme", type=str, default="t2t", help="""\
      How to warmup learning rates. Options include:
        t2t: Tensor2Tensor's way, start with lr 100 times smaller, then
             exponentiate until the specified lr.\
      """)
  parser.add_argument(
      "--decay_scheme", type=str, default="luong234", help="""\
      How we decay learning rate. Options include:
        luong234: after 2/3 num train steps, we start halving the learning rate
          for 4 times before finishing.
        luong5: after 1/2 num train steps, we start halving the learning rate
          for 5 times before finishing.\
        luong10: after 1/2 num train steps, we start halving the learning rate
          for 10 times before finishing.\
      """)

  parser.add_argument(
      "--max_train_epochs", type=int, default=6, help="Max number of epochs.")
  parser.add_argument(
      "--target_bleu", type=float, default=None, help="Target bleu.")
  parser.add_argument("--colocate_gradients_with_ops", type="bool", nargs="?",
                      const=True,
                      default=True,
                      help=("Whether try colocating gradients with "
                            "corresponding op"))
  parser.add_argument("--label_smoothing", type=float, default=0.1,
                      help=("If nonzero, smooth the labels towards "
                            "1/num_classes."))

  # initializer
  parser.add_argument("--init_op", type=str, default="uniform",
                      help="uniform | glorot_normal | glorot_uniform")
  parser.add_argument("--init_weight", type=float, default=0.1,
                      help=("for uniform init_op, initialize weights "
                            "between [-this, this]."))

  # data
  parser.add_argument(
      "--src", type=str, default="en", help="Source suffix, e.g., en.")
  parser.add_argument(
      "--tgt", type=str, default="de", help="Target suffix, e.g., de.")
  parser.add_argument(
      "--data_dir", type=str, default="data/wmt16_de_en",
      help="Training/eval data directory.")

  parser.add_argument(
      "--train_prefix",
      type=str,
      default="train.tok.clean.bpe.32000",
      help="Train prefix, expect files with src/tgt suffixes.")
  parser.add_argument(
      "--test_prefix",
      type=str,
      default="newstest2014.tok.bpe.32000",
      help="Test prefix, expect files with src/tgt suffixes.")

  parser.add_argument(
      "--translate_file",
      type=str,
      help="File to translate, works only with translate mode")

  parser.add_argument(
      "--output_dir", type=str, default="results",
      help="Store log/model files.")

  # Vocab
  parser.add_argument(
      "--vocab_prefix",
      type=str,
      default="vocab.bpe.32000",
      help="""\
      Vocab prefix, expect files with src/tgt suffixes.\
      """)

  parser.add_argument(
      "--embed_prefix",
      type=str,
      default=None,
      help="""\
      Pretrained embedding prefix, expect files with src/tgt suffixes.
      The embedding files should be Glove formatted txt files.\
      """)
  parser.add_argument("--sos", type=str, default="<s>",
                      help="Start-of-sentence symbol.")
  parser.add_argument("--eos", type=str, default="</s>",
                      help="End-of-sentence symbol.")
  parser.add_argument(
      "--share_vocab",
      type="bool",
      nargs="?",
      const=True,
      default=True,
      help="""\
      Whether to use the source vocab and embeddings for both source and
      target.\
      """)
  parser.add_argument("--check_special_token", type="bool", default=True,
                      help="""\
                      Whether check special sos, eos, unk tokens exist in the
                      vocab files.\
                      """)

  # Sequence lengths
  parser.add_argument(
      "--src_max_len",
      type=int,
      default=50,
      help="Max length of src sequences during training (including EOS).")
  parser.add_argument(
      "--tgt_max_len",
      type=int,
      default=50,
      help="Max length of tgt sequences during training (including BOS).")
  parser.add_argument("--src_max_len_infer", type=int, default=None,
                      help="Max length of src sequences during inference (including EOS).")
  parser.add_argument("--tgt_max_len_infer", type=int, default=80,
                      help="""\
      Max length of tgt sequences during inference (including BOS).  Also use to restrict the
      maximum decoding length.\
      """)

  # Default settings works well (rarely need to change)
  parser.add_argument("--unit_type", type=str, default="lstm",
                      help="lstm | gru | layer_norm_lstm | nas")
  parser.add_argument("--forget_bias", type=float, default=0.0,
                      help="Forget bias for BasicLSTMCell.")
  parser.add_argument("--dropout", type=float, default=0.2,
                      help="Dropout rate (not keep_prob)")
  parser.add_argument("--max_gradient_norm", type=float, default=5.0,
                      help="Clip gradients to this norm.")
  parser.add_argument("--batch_size", type=int, default=128, help="Total batch size.")

  parser.add_argument(
      "--num_buckets",
      type=int,
      default=5,
      help="Put data into similar-length buckets (only for training).")

  # SPM
  parser.add_argument("--subword_option", type=str, default="bpe",
                      choices=["", "bpe", "spm"],
                      help="""\
                      Set to bpe or spm to activate subword desegmentation.\
                      """)

  # Experimental encoding feature.
  parser.add_argument("--use_char_encode", type="bool", default=False,
                      help="""\
                      Whether to split each word or bpe into character, and then
                      generate the word-level representation from the character
                      reprentation.
                      """)

  # Misc
  parser.add_argument(
      "--save_checkpoints_steps", type=int, default=2000,
      help="save_checkpoints_steps")
  parser.add_argument(
      "--log_step_count_steps", type=int, default=10,
      help=("The frequency, in number of global steps, that the global step "
            "and the loss will be logged during training"))
  parser.add_argument(
      "--num_gpus", type=int, default=1, help="Number of gpus in each worker.")
  parser.add_argument("--hparams_path", type=str, default=None,
      help=("Path to standard hparams json file that overrides"
            "hparams values from FLAGS."))
  parser.add_argument(
      "--random_seed",
      type=int,
      default=1,
      help="Random seed (>0, set a specific seed).")
  parser.add_argument("--language_model", type="bool", nargs="?",
                      const=True, default=False,
                      help="True to train a language model, ignoring encoder")

  # Inference
  parser.add_argument("--ckpt", type=str, default=None,
                      help="Checkpoint file to load a model for inference. (defaults to newest checkpoint)")
  parser.add_argument(
      "--infer_batch_size",
      type=int,
      default=128,
      help="Batch size for inference mode.")
  parser.add_argument("--detokenizer_file", type=str,
                      default=None,
                      help=("""Detokenizer script file. Default: DATA_DIR/mosesdecoder/scripts/tokenizer/detokenizer.perl"""))
  parser.add_argument("--tokenizer_file", type=str,
                      default=None,
                      help=("""Tokenizer script file. Default: DATA_DIR/mosesdecoder/scripts/tokenizer/tokenizer.perl"""))

  # Advanced inference arguments
  parser.add_argument("--infer_mode", type=str, default="beam_search",
                      choices=["greedy", "beam_search"],
                      help="Which type of decoder to use during inference.")
  parser.add_argument("--beam_width", type=int, default=5,
                      help=("""\
      beam width when using beam search decoder. If 0, use standard
      decoder with greedy helper.\
      """))
  parser.add_argument(
      "--length_penalty_weight",
      type=float,
      default=0.6,
      help="Length penalty for beam search.")
  parser.add_argument(
      "--coverage_penalty_weight",
      type=float,
      default=0.1,
      help="Coverage penalty for beam search.")

  # Job info
  parser.add_argument("--num_workers", type=int, default=1,
                      help="Number of workers (inference only).")

  parser.add_argument("--amp", action='store_true',
                      help="use amp for training and inference")
  parser.add_argument("--use_fastmath", type="bool", default=False,
                      help="use_fastmath for training and inference")
  parser.add_argument("--use_fp16", type="bool", default=False,
                      help="use_fp16 for training and inference")
  parser.add_argument(
      "--fp16_loss_scale",
      type=float,
      default=128,
      help="If fp16 is enabled, the loss is multiplied by this amount "
      "right before gradients are computed, then each gradient "
      "is divided by this amount. Mathematically, this has no "
      "effect, but it helps avoid fp16 underflow. Set to 1 to "
      "effectively disable.")
  parser.add_argument(
      "--enable_auto_loss_scale",
      type="bool",
      default=True,
      help="If True and use_fp16 is True, automatically adjust the "
      "loss scale during training.")
  parser.add_argument(
      "--fp16_inc_loss_scale_every_n",
      type=int,
      default=128,
      help="If fp16 is enabled and enable_auto_loss_scale is "
      "True, increase the loss scale every n steps.")
  parser.add_argument(
      "--check_tower_loss_numerics",
      type="bool",
      default=False,  # Set to false for xla.compile()
      help="whether to check tower loss numerics")
  parser.add_argument(
      "--use_fp32_batch_matmul",
      type="bool",
      default=False,
      help="Whether to use fp32 batch matmul")

  # Performance
  # XLA
  parser.add_argument(
      "--force_inputs_padding",
      type="bool",
      default=False,
      help="Force padding input batch to src_max_len and tgt_max_len")
  parser.add_argument(
      "--use_xla",
      type="bool",
      default=False,
      help="Use xla to compile a few selected locations, mostly Defuns.")
  parser.add_argument(
      "--xla_compile",
      type="bool",
      default=False,
      help="Use xla.compile() for each tower's fwd and bak pass.")
  parser.add_argument(
      "--use_autojit_xla",
      type="bool",
      default=False,
      help="Use auto jit xla.")
  # GPU knobs
  parser.add_argument(
      "--use_pintohost_optimizer",
      type="bool",
      default=False,
      help="whether to use PinToHost optimizer")
  parser.add_argument(
      "--use_cudnn_lstm",
      type="bool",
      default=False,
      help="whether to use cudnn_lstm for encoder, non residual layers")
  parser.add_argument(
      "--use_loose_bidi_cudnn_lstm",
      type="bool",
      default=False,
      help="whether to use loose bidi cudnn_lstm")
  parser.add_argument(
      "--use_fused_lstm",
      type="bool",
      default=True,
      help="whether to use fused lstm and variant. If enabled, training will "
      "use LSTMBlockFusedCell, infer will use LSTMBlockCell when appropriate.")
  parser.add_argument(
      "--use_fused_lstm_dec",
      type="bool",
      default=False,
      help="whether to use fused lstm for decoder (training only).")
  parser.add_argument(
      "--gpu_indices",
      type=str,
      default="",
      help="Indices of worker GPUs in ring order")
  # Graph knobs
  parser.add_argument("--parallel_iterations", type=int, default=10,
                      help="number of parallel iterations in dynamic_rnn")
  parser.add_argument("--use_dist_strategy", type="bool", default=False,
                      help="whether to use distribution strategy")
  parser.add_argument(
      "--hierarchical_copy",
      type="bool",
      default=False,
      help="Use hierarchical copies. Currently only optimized for "
      "use on a DGX-1 with 8 GPUs and may perform poorly on "
      "other hardware. Requires --num_gpus > 1, and only "
      "recommended when --num_gpus=8")
  parser.add_argument(
      "--network_topology",
      type=constants.NetworkTopology,
      default=constants.NetworkTopology.DGX1,
      choices=list(constants.NetworkTopology))
  parser.add_argument(
      "--use_block_lstm",
      type="bool",
      default=False,
      help="whether to use block lstm")
  parser.add_argument(
      "--use_defun",
      type="bool",
      default=False,
      help="whether to use Defun")

  # Gradient tricks
  parser.add_argument(
      "--gradient_repacking",
      type=int,
      default=0,
      help="Use gradient repacking. It"
      "currently only works with replicated mode. At the end of"
      "of each step, it repacks the gradients for more efficient"
      "cross-device transportation. A non-zero value specifies"
      "the number of split packs that will be formed.")
  parser.add_argument(
      "--compact_gradient_transfer",
      type="bool",
      default=True,
      help="Compact gradient as much as possible for cross-device transfer and "
      "aggregation.")
  parser.add_argument(
      "--all_reduce_spec",
      type=str,
      default="nccl",
      help="A specification of the all_reduce algorithm to be used "
      "for reducing gradients.  For more details, see "
      "parse_all_reduce_spec in variable_mgr.py.  An "
      "all_reduce_spec has BNF form:\n"
      "int ::= positive whole number\n"
      "g_int ::= int[KkMGT]?\n"
      "alg_spec ::= alg | alg#int\n"
      "range_spec ::= alg_spec | alg_spec/alg_spec\n"
      "spec ::= range_spec | range_spec:g_int:range_spec\n"
      "NOTE: not all syntactically correct constructs are "
      "supported.\n\n"
      "Examples:\n "
      "\"xring\" == use one global ring reduction for all "
      "tensors\n"
      "\"pscpu\" == use CPU at worker 0 to reduce all tensors\n"
      "\"nccl\" == use NCCL to locally reduce all tensors.  "
      "Limited to 1 worker.\n"
      "\"nccl/xring\" == locally (to one worker) reduce values "
      "using NCCL then ring reduce across workers.\n"
      "\"pscpu:32k:xring\" == use pscpu algorithm for tensors of "
      "size up to 32kB, then xring for larger tensors.")
  parser.add_argument(
      "--agg_small_grads_max_bytes",
      type=int,
      default=0,
      help="If > 0, try to aggregate tensors of less than this "
      "number of bytes prior to all-reduce.")
  parser.add_argument(
      "--agg_small_grads_max_group",
      type=int,
      default=10,
      help="When aggregating small tensors for all-reduce do not "
      "aggregate more than this many into one new tensor.")
  parser.add_argument(
      "--allreduce_merge_scope",
      type=int,
      default=1,
      help="Establish a name scope around this many "
      "gradients prior to creating the all-reduce operations. "
      "It may affect the ability of the backend to merge "
      "parallel ops.")
  # Other knobs
  parser.add_argument(
      "--local_parameter_device",
      type=str,
      default="gpu",
      help="Device to use as parameter server: cpu or gpu. For "
      "distributed training, it can affect where caching of "
      "variables happens.")
  parser.add_argument(
      "--use_resource_vars",
      type="bool",
      default=False,
      help="Use resource variables instead of normal variables. "
      "Resource variables are slower, but this option is useful "
      "for debugging their performance.")


  parser.add_argument("--debug", type="bool", default=False,
                      help="Debug train and eval")
  parser.add_argument(
      "--debug_num_train_steps", type=int, default=None, help="Num steps to train.")
  parser.add_argument("--show_metrics", type="bool", default=True,
                      help="whether to show detailed metrics")
  parser.add_argument("--clip_grads", type="bool", default=True,
                      help="whether to clip gradients")
  parser.add_argument("--profile", type="bool", default=False,
                      help="If generate profile")
  parser.add_argument("--profile_save_steps", type=int, default=10,
                      help="Save timeline every N steps.")

  parser.add_argument("--use_dynamic_rnn", type="bool", default=True)
  parser.add_argument("--use_synthetic_data", type="bool", default=False)
  parser.add_argument(
      "--mode", type=str, default="train_and_eval",
      choices=("train_and_eval", "infer", "translate"))


def create_hparams(flags):
  """Create training hparams."""
  return tf.contrib.training.HParams(
      # Data
      src=flags.src,
      tgt=flags.tgt,
      train_prefix=os.path.join(flags.data_dir, flags.train_prefix),
      test_prefix=os.path.join(flags.data_dir, flags.test_prefix),
      translate_file=flags.translate_file,
      vocab_prefix=os.path.join(flags.data_dir, flags.vocab_prefix),
      embed_prefix=flags.embed_prefix,
      output_dir=flags.output_dir,

      # Networks
      num_units=flags.num_units,
      num_encoder_layers=(flags.num_encoder_layers or flags.num_layers),
      num_decoder_layers=(flags.num_decoder_layers or flags.num_layers),
      dropout=flags.dropout,
      unit_type=flags.unit_type,
      encoder_type=flags.encoder_type,
      residual=flags.residual,
      time_major=flags.time_major,
      num_embeddings_partitions=flags.num_embeddings_partitions,

      # Attention mechanisms
      attention=flags.attention,
      attention_architecture=flags.attention_architecture,
      output_attention=flags.output_attention,
      pass_hidden_state=flags.pass_hidden_state,

      # Train
      optimizer=flags.optimizer,
      max_train_epochs=flags.max_train_epochs,
      target_bleu=flags.target_bleu,
      label_smoothing=flags.label_smoothing,
      batch_size=flags.batch_size,
      init_op=flags.init_op,
      init_weight=flags.init_weight,
      max_gradient_norm=flags.max_gradient_norm,
      learning_rate=flags.learning_rate,
      warmup_steps=flags.warmup_steps,
      warmup_scheme=flags.warmup_scheme,
      decay_scheme=flags.decay_scheme,
      colocate_gradients_with_ops=flags.colocate_gradients_with_ops,

      # Data constraints
      num_buckets=flags.num_buckets,
      src_max_len=flags.src_max_len,
      tgt_max_len=flags.tgt_max_len,

      # Inference
      src_max_len_infer=flags.src_max_len_infer,
      tgt_max_len_infer=flags.tgt_max_len_infer,
      ckpt=flags.ckpt,
      infer_batch_size=flags.infer_batch_size,
      detokenizer_file=flags.detokenizer_file if flags.detokenizer_file is not None \
                       else os.path.join(flags.data_dir, 'mosesdecoder/scripts/tokenizer/detokenizer.perl'),
      tokenizer_file=flags.tokenizer_file if flags.tokenizer_file is not None \
                       else os.path.join(flags.data_dir, 'mosesdecoder/scripts/tokenizer/tokenizer.perl'),

      # Advanced inference arguments
      infer_mode=flags.infer_mode,
      beam_width=flags.beam_width,
      length_penalty_weight=flags.length_penalty_weight,
      coverage_penalty_weight=flags.coverage_penalty_weight,

      # Vocab
      sos=flags.sos if flags.sos else vocab_utils.SOS,
      eos=flags.eos if flags.eos else vocab_utils.EOS,
      subword_option=flags.subword_option,
      check_special_token=flags.check_special_token,
      use_char_encode=flags.use_char_encode,

      # Misc
      forget_bias=flags.forget_bias,
      num_gpus=flags.num_gpus,
      save_checkpoints_steps=flags.save_checkpoints_steps,
      log_step_count_steps=flags.log_step_count_steps,
      epoch_step=0,  # record where we were within an epoch.
      share_vocab=flags.share_vocab,
      random_seed=flags.random_seed,
      language_model=flags.language_model,

      amp=flags.amp,
      use_fastmath=flags.use_fastmath,
      use_fp16=flags.use_fp16,
      fp16_loss_scale=flags.fp16_loss_scale,
      enable_auto_loss_scale=flags.enable_auto_loss_scale,
      fp16_inc_loss_scale_every_n=flags.fp16_inc_loss_scale_every_n,
      check_tower_loss_numerics=flags.check_tower_loss_numerics,
      use_fp32_batch_matmul=flags.use_fp32_batch_matmul,

      # Performance
      # GPU knbs
      force_inputs_padding=flags.force_inputs_padding,
      use_xla=flags.use_xla,
      xla_compile=flags.xla_compile,
      use_autojit_xla=flags.use_autojit_xla,
      use_pintohost_optimizer=flags.use_pintohost_optimizer,
      use_cudnn_lstm=flags.use_cudnn_lstm,
      use_loose_bidi_cudnn_lstm=flags.use_loose_bidi_cudnn_lstm,
      use_fused_lstm=flags.use_fused_lstm,
      use_fused_lstm_dec=flags.use_fused_lstm_dec,
      gpu_indices=flags.gpu_indices,
      # Graph knobs
      parallel_iterations=flags.parallel_iterations,
      use_dynamic_rnn=flags.use_dynamic_rnn,
      use_dist_strategy=flags.use_dist_strategy,
      hierarchical_copy=flags.hierarchical_copy,
      network_topology=flags.network_topology,
      use_block_lstm=flags.use_block_lstm,
      # Grad tricks
      gradient_repacking=flags.gradient_repacking,
      compact_gradient_transfer=flags.compact_gradient_transfer,
      all_reduce_spec=flags.all_reduce_spec,
      agg_small_grads_max_bytes=flags.agg_small_grads_max_bytes,
      agg_small_grads_max_group=flags.agg_small_grads_max_group,
      allreduce_merge_scope=flags.allreduce_merge_scope,
      # Other knobs
      local_parameter_device=("cpu" if flags.num_gpus ==0
                              else flags.local_parameter_device),
      use_resource_vars=flags.use_resource_vars,

      debug=flags.debug,
      debug_num_train_steps=flags.debug_num_train_steps,
      clip_grads=flags.clip_grads,
      profile=flags.profile,
      profile_save_steps=flags.profile_save_steps,
      show_metrics=flags.show_metrics,

      use_synthetic_data=flags.use_synthetic_data,
      mode=flags.mode,
  )


def _add_argument(hparams, key, value, update=True):
  """Add an argument to hparams; if exists, change the value if update==True."""
  if hasattr(hparams, key):
    if update:
      setattr(hparams, key, value)
  else:
    hparams.add_hparam(key, value)


def extend_hparams(hparams):
  """Add new arguments to hparams."""
  # Sanity checks
  if hparams.encoder_type == "bi" and hparams.num_encoder_layers % 2 != 0:
    raise ValueError("For bi, num_encoder_layers %d should be even" %
                     hparams.num_encoder_layers)
  if (hparams.attention_architecture in ["gnmt"] and
      hparams.num_encoder_layers < 2):
    raise ValueError("For gnmt attention architecture, "
                     "num_encoder_layers %d should be >= 2" %
                     hparams.num_encoder_layers)
  if hparams.subword_option and hparams.subword_option not in ["spm", "bpe"]:
    raise ValueError("subword option must be either spm, or bpe")
  if hparams.infer_mode == "beam_search" and hparams.beam_width <= 0:
    raise ValueError("beam_width must greater than 0 when using beam_search"
                     "decoder.")
  if hparams.mode == "translate" and not hparams.translate_file:
    raise ValueError("--translate_file flag must be specified in translate mode")

  # Different number of encoder / decoder layers
  assert hparams.num_encoder_layers and hparams.num_decoder_layers
  if hparams.num_encoder_layers != hparams.num_decoder_layers:
    hparams.pass_hidden_state = False
    utils.print_out("Num encoder layer %d is different from num decoder layer"
                    " %d, so set pass_hidden_state to False" % (
                        hparams.num_encoder_layers,
                        hparams.num_decoder_layers))

  # Set residual layers
  num_encoder_residual_layers = 0
  num_decoder_residual_layers = 0
  if hparams.residual:
    if hparams.num_encoder_layers > 1:
      num_encoder_residual_layers = hparams.num_encoder_layers - 1
    if hparams.num_decoder_layers > 1:
      num_decoder_residual_layers = hparams.num_decoder_layers - 1

    if hparams.encoder_type == "gnmt":
      # The first unidirectional layer (after the bi-directional layer) in
      # the GNMT encoder can't have residual connection due to the input is
      # the concatenation of fw_cell and bw_cell's outputs.
      num_encoder_residual_layers = hparams.num_encoder_layers - 2

      # Compatible for GNMT models
      if hparams.num_encoder_layers == hparams.num_decoder_layers:
        num_decoder_residual_layers = num_encoder_residual_layers
  _add_argument(hparams, "num_encoder_residual_layers",
                num_encoder_residual_layers)
  _add_argument(hparams, "num_decoder_residual_layers",
                num_decoder_residual_layers)

  # Language modeling
  if hparams.language_model:
    hparams.attention = ""
    hparams.attention_architecture = ""
    hparams.pass_hidden_state = False
    hparams.share_vocab = True
    hparams.src = hparams.tgt
    utils.print_out("For language modeling, we turn off attention and "
                    "pass_hidden_state; turn on share_vocab; set src to tgt.")

  ## Vocab
  # Get vocab file names first
  if hparams.vocab_prefix:
    src_vocab_file = hparams.vocab_prefix + "." + hparams.src
    tgt_vocab_file = hparams.vocab_prefix + "." + hparams.tgt
  else:
    raise ValueError("hparams.vocab_prefix must be provided.")

  # Source vocab
  src_vocab_size, src_vocab_file = vocab_utils.check_vocab(
      src_vocab_file,
      hparams.output_dir,
      check_special_token=hparams.check_special_token,
      sos=hparams.sos,
      eos=hparams.eos,
      unk=vocab_utils.UNK,
      pad_vocab=True)

  # Target vocab
  if hparams.share_vocab:
    utils.print_out("  using source vocab for target")
    tgt_vocab_file = src_vocab_file
    tgt_vocab_size = src_vocab_size
  else:
    tgt_vocab_size, tgt_vocab_file = vocab_utils.check_vocab(
        tgt_vocab_file,
        hparams.output_dir,
        check_special_token=hparams.check_special_token,
        sos=hparams.sos,
        eos=hparams.eos,
        unk=vocab_utils.UNK)
  _add_argument(hparams, "src_vocab_size", src_vocab_size)
  _add_argument(hparams, "tgt_vocab_size", tgt_vocab_size)
  _add_argument(hparams, "src_vocab_file", src_vocab_file)
  _add_argument(hparams, "tgt_vocab_file", tgt_vocab_file)

  # Num embedding partitions
  _add_argument(
      hparams, "num_enc_emb_partitions", hparams.num_embeddings_partitions)
  _add_argument(
      hparams, "num_dec_emb_partitions", hparams.num_embeddings_partitions)

  # Pretrained Embeddings
  _add_argument(hparams, "src_embed_file", "")
  _add_argument(hparams, "tgt_embed_file", "")
  if hparams.embed_prefix:
    src_embed_file = hparams.embed_prefix + "." + hparams.src
    tgt_embed_file = hparams.embed_prefix + "." + hparams.tgt

    if tf.gfile.Exists(src_embed_file):
      utils.print_out("  src_embed_file %s exist" % src_embed_file)
      hparams.src_embed_file = src_embed_file

      utils.print_out(
          "For pretrained embeddings, set num_enc_emb_partitions to 1")
      hparams.num_enc_emb_partitions = 1
    else:
      utils.print_out("  src_embed_file %s doesn't exist" % src_embed_file)

    if tf.gfile.Exists(tgt_embed_file):
      utils.print_out("  tgt_embed_file %s exist" % tgt_embed_file)
      hparams.tgt_embed_file = tgt_embed_file

      utils.print_out(
          "For pretrained embeddings, set num_dec_emb_partitions to 1")
      hparams.num_dec_emb_partitions = 1
    else:
      utils.print_out("  tgt_embed_file %s doesn't exist" % tgt_embed_file)

  # Evaluation
  metric = "bleu"
  best_metric_dir = os.path.join(hparams.output_dir, "best_" + metric)
  tf.gfile.MakeDirs(best_metric_dir)
  _add_argument(hparams, "best_" + metric, 0, update=False)
  _add_argument(hparams, "best_" + metric + "_dir", best_metric_dir)

  return hparams


def create_or_load_hparams(default_hparams, hparams_path):
  """Create hparams or load hparams from output_dir."""
  hparams = utils.maybe_parse_standard_hparams(default_hparams, hparams_path)
  hparams = extend_hparams(hparams)
  # Print HParams
  utils.print_hparams(hparams)
  return hparams


def run_main(flags, default_hparams, estimator_fn):
  """Run main."""
  # Random
  random_seed = flags.random_seed
  if random_seed is not None and random_seed > 0:
    utils.print_out("# Set random seed to %d" % random_seed)
    random.seed(random_seed)
    np.random.seed(random_seed)
    tf.set_random_seed(random_seed)

  # Model output directory
  output_dir = flags.output_dir
  if output_dir and not tf.gfile.Exists(output_dir):
    utils.print_out("# Creating output directory %s ..." % output_dir)
    tf.gfile.MakeDirs(output_dir)

  # Load hparams.
  hparams = create_or_load_hparams(default_hparams, flags.hparams_path)

  # Train or Evaluation
  estimator_fn(hparams)
  return hparams

def tokenize(hparams, file, tokenized_file):
  utils.print_out("tokenizing {} -> {}".format(file, tokenized_file))
  with open(file, 'rb') as input_file:
    with open(tokenized_file, 'wb') as output_file:
      subprocess.run([hparams.tokenizer_file, '-l', hparams.src], stdin=input_file, stdout=output_file)

def detokenize(hparams, file, detokenized_file):
  utils.print_out("detokenizing {} -> {}".format(file, detokenized_file))
  with open(file, 'rb') as input_file:
    with open(detokenized_file, 'wb') as output_file:
      subprocess.run([hparams.detokenizer_file, '-l', hparams.tgt], stdin=input_file, stdout=output_file)

def main(unused_argv):
  experiment_start = time.time()

  tf.logging.set_verbosity(tf.logging.INFO)

  if FLAGS.use_fp16 and FLAGS.use_dist_strategy:
    raise ValueError("use_fp16 and use_dist_strategy aren't compatible")

  if FLAGS.use_fp16 + FLAGS.amp + FLAGS.use_fastmath > 1:
    raise ValueError("Only one of use_fp16, amp, use_fastmath can be set")

  if FLAGS.amp:
    utils.print_out('Enabling TF-AMP')

    os.environ['TF_ENABLE_AUTO_MIXED_PRECISION'] = '1'

  if FLAGS.use_fastmath:
    utils.print_out('Enabling FastMath')

    os.environ["TF_ENABLE_CUBLAS_TENSOR_OP_MATH_FP32"] = '1'
    os.environ["TF_ENABLE_CUDNN_TENSOR_OP_MATH_FP32"] = '1'
    os.environ["TF_ENABLE_CUDNN_RNN_TENSOR_OP_MATH_FP32"] = '1'

  # Set up hacky envvars.
  # Hack that affects Defun in attention_wrapper.py
  active_xla_option_nums = np.sum([FLAGS.use_xla, FLAGS.use_autojit_xla,
                                   FLAGS.xla_compile])
  if active_xla_option_nums > 1:
    raise ValueError(
        "Only one of use_xla, xla_compile, use_autojit_xla can be set")

  os.environ["use_xla"] = str(FLAGS.use_xla).lower()
  if FLAGS.use_xla:
    os.environ["use_defun"] = str(True).lower()
  else:
    os.environ["use_defun"] = str(FLAGS.use_defun).lower()
  utils.print_out("use_defun is %s for attention" % os.environ["use_defun"])

  # TODO(jamesqin): retire this config after Cuda9.1
  os.environ["use_fp32_batch_matmul"] = ("true" if FLAGS.use_fp32_batch_matmul
                                         else "false")
  os.environ["xla_compile"] = "true" if FLAGS.xla_compile else "false"
  os.environ["force_inputs_padding"] = (
      "true" if FLAGS.force_inputs_padding else "false")

  if FLAGS.mode == "train":
    utils.print_out("Running training mode.")
    default_hparams = create_hparams(FLAGS)
    run_main(FLAGS, default_hparams, estimator.train_fn)
  elif FLAGS.mode == "infer" or FLAGS.mode == "translate":
    if FLAGS.mode == "infer":
        utils.print_out("Running inference mode.")
        translate_mode = False
    else:
        utils.print_out("Running translate mode on file {}.".format(FLAGS.translate_file))
        translate_mode = True

    # Random
    random_seed = FLAGS.random_seed
    if random_seed is not None and random_seed > 0:
      utils.print_out("# Set random seed to %d" % random_seed)
      random.seed(random_seed)
      np.random.seed(random_seed)
      tf.set_random_seed(random_seed)

    # Model output directory
    output_dir = FLAGS.output_dir
    if output_dir and not tf.gfile.Exists(output_dir):
      utils.print_out("# Creating output directory %s ..." % output_dir)
      tf.gfile.MakeDirs(output_dir)

    dllogger.init(backends=[
        dllogger.StdOutBackend(dllogger.Verbosity.DEFAULT),
        dllogger.JSONStreamBackend(dllogger.Verbosity.VERBOSE, os.path.join(FLAGS.output_dir, FLAGS.mode + '-report.json')),
    ])
    dllogger.log('PARAMETER', vars(FLAGS))

    # Load hparams.
    default_hparams = create_hparams(FLAGS)
    default_hparams.num_buckets = 1
    # The estimator model_fn is written in a way allowing train hparams to be
    # passed in infer mode.
    hparams = create_or_load_hparams(default_hparams, FLAGS.hparams_path)
    utils.print_out("infer_hparams:")
    utils.print_hparams(hparams)

    if translate_mode:
      tokenize(hparams, hparams.translate_file, hparams.translate_file + ".tok")

    eval_sentences, eval_src_tokens, _ = iterator_utils.get_effective_epoch_size(hparams, train=False)

    # Run evaluation when there's a new checkpoint
    tf.logging.info("Starting to evaluate...")
    eval_start = time.time()
    _, (eval_speed, eval_latencies), eval_output_tokens = estimator.eval_fn(hparams, hparams.ckpt, only_translate=translate_mode)
    eval_end = time.time()
    eval_delta = eval_end - eval_start
    utils.print_out("eval time for ckpt: %.2f mins (%.2f sent/sec, %.2f tokens/sec)" %
                    (eval_delta / 60., eval_speed, eval_speed * (eval_src_tokens + eval_output_tokens) / eval_sentences), f=sys.stderr)
    logging_data = {
      'infer_speed_sent': eval_speed,
      'infer_speed_toks': eval_speed * (eval_src_tokens + eval_output_tokens) / eval_sentences,
    }
    for lat in sorted(eval_latencies):
      utils.print_out("eval latency_%s for ckpt: %.2f ms" % (lat, eval_latencies[lat] * 1000))
      logging_data['infer_latency_{}'.format(lat)] = eval_latencies[lat] * 1000

    dllogger.log((), logging_data)
    dllogger.flush()


    if translate_mode:
      detokenize(hparams, hparams.translate_file + ".trans.tok", hparams.translate_file + ".trans")

  else:
    assert FLAGS.mode == "train_and_eval"
    utils.print_out("Running train and eval mode.")

    # Random
    random_seed = FLAGS.random_seed
    if random_seed is not None and random_seed > 0:
      utils.print_out("# Set random seed to %d" % random_seed)
      random.seed(random_seed)
      np.random.seed(random_seed)
      tf.set_random_seed(random_seed)

    # Model output directory
    output_dir = FLAGS.output_dir
    if output_dir and not tf.gfile.Exists(output_dir):
      utils.print_out("# Creating output directory %s ..." % output_dir)
      tf.gfile.MakeDirs(output_dir)

    dllogger.init(backends=[
        dllogger.StdOutBackend(dllogger.Verbosity.DEFAULT),
        dllogger.JSONStreamBackend(dllogger.Verbosity.VERBOSE, os.path.join(FLAGS.output_dir, FLAGS.mode + '-report.json')),
    ])
    dllogger.log('PARAMETER', vars(FLAGS))
    dllogger.metadata("bleu", {"unit": None})
    dllogger.metadata("train_speed_sent", {"unit": "sequences/s"})
    dllogger.metadata("train_speed_toks", {"unit": "tokens/s"})

    # Load hparams.
    default_hparams = create_hparams(FLAGS)

    hparams = create_or_load_hparams(default_hparams, FLAGS.hparams_path)
    utils.print_out("training hparams:")
    utils.print_hparams(hparams)
    with tf.gfile.GFile(os.path.join(output_dir, "train_hparams.txt"), "w") as f:
      f.write(utils.serialize_hparams(hparams) + "\n")

    # The estimator model_fn is written in a way allowing train hparams to be
    # passed in infer mode.
    infer_hparams = tf.contrib.training.HParams(**hparams.values())
    infer_hparams.num_buckets = 1
    utils.print_out("infer_hparams:")
    utils.print_hparams(infer_hparams)
    with tf.gfile.GFile(os.path.join(output_dir, "infer_hparams.txt"), "w") as f:
      f.write(utils.serialize_hparams(infer_hparams) + "\n")

    epochs = 0
    should_stop = epochs >= FLAGS.max_train_epochs

    train_sentences, train_src_tokens, train_tgt_tokens = iterator_utils.get_effective_epoch_size(hparams)
    eval_sentences, eval_src_tokens, _ = iterator_utils.get_effective_epoch_size(hparams, train=False)

    while not should_stop:
      utils.print_out("Starting epoch %d" % epochs)
      try:
        train_start = time.time()
        train_speed, _ = estimator.train_fn(hparams)
      except tf.errors.OutOfRangeError:
        utils.print_out("training hits OutOfRangeError", f=sys.stderr)

      train_end = time.time()
      train_delta = train_end - train_start
      utils.print_out("training time for epoch %d: %.2f mins (%.2f sent/sec, %.2f tokens/sec)" %
                      (epochs + 1, train_delta / 60., train_speed, train_speed * (train_src_tokens + train_tgt_tokens) / train_sentences), f=sys.stderr)
      logging_data = {
        'train_speed_sent': train_speed,
        'train_speed_toks': train_speed * (train_src_tokens + train_tgt_tokens) / train_sentences,
      }

      # This is probably sub-optimal, doing eval per-epoch
      eval_start = time.time()
      bleu_score, (eval_speed, eval_latencies), eval_output_tokens = estimator.eval_fn(infer_hparams)
      eval_end = time.time()
      eval_delta = eval_end - eval_start
      utils.print_out("eval time for epoch %d: %.2f mins (%.2f sent/sec, %.2f tokens/sec)" %
                      (epochs + 1, eval_delta / 60., eval_speed, eval_speed * (eval_src_tokens + eval_output_tokens) / eval_sentences), f=sys.stderr)
      logging_data.update({
        'bleu': bleu_score,
        'infer_speed_sent': eval_speed,
        'infer_speed_toks': eval_speed * (eval_src_tokens + eval_output_tokens) / eval_sentences,
      })
      for lat in sorted(eval_latencies):
        utils.print_out("eval latency_%s for epoch %d: %.2f ms" % (lat, epochs + 1, eval_latencies[lat] * 1000))
        logging_data['eval_latency_{}'.format(lat)] = eval_latencies[lat] * 1000

      dllogger.log((epochs,), logging_data)
      dllogger.flush()


      if FLAGS.debug or (FLAGS.target_bleu is not None and bleu_score > FLAGS.target_bleu):
        should_stop = True
        utils.print_out(
            "Stop job since target bleu is reached at epoch %d ." % epochs,
            f=sys.stderr)

      epochs += 1
      if epochs >= FLAGS.max_train_epochs:
        should_stop = True
        utils.print_out("Stop job since max_train_epochs is reached.",
                        f=sys.stderr)

    dllogger.log((), logging_data)
    dllogger.flush()

  experiment_end = time.time()
  utils.print_out('Experiment took {} min'.format((experiment_end - experiment_start) / 60))


if __name__ == "__main__":
  nmt_parser = argparse.ArgumentParser()
  add_arguments(nmt_parser)
  FLAGS, unparsed = nmt_parser.parse_known_args()
  tf.app.run(main=main, argv=[sys.argv[0]] + unparsed)
