Repository: keras-team/keras-nlp

Issue:
Add a start and end token packer layer
```
class StartEndPacker(keras.layers.Layer):
    """Adds start and end tokens to a sequence and pads to a fixed length.

    If inputs are batched, input should be a `tf.RaggedTensor`s with shape
    `[batch_size, None]` and will be packed and converted to a dense tensor with
    shape `[batch_size, sequence_length]`.
    If inputs are unbatched, inputs should be dense rank-1 tensors of any shape,
    and will be packed to shape `[sequence_length]`.

    Args:
        sequence_length: The desired output length.
        start_value: The id or token that is to be placed at the start of each
            sequence. The dtype must mach the dtype of the input tensors to
            the layer. If None, no start value will be added.
        end_value: The id or token that is to be placed at the end of each
            input segment. The dtype must mach the dtype of the input tensors to
            the layer. If None, no end value will be added.
        pad_value: The id or token that is to be placed into the unused
            positions after the last segment in the sequence.
    """
    def __init__(
        self,
        sequence_length,
        start_value=None,
        end_value=None,
        pad_value=None,
    )
```


Discussion:
@abheesht17 rough docstring above. Feel free to give any feedback!

If start_value is not provided, don't add it.
If end_value is not provided, don't add it.
If pad_value is not provided, pad with the defaults of `ragged.to_dense()`. E.g. `0` or `""` depending on type.

Ideally this could replace your preprocessing for the translation example so preprocessing looks like `packer(tokenizer(inputs))`. Though we can submit the example first, and do this as a follow up.
The doc-string LGTM! Will open a PR for this ASAP.

Resolving patch:
diff --git a/keras_nlp/layers/__init__.py b/keras_nlp/layers/__init__.py
--- a/keras_nlp/layers/__init__.py
+++ b/keras_nlp/layers/__init__.py
@@ -17,6 +17,7 @@
 from keras_nlp.layers.mlm_mask_generator import MLMMaskGenerator
 from keras_nlp.layers.position_embedding import PositionEmbedding
 from keras_nlp.layers.sine_position_encoding import SinePositionEncoding
+from keras_nlp.layers.start_end_packer import StartEndPacker
 from keras_nlp.layers.token_and_position_embedding import (
     TokenAndPositionEmbedding,
 )
diff --git a/keras_nlp/layers/start_end_packer.py b/keras_nlp/layers/start_end_packer.py
new file mode 100644
--- /dev/null
+++ b/keras_nlp/layers/start_end_packer.py
@@ -0,0 +1,166 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Start End Packer implementation based on `keras.layers.Layer`."""
+
+import tensorflow as tf
+from tensorflow import keras
+
+
+class StartEndPacker(keras.layers.Layer):
+    """Adds start and end tokens to a sequence and pads to a fixed length.
+
+    This layer is useful when tokenizing inputs for tasks like translation,
+    where each sequence should include a start and end marker. It should
+    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]`.
+
+    Args:
+        sequence_length: int. The desired output length.
+        start_value: int/str. The ID or token that is to be placed at the start
+            of each sequence. The dtype must match the dtype of the input
+            tensors to the layer. If None, no start value will be added.
+        end_value: int/str. The ID or token that is to be placed at the end of
+            each input segment. The dtype must match the dtype of the input
+            tensors to the layer. If None, no end value will be added.
+        pad_value: int/str. The ID or token that is to be placed into the
+            unused positions after the last segment in the sequence. If None,
+            0 or "" will be added depending on the dtype of the input tensor.
+
+    Examples:
+
+    Unbatched input (int).
+    >>> input_data = tf.constant([5, 6, 7])
+    >>> start_end_packer = keras_nlp.layers.StartEndPacker(
+    ...     sequence_length=7, start_value=1, end_value=2,
+    ... )
+    >>> start_end_packer(input_data)
+    <tf.Tensor: shape=(7,), dtype=int32, numpy=
+    array([1, 5, 6, 7, 2, 0, 0], dtype=int32)>
+
+    Batched input (int).
+    >>> input_data = tf.ragged.constant([[5, 6, 7], [8, 9, 10, 11, 12, 13, 14]])
+    >>> start_end_packer = keras_nlp.layers.StartEndPacker(
+    ...     sequence_length=6, start_value=1, end_value=2,
+    ... )
+    >>> start_end_packer(input_data)
+    <tf.Tensor: shape=(2, 6), dtype=int32, numpy=
+    array([[ 1,  5,  6,  7,  2,  0],
+           [ 1,  8,  9, 10, 11,  2]], dtype=int32)>
+
+    Unbatched input (str).
+    >>> input_data = tf.constant(["this", "is", "fun"])
+    >>> start_end_packer = keras_nlp.layers.StartEndPacker(
+    ...     sequence_length=6, start_value="
", end_value="
",
+    ...     pad_value="<pad>"
+    ... )
+    >>> start_end_packer(input_data)
+    <tf.Tensor: shape=(6,), dtype=string, numpy=
+    array([b'
', b'this', b'is', b'fun', b'
', b'<pad>'], dtype=object)>
+
+    Batched input (str).
+    >>> input_data = tf.ragged.constant([["this", "is", "fun"], ["awesome"]])
+    >>> start_end_packer = keras_nlp.layers.StartEndPacker(
+    ...     sequence_length=6, start_value="
", end_value="
",
+    ...     pad_value="<pad>"
+    ... )
+    >>> start_end_packer(input_data)
+    <tf.Tensor: shape=(2, 6), dtype=string, numpy=
+    array([[b'
', b'this', b'is', b'fun', b'
', b'<pad>'],
+           [b'
', b'awesome', b'
', b'<pad>', b'<pad>', b'<pad>']],
+          dtype=object)>
+    """
+
+    def __init__(
+        self,
+        sequence_length,
+        start_value=None,
+        end_value=None,
+        pad_value=None,
+        name=None,
+        **kwargs,
+    ):
+        super().__init__(name=name, **kwargs)
+
+        self.sequence_length = sequence_length
+        self.start_value = start_value
+        self.end_value = end_value
+        self.pad_value = pad_value
+
+    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}"
+                )
+
+        batch_size = tf.shape(inputs)[0]
+
+        # Concatenate start and end tokens.
+        if self.start_value is not None:
+            start_token_id_tensor = tf.fill((batch_size, 1), self.start_value)
+            inputs = tf.concat([start_token_id_tensor, inputs], axis=-1)
+        if self.end_value is not None:
+            end_token_id_tensor = tf.fill((batch_size, 1), self.end_value)
+
+            # Trim to leave room for end token.
+            inputs = inputs[..., : self.sequence_length - 1]
+            inputs = tf.concat([inputs, end_token_id_tensor], axis=-1)
+
+        # Pad to desired length.
+        inputs = inputs.to_tensor(
+            default_value=self.pad_value,
+            shape=(batch_size, self.sequence_length),
+        )
+
+        if input_is_dense:
+            inputs = tf.squeeze(inputs, axis=0)
+
+        return inputs
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "sequence_length": self.sequence_length,
+                "start_value": self.start_value,
+                "end_value": self.end_value,
+                "pad_value": self.pad_value,
+            }
+        )
+        return config


Test patch:
diff --git a/keras_nlp/layers/start_end_packer_test.py b/keras_nlp/layers/start_end_packer_test.py
new file mode 100644
--- /dev/null
+++ b/keras_nlp/layers/start_end_packer_test.py
@@ -0,0 +1,151 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for Start End Packer layer."""
+
+
+import tensorflow as tf
+from tensorflow import keras
+
+from keras_nlp.layers.start_end_packer import StartEndPacker
+
+
+class StartEndPackerTest(tf.test.TestCase):
+    def test_dense_input(self):
+        input_data = tf.constant([5, 6, 7])
+        start_end_packer = StartEndPacker(sequence_length=5)
+        output = start_end_packer(input_data)
+        expected_output = [5, 6, 7, 0, 0]
+        self.assertAllEqual(output, expected_output)
+
+    def test_dense_input_error(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)
+
+    def test_ragged_input(self):
+        input_data = tf.ragged.constant([[5, 6, 7], [8, 9, 10, 11]])
+        start_end_packer = StartEndPacker(sequence_length=5)
+        output = start_end_packer(input_data)
+        expected_output = [[5, 6, 7, 0, 0], [8, 9, 10, 11, 0]]
+        self.assertAllEqual(output, expected_output)
+
+    def test_ragged_input_error(self):
+        input_data = tf.ragged.constant([[[5, 6, 7], [8, 9, 10, 11]]])
+        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 "
+            "rank=3",
+        ):
+            start_end_packer(input_data)
+
+    def test_start_end_token(self):
+        input_data = tf.ragged.constant([[5, 6, 7], [8, 9, 10, 11]])
+        start_end_packer = StartEndPacker(
+            sequence_length=6, start_value=1, end_value=2
+        )
+        output = start_end_packer(input_data)
+        expected_output = [[1, 5, 6, 7, 2, 0], [1, 8, 9, 10, 11, 2]]
+        self.assertAllEqual(output, expected_output)
+
+    def test_start_end_padding_value(self):
+        input_data = tf.ragged.constant([[5, 6, 7], [8, 9, 10, 11]])
+        start_end_packer = StartEndPacker(
+            sequence_length=7, start_value=1, end_value=2, pad_value=3
+        )
+        output = start_end_packer(input_data)
+        expected_output = [[1, 5, 6, 7, 2, 3, 3], [1, 8, 9, 10, 11, 2, 3]]
+        self.assertAllEqual(output, expected_output)
+
+    def test_end_token_value_during_truncation(self):
+        input_data = tf.ragged.constant([[5, 6], [8, 9, 10, 11, 12, 13]])
+        start_end_packer = StartEndPacker(
+            sequence_length=5, start_value=1, end_value=2, pad_value=0
+        )
+        output = start_end_packer(input_data)
+        expected_output = [[1, 5, 6, 2, 0], [1, 8, 9, 10, 2]]
+        self.assertAllEqual(output, expected_output)
+
+    def test_string_input(self):
+        input_data = tf.ragged.constant(
+            [["KerasNLP", "is", "awesome"], ["amazing"]]
+        )
+        start_end_packer = StartEndPacker(
+            sequence_length=5,
+            start_value="[START]",
+            end_value="[END]",
+            pad_value="[PAD]",
+        )
+        output = start_end_packer(input_data)
+        expected_output = [
+            ["[START]", "KerasNLP", "is", "awesome", "[END]"],
+            ["[START]", "amazing", "[END]", "[PAD]", "[PAD]"],
+        ]
+        self.assertAllEqual(output, expected_output)
+
+    def test_functional_model(self):
+        input_data = tf.ragged.constant([[5, 6, 7], [8, 9, 10, 11]])
+        start_end_packer = StartEndPacker(
+            sequence_length=7, start_value=1, end_value=2, pad_value=3
+        )
+
+        inputs = keras.Input(dtype=tf.int32, shape=())
+        outputs = start_end_packer(inputs)
+        model = keras.Model(inputs, outputs)
+        model_output = model(input_data)
+
+        expected_output = [[1, 5, 6, 7, 2, 3, 3], [1, 8, 9, 10, 11, 2, 3]]
+        self.assertAllEqual(model_output, expected_output)
+
+    def test_batch(self):
+        start_end_packer = StartEndPacker(
+            sequence_length=7, start_value=1, end_value=2, pad_value=3
+        )
+
+        ds = tf.data.Dataset.from_tensor_slices(
+            tf.ragged.constant([[5, 6, 7], [8, 9, 10, 11]])
+        )
+        ds = ds.batch(2).map(start_end_packer)
+        output = ds.take(1).get_single_element()
+
+        exp_output = [[1, 5, 6, 7, 2, 3, 3], [1, 8, 9, 10, 11, 2, 3]]
+
+        for i in range(output.shape[0]):
+            self.assertAllEqual(output[i], exp_output[i])
+
+    def test_get_config(self):
+        start_end_packer = StartEndPacker(
+            sequence_length=512,
+            start_value=10,
+            end_value=20,
+            pad_value=100,
+            name="start_end_packer_test",
+        )
+
+        config = start_end_packer.get_config()
+        expected_config_subset = {
+            "sequence_length": 512,
+            "start_value": 10,
+            "end_value": 20,
+            "pad_value": 100,
+        }
+
+        self.assertEqual(config, {**config, **expected_config_subset})
diff --git a/keras_nlp/tokenizers/byte_tokenizer_test.py b/keras_nlp/tokenizers/byte_tokenizer_test.py
--- a/keras_nlp/tokenizers/byte_tokenizer_test.py
+++ b/keras_nlp/tokenizers/byte_tokenizer_test.py
@@ -130,7 +130,6 @@ def test_tokenize_first_batch_second(self):
             [104, 97, 104, 97],
         ]
         for i in range(output.shape[0]):
-            print(output[i])
             self.assertAllEqual(output[i], exp_output[i])
 
     def test_tokenize_first_batch_second_with_sequence_length(self):
@@ -150,7 +149,6 @@ def test_tokenize_first_batch_second_with_sequence_length(self):
             [104, 97, 104, 97, 0, 0, 0, 0, 0, 0],
         ]
         for i in range(output.shape[0]):
-            print(output[i])
             self.assertAllEqual(output[i], exp_output[i])
 
     def test_batch_first_tokenize_second(self):
@@ -169,7 +167,6 @@ def test_batch_first_tokenize_second(self):
             [104, 97, 104, 97],
         ]
         for i in range(output.shape[0]):
-            print(output[i])
             self.assertAllEqual(output[i], exp_output[i])
 
     def test_batch_first_tokenize_second_with_sequence_length(self):
@@ -188,7 +185,6 @@ def test_batch_first_tokenize_second_with_sequence_length(self):
             [104, 97, 104, 97, 0, 0, 0, 0, 0, 0],
         ]
         for i in range(output.shape[0]):
-            print(output[i])
             self.assertAllEqual(output[i], exp_output[i])
 
     def test_functional_model(self):