Repository: OpenNMT/OpenNMT-tf
Pull request #513 (2019-10-09 10:54:09)

Extend tokenization configuration to in-graph tokenizers


Patch:
diff --git a/docs/configuration.md b/docs/configuration.md
index c74d2ea2..f9474994 100644
--- a/docs/configuration.md
+++ b/docs/configuration.md
@@ -41,13 +41,15 @@ data:
   source_vocabulary: data/toy-ende/src-vocab.txt
   target_vocabulary: data/toy-ende/tgt-vocab.txt
 
-  # (optional) OpenNMT tokenization configuration (or path to a configuration file).
+  # (optional) Tokenization configuration (or path to a configuration file).
   # See also: https://github.com/OpenNMT/Tokenizer/blob/master/docs/options.md
   source_tokenization:
-    mode: aggressive
-    joiner_annotate: true
-    segment_numbers: true
-    segment_alphabet_change: true
+    type: OpenNMTTokenizer
+    params:
+      mode: aggressive
+      joiner_annotate: true
+      segment_numbers: true
+      segment_alphabet_change: true
   target_tokenization: config/tokenization/aggressive.yml
 
   # (optional) Pretrained embedding configuration.
diff --git a/docs/serving.md b/docs/serving.md
index 3f4b23b4..329105ea 100644
--- a/docs/serving.md
+++ b/docs/serving.md
@@ -36,9 +36,9 @@ saved_model_cli show --dir ~/my-models/ende \
 
 Some examples using exported models are available in the [`examples/serving`](https://github.com/OpenNMT/OpenNMT-tf/tree/master/examples/serving) directory.
 
-## Input preprocessing
+## Input preprocessing and tokenization
 
-TensorFlow Serving only runs TensorFlow operations. Preprocessing functions such as the tokenization is usually not implemented in terms of TensorFlow operators. These functions should be run outside of the TensorFlow engine, either by the client or a proxy server.
+TensorFlow Serving only runs TensorFlow operations. Preprocessing functions such as the tokenization is sometimes not implemented in terms of TensorFlow ops (see the *Tokenization* page for more details). In this case, these functions should be run outside of the TensorFlow engine, either by the client or a proxy server.
 
 * The OpenNMT-tf [serving example](https://github.com/OpenNMT/OpenNMT-tf/tree/master/examples/serving) uses the client approach to implement a simple interactive translation loop
 * The project [nmt-wizard-docker](https://github.com/OpenNMT/nmt-wizard-docker) uses the proxy server approach to wrap a TensorFlow Serving instance with a custom processing layer and REST API. Exported OpenNMT-tf models can integrated with this tool by following these [instructions](https://github.com/OpenNMT/nmt-wizard-docker/issues/46#issuecomment-456795844).
diff --git a/docs/tokenization.md b/docs/tokenization.md
index 6a7aa69e..c4db3559 100644
--- a/docs/tokenization.md
+++ b/docs/tokenization.md
@@ -14,14 +14,21 @@ However, OpenNMT-tf provides tokenization tools based on the C++ OpenNMT [Tokeni
 YAML files are used to set the tokenizer options to ensure consistency during data preparation and training. For example, this configuration defines a simple word-based tokenization using the OpenNMT tokenizer:
 
 ```yaml
-mode: aggressive
-joiner_annotate: true
-segment_numbers: true
-segment_alphabet_change: true
+type: OpenNMTTokenizer
+params:
+  mode: aggressive
+  joiner_annotate: true
+  segment_numbers: true
+  segment_alphabet_change: true
 ```
 
 *For a complete list of available options, see the <a href="https://github.com/OpenNMT/Tokenizer/blob/master/docs/options.md">Tokenizer documentation</a>).*
 
+OpenNMT-tf also defines additional tokenizers:
+
+* `CharacterTokenizer`
+* `SpaceTokenizer`
+
 ## Offline usage
 
 You can invoke the `onmt-tokenize-text` script directly and pass the tokenizer configuration:
@@ -54,6 +61,21 @@ data:
   target_tokenization: config/tokenization/aggressive.yml
 ```
 
-## Notes
+## Exported graph
+
+Only TensorFlow ops can be exported to graphs and used for serving. When a tokenizer is not implemented in terms of TensorFlow ops such as the OpenNMT tokenizer, it will not be part of the exported graph. The model will then expects tokenized inputs during serving.
+
+**In-graph tokenizers:**
+
+* `CharacterTokenizer`
+* `SpaceTokenizer`
+
+Model inputs: `text` (1D string tensor)
+
+**Out-of-graph tokenizers:**
+
+* `OpenNMTTokenizer` (\*)
+
+Model inputs: `tokens` (2D string tensor), `length` (1D int32 tensor)
 
-* As of now, tokenizers are not part of the exported graph. However, all tokenization resources (configuration, subword models, etc.) are registered as additional assets in the `SavedModel` bundle in the `assets.extra` directory.
+(\*) During model export, tokenization resources used by the OpenNMT tokenizer (configuration, subword models, etc.) are registered as additional assets in the `SavedModel`'s `assets.extra` directory.
diff --git a/opennmt/inputters/text_inputter.py b/opennmt/inputters/text_inputter.py
index bee7a665..1b38d6af 100644
--- a/opennmt/inputters/text_inputter.py
+++ b/opennmt/inputters/text_inputter.py
@@ -238,8 +238,14 @@ class TextInputter(Inputter):
       features = {}
     if "tokens" in features:
       return features
+    if "text" in features:
+      element = features.pop("text")
     tokens = self.tokenizer.tokenize(element)
-    length = tf.shape(tokens)[0]
+    if isinstance(tokens, tf.RaggedTensor):
+      length = tokens.row_lengths()
+      tokens = tokens.to_tensor()
+    else:
+      length = tf.shape(tokens)[0]
     if training and self.noiser is not None:
       noisy_tokens, noisy_length = self.noiser(tokens, keep_shape=False)
       if self.in_place_noise:
@@ -258,10 +264,15 @@ class TextInputter(Inputter):
     return features
 
   def input_signature(self):
-    return {
-        "tokens": tf.TensorSpec([None, None], tf.string),
-        "length": tf.TensorSpec([None], tf.int32)
-    }
+    if self.tokenizer.in_graph:
+      return {
+          "text": tf.TensorSpec([None], tf.string)
+      }
+    else:
+      return {
+          "tokens": tf.TensorSpec([None, None], tf.string),
+          "length": tf.TensorSpec([None], tf.int32)
+      }
 
 
 class WordEmbedder(TextInputter):
diff --git a/opennmt/tokenizers/__init__.py b/opennmt/tokenizers/__init__.py
index 24a729ab..fbfa74f4 100644
--- a/opennmt/tokenizers/__init__.py
+++ b/opennmt/tokenizers/__init__.py
@@ -20,7 +20,7 @@ from opennmt.tokenizers.tokenizer import SpaceTokenizer
 from opennmt.tokenizers.tokenizer import CharacterTokenizer
 
 
-def make_tokenizer(config):
+def make_tokenizer(config=None):
   """Creates a tokenizer instance from the configuration.
 
   Args:
@@ -28,12 +28,28 @@ def make_tokenizer(config):
 
   Returns:
     A :class:`opennmt.tokenizers.Tokenizer` instance.
+
+  Raises:
+    ValueError: if :obj:`config` is invalid.
   """
   if config:
     if isinstance(config, six.string_types) and tf.io.gfile.exists(config):
       with tf.io.gfile.GFile(config, mode="rb") as config_file:
         config = yaml.load(config_file, Loader=yaml.UnsafeLoader)
-    tokenizer = OpenNMTTokenizer(**config)
+    if isinstance(config, dict):
+      tokenizer_type = config.get("type")
+      tokenizer_params = config.get("params", {})
+      if tokenizer_type is None:
+        tokenizer = OpenNMTTokenizer(**config)
+      else:
+        tokenizer_class = getattr(sys.modules[__name__], tokenizer_type, None)
+        if tokenizer_class is None:
+          raise ValueError("Invalid tokenizer type: %s" % tokenizer_type)
+        tokenizer = tokenizer_class(**tokenizer_params)
+    else:
+      raise ValueError("Invalid tokenization configuration: %s" % str(config))
   else:
-    tokenizer = SpaceTokenizer()
+    # If the tokenization was not configured, we assume that an external tokenization
+    # was used and we don't include the tokenizer in the exprted graph.
+    tokenizer = SpaceTokenizer(in_graph=False)
   return tokenizer
diff --git a/opennmt/tokenizers/tokenizer.py b/opennmt/tokenizers/tokenizer.py
index 82cb2426..f256acf5 100644
--- a/opennmt/tokenizers/tokenizer.py
+++ b/opennmt/tokenizers/tokenizer.py
@@ -15,6 +15,11 @@ from opennmt.utils.misc import print_bytes
 class Tokenizer(object):
   """Base class for tokenizers."""
 
+  @property
+  def in_graph(self):
+    """Returns ``True`` if this tokenizer can be run in graph (i.e. uses TensorFlow ops)."""
+    return False
+
   def export_assets(self, asset_dir, asset_prefix=""):  # pylint: disable=unused-argument
     """Exports assets for this tokenizer.
 
@@ -253,6 +258,18 @@ class Tokenizer(object):
 class SpaceTokenizer(Tokenizer):
   """A tokenizer that splits on spaces."""
 
+  def __init__(self, in_graph=True):
+    """Initializes the tokenizer.
+
+    Args:
+      in_graph: If ``True``, this tokenizer should be integrated in the exported graph.
+    """
+    self._in_graph = in_graph
+
+  @property
+  def in_graph(self):
+    return self._in_graph
+
   def _tokenize_tensor(self, text):
     return self._tokenize_batch_tensor(text)
 
@@ -279,6 +296,10 @@ class SpaceTokenizer(Tokenizer):
 class CharacterTokenizer(Tokenizer):
   """A tokenizer that splits unicode characters."""
 
+  @property
+  def in_graph(self):
+    return True
+
   def _tokenize_tensor(self, text):
     return self._tokenize_batch_tensor(text)
 


Test patch:
diff --git a/opennmt/tests/inputter_test.py b/opennmt/tests/inputter_test.py
index 50187730..5654db1e 100644
--- a/opennmt/tests/inputter_test.py
+++ b/opennmt/tests/inputter_test.py
@@ -13,6 +13,7 @@ from tensorboard.plugins import projector
 from google.protobuf import text_format
 from parameterized import parameterized
 
+from opennmt import inputters
 from opennmt import tokenizers
 from opennmt.data import dataset as dataset_util
 from opennmt.data import noise
@@ -135,6 +136,14 @@ class InputterTest(tf.test.TestCase):
       self.assertIn(name, features)
       self.assertTrue(features[name].shape.is_compatible_with(expected_shape))
 
+  def _testServing(self, inputter):
+    @tf.function(input_signature=(inputter.input_signature(),))
+    def _serving_fun(features):
+      features = inputter.make_features(features=features.copy())
+      inputs = inputter(features)
+      return inputs
+    _serving_fun.get_concrete_function()
+
   def _makeDataset(self, inputter, data_file, data_config=None, dataset_size=1, shapes=None):
     if data_config is not None:
       inputter.initialize(data_config)
@@ -145,6 +154,8 @@ class InputterTest(tf.test.TestCase):
     if shapes is not None:
       self._checkFeatures(features, shapes)
     inputs = inputter(features, training=True)
+    if not isinstance(inputter, inputters.ExampleInputter):
+      self._testServing(inputter)
     return self.evaluate((features, inputs))
 
   def testWordEmbedder(self):
@@ -188,6 +199,17 @@ class InputterTest(tf.test.TestCase):
     self.assertAllEqual([4], features["length"])
     self.assertAllEqual([[2, 1, 3, 4]], features["ids"])
 
+  def testWordEmbedderWithInGraphTokenizer(self):
+    vocab_file = self._makeTextFile("vocab.txt", ["the", "world", "hello", "￭"])
+    embedder = text_inputter.WordEmbedder(embedding_size=10)
+    data_config = {
+        "vocabulary": vocab_file,
+        "tokenization": {"type": "CharacterTokenizer"}
+    }
+    embedder.initialize(data_config)
+    self.assertIn("text", embedder.input_signature())
+    self._testServing(embedder)
+
   def testWordEmbedderWithCompression(self):
     vocab_file = self._makeTextFile("vocab.txt", ["the", "world", "hello", "￭"])
     data_file = self._makeTextFile("data.txt", ["hello world !", "how are you ?"], compress=True)
diff --git a/opennmt/tests/tokenizer_test.py b/opennmt/tests/tokenizer_test.py
index 2b9fef3e..f39bad6f 100644
--- a/opennmt/tests/tokenizer_test.py
+++ b/opennmt/tests/tokenizer_test.py
@@ -5,7 +5,7 @@ import yaml
 
 import tensorflow as tf
 
-from opennmt.tokenizers import SpaceTokenizer, CharacterTokenizer, OpenNMTTokenizer
+from opennmt import tokenizers
 
 
 class TokenizerTest(tf.test.TestCase):
@@ -74,37 +74,41 @@ class TokenizerTest(tf.test.TestCase):
 
   def testSpaceTokenizer(self):
     self._testTokenizer(
-        SpaceTokenizer(),
+        tokenizers.SpaceTokenizer(),
         ["Hello world !", "How are you ?", "Good !"],
         [["Hello", "world", "!"], ["How", "are", "you", "?"], ["Good", "!"]])
     self._testDetokenizer(
-        SpaceTokenizer(),
+        tokenizers.SpaceTokenizer(),
         [["Hello", "world", "!"], ["Test"], ["My", "name"]],
         ["Hello world !", "Test", "My name"])
 
   def testCharacterTokenizer(self):
     self._testTokenizer(
-        CharacterTokenizer(),
+        tokenizers.CharacterTokenizer(),
         ["a b", "cd e"],
         [["a", "▁", "b"], ["c", "d", "▁", "e"]])
     self._testDetokenizer(
-        CharacterTokenizer(),
+        tokenizers.CharacterTokenizer(),
         [["a", "▁", "b"], ["c", "d", "▁", "e"]],
         ["a b", "cd e"])
-    self._testTokenizer(CharacterTokenizer(), ["你好，世界！"], [["你", "好", "，", "世", "界", "！"]])
+    self._testTokenizer(
+        tokenizers.CharacterTokenizer(),
+        ["你好，世界！"],
+        [["你", "好", "，", "世", "界", "！"]])
 
   def testOpenNMTTokenizer(self):
     self._testTokenizer(
-        OpenNMTTokenizer(),
+        tokenizers.OpenNMTTokenizer(),
         ["Hello world!", "How are you?"],
         [["Hello", "world", "!"], ["How", "are", "you", "?"]])
     self._testDetokenizer(
-        OpenNMTTokenizer(),
+        tokenizers.OpenNMTTokenizer(),
         [["Hello", "world", "￭!"], ["Test"], ["My", "name"]],
         ["Hello world!", "Test", "My name"])
 
   def testOpenNMTTokenizerArguments(self):
-    tokenizer = OpenNMTTokenizer(mode="aggressive", spacer_annotate=True, spacer_new=True)
+    tokenizer = tokenizers.OpenNMTTokenizer(
+        mode="aggressive", spacer_annotate=True, spacer_new=True)
     self._testTokenizer(tokenizer, ["Hello World-s"], [["Hello", "▁", "World", "-", "s"]])
 
   def testOpenNMTTokenizerAssets(self):
@@ -114,7 +118,7 @@ class TokenizerTest(tf.test.TestCase):
     with open(bpe_model_path, "wb") as bpe_model_file:
       bpe_model_file.write(b"#version: 0.2\ne s</w>\n")
 
-    tokenizer = OpenNMTTokenizer(mode="conservative", bpe_model_path=bpe_model_path)
+    tokenizer = tokenizers.OpenNMTTokenizer(mode="conservative", bpe_model_path=bpe_model_path)
 
     # Generated assets are prefixed but not existing resources.
     assets = tokenizer.export_assets(asset_dir, asset_prefix="source_")
@@ -128,6 +132,28 @@ class TokenizerTest(tf.test.TestCase):
       asset_config = yaml.load(config_file.read(), Loader=yaml.UnsafeLoader)
     self.assertDictEqual(asset_config, {"mode": "conservative", "bpe_model_path": "model.bpe"})
 
+  def testMakeTokenizer(self):
+    tokenizer = tokenizers.make_tokenizer()
+    self.assertIsInstance(tokenizer, tokenizers.SpaceTokenizer)
+    self.assertFalse(tokenizer.in_graph)
+    tokenizer = tokenizers.make_tokenizer({"type": "SpaceTokenizer"})
+    self.assertIsInstance(tokenizer, tokenizers.SpaceTokenizer)
+    self.assertTrue(tokenizer.in_graph)
+    self.assertIsInstance(
+        tokenizers.make_tokenizer({"mode": "conservative"}),
+        tokenizers.OpenNMTTokenizer)
+    self.assertIsInstance(
+        tokenizers.make_tokenizer({"type": "OpenNMTTokenizer", "params": {"mode": "conservative"}}),
+        tokenizers.OpenNMTTokenizer)
+    config_path = os.path.join(self.get_temp_dir(), "tok_config.yml")
+    with open(config_path, "w") as config_file:
+      yaml.dump({"mode": "conservative"}, config_file)
+    self.assertIsInstance(
+        tokenizers.make_tokenizer(config_path),
+        tokenizers.OpenNMTTokenizer)
+    with self.assertRaises(ValueError):
+      tokenizers.make_tokenizer({"type": "UnknownTokenizer"})
+
 
 if __name__ == "__main__":
   tf.test.main()