Python gin.tf() Examples

def get_eval_hooks(self, config, params):
    """Get eval_hooks to be passed to estimator spec."""
    logging.warning('This function is deprecated and will be replaced.')
    hooks = []
    summary_op = tf.summary.merge_all()
    if summary_op is not None:
      eval_name = 'eval'
      if params is not None:
        eval_name = params.get('eval_name', eval_name)
      hooks = [
          tf.train.SummarySaverHook(
              output_dir=os.path.join(config.model_dir, eval_name),
              save_steps=config.save_summary_steps,
              summary_op=summary_op),
      ]
    return hooks

  #############################################################################
  # END DEPRECATED functions which will be removed soon.
  ############################################################################# 

def _get_latest_checkpoint_from_dir(model_dir):
  """Helper function to return the latest checkpoint number from a directory.

  Args:
    model_dir: str, Directory with checkpoint files.

  Returns:
    an int, latest checkpoint number.

  Raises:
    ValueError: if no checkpoints are found.
  """
  ckpt = tf.train.latest_checkpoint(model_dir)
  if ckpt is None:
    raise ValueError("No checkpoints found in model directory: %s" % model_dir)
  return int(re.sub(".*ckpt-", "", ckpt)) 

def clean_decodes(ids, eos_id=1, pad_id=0, length_axis=-1):
  """Replaces everything after EOS with PAD (along last axis).

  Args:
    ids: a d Tensor of type int.
    eos_id: int, EOS id.
    pad_id: int, PAD id.
    length_axis: an integer.

  Returns:
    a Tensor of type int of ids.
  """
  eos_and_after = tf.cumsum(tf.cast(tf.equal(ids, eos_id), tf.int32),
                            exclusive=True, axis=length_axis)
  valid_ids = tf.equal(eos_and_after, 0)
  return tf.where_v2(valid_ids, ids, pad_id) 

def get_variable_dtype(
    master_dtype=tf.bfloat16,
    slice_dtype=tf.float32,
    activation_dtype=tf.float32):
  """Datatypes to use for the run.

  Args:
    master_dtype: string, datatype for checkpoints
      keep this the same between training and eval/inference
    slice_dtype: string, datatype for variables in memory
      must be tf.float32 for training
    activation_dtype: string, datatype for activations
      less memory usage if tf.bfloat16 but possible numerical issues
  Returns:
    a mtf.VariableDtype
  """
  return mtf.VariableDType(
      master_dtype=tf.as_dtype(master_dtype),
      slice_dtype=tf.as_dtype(slice_dtype),
      activation_dtype=tf.as_dtype(activation_dtype)) 

def metric_sum(values, name=None, **kwargs):
  del kwargs
  with tf.variable_scope(name, "metric_sum", [values]):
    accum = tf.get_variable(
        "accum", shape=[], dtype=tf.float32, trainable=False,
        collections=[tf.GraphKeys.LOCAL_VARIABLES],
        initializer=tf.zeros_initializer())
    update_op = tf.assign_add(accum, tf.reduce_sum(tf.cast(values, tf.float32)))
    return accum, update_op 

def record_operative_gin_configurations(operative_config_dir):
  """Record operative Gin configurations in the given directory."""
  gin_log_file = operative_config_path(operative_config_dir)
  # If it exists already, rename it instead of overwriting it.
  # This just saves the previous one, not all the ones before.
  if tf.io.gfile.exists(gin_log_file):
    tf.io.gfile.rename(gin_log_file, gin_log_file + '.old', overwrite=True)
  with tf.io.gfile.GFile(gin_log_file, 'w') as f:
    f.write(gin.operative_config_str()) 

def main(argv):
  tf.compat.v1.enable_resource_variables()
  FLAGS(argv)  # raises UnrecognizedFlagError for undefined flags
  tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.INFO)
  gin.parse_config_files_and_bindings(FLAGS.gin_file, FLAGS.gin_param,
                                      skip_unknown=False)
  train_eval(FLAGS.root_dir, FLAGS.experiment_name,
             train_eval_dir=FLAGS.train_eval_dir) 

def _score_with_estimator(estimator, input_fn, eval_checkpoint_step, model_dir,
                          scores_filename, num_examples=None):
  """For each example returned by input_fn, compute log likelihood.

  Args:
    estimator: a TPUEstimator
    input_fn: a function that that returns a tf.data.Dataset with examples
      containing the string field 'targets' and optionally the field 'inputs'
    eval_checkpoint_step: int, list of ints, or None, see `eval_model`
      docstring.
    model_dir: string, estimator model_dir
    scores_filename: a string (path of file to write scores to)
    num_examples: int, the total # of examples being scored, None if unknown

  Returns:
    a list of floats
  """
  checkpoint_path, = get_checkpoint_iterator(eval_checkpoint_step, model_dir)

  result_iter = estimator.predict(input_fn, checkpoint_path=checkpoint_path)
  scores = [m["scores"] for m in result_iter]
  # Remove any padding examples
  scores = scores[:num_examples]
  if scores_filename is not None:
    write_lines_to_file(["%f" % f for f in scores], scores_filename)
  return scores 

def write_lines_to_file(lines, filename):
  """Write each line to a filename, replacing the file if it exists.

  Args:
    lines: list of str, lines to write out.
    filename: str, path to filename.
  """
  if tf.io.gfile.exists(filename):
    tf.io.gfile.remove(filename)
  with tf.io.gfile.GFile(filename, "w") as output_file:
    for line in lines:
      output_file.write("{}\n".format(line)) 

def decode(estimator,
           input_fn,
           vocabulary,
           checkpoint_path=None):
  """Decode from an input_fn.

  Args:
    estimator: a TPUEstimator
    input_fn: function that returns a tf.Dataset
    vocabulary: a vocabulary.Vocabulary or (inputs_vocabulary,
      targets_vocabulary) tuple
    checkpoint_path: an optional string

  Returns:
    list of decoded strings
  """
  result_iter = estimator.predict(
      input_fn, checkpoint_path=checkpoint_path)

  def _maybe_detokenize(value, vocab):
    if isinstance(value, six.binary_type):
      return value
    return vocab.decode([int(x) for x in value])

  decodes = []
  for i, result in enumerate(result_iter):
    input_string = _maybe_detokenize(
        result["inputs"], inputs_vocabulary(vocabulary))
    output_string = _maybe_detokenize(
        result["outputs"], targets_vocabulary(vocabulary))
    decodes.append(output_string)
    if i & (i - 1) == 0:
      # LOG every power of 2.
      tf.logging.info("decoded {}: {}".format(i, input_string))
      tf.logging.info("            -> {}".format(output_string))
  return decodes 

def get_inputs_from_file(input_filename, ignore_comments=False):
  """Read data from file and strip new lines."""
  inputs = [line.rstrip() for line in tf.io.gfile.GFile(input_filename)]

  # Strip the last empty line.
  if not inputs[-1]:
    inputs.pop()

  if ignore_comments:
    inputs = [l for l in inputs if not l.startswith("#")]

  return inputs 

def metric_max(values, name=None, **kwargs):
  del kwargs
  with tf.variable_scope(name, "metric_max", [values]):
    accum = tf.get_variable(
        "accum", shape=[], dtype=tf.float32, trainable=False,
        collections=[tf.GraphKeys.LOCAL_VARIABLES],
        initializer=tf.zeros_initializer())
    update_op = tf.assign(
        accum, tf.maximum(accum, tf.reduce_max(tf.cast(values, tf.float32))))
    return accum, update_op 

def default_init_from_checkpoint_fn(checkpoint,
                                    allow_partial_restore = False):
  """init_from_checkpoint_fn that can be used to init a model from a checkpoint.

  Args:
    checkpoint: String pointing to path of TF checkpoint.
    allow_partial_restore: If True, we allow partial restore, otherwise we raise
      an error if a variable cannot be restored.

  Raises:
    A ValueError if a variable(s) is missing and partial restore is not
    explicitly enabled.
  """
  logging.info('Initializing model weights from %s', checkpoint)
  reader = tf.train.load_checkpoint(checkpoint)
  variables_to_restore = contrib_framework.get_variables()
  assignment_map = {}
  for v in variables_to_restore:
    op_name = v.op.name
    if reader.has_tensor(op_name):
      logging.info('Loading variable %s from checkpoint', op_name)
      assignment_map[op_name] = v
    elif allow_partial_restore:
      logging.warning('Variable %s is not in the checkpoint, skipping.',
                      op_name)
    else:
      raise ValueError('Attempting to restore variable {} which is '
                       'not in the checkpoint.'.format(op_name))

  tf.train.init_from_checkpoint(checkpoint, assignment_map) 

def get_feature_specification(
      self, mode):
    """Returns the feature specification with bfloat16 replacing float32."""
    return tensorspec_utils.replace_dtype(
        self._t2r_model.get_feature_specification(mode),
        from_dtype=tf.float32,
        to_dtype=tf.bfloat16) 

def get_session_config(self):
    """Get the session config for Estimator model.

    Defaults to None which tells tf.Estimator to use its default session config.
    Not used in TPU jobs at the moment.

    Returns:
      None, or the desired session config.
    """
    return None 

def get_run_config(self):
    """Get the RunConfig for Estimator model.

    Returns:
      tf.estimator.RunConfig() for this model.
    """
    return gin_configurable_run_config_cls(
        session_config=self.get_session_config()) 

def get_label_specification(self, mode):
    """Returns the label specification with bfloat16 replacing float32."""
    return tensorspec_utils.replace_dtype(
        self._t2r_model.get_label_specification(mode),
        from_dtype=tf.float32,
        to_dtype=tf.bfloat16) 

def get_feature_specification(
      self, mode):
    """Required features for the model_fn/model_inference_fn.

    Note, the model_fn might use additional features for debugging/development
    purposes. The create_export_outputs_fn will however only require the
    specified required features. Only this subset of features will be used to
    generate automatic tf.Example extractors and numpy placeholders for the
    serving models.

    Args:
      mode: The mode for feature specifications
    """ 

def __init__(self,
               preprocessor_cls=None,
               create_optimizer_fn=optimizers.default_create_optimizer_fn,
               device_type = DEVICE_TYPE_CPU,
               summarize_gradients = True,
               use_sync_replicas_optimizer = False,
               use_avg_model_params = False,
               init_from_checkpoint_fn=None):
    """Base constructor to be used by subclass.

    Args:
      preprocessor_cls: (Optional) A class derived from
        preprocessors.AbstractPreprocessor.
      create_optimizer_fn: A callable function which returns an instance of a
        subclass of tf.train.Optimizer. We couldn't take an optimizer instance
        here because some optimizer's constructor may need to access the
        variables in the graph, which will be created by Estimator when calling
        model_fn. More precisely we will only create an instance during training
        (mode == ModeKeys.TRAIN) within the _optimizer property call which will
        wrap the optimizer instance for GPU towers or TPUs if necessary. The
        _optimizer property is only used within create_train_op function.
      device_type: The device type this model will be deployed on (
        DEVICE_TYPE_CPU, DEVICE_TYPE_GPU, DEVICE_TYPE_TPU).
      summarize_gradients: If True summaries for the gradients produced by the
        train_op will be created. Note, we will automatically disable these
        summaries in case of DEVICE_TYPE_TPU.
      use_sync_replicas_optimizer: If True, synchronize gradient updates from
        the different replicas. (GPU-only, since TPUs are already synchronous).
      use_avg_model_params: During training use a MovingAverageOptimizer and
        swapping saver to compute a running average of the model variables for
        inference.
      init_from_checkpoint_fn: A function that calls
        tf.train.init_from_checkpoint.
    """
    self._preprocessor_cls = preprocessor_cls
    self._create_optimizer_fn = create_optimizer_fn
    self._device_type = device_type
    self._summarize_gradients = summarize_gradients
    self._use_sync_replicas_optimizer = use_sync_replicas_optimizer
    self._sync_replicas_optimizer = None
    self._use_avg_model_params = use_avg_model_params
    self._init_from_checkpoint_fn = init_from_checkpoint_fn
    self._optimizer = None  # type: Optional[tf.train.Optimizer]
    self._scaffold_fn = tf.train.Scaffold 

def compute_summaries(metrics,
                      environment,
                      policy,
                      num_episodes=1,
                      num_episodes_to_render=1,
                      images_ph=None,
                      images_summary=None,
                      render_images_summary=None):
  for metric in metrics:
    metric.reset()

  if num_episodes_to_render:
    environment.start_rendering()

  time_step = environment.reset()
  policy_state = policy.get_initial_state(environment.batch_size)

  render_images = []
  if num_episodes_to_render and 'pixels' in time_step.observation:
    images = [[time_step.observation['pixels']]]
  else:
    images = []
  step = 0
  episode = 0
  while episode < num_episodes:
    action_step = policy.action(time_step, policy_state)
    next_time_step = environment.step(action_step.action)

    traj = trajectory.from_transition(time_step, action_step, next_time_step)
    for observer in metrics:
      observer(traj)

    if episode < num_episodes_to_render:
      if traj.is_last():
        render_images.append(list(environment.frames))
        environment.frames[:] = []
        if episode + 1 >= num_episodes_to_render:
          environment.stop_rendering()

      if 'pixels' in time_step.observation:
        if traj.is_boundary():
          images.append([])
        images[-1].append(next_time_step.observation['pixels'])

    episode += np.sum(traj.is_last())
    step += np.sum(~traj.is_boundary())

    time_step = next_time_step
    policy_state = action_step.state

  py_metric.run_summaries(metrics)

  if render_images:
    render_images = pad_and_concatenate_videos(render_images)
    session = tf.compat.v1.get_default_session()
    session.run(render_images_summary, feed_dict={images_ph: [render_images]})
  if images:
    images = pad_and_concatenate_videos(images)
    session = tf.compat.v1.get_default_session()
    session.run(images_summary, feed_dict={images_ph: [images]}) 

def create_train_op(
      self,
      loss,
      optimizer,
      update_ops = None,
      train_outputs = None,
      filter_trainables_fn = None):  # pylint: disable=line-too-long
    """Create the train_op of from the loss obtained from model_train_fn.

    Args:
      loss: The loss we compute within model_train_fn.
      optimizer: An instance of `tf.train.Optimizer`.
      update_ops: List of update ops to execute alongside the training op.
      train_outputs: (Optional) A dict with additional tensors the training
        model generates.
      filter_trainables_fn: (Optional) A function that takes a trainable
        TensorFlow variable and returns whether it should be updated or not.
        By default, all trainable variables are updated.

    Returns:
      train_op: Op for the training step.
    """
    summarize_gradients = self._summarize_gradients
    if self.is_device_tpu:
      # TPUs don't support summaries up until now. Hence, we overwrite the user
      # provided summarize_gradients option to False.
      if self._summarize_gradients:
        logging.info('We cannot use summarize_gradients on TPUs.')
      summarize_gradients = False
    variables_to_train = None
    if filter_trainables_fn is not None:
      logging.info('Filtering trainable variables')
      variables_to_train = [
          var for var in tf.trainable_variables() if filter_trainables_fn(var)]
      logging.info('Only updating the following trainables:')
      for var in variables_to_train:
        logging.info('  %s', var.name)
    return contrib_training.create_train_op(
        loss,
        optimizer,
        summarize_gradients=summarize_gradients,
        update_ops=update_ops,
        variables_to_train=variables_to_train) 

def serialize_num_microbatches(batch_dim,
                               sequence_length,
                               mesh_shape,
                               layout_rules,
                               tokens_per_microbatch_per_replica=None):
  """Number of microbatches per batch for serialized training.

  We want to split each training step into multiple sequential steps
  to limit memory usage.  Gradients are accumulated locally and reduced once.

  This function determines the number of microbatches per batch.
  If tokens_per_microbatch_per_replica=None, then the batch is not split.

  Args:
    batch_dim: a mtf.Dimension
    sequence_length: an integer or a dict from feature-key to integer
      the (packed) sequence length, e.g. {"inputs": 512, "targets": 128}
    mesh_shape: an input to mtf.convert_to_shape()
    layout_rules: an input to mtf.convert_to_layout_rules()
    tokens_per_microbatch_per_replica: an optional integer, e.g. 2048
  Returns:
    an integer
  """
  if not tokens_per_microbatch_per_replica:
    return 1
  batch_per_replica = mtf.tensor_dim_to_size_per_split(
      layout_rules, mesh_shape, batch_dim)
  # number of sequences per microbatch
  microbatch_size = max(
      1, tokens_per_microbatch_per_replica // max(sequence_length.values()))
  # decrease microbatch_size until it is a divisor of batch_per_replica
  # This is guaranteed to stop at microbatch_size=1 if not earlier.
  while batch_per_replica % microbatch_size:
    microbatch_size -= 1
  num_microbatches = batch_per_replica // microbatch_size
  tf.logging.info(
      "serialize_num_microbatches: "
      "tokens_per_microbatch_per_replica=%d "
      "batch_dim=%s "
      "sequence_length=%s "
      "batch_per_replica=%d "
      "num_microbatches=%d",
      tokens_per_microbatch_per_replica,
      batch_dim,
      sequence_length,
      batch_per_replica,
      num_microbatches)
  return num_microbatches 

def export_model(estimator, export_dir, vocabulary, sequence_length,
                 batch_size=1, checkpoint_path=None):
  """Export a model in TF SavedModel format to be used for inference on CPUs.

  Args:
    estimator: Estimator object, estimator created with the appropriate
      model_fn.
    export_dir: str, a directory in which to create timestamped subdirectories
      containing exported SavedModels.
    vocabulary: sentencepiece vocab, vocabulary instance to use for encoding.
    sequence_length: an integer or a dict from feature-key to integer
      the (packed) sequence length, e.g. {"inputs": 512, "targets": 128}
    batch_size: int, number of sequences per batch. Should match estimator.
    checkpoint_path: str, path to checkpoint. If None (default), use the most
      recent in the model directory.

  Returns:
    The string path to the exported directory.
  """

  def serving_input_fn():
    """Constructs input portion of Graph in serving.

    Input is a batch of strings.

    Returns:
      a ServingInputReceiver
    """
    inputs = tf.placeholder(
        dtype=tf.string,
        shape=[None],
        name="inputs")

    padded_inputs = tf.pad(inputs, [(0, tf.mod(-tf.size(inputs), batch_size))])

    dataset = tf.data.Dataset.from_tensor_slices(padded_inputs)
    dataset = dataset.map(lambda x: {"inputs": x})
    dataset = transformer_dataset.encode_all_features(dataset, vocabulary)
    dataset = transformer_dataset.pack_or_pad(
        dataset=dataset,
        length=sequence_length,
        pack=False,
        feature_keys=["inputs"]
    )

    dataset = dataset.batch(batch_size)

    features = tf.data.experimental.get_single_element(dataset)
    return tf.estimator.export.ServingInputReceiver(
        features=features, receiver_tensors=inputs)

  return estimator.export_saved_model(
      export_dir, serving_input_fn, checkpoint_path=checkpoint_path) 

def score_from_dataset(estimator, vocabulary, batch_size, sequence_length,
                       model_dir, eval_checkpoint_step, dataset_split,
                       score_dataset_fn=None, scores_filename=gin.REQUIRED):
  """Compute log likelihoods per example and write to a text file.



  The function returns a list of floats represnenting the log-liekelihood of the
  target given the input.  If `scores_filename` is present, then these are also
  written out as a text file, one per line.

  Args:
    estimator: a TPUEstimator
    vocabulary: a mtf.transformer.vocabulary.Vocabulary
    batch_size: an integer
    sequence_length: an integer or a dict from feature-key to integer
      the (packed) sequence length, e.g. {"inputs": 512, "targets": 128}
    model_dir: string, estimator model_dir
    eval_checkpoint_step: int, list of ints, or None, see `eval_model`
      docstring.
        dataset_split: a string
    score_dataset_fn: A function returning a list of dataset.EvalDataset tuples.
      See `eval_dataset_fn` argument to `eval_model` for details.
    scores_filename: a string (path of file to write)

  Returns:
    a list of floats
  """
  scoring_datasets = score_dataset_fn(
      sequence_length=sequence_length,
      vocabulary=vocabulary,
      dataset_split=dataset_split)
  if len(scoring_datasets) != 1:
    raise ValueError("Only scoring from a single dataset supported.")
  scoring_dataset = scoring_datasets[0]

  def input_fn(params):
    """Eval input function for estimator."""
    del params
    dataset = scoring_dataset.dataset_fn()
    dataset = dataset.map(_filter_features)
    dataset = dataset.batch(batch_size, drop_remainder=False)
    # Pad the final batch.
    dataset = transformer_dataset.trim_and_pad_dataset(
        dataset, length=batch_size)
    dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
    return dataset

  # TODO(dei): Since we pass in the num_examples as None, scores for the
  # padding examples will get written to the output file. Should fix this.
  return _score_with_estimator(estimator, input_fn, eval_checkpoint_step,
                               model_dir, scores_filename, None) 

def inference_network_fn(
      self,
      features,
      labels,
      mode,
      config = None,
      params = None):
    """The inference network implementation.

    This creates the main network based on features.
    Optionally (mode=ModeKeys.TRAIN or ModeKeys.EVAL) the model can do
    additional processing on labels, however, it has to be ensured that this is
    optional and the graph is fully operational without labels. At inference
    time we will have no access to labels. Tensors which are required for loss
    computation or debugging must be put into the inference_outputs dict.
    Having a dedicated inference_network_fn allows to compose new networks by
    using other TFModels.

    Please, use the following pattern to add not supported tpu model components
    such as tf.summary.*
    if self.use_summaries(params):
      # Do operations which are not supported on tpus.

    If your model does not support TPUs at all, please call the following
    function.
    self.raise_no_tpu_support()

    Args:
      features: This is the first item returned from the input_fn and parsed by
        tensorspec_utils.validate_and_pack. A spec_structure which fulfills the
        requirements of the self.get_feature_specification.
      labels: This is the second item returned from the input_fn and parsed by
        tensorspec_utils.validate_and_pack. A spec_structure which fulfills the
        requirements of the self.get_feature_specification.
      mode: (ModeKeys) Specifies if this is training, evaluation or prediction.
      config: (Optional tf.estimator.RunConfig or contrib_tpu.RunConfig) Will
        receive what is passed to Estimator in config parameter, or the default
        config (tf.estimator.RunConfig). Allows updating things in your model_fn
        based on  configuration such as num_ps_replicas, or model_dir.
      params: An optional dict of hyper parameters that will be passed into
        input_fn and model_fn. Keys are names of parameters, values are basic
        python types. There are reserved keys for TPUEstimator, including
        'batch_size'.

    Returns:
      inference_outputs: A dict with output tensors.
    """ 

def model_train_fn(self,
                     features,
                     labels,
                     inference_outputs,
                     mode,
                     config = None,
                     params = None):
    """The training model implementation.

    This model_fn should add the loss computation based on the inference_outputs
    and labels. For better debugging we also provide access to the input
    features. Note, no new variables should be generated in this model_fn since
    the model_inference_fn and the maybe_init_from_checkpoint function would
    not have access to these variables. We output the final loss (scalar) and
    a dict of optional train_outputs which might be useful for the
    model_eval_fn.

    Please, use the following pattern to add not supported tpu model components
    such as tf.summary.*
    if self.use_summaries(params):
      # Do operations which are not supported on tpus.

    If your model does not support TPUs at all, please call the following
    function.
    self.raise_no_tpu_support()

    Args:
      features: This is the first item returned from the input_fn and parsed by
        tensorspec_utils.validate_and_pack. A spec_structure which fulfills the
        requirements of the self.get_feature_specification.
      labels: This is the second item returned from the input_fn and parsed by
        tensorspec_utils.validate_and_pack. A spec_structure which fulfills the
        requirements of the self.get_feature_specification.
      inference_outputs: A dict containing the output tensors of
        model_inference_fn.
      mode: (ModeKeys) Specifies if this is training, evaluation or prediction.
      config: (Optional tf.estimator.RunConfig or contrib_tpu.RunConfig) Will
        receive what is passed to Estimator in config parameter, or the default
        config (tf.estimator.RunConfig). Allows updating things in your model_fn
        based on  configuration such as num_ps_replicas, or model_dir.
      params: An optional dict of hyper parameters that will be passed into
        input_fn and model_fn. Keys are names of parameters, values are basic
        python types. There are reserved keys for TPUEstimator, including
        'batch_size'.

    Returns:
      loss: The loss we will optimize.
      train_outputs: (Optional) A dict with additional tensors the training
        model generates. We output these tensors such that model_eval_fn could
        introspect these tensors.
    """ 

def decode_from_file(estimator,
                     vocabulary,
                     model_type,
                     batch_size,
                     sequence_length,
                     checkpoint_path=None,
                     input_filename=gin.REQUIRED,
                     output_filename=gin.REQUIRED,
                     eos_id=1,
                     repeats=1):
  """Decode from a text file and write to output_filename.

  Args:
    estimator: a TPUEstimator
    vocabulary: a mtf.transformer.vocabulary.Vocabulary
    model_type: a string
    batch_size: an integer
    sequence_length: an integer or a dict from feature-key to integer
      the (packed) sequence length, e.g. {"inputs": 512, "targets": 128}
    checkpoint_path: an optional string
    input_filename: a string
    output_filename: a string
    eos_id: EOS id
    repeats: an integer, the number of times to repeat each input.
  """
  inputs = get_inputs_from_file(input_filename)

  all_input_ids = encode_inputs(inputs, vocabulary, model_type, batch_size,
                                sequence_length["inputs"], eos_id=eos_id)
  def input_fn(params):
    del params
    dataset = tf.data.Dataset.from_tensor_slices({"inputs": all_input_ids})
    dataset = dataset.flat_map(
        lambda x: tf.data.Dataset.from_tensors(x).repeat(repeats))
    dataset = dataset.batch(batch_size, drop_remainder=True)
    dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
    return dataset

  checkpoint_step = get_step_from_checkpoint_path(checkpoint_path)
  decodes = decode(
      estimator, input_fn, vocabulary, checkpoint_path=checkpoint_path)
  # Remove any padded examples
  dataset_size = len(inputs) * repeats
  decodes = decodes[:dataset_size]
  output_filename = "{}-{}".format(output_filename, checkpoint_step)
  write_lines_to_file(decodes, output_filename) 

def model_eval_fn(self,
                    features,
                    labels,
                    inference_outputs,
                    train_loss,
                    train_outputs,
                    mode,
                    config = None,
                    params = None):
    """The eval model implementation, by default we report the loss for eval.

    This function should add the eval_metrics computation based on the
    inference_outputs, labels and the train_loss. For better debugging we also
    provide access to the input features and the train_outputs. Note, no new
    variables should be generated in this model_fn since the model_inference_fn
    and the maybe_init_from_checkpoint function would not have access to these
    variables.

    Please, use the following pattern to add not supported tpu model components
    such as tf.summary.*
    if self.use_summaries(params):
      # Do operations which are not supported on tpus.

    If your model does not support TPUs at all, please call the following
    function.
    self.raise_no_tpu_support()

    Args:
      features: This is the first item returned from the input_fn and parsed by
        tensorspec_utils.validate_and_pack. A spec_structure which fulfills the
        requirements of the self.get_feature_specification.
      labels: This is the second item returned from the input_fn and parsed by
        tensorspec_utils.validate_and_pack. A spec_structure which fulfills the
        requirements of the self.get_feature_specification.
      inference_outputs: A dict containing the output tensors of
        model_inference_fn.
      train_loss: The final loss from model_train_fn.
      train_outputs: A dict containing the output tensors (dict) of
        model_train_fn.
      mode: (ModeKeys) Specifies if this is training, evaluation or prediction.
      config: (Optional tf.estimator.RunConfig or contrib_tpu.RunConfig) Will
        receive what is passed to Estimator in config parameter, or the default
        config (tf.estimator.RunConfig). Allows updating things in your model_fn
        based on  configuration such as num_ps_replicas, or model_dir.
      params: An optional dict of hyper parameters that will be passed into
        input_fn and model_fn. Keys are names of parameters, values are basic
        python types. There are reserved keys for TPUEstimator, including
        'batch_size'.

    Returns:
      eval_metrics: A tuple of (metric_fn, metric_fn_inputs) where metric_fn
        is a dict with {metric_description: tf.metrics.*}.
    """
    del features, labels, inference_outputs, train_loss, train_outputs
    del mode, config, params
    # By default we don't have any eval_metrics. The loss computation used
    # to optimize the model_fn will be reported for the model_eval_fn as well.
    # Hence, by default the EVAL mode can be used to determine the loss
    # performance on the eval dataset or even a larger train dataset.
    return None 

def add_summaries(self,
                    features,
                    labels,
                    inference_outputs,
                    train_loss,
                    train_outputs,
                    mode,
                    config = None,
                    params = None):
    """Add summaries to the graph.

    Having a central place to add all summaries to the graph is helpful in order
    to compose models. For example, if an inference_network_fn is used within
    a while loop no summaries can be added. This function will allow to add
    summaries after the while loop has been processed.

    Args:
      features: This is the first item returned from the input_fn and parsed by
        tensorspec_utils.validate_and_pack. A spec_structure which fulfills the
        requirements of the self.get_feature_specification.
      labels: This is the second item returned from the input_fn and parsed by
        tensorspec_utils.validate_and_pack. A spec_structure which fulfills the
        requirements of the self.get_feature_specification.
      inference_outputs: A dict containing the output tensors of
        model_inference_fn.
      train_loss: The final loss from model_train_fn.
      train_outputs: A dict containing the output tensors (dict) of
        model_train_fn.
      mode: (ModeKeys) Specifies if this is training, evaluation or prediction.
      config: (Optional tf.estimator.RunConfig or contrib_tpu.RunConfig) Will
        receive what is passed to Estimator in config parameter, or the default
        config (tf.estimator.RunConfig). Allows updating things in your model_fn
        based on  configuration such as num_ps_replicas, or model_dir.
      params: An optional dict of hyper parameters that will be passed into
        input_fn and model_fn. Keys are names of parameters, values are basic
        python types. There are reserved keys for TPUEstimator, including
        'batch_size'.
    """
    del features, labels, inference_outputs, train_loss, train_outputs, mode
    del config
    if not self.use_summaries(params):
      return 

def create_export_outputs_fn(self,
                               features,
                               inference_outputs,
                               mode,
                               config = None,
                               params = None):
    """We export the final output used for model inference.

    This model_fn should create the optional export_outputs, see
    tf.estimator.EstimatorSpec for a more in depth description, and the
    required predictions dict. Note, the predictions dict should more often
    than not be a small subset of the inference_outputs.

    Please, use the following pattern to add not supported tpu model components
    such as tf.summary.*
    if self.use_summaries(params):
      # Do operations which are not supported on tpus.

    If your model does not support TPUs at all, please call the following
    function.
    self.raise_no_tpu_support()

    Args:
      features: This is the first item returned from the input_fn and parsed by
        tensorspec_utils.validate_and_pack. A spec_structure which fulfills the
        requirements of the self.get_feature_specification.
      inference_outputs: A dict containing the output tensors of
        model_inference_fn.
      mode: (ModeKeys) Specifies if this is training, evaluation or prediction.
      config: (Optional tf.estimator.RunConfig or contrib_tpu.RunConfig) Will
        receive what is passed to Estimator in config parameter, or the default
        config (tf.estimator.RunConfig). Allows updating things in your model_fn
        based on  configuration such as num_ps_replicas, or model_dir.
      params: An optional dict of hyper parameters that will be passed into
        input_fn and model_fn. Keys are names of parameters, values are basic
        python types. There are reserved keys for TPUEstimator, including
        'batch_size'.

    Returns:
      predictions: A dict of tensors.
      export_outputs: (Optional) A dict containing an arbitrary name for the
        output and tf.estimator.export.PredictOutput(output_dict) as value.
        The output dict is a {name: tensor} mapping. If None, the default
        mapping for predictions is generated. The export_outputs are used
        for the serving model. Multi-headed models should have one name
        per head.
    """
    del features, mode, config, params
    # By default we will export all outputs generated by the
    # inference_network_fn.
    return inference_outputs