Repository: OpenNMT/OpenNMT-tf
Pull request #934 (2022-03-16 10:02:30)

Fix WordNoiser on inputs with rank > 2


Patch:
diff --git a/opennmt/data/noise.py b/opennmt/data/noise.py
index b354031f..8352d6fa 100644
--- a/opennmt/data/noise.py
+++ b/opennmt/data/noise.py
@@ -71,6 +71,7 @@ class WordNoiser(object):
     def _call(self, tokens, sequence_length, keep_shape, probability):
         input_rank = tokens.shape.rank
         input_is_ragged = isinstance(tokens, tf.RaggedTensor)
+        batch_shape = None
 
         if input_is_ragged:
             if keep_shape:
@@ -79,12 +80,15 @@ class WordNoiser(object):
         elif input_rank == 1:
             tokens = tf.expand_dims(tokens, 0)
             ragged_tokens = tf.RaggedTensor.from_tensor(tokens)
-        elif input_rank == 2:
+        elif input_rank >= 2:
             if sequence_length is None:
-                raise ValueError("sequence_length must be passed for 2D dense inputs")
+                raise ValueError("sequence_length must be passed for ND dense inputs")
+            if input_rank > 2:
+                input_shape = misc.shape_list(tokens)
+                batch_shape = input_shape[:-1]
+                tokens = tf.reshape(tokens, [-1, input_shape[-1]])
+                sequence_length = tf.reshape(sequence_length, [-1])
             ragged_tokens = tf.RaggedTensor.from_tensor(tokens, lengths=sequence_length)
-        else:
-            raise ValueError("unsupported rank %d for WordNoiser input" % input_rank)
 
         noisy_tokens = tf.map_fn(
             lambda tokens: self._maybe_apply_noise(tokens, probability),
@@ -104,6 +108,9 @@ class WordNoiser(object):
         if input_rank == 1:
             new_lengths = new_lengths[0]
             noisy_tokens = noisy_tokens[0]
+        elif batch_shape is not None:
+            noisy_tokens = tf.reshape(noisy_tokens, batch_shape + [-1])
+            new_lengths = tf.reshape(new_lengths, batch_shape)
 
         return noisy_tokens, new_lengths
 


Test patch:
diff --git a/opennmt/tests/model_test.py b/opennmt/tests/model_test.py
index dca162c6..56f02a53 100644
--- a/opennmt/tests/model_test.py
+++ b/opennmt/tests/model_test.py
@@ -329,6 +329,21 @@ class ModelTest(tf.test.TestCase):
         features = next(iter(dataset))
         _, predictions = model(features)
 
+    def testSequenceToSequenceWithNoisyDecoding(self):
+        model, params = _seq2seq_model()
+        params["maximum_decoding_length"] = 20
+        params["beam_width"] = 2
+        params["decoding_noise"] = [
+            {"dropout": 0.1},
+            {"replacement": [0.1, "<unk>"]},
+            {"permutation": 3},
+        ]
+        features_file, labels_file, data_config = self._makeToyEnDeData()
+        model.initialize(data_config, params=params)
+        dataset = model.examples_inputter.make_inference_dataset(features_file, 16)
+        features = next(iter(dataset))
+        _, predictions = model(features)
+
     def testSequenceToSequenceWithScheduledSampling(self):
         model = models.SequenceToSequence(
             inputters.WordEmbedder(16),
diff --git a/opennmt/tests/noise_test.py b/opennmt/tests/noise_test.py
index 686514d7..2d3f0951 100644
--- a/opennmt/tests/noise_test.py
+++ b/opennmt/tests/noise_test.py
@@ -69,6 +69,14 @@ class NoiseTest(tf.test.TestCase):
             [True, [["a￭", "b", "c￭", "d", "￭e"], ["a", "b", "c", "", ""]], [5, 3]],
             [False, [["a￭", "b", "c￭", "d", "￭e"], ["a", "b", "c", "", ""]], [5, 3]],
             [False, ["a￭", "b", "c￭", "d", "￭e"], None],
+            [
+                False,
+                [
+                    [["a￭", "b", "c￭", "d", "￭e"], ["a￭", "b", "c￭", "d", "￭e"]],
+                    [["a", "b", "c", "", ""], ["a", "b", "c", "", ""]],
+                ],
+                [[5, 5], [3, 3]],
+            ],
         ]
     )
     def testWordNoising(self, as_function, tokens, lengths):