text: """Training related classes and functions."""
  
  import contextlib
  import itertools
  import time
  
  import tensorflow as tf
  
  from opennmt.optimizers import utils as optimizer_util
  
  
  class Trainer:
    """Base class for model trainer, implementing single-GPU training."""
  
    def __init__(self, model, optimizer, checkpoint=None, is_master=True):
      """Initializes the trainer.
  
      Args:
        model: A :class:`opennmt.models.Model` instance to train.
        optimizer: A ``tf.keras.optimizers.Optimizer`` instance.
        checkpoint: A :class:`opennmt.utils.checkpoint.Checkpoint` instance. If
          not set, no checkpoints will be saved.
        is_master: Whether this trainer instance is the master trainer.
      """
      self._checkpoint = checkpoint
      self._is_master = is_master
      self._model = model
      if checkpoint is not None:
        self._summary_writer = tf.summary.create_file_writer(checkpoint.model_dir)
      else:
        self._summary_writer = tf.summary.create_noop_writer()
      self._training_stats = None
      self._gradient_accumulator = optimizer_util.GradientAccumulator()
  
      if optimizer is None:
        raise ValueError("No optimizer is defined")
      graph_optimizer_options = tf.config.optimizer.get_experimental_options()
      mixed_precision_enabled = graph_optimizer_options.get("auto_mixed_precision")
      if (mixed_precision_enabled
          and not isinstance(optimizer, tf.keras.mixed_precision.experimental.LossScaleOptimizer)):
        optimizer = tf.keras.mixed_precision.experimental.LossScaleOptimizer(optimizer, "dynamic")
      self._optimizer = optimizer
  
    @property
    def num_replicas(self):
      """Number of synchronous training replicas."""
      return 1
  
    def __call__(self,
                 dataset,
                 max_step=None,
                 accum_steps=1,
                 report_steps=100,
                 save_steps=5000,
                 evaluator=None,
                 eval_steps=5000,
                 moving_average_decay=None):
      """Runs the training.
  
      Args:
        dataset: A ``tf.data.Dataset`` or a function taking a ``tf.distribute.InputContext``
          instance and returning a ``tf.data.Dataset``.
        max_step: The final training step.
        accum_steps: The number of gradient accumulation steps.
        report_steps: Report status every this many steps.
        save_steps: Save a checkpoint every this many steps.
        evaluator: A :class:`opennmt.evaluation.Evaluator` instance to call for
          evaluation.
        eval_steps: Evaluate every this many steps.
        moving_average_decay: If set, maintain an exponential moving average of the model
          variables using this decay value (usually close to 1, e.g. 0.9999). See
          https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage.
  
      Returns:
        A dictionary with various training statistics.
      """
      if max_step is not None and self._optimizer.iterations.numpy() >= max_step:
        raise RuntimeError("The training already reached max_step (%d). If you "
                           "want to continue the training, you should increase the "
                           "max_step value in the training parameters." % max_step)
      if evaluator is not None and evaluator.should_stop():
        raise RuntimeError("The early stopping conditions are already met. If you "
                           "want to continue the training, you should update your "
                           "early stopping parameters.")
  
      self._gradient_accumulator.reset()
  
      with self._summary_writer.as_default():
        self._training_stats = TrainingStats(
            self._model, self._optimizer, reduce_fn=self._all_reduce_sum)
        iterations = self._optimizer.iterations
        tf.summary.experimental.set_step(iterations)
  
        step = None
        moving_average = None
        for loss in self._steps(dataset, accum_steps=accum_steps, report_steps=report_steps):
          if tf.math.is_nan(loss):
            raise RuntimeError("Model diverged with loss = NaN.")
  
          if moving_average_decay is not None and self._is_master:
            if moving_average is None:
              moving_average = MovingAverage(
                  self._model.trainable_variables,
                  iterations,
                  decay=moving_average_decay)
            self._update_moving_average(moving_average)
  
          step = iterations.numpy()
          reset_throughput = False
          self._training_stats.update_on_step(step, loss)
          if step % report_steps == 0:
            self._training_stats.log(self._is_master)
            reset_throughput = True
          if step == 1 or (save_steps is not None and step % save_steps == 0):
            self._save_checkpoint(step, moving_average=moving_average)
            reset_throughput = True
          if eval_steps is not None and step % eval_steps == 0:
            early_stop = self._evaluate(evaluator, step, moving_average=moving_average)
            reset_throughput = True
            if early_stop:
              tf.get_logger().warning("Early stopping conditions are met. Exiting.")
              break
          if step == max_step:
            break
          if reset_throughput:
            self._training_stats.reset_throughput()
  
        if step is None:
          raise RuntimeError("No training steps were executed. This usually means the "
                             "training file is empty or all examples were filtered out. "
                             "For the latter, verify that the maximum_*_length values are "
                             "consistent with your data.")
        summary = self._training_stats.get_global_summary()
        self._save_checkpoint(step, moving_average=moving_average)
        self._evaluate(evaluator, step, moving_average=moving_average)
        return summary
  
    def _save_checkpoint(self, step, moving_average=None):
      """Saves a checkpoint for step."""
      if not self._is_master or self._checkpoint is None or step == self._checkpoint.last_saved_step:
        return
      with moving_average.shadow_variables() if moving_average is not None else contextlib.suppress():
        self._checkpoint.save(step)
  
    def _evaluate(self, evaluator, step, moving_average=None):
      """Runs evaluation for step. Returns ``True`` is early conditions are met."""
      if not self._is_master or evaluator is None or step == evaluator.last_evaluated_step:
        return False
      with moving_average.shadow_variables() if moving_average is not None else contextlib.suppress():
        evaluator(step)
        return evaluator.should_stop()
  
    def _finalize_dataset(self, dataset):
      """Returns the final dataset instance to be used for training.
  
      Args:
        dataset: A ``tf.data.Dataset`` or a function taking a ``tf.distribute.InputContext``
          instance and returning a ``tf.data.Dataset``.
  
      Returns:
        A ``tf.data.Dataset``.
      """
      if callable(dataset):
        dataset = dataset(tf.distribute.InputContext())
      return dataset
  
    def _steps(self, dataset, accum_steps=1, report_steps=None):
      """Returns a generator over training steps (i.e. parameters update).
  
      Args:
        dataset: The training dataset.
        accum_steps: Accumulate the gradients of this many steps/batches.
        report_steps: Report summary statistics every this many steps. This should
          typically be used in a ``tf.summary.record_if`` context.
  
      Returns:
        A generator that yields a loss value to report for this step.
      """
      dataset = self._finalize_dataset(dataset)
      iterator = iter(dataset)
      with_accum = accum_steps > 1
  
      # We define 2 separate functions to support gradient accumulation:
      #  * forward: compute and accumulate the gradients
      #  * step: apply the gradients
      # When gradient accumulation is disabled, the forward function also applies the gradients.
  
      def _forward():
        # We get the next dataset element within the function for increased efficiency
        # and avoid dealing with tf.function input signatures.
        source, target = next(iterator)
        record_summaries = self._should_record_summaries(report_steps, with_accum=with_accum)
        forward_fn = self._forward if with_accum else self._forward_and_step
        return forward_fn(source, target, record_summaries=record_summaries)
  
      def _step():
        return self._step()
  
      # Wrap forward and step with tf.function to run in graph mode.
      forward_fn = tf.function(_forward)
      step_fn = tf.function(_step) if with_accum else lambda: None
  
      for i in itertools.count():
        try:
          loss = forward_fn()
        except (StopIteration, tf.errors.OutOfRangeError):  # Dataset iterator exhausted.
          break
  
        if i == 0 or (i + 1) % accum_steps == 0:
          step_fn()
          if i == 0:
            self._broadcast_variables()
          yield loss
  
    def _run_model(self, source, target):
      """Computes the loss of the given source and target pair.
  
      Args:
        source: A nested structure of tensors.
        target: A nested structure of tensors.
  
      Returns:
        A tuple containing,
  
        - The loss to compute the gradients.
        - The loss to report.
      """
      first_call = not self._model.built
      outputs, _ = self._model(
          source,
          labels=target,
          training=True,
          step=self._optimizer.iterations)
      loss = self._model.compute_loss(outputs, target, training=True)
      if isinstance(loss, tuple):
        training_loss = loss[0] / loss[1]
        reported_loss = loss[0] / loss[2] if len(loss) > 2 else training_loss
      else:
        training_loss, reported_loss = loss, loss
      training_loss = self._model.regularize_loss(
          training_loss, variables=self._model.trainable_variables)
      self._training_stats.update_on_example(source, target)
      if first_call and self._is_master:
        if self._checkpoint is not None:
          self._model.visualize(self._checkpoint.model_dir)
        tf.get_logger().info("Number of model parameters: %d", self._model.count_params())
        tf.get_logger().info(
            "Number of model weights: %d (trainable = %d, non trainable = %d)",
            len(self._model.weights),
            len(self._model.trainable_weights),
            len(self._model.non_trainable_weights))
      return training_loss, reported_loss
  
    def _should_record_summaries(self, report_steps, with_accum=False):
      """Returns a boolean tensor to be used in tf.summary.record_if."""
      if report_steps is None or not self._is_master:
        return False
      record_summaries = tf.equal(self._optimizer.iterations % report_steps, 0)
      if with_accum:
        record_summaries = tf.logical_and(
            record_summaries,
            tf.equal(self._gradient_accumulator.step, 0))
      return record_summaries
  
    def _compute_gradients(self, source, target, record_summaries=False):
      """Computes the gradient of a training example."""
      with tf.summary.record_if(record_summaries):
        if tf.executing_eagerly():
          with tf.GradientTape() as tape:
            training_loss, reported_loss = self._run_model(source, target)
          gradients = tape.gradient(training_loss, self._model.trainable_variables)
        else:
          training_loss, reported_loss = self._run_model(source, target)
          gradients = self._optimizer.get_gradients(training_loss, self._model.trainable_variables)
        _summarize_gradients(gradients, record_summaries)
      return reported_loss, gradients
  
    def _apply_gradients(self, gradients, scale=1):
      """Applies the gradients."""
      gradient_scale = scale * self.num_replicas
      if tf.is_tensor(gradient_scale) or gradient_scale != 1:
        gradients = [
            self._all_reduce_sum(gradient / tf.cast(gradient_scale, gradient.dtype))
            for gradient in gradients]
      self._optimizer.apply_gradients(list(zip(gradients, self._model.trainable_variables)))
  
    def _forward_and_step(self, source, target, record_summaries=False):
      """Forwards a training example and applies the gradients, without accumulation."""
      loss, gradients = self._compute_gradients(source, target, record_summaries=record_summaries)
      self._apply_gradients(gradients)
      return loss
  
    def _forward(self, source, target, record_summaries=False):
      """Forwards a training example and accumulates the gradients."""
      loss, gradients = self._compute_gradients(source, target, record_summaries=record_summaries)
      self._gradient_accumulator(gradients)
      return loss
  
    def _step(self):
      """Applies gradients and resets accumulation."""
      self._apply_gradients(
          self._gradient_accumulator.gradients,
          scale=self._gradient_accumulator.step)
      self._gradient_accumulator.reset()
  
    def _update_moving_average(self, moving_average):
      """Updates the moving average of variables."""
      moving_average.update()
  
    def _broadcast_variables(self):
      """Broadcasts variables to other replicas, if required."""
      return
  
    def _all_reduce_sum(self, value):
      """Reduces the value across all replicas."""
      return value
  
  
  class HorovodTrainer(Trainer):
    """Trainer compatible with Horovod distributed training."""
  
    def __init__(self, model, optimizer, hvd, checkpoint=None):
      """Initializes the Horovod trainer.
  
      Args:
        model: A :class:`opennmt.models.Model` instance to train.
        optimizer: A ``tf.keras.optimizers.Optimizer`` instance.
        hvd: The global Horovod object.
        checkpoint: A :class:`opennmt.utils.checkpoint.Checkpoint` instance. If
          not set, no checkpoints will be saved.
      """
      super().__init__(model, optimizer, checkpoint=checkpoint, is_master=hvd.rank() == 0)
      self._hvd = hvd
  
    @property
    def num_replicas(self):
      return self._hvd.size()
  
    def _finalize_dataset(self, dataset):
      if callable(dataset):
        dataset = dataset(tf.distribute.InputContext(
            num_input_pipelines=self._hvd.size(),
            input_pipeline_id=self._hvd.rank(),
            num_replicas_in_sync=self._hvd.size()))
      return dataset
  
    def _broadcast_variables(self):
      self._hvd.broadcast_variables(self._model.variables, root_rank=0)
      self._hvd.broadcast_variables(self._optimizer.variables(), root_rank=0)
  
    def _all_reduce_sum(self, value):
      return self._hvd.allreduce(value, op=self._hvd.Sum)
  
  
  class MirroredStrategyTrainer(Trainer):
    """Trainer based on ``tf.distribute.MirroredStrategy`` for local multi-GPU training."""
  
    def __init__(self, model, optimizer, checkpoint=None, devices=None):
      """Initializes the MirroredStrategy trainer.
  
      Args:
        model: A :class:`opennmt.models.Model` instance to train.
        optimizer: A ``tf.keras.optimizers.Optimizer`` instance.
        checkpoint: A :class:`opennmt.utils.checkpoint.Checkpoint` instance. If
          not set, no checkpoints will be saved.
        devices: List of device strings to use for training. If not set, all
          visible GPUs are used.
      """
      super().__init__(model, optimizer, checkpoint=checkpoint)
      self._strategy = tf.distribute.MirroredStrategy(devices=devices)
      with self._strategy.scope():
        # Create some variables under the strategy scope.
        _ = self._optimizer.iterations
  
    @property
    def num_replicas(self):
      return self._strategy.num_replicas_in_sync
  
    def _finalize_dataset(self, dataset):
      # We prefer not to use experimental_distribute_dataset here because it
      # sometimes fails to split the batches (noticed with tokens batch type).
      dataset_fn = dataset if callable(dataset) else lambda _: dataset
      return self._strategy.experimental_distribute_datasets_from_function(dataset_fn)
  
    def _forward_and_step(self, source, target, record_summaries=False):
      per_replica_loss = self._strategy.run(
          super()._forward_and_step,
          args=(source, target),
          kwargs=dict(record_summaries=record_summaries))
      return self._strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_loss, None)
  
    def _forward(self, source, target, record_summaries=False):
      per_replica_loss = self._strategy.run(
          super()._forward,
          args=(source, target),
          kwargs=dict(record_summaries=record_summaries))
      # TODO: this reduction could be delayed until _step is called.
      return self._strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_loss, None)
  
    def _step(self):
      self._strategy.run(super()._step)
  
    def _update_moving_average(self, moving_average):
      with self._strategy.scope():
        super()._update_moving_average(moving_average)
  
  
  def _summarize_gradients(gradients, should_record):
    # Only compute the gradients global norm when the value is actually recorded.
    if isinstance(should_record, bool) and not should_record:
      return
    tf.summary.scalar(
        "gradients/global_norm",
        tf.cond(
            should_record,
            true_fn=lambda: tf.linalg.global_norm(gradients),
            false_fn=lambda: tf.constant(0, dtype=gradients[0].dtype)))
  
  
  class MovingAverage(object):
    """Object holding an exponential moving average of variables."""
  
    def __init__(self, variables, step, decay=0.9999):
      """Initializes the moving average object.
  
      Args:
        variables: The list of variable for which to maintain a moving average.
        step: The training step counter as a ``tf.Variable``.
        decay: The decay rate of the exponential moving average. Usually close to
          1, e.g. 0.9999, see the complete formula on
          https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage.
  
      Raises:
        TypeError: is :obj:`step` is not a ``tf.Variable``.
      """
      if not isinstance(step, tf.Variable):
        raise TypeError("step should be a tf.Variable")
      if decay < 0.9 or decay > 1:
        tf.get_logger().warning("Moving average decay should be close to 1 (e.g. 0.9999) but you "
                                "passed %f, is it correct? See https://www.tensorflow.org/api_docs"
                                "/python/tf/train/ExponentialMovingAverage for details about the "
                                "formula and recommended decay values.")
      self._ema = tf.train.ExponentialMovingAverage(decay, num_updates=step)
      self._variables = variables
  
    @tf.function
    def update(self):
      """Updates the moving average of the variables."""
      self._ema.apply(self._variables)
  
    @contextlib.contextmanager
    def shadow_variables(self):
      """Returns a context manager that assigns the variables to their moving
      average value on enter and restores the previous value on exit.
  
      Returns:
        A context manager.
      """
      previous_values = []
      for variable in self._variables:
        previous_values.append(variable.value())
        variable.assign(self._ema.average(variable))
      yield
      for previous_value, variable in zip(previous_values, self._variables):
        variable.assign(previous_value)
  
  
  class TrainingStats:
    """Aggregate and summarize training statistics."""
  
    def __init__(self, model, optimizer, reduce_fn=None, warmup_steps=2):
      """Initializes the statistics.
  
      Args:
        model: The model.
        optimizer: The optimizer.
        reduce_fn: In case of distributed training, a function to sum reduce
          distributed values.
        warmup_steps: Throughput values are ignored for this many steps at the
          start of the training.
      """
      self._model = model
      self._optimizer = optimizer
      self._reduce_fn = reduce_fn
      self._warmup_steps = warmup_steps
      self._words_counters = {}
      self._num_updates = 0
      self._average_loss = 0
      self._last_loss = None
      self._last_step = optimizer.iterations.numpy()
      self._last_logged_step = self._last_step
      self._last_logged_time = time.time()
  
    def update_on_example(self, source, target):
      """Updates the training statistics on a new training example.
  
      This may be called within a tf.function.
  
      Args:
        source: A dictionary of source features.
        tagret: A dictionary of target features.
      """
      self._update_words_counter("source", source)
      if not self._model.unsupervised:
        self._update_words_counter("target", target)
  
    def update_on_step(self, step, loss):
      """Updates the training statistics on a new training step.
  
      Args:
        step: The current training step.
        loss: The loss for this step.
      """
      # Convert Numpy or Tensor values to Python.
      step = int(step)
      loss = float(loss)
  
      self._last_step = step
      self._last_loss = loss
      self._average_loss = (self._average_loss * self._num_updates + loss) / (self._num_updates + 1)
  
      if self._num_updates < self._warmup_steps:
        self.reset_throughput()
      self._num_updates += 1
  
    def get_last_summary(self):
      """Returns a summary of the training since the last log.
  
      Returns:
        A dictionary containing various statistics.
      """
      elapsed_time = time.time() - self._last_logged_time
      return {
          "learning_rate": self._get_learning_rate(),
          "loss": self._last_loss,
          "step": self._last_step,
          "steps_per_sec": (self._last_step - self._last_logged_step) / elapsed_time,
          "words_per_sec": {
              name:int(value.numpy() / elapsed_time)
              for name, value in self._get_words_counters().items()},
      }
  
    def get_global_summary(self):
      """Returns a summary of the training since the beginning.
  
      Returns:
        A dictionary containing various statistics.
      """
      return {
          "average_loss": self._average_loss,
          "last_learning_rate": self._get_learning_rate(),
          "last_loss": self._last_loss,
          "last_step": self._last_step,
          "num_steps": self._num_updates,
      }
  
    def log(self, is_master=True):
      """Logs the last training statistics.
  
      Args:
        is_master: Whether this process is the master worker or not.
      """
  
      # Only the master should log the training statistics but we build the
      # summary on all workers since it may reduce distributed values.
      summary = self.get_last_summary()
  
      if not is_master:
        return
  
      tf.summary.scalar(
          "steps_per_sec", summary["steps_per_sec"], description="Training steps per second")
      steps_per_sec_fmt = "steps/s = %0.2f" % summary["steps_per_sec"]
  
      words_per_sec_fmt = []
      for name, avg in summary["words_per_sec"].items():
        tf.summary.scalar(
            "words_per_sec/%s" % name,
            avg,
            description="%s words per second" % name.capitalize())
        words_per_sec_fmt.append("%s words/s = %d" % (name, avg))
  
      tf.get_logger().info(
          "Step = %d ; %s ; Learning rate = %f ; Loss = %f",
          summary["step"],
          ", ".join([steps_per_sec_fmt] + list(sorted(words_per_sec_fmt))),
          summary["learning_rate"],
          summary["loss"])
      tf.summary.scalar("loss", summary["loss"], description="Training loss")
      tf.summary.scalar("optim/learning_rate", summary["learning_rate"], description="Learning rate")
  
    def reset_throughput(self):
      """Resets the accumulated values since the last log."""
      self._reset_words_counters()
      self._last_logged_step = self._last_step
      self._last_logged_time = time.time()
  
    def _get_learning_rate(self):
      learning_rate = self._optimizer.learning_rate
      if isinstance(learning_rate, tf.optimizers.schedules.LearningRateSchedule):
        learning_rate = learning_rate(self._last_step)
      return float(learning_rate)
  
    def _update_words_counter(self, name, features):
      """Accumulates the number of source and target tokens to report throughput."""
      length = features.get("length")
      if length is None:
        return
      num_words = tf.reduce_sum(length)
      counter = self._words_counters.get(name)
      if counter is None:
        counter = tf.Variable(
            tf.constant(0, dtype=tf.int64),
            trainable=False,
            synchronization=tf.VariableSynchronization.ON_READ,
            aggregation=tf.VariableAggregation.SUM)
        self._words_counters[name] = counter
      counter.assign_add(tf.cast(num_words, tf.int64))
  
    @tf.function
    def _get_words_counters(self):
      """Returns the accumulated words counters.
  
      Returns:
        A dictionary mapping a counter name to a value.
      """
      counters = {}
      for name, counter in self._words_counters.items():
        counter = counter.read_value()
        if self._reduce_fn is not None:
          counter = self._reduce_fn(counter)
        counters[name] = counter
      return counters
  
    @tf.function
    def _reset_words_counters(self):
      """Resets the accumulated words counters."""
      for counter in self._words_counters.values():
        counter.assign(tf.constant(0, dtype=tf.int64))
  