Exports the last checkpoint from the directory as tfhub module.

  It creates the Module spec and signature (based on T2T problem information),
  which is later used to create and export the hub module.
  Module will be saved inside the ckpt_dir.

  Args:
    model_name: name of the model to be exported.
    hparams: T2T parameters, model graph will be based on them.
    decode_hparams: T2T parameters for decoding.
    problem: the name of the problem
    checkpoint_path: path to the checkpoint to be exported.
    export_dir: Directory to write the exported model to.

def export_as_tfhub_module(model_name,
                           hparams,
                           decode_hparams,
                           problem,
                           checkpoint_path,
                           export_dir):
  """Exports the last checkpoint from the directory as tfhub module.

  It creates the Module spec and signature (based on T2T problem information),
  which is later used to create and export the hub module.
  Module will be saved inside the ckpt_dir.

  Args:
    model_name: name of the model to be exported.
    hparams: T2T parameters, model graph will be based on them.
    decode_hparams: T2T parameters for decoding.
    problem: the name of the problem
    checkpoint_path: path to the checkpoint to be exported.
    export_dir: Directory to write the exported model to.
  """

  def hub_module_fn():
    """Creates the TF graph for the hub module."""
    model_fn = t2t_model.T2TModel.make_estimator_model_fn(
        model_name,
        hparams,
        decode_hparams=decode_hparams,
        use_tpu=FLAGS.use_tpu)
    features = problem.serving_input_fn(
        hparams, decode_hparams, use_tpu=FLAGS.use_tpu).features

    # we must do a copy of the features, as the model_fn can add additional
    # entries there (like hyperparameter settings etc).
    original_features = features.copy()
    spec = model_fn(features, labels=None, mode=tf.estimator.ModeKeys.PREDICT)

    hub.add_signature(
        inputs=original_features,
        outputs=spec.export_outputs["serving_default"].outputs)

  # TFHub doesn't support the following collections.
  drop_collections = [tf.GraphKeys.LOSSES,
                      tf.GraphKeys.SUMMARIES, tf.GraphKeys.LOCAL_VARIABLES]
  module_spec = hub.create_module_spec(
      hub_module_fn, drop_collections=drop_collections)
  # Loads the weights from the checkpoint using the model above
  # and saves it in the export_path.
  export_module_spec_with_checkpoint(
      module_spec,
      checkpoint_path=checkpoint_path,
      export_path=export_dir,
      scope_prefix="")