Repository: OpenNMT/OpenNMT-tf
Pull request #520 (2019-10-16 09:35:44)

Enable in-graph detokenization when supported


Patch:
diff --git a/opennmt/models/sequence_to_sequence.py b/opennmt/models/sequence_to_sequence.py
index 177bfcdb..c309ca89 100644
--- a/opennmt/models/sequence_to_sequence.py
+++ b/opennmt/models/sequence_to_sequence.py
@@ -277,13 +277,17 @@ class SequenceToSequence(model.SequenceGenerator):
           params.get("decoding_subword_token_is_spacer"))
       alignment = None  # Invalidate alignments.
 
-    predictions = {
-        "tokens": target_tokens,
-        "length": sampled_length,
-        "log_probs": log_probs
-    }
-    if alignment is not None:
-      predictions["alignment"] = alignment
+    predictions = {"log_probs": log_probs}
+    if self.labels_inputter.tokenizer.in_graph:
+      detokenized_text = self.labels_inputter.tokenizer.detokenize(
+          tf.reshape(target_tokens, [batch_size * beam_size, -1]),
+          sequence_length=tf.reshape(sampled_length, [batch_size * beam_size]))
+      predictions["text"] = tf.reshape(detokenized_text, [batch_size, beam_size])
+    else:
+      predictions["tokens"] = target_tokens
+      predictions["length"] = sampled_length
+      if alignment is not None:
+        predictions["alignment"] = alignment
 
     # Maybe restrict the number of returned hypotheses based on the user parameter.
     num_hypotheses = params.get("num_hypotheses", 1)
@@ -337,19 +341,24 @@ class SequenceToSequence(model.SequenceGenerator):
   def print_prediction(self, prediction, params=None, stream=None):
     if params is None:
       params = {}
-    num_hypotheses = len(prediction["tokens"])
+    with_scores = params.get("with_scores")
+    alignment_type = params.get("with_alignments")
+    if alignment_type and "alignment" not in prediction:
+      raise ValueError("with_alignments is set but the model did not return alignment information")
+    num_hypotheses = len(prediction["log_probs"])
     for i in range(num_hypotheses):
-      target_length = prediction["length"][i]
-      tokens = prediction["tokens"][i][:target_length]
-      sentence = self.labels_inputter.tokenizer.detokenize(tokens)
+      if "tokens" in prediction:
+        target_length = prediction["length"][i]
+        tokens = prediction["tokens"][i][:target_length]
+        sentence = self.labels_inputter.tokenizer.detokenize(tokens)
+      else:
+        sentence = prediction["text"][i]
       score = None
       attention = None
-      alignment_type = None
-      if params.get("with_scores"):
+      if with_scores:
         score = prediction["log_probs"][i]
-      if params.get("with_alignments"):
+      if alignment_type:
         attention = prediction["alignment"][i][:target_length]
-        alignment_type = params["with_alignments"]
       sentence = format_translation_output(
           sentence,
           score=score,


Test patch:
diff --git a/opennmt/tests/model_test.py b/opennmt/tests/model_test.py
index 4197b4eb..200b2499 100644
--- a/opennmt/tests/model_test.py
+++ b/opennmt/tests/model_test.py
@@ -12,6 +12,7 @@ from opennmt import encoders
 from opennmt import inputters
 from opennmt import models
 from opennmt.tests import test_util
+from opennmt.utils import misc
 
 
 def _seq2seq_model(training=None):
@@ -114,6 +115,7 @@ class ModelTest(tf.test.TestCase):
     if data_config is None:
       data_config = {}
     model.initialize(data_config, params=params)
+    # Build a dataset for mode.
     if mode == tf.estimator.ModeKeys.PREDICT:
       dataset = model.examples_inputter.make_inference_dataset(
           features_file, batch_size)
@@ -123,6 +125,7 @@ class ModelTest(tf.test.TestCase):
     elif mode == tf.estimator.ModeKeys.TRAIN:
       dataset = model.examples_inputter.make_training_dataset(
           features_file, labels_file, batch_size)
+    # Forward first batch into the model.
     data = iter(dataset).next()
     if mode != tf.estimator.ModeKeys.PREDICT:
       features, labels = data
@@ -133,16 +136,34 @@ class ModelTest(tf.test.TestCase):
     if mode != tf.estimator.ModeKeys.PREDICT:
       loss = model.compute_loss(outputs, labels, training=training)
       if mode == tf.estimator.ModeKeys.EVAL:
+        # Check that returned evaluation metrics are expected.
         eval_metrics = model.get_metrics()
         if eval_metrics is not None:
           model.update_metrics(eval_metrics, predictions, labels)
           for metric in metrics:
             self.assertIn(metric, eval_metrics)
+        try:
+          # Check that scores can be computed and printed without errors.
+          scores = model.score(features, labels)
+          first_score = tf.nest.map_structure(
+              lambda x: x.numpy(),
+              next(misc.extract_batches(scores)))
+          with open(os.devnull, "w") as devnull:
+            model.print_score(first_score, stream=devnull)
+        except NotImplementedError:
+          pass
     else:
+      # Check that all prediction heads are returned.
       self.assertIsInstance(predictions, dict)
       if prediction_heads is not None:
         for head in prediction_heads:
           self.assertIn(head, predictions)
+      # Check that the prediction can be printed without errors.
+      first_prediction = tf.nest.map_structure(
+          lambda x: x.numpy(),
+          next(misc.extract_batches(predictions)))
+      with open(os.devnull, "w") as devnull:
+        model.print_prediction(first_prediction, stream=devnull)
 
   @parameterized.expand([
       [tf.estimator.ModeKeys.TRAIN],
@@ -160,6 +181,24 @@ class ModelTest(tf.test.TestCase):
         prediction_heads=["tokens", "length", "log_probs"],
         params=params)
 
+  @parameterized.expand([
+      [tf.estimator.ModeKeys.EVAL],
+      [tf.estimator.ModeKeys.PREDICT]])
+  def testSequenceToSequenceWithInGraphTokenizer(self, mode):
+    model, params = _seq2seq_model(mode)
+    features_file, labels_file, data_config = self._makeToyEnDeData()
+    tokenization_config = {"type": "SpaceTokenizer"}
+    data_config["source_tokenization"] = tokenization_config
+    data_config["target_tokenization"] = tokenization_config
+    self._testGenericModel(
+        model,
+        mode,
+        features_file,
+        labels_file,
+        data_config,
+        prediction_heads=["text", "log_probs"],
+        params=params)
+
   @parameterized.expand([["ce"], ["mse"]])
   def testSequenceToSequenceWithGuidedAlignment(self, ga_type):
     model, params = _seq2seq_model(training=True)