Repository: keras-team/keras-nlp
Pull request #339 (2022-09-03 16:55:28)

Fixes for Random Deletion Layer


Patch:
diff --git a/keras_nlp/layers/random_deletion.py b/keras_nlp/layers/random_deletion.py
index 11314721..641d6910 100644
--- a/keras_nlp/layers/random_deletion.py
+++ b/keras_nlp/layers/random_deletion.py
@@ -25,8 +25,9 @@ class RandomDeletion(keras.layers.Layer):
     augmentation as described in the paper [EDA: Easy Data Augmentation
     Techniques for Boosting Performance on Text Classification Tasks]
     (https://arxiv.org/pdf/1901.11196.pdf). The layer expects the inputs to be
-    pretokenized so that each token can be individually treated as a possible
-    candidate for deletion.
+    pre-split into token level inputs. This allows control over the level of
+    augmentation, you can split by character for character level swaps, or by
+    word for word level swaps.
 
     Input should be either a `tf.RaggedTensor` or a dense `tf.Tensor`, and
     either rank-1 or rank-2.
@@ -99,8 +100,8 @@ class RandomDeletion(keras.layers.Layer):
     ...     skip_py_fn=skip_py_fn, seed=42)
     >>> augmented=augmenter(inputs)
     >>> tf.strings.reduce_join(augmented, separator=" ", axis=-1)
-    <tf.Tensor: shape=(2,), dtype=string, numpy=array([b'Hey I', b'and'],
-    dtype=object)>
+    <tf.Tensor: shape=(2,), dtype=string,
+    numpy=array([b'Hey I', b'and Tensorflow'], dtype=object)>
     """
 
     def __init__(
@@ -133,6 +134,11 @@ class RandomDeletion(keras.layers.Layer):
         self.skip_list = skip_list
         self.skip_fn = skip_fn
         self.skip_py_fn = skip_py_fn
+        if self.max_deletions is not None and self.max_deletions < 0:
+            raise ValueError(
+                "max_deletions must be non-negative."
+                f"Received max_deletions={max_deletions}."
+            )
 
         if self.rate > 1 or self.rate < 0:
             raise ValueError(
@@ -206,7 +212,7 @@ class RandomDeletion(keras.layers.Layer):
             )
 
         # Figure out how many we are going to select.
-        token_counts = tf.cast(inputs.row_lengths(), "float32")
+        token_counts = tf.cast(positions.row_lengths(), "float32")
         num_to_select = tf.random.stateless_binomial(
             shape=tf.shape(token_counts),
             seed=self._generator.make_seeds()[:, 0],


Test patch:
diff --git a/keras_nlp/layers/random_deletion_test.py b/keras_nlp/layers/random_deletion_test.py
index b345400a..1582c5f4 100644
--- a/keras_nlp/layers/random_deletion_test.py
+++ b/keras_nlp/layers/random_deletion_test.py
@@ -16,8 +16,7 @@
 import tensorflow as tf
 from tensorflow import keras
 
-from keras_nlp.layers import random_deletion
-from keras_nlp.tokenizers import UnicodeCodepointTokenizer
+from keras_nlp.layers import RandomDeletion
 
 
 class RandomDeletionTest(tf.test.TestCase):
@@ -25,48 +24,36 @@ class RandomDeletionTest(tf.test.TestCase):
         keras.utils.set_random_seed(1337)
         inputs = ["Hey I like", "Keras and Tensorflow"]
         split = tf.strings.split(inputs)
-        augmenter = random_deletion.RandomDeletion(
-            rate=0.4, max_deletions=1, seed=42
-        )
+        augmenter = RandomDeletion(rate=0.4, max_deletions=1, seed=42)
         augmented = augmenter(split)
         output = tf.strings.reduce_join(augmented, separator=" ", axis=-1)
         self.assertAllEqual(output.shape, tf.convert_to_tensor(inputs).shape)
         exp_output = [b"I like", b"and Tensorflow"]
-        for i in range(output.shape[0]):
-            self.assertAllEqual(output[i], exp_output[i])
+        self.assertAllEqual(output, exp_output)
 
     def test_shape_and_output_from_character_swaps(self):
         keras.utils.set_random_seed(1337)
         inputs = ["Hey I like", "Keras and Tensorflow"]
         split = tf.strings.unicode_split(inputs, "UTF-8")
-        augmenter = random_deletion.RandomDeletion(
-            rate=0.4, max_deletions=1, seed=42
-        )
+        augmenter = RandomDeletion(rate=0.4, max_deletions=1, seed=42)
         augmented = augmenter(split)
         output = tf.strings.reduce_join(augmented, axis=-1)
         self.assertAllEqual(output.shape, tf.convert_to_tensor(inputs).shape)
         exp_output = [b"Hey I lie", b"Keras and Tensoflow"]
-        for i in range(output.shape[0]):
-            self.assertAllEqual(output[i], exp_output[i])
+        self.assertAllEqual(output, exp_output)
 
     def test_with_integer_tokens(self):
         keras.utils.set_random_seed(1337)
-        inputs = ["Hey I like", "Keras and Tensorflow"]
-        tokenizer = UnicodeCodepointTokenizer(lowercase=False)
-        tokenized = tokenizer.tokenize(inputs)
-        augmenter = random_deletion.RandomDeletion(
-            rate=0.4, max_deletions=4, seed=42
-        )
-        augmented = augmenter(tokenized)
-        output = tokenizer.detokenize(augmented)
-        self.assertAllEqual(output.shape, tf.convert_to_tensor(inputs).shape)
-        exp_output = [b"Hey Ile", b"Keas and ensoflw"]
-        for i in range(output.shape[0]):
-            self.assertAllEqual(output[i], exp_output[i])
+        inputs = tf.constant([[1, 2], [3, 4]])
+        augmenter = RandomDeletion(rate=0.4, max_deletions=4, seed=42)
+        output = augmenter(inputs)
+        self.assertAllEqual(output.to_tensor().shape[0], inputs.shape[0])
+        exp_output = [[2], [4]]
+        self.assertAllEqual(output, exp_output)
 
     def test_skip_options(self):
         keras.utils.set_random_seed(1337)
-        augmenter = random_deletion.RandomDeletion(
+        augmenter = RandomDeletion(
             rate=0.4, max_deletions=1, seed=42, skip_list=["Tensorflow", "like"]
         )
         inputs = ["Hey I like", "Keras and Tensorflow"]
@@ -75,43 +62,37 @@ class RandomDeletionTest(tf.test.TestCase):
         output = tf.strings.reduce_join(augmented, separator=" ", axis=-1)
         self.assertAllEqual(output.shape, tf.convert_to_tensor(inputs).shape)
         exp_output = [b"I like", b"and Tensorflow"]
-        for i in range(output.shape[0]):
-            self.assertAllEqual(output[i], exp_output[i])
+        self.assertAllEqual(output, exp_output)
 
         def skip_fn(word):
             if word == "Tensorflow" or word == "like":
                 return True
             return False
 
-        augmenter = random_deletion.RandomDeletion(
+        augmenter = RandomDeletion(
             rate=0.4, max_deletions=1, seed=42, skip_fn=skip_fn
         )
         augmented = augmenter(split)
         output = tf.strings.reduce_join(augmented, separator=" ", axis=-1)
         self.assertAllEqual(output.shape, tf.convert_to_tensor(inputs).shape)
         exp_output = [b"Hey like", b"Keras Tensorflow"]
-        for i in range(output.shape[0]):
-            self.assertAllEqual(output[i], exp_output[i])
+        self.assertAllEqual(output, exp_output)
 
         def skip_py_fn(word):
             if word == "Tensorflow" or word == "like":
                 return True
             return False
 
-        augmenter = random_deletion.RandomDeletion(
+        augmenter = RandomDeletion(
             rate=0.4, max_deletions=1, seed=42, skip_py_fn=skip_py_fn
         )
         augmented = augmenter(split)
         output = tf.strings.reduce_join(augmented, separator=" ", axis=-1)
         self.assertAllEqual(output.shape, tf.convert_to_tensor(inputs).shape)
         exp_output = [b"Hey like", b"Keras Tensorflow"]
-        for i in range(output.shape[0]):
-            self.assertAllEqual(output[i], exp_output[i])
 
     def test_get_config_and_from_config(self):
-        augmenter = random_deletion.RandomDeletion(
-            rate=0.4, max_deletions=1, seed=42
-        )
+        augmenter = RandomDeletion(rate=0.4, max_deletions=1, seed=42)
 
         expected_config_subset = {"max_deletions": 1, "rate": 0.4, "seed": 42}
 
@@ -119,7 +100,7 @@ class RandomDeletionTest(tf.test.TestCase):
 
         self.assertEqual(config, {**config, **expected_config_subset})
 
-        restored_augmenter = random_deletion.RandomDeletion.from_config(
+        restored_augmenter = RandomDeletion.from_config(
             config,
         )
 
@@ -130,9 +111,7 @@ class RandomDeletionTest(tf.test.TestCase):
 
     def test_augment_first_batch_second(self):
         keras.utils.set_random_seed(1337)
-        augmenter = random_deletion.RandomDeletion(
-            rate=0.4, max_deletions=1, seed=42
-        )
+        augmenter = RandomDeletion(rate=0.4, max_deletions=1, seed=42)
         inputs = ["Hey I like", "Keras and Tensorflow"]
         split = tf.strings.split(inputs)
         ds = tf.data.Dataset.from_tensor_slices(split)
@@ -141,8 +120,7 @@ class RandomDeletionTest(tf.test.TestCase):
         output = ds.take(1).get_single_element()
 
         exp_output = [[b"I", b"like"], [b"Keras", b"and", b"Tensorflow"]]
-        for i in range(output.shape[0]):
-            self.assertAllEqual(output[i], exp_output[i])
+        self.assertAllEqual(output, exp_output)
 
         def skip_fn(word):
             return tf.strings.regex_full_match(word, r"\pP")
@@ -150,7 +128,7 @@ class RandomDeletionTest(tf.test.TestCase):
         def skip_py_fn(word):
             return len(word) < 4
 
-        augmenter = random_deletion.RandomDeletion(
+        augmenter = RandomDeletion(
             rate=0.8, max_deletions=1, seed=42, skip_fn=skip_fn
         )
         ds = tf.data.Dataset.from_tensor_slices(split)
@@ -158,25 +136,21 @@ class RandomDeletionTest(tf.test.TestCase):
         ds = ds.apply(tf.data.experimental.dense_to_ragged_batch(2))
         output = ds.take(1).get_single_element()
         exp_output = [[b"I", b"like"], [b"and", b"Tensorflow"]]
-        for i in range(output.shape[0]):
-            self.assertAllEqual(output[i], exp_output[i])
+        self.assertAllEqual(output, exp_output)
 
-        augmenter = random_deletion.RandomDeletion(
+        augmenter = RandomDeletion(
             rate=0.8, max_deletions=1, seed=42, skip_py_fn=skip_py_fn
         )
         ds = tf.data.Dataset.from_tensor_slices(split)
         ds = ds.map(augmenter)
         ds = ds.apply(tf.data.experimental.dense_to_ragged_batch(2))
         output = ds.take(1).get_single_element()
-        exp_output = [[b"Hey", b"I"], [b"and", b"Tensorflow"]]
-        for i in range(output.shape[0]):
-            self.assertAllEqual(output[i], exp_output[i])
+        exp_output = [[b"Hey", b"I", b"like"], [b"and", b"Tensorflow"]]
+        self.assertAllEqual(output, exp_output)
 
     def test_batch_first_augment_second(self):
         keras.utils.set_random_seed(1337)
-        augmenter = random_deletion.RandomDeletion(
-            rate=0.4, max_deletions=1, seed=42
-        )
+        augmenter = RandomDeletion(rate=0.4, max_deletions=1, seed=42)
         inputs = ["Hey I like", "Keras and Tensorflow"]
         split = tf.strings.split(inputs)
         ds = tf.data.Dataset.from_tensor_slices(split)
@@ -184,8 +158,7 @@ class RandomDeletionTest(tf.test.TestCase):
         output = ds.take(1).get_single_element()
 
         exp_output = [[b"I", b"like"], [b"and", b"Tensorflow"]]
-        for i in range(output.shape[0]):
-            self.assertAllEqual(output[i], exp_output[i])
+        self.assertAllEqual(output, exp_output)
 
         def skip_fn(word):
             return tf.strings.regex_full_match(word, r"\pP")
@@ -193,32 +166,28 @@ class RandomDeletionTest(tf.test.TestCase):
         def skip_py_fn(word):
             return len(word) < 4
 
-        augmenter = random_deletion.RandomDeletion(
+        augmenter = RandomDeletion(
             rate=0.8, max_deletions=1, seed=42, skip_fn=skip_fn
         )
         ds = tf.data.Dataset.from_tensor_slices(split)
         ds = ds.batch(5).map(augmenter)
         output = ds.take(1).get_single_element()
         exp_output = [[b"I", b"like"], [b"and", b"Tensorflow"]]
-        for i in range(output.shape[0]):
-            self.assertAllEqual(output[i], exp_output[i])
+        self.assertAllEqual(output, exp_output)
 
-        augmenter = random_deletion.RandomDeletion(
+        augmenter = RandomDeletion(
             rate=0.8, max_deletions=1, seed=42, skip_py_fn=skip_py_fn
         )
         ds = tf.data.Dataset.from_tensor_slices(split)
         ds = ds.batch(5).map(augmenter)
         output = ds.take(1).get_single_element()
-        exp_output = [[b"Hey", b"I"], [b"and", b"Tensorflow"]]
-        for i in range(output.shape[0]):
-            self.assertAllEqual(output[i], exp_output[i])
+        exp_output = [[b"Hey", b"I", b"like"], [b"and", b"Tensorflow"]]
+        self.assertAllEqual(output, exp_output)
 
     def test_functional_model(self):
         keras.utils.set_random_seed(1337)
         input_data = tf.constant(["Hey I like", "Keras and Tensorflow"])
-        augmenter = random_deletion.RandomDeletion(
-            rate=0.4, max_deletions=1, seed=42
-        )
+        augmenter = RandomDeletion(rate=0.4, max_deletions=1, seed=42)
         inputs = tf.keras.Input(dtype="string", shape=())
         outputs = augmenter(tf.strings.split(inputs))
         model = tf.keras.Model(inputs, outputs)