Repository: keras-team/keras-nlp
Pull request #116 (2022-04-14 21:48:22)

Add docstring testing
We will use the same docstring checker as Tensorflow and Keras. We may want to introduce a canonical tensorflow or keras version of this we can rely on. But in the interest of getting these tests up sooner rather than later, we will just duplicate it for now.

Fixes https://github.com/keras-team/keras-nlp/issues/82


Patch:
diff --git a/README.md b/README.md
index 461f2ec7..6d21abbf 100644
--- a/README.md
+++ b/README.md
@@ -26,13 +26,18 @@ If you'd like to contribute, please see our [contributing guide](CONTRIBUTING.md
 
 The fastest way to find a place to contribute is to browse our
 [open issues](https://github.com/keras-team/keras-nlp/issues) and find an
-unclaimed issue to work on. We tag good introductory issues with the
-`good first issue` label.
+unclaimed issue to work on. Issues with a [contributions welcome](
+https://github.com/keras-team/keras-nlp/issues?q=is%3Aissue+is%3Aopen+label%3A%22contributions+welcome%22)
+tag are places where we are actively looking for support, and a
+[good first issue](https://github.com/keras-team/keras-nlp/issues?q=is%3Aissue+is%3Aopen+label%3A%22good+first+issue%22)
+tag means we think this could be a accessible a first time contributor.
 
 If you would like to propose a new symbol or feature, please open an issue to
-discuss. Be aware the design for new features may take longer at this early
-stage. If you have a design in mind, please include a colab notebook showing the
-proposed design in a end-to-end example.
+discuss. Be aware the design for new features may take longer than contributing
+pre-planned features. If you have a design in mind, please include a colab
+notebook showing the proposed design in a end-to-end example. Make sure to
+follow the [Keras API design guidelines](
+https://github.com/keras-team/governance/blob/master/keras_api_design_guidelines.md).
 
 ## Roadmap
 
@@ -49,6 +54,19 @@ on components that follow a established Keras interface (e.g.
 As we progress further with the library, we will attempt to cover an ever
 expanding list of widely cited model architectures.
 
+## Releases
+
+KerasNLP release are documented on our
+[github release page](https://github.com/keras-team/keras-nlp/releases) and
+available to download from our [PyPI project](
+https://pypi.org/project/keras-nlp/).
+
+To install KerasNLP and all it's dependencies, simply run:
+
+```
+pip install keras-nlp
+```
+
 ## Compatibility
 
 We follow [Semantic Versioning](https://semver.org/), and plan to
diff --git a/examples/machine_translation/train.py b/examples/machine_translation/train.py
index c65125c6..648585bc 100644
--- a/examples/machine_translation/train.py
+++ b/examples/machine_translation/train.py
@@ -77,9 +77,12 @@ def run_training(model, train_ds, val_ds):
 
 
 def main(_):
-    (train_ds, val_ds, test_ds), (
-        eng_tokenizer,
-        spa_tokenizer,
+    (
+        (train_ds, val_ds, test_ds),
+        (
+            eng_tokenizer,
+            spa_tokenizer,
+        ),
     ) = get_dataset_and_tokenizer(
         FLAGS.sequence_length, FLAGS.vocab_size, FLAGS.batch_size
     )
diff --git a/keras_nlp/layers/__init__.py b/keras_nlp/layers/__init__.py
index fe677bbe..9e246621 100644
--- a/keras_nlp/layers/__init__.py
+++ b/keras_nlp/layers/__init__.py
@@ -14,6 +14,7 @@
 
 from keras_nlp.layers.fnet_encoder import FNetEncoder
 from keras_nlp.layers.position_embedding import PositionEmbedding
+from keras_nlp.layers.preprocessing import MLMMaskGenerator
 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/preprocessing/__init__.py b/keras_nlp/layers/preprocessing/__init__.py
new file mode 100644
index 00000000..c8059499
--- /dev/null
+++ b/keras_nlp/layers/preprocessing/__init__.py
@@ -0,0 +1,15 @@
+# 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.
+
+from keras_nlp.layers.preprocessing.mlm_mask_generator import MLMMaskGenerator
diff --git a/keras_nlp/layers/preprocessing/mlm_mask_generator.py b/keras_nlp/layers/preprocessing/mlm_mask_generator.py
new file mode 100644
index 00000000..59daa236
--- /dev/null
+++ b/keras_nlp/layers/preprocessing/mlm_mask_generator.py
@@ -0,0 +1,202 @@
+# 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
+import tensorflow_text as tf_text
+from tensorflow import keras
+
+
+class MLMMaskGenerator(keras.layers.Layer):
+    """Layer that applies language model masking.
+
+    This layer is useful for preparing inputs for masked languaged modeling
+    (MLM) tasks. It follows the masking strategy described in the [original BERT
+    paper](https://arxiv.org/abs/1810.04805). Given tokenized text,
+    it randomly selects certain number of tokens for masking. Then for each
+    selected token, it has a chance (configurable) to be replaced by
+    "mask token" or random token, or stay unchanged.
+
+    Users should use this layer with `tf.data` to generate masks.
+
+    Args:
+        vocabulary_size: int, the size of the vocabulary.
+        mask_selection_rate: float, the probability of a token is selected for
+            masking.
+        mask_token_id: int. The id of mask token.
+        mask_selection_length: int, defaults to None. Maximum number of tokens
+            selected for  masking in each sequence. If set, the output
+            `mask_positions`, `mask_ids` and `mask_weights` will be padded
+            to dense tensors of length `mask_selection_length`,
+            otherwise the output will be a RaggedTensor.
+        unselectable_token_ids: A list of tokens, defaults to [0] (the default
+            `padding_token_id`).
+        mask_token_rate: float, defaults to 0.8. `mask_token_rate` must be
+            between 0 and 1 which indicates how often the mask_token is
+            substituted for tokens selected for masking.
+        random_token_rate: float, defaults to 0.1. `random_token_rate` must be
+            between 0 and 1 which indicates how often a random token is
+            substituted for tokens selected for masking. Default is 0.1.
+            Note: mask_token_rate + random_token_rate <= 1,  and for
+            (1 - mask_token_rate - random_token_rate), the token will not be
+            changed.
+
+    Input:
+        A 1D integer tensor of shape [sequence_length] or a 2D integer tensor
+        of shape [batch_size, sequence_length], or a 2D integer RaggedTensor.
+        Represents the sequence to mask.
+
+    Returns:
+        A Dict with 4 keys:
+            tokens: Tensor or RaggedTensor, has the same type and shape of
+                input. Sequence after getting masked.
+            mask_positions: Tensor, or RaggedTensor if `mask_selection_length`
+                is None. The positions of tokens getting masked.
+            mask_ids: Tensor, or RaggedTensor if  `mask_selection_length` is
+                None. The original token ids at masked positions.
+            mask_weights: Tensor, or RaggedTensor if `mask_selection_length` is
+                None. `mask_weights` has the same shape as `mask_positions` and
+                `mask_ids`. Each element in `mask_weights` should be 0 or 1,
+                1 means the corresponding position in `mask_positions` is an
+                actual mask, 0 means it is a pad.
+
+    Examples:
+
+    Basic usage.
+    >>> masker = keras_nlp.preprocessing.MLMMaskGenerator( \
+    ...     vocaublary_size=10, mask_selection_rate=0.2, mask_token_id=0)
+    >>> masker(tf.constant([1, 2, 3, 4, 5]))
+    {'tokens': <tf.Tensor: shape=(5,), dtype=int32,
+        numpy=array([1, 2, 3, 4, 0], dtype=int32)>,
+    'mask_positions': <tf.Tensor: shape=(5,), dtype=int64,
+        numpy=array([4, 0, 0, 0, 0])>,
+    'mask_ids': <tf.Tensor: shape=(5,), dtype=int32,
+        numpy=array([5, 0, 0, 0, 0], dtype=int32)>,
+    'mask_weights': <tf.Tensor: shape=(1, 5), dtype=int64,
+        numpy=array([[1, 0, 0, 0, 0]])>}
+
+    Ragged Input:
+    >>> masker = keras_nlp.preprocessing.MLMMaskGenerator( \
+    ...     vocaublary_size=10, mask_selection_rate=0.2, mask_token_id=0)
+    >>> masker(tf.ragged.constant([[1, 2], [1, 2, 3, 4]]))
+    {'tokens': <tf.RaggedTensor [[1, 4], [0, 2, 3, 0]]>,
+    'mask_positions': <tf.Tensor: shape=(2, 5), dtype=int64, numpy=
+        array([[1, 0, 0, 0, 0],
+               [0, 3, 0, 0, 0]])>,
+    'mask_ids': <tf.Tensor: shape=(2, 5), dtype=int32, numpy=
+        array([[2, 0, 0, 0, 0],
+               [1, 4, 0, 0, 0]], dtype=int32)>,
+    'mask_weights': <tf.Tensor: shape=(2, 5), dtype=int64, numpy=
+        array([[1, 0, 0, 0, 0],
+               [1, 1, 0, 0, 0]])>}
+    """
+
+    def __init__(
+        self,
+        vocabulary_size,
+        mask_selection_rate,
+        mask_token_id,
+        mask_selection_length=None,
+        unselectable_token_ids=[0],
+        mask_token_rate=0.8,
+        random_token_rate=0.1,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.vocabulary_size = vocabulary_size
+        self.unselectable_token_ids = unselectable_token_ids
+        self.mask_selection_rate = mask_selection_rate
+        self.mask_selection_length = mask_selection_length
+        self.mask_token_rate = mask_token_rate
+        self.random_token_rate = random_token_rate
+
+        if mask_token_id >= vocabulary_size:
+            raise ValueError(
+                f"Mask token id should be in range [0, vocabulary_size - 1], "
+                f"but received mask_token_id={mask_token_id}."
+            )
+        self.mask_token_id = mask_token_id
+
+        max_selections = self.mask_selection_length
+        if max_selections is None:
+            # Set a large number to remove the `max_selections_per_batch` cap.
+            max_selections = 2**31 - 1
+        self._random_selector = tf_text.RandomItemSelector(
+            max_selections_per_batch=max_selections,
+            selection_rate=self.mask_selection_rate,
+            unselectable_ids=self.unselectable_token_ids,
+        )
+        self._mask_values_chooser = tf_text.MaskValuesChooser(
+            self.vocabulary_size,
+            self.mask_token_id,
+            mask_token_rate=self.mask_token_rate,
+            random_token_rate=self.random_token_rate,
+        )
+
+    def call(self, inputs):
+        input_is_ragged = isinstance(inputs, tf.RaggedTensor)
+        input_is_1d = tf.rank(inputs) == 1
+        if input_is_1d:
+            # If inputs is of rank 1, we manually add the batch axis.
+            inputs = inputs[tf.newaxis, :]
+        if not input_is_ragged:
+            # `tf_text.mask_language_model` requires a ragged tensor, so
+            # convert dense to ragged.
+            inputs = tf.RaggedTensor.from_tensor(inputs)
+        (tokens, mask_positions, mask_ids,) = tf_text.mask_language_model(
+            inputs,
+            item_selector=self._random_selector,
+            mask_values_chooser=self._mask_values_chooser,
+        )
+
+        if not input_is_ragged:
+            # If we converted the input from dense to ragged, convert back.
+            tokens = tokens.to_tensor()
+
+        mask_weights = tf.ones_like(mask_positions, self.compute_dtype)
+        # If mask_selection_length is set, covert to raggeds to dense.
+        if self.mask_selection_length:
+            target_shape = tf.cast([-1, self.mask_selection_length], tf.int64)
+            mask_positions = mask_positions.to_tensor(shape=target_shape)
+            mask_ids = mask_ids.to_tensor(shape=target_shape)
+            mask_weights = mask_weights.to_tensor(shape=target_shape)
+
+        if input_is_1d:
+            # If inputs is 1D, we format the output to be 1D as well.
+            tokens = tf.squeeze(tokens, axis=0)
+            mask_positions = tf.squeeze(mask_positions, axis=0)
+            mask_ids = tf.squeeze(mask_ids, axis=0)
+            mask_weights = tf.squeeze(mask_weights, axis=0)
+
+        output_dict = {
+            "tokens": tokens,
+            "mask_positions": mask_positions,
+            "mask_ids": mask_ids,
+            "mask_weights": mask_weights,
+        }
+        return output_dict
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "vocabulary_size": self.vocabulary_size,
+                "mask_selection_rate": self.mask_selection_rate,
+                "mask_selection_length": self.mask_selection_length,
+                "unselectable_token_ids": self.unselectable_token_ids,
+                "mask_token_id": self.mask_token_id,
+                "mask_token_rate": self.mask_token_rate,
+                "random_token_rate": self.random_token_rate,
+            }
+        )
+        return config
diff --git a/keras_nlp/tokenizers/byte_tokenizer.py b/keras_nlp/tokenizers/byte_tokenizer.py
index 2f142004..033699d3 100644
--- a/keras_nlp/tokenizers/byte_tokenizer.py
+++ b/keras_nlp/tokenizers/byte_tokenizer.py
@@ -86,8 +86,8 @@ class ByteTokenizer(tokenizer.Tokenizer):
     >>> 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]])>
+    <tf.Tensor: shape=(1, 8), dtype=int32, numpy=
+    array([[104, 101, 108, 108, 111,   0,   0,   0]], dtype=int32)>
 
     Tokenize first, then batch the dataset up.
     >>> tokenizer = keras_nlp.tokenizers.ByteTokenizer()
@@ -112,7 +112,7 @@ class ByteTokenizer(tokenizer.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]])>
+           [102, 117, 110,   0,   0]], dtype=int32)>
 
     Batch up the inputs and then tokenize (`sequence_length` provided).
     >>> tokenizer = keras_nlp.tokenizers.ByteTokenizer(sequence_length=5)
@@ -121,7 +121,7 @@ class ByteTokenizer(tokenizer.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]])>
+           [102, 117, 110,   0,   0]], dtype=int32)>
 
     Detokenization.
     >>> inputs = tf.constant([104, 101, 108, 108, 111], dtype=tf.int32)
diff --git a/keras_nlp/tokenizers/word_piece_tokenizer.py b/keras_nlp/tokenizers/word_piece_tokenizer.py
index 5350aff8..35f2bf81 100644
--- a/keras_nlp/tokenizers/word_piece_tokenizer.py
+++ b/keras_nlp/tokenizers/word_piece_tokenizer.py
@@ -132,8 +132,8 @@ class WordPieceTokenizer(tokenizer.Tokenizer):
     >>> tokenizer = keras_nlp.tokenizers.WordPieceTokenizer(
     ...     vocabulary=vocab, sequence_length=10)
     >>> tokenizer(inputs)
-    <tf.Tensor: shape=(1, 10), dtype=int32,
-        numpy=array([[1, 2, 3, 4, 5, 6, 7, 0, 0, 0]]>
+    <tf.Tensor: shape=(1, 10), dtype=int32, numpy=
+    array([[1, 2, 3, 4, 5, 6, 7, 0, 0, 0]], dtype=int32)>
 
     String output.
     >>> vocab = ["[UNK]", "the", "qu", "##ick", "br", "##own", "fox", "."]
@@ -145,16 +145,15 @@ class WordPieceTokenizer(tokenizer.Tokenizer):
 
     Detokenization.
     >>> vocab = ["[UNK]", "the", "qu", "##ick", "br", "##own", "fox", "."]
-    >>> inputs = ["The quick brown fox."]
+    >>> inputs = "The quick brown fox."
     >>> tokenizer = keras_nlp.tokenizers.WordPieceTokenizer(vocabulary=vocab)
-    >>> tokenizer.detokenize(tokenizer.tokenize(inputs))
-    <tf.Tensor: shape=(1,), dtype=string,
-        numpy=array([b"the quick brown fox ."], dtype=object)>
+    >>> tokenizer.detokenize(tokenizer.tokenize(inputs)).numpy().decode('utf-8')
+    'the quick brown fox .'
 
     Custom splitting.
     >>> vocab = ["[UNK]", "fox", ","]
     >>> inputs = ["fox,,fox,fox"]
-    >>> tokenizer = keras_nlp.tokenizers.WordPieceTokenizer(vocabulary=vocab,
+    >>> keras_nlp.tokenizers.WordPieceTokenizer(vocabulary=vocab,
     ...     split_pattern=",", keep_pattern=",", dtype='string')(inputs)
     <tf.RaggedTensor [[b'fox', b',', b',', b'fox', b',', b'fox']]>
     >>> keras_nlp.tokenizers.WordPieceTokenizer(vocabulary=vocab,


Test patch:
diff --git a/keras_nlp/layers/preprocessing/mlm_mask_generator_test.py b/keras_nlp/layers/preprocessing/mlm_mask_generator_test.py
new file mode 100644
index 00000000..66547984
--- /dev/null
+++ b/keras_nlp/layers/preprocessing/mlm_mask_generator_test.py
@@ -0,0 +1,234 @@
+# 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 keras_nlp.layers.preprocessing.mlm_mask_generator import MLMMaskGenerator
+
+
+class MLMMaskGeneratorTest(tf.test.TestCase):
+    def setUp(self):
+        super().setUp()
+        self.VOCAB = [
+            "[UNK]",
+            "[MASK]",
+            "[RANDOM]",
+            "[CLS]",
+            "[SEP]",
+            "do",
+            "you",
+            "like",
+            "machine",
+            "learning",
+            "welcome",
+            "to",
+            "keras",
+        ]
+        self.mask_token_id = self.VOCAB.index("[MASK]")
+        self.vocabulary_size = len(self.VOCAB)
+
+    def test_mask_ragged_tensor(self):
+        mlm_masker = MLMMaskGenerator(
+            vocabulary_size=self.vocabulary_size,
+            mask_selection_rate=0.5,
+            mask_selection_length=5,
+            mask_token_id=self.mask_token_id,
+            mask_token_rate=1,
+            random_token_rate=0,
+        )
+        inputs = tf.ragged.constant([[5, 3, 2], [1, 2, 3, 4, 5]])
+        outputs = mlm_masker(inputs)
+        tokens, mask_positions, mask_ids = (
+            outputs["tokens"],
+            outputs["mask_positions"],
+            outputs["mask_ids"],
+        )
+        self.assertEqual(type(tokens), type(inputs))
+        self.assertAllEqual(tokens.shape, inputs.shape)
+        self.assertAllEqual(mask_positions.shape, mask_ids.shape)
+
+        # Test all selected tokens are correctly masked.
+        masked_values = tf.gather(
+            tokens,
+            mask_positions,
+            batch_dims=1,
+        )
+        self.assertEqual(tf.reduce_mean(masked_values), self.mask_token_id)
+
+    def test_mask_tensor(self):
+        mlm_masker = MLMMaskGenerator(
+            vocabulary_size=self.vocabulary_size,
+            mask_selection_rate=0.5,
+            mask_selection_length=5,
+            mask_token_id=self.mask_token_id,
+            mask_token_rate=1,
+            random_token_rate=0,
+        )
+        inputs = tf.random.uniform(
+            shape=[5, 10],
+            maxval=len(self.VOCAB),
+            dtype=tf.int32,
+        )
+        outputs = mlm_masker(inputs)
+        tokens, mask_positions, mask_ids = (
+            outputs["tokens"],
+            outputs["mask_positions"],
+            outputs["mask_ids"],
+        )
+        self.assertEqual(type(tokens), type(inputs))
+        self.assertAllEqual(tokens.shape, inputs.shape)
+        self.assertAllEqual(mask_positions.shape, mask_ids.shape)
+        # Test all selected tokens are correctly masked.
+        masked_values = tf.gather(
+            tokens,
+            mask_positions,
+            batch_dims=1,
+        )
+        self.assertEqual(tf.reduce_mean(masked_values), self.mask_token_id)
+
+    def test_mask_1d_input(self):
+        mlm_masker = MLMMaskGenerator(
+            vocabulary_size=self.vocabulary_size,
+            mask_selection_rate=0.5,
+            mask_selection_length=5,
+            mask_token_id=self.mask_token_id,
+            mask_token_rate=1,
+            random_token_rate=0,
+        )
+        inputs = tf.constant([1, 2, 3, 4, 5])
+        outputs = mlm_masker(inputs)
+        self.assertAllEqual(outputs["tokens"].shape, inputs.shape)
+
+    def test_number_of_masked_position_as_expected(self):
+        mask_selection_rate = 0.5
+        mask_selection_length = 5
+        mlm_masker = MLMMaskGenerator(
+            vocabulary_size=self.vocabulary_size,
+            mask_selection_rate=mask_selection_rate,
+            mask_token_id=self.mask_token_id,
+            unselectable_token_ids=None,
+        )
+        inputs = tf.ragged.constant(
+            [[0, 1, 2], [0, 1, 2, 3, 4, 5], [0, 1, 2, 3, 4]]
+        )
+        outputs = mlm_masker(inputs)
+        expected_number_of_masked_tokens = tf.cast(
+            tf.math.ceil(
+                tf.cast(inputs.row_lengths(), dtype=tf.float32)
+                * mask_selection_rate,
+            ),
+            dtype=tf.int64,
+        )
+
+        self.assertAllEqual(
+            outputs["mask_positions"].row_lengths(),
+            expected_number_of_masked_tokens,
+        )
+
+        # Cap the number of masked tokens at 0, so we can test if
+        # mask_selection_length takes effect.
+        mask_selection_length = 0
+        mlm_masker = MLMMaskGenerator(
+            vocabulary_size=self.vocabulary_size,
+            mask_selection_rate=mask_selection_rate,
+            mask_token_id=self.mask_token_id,
+            mask_selection_length=mask_selection_length,
+        )
+        outputs = mlm_masker(inputs)
+        self.assertEqual(tf.reduce_sum(outputs["mask_positions"]), 0)
+
+    def test_apply_random_token_not_mask(self):
+        mlm_masker = MLMMaskGenerator(
+            vocabulary_size=self.vocabulary_size,
+            mask_selection_rate=0.5,
+            mask_token_id=self.mask_token_id,
+            mask_token_rate=0,
+            random_token_rate=1,
+        )
+        inputs = tf.random.uniform(
+            shape=[5, 10],
+            maxval=len(self.VOCAB),
+            dtype=tf.int32,
+        )
+        outputs = mlm_masker(inputs)
+        tokens, mask_positions, mask_ids = (
+            outputs["tokens"],
+            outputs["mask_positions"],
+            outputs["mask_ids"],
+        )
+        self.assertAllEqual(tokens.shape, inputs.shape)
+        self.assertAllEqual(
+            mask_positions.row_lengths(), mask_ids.row_lengths()
+        )
+        masked_values = tf.gather(
+            tokens,
+            mask_positions,
+            batch_dims=1,
+        )
+        # Verify that selected tokens are replaced by random tokens.
+        self.assertNotEqual(tf.reduce_mean(masked_values), self.mask_token_id)
+
+    def test_invalid_mask_token(self):
+        with self.assertRaisesRegex(ValueError, "Mask token id should be*"):
+            _ = MLMMaskGenerator(
+                vocabulary_size=self.vocabulary_size,
+                mask_selection_rate=0.5,
+                mask_token_id=self.vocabulary_size,
+                mask_selection_length=5,
+            )
+
+    def test_unselectable_tokens(self):
+        unselectable_token_ids = [
+            self.vocabulary_size - 1,
+            self.vocabulary_size - 2,
+        ]
+        mlm_masker = MLMMaskGenerator(
+            vocabulary_size=self.vocabulary_size,
+            mask_selection_rate=1,
+            mask_token_id=self.mask_token_id,
+            mask_selection_length=5,
+            unselectable_token_ids=unselectable_token_ids,
+            mask_token_rate=1,
+            random_token_rate=0,
+        )
+        inputs = [unselectable_token_ids]
+        outputs = mlm_masker(inputs)
+        # Verify that no token is masked out.
+        self.assertEqual(tf.reduce_sum(outputs["mask_positions"]), 0)
+
+    def test_config(self):
+        unselectable_token_ids = [
+            self.vocabulary_size - 1,
+            self.vocabulary_size - 2,
+        ]
+        mlm_masker = MLMMaskGenerator(
+            vocabulary_size=self.vocabulary_size,
+            mask_selection_rate=0.5,
+            mask_token_id=self.mask_token_id,
+            mask_selection_length=5,
+            unselectable_token_ids=unselectable_token_ids,
+        )
+        config = mlm_masker.get_config()
+        expected_config = {
+            "vocabulary_size": self.vocabulary_size,
+            "unselectable_token_ids": unselectable_token_ids,
+        }
+        self.assertDictContainsSubset(expected_config, config)
+
+        # Test cloned mlm_masker can be run.
+        cloned_mlm_masker = MLMMaskGenerator.from_config(config)
+        inputs = tf.ragged.constant(
+            [[0, 1, 2], [0, 1, 2, 3, 4, 5], [0, 1, 2, 3, 4]]
+        )
+        cloned_mlm_masker(inputs)
diff --git a/keras_nlp/tests/__init__.py b/keras_nlp/tests/__init__.py
new file mode 100644
index 00000000..6e4df4e7
--- /dev/null
+++ b/keras_nlp/tests/__init__.py
@@ -0,0 +1,13 @@
+# 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.
diff --git a/keras_nlp/tests/docstring_lib.py b/keras_nlp/tests/docstring_lib.py
new file mode 100644
index 00000000..b3f657d5
--- /dev/null
+++ b/keras_nlp/tests/docstring_lib.py
@@ -0,0 +1,221 @@
+# 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.
+"""Doctests lib for KerasNLP."""
+
+import doctest
+import re
+import textwrap
+
+import numpy as np
+
+
+class _FloatExtractor(object):
+    """Class for extracting floats from a string.
+
+    For example:
+
+    >>> text_parts, floats = _FloatExtractor()("Text 1.0 Text")
+    >>> text_parts
+    ['Text ', ' Text']
+    >>> floats
+    array([1.])
+    """
+
+    # Note: non-capturing groups "(?" are not returned in matched groups, or by
+    # re.split.
+    _FLOAT_RE = re.compile(
+        r"""
+        (                          # Captures the float value.
+            (?:
+            [-+]|                 # Start with a sign is okay anywhere.
+            (?:                   # Otherwise:
+                ^|                # Start after the start of string
+                (?<=[^\w.])       # Not after a word char, or a .
+            )
+            )
+            (?:                      # Digits and exponent - something like:
+            {digits_dot_maybe_digits}{exponent}?|   # "1.0" "1." "1.0e3", "1.e3"
+            {dot_digits}{exponent}?|                # ".1" ".1e3"
+            {digits}{exponent}|                     # "1e3"
+            {digits}(?=j)                           # "300j"
+            )
+        )
+        j?                         # Optional j for cplx numbers, not captured.
+        (?=                        # Only accept the match if
+            $|                       # * At the end of the string, or
+            [^\w.]                   # * Next char is not a word char or "."
+        )
+        """.format(
+            # Digits, a "." and optional more digits: "1.1".
+            digits_dot_maybe_digits=r"(?:[0-9]+\.(?:[0-9]*))",
+            # A "." with trailing digits ".23"
+            dot_digits=r"(?:\.[0-9]+)",
+            # digits: "12"
+            digits=r"(?:[0-9]+)",
+            # The exponent: An "e" or "E", optional sign, and at least one
+            # digit. "e-123", "E+12", "e12"
+            exponent=r"(?:[eE][-+]?[0-9]+)",
+        ),
+        re.VERBOSE,
+    )
+
+    def __call__(self, string):
+        """Extracts floats from a string.
+
+        >>> text_parts, floats = _FloatExtractor()("Text 1.0 Text")
+        >>> text_parts
+        ['Text ', ' Text']
+        >>> floats
+        array([1.])
+
+        Args:
+            string: the string to extract floats from.
+
+        Returns:
+            A (string, array) pair, where `string` has each float replaced by
+            "..." and `array` is a `float32` `numpy.array` containing the
+            extracted floats.
+        """
+        texts = []
+        floats = []
+        for i, part in enumerate(self._FLOAT_RE.split(string)):
+            if i % 2 == 0:
+                texts.append(part)
+            else:
+                floats.append(float(part))
+
+        return texts, np.array(floats)
+
+
+class DoctestOutputChecker(doctest.OutputChecker, object):
+    """Customizes how `want` and `got` are compared, see `check_output`."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.extract_floats = _FloatExtractor()
+        self.text_good = None
+        self.float_size_good = None
+
+    _ADDRESS_RE = re.compile(r"\bat 0x[0-9a-f]*?>")
+    # TODO(yashkatariya): Add other tensor's string substitutions too.
+    # tf.RaggedTensor doesn't need one.
+    _NUMPY_OUTPUT_RE = re.compile(r"<tf.Tensor.*?numpy=(.*?)>", re.DOTALL)
+
+    def _allclose(self, want, got, rtol=1e-3, atol=1e-3):
+        return np.allclose(want, got, rtol=rtol, atol=atol)
+
+    def _tf_tensor_numpy_output(self, string):
+        modified_string = self._NUMPY_OUTPUT_RE.sub(r"\1", string)
+        return modified_string, modified_string != string
+
+    MESSAGE = textwrap.dedent(
+        """\n
+        #############################################################
+        Check the documentation on how to write testable docstrings.
+        https://www.tensorflow.org/community/contribute/docs_ref
+        #############################################################"""
+    )
+
+    def check_output(self, want, got, optionflags):
+        """Compares the docstring output to the output from running the code.
+
+        Python addresses in the output are replaced with wildcards.
+
+        Float values in the output compared as using `np.allclose`:
+
+            * Float values are extracted from the text and replaced with
+              wildcards.
+            * The wildcard text is compared to the actual output.
+            * The float values are compared using `np.allclose`.
+
+        The method returns `True` if both the text comparison and the numeric
+        comparison are successful.
+
+        The numeric comparison will fail if either:
+
+            * The wrong number of floats are found.
+            * The float values are not within tolerence.
+
+        Args:
+            want: The output in the docstring.
+            got: The output generated after running the snippet.
+            optionflags: Flags passed to the doctest.
+
+        Returns:
+            A bool, indicating if the check was successful or not.
+        """
+
+        # If the docstring's output is empty and there is some output generated
+        # after running the snippet, return True. This is because if the user
+        # doesn't want to display output, respect that over what the doctest
+        # wants.
+        if got and not want:
+            return True
+
+        if want is None:
+            want = ""
+
+        # Replace python's addresses with ellipsis (`...`) since it can change
+        # on each execution.
+        want = self._ADDRESS_RE.sub("at ...>", want)
+
+        # Replace tf.Tensor strings with only their numpy field values.
+        want, want_changed = self._tf_tensor_numpy_output(want)
+        if want_changed:
+            got, _ = self._tf_tensor_numpy_output(got)
+
+        # Separate out the floats, and replace `want` with the wild-card version
+        # "result=7.0" => "result=..."
+        want_text_parts, self.want_floats = self.extract_floats(want)
+        want_text_wild = "...".join(want_text_parts)
+
+        # Find the floats in the string returned by the test
+        _, self.got_floats = self.extract_floats(got)
+
+        self.text_good = super().check_output(
+            want=want_text_wild, got=got, optionflags=optionflags
+        )
+        if not self.text_good:
+            return False
+
+        if self.want_floats.size == 0:
+            # If there are no floats in the "want" string, ignore all the floats
+            # in the result. "np.array([ ... ])" matches
+            # "np.array([ 1.0, 2.0 ])"
+            return True
+
+        self.float_size_good = self.want_floats.size == self.got_floats.size
+
+        if self.float_size_good:
+            return self._allclose(self.want_floats, self.got_floats)
+        else:
+            return False
+
+    def output_difference(self, example, got, optionflags):
+        got = [got]
+
+        # If the some of the float output is hidden with `...`,
+        # `float_size_good` will be False. This is because the floats extracted
+        # from the string is converted into a 1-D numpy array. Hence hidding
+        # floats is not allowed anymore.
+        if self.text_good:
+            if not self.float_size_good:
+                got.append(
+                    "\n\nCAUTION: tf_doctest doesn't work if *some* of the "
+                    '*float output* is hidden with a "...".'
+                )
+
+        got.append(self.MESSAGE)
+        got = "\n".join(got)
+        return super().output_difference(example, got, optionflags)
diff --git a/keras_nlp/tests/docstring_test.py b/keras_nlp/tests/docstring_test.py
new file mode 100644
index 00000000..35ec2666
--- /dev/null
+++ b/keras_nlp/tests/docstring_test.py
@@ -0,0 +1,75 @@
+# 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 doctest
+import os
+import sys
+import unittest
+
+import numpy as np
+import pytest
+import tensorflow as tf
+from tensorflow import keras
+
+import keras_nlp
+from keras_nlp.tests import docstring_lib
+
+PACKAGE = "keras_nlp."
+
+
+def find_modules():
+    keras_nlp_modules = []
+    for name, module in sys.modules.items():
+        if name.startswith(PACKAGE):
+            keras_nlp_modules.append(module)
+
+    return keras_nlp_modules
+
+
+@pytest.mark.skipif(
+    sys.platform == "win32", reason="Numpy prints differently on windows"
+)
+def test_docstrings():
+    keras_nlp_modules = find_modules()
+    # As of this writing, it doesn't seem like pytest support load_tests
+    # protocol for unittest:
+    #     https://docs.pytest.org/en/7.1.x/how-to/unittest.html
+    # So we run the unittest.TestSuite manually and report the results back.
+    runner = unittest.TextTestRunner()
+    suite = unittest.TestSuite()
+    for module in keras_nlp_modules:
+        suite.addTest(
+            doctest.DocTestSuite(
+                module,
+                test_finder=doctest.DocTestFinder(exclude_empty=False),
+                extraglobs={
+                    "tf": tf,
+                    "np": np,
+                    "os": os,
+                    "keras": keras,
+                    "keras_nlp": keras_nlp,
+                },
+                checker=docstring_lib.DoctestOutputChecker(),
+                optionflags=(
+                    doctest.ELLIPSIS
+                    | doctest.NORMALIZE_WHITESPACE
+                    | doctest.IGNORE_EXCEPTION_DETAIL
+                    | doctest.DONT_ACCEPT_BLANKLINE
+                ),
+            )
+        )
+    result = runner.run(suite)
+    if not result.wasSuccessful():
+        print(result)
+    assert result.wasSuccessful()