Repository: OpenNMT/OpenNMT-tf

Issue:
GPT2 Model export error
When I want to use onmt-main command or runner.export to export a GPT2Small model.

```python
print(tf.__version__)
print(onmt.__version__)

2.1.1
2.9.3
```

I got:

```
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
```

I think the problem is the function 'serve_function' in opennmt/models/model.py

When I use custom serve_function, it works:

```python
def serve_function(model):
  """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 model.features_inputter.input_signature().items()
  }

  @tf.function(input_signature=(input_signature,))
  def _run(features):
    # I THINK THIS IS THE PROBLEM
    # OLD
    # features = self.features_inputter.make_features(features=features.copy())
    # NEW
    features, _ = model.features_inputter.make_features(features=features.copy())
    _, predictions = model(features)
    return predictions

  return _run

# USE custom method to build model and export with custom serve_function

config = runner._finalize_config()
model = runner._init_model(config)
checkpoint = checkpoint_util.Checkpoint.from_config(config, model)
checkpoint.restore(checkpoint_path=None, weights_only=True)
tf.saved_model.save(model, '/gdrive/My Drive/nmt/chat50w/model_export', signatures=serve_function(model))
```



Discussion:

Resolving patch:
diff --git a/opennmt/models/model.py b/opennmt/models/model.py
--- a/opennmt/models/model.py
+++ b/opennmt/models/model.py
@@ -235,6 +235,9 @@ def serve_function(self):
     @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
 


Test patch:
diff --git a/opennmt/tests/model_test.py b/opennmt/tests/model_test.py
--- a/opennmt/tests/model_test.py
+++ b/opennmt/tests/model_test.py
@@ -318,6 +318,15 @@ def testLanguageModel(self, mode):
         prediction_heads=["tokens", "length"],
         params=params)
 
+  def testLanguageModelServing(self):
+    _, data_config = self._makeToyLMData()
+    decoder = decoders.SelfAttentionDecoder(
+        2, num_units=16, num_heads=4, ffn_inner_dim=32, num_sources=0)
+    model = models.LanguageModel(decoder, embedding_size=16)
+    model.initialize(data_config)
+    function = model.serve_function()
+    function.get_concrete_function()
+
   def testLanguageModelInputter(self):
     vocabulary_path = test_util.make_vocab(
         os.path.join(self.get_temp_dir(), "vocab.txt"), ["a", "b", "c"])