Repository: keras-team/keras-nlp

Issue:
Add a ByteTokenizer
This would be a simple tokenizer which has no vocabulary, and simply converts text to raw bytes.

## Potential docstring and usage

```python
"""
Raw byte tokenizer.

This tokenizer is a vocabulary free tokenizer which will tokenize text as
as raw bytes from [0, 256).
 
Args:
   lowercase: if true, lowercase text before tokenizing.
   sequence_length: If set, the output will be converted to a dense
     tensor and padded/trimmed so all outputs are of sequence_length.
   normalization_form: One of the following string values (None, 'NFC',
     'NFKC', 'NFD', 'NFKD'). If set will normalize unicode to the given form
     before tokenizing.
   errors: One of ('replace', 'remove', 'strict'). Specifies the `detokenize()`
     behavior when an invalid byte sequence is encountered. (same behavior as
     https://www.tensorflow.org/api_docs/python/tf/strings/unicode_transcode)
"""

tokenizer = keras_nlp.tokenizers.ByteTokenizer()

inputs = ["▀▁▂▃▄", "hello"]
tokens = tokenizer.tokenize(inputs)
>>> <tf.RaggedTensor [
      [226, 150, 128, 226, 150, 129, 226, 150, 130, 226, 150, 131, 226, 150, 132],
      [104, 101, 108, 108, 111]]>


outputs = tokenizer.detokenize(tokens)
outputs = [x.decode("utf-8") for x in strings.numpy().tolist()]
>>> ['▀▁▂▃▄', 'hello']
```

## Colab demonstrating basic functionality:

https://colab.sandbox.google.com/gist/mattdangerw/99e8f3795e37fe539731dcbfc1c09d47/bytetokenizer.ipynb

Discussion:
abheesht17: @mattdangerw, I've made the required changes. I've done something for https://github.com/keras-team/keras-nlp/pull/80#discussion_r840098827 but I am not sure if it is correct. 

Secondly, I am not quite sure how the test cases for the fix for the above comment should be written. In particular, I am confused about this: https://github.com/keras-team/keras-nlp/pull/80/files#diff-0356ce59821651a89b22329ffcf6ca02833133960abcee141a0c7eeee5623ccaR56-R58.

P.S. Sorry for the delay, took some time to figure it out.
mattdangerw: @abheesht17 thanks!

The way I would test this is by writing unit tests for two flows--a "batch then tokenize" flow and a "tokenize then batch" flow. I would actually add these as unit tests using `tf.data`. There's four cases considering both seq length and batching order. I made a colab showing examples.

https://colab.sandbox.google.com/gist/mattdangerw/b163a5798ed20d194e209f83678c25f3/rank1tokenizer.ipynb

Playing around in the colab, it looks like you might be able to simplify your code a lot if you
 - uprank scalar inputs with a leading axis at the beginning of call
 - run unified logic for tokenizing  as raggeds and optionally converting to dense
 - downrank outputs in the scalar case by squeezing off the leading axis

It looks like you don't actually need any `set_shape` call in this case. As `ragged.to_tensor()` will to the right thing here.

Hope that helps!
mattdangerw: @abheesht17 thanks!

The way I would test this is by writing unit tests for two flows--a "batch then tokenize" flow and a "tokenize then batch" flow. I would actually add these as unit tests using `tf.data`. There's four cases considering both seq length and batching order. I made a colab showing examples.

https://colab.sandbox.google.com/gist/mattdangerw/b163a5798ed20d194e209f83678c25f3/rank1tokenizer.ipynb

Playing around in the colab, it looks like you might be able to simplify your code a lot if you
 - uprank scalar inputs with a leading axis at the beginning of call
 - run unified logic for tokenizing  as raggeds and optionally converting to dense
 - downrank outputs in the scalar case by squeezing off the leading axis

It looks like you don't actually need any `set_shape` call in this case. As `ragged.to_tensor()` will to the right thing here.

Hope that helps!
abheesht17: @mattdangerw, I've incorporated your suggestions. Let me know if there are further changes to be made. Thank you!

Patch:
diff --git a/examples/bert/bert_model.py b/examples/bert/bert_model.py
index 3283d8b5..4569376a 100644
--- a/examples/bert/bert_model.py
+++ b/examples/bert/bert_model.py
@@ -19,6 +19,8 @@ We should work to replace this with components from the keras-nlp library.
 import tensorflow as tf
 from tensorflow import keras
 
+import keras_nlp
+
 
 def make_attention_mask(inputs, mask):
     """Make a 3D attention mask from a 2D input mask.
@@ -331,72 +333,6 @@ class TransformerEncoderBlock(keras.layers.Layer):
         return self._output_layer_norm(layer_output + attention_output)
 
 
-class PositionEmbedding(keras.layers.Layer):
-    """Creates a positional embedding.
-
-    Example:
-    ```python
-    position_embedding = PositionEmbedding(max_length=100)
-    inputs = keras.Input((100, 32), dtype=tf.float32)
-    outputs = position_embedding(inputs)
-    ```
-
-
-    Args:
-        max_length: The maximum size of the dynamic sequence.
-        initializer: The initializer to use for the embedding weights. Defaults
-            to "glorot_uniform".
-        seq_axis: The axis of the input tensor where we add the embeddings.
-
-    Reference: This layer creates a positional embedding as described in
-    [BERT: Pre-training of Deep Bidirectional Transformers for Language
-    Understanding](https://arxiv.org/abs/1810.04805).
-    """
-
-    def __init__(
-        self, max_length, initializer="glorot_uniform", seq_axis=1, **kwargs
-    ):
-
-        super().__init__(**kwargs)
-        if max_length is None:
-            raise ValueError("`max_length` must be an Integer, not `None`.")
-        self._max_length = max_length
-        self._initializer = keras.initializers.get(initializer)
-        self._seq_axis = seq_axis
-
-    def get_config(self):
-        config = {
-            "max_length": self._max_length,
-            "initializer": keras.initializers.serialize(self._initializer),
-            "seq_axis": self._seq_axis,
-        }
-        base_config = super().get_config()
-        return dict(list(base_config.items()) + list(config.items()))
-
-    def build(self, input_shape):
-        dimension_list = input_shape.as_list()
-        width = dimension_list[-1]
-        weight_sequence_length = self._max_length
-
-        self._position_embeddings = self.add_weight(
-            "embeddings",
-            shape=[weight_sequence_length, width],
-            initializer=self._initializer,
-        )
-
-        super().build(input_shape)
-
-    def call(self, inputs):
-        input_shape = tf.shape(inputs)
-        actual_seq_len = input_shape[self._seq_axis]
-        position_embeddings = self._position_embeddings[:actual_seq_len, :]
-        new_shape = [1 for _ in inputs.get_shape().as_list()]
-        new_shape[self._seq_axis] = actual_seq_len
-        new_shape[-1] = position_embeddings.get_shape().as_list()[-1]
-        position_embeddings = tf.reshape(position_embeddings, new_shape)
-        return tf.broadcast_to(position_embeddings, input_shape)
-
-
 # TODO(mattdangerw): This class is needed for TPU friendly embeddings, we should
 # remove it entirely and fix tf.keras.layers.Embedding as needed.
 class OnDeviceEmbedding(keras.layers.Layer):
@@ -546,7 +482,7 @@ class BertModel(keras.Model):
             name="word_embeddings",
         )
 
-        self._position_embedding_layer = PositionEmbedding(
+        self._position_embedding_layer = keras_nlp.layers.PositionEmbedding(
             initializer=initializer,
             max_length=max_sequence_length,
             name="position_embedding",
diff --git a/keras_nlp/__init__.py b/keras_nlp/__init__.py
index 46b298c1..2c9463db 100644
--- a/keras_nlp/__init__.py
+++ b/keras_nlp/__init__.py
@@ -15,4 +15,4 @@
 from keras_nlp import layers
 from keras_nlp import tokenizers
 
-__version__ = "0.1.0"
+__version__ = "0.1.1"
diff --git a/keras_nlp/layers/__init__.py b/keras_nlp/layers/__init__.py
index d2f5b578..fe677bbe 100644
--- a/keras_nlp/layers/__init__.py
+++ b/keras_nlp/layers/__init__.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 from keras_nlp.layers.fnet_encoder import FNetEncoder
+from keras_nlp.layers.position_embedding import PositionEmbedding
 from keras_nlp.layers.sine_position_encoding import SinePositionEncoding
 from keras_nlp.layers.transformer_decoder import TransformerDecoder
 from keras_nlp.layers.transformer_encoder import TransformerEncoder
diff --git a/keras_nlp/layers/position_embedding.py b/keras_nlp/layers/position_embedding.py
new file mode 100644
index 00000000..622d0e3d
--- /dev/null
+++ b/keras_nlp/layers/position_embedding.py
@@ -0,0 +1,113 @@
+# 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.
+
+"""Position embedding implementation based on `keras.layers.Layer`."""
+
+import tensorflow as tf
+from tensorflow import keras
+
+SEQUENCE_AXIS = -2
+
+
+class PositionEmbedding(keras.layers.Layer):
+    """Creates a layer which learns a position embedding for inputs sequences.
+
+    This class assumes that in the input tensor, the last dimension corresponds
+    to the features, and the dimension before the last corresponds to the
+    sequence.
+
+    This class accepts `RaggedTensor`s as inputs to process batches of sequences
+    of different lengths. The one ragged dimension must be the dimension that
+    corresponds to the sequence, that is, the penultimate dimension.
+
+    Args:
+        max_length: The maximum length of the dynamic sequence.
+        initializer: The initializer to use for the embedding weights. Defaults
+            to "glorot_uniform".
+        seq_axis: The axis of the input tensor where we add the embeddings.
+
+    Example:
+    ```python
+    token_embeddings = layers.Embedding(
+        input_dim=vocab_size, output_dim=embed_dim
+    )
+    position_embeddings = keras_nlp.layers.PositionEmbedding(
+        max_length=max_length
+    )
+
+    embedded_tokens = token_embeddings(inputs)
+    embedded_positions = position_embeddings(embedded_tokens)
+    outputs = embedded_tokens + embedded_positions
+    ```
+
+    Reference:
+        [BERT: Pre-training of Deep Bidirectional Transformers for Language
+        Understanding](https://arxiv.org/abs/1810.04805).
+    """
+
+    def __init__(
+        self,
+        max_length,
+        initializer="glorot_uniform",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        if max_length is None:
+            raise ValueError("`max_length` must be an Integer, not `None`.")
+        self.max_length = int(max_length)
+        self.initializer = keras.initializers.get(initializer)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "max_length": self.max_length,
+                "initializer": keras.initializers.serialize(self.initializer),
+            }
+        )
+        return config
+
+    def build(self, input_shape):
+        feature_size = input_shape[-1]
+        self.position_embeddings = self.add_weight(
+            "embeddings",
+            shape=[self.max_length, feature_size],
+            initializer=self.initializer,
+            trainable=True,
+        )
+
+        super().build(input_shape)
+
+    def call(self, inputs):
+        if isinstance(inputs, tf.RaggedTensor):
+            bounding_shape = inputs.bounding_shape()
+            position_embeddings = self._trim_and_broadcast_position_embeddings(
+                bounding_shape,
+            )
+            # then apply row lengths to recreate the same ragged shape as inputs
+            return tf.RaggedTensor.from_tensor(
+                position_embeddings,
+                inputs.nested_row_lengths(),
+            )
+        else:
+            return self._trim_and_broadcast_position_embeddings(
+                tf.shape(inputs),
+            )
+
+    def _trim_and_broadcast_position_embeddings(self, shape):
+        sequence_length = shape[SEQUENCE_AXIS]
+        # trim to match the length of the sequence
+        position_embeddings = self.position_embeddings[:sequence_length, :]
+        # then broadcast to add the missing dimensions to match "shape"
+        return tf.broadcast_to(position_embeddings, shape)
diff --git a/keras_nlp/tokenizers/__init__.py b/keras_nlp/tokenizers/__init__.py
index d6eb0b63..7206410a 100644
--- a/keras_nlp/tokenizers/__init__.py
+++ b/keras_nlp/tokenizers/__init__.py
@@ -12,5 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from keras_nlp.tokenizers.byte_tokenizer import ByteTokenizer
 from keras_nlp.tokenizers.tokenizer import Tokenizer
 from keras_nlp.tokenizers.word_piece_tokenizer import WordPieceTokenizer
diff --git a/keras_nlp/tokenizers/byte_tokenizer.py b/keras_nlp/tokenizers/byte_tokenizer.py
new file mode 100644
index 00000000..2f142004
--- /dev/null
+++ b/keras_nlp/tokenizers/byte_tokenizer.py
@@ -0,0 +1,256 @@
+# 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.
+
+"""Byte Tokenizer."""
+
+from typing import Any
+from typing import Dict
+
+import numpy as np
+import tensorflow as tf
+import tensorflow_text as tf_text
+
+from keras_nlp.tokenizers import tokenizer
+
+
+class ByteTokenizer(tokenizer.Tokenizer):
+    """Raw byte tokenizer.
+
+    This tokenizer is a vocabulary-free tokenizer which will tokenize text as
+    as raw bytes from [0, 256).
+
+    If input is a batch of strings:
+    By default, the layer will output a `tf.RaggedTensor` where the last
+    dimension of the output is ragged. If `sequence_length` is set, the layer
+    will output a dense `tf.Tensor` where all inputs have been padded or
+    truncated to `sequence_length`. The output dtype can be controlled via the
+    `dtype` argument, which should be an integer type
+    (tf.int16, tf.int32, etc.).
+
+    If input is a scalar string:
+    There are two cases here. If `sequence_length` is set, the output will be
+    a dense `tf.Tensor` of shape `[sequence_length]`. Otherwise, the output will
+    be a dense `tf.Tensor` of shape `[None]`.
+
+    Args:
+        lowercase: boolean. If True, the input text will be converted to
+            lowercase before tokenization.
+        sequence_length: int. If set, the output will be converted to a dense
+            tensor and padded/trimmed so all outputs are of sequence_length.
+        normalization_form: string. One of the following values: (None, "NFC",
+            "NFKC", "NFD", "NFKD"). If set, every UTF-8 string in the input
+            tensor text will be normalized to the given form before tokenizing.
+        errors: string. One of ("strict", "replace", "ignore"). Defaults to
+            "replace". Specifies the `detokenize()` behaviour when an invalid
+            byte sequence is encountered (same behaviour as
+            https://www.tensorflow.org/api_docs/python/tf/strings/unicode_transcode).
+        replacement_char: int. Defaults to 65533. The replacement character to
+            use when an invalid byte sequence is encountered and when `errors`
+            is set to "replace" (same behaviour as
+            https://www.tensorflow.org/api_docs/python/tf/strings/unicode_transcode).
+
+    Examples:
+
+    Basic usage.
+    >>> tokenizer = keras_nlp.tokenizers.ByteTokenizer()
+    >>> tokenizer("hello")
+    <tf.Tensor: shape=(5,), dtype=int32, numpy=
+    array([104, 101, 108, 108, 111], dtype=int32)>
+
+    Ragged outputs.
+    >>> inputs = tf.constant(["hello", "hi"])
+    >>> tokenizer = keras_nlp.tokenizers.ByteTokenizer()
+    >>> tokenizer(inputs)
+    <tf.RaggedTensor [[104, 101, 108, 108, 111], [104, 105]]>
+
+    Dense outputs.
+    >>> inputs = tf.constant(["hello", "hi"])
+    >>> tokenizer = keras_nlp.tokenizers.ByteTokenizer(sequence_length=8)
+    >>> tokenizer(inputs)
+    <tf.Tensor: shape=(2, 8), dtype=int32, numpy=
+    array([[104, 101, 108, 108, 111,   0,   0,   0],
+           [104, 105,   0,   0,   0,   0,   0,   0]], dtype=int32)>
+
+    Dense outputs.
+    >>> inputs = tf.constant(["hello"])
+    >>> tokenizer = keras_nlp.tokenizers.ByteTokenizer(sequence_length=8)
+    >>> tokenizer(inputs)
+    <tf.Tensor: shape=(1, 8), dtype=int32,
+                numpy=array([[104, 101, 108, 108, 111,   0,   0,   0]])>
+
+    Tokenize first, then batch the dataset up.
+    >>> tokenizer = keras_nlp.tokenizers.ByteTokenizer()
+    >>> ds = tf.data.Dataset.from_tensor_slices(["hello", "fun"])
+    >>> ds = ds.map(tokenizer)
+    >>> ds = ds.apply(tf.data.experimental.dense_to_ragged_batch(2))
+    >>> ds.take(1).get_single_element()
+    <tf.RaggedTensor [[104, 101, 108, 108, 111], [102, 117, 110]]>
+
+    Batch up the inputs and then tokenize.
+    >>> tokenizer = keras_nlp.tokenizers.ByteTokenizer()
+    >>> ds = tf.data.Dataset.from_tensor_slices(["hello", "fun"])
+    >>> ds = ds.batch(2).map(tokenizer)
+    >>> ds.take(1).get_single_element()
+    <tf.RaggedTensor [[104, 101, 108, 108, 111], [102, 117, 110]]>
+
+    Tokenize first, then batch the dataset up (`sequence_length` provided).
+    >>> tokenizer = keras_nlp.tokenizers.ByteTokenizer(sequence_length=5)
+    >>> ds = tf.data.Dataset.from_tensor_slices(["hello", "fun"])
+    >>> ds = ds.map(tokenizer)
+    >>> ds = ds.apply(tf.data.experimental.dense_to_ragged_batch(2))
+    >>> ds.take(1).get_single_element()
+    <tf.Tensor: shape=(2, 5), dtype=int32, numpy=
+    array([[104, 101, 108, 108, 111],
+        [102, 117, 110,   0,   0]])>
+
+    Batch up the inputs and then tokenize (`sequence_length` provided).
+    >>> tokenizer = keras_nlp.tokenizers.ByteTokenizer(sequence_length=5)
+    >>> ds = tf.data.Dataset.from_tensor_slices(["hello", "fun"])
+    >>> ds = ds.batch(2).map(tokenizer)
+    >>> ds.take(1).get_single_element()
+    <tf.Tensor: shape=(2, 5), dtype=int32, numpy=
+    array([[104, 101, 108, 108, 111],
+        [102, 117, 110,   0,   0]])>
+
+    Detokenization.
+    >>> inputs = tf.constant([104, 101, 108, 108, 111], dtype=tf.int32)
+    >>> tokenizer = keras_nlp.tokenizers.ByteTokenizer()
+    >>> tokenizer.detokenize(inputs)
+    <tf.Tensor: shape=(), dtype=string, numpy=b'hello'>
+
+    Detokenization with invalid bytes.
+    >>> # The 255 below is invalid utf-8.
+    >>> inputs = tf.constant([104, 101, 255, 108, 108, 111], dtype=tf.int32)
+    >>> tokenizer = keras_nlp.tokenizers.ByteTokenizer(
+    ...     errors="replace", replacement_char=88)
+    >>> tokenizer.detokenize(inputs).numpy().decode('utf-8')
+    'heXllo'
+    """
+
+    def __init__(
+        self,
+        lowercase: bool = True,
+        sequence_length: int = None,
+        normalization_form: str = None,
+        errors: str = "replace",
+        replacement_char: int = 65533,
+        **kwargs,
+    ):
+        # Check dtype and provide a default.
+        if "dtype" not in kwargs or kwargs["dtype"] is None:
+            kwargs["dtype"] = tf.int32
+        else:
+            dtype = tf.dtypes.as_dtype(kwargs["dtype"])
+            if not dtype.is_integer:
+                raise ValueError(
+                    "Output dtype must be an integer type. "
+                    f"Received: dtype={dtype}"
+                )
+
+        # Check normalization_form.
+        if normalization_form not in (None, "NFC", "NFKC", "NFD", "NFKD"):
+            raise ValueError(
+                '`normalization_form` must be one of None, "NFC", "NFKC", '
+                '"NFD", "NFKD". Received: normalization_form='
+                f"{normalization_form}"
+            )
+
+        # Check errors.
+        if errors not in ("strict", "replace", "ignore"):
+            raise ValueError(
+                '`errors` must be one of "strict", "replace", "ignore" '
+                f"Received: errors={errors}"
+            )
+
+        super().__init__(**kwargs)
+
+        self.lowercase = lowercase
+        self.sequence_length = sequence_length
+        self.normalization_form = normalization_form
+        self.errors = errors
+        self.replacement_char = replacement_char
+
+        self._char_lst = tf.constant(
+            [i.tobytes() for i in np.arange(256, dtype=np.uint8)]
+        )
+
+    def vocabulary_size(self) -> int:
+        """Get the size of the tokenizer vocabulary."""
+        return 256
+
+    def tokenize(self, inputs):
+
+        if not isinstance(inputs, (tf.Tensor, tf.RaggedTensor)):
+            inputs = tf.convert_to_tensor(inputs)
+
+        scalar_input = inputs.shape.rank == 0
+        if scalar_input:
+            inputs = tf.expand_dims(inputs, 0)
+
+        # Optional: Lowercase the input.
+        if self.lowercase:
+            inputs = tf_text.case_fold_utf8(inputs)
+
+        # Optional: Normalize unicode.
+        if self.normalization_form is not None:
+            inputs = tf_text.normalize_utf8(inputs, self.normalization_form)
+
+        # Tokenize input strings.
+        tokens = tf.strings.bytes_split(inputs)
+        tokens = tf.squeeze(
+            tf.ragged.map_flat_values(tf.io.decode_raw, tokens, tf.uint8), -1
+        )
+        tokens = tf.cast(tokens, self.compute_dtype)
+
+        # Convert to a dense output if `sequence_length` is set.
+        if self.sequence_length:
+            output_shape = tokens.shape.as_list()
+            output_shape[-1] = self.sequence_length
+            tokens = tokens.to_tensor(shape=output_shape)
+
+        if scalar_input:
+            tokens = tf.squeeze(tokens, 0)
+        return tokens
+
+    def detokenize(self, inputs):
+        # Remove trailing padding tokens, so that trailing "\x00" bytes don't
+        # show up in the detokenized output.
+        inputs = tf.ragged.boolean_mask(inputs, tf.not_equal(inputs, 0))
+
+        decoded = tf.strings.reduce_join(
+            tf.gather(self._char_lst, inputs), axis=-1
+        )
+
+        # Handle errors if an invalid byte sequence is encountered.
+        decoded = tf.strings.unicode_transcode(
+            decoded,
+            "UTF-8",
+            "UTF-8",
+            errors=self.errors,
+            replacement_char=self.replacement_char,
+        )
+        return decoded
+
+    def get_config(self) -> Dict[str, Any]:
+        config = super().get_config()
+        config.update(
+            {
+                "lowercase": self.lowercase,
+                "sequence_length": self.sequence_length,
+                "normalization_form": self.normalization_form,
+                "errors": self.errors,
+                "replacement_char": self.replacement_char,
+            }
+        )
+        return config
diff --git a/setup.py b/setup.py
index bf41657d..e05c744e 100644
--- a/setup.py
+++ b/setup.py
@@ -29,7 +29,7 @@ setup(
     ),
     long_description=README,
     long_description_content_type="text/markdown",
-    version="0.1.0",
+    version="0.1.1",
     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/position_embedding_test.py b/keras_nlp/layers/position_embedding_test.py
new file mode 100644
index 00000000..36befb35
--- /dev/null
+++ b/keras_nlp/layers/position_embedding_test.py
@@ -0,0 +1,328 @@
+# 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 position embedding layer."""
+
+import os
+
+import tensorflow as tf
+from tensorflow import keras
+
+from keras_nlp.layers import position_embedding
+
+
+def custom_init(shape, dtype=None):
+    count = 1
+    for length in shape:
+        count *= length
+    return tf.reshape(tf.range(count, dtype=dtype), shape)
+
+
+class PositionEmbeddingLayerTest(tf.test.TestCase):
+    def test_static_layer_output_shape(self):
+        # Create a 3-dimensional input (the first dimension is implicit).
+        sequence_length = 21
+        feature_size = 30
+        test_layer = position_embedding.PositionEmbedding(
+            max_length=sequence_length
+        )
+        input_tensor = tf.keras.Input(shape=(sequence_length, feature_size))
+        output_tensor = test_layer(input_tensor)
+
+        # When using static position embedding shapes, the output is expected
+        # to be the same as the input shape in all dimensions save batch.
+        expected_output_shape = [None, sequence_length, feature_size]
+        self.assertEqual(expected_output_shape, output_tensor.shape.as_list())
+        # The default output dtype for this layer should be tf.float32.
+        self.assertEqual(tf.float32, output_tensor.dtype)
+
+    def test_more_than_3_dimensions_static(self):
+        # Create a 4-dimensional input (the first dimension is implicit).
+        sequence_length = 21
+        feature_size = 30
+        test_layer = position_embedding.PositionEmbedding(
+            max_length=sequence_length
+        )
+        input_tensor = tf.keras.Input(
+            shape=(feature_size, sequence_length, feature_size)
+        )
+        output_tensor = test_layer(input_tensor)
+
+        # When using static position embedding shapes, the output is expected
+        # to be the same as the input shape in all dimensions save batch.
+        expected_output_shape = [
+            None,
+            feature_size,
+            sequence_length,
+            feature_size,
+        ]
+        self.assertEqual(expected_output_shape, output_tensor.shape.as_list())
+        # The default output dtype for this layer should be tf.float32.
+        self.assertEqual(tf.float32, output_tensor.dtype)
+
+    def test_float16_dtype(self):
+        # Create a 3-dimensional input (the first dimension is implicit).
+        sequence_length = 21
+        feature_size = 30
+        test_layer = position_embedding.PositionEmbedding(
+            max_length=sequence_length, dtype="float16"
+        )
+        input_tensor = tf.keras.Input(shape=(sequence_length, feature_size))
+        output_tensor = test_layer(input_tensor)
+
+        # When using static position embedding shapes, the output is expected
+        # to be the same as the input shape in all dimensions save batch.
+        expected_output_shape = [None, sequence_length, feature_size]
+        self.assertEqual(expected_output_shape, output_tensor.shape.as_list())
+        # The default output dtype for this layer should be tf.float32.
+        self.assertEqual(tf.float16, output_tensor.dtype)
+
+    def test_dynamic_layer_output_shape(self):
+        max_sequence_length = 40
+        feature_size = 30
+        test_layer = position_embedding.PositionEmbedding(
+            max_length=max_sequence_length
+        )
+        # Create a 3-dimensional input (the first dimension is implicit).
+        input_tensor = tf.keras.Input(shape=(None, feature_size))
+        output_tensor = test_layer(input_tensor)
+
+        # When using dynamic position embedding shapes, the output is expected
+        # to be the same as the input shape in all dimensions - but may be None
+        # if the input shape is None there.
+        expected_output_shape = [None, None, feature_size]
+        self.assertEqual(expected_output_shape, output_tensor.shape.as_list())
+
+    def test_more_than_3_dimensions_dynamic(self):
+        max_sequence_length = 60
+        feature_size = 30
+        test_layer = position_embedding.PositionEmbedding(
+            max_length=max_sequence_length
+        )
+        # Create a 4-dimensional input (the first dimension is implicit).
+        input_tensor = tf.keras.Input(shape=(None, None, feature_size))
+        output_tensor = test_layer(input_tensor)
+
+        # When using dynamic position embedding shapes, the output is expected
+        # to be the same as the input shape in all dimensions - but may be None
+        # if the input shape is None there.
+        expected_output_shape = [None, None, None, feature_size]
+        self.assertEqual(expected_output_shape, output_tensor.shape.as_list())
+
+    def test_dynamic_layer_slicing(self):
+        max_sequence_length = 40
+        feature_size = 30
+        test_layer = position_embedding.PositionEmbedding(
+            max_length=max_sequence_length
+        )
+        # Create a 3-dimensional input (the first dimension is implicit).
+        input_tensor = tf.keras.Input(shape=(None, feature_size))
+        output_tensor = test_layer(input_tensor)
+
+        model = tf.keras.Model(input_tensor, output_tensor)
+
+        # Create input data that is shorter than max_sequence_length, which
+        # should trigger a down-slice.
+        input_length = 17
+        # Note: This test explicitly uses a batch size of 1. This is to get
+        # around Keras' restriction on Model invocations: inputs are expected to
+        # have the same batch cardinality as outputs. In practice, this layer
+        # should be used inside a model, where it can be projected when added to
+        # another tensor.
+        input_data = tf.ones(shape=[1, input_length, feature_size])
+        output_data = model.predict(input_data)
+
+        self.assertAllEqual([1, input_length, feature_size], output_data.shape)
+
+    def test_callable_initializer(self):
+        max_sequence_length = 4
+        feature_size = 3
+        test_layer = position_embedding.PositionEmbedding(
+            max_length=max_sequence_length,
+            initializer=custom_init,
+        )
+        inputs = tf.keras.Input(shape=(max_sequence_length, feature_size))
+        outputs = test_layer(inputs)
+        model = keras.Model(inputs=inputs, outputs=outputs)
+
+        batch_size = 2
+        data = tf.zeros(shape=[batch_size, max_sequence_length, feature_size])
+        model(data)
+        model_output = model.predict(data)
+        expected_output = tf.broadcast_to(
+            tf.reshape(
+                tf.range(max_sequence_length * feature_size),
+                [max_sequence_length, feature_size],
+            ),
+            [batch_size, max_sequence_length, feature_size],
+        )
+        self.assertAllClose(model_output, expected_output)
+
+    def test_ragged_tensor_with_3_dimensions(self):
+        max_sequence_length = 4
+        feature_size = 2
+        test_layer = position_embedding.PositionEmbedding(
+            max_length=max_sequence_length,
+            initializer=custom_init,
+        )
+        # Create a 3-dimensional ragged input (the first dimension is implicit).
+        input_tensor = tf.keras.Input(
+            shape=(None, feature_size), dtype=tf.float32, ragged=True
+        )
+        output_tensor = test_layer(input_tensor)
+        model = tf.keras.Model(input_tensor, output_tensor)
+
+        input_data = tf.ragged.constant(
+            [
+                [[1.0, 1.0], [1.0, 1.0]],
+                [],
+                [[1.0, 1.0], [1.0, 1.0], [1.0, 1.0]],
+                [[1.0, 1.0]],
+            ],
+            ragged_rank=1,
+            inner_shape=(2,),
+        )
+        expected_output_data = tf.ragged.constant(
+            [
+                [[0.0, 1.0], [2.0, 3.0]],
+                [],
+                [[0.0, 1.0], [2.0, 3.0], [4.0, 5.0]],
+                [[0.0, 1.0]],
+            ],
+            ragged_rank=1,
+            inner_shape=(2,),
+        )
+        output_data = model.predict(input_data)
+        self.assertAllClose(output_data, expected_output_data)
+
+    def test_ragged_tensor_with_4_dimensions(self):
+        max_sequence_length = 4
+        feature_size = 2
+        test_layer = position_embedding.PositionEmbedding(
+            max_length=max_sequence_length,
+            initializer=custom_init,
+        )
+        # Create a 4-dimensional ragged input (the first dimension is implicit).
+        input_tensor = tf.keras.Input(
+            shape=(None, None, feature_size), dtype=tf.float32, ragged=True
+        )
+        output_tensor = test_layer(input_tensor)
+        model = tf.keras.Model(input_tensor, output_tensor)
+
+        input_data = tf.ragged.constant(
+            [
+                [
+                    [[1.0, 1.0], [1.0, 1.0]],
+                    [],
+                ],
+                [
+                    [[1.0, 1.0], [1.0, 1.0], [1.0, 1.0]],
+                    [[1.0, 1.0]],
+                ],
+            ],
+            ragged_rank=2,
+            inner_shape=(2,),
+        )
+        expected_output_data = tf.ragged.constant(
+            [
+                [
+                    [[0.0, 1.0], [2.0, 3.0]],
+                    [],
+                ],
+                [
+                    [[0.0, 1.0], [2.0, 3.0], [4.0, 5.0]],
+                    [[0.0, 1.0]],
+                ],
+            ],
+            ragged_rank=2,
+            inner_shape=(2,),
+        )
+        output_data = model.predict(input_data)
+        self.assertAllClose(output_data, expected_output_data)
+
+    def test_one_training_step(self):
+        max_sequence_length = 4
+        feature_size = 3
+        test_layer = position_embedding.PositionEmbedding(
+            max_length=max_sequence_length
+        )
+        inputs = tf.keras.Input(shape=(max_sequence_length, feature_size))
+        outputs = test_layer(inputs)
+        model = keras.Model(inputs=inputs, outputs=outputs)
+
+        batch_size = 2
+        data = tf.random.uniform(
+            shape=[batch_size, max_sequence_length, feature_size]
+        )
+        label = tf.random.uniform(shape=[max_sequence_length, feature_size])
+
+        loss_fn = keras.losses.MeanSquaredError()
+        optimizer = keras.optimizers.Adam()
+        with tf.GradientTape() as tape:
+            pred = model(data)
+            loss = loss_fn(label, pred)
+        grad = tape.gradient(loss, model.trainable_variables)
+        self.assertEquals(len(grad), 1)
+
+        trainable_variables_before = tf.Variable(model.trainable_variables[0])
+        optimizer.apply_gradients(zip(grad, model.trainable_variables))
+        self.assertNotAllClose(
+            trainable_variables_before, model.trainable_variables[0]
+        )
+
+    def test_get_config_and_from_config(self):
+        max_sequence_length = 40
+        test_layer = position_embedding.PositionEmbedding(
+            max_length=max_sequence_length
+        )
+        config = test_layer.get_config()
+        expected_config_subset = {
+            "max_length": max_sequence_length,
+            "initializer": {
+                "class_name": "GlorotUniform",
+                "config": {"seed": None},
+            },
+        }
+        self.assertEqual(config, {**config, **expected_config_subset})
+
+        restored_encoder = position_embedding.PositionEmbedding.from_config(
+            config,
+        )
+        self.assertEqual(
+            restored_encoder.get_config(), {**config, **expected_config_subset}
+        )
+
+    def test_save_model(self):
+        max_sequence_length = 4
+        feature_size = 6
+        test_layer = position_embedding.PositionEmbedding(
+            max_length=max_sequence_length
+        )
+        inputs = tf.keras.Input(shape=(max_sequence_length, feature_size))
+        outputs = test_layer(inputs)
+        model = keras.Model(inputs=inputs, outputs=outputs)
+
+        data = tf.zeros(shape=[2, max_sequence_length, feature_size])
+        model(data)
+
+        path = os.path.join(self.get_temp_dir(), "model")
+        model.save(path)
+        loaded_model = keras.models.load_model(path)
+
+        model_output = model.predict(data)
+        loaded_model_output = loaded_model.predict(data)
+        self.assertAllClose(model_output, loaded_model_output)
+
+
+if __name__ == "__main__":
+    tf.test.main()
diff --git a/keras_nlp/tokenizers/byte_tokenizer_test.py b/keras_nlp/tokenizers/byte_tokenizer_test.py
new file mode 100644
index 00000000..80b86bc7
--- /dev/null
+++ b/keras_nlp/tokenizers/byte_tokenizer_test.py
@@ -0,0 +1,266 @@
+# 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.
+
+import tensorflow as tf
+from tensorflow import keras
+
+from keras_nlp.tokenizers.byte_tokenizer import ByteTokenizer
+
+
+class ByteTokenizerTest(tf.test.TestCase):
+    def test_tokenize(self):
+        input_data = tf.constant(["hello", "fun", "▀▁▂▃"])
+        tokenizer = ByteTokenizer()
+        call_output = tokenizer(input_data)
+        tokenize_output = tokenizer.tokenize(input_data)
+        self.assertIsInstance(call_output, tf.RaggedTensor)
+        exp_outputs = [
+            [104, 101, 108, 108, 111],
+            [102, 117, 110],
+            [226, 150, 128, 226, 150, 129, 226, 150, 130, 226, 150, 131],
+        ]
+        for i in range(call_output.shape[0]):
+            self.assertAllEqual(call_output[i], exp_outputs[i])
+            self.assertAllEqual(tokenize_output[i], exp_outputs[i])
+
+    def test_tokenize_scalar(self):
+        input_data = "hello"
+        tokenizer = ByteTokenizer()
+        call_output = tokenizer(input_data)
+        tokenize_output = tokenizer.tokenize(input_data)
+
+        self.assertAllEqual(call_output, [104, 101, 108, 108, 111])
+        self.assertAllEqual(tokenize_output, [104, 101, 108, 108, 111])
+
+    def test_dense_output(self):
+        input_data = tf.constant(["hello", "fun", "▀▁▂▃"])
+        tokenizer = ByteTokenizer(sequence_length=10)
+        call_output = tokenizer(input_data)
+        self.assertIsInstance(call_output, tf.Tensor)
+        self.assertAllEqual(
+            call_output,
+            [
+                [104, 101, 108, 108, 111, 0, 0, 0, 0, 0],
+                [102, 117, 110, 0, 0, 0, 0, 0, 0, 0],
+                [226, 150, 128, 226, 150, 129, 226, 150, 130, 226],
+            ],
+        )
+
+    def test_detokenize(self):
+        input_data = tf.ragged.constant(
+            [
+                [104, 101, 108, 108, 111],
+                [102, 117, 110],
+                [226, 150, 128, 226, 150, 129, 226, 150, 130, 226, 150, 131],
+            ]
+        )
+
+        tokenizer = ByteTokenizer()
+        detokenize_output = tokenizer.detokenize(input_data)
+        self.assertAllEqual(detokenize_output, ["hello", "fun", "▀▁▂▃"])
+
+    def test_detokenize_replace_error(self):
+        # 226 is an invalid UTF-8 byte.
+        input_data = tf.ragged.constant([[104, 101, 226, 150, 108, 108, 111]])
+
+        tokenizer = ByteTokenizer(errors="replace", replacement_char=341)
+        detokenize_output = tokenizer.detokenize(input_data)
+        self.assertAllEqual(detokenize_output, [b"he\xc5\x95llo"])
+
+    def test_detokenize_ignore_error(self):
+        input_data = tf.ragged.constant([[104, 101, 226, 150, 108, 108, 111]])
+
+        tokenizer = ByteTokenizer(errors="ignore")
+        detokenize_output = tokenizer.detokenize(input_data)
+        self.assertAllEqual(detokenize_output, [b"hello"])
+
+    def test_detokenize_strict_error(self):
+        input_data = tf.ragged.constant([[104, 101, 226, 150, 108, 108, 111]])
+
+        tokenizer = ByteTokenizer(errors="strict")
+        with self.assertRaises(tf.errors.InvalidArgumentError):
+            _ = tokenizer.detokenize(input_data)
+
+    def test_vocab_size(self):
+        tokenizer = ByteTokenizer()
+        self.assertEqual(tokenizer.vocabulary_size(), 256)
+
+    def test_lowercase(self):
+        input_data = tf.constant(["HeLlO wOrLd"])
+        tokenizer = ByteTokenizer()
+        call_output = tokenizer(input_data)
+        self.assertAllEqual(
+            call_output,
+            [[104, 101, 108, 108, 111, 32, 119, 111, 114, 108, 100]],
+        )
+
+    def test_skip_lowercase(self):
+        input_data = tf.constant(["HeLlO wOrLd"])
+        tokenizer = ByteTokenizer(lowercase=False)
+        call_output = tokenizer(input_data)
+        self.assertAllEqual(
+            call_output, [[72, 101, 76, 108, 79, 32, 119, 79, 114, 76, 100]]
+        )
+
+    def test_tokenize_first_batch_second(self):
+        tokenizer = ByteTokenizer()
+
+        ds = tf.data.Dataset.from_tensor_slices(
+            ["hello", "fun", "▀▁▂▃", "haha"]
+        )
+        ds = ds.map(tokenizer)
+        ds = ds.apply(tf.data.experimental.dense_to_ragged_batch(4))
+        output = ds.take(1).get_single_element()
+
+        exp_output = [
+            [104, 101, 108, 108, 111],
+            [102, 117, 110],
+            [226, 150, 128, 226, 150, 129, 226, 150, 130, 226, 150, 131],
+            [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):
+        tokenizer = ByteTokenizer(sequence_length=10)
+
+        ds = tf.data.Dataset.from_tensor_slices(
+            ["hello", "fun", "▀▁▂▃", "haha"]
+        )
+        ds = ds.map(tokenizer)
+        ds = ds.apply(tf.data.experimental.dense_to_ragged_batch(4))
+        output = ds.take(1).get_single_element()
+
+        exp_output = [
+            [104, 101, 108, 108, 111, 0, 0, 0, 0, 0],
+            [102, 117, 110, 0, 0, 0, 0, 0, 0, 0],
+            [226, 150, 128, 226, 150, 129, 226, 150, 130, 226],
+            [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):
+        tokenizer = ByteTokenizer()
+
+        ds = tf.data.Dataset.from_tensor_slices(
+            ["hello", "fun", "▀▁▂▃", "haha"]
+        )
+        ds = ds.batch(4).map(tokenizer)
+        output = ds.take(1).get_single_element()
+
+        exp_output = [
+            [104, 101, 108, 108, 111],
+            [102, 117, 110],
+            [226, 150, 128, 226, 150, 129, 226, 150, 130, 226, 150, 131],
+            [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):
+        tokenizer = ByteTokenizer(sequence_length=10)
+
+        ds = tf.data.Dataset.from_tensor_slices(
+            ["hello", "fun", "▀▁▂▃", "haha"]
+        )
+        ds = ds.batch(4).map(tokenizer)
+        output = ds.take(1).get_single_element()
+
+        exp_output = [
+            [104, 101, 108, 108, 111, 0, 0, 0, 0, 0],
+            [102, 117, 110, 0, 0, 0, 0, 0, 0, 0],
+            [226, 150, 128, 226, 150, 129, 226, 150, 130, 226],
+            [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):
+        input_data = tf.constant(["hello", "fun", "▀▁▂▃"])
+        tokenizer = ByteTokenizer()
+        inputs = keras.Input(dtype="string", shape=())
+        outputs = tokenizer.detokenize(tokenizer.tokenize(inputs))
+        model = keras.Model(inputs, outputs)
+        model_output = model(input_data)
+        self.assertAllEqual(model_output, ["hello", "fun", "▀▁▂▃"])
+
+    def test_load_model_with_config(self):
+        input_data = tf.constant(["hello"])
+
+        original_tokenizer = ByteTokenizer(
+            lowercase=False,
+            sequence_length=8,
+            normalization_form="NFC",
+            errors="ignore",
+        )
+        cloned_tokenizer = ByteTokenizer.from_config(
+            original_tokenizer.get_config()
+        )
+        self.assertAllEqual(
+            original_tokenizer(input_data),
+            cloned_tokenizer(input_data),
+        )
+
+        decoded_input = [[104, 101, 226, 150, 108, 108, 111]]
+        self.assertAllEqual(
+            original_tokenizer.detokenize(decoded_input),
+            cloned_tokenizer.detokenize(decoded_input),
+        )
+
+    def test_config(self):
+
+        tokenizer = ByteTokenizer(
+            name="byte_tokenizer_config_test",
+            lowercase=False,
+            sequence_length=8,
+            normalization_form="NFC",
+            errors="ignore",
+            replacement_char=0,
+        )
+        exp_config = {
+            "dtype": "int32",
+            "errors": "ignore",
+            "lowercase": False,
+            "name": "byte_tokenizer_config_test",
+            "normalization_form": "NFC",
+            "replacement_char": 0,
+            "sequence_length": 8,
+            "trainable": True,
+        }
+        self.assertEqual(tokenizer.get_config(), exp_config)
+
+    def test_saving(self):
+        input_data = tf.constant(["this is fun"])
+
+        tokenizer = ByteTokenizer(
+            name="byte_tokenizer_config_test",
+            lowercase=False,
+            sequence_length=20,
+            normalization_form="NFKC",
+            errors="replace",
+        )
+        inputs = keras.Input(dtype="string", shape=())
+        outputs = tokenizer(inputs)
+        model = keras.Model(inputs, outputs)
+        model.save(self.get_temp_dir())
+        restored_model = keras.models.load_model(self.get_temp_dir())
+        self.assertAllEqual(
+            model(input_data),
+            restored_model(input_data),
+        )