Repository: OpenNMT/OpenNMT-tf
Pull request #511 (2019-10-08 16:56:44)

Support RaggedTensor in tokenizers


Patch:
diff --git a/opennmt/tokenizers/tokenizer.py b/opennmt/tokenizers/tokenizer.py
index 5224d0b8..82cb2426 100644
--- a/opennmt/tokenizers/tokenizer.py
+++ b/opennmt/tokenizers/tokenizer.py
@@ -58,21 +58,29 @@ class Tokenizer(object):
     """Tokenizes text.
 
     Args:
-      text: The text to tokenize as a ``tf.Tensor`` or Python string.
+      text: A string or batch of strings to tokenize as a ``tf.Tensor`` or
+        Python values.
 
     Returns:
-      A 1-D string ``tf.Tensor`` if :obj:`text` is a ``tf.Tensor`` or a list of
-      Python unicode strings otherwise.
+      - If :obj:`text` is a Python string, a list of Python strings.
+      - If :obj:`text` is a list of Python strings, a list of list of Python
+        strings.
+      - If :obj:`text` is a 0-D ``tf.Tensor``, a 1-D ``tf.Tensor``.
+      - If :obj:`text` is a 1-D ``tf.Tensor``, a 2-D ``tf.RaggedTensor``.
 
     Raises:
-      ValueError: if the rank of :obj:`text` is greater than 0.
+      ValueError: if the rank of :obj:`text` is greater than 1.
     """
     if tf.is_tensor(text):
       rank = len(text.shape)
       if rank == 0:
         return self._tokenize_tensor(text)
+      elif rank == 1:
+        return self._tokenize_batch_tensor(text)
       else:
-        raise ValueError("Unsupported tensor rank for tokenization: {}".format(rank))
+        raise ValueError("Unsupported tensor rank %d for tokenization" % rank)
+    elif isinstance(text, list):
+      return list(map(self.tokenize, text))
     else:
       text = tf.compat.as_text(text)
       return self._tokenize_string(text)
@@ -83,18 +91,20 @@ class Tokenizer(object):
     The Tensor version supports batches of tokens.
 
     Args:
-      tokens: The tokens as a 1-D or 2-D ``tf.Tensor`` or list of Python
-        strings.
+      tokens: Tokens or batch of tokens as a ``tf.Tensor``, ``tf.RaggedTensor``,
+        or Python values.
       sequence_length: The length of each sequence. Required if :obj:`tokens`
-        is a ``tf.Tensor``.
+        is a dense 2-D ``tf.Tensor``.
 
     Returns:
-      A 0-D or 1-D string ``tf.Tensor`` if :obj:`tokens` is a ``tf.Tensor`` or a
-      Python unicode strings otherwise.
+      - If :obj:`tokens` is a list of list of Python strings, a list of Python strings.
+      - If :obj:`tokens` is a list of Python strings, a Python string.
+      - If :obj:`tokens` is a N-D ``tf.Tensor`` (or ``tf.RaggedTensor``), a
+        (N-1)-D ``tf.Tensor``.
 
     Raises:
       ValueError: if the rank of :obj:`tokens` is greater than 2.
-      ValueError: if :obj:`tokens` is a 2-D ``tf.Tensor`` and
+      ValueError: if :obj:`tokens` is a 2-D dense ``tf.Tensor`` and
         :obj:`sequence_length` is not set.
     """
     if tf.is_tensor(tokens):
@@ -106,7 +116,17 @@ class Tokenizer(object):
           raise ValueError("sequence_length is required for Tensor detokenization")
         return self._detokenize_batch_tensor(tokens, sequence_length)
       else:
-        raise ValueError("Unsupported tensor rank for detokenization: {}".format(rank))
+        raise ValueError("Unsupported tensor rank %d for detokenization" % rank)
+    elif isinstance(tokens, tf.RaggedTensor):
+      rank = len(tokens.shape)
+      if rank == 1:
+        return self._detokenize_tensor(tokens.values)
+      elif rank == 2:
+        return self._detokenize_ragged_tensor(tokens)
+      else:
+        raise ValueError("Unsupported RaggedTensor rank %d for detokenization" % rank)
+    elif isinstance(tokens, list) and tokens and isinstance(tokens[0], list):
+      return list(map(self.detokenize, tokens))
     else:
       tokens = [tf.compat.as_text(token) for token in tokens]
       return self._detokenize_string(tokens)
@@ -132,6 +152,24 @@ class Tokenizer(object):
     tokens.set_shape([None])
     return tokens
 
+  def _tokenize_batch_tensor(self, text):
+    """Tokenizes a batch of texts.
+
+    When not overriden, this default implementation calls _tokenize_tensor on
+    each tensor within the batch.
+
+    Args:
+      text: A 1-D string ``tf.Tensor``.
+
+    Returns:
+      A 2-D string ``tf.RaggedTensor``.
+    """
+    # map_fn expects each output element to have the same shape, so join tokens with
+    # spaces first and then split on spaces with a function returning a RaggedTensor.
+    tokens = tf.map_fn(
+        lambda x: tf.strings.reduce_join(self._tokenize_tensor(x), axis=0, separator=" "), text)
+    return tf.strings.split(tokens, sep=" ")
+
   @tf.autograph.experimental.do_not_convert
   def _detokenize_tensor(self, tokens):
     """Detokenizes tokens.
@@ -171,6 +209,20 @@ class Tokenizer(object):
         dtype=tf.string,
         back_prop=False)
 
+  def _detokenize_ragged_tensor(self, tokens):
+    """Detokenizes a batch of tokens as a ``tf.RaggedTensor``
+
+    When not overriden, this default implementation calls _detokenize_batch_tensor
+    on the dense representation.
+
+    Args:
+      tokens: A 2-D ``tf.RaggedTensor``.
+
+    Returns:
+      A 1-D string ``tf.Tensor``.
+    """
+    return self._detokenize_batch_tensor(tokens.to_tensor(), tokens.row_lengths())
+
   @abc.abstractmethod
   def _tokenize_string(self, text):
     """Tokenizes a Python unicode string.
@@ -202,10 +254,20 @@ class SpaceTokenizer(Tokenizer):
   """A tokenizer that splits on spaces."""
 
   def _tokenize_tensor(self, text):
-    return tf.strings.split([text]).values
+    return self._tokenize_batch_tensor(text)
+
+  def _tokenize_batch_tensor(self, text):
+    return tf.strings.split(text)
 
   def _detokenize_tensor(self, tokens):
-    return tf.strings.reduce_join(tokens, axis=0, separator=" ")
+    return self._detokenize_ragged_tensor(tokens)
+
+  def _detokenize_batch_tensor(self, tokens, sequence_length):
+    ragged = tf.RaggedTensor.from_tensor(tokens, lengths=sequence_length)
+    return self._detokenize_ragged_tensor(ragged)
+
+  def _detokenize_ragged_tensor(self, tokens):
+    return tf.strings.reduce_join(tokens, axis=tokens.shape.rank - 1, separator=" ")
 
   def _tokenize_string(self, text):
     return text.split()
@@ -218,11 +280,21 @@ class CharacterTokenizer(Tokenizer):
   """A tokenizer that splits unicode characters."""
 
   def _tokenize_tensor(self, text):
+    return self._tokenize_batch_tensor(text)
+
+  def _tokenize_batch_tensor(self, text):
     text = tf.strings.regex_replace(text, " ", "▁")
     return tf.strings.unicode_split(text, "UTF-8")
 
   def _detokenize_tensor(self, tokens):
-    text = tf.strings.reduce_join(tokens, axis=0)
+    return self._detokenize_ragged_tensor(tokens)
+
+  def _detokenize_batch_tensor(self, tokens, sequence_length):
+    _ = sequence_length
+    return self._detokenize_ragged_tensor(tokens)
+
+  def _detokenize_ragged_tensor(self, tokens):
+    text = tf.strings.reduce_join(tokens, axis=tokens.shape.rank - 1)
     return tf.strings.regex_replace(text, "▁", " ")
 
   def _tokenize_string(self, text):


Test patch:
diff --git a/opennmt/tests/tokenizer_test.py b/opennmt/tests/tokenizer_test.py
index 1c725bfe..2b9fef3e 100644
--- a/opennmt/tests/tokenizer_test.py
+++ b/opennmt/tests/tokenizer_test.py
@@ -14,16 +14,26 @@ class TokenizerTest(tf.test.TestCase):
     ref_tokens = [tf.compat.as_bytes(token) for token in ref_tokens]
     text = tf.constant(text)
     tokens = tokenizer.tokenize(text)
+    self.assertIsInstance(tokens, tf.Tensor)
     self.assertAllEqual(ref_tokens, self.evaluate(tokens))
 
+  def _testTokenizerOnBatchTensor(self, tokenizer, text, ref_tokens):
+    text = tf.constant(text)
+    tokens = tokenizer.tokenize(text)
+    self.assertIsInstance(tokens, tf.RaggedTensor)
+    self.assertAllEqual(tokens.to_list(), tf.nest.map_structure(tf.compat.as_bytes, ref_tokens))
+
   def _testTokenizerOnString(self, tokenizer, text, ref_tokens):
     ref_tokens = [tf.compat.as_text(token) for token in ref_tokens]
     tokens = tokenizer.tokenize(text)
     self.assertAllEqual(ref_tokens, tokens)
 
   def _testTokenizer(self, tokenizer, text, ref_tokens):
-    self._testTokenizerOnTensor(tokenizer, text, ref_tokens)
-    self._testTokenizerOnString(tokenizer, text, ref_tokens)
+    self.assertAllEqual(tokenizer.tokenize(text), ref_tokens)
+    self._testTokenizerOnBatchTensor(tokenizer, text, ref_tokens)
+    for txt, ref in zip(text, ref_tokens):
+      self._testTokenizerOnTensor(tokenizer, txt, ref)
+      self._testTokenizerOnString(tokenizer, txt, ref)
 
   def _testDetokenizerOnTensor(self, tokenizer, tokens, ref_text):
     ref_text = tf.compat.as_bytes(ref_text)
@@ -41,6 +51,13 @@ class TokenizerTest(tf.test.TestCase):
     text = tokenizer.detokenize(tokens, sequence_length=sequence_length)
     self.assertAllEqual(ref_text, self.evaluate(text))
 
+  def _testDetokenizerOnBatchRaggedTensor(self, tokenizer, tokens, ref_text):
+    lengths = tf.constant([len(x) for x in tokens])
+    flat_tokens = tf.constant(tf.nest.flatten(tokens))
+    ragged_tokens = tf.RaggedTensor.from_row_lengths(flat_tokens, lengths)
+    text = tokenizer.detokenize(ragged_tokens)
+    self.assertAllEqual(self.evaluate(text), tf.nest.map_structure(tf.compat.as_bytes, ref_text))
+
   def _testDetokenizerOnString(self, tokenizer, tokens, ref_text):
     tokens = [tf.compat.as_text(token) for token in tokens]
     ref_text = tf.compat.as_text(ref_text)
@@ -48,25 +65,39 @@ class TokenizerTest(tf.test.TestCase):
     self.assertAllEqual(ref_text, text)
 
   def _testDetokenizer(self, tokenizer, tokens, ref_text):
+    self.assertAllEqual(tokenizer.detokenize(tokens), ref_text)
     self._testDetokenizerOnBatchTensor(tokenizer, tokens, ref_text)
+    self._testDetokenizerOnBatchRaggedTensor(tokenizer, tokens, ref_text)
     for tok, ref in zip(tokens, ref_text):
       self._testDetokenizerOnTensor(tokenizer, tok, ref)
       self._testDetokenizerOnString(tokenizer, tok, ref)
 
   def testSpaceTokenizer(self):
-    self._testTokenizer(SpaceTokenizer(), "Hello world !", ["Hello", "world", "!"])
+    self._testTokenizer(
+        SpaceTokenizer(),
+        ["Hello world !", "How are you ?", "Good !"],
+        [["Hello", "world", "!"], ["How", "are", "you", "?"], ["Good", "!"]])
     self._testDetokenizer(
         SpaceTokenizer(),
         [["Hello", "world", "!"], ["Test"], ["My", "name"]],
         ["Hello world !", "Test", "My name"])
 
   def testCharacterTokenizer(self):
-    self._testTokenizer(CharacterTokenizer(), "a b", ["a", "▁", "b"])
-    self._testDetokenizer(CharacterTokenizer(), [["a", "▁", "b"]], ["a b"])
-    self._testTokenizer(CharacterTokenizer(), "你好，世界！", ["你", "好", "，", "世", "界", "！"])
+    self._testTokenizer(
+        CharacterTokenizer(),
+        ["a b", "cd e"],
+        [["a", "▁", "b"], ["c", "d", "▁", "e"]])
+    self._testDetokenizer(
+        CharacterTokenizer(),
+        [["a", "▁", "b"], ["c", "d", "▁", "e"]],
+        ["a b", "cd e"])
+    self._testTokenizer(CharacterTokenizer(), ["你好，世界！"], [["你", "好", "，", "世", "界", "！"]])
 
   def testOpenNMTTokenizer(self):
-    self._testTokenizer(OpenNMTTokenizer(), "Hello world!", ["Hello", "world", "!"])
+    self._testTokenizer(
+        OpenNMTTokenizer(),
+        ["Hello world!", "How are you?"],
+        [["Hello", "world", "!"], ["How", "are", "you", "?"]])
     self._testDetokenizer(
         OpenNMTTokenizer(),
         [["Hello", "world", "￭!"], ["Test"], ["My", "name"]],
@@ -74,7 +105,7 @@ class TokenizerTest(tf.test.TestCase):
 
   def testOpenNMTTokenizerArguments(self):
     tokenizer = OpenNMTTokenizer(mode="aggressive", spacer_annotate=True, spacer_new=True)
-    self._testTokenizer(tokenizer, "Hello World-s", ["Hello", "▁", "World", "-", "s"])
+    self._testTokenizer(tokenizer, ["Hello World-s"], [["Hello", "▁", "World", "-", "s"]])
 
   def testOpenNMTTokenizerAssets(self):
     asset_dir = self.get_temp_dir()