# Buggy code
def serve_function(self):
    """Returns a function for serving this model.

    Returns:
      A ``tf.function``.
    """
    # Set name attribute of the input TensorSpec.
    input_signature = {
        name: tf.TensorSpec.from_spec(spec, name=name)
        for name, spec in self.features_inputter.input_signature().items()
    }

    @tf.function(input_signature=(input_signature,))
    def _run(features):
        features = self.features_inputter.make_features(features=features.copy())
        _, predictions = self(features)
        return predictions

    return _run


# Traceback (TypeError)
runner.export(export_dir='/gdrive/My Drive/nmt/chat50w/models')

INFO:tensorflow:Using parameters:
data:
train_features_file: /gdrive/My Drive/nmt/chat50w/train_tokens.txt
vocabulary: /gdrive/My Drive/nmt/chat50w/vocab.txt
eval:
batch_size: 32
infer:
batch_size: 16
length_bucket_width: 1
model_dir: /gdrive/My Drive/nmt/chat50w/run/
params:
average_loss_in_time: true
decay_params:
max_step: 1000000
warmup_steps: 2000
decay_type: CosineAnnealing
learning_rate: 0.00025
num_hypotheses: 1
optimizer: Adam
score:
batch_size: 64
train:
batch_size: 32
batch_type: examples
length_bucket_width: 1
maximum_features_length: 512
sample_buffer_size: 500000
save_summary_steps: 100

INFO:tensorflow:Restored checkpoint /gdrive/My Drive/nmt/chat50w/run/ckpt-615000
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-35-9a691167a69a> in <module>()
----> 1 runner.export(export_dir='/gdrive/My Drive/nmt/chat50w/models')

14 frames
/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/framework/func_graph.py in wrapper(*args, **kwargs)
966           except Exception as e:  # pylint:disable=broad-except
967             if hasattr(e, "ag_error_metadata"):
--> 968               raise e.ag_error_metadata.to_exception(e)
969             else:
970               raise

TypeError: in converted code:

/usr/local/lib/python3.6/dist-packages/opennmt/models/model.py:238 _run  *
_, predictions = self(features)
/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/keras/engine/base_layer.py:778 __call__
outputs = call_fn(cast_inputs, *args, **kwargs)
/usr/local/lib/python3.6/dist-packages/opennmt/models/language_model.py:57 call  *
ids, length = features["ids"], features["length"]

TypeError: tuple indices must be integers or slices, not str

# Fixed code
def serve_function(self):
    """Returns a function for serving this model.

    Returns:
      A ``tf.function``.
    """
    # Set name attribute of the input TensorSpec.
    input_signature = {
        name: tf.TensorSpec.from_spec(spec, name=name)
        for name, spec in self.features_inputter.input_signature().items()
    }

    @tf.function(input_signature=(input_signature,))
    def _run(features):
        features = self.features_inputter.make_features(features=features.copy())
        if isinstance(features, (list, tuple)):
            # Special case for unsupervised inputters that always return a tuple (features, labels).
            features = features[0]
        _, predictions = self(features)
        return predictions

    return _run