Add support for StartEndPacker packing 2D tensor (#240)

* Added support for packing 2D tensor.

* minor edit to test

* minor fixes

* minor name changes

* input shape check fixes and docstring simplification

* minor changes

Author: jessechancy

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 @@ def call(self, inputs):
         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 @@ def call(self, inputs):
             shape=(batch_size, self.sequence_length),
         )
 
-        if input_is_dense:
+        if input_is_1d:
             inputs = tf.squeeze(inputs, axis=0)
 
         return inputs

keras_nlp/layers/start_end_packer_test.py

@@ -29,15 +29,12 @@ def test_dense_input(self):
         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 @@ def test_ragged_input_error(self):
         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)