Fix bug in T5 model export for eval_with_score mode -- handle case where vocabulary is tuple.

The fix is to encode inputs and targets each with their respective vocabularies.

PiperOrigin-RevId: 350812757

Author: Mesh TensorFlow Team

mesh_tensorflow/transformer/utils.py

@@ -1991,9 +1991,18 @@ def str_placeholder(name):
       targets = str_placeholder("targets")
 
       predict_batch_size = tf.shape(inputs)[0]
-      dataset = tf.data.Dataset.from_tensor_slices(
-          {"inputs": inputs, "targets": targets})
-      dataset = transformer_dataset.encode_all_features(dataset, vocabulary)
+
+      inputs_dataset = transformer_dataset.encode_all_features(
+          tf.data.Dataset.from_tensor_slices({"inputs": inputs}),
+          inputs_vocabulary(vocabulary))
+
+      targets_dataset = transformer_dataset.encode_all_features(
+          tf.data.Dataset.from_tensor_slices({"targets": targets}),
+          targets_vocabulary(vocabulary))
+
+      dataset = tf.data.Dataset.zip((inputs_dataset, targets_dataset))
+      dataset = dataset.map(lambda x, y: {**x, **y},
+                            num_parallel_calls=tf.data.experimental.AUTOTUNE)
 
       receiver_tensors = {"inputs": inputs, "targets": targets}