Repository: google-research/text-to-text-transfer-transformer

Issue:
Issues running exported SavedModel
After using the collab to export a SavedModel. I attempted to inference it locally with the following code.

```
from __future__ import absolute_import, division, print_function, unicode_literals
import tensorflow as tf
import tensorflow_text
import numpy as np
physical_devices = tf.config.experimental.list_physical_devices('GPU')
if physical_devices:
  tf.config.experimental.set_memory_growth(physical_devices[0], True)

# %%
tf.compat.v1.enable_resource_variables()
path ="PATH_TO_SAVEDMODEL"
loaded = tf.saved_model.load(path, tags='serve')
print(list(loaded.signatures.keys()))
# %%

x = tf.constant(["translate English to German: this is a test.","translate English to German: this is a test."])
print(x)
# loaded.signatures['serving_default'](x)
print(loaded.signatures['serving_default'](x))
```

and got the following error:

```
InvalidArgumentError:  Could not parse example input, value: 'translate English to German: this is a test.'
	 [[{{node ParseSingleExample/ParseSingleExample}}]]
	 [[DatasetToSingleElement]] [Op:__inference_pruned_275843]

Function call stack:
pruned
```

for inference I am using:

ubuntu: 18.04
CUDA Version: 10.2
tensorflow: 2.1

I'm sure that I'm missing something simple, but any advice on how to solve this would be greatly appreciated.


Discussion:
I believe it is expecting a serialized tf.train.Example proto, not a raw string.

You should try:
```
feature = {"input": tf.train.Feature(bytes_list=tf.train.BytesList(value=[b"test"]))}
input_str = tf.train.Example(features=tf.train.Features(feature=feature)).SerializeAsString()
```

Clearly this is not a great API. We should look into being able to pass strings directly.
I also tried this.
```
feature = {"inputs": tf.train.Feature(bytes_list=tf.train.BytesList(value=[b"test"]))}
input_str = tf.train.Example(features=tf.train.Features(feature=feature)).SerializeToString()
x = tf.constant([input_str])
```
but when running, got this error:
```
InvalidArgumentError:  Input to reshape is a tensor with 2056192 values, but the requested shape has 131596288
	 [[node encoder/einsum/parallel_4_1/einsum/Reshape (defined at <ipython-input-4-c4be6460dc7f>:3) ]] [Op:__inference_pruned_363594]

Function call stack:
pruned
```
It worked for me (thank you @adarob).

@NaxAlpha did you put multiple examples in at once? That triggered a similar error for me.

Also, is there a way to get the SavedModel to return multiple beams and the loss for each by just modifying gin parameters? I was hoping not to have to mess with the meshtensorflow to get it out but its seeming to not be an option.

Again thanks a ton for the timely response.

Resolving patch:
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -22,7 +22,7 @@
 
 setuptools.setup(
     name='t5',
-    version='0.1.8',
+    version='0.2.0',
     description='Text-to-text transfer transformer',
     long_description=_LONG_DESCRIPTION,
     long_description_content_type='text/markdown',
@@ -40,7 +40,7 @@
         'allennlp',
         'babel',
         'gin-config',
-        'mesh-tensorflow[transformer]>=0.1.9',
+        'mesh-tensorflow[transformer]>=0.1.10',
         'nltk',
         'numpy',
         'pandas',
diff --git a/t5/models/mtf_model.py b/t5/models/mtf_model.py
--- a/t5/models/mtf_model.py
+++ b/t5/models/mtf_model.py
@@ -176,11 +176,10 @@ def batch_size(self, batch_size):
     else:
       self._batch_size = batch_size
 
-  def estimator(self, vocabulary, init_checkpoint=None):
+  def estimator(self, vocabulary, init_checkpoint=None, disable_tpu=False):
     return utils.get_estimator(
         model_type=self._model_type,
-        input_vocab_size=utils.inputs_vocabulary(vocabulary).vocab_size,
-        output_vocab_size=utils.targets_vocabulary(vocabulary).vocab_size,
+        vocabulary=vocabulary,
         layout_rules=self._layout_rules,
         mesh_shape=self._mesh_shape,
         model_dir=self._model_dir,
@@ -194,7 +193,7 @@ def estimator(self, vocabulary, init_checkpoint=None):
         predict_fn=self._predict_fn,
         variable_filter=self._variable_filter,
         ensemble_inputs=self._ensemble_inputs,
-        use_tpu=self._tpu,
+        use_tpu=None if disable_tpu else self._tpu,
         tpu_job_name=self._tpu_job_name,
         iterations_per_loop=self._iterations_per_loop,
         cluster=self._cluster,
@@ -311,10 +310,10 @@ def predict(self, input_file, output_file, checkpoint_steps=-1,
       gin.bind_parameter("Bitransformer.decode.temperature", temperature)
 
     vocabulary = t5.data.SentencePieceVocabulary(sentencepiece_model_path)
-    utils.infer_model(self.estimator(vocabulary), vocabulary,
-                      self._sequence_length, self.batch_size,
-                      self._model_type, self._model_dir, checkpoint_steps,
-                      input_file, output_file)
+    utils.infer_model(
+        self.estimator(vocabulary), vocabulary, self._sequence_length,
+        self.batch_size, self._model_type, self._model_dir, checkpoint_steps,
+        input_file, output_file)
 
   def export(self, export_dir=None, checkpoint_step=-1, beam_size=1,
              temperature=1.0,
@@ -345,7 +344,8 @@ def export(self, export_dir=None, checkpoint_step=-1, beam_size=1,
     vocabulary = t5.data.SentencePieceVocabulary(sentencepiece_model_path)
     model_ckpt = "model.ckpt-" + str(checkpoint_step)
     export_dir = export_dir or self._model_dir
-    utils.export_model(self.estimator(vocabulary), export_dir, vocabulary,
-                       self._sequence_length,
-                       os.path.join(self._model_dir, model_ckpt))
+    utils.export_model(
+        self.estimator(vocabulary, disable_tpu=True), export_dir, vocabulary,
+        self._sequence_length, batch_size=self._batch_size,
+        checkpoint_path=os.path.join(self._model_dir, model_ckpt))
 


Test patch: