Repository: keras-team/keras-nlp
Pull request #240 (2022-06-28 23:16:21)

Add support for StartEndPacker packing 2D tensor
StartEndPacker should support rank-2 dense tensors.

mattdangerw: also please fix failing tests!

Patch:
diff --git a/keras_nlp/__init__.py b/keras_nlp/__init__.py
index 7868a8a9..506d071c 100644
--- a/keras_nlp/__init__.py
+++ b/keras_nlp/__init__.py
@@ -17,4 +17,4 @@ from keras_nlp import metrics
 from keras_nlp import tokenizers
 from keras_nlp import utils
 
-__version__ = "0.2.0"
+__version__ = "0.3.0"
diff --git a/keras_nlp/layers/start_end_packer.py b/keras_nlp/layers/start_end_packer.py
index 704f4892..ef8ba01f 100644
--- a/keras_nlp/layers/start_end_packer.py
+++ b/keras_nlp/layers/start_end_packer.py
@@ -26,11 +26,8 @@ class StartEndPacker(keras.layers.Layer):
     be called after tokenization. The layer will first trim inputs to fit, then
     add start/end tokens, and finally pad, if necessary, to `sequence_length`.
 
-    If input is batched, input should be a `tf.RaggedTensor` with shape
-    `[batch_size, None]` and will be packed and converted to a dense tensor with
-    shape `[batch_size, sequence_length]`.
-    If input is unbatched, input should be a dense rank-1 tensor of any shape,
-    and will be packed to shape `[sequence_length]`.
+    Input should be either a `tf.RaggedTensor` or a dense `tf.Tensor`, and
+    either rank-1 or rank-2.
 
     Args:
         sequence_length: int. The desired output length.
@@ -108,26 +105,19 @@ class StartEndPacker(keras.layers.Layer):
         if not isinstance(inputs, (tf.Tensor, tf.RaggedTensor)):
             inputs = tf.convert_to_tensor(inputs)
 
-        input_is_dense = isinstance(inputs, tf.Tensor)
-        input_is_ragged = isinstance(inputs, tf.RaggedTensor)
-
-        if input_is_dense:
-            if inputs.shape.rank != 1:
-                raise ValueError(
-                    "Input must either be dense with rank 1, or ragged with "
-                    "rank 2. Received dense input with "
-                    f"rank={inputs.shape.rank}"
-                )
-
-            # Add a new axis at the beginning and convert to ragged tensor.
-            inputs = tf.RaggedTensor.from_tensor(tf.expand_dims(inputs, axis=0))
-        elif input_is_ragged:
-            if inputs.shape.rank != 2:
-                raise ValueError(
-                    "Input must either be dense with rank 1, or ragged with "
-                    "rank 2. Received ragged input with "
-                    f"rank={inputs.shape.rank}"
-                )
+        input_is_1d = False
+        if inputs.shape.rank < 1 or inputs.shape.rank > 2:
+            raise ValueError(
+                "Input must either be rank 1 or rank 2. Received input with "
+                f"rank={inputs.shape.rank}"
+            )
+        elif inputs.shape.rank == 1:
+            input_is_1d = True
+            # Add a new axis at the beginning.
+            inputs = tf.expand_dims(inputs, axis=0)
+        if isinstance(inputs, tf.Tensor):
+            # Convert to ragged tensor.
+            inputs = tf.RaggedTensor.from_tensor(inputs)
 
         batch_size = tf.shape(inputs)[0]
 
@@ -148,7 +138,7 @@ class StartEndPacker(keras.layers.Layer):
             shape=(batch_size, self.sequence_length),
         )
 
-        if input_is_dense:
+        if input_is_1d:
             inputs = tf.squeeze(inputs, axis=0)
 
         return inputs
diff --git a/setup.py b/setup.py
index 9759ee89..9a2da4b4 100644
--- a/setup.py
+++ b/setup.py
@@ -29,7 +29,7 @@ setup(
     ),
     long_description=README,
     long_description_content_type="text/markdown",
-    version="0.2.0",
+    version="0.3.0",
     url="https://github.com/keras-team/keras-nlp",
     author="Keras team",
     author_email="keras-nlp@google.com",


Test patch:
diff --git a/keras_nlp/layers/start_end_packer_test.py b/keras_nlp/layers/start_end_packer_test.py
index b36bc3b1..efc29747 100644
--- a/keras_nlp/layers/start_end_packer_test.py
+++ b/keras_nlp/layers/start_end_packer_test.py
@@ -29,15 +29,12 @@ class StartEndPackerTest(tf.test.TestCase):
         expected_output = [5, 6, 7, 0, 0]
         self.assertAllEqual(output, expected_output)
 
-    def test_dense_input_error(self):
+    def test_dense_2D_input(self):
         input_data = tf.constant([[5, 6, 7]])
         start_end_packer = StartEndPacker(sequence_length=5)
-        with self.assertRaisesRegex(
-            ValueError,
-            "Input must either be dense with rank 1, or ragged with rank 2. "
-            "Received dense input with rank=2",
-        ):
-            start_end_packer(input_data)
+        output = start_end_packer(input_data)
+        expected_output = [[5, 6, 7, 0, 0]]
+        self.assertAllEqual(output, expected_output)
 
     def test_ragged_input(self):
         input_data = tf.ragged.constant([[5, 6, 7], [8, 9, 10, 11]])
@@ -51,8 +48,7 @@ class StartEndPackerTest(tf.test.TestCase):
         start_end_packer = StartEndPacker(sequence_length=5)
         with self.assertRaisesRegex(
             ValueError,
-            "Input must either be dense with rank 1, or ragged with rank 2. "
-            "Received ragged input with "
+            "Input must either be rank 1 or rank 2. Received input with "
             "rank=3",
         ):
             start_end_packer(input_data)