Repository: keras-team/keras-nlp
Pull request #132 (2022-04-19 16:25:11)

Add a MLMHead layer
Fixes https://github.com/keras-team/keras-nlp/issues/129


Patch:
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index 4ed655e0..fa08c85f 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -26,12 +26,14 @@ development environment and run the unit tests. This is covered in section
 ### Step 3. Create a pull request
 
 Once the change is ready, open a pull request from your branch in your fork to
-the master branch in [keras-team/keras-nlp](https://github.com/keras-team/keras-nlp).
+the master branch in 
+[keras-team/keras-nlp](https://github.com/keras-team/keras-nlp).
 
 ### Step 4. Sign the Contributor License Agreement
 
 After creating the pull request, you will need to sign the Google CLA agreement.
-The agreement can be foiund at [https://cla.developers.google.com/clas](https://cla.developers.google.com/clas).
+The agreement can be found at
+[https://cla.developers.google.com/clas](https://cla.developers.google.com/clas).
 
 ### Step 5. Code review
 
@@ -47,20 +49,24 @@ Once the pull request is approved, a team member will take care of merging.
 
 ## Setup environment
 
+Python 3.7 or later is required.
+
 Setting up your KerasNLP development environment requires you to fork the
-KerasNLP repository, clone the repository, install dependencies, and execute
-`python setup.py develop`.
+KerasNLP repository, clone the repository, create a virtual environment, and 
+install dependencies.
 
 You can achieve this by running the following commands:
 
 ```shell
 gh repo fork keras-team/keras-nlp --clone --remote
 cd keras-nlp
-pip install ".[tests]"
-python setup.py develop
+python -m venv ~/keras-nlp-venv
+source ~/keras-nlp-venv/bin/activate
+pip install -e ".[tests]"
 ```
 
-The first line relies on having an installation of [the GitHub CLI](https://github.com/cli/cli).
+The first line relies on having an installation of
+[the GitHub CLI](https://github.com/cli/cli).
 
 Following these commands you should be able to run the tests using
 `pytest keras_nlp`. Please report any issues running tests following these
@@ -114,4 +120,5 @@ like - `: invalid option`
 
 ## Community Guidelines
 
-This project follows [Google's Open Source Community Guidelines](https://opensource.google/conduct/).
+This project follows 
+[Google's Open Source Community Guidelines](https://opensource.google/conduct/).
diff --git a/keras_nlp/layers/__init__.py b/keras_nlp/layers/__init__.py
index 4d1f5ddf..719b009e 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.mlm_head import MLMHead
 from keras_nlp.layers.position_embedding import PositionEmbedding
 from keras_nlp.layers.preprocessing import MLMMaskGenerator
 from keras_nlp.layers.sine_position_encoding import SinePositionEncoding
diff --git a/keras_nlp/layers/mlm_head.py b/keras_nlp/layers/mlm_head.py
new file mode 100644
index 00000000..09644205
--- /dev/null
+++ b/keras_nlp/layers/mlm_head.py
@@ -0,0 +1,213 @@
+# 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.
+
+"""Masked Language Model (MLM) head."""
+
+import tensorflow as tf
+from tensorflow import keras
+
+
+class MLMHead(keras.layers.Layer):
+    """Masked Language Model (MLM) head.
+
+    This layer takes two inputs:
+     - `inputs`: which should be a tensor of encoded tokens with shape
+            `(batch_size, sequence_length, encoding_dim)`.
+     - `mask_positions`: which should be a tensor of integer positions to
+            predict with shape `(batch_size, masks_per_sequence)`.
+
+    The token encodings should usually be the last output of an encoder model,
+    and mask positions should be the interger positions you would like to
+    predict for the MLM task.
+
+    The layer will first gather the token encodings at the mask positions. These
+    gathered tokens will be passed through a dense layer the same size as
+    encoding dimension, then transformed to predictions the same size as the
+    input vocabulary. This layer will produce a single output with shape
+    `(batch_size, masks_per_sequence, vocabulary_size)`, which can be used to
+    compute an MLM loss function.
+
+    This layer is often be paired with `keras_nlp.layers.MLMMaskGenerator`,
+    which will help prepare inputs for the MLM task.
+
+    Args:
+        vocabulary_size: The total size of the vocabulary for predictions.
+        embedding_weights: Optional. The weights of the word embedding used
+            to transform input token ids. The transpose of this weight matrix
+            will be used to project a token embedding vector to a prediction
+            over all input words, as described in [1].
+        intermediate_activation: The activation function of inner dense layer.
+        activation: The activation function for the outputs of the layer.
+            Usually either `None` (return logits), or `"softmax"`
+            (return probabilities).
+        layer_norm_epsilon: float, defaults to 1e-5. The epsilon value in layer
+            normalization components.
+        kernel_initializer: string or tf.keras.initializers initializer,
+            defaults to "glorot_uniform". The kernel initializer for
+            the dense and multiheaded attention layers.
+        bias_initializer: string or tf.keras.initializers initializer,
+            defaults to "zeros". The bias initializer for
+            the dense and multiheaded attention layers.
+        name: string, defaults to None. The name of the layer.
+        **kwargs: other keyword arguments.
+
+    Examples:
+
+    ```python
+    batch_size = 32
+    vocab_size = 100
+    encoding_size = 32
+    seq_length = 50
+    mask_length = 10
+
+    # Generate a random encoding.
+    encoded_tokens = tf.random.normal([batch_size, seq_length, encoding_size])
+    # Generate random positions and labels
+    mask_positions = tf.random.uniform(
+        [batch_size, mask_length], maxval=seq_length, dtype="int32"
+    )
+    mask_ids = tf.random.uniform(
+        [batch_size, mask_length], maxval=vocab_size, dtype="int32"
+    )
+
+    # Predict an output word for each masked input token.
+    mask_preds = keras_nlp.layers.MLMHead(
+        vocabulary_size=vocab_size,
+        activation="softmax",
+    )(encoded_tokens, mask_positions=mask_positions)
+    # Calculate a loss.
+    keras.losses.sparse_categorical_crossentropy(mask_ids, mask_preds)
+    ```
+
+    References:
+        [1] [Press and Wolf, 2016](https://arxiv.org/abs/1608.05859)
+    """
+
+    def __init__(
+        self,
+        vocabulary_size=None,
+        embedding_weights=None,
+        intermediate_activation="relu",
+        activation=None,
+        layer_norm_epsilon=1e-05,
+        kernel_initializer="glorot_uniform",
+        bias_initializer="zeros",
+        name=None,
+        **kwargs,
+    ):
+        super().__init__(name=name, **kwargs)
+
+        self.vocabulary_size = vocabulary_size
+        self.embedding_weights = embedding_weights
+        self.intermediate_activation = keras.activations.get(
+            intermediate_activation
+        )
+        self.activation = keras.activations.get(activation)
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.kernel_initializer = keras.initializers.get(kernel_initializer)
+        self.bias_initializer = keras.initializers.get(bias_initializer)
+        self._built = False
+
+        if vocabulary_size is None and embedding_weights is None:
+            raise ValueError(
+                "One of `vocabulary_size` or `embedding_weights` must be set. "
+                "Received: `vocabulary_size=None`, `embedding_weights=None`"
+            )
+
+        if embedding_weights is not None:
+            shape = embedding_weights.shape
+            if vocabulary_size is not None and vocabulary_size != shape[0]:
+                raise ValueError(
+                    "`vocabulary_size` should match the first dimension of the "
+                    "shape of `embedding_weights`. Received: "
+                    f"`vocabulary_size={vocabulary_size}`, "
+                    f"`embedding_weights.shape={shape}`"
+                )
+            self.vocabulary_size = shape[0]
+
+    def _build(self, input_shape):
+        # Create layers based on input shape.
+        self._built = True
+        feature_size = input_shape[-1]
+
+        self._dense = keras.layers.Dense(
+            feature_size,
+            activation=self.intermediate_activation,
+            kernel_initializer=self.kernel_initializer,
+            bias_initializer=self.bias_initializer,
+        )
+        self._layer_norm = tf.keras.layers.LayerNormalization(
+            epsilon=self.layer_norm_epsilon,
+        )
+        if self.embedding_weights is None:
+            self._kernel = self.add_weight(
+                name="output_kernel",
+                shape=[feature_size, self.vocabulary_size],
+                initializer=self.kernel_initializer,
+                dtype=self.dtype,
+            )
+        self._bias = self.add_weight(
+            name="output_bias",
+            shape=[self.vocabulary_size],
+            initializer=self.bias_initializer,
+            dtype=self.dtype,
+        )
+
+    def call(self, inputs, mask_positions):
+        if not self._built:
+            self._build(inputs.shape)
+
+        # Gather the encoded tokens at the masked indices.
+        x = tf.gather(inputs, mask_positions, axis=1, batch_dims=1)
+
+        # Apply a trainable linear transformation and a layer norm.
+        x = self._dense(x)
+        x = self._layer_norm(x)
+
+        # Transform encodings to vocabulary_size predictions.
+        if self.embedding_weights is None:
+            outputs = tf.matmul(x, self._kernel)
+        else:
+            outputs = tf.matmul(
+                x,
+                tf.cast(self.embedding_weights, self.compute_dtype),
+                transpose_b=True,
+            )
+        outputs = outputs + self._bias
+
+        # Apply a final activation.
+        if self.activation is not None:
+            outputs = self.activation(outputs)
+
+        return outputs
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "vocabulary_size": self.vocabulary_size,
+                "intermediate_activation": keras.activations.serialize(
+                    self.intermediate_activation
+                ),
+                "activation": keras.activations.serialize(self.activation),
+                "layer_norm_epsilon": self.layer_norm_epsilon,
+                "kernel_initializer": keras.initializers.serialize(
+                    self.kernel_initializer
+                ),
+                "bias_initializer": keras.initializers.serialize(
+                    self.bias_initializer
+                ),
+            }
+        )
+        return config
diff --git a/keras_nlp/layers/token_and_position_embedding.py b/keras_nlp/layers/token_and_position_embedding.py
index 72654359..035ddd02 100644
--- a/keras_nlp/layers/token_and_position_embedding.py
+++ b/keras_nlp/layers/token_and_position_embedding.py
@@ -91,6 +91,7 @@ class TokenAndPositionEmbedding(keras.layers.Layer):
             max_length=max_length,
             initializer=embeddings_initializer,
         )
+        self.supports_masking = self.token_embedding.supports_masking
 
     def get_config(self):
         config = super().get_config()
@@ -112,3 +113,6 @@ class TokenAndPositionEmbedding(keras.layers.Layer):
         embedded_positions = self.position_embedding(embedded_tokens)
         outputs = embedded_tokens + embedded_positions
         return outputs
+
+    def compute_mask(self, inputs, mask=None):
+        return self.token_embedding.compute_mask(inputs, mask=mask)
diff --git a/keras_nlp/layers/transformer_decoder.py b/keras_nlp/layers/transformer_decoder.py
index 2df066ac..92d59c2e 100644
--- a/keras_nlp/layers/transformer_decoder.py
+++ b/keras_nlp/layers/transformer_decoder.py
@@ -94,6 +94,7 @@ class TransformerDecoder(keras.layers.Layer):
         self.kernel_initializer = keras.initializers.get(kernel_initializer)
         self.bias_initializer = keras.initializers.get(bias_initializer)
         self._built = False
+        self.supports_masking = True
 
     def _build(self, input_shape):
         # Create layers based on input shape.
@@ -117,9 +118,15 @@ class TransformerDecoder(keras.layers.Layer):
             bias_initializer=self.bias_initializer,
         )
 
-        self._decoder_attention_layernorm = keras.layers.LayerNormalization()
-        self._enc_dec_attention_layernorm = keras.layers.LayerNormalization()
-        self._feedforward_layernorm = keras.layers.LayerNormalization()
+        self._decoder_attention_layernorm = keras.layers.LayerNormalization(
+            epsilon=self.layer_norm_epsilon,
+        )
+        self._enc_dec_attention_layernorm = keras.layers.LayerNormalization(
+            epsilon=self.layer_norm_epsilon,
+        )
+        self._feedforward_layernorm = keras.layers.LayerNormalization(
+            epsilon=self.layer_norm_epsilon,
+        )
 
         self._self_attention_dropout = keras.layers.Dropout(rate=self.dropout)
         self._enc_dec_attentiondropout = keras.layers.Dropout(
diff --git a/keras_nlp/layers/transformer_encoder.py b/keras_nlp/layers/transformer_encoder.py
index ef057aeb..870228be 100644
--- a/keras_nlp/layers/transformer_encoder.py
+++ b/keras_nlp/layers/transformer_encoder.py
@@ -89,6 +89,7 @@ class TransformerEncoder(keras.layers.Layer):
         self.kernel_initializer = keras.initializers.get(kernel_initializer)
         self.bias_initializer = keras.initializers.get(bias_initializer)
         self._built = False
+        self.supports_masking = True
 
     def _build(self, input_shape):
         # Create layers based on input shape.
@@ -104,8 +105,12 @@ class TransformerEncoder(keras.layers.Layer):
             bias_initializer=self.bias_initializer,
         )
 
-        self._attention_layernorm = keras.layers.LayerNormalization()
-        self._feedforward_layernorm = keras.layers.LayerNormalization()
+        self._attention_layernorm = keras.layers.LayerNormalization(
+            epsilon=self.layer_norm_epsilon,
+        )
+        self._feedforward_layernorm = keras.layers.LayerNormalization(
+            epsilon=self.layer_norm_epsilon,
+        )
 
         self._attention_dropout = keras.layers.Dropout(rate=self.dropout)
 
diff --git a/keras_nlp/tokenizers/word_piece_tokenizer.py b/keras_nlp/tokenizers/word_piece_tokenizer.py
index 35f2bf81..900692d8 100644
--- a/keras_nlp/tokenizers/word_piece_tokenizer.py
+++ b/keras_nlp/tokenizers/word_piece_tokenizer.py
@@ -198,6 +198,8 @@ class WordPieceTokenizer(tokenizer.Tokenizer):
                 "Vocabulary must be an file path or list of terms. "
                 f"Received: vocabulary={vocabulary}"
             )
+        if oov_token is None:
+            raise ValueError("`oov_token` cannot be None.")
 
         self.sequence_length = sequence_length
         self.lowercase = lowercase
@@ -207,6 +209,15 @@ class WordPieceTokenizer(tokenizer.Tokenizer):
         self.suffix_indicator = suffix_indicator
         self.oov_token = oov_token
 
+        if oov_token not in self.vocabulary:
+            raise RuntimeError(
+                f'Cannot find `oov_token="{self.oov_token}"` in the '
+                "vocabulary.\n"
+                "You can either update the vocabulary to include "
+                f'`"{self.oov_token}"`, or pass a different value for '
+                "the `oov_token` argument when creating the tokenizer."
+            )
+
         self._fast_word_piece = tf_text.FastWordpieceTokenizer(
             vocab=self.vocabulary,
             token_out_type=self.compute_dtype,
diff --git a/shell/format.sh b/shell/format.sh
index 691a72ba..0cdf6c79 100755
--- a/shell/format.sh
+++ b/shell/format.sh
@@ -1,15 +1,16 @@
 #!/bin/bash -e
 
 base_dir=$(dirname $(dirname $0))
+targets="${base_dir}/*.py ${base_dir}/examples/ ${base_dir}/keras_nlp/"
 
-isort --sl ${base_dir}
-black --line-length 80 ${base_dir}
+isort --sp "${base_dir}/setup.cfg" --sl ${targets}
+black --line-length 80 ${targets}
 
-for i in $(find ${base_dir} -name '*.py'); do
+for i in $(find ${targets} -name '*.py'); do
   if ! grep -q Copyright $i; then
     echo $i
     cat shell/copyright.txt $i >$i.new && mv $i.new $i
   fi
 done
 
-flake8 ${base_dir}
+flake8 --config "${base_dir}/setup.cfg" ${targets}
diff --git a/shell/lint.sh b/shell/lint.sh
index f20cba65..5fdcb68b 100755
--- a/shell/lint.sh
+++ b/shell/lint.sh
@@ -1,23 +1,26 @@
 #!/bin/bash -e
 
 base_dir=$(dirname $(dirname $0))
+targets="${base_dir}/*.py ${base_dir}/examples/ ${base_dir}/keras_nlp/"
 
-isort --sl -c "${base_dir}"
+isort --sp "${base_dir}/setup.cfg" --sl -c ${targets}
 if ! [ $? -eq 0 ]; then
   echo "Please run \"./shell/format.sh\" to format the code."
   exit 1
 fi
-flake8 "${base_dir}"
+
+flake8 --config "${base_dir}/setup.cfg" ${targets}
 if ! [ $? -eq 0 ]; then
   echo "Please fix the code style issue."
   exit 1
 fi
-black --check --line-length 80 "${base_dir}"
+
+black --check --line-length 80 ${targets}
 if ! [ $? -eq 0 ]; then
   echo "Please run \"./shell/format.sh\" to format the code."
     exit 1
 fi
-for i in $(find "${base_dir}" -name '*.py'); do
+for i in $(find ${targets} -name '*.py'); do
   if ! grep -q Copyright $i; then
     echo "Please run \"./shell/format.sh\" to format the code."
     exit 1


Test patch:
diff --git a/keras_nlp/layers/mlm_head_test.py b/keras_nlp/layers/mlm_head_test.py
new file mode 100644
index 00000000..0271d915
--- /dev/null
+++ b/keras_nlp/layers/mlm_head_test.py
@@ -0,0 +1,179 @@
+# 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 Transformer Encoder."""
+
+import os
+
+import tensorflow as tf
+from tensorflow import keras
+
+from keras_nlp.layers import mlm_head
+
+
+class MLMHeadTest(tf.test.TestCase):
+    def test_valid_call(self):
+        head = mlm_head.MLMHead(
+            vocabulary_size=100,
+            activation="softmax",
+        )
+        encoded_tokens = keras.Input(shape=(10, 16))
+        positions = keras.Input(shape=(5,), dtype="int32")
+        outputs = head(encoded_tokens, mask_positions=positions)
+        model = keras.Model((encoded_tokens, positions), outputs)
+
+        token_data = tf.random.uniform(shape=(4, 10, 16))
+        position_data = tf.random.uniform(
+            shape=(4, 5), maxval=10, dtype="int32"
+        )
+        model((token_data, position_data))
+
+    def test_valid_call_with_embedding_weights(self):
+        embedding = keras.layers.Embedding(100, 16)
+        embedding.build((4, 10))
+        head = mlm_head.MLMHead(
+            vocabulary_size=100,
+            embedding_weights=embedding.embeddings,
+            activation="softmax",
+        )
+        encoded_tokens = keras.Input(shape=(10, 16))
+        positions = keras.Input(shape=(5,), dtype="int32")
+        outputs = head(encoded_tokens, mask_positions=positions)
+        model = keras.Model((encoded_tokens, positions), outputs)
+        token_data = tf.random.uniform(shape=(4, 10, 16))
+        position_data = tf.random.uniform(
+            shape=(4, 5), maxval=10, dtype="int32"
+        )
+        model((token_data, position_data))
+
+    def test_get_config_and_from_config(self):
+        head = mlm_head.MLMHead(
+            vocabulary_size=100,
+            kernel_initializer="HeNormal",
+            bias_initializer="Zeros",
+            activation="softmax",
+        )
+
+        config = head.get_config()
+
+        expected_params = {
+            "vocabulary_size": 100,
+            "kernel_initializer": keras.initializers.serialize(
+                keras.initializers.HeNormal()
+            ),
+            "bias_initializer": keras.initializers.serialize(
+                keras.initializers.Zeros()
+            ),
+            "activation": keras.activations.serialize(
+                keras.activations.softmax
+            ),
+        }
+
+        self.assertEqual(config, {**config, **expected_params})
+
+        restored = mlm_head.MLMHead.from_config(config)
+        restored_config = restored.get_config()
+
+        self.assertEqual(
+            restored_config, {**restored_config, **expected_params}
+        )
+        self.assertEqual(restored_config, config)
+
+    def test_value_error_when_neither_embedding_or_vocab_size_set(self):
+        with self.assertRaises(ValueError):
+            mlm_head.MLMHead()
+
+    def test_value_error_when_vocab_size_mismatch(self):
+        embedding = keras.layers.Embedding(100, 16)
+        embedding.build((4, 10))
+        with self.assertRaises(ValueError):
+            mlm_head.MLMHead(
+                vocabulary_size=101,
+                embedding_weights=embedding.embeddings,
+            )
+
+    def test_one_train_step(self):
+        head = mlm_head.MLMHead(
+            vocabulary_size=100,
+        )
+        encoded_tokens = keras.Input(shape=(10, 16))
+        positions = keras.Input(shape=(5,), dtype="int32")
+        outputs = head(encoded_tokens, mask_positions=positions)
+        model = keras.Model((encoded_tokens, positions), outputs)
+
+        token_data = tf.random.uniform(shape=(4, 10, 16))
+        position_data = tf.random.uniform(
+            shape=(4, 5), maxval=10, dtype="int32"
+        )
+        label_data = tf.random.uniform(shape=(4, 5), maxval=100, dtype="int32")
+
+        loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(
+            from_logits=False
+        )
+        optimizer = keras.optimizers.Adam()
+        with tf.GradientTape() as tape:
+            model((token_data, position_data))
+            pred = model((token_data, position_data))
+            loss = loss_fn(label_data, pred)
+        grad = tape.gradient(loss, model.trainable_variables)
+        self.assertGreater(len(grad), 1)
+        optimizer.apply_gradients(zip(grad, model.trainable_variables))
+
+    def test_checkpointing(self):
+        head1 = mlm_head.MLMHead(
+            vocabulary_size=100,
+            activation="softmax",
+        )
+        head2 = mlm_head.MLMHead(
+            vocabulary_size=100,
+            activation="softmax",
+        )
+        token_data = tf.random.uniform(shape=(4, 10, 16))
+        position_data = tf.random.uniform(
+            shape=(4, 5), maxval=10, dtype="int32"
+        )
+        # The weights of head1 and head2 are different.
+        head1_output = head1(token_data, mask_positions=position_data)
+        head2_output = head2(token_data, mask_positions=position_data)
+        self.assertNotAllClose(head1_output, head2_output)
+
+        checkpoint = tf.train.Checkpoint(head1)
+        checkpoint2 = tf.train.Checkpoint(head2)
+        save_path = checkpoint.save(self.get_temp_dir())
+        checkpoint2.restore(save_path)
+
+        head1_output = head1(token_data, mask_positions=position_data)
+        head2_output = head2(token_data, mask_positions=position_data)
+        self.assertAllClose(head1_output, head2_output)
+
+    def test_saving_model(self):
+        head = mlm_head.MLMHead(
+            vocabulary_size=100,
+            activation="softmax",
+        )
+        encoded_tokens = keras.Input(shape=(10, 16))
+        positions = keras.Input(shape=(5,), dtype="int32")
+        outputs = head(encoded_tokens, mask_positions=positions)
+        model = keras.Model((encoded_tokens, positions), outputs)
+
+        token_data = tf.random.uniform(shape=(4, 10, 16))
+        position_data = tf.random.uniform(
+            shape=(4, 5), maxval=10, dtype="int32"
+        )
+        model_output = model((token_data, position_data))
+        save_path = os.path.join(self.get_temp_dir(), "model")
+        model.save(save_path)
+        restored = keras.models.load_model(save_path)
+
+        restored_output = restored((token_data, position_data))
+        self.assertAllClose(model_output, restored_output)
diff --git a/keras_nlp/layers/token_and_position_embedding_test.py b/keras_nlp/layers/token_and_position_embedding_test.py
index 48149299..a9916ab9 100644
--- a/keras_nlp/layers/token_and_position_embedding_test.py
+++ b/keras_nlp/layers/token_and_position_embedding_test.py
@@ -112,51 +112,27 @@ class TokenAndPositionEmbeddingTest(tf.test.TestCase):
         )
         # Create a 2-dimensional input
         # (the first dimension is implicit).
-        input_tensor = tf.keras.Input(
-            shape=(sequence_length,), dtype=tf.float32, ragged=True
-        )
-        output_tensor = test_layer(input_tensor)
-        model = tf.keras.Model(input_tensor, output_tensor)
+        inputs = tf.keras.Input(shape=(sequence_length,), dtype="int32")
+        outputs = test_layer(inputs)
+        model = tf.keras.Model(inputs, outputs)
 
-        input_data = tf.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],
-                [1.0, 1.0, 1.0, 1.0],
-            ],
-        )
-        expected_output_data = tf.constant(
-            [
-                [
-                    [2.0, 2.0, 2.0],
-                    [2.0, 2.0, 2.0],
-                    [2.0, 2.0, 2.0],
-                    [2.0, 2.0, 2.0],
-                ],
-                [
-                    [2.0, 2.0, 2.0],
-                    [2.0, 2.0, 2.0],
-                    [2.0, 2.0, 2.0],
-                    [2.0, 2.0, 2.0],
-                ],
-                [
-                    [2.0, 2.0, 2.0],
-                    [2.0, 2.0, 2.0],
-                    [2.0, 2.0, 2.0],
-                    [2.0, 2.0, 2.0],
-                ],
-                [
-                    [2.0, 2.0, 2.0],
-                    [2.0, 2.0, 2.0],
-                    [2.0, 2.0, 2.0],
-                    [2.0, 2.0, 2.0],
-                ],
-            ],
-        )
+        input_data = tf.ones((2, sequence_length), dtype="int32")
+        expected_output_data = tf.ones((2, sequence_length, embedding_dim)) * 2
         output_data = model.predict(input_data)
         self.assertAllClose(output_data, expected_output_data)
 
+    def test_mask_propagation(self):
+        test_layer = TokenAndPositionEmbedding(
+            vocabulary_size=5,
+            max_length=4,
+            embedding_dim=3,
+            mask_zero=True,
+        )
+        input_data = tf.constant([[1, 0], [1, 0]])
+        mask = input_data != 0
+        outputs = test_layer(input_data)
+        self.assertAllEqual(outputs._keras_mask, mask)
+
     def test_save_model(self):
         vocabulary_size = 5
         sequence_length = 4
diff --git a/keras_nlp/layers/transformer_decoder_test.py b/keras_nlp/layers/transformer_decoder_test.py
index 6cbfd5fd..43460892 100644
--- a/keras_nlp/layers/transformer_decoder_test.py
+++ b/keras_nlp/layers/transformer_decoder_test.py
@@ -162,6 +162,18 @@ class TransformerDecoderTest(tf.test.TestCase):
         )
         self.assertAllClose(decoder1_output, decoder2_output)
 
+    def test_mask_propagation(self):
+        decoder = transformer_decoder.TransformerDecoder(
+            intermediate_dim=4,
+            num_heads=2,
+        )
+        decoder_sequence = tf.random.uniform(shape=[1, 4, 6])
+        encoder_sequence = tf.random.uniform(shape=[1, 4, 6])
+        mask = tf.constant([[True, True, False, False]])
+        decoder_sequence._keras_mask = mask
+        outputs = decoder(decoder_sequence, encoder_sequence)
+        self.assertAllEqual(outputs._keras_mask, mask)
+
     def test_save_model(self):
         encoder_input = keras.Input(shape=[4, 6])
         decoder_input = keras.Input(shape=[4, 6])
diff --git a/keras_nlp/layers/transformer_encoder_test.py b/keras_nlp/layers/transformer_encoder_test.py
index e1318989..8e5d13c9 100644
--- a/keras_nlp/layers/transformer_encoder_test.py
+++ b/keras_nlp/layers/transformer_encoder_test.py
@@ -111,6 +111,17 @@ class TransformerEncoderTest(tf.test.TestCase):
         self.assertGreater(len(grad), 1)
         optimizer.apply_gradients(zip(grad, model.trainable_variables))
 
+    def test_mask_propagation(self):
+        encoder = transformer_encoder.TransformerEncoder(
+            intermediate_dim=4,
+            num_heads=2,
+        )
+        inputs = tf.random.uniform(shape=[1, 4, 6])
+        mask = tf.constant([[True, True, False, False]])
+        inputs._keras_mask = mask
+        outputs = encoder(inputs)
+        self.assertAllEqual(outputs._keras_mask, mask)
+
     def test_checkpointing_transformer_encoder(self):
         encoder1 = transformer_encoder.TransformerEncoder(
             intermediate_dim=4,
diff --git a/keras_nlp/tokenizers/word_piece_tokenizer_test.py b/keras_nlp/tokenizers/word_piece_tokenizer_test.py
index 9a54f857..c7430ae3 100644
--- a/keras_nlp/tokenizers/word_piece_tokenizer_test.py
+++ b/keras_nlp/tokenizers/word_piece_tokenizer_test.py
@@ -223,3 +223,25 @@ class WordPieceTokenizerTest(tf.test.TestCase):
             model(input_data),
             restored_model(input_data),
         )
+
+    def test_no_oov_token_in_vocabulary(self):
+        vocab_data = ["qu", "@@ick", "br", "@@OWN", "fox"]
+        with self.assertRaises(RuntimeError):
+            WordPieceTokenizer(
+                vocabulary=vocab_data,
+            )
+
+        vocab_data = ["@UNK@", "qu", "@@ick", "br", "@@OWN", "fox"]
+        with self.assertRaises(RuntimeError):
+            WordPieceTokenizer(
+                vocabulary=vocab_data,
+            )
+
+        vocab_data = ["UNK", "qu", "@@ick", "br", "@@OWN", "fox"]
+        with self.assertRaises(RuntimeError):
+            WordPieceTokenizer(
+                vocabulary=vocab_data,
+            )
+
+        with self.assertRaises(ValueError):
+            WordPieceTokenizer(vocabulary=vocab_data, oov_token=None)