Repository: keras-team/keras-nlp

Issue:
Add a UnicodeCharacterTokenizer
This would be a simple tokenizer which has no vocabulary, and simply converts text to unicode codepoints.

## Potential docstring and usage

```python
"""
Unicode character codepoint tokenizer.
 
This tokenizer is a vocabulary free tokenizer which will tokenize text as
unicode character codepoints.
 
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 codepoint is encountered. (same behavior as
     https://www.tensorflow.org/api_docs/python/tf/strings/unicode_transcode)
"""

tokenizer = keras_nlp.tokenizers.UnicodeCharacterTokenizer()

inputs = ["▀▁▂▃▄", "hello"]
tokens = tokenizer.tokenize(inputs)
>>> <tf.RaggedTensor [
      [9600, 9601, 9602, 9603, 9604],
      [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/9d9dcb8b640d8183a614ddd6f92fe368/unicodecharactertokenizer.ipynb

Discussion:
mattdangerw: @aflah02 thanks! Will take a look soon.

Re decoding, yeah the way to go from a tensor scalar string to a python unicode strings is `.numpy().decode('utf-8')`. That's totally correct and you can use that for code examples/tests. We might want to provide some extra helper for this on the tokenizer base class (e.g. `tokenizer.detokenize_as_string`) as I suspect this will be a common point of confusion, but I don't think we need to worry about that on this PR.
aflah02: @mattdangerw I've made most changes however I'm facing 2 issues

- lint keeps raising an issue that the imports are ordered incorrectly even after I run format.sh, the exact error is:
```
ERROR: /mnt/c/Users/ASUS/Desktop/keras-nlp/keras_nlp/tokenizers/unicode_character_tokenizer_test.py Imports are incorrectly sorted and/or formatted.
ERROR: /mnt/c/Users/ASUS/Desktop/keras-nlp/keras_nlp/tokenizers/__init__.py Imports are incorrectly sorted and/or formatted.
Skipped 1 files
```
- I'm unable to figure out how to decode the inputs to UTF-8 I tried calling .numpy() and other variants but it's not working as if I try to do something like a .numpy() .decode() call after unicode_encode there's an error that the object has no .numpy() attribute. Any clues as to what might be going wrong? I've tried several variations but none seem to be working out
mattdangerw: @aflah02 thanks!

Re first issue, I think someone else was hitting that too. I will take a look later today.

Re second issue, do you have a code snippet that's not working? In terms of a utility to detokenize to python strings (for all tokenizers) let's do that on a separate issue. I'll open one today. Are you interested in taking that on?
aflah02: @mattdangerw 
Cool 
Yeah sure I can take it up and I'll share the code snippets that I tried in a bit my commits would have them but I'll just list out the main ones
aflah02: @mattdangerw Thanks!
These are the snippets I've tried:
- Doing a `.numpy() `followed by .decode('utf-8') before returning [here](https://github.com/aflah02/keras-nlp/blob/1ec59dfb6ab3f8376e8111da2c7773895eebcd1b/keras_nlp/tokenizers/unicode_character_tokenizer.py#L276)
- Since `.numpy()` raised errors I tried to list all available methods with the object and spotted `._numpy` but that also raised issues after that
- Using exactly what Byte Tokenizer does however it uses a character list which would be very large for unicode characters to always generate for detokenization

So essentially the only remaining issue is `"▀▁▂▃"` gets tokenized to `[9600, 9601, 9602, 9603]` but during detokenization it becomes `b"\xe2\x96\x80\xe2\x96\x81\xe2\x96\x82\xe2\x96\x83"` and not `"▀▁▂▃"` again.

mattdangerw: @aflah02 I think your return types are actually correct. b"\xe2\x96\x80\xe2\x96\x81\xe2\x96\x82\xe2\x96\x83" is a byte string representation for "▀▁▂▃" in a utf-8 encoding. And by default when you are working with tf.Tensor strings, you get byte strings with a utf-8 encoding. https://www.tensorflow.org/text/guide/unicode

So the issue is not that this representation of the return type is wrong, it's just that most users will want eventually to go from tf.Tensors to regular python strings and it is confusing how to reverse this process.

So that's where I was proposing we add a helper to all tokenizers, e.g. `tokenizer.detokenize_to_strings()` that will convert all tensors to lists of lists, and then convert all list elements to python strings. https://github.com/keras-team/keras-nlp/issues/113
aflah02: @mattdangerw Thanks! That makes total sense now

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..19c6dfc8
--- /dev/null
+++ b/keras_nlp/layers/preprocessing/mlm_mask_generator.py
@@ -0,0 +1,186 @@
+# 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.layers.preprocessing.MLMMaskGenerator( \
+            vocabulary_size=10, mask_selection_rate=0.2, mask_token_id=0, \
+            mask_selection_length=5)
+    >>> masker(tf.constant([1, 2, 3, 4, 5]))
+
+    Ragged Input:
+    >>> masker = keras_nlp.layers.preprocessing.MLMMaskGenerator( \
+            vocabulary_size=10, mask_selection_rate=0.5, mask_token_id=0, \
+            mask_selection_length=5)
+    >>> masker(tf.ragged.constant([[1, 2], [1, 2, 3, 4]]))
+    """
+
+    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/__init__.py b/keras_nlp/tokenizers/__init__.py
index 7206410a..36c530ee 100644
--- a/keras_nlp/tokenizers/__init__.py
+++ b/keras_nlp/tokenizers/__init__.py
@@ -14,4 +14,7 @@
 
 from keras_nlp.tokenizers.byte_tokenizer import ByteTokenizer
 from keras_nlp.tokenizers.tokenizer import Tokenizer
+from keras_nlp.tokenizers.unicode_character_tokenizer import (
+    UnicodeCharacterTokenizer,
+)
 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
index 2f142004..1a106cd9 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()
@@ -97,7 +97,7 @@ class ByteTokenizer(tokenizer.Tokenizer):
     >>> ds.take(1).get_single_element()
     <tf.RaggedTensor [[104, 101, 108, 108, 111], [102, 117, 110]]>
 
-    Batch up the inputs and then tokenize.
+    Batch 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)
@@ -112,16 +112,16 @@ 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).
+    Batch 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]])>
+           [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/unicode_character_tokenizer.py b/keras_nlp/tokenizers/unicode_character_tokenizer.py
new file mode 100644
index 00000000..ade085e3
--- /dev/null
+++ b/keras_nlp/tokenizers/unicode_character_tokenizer.py
@@ -0,0 +1,272 @@
+# 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 typing import Any
+from typing import Dict
+
+import tensorflow as tf
+import tensorflow_text as tf_text
+
+from keras_nlp.tokenizers import tokenizer
+
+
+class UnicodeCharacterTokenizer(tokenizer.Tokenizer):
+    """A unicode character tokenizer layer.
+
+    This tokenizer is a vocabulary free tokenizer which tokenizes text as
+    unicode characters codepoints.
+
+    Args:
+        lowercase: If true, the input text will be first lowered before
+            tokenization.
+        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 codepoint is encountered.
+            (same behavior as
+            https://www.tensorflow.org/api_docs/python/tf/strings/unicode_transcode)
+        replacement_char: The unicode codepoint to use in place of invalid
+            codepoints. Defaults to 65533 (U+FFFD).
+        input_encoding: One of ("UTF-8", "UTF-16-BE", or "UTF-32-BE").
+            One of The encoding of the input text. Defaults to "UTF-8".
+        output_encoding: One of ("UTF-8", "UTF-16-BE", or "UTF-32-BE").
+            The encoding of the output text. Defaults to "UTF-8".
+
+    Examples:
+
+    Basic Usage.
+    >>> inputs = "Unicode Tokenizer"
+    >>> tokenizer = keras_nlp.tokenizers.UnicodeCharacterTokenizer()
+    >>> tokenizer(inputs)
+    <tf.Tensor: shape=(17,), dtype=int32, numpy=
+    array([117, 110, 105,  99, 111, 100, 101,  32, 116, 111, 107, 101, 110,
+        105, 122, 101, 114], dtype=int32)>
+
+    Ragged outputs.
+    >>> inputs = ["Book", "पुस्तक", "کتاب"]
+    >>> tokenizer = keras_nlp.tokenizers.UnicodeCharacterTokenizer()
+    >>> tokenizer(inputs)
+    <tf.RaggedTensor [[98, 111, 111, 107],
+        [2346, 2369, 2360, 2381, 2340, 2325],
+        [1705, 1578, 1575, 1576]]>
+
+    Dense outputs.
+    >>> inputs = ["Book", "पुस्तक", "کتاب"]
+    >>> tokenizer = keras_nlp.tokenizers.UnicodeCharacterTokenizer(
+    ...     sequence_length=8)
+    >>> tokenizer(inputs)
+    <tf.Tensor: shape=(3, 8), dtype=int32, numpy=
+    array([[  98,  111,  111,  107,    0,    0,    0,    0],
+        [2346, 2369, 2360, 2381, 2340, 2325,    0,    0],
+        [1705, 1578, 1575, 1576,    0,    0,    0,    0]], dtype=int32)>
+
+    Tokenize first, then batch the dataset.
+    >>> inputs = ["Book", "पुस्तक", "کتاب"]
+    >>> tokenizer = keras_nlp.tokenizers.UnicodeCharacterTokenizer()
+    >>> ds = tf.data.Dataset.from_tensor_slices(inputs)
+    >>> ds = ds.map(tokenizer)
+    >>> ds = ds.apply(tf.data.experimental.dense_to_ragged_batch(3))
+    >>> ds.take(1).get_single_element()
+    <tf.RaggedTensor [[98, 111, 111, 107],
+        [2346, 2369, 2360, 2381, 2340, 2325],
+        [1705, 1578, 1575, 1576]]>
+
+    Batch the inputs and then tokenize.
+    >>> inputs = ["Book", "पुस्तक", "کتاب"]
+    >>> tokenizer = keras_nlp.tokenizers.UnicodeCharacterTokenizer()
+    >>> ds = tf.data.Dataset.from_tensor_slices(inputs)
+    >>> ds = ds.batch(3).map(tokenizer)
+    >>> ds.take(1).get_single_element()
+    <tf.RaggedTensor [[98, 111, 111, 107],
+        [2346, 2369, 2360, 2381, 2340, 2325],
+        [1705, 1578, 1575, 1576]]>
+
+    Tokenize first, then batch for dense outputs (`sequence_length` provided).
+    >>> inputs = ["Book", "पुस्तक", "کتاب"]
+    >>> tokenizer = keras_nlp.tokenizers.UnicodeCharacterTokenizer(
+    ...     sequence_length=5)
+    >>> ds = tf.data.Dataset.from_tensor_slices(inputs)
+    >>> ds = ds.map(tokenizer)
+    >>> ds = ds.apply(tf.data.experimental.dense_to_ragged_batch(3))
+    >>> ds.take(1).get_single_element()
+    <tf.Tensor: shape=(3, 5), dtype=int32, numpy=
+    array([[  98,  111,  111,  107,    0],
+        [2346, 2369, 2360, 2381, 2340],
+        [1705, 1578, 1575, 1576,    0]], dtype=int32)>
+
+    Batch first, then tokenize for dense outputs (`sequence_length` provided).
+    (`sequence_length` provided).
+    >>> inputs = ["Book", "पुस्तक", "کتاب"]
+    >>> tokenizer = keras_nlp.tokenizers.UnicodeCharacterTokenizer(
+    ...     sequence_length=5)
+    >>> ds = tf.data.Dataset.from_tensor_slices(inputs)
+    >>> ds = ds.batch(3).map(tokenizer)
+    >>> ds.take(1).get_single_element()
+    <tf.Tensor: shape=(3, 5), dtype=int32, numpy=
+    array([[  98,  111,  111,  107,    0],
+        [2346, 2369, 2360, 2381, 2340],
+        [1705, 1578, 1575, 1576,    0]], dtype=int32)>
+
+    Tokenization showcasing truncation of long sequences.
+    >>> inputs = ["I Like to Travel a Lot", "मैं किताबें पढ़ना पसंद करता हूं"]
+    >>> tokenizer = keras_nlp.tokenizers.UnicodeCharacterTokenizer(
+    ...     sequence_length=5)
+    >>> tokenizer(inputs)
+    <tf.Tensor: shape=(5,), dtype=int32,
+        numpy=array([[ 105,   32,  108,  105,  107],
+       [2350, 2376, 2306,   32, 2325]], dtype=int32)>
+
+    Detokenization.
+    >>> inputs = tf.constant([110, 105, 110, 106,  97], dtype=tf.int32)
+    >>> tokenizer = keras_nlp.tokenizers.UnicodeCharacterTokenizer()
+    >>> tokenizer.detokenize(inputs)
+    <tf.Tensor: shape=(), dtype=string, numpy=b'ninja'>
+
+    Detokenization while showcasing padded characters being removed
+    >>> tokenizer = keras_nlp.tokenizers.UnicodeCharacterTokenizer(
+    ...     sequence_length=7)
+    >>> dataset = tf.data.Dataset.from_tensor_slices(["a b c", "b c", "a"])
+    >>> dataset = dataset.map(tokenizer)
+    >>> dataset.take(1).get_single_element()
+    <tf.Tensor: shape=(7,), dtype=int32,
+        numpy=array([97, 32, 98, 32, 99,  0,  0], dtype=int32)>
+    >>> detokunbatched = dataset.map(tokenizer.detokenize)
+    >>> detokunbatched = dataset.map(tokenizer.detokenize)
+    >>> detokunbatched.take(1).get_single_element()
+    <tf.Tensor: shape=(), dtype=string, numpy=b'a b c'>
+
+    Detokenization with invalid bytes.
+    >>> # The 10000000 in the inputs tensor below is an invalid value
+    >>> # Hence it replaces to the replacement_char 75 which represents 'K'
+    >>> inputs = tf.constant([110, 105, 10000000, 110, 106,  97])
+    >>> tokenizer = keras_nlp.tokenizers.UnicodeCharacterTokenizer(
+    ...     errors="replace", replacement_char=75)
+    >>> tokenizer.detokenize(inputs).numpy().decode('utf-8')
+    'niKnja'
+    """
+
+    def __init__(
+        self,
+        sequence_length: int = None,
+        lowercase: bool = True,
+        normalization_form: str = None,
+        errors: str = "replace",
+        replacement_char: int = 65533,
+        input_encoding: str = "UTF-8",
+        output_encoding: str = "UTF-8",
+        **kwargs,
+    ) -> None:
+        # 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 and dtype != tf.string:
+                raise ValueError(
+                    "Output dtype must be an integer type of a string. "
+                    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}"
+            )
+
+        # Check normalization_form matches input_encoding.
+        if normalization_form:
+            if input_encoding != "UTF-8":
+                raise ValueError(
+                    """Normalization Forms are Only Supported for Input Encoding
+                     UTF-8"""
+                )
+
+        super().__init__(**kwargs)
+
+        self.sequence_length = sequence_length
+        self.lowercase = lowercase
+        self.normalization_form = normalization_form
+        self.errors = errors
+        self.replacement_char = replacement_char
+        self.input_encoding = input_encoding
+        self.output_encoding = output_encoding
+
+    def get_config(self) -> Dict[str, Any]:
+        config = super().get_config()
+        config.update(
+            {
+                "sequence_length": self.sequence_length,
+                "lowercase": self.lowercase,
+                "normalization_form": self.normalization_form,
+                "errors": self.errors,
+                "replacement_char": self.replacement_char,
+                "input_encoding": self.input_encoding,
+                "output_encoding": self.output_encoding,
+            }
+        )
+        return config
+
+    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)
+
+        # Optionally lowercase the text
+        if self.lowercase:
+            inputs = tf_text.case_fold_utf8(inputs)
+
+        # Optionally normalize the text to a given form
+        if self.normalization_form:
+            inputs = tf_text.normalize_utf8(inputs, self.normalization_form)
+
+        tokens = tf.strings.unicode_decode(
+            inputs,
+            errors=self.errors,
+            replacement_char=self.replacement_char,
+            input_encoding=self.input_encoding,
+        )
+
+        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):
+        inputs = tf.ragged.boolean_mask(inputs, tf.not_equal(inputs, 0))
+        encoded_string = tf.strings.unicode_encode(
+            inputs,
+            errors=self.errors,
+            replacement_char=self.replacement_char,
+            output_encoding=self.output_encoding,
+        )
+        return encoded_string
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()
diff --git a/keras_nlp/tokenizers/unicode_character_tokenizer_test.py b/keras_nlp/tokenizers/unicode_character_tokenizer_test.py
new file mode 100644
index 00000000..34df6a50
--- /dev/null
+++ b/keras_nlp/tokenizers/unicode_character_tokenizer_test.py
@@ -0,0 +1,317 @@
+# 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.unicode_character_tokenizer import (
+    UnicodeCharacterTokenizer,
+)
+
+
+class UnicodeCharacterTokenizerTest(tf.test.TestCase):
+    def test_tokenize(self):
+        input_data = tf.constant(["ninja", "samurai", "▀▁▂▃"])
+        tokenizer = UnicodeCharacterTokenizer()
+        call_output = tokenizer(input_data)
+        tokenize_output = tokenizer.tokenize(input_data)
+        self.assertIsInstance(call_output, tf.RaggedTensor)
+        exp_outputs = [
+            [110, 105, 110, 106, 97],
+            [115, 97, 109, 117, 114, 97, 105],
+            [9600, 9601, 9602, 9603],
+        ]
+        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 = "ninja"
+        tokenizer = UnicodeCharacterTokenizer()
+        call_output = tokenizer(input_data)
+        tokenize_output = tokenizer.tokenize(input_data)
+
+        self.assertAllEqual(call_output, [110, 105, 110, 106, 97])
+        self.assertAllEqual(tokenize_output, [110, 105, 110, 106, 97])
+
+    def test_dense_output(self):
+        input_data = tf.constant(["ninja", "samurai", "▀▁▂▃"])
+        tokenizer = UnicodeCharacterTokenizer(sequence_length=10)
+        call_output = tokenizer(input_data)
+        self.assertIsInstance(call_output, tf.Tensor)
+        self.assertAllEqual(
+            call_output,
+            [
+                [110, 105, 110, 106, 97, 0, 0, 0, 0, 0],
+                [115, 97, 109, 117, 114, 97, 105, 0, 0, 0],
+                [9600, 9601, 9602, 9603, 0, 0, 0, 0, 0, 0],
+            ],
+        )
+
+    def test_detokenize(self):
+        input_data = tf.ragged.constant(
+            [
+                [110, 105, 110, 106, 97],
+                [115, 97, 109, 117, 114, 97, 105],
+                [9600, 9601, 9602, 9603],
+            ]
+        )
+
+        tokenizer = UnicodeCharacterTokenizer()
+        detokenize_output = tokenizer.detokenize(input_data)
+        self.assertAllEqual(
+            detokenize_output,
+            [
+                b"ninja",
+                b"samurai",
+                b"\xe2\x96\x80\xe2\x96\x81\xe2\x96\x82\xe2\x96\x83",
+            ],
+        )
+
+    def test_detokenize_replace_error(self):
+        # 10000000 is an invalid value
+        input_data = tf.ragged.constant([[110, 105, 10000000, 110, 106, 97]])
+        tokenizer = UnicodeCharacterTokenizer(
+            errors="replace", replacement_char=75
+        )
+        detokenize_output = tokenizer.detokenize(input_data)
+        self.assertAllEqual(detokenize_output, [b"niKnja"])
+
+    def test_detokenize_ignore_error(self):
+        input_data = tf.ragged.constant([[110, 105, 10000000, 110, 106, 97]])
+        tokenizer = UnicodeCharacterTokenizer(errors="ignore")
+        detokenize_output = tokenizer.detokenize(input_data)
+        self.assertAllEqual(detokenize_output, [b"ninja"])
+
+    def test_detokenize_strict_error(self):
+        input_data = tf.ragged.constant([[110, 105, 10000000, 110, 106, 97]])
+        tokenizer = UnicodeCharacterTokenizer(errors="strict")
+        with self.assertRaises(tf.errors.InvalidArgumentError):
+            _ = tokenizer.detokenize(input_data)
+
+    def test_normalization_without_UTF8_valueerror(self):
+        with self.assertRaises(ValueError):
+            _ = UnicodeCharacterTokenizer(
+                errors="strict",
+                input_encoding="UTF-16",
+                normalization_form="NFC",
+            )
+
+    def test_lowercase(self):
+        input_data = tf.constant(["NiNJaS"])
+        tokenizer = UnicodeCharacterTokenizer()
+        call_output = tokenizer(input_data)
+        self.assertAllEqual(
+            call_output,
+            [[110, 105, 110, 106, 97, 115]],
+        )
+
+    def test_skip_lowercase(self):
+        input_data = tf.constant(["NiNJaS"])
+        tokenizer = UnicodeCharacterTokenizer(lowercase=False)
+        call_output = tokenizer(input_data)
+        self.assertAllEqual(
+            call_output,
+            [[78, 105, 78, 74, 97, 83]],
+        )
+
+    def test_tokenize_first_batch_second(self):
+        tokenizer = UnicodeCharacterTokenizer()
+
+        ds = tf.data.Dataset.from_tensor_slices(
+            ["ninja", "samurai", "▀▁▂▃", "keras", "tensorflow"]
+        )
+        ds = ds.map(tokenizer)
+        ds = ds.apply(tf.data.experimental.dense_to_ragged_batch(5))
+        output = ds.take(1).get_single_element()
+
+        exp_output = [
+            [110, 105, 110, 106, 97],
+            [115, 97, 109, 117, 114, 97, 105],
+            [9600, 9601, 9602, 9603],
+            [107, 101, 114, 97, 115],
+            [116, 101, 110, 115, 111, 114, 102, 108, 111, 119],
+        ]
+        for i in range(output.shape[0]):
+            self.assertAllEqual(output[i], exp_output[i])
+
+    def test_tokenize_first_batch_second_with_sequence_length(self):
+        tokenizer = UnicodeCharacterTokenizer(sequence_length=10)
+
+        ds = tf.data.Dataset.from_tensor_slices(
+            ["ninja", "samurai", "▀▁▂▃", "keras", "tensorflow"]
+        )
+        ds = ds.map(tokenizer)
+        ds = ds.apply(tf.data.experimental.dense_to_ragged_batch(5))
+        output = ds.take(1).get_single_element()
+
+        exp_output = [
+            [110, 105, 110, 106, 97, 0, 0, 0, 0, 0],
+            [115, 97, 109, 117, 114, 97, 105, 0, 0, 0],
+            [9600, 9601, 9602, 9603, 0, 0, 0, 0, 0, 0],
+            [107, 101, 114, 97, 115, 0, 0, 0, 0, 0],
+            [116, 101, 110, 115, 111, 114, 102, 108, 111, 119],
+        ]
+        for i in range(output.shape[0]):
+            self.assertAllEqual(output[i], exp_output[i])
+
+    def test_batch_first_tokenize_second(self):
+        tokenizer = UnicodeCharacterTokenizer()
+
+        ds = tf.data.Dataset.from_tensor_slices(
+            ["ninja", "samurai", "▀▁▂▃", "keras", "tensorflow"]
+        )
+        ds = ds.batch(5).map(tokenizer)
+        output = ds.take(1).get_single_element()
+
+        exp_output = [
+            [110, 105, 110, 106, 97],
+            [115, 97, 109, 117, 114, 97, 105],
+            [9600, 9601, 9602, 9603],
+            [107, 101, 114, 97, 115],
+            [116, 101, 110, 115, 111, 114, 102, 108, 111, 119],
+        ]
+        for i in range(output.shape[0]):
+            self.assertAllEqual(output[i], exp_output[i])
+
+    def test_batch_first_tokenize_second_with_sequence_length(self):
+        tokenizer = UnicodeCharacterTokenizer(sequence_length=10)
+
+        ds = tf.data.Dataset.from_tensor_slices(
+            ["ninja", "samurai", "▀▁▂▃", "keras", "tensorflow"]
+        )
+        ds = ds.batch(5).map(tokenizer)
+        output = ds.take(1).get_single_element()
+
+        exp_output = [
+            [110, 105, 110, 106, 97, 0, 0, 0, 0, 0],
+            [115, 97, 109, 117, 114, 97, 105, 0, 0, 0],
+            [9600, 9601, 9602, 9603, 0, 0, 0, 0, 0, 0],
+            [107, 101, 114, 97, 115, 0, 0, 0, 0, 0],
+            [116, 101, 110, 115, 111, 114, 102, 108, 111, 119],
+        ]
+        for i in range(output.shape[0]):
+            self.assertAllEqual(output[i], exp_output[i])
+
+    def test_functional_model(self):
+        input_data = tf.constant(
+            ["ninja", "samurai", "▀▁▂▃", "keras", "tensorflow"]
+        )
+        tokenizer = UnicodeCharacterTokenizer()
+        inputs = tf.keras.Input(dtype="string", shape=())
+        outputs = tokenizer.detokenize(tokenizer.tokenize(inputs))
+        model = tf.keras.Model(inputs, outputs)
+        model_output = model(input_data)
+        self.assertAllEqual(
+            model_output,
+            [
+                b"ninja",
+                b"samurai",
+                b"\xe2\x96\x80\xe2\x96\x81\xe2\x96\x82\xe2\x96\x83",
+                b"keras",
+                b"tensorflow",
+            ],
+        )
+
+    def test_load_model_with_config(self):
+        input_data = tf.constant(["hello"])
+
+        original_tokenizer = UnicodeCharacterTokenizer(
+            lowercase=False,
+            sequence_length=11,
+            normalization_form="NFC",
+            errors="strict",
+        )
+        cloned_tokenizer = UnicodeCharacterTokenizer.from_config(
+            original_tokenizer.get_config()
+        )
+        self.assertAllEqual(
+            original_tokenizer(input_data),
+            cloned_tokenizer(input_data),
+        )
+
+        decoded_input = [107, 101, 114, 97, 115]
+        self.assertAllEqual(
+            original_tokenizer.detokenize(decoded_input),
+            cloned_tokenizer.detokenize(decoded_input),
+        )
+
+    def test_config(self):
+        tokenizer = UnicodeCharacterTokenizer(
+            name="unicode_character_tokenizer_config_gen",
+            lowercase=False,
+            sequence_length=8,
+            normalization_form="NFC",
+            errors="ignore",
+            replacement_char=0,
+        )
+        exp_config = {
+            "dtype": "int32",
+            "errors": "ignore",
+            "lowercase": False,
+            "name": "unicode_character_tokenizer_config_gen",
+            "normalization_form": "NFC",
+            "replacement_char": 0,
+            "sequence_length": 8,
+            "input_encoding": "UTF-8",
+            "output_encoding": "UTF-8",
+            "trainable": True,
+        }
+        self.assertEqual(tokenizer.get_config(), exp_config)
+
+        tokenize_different_encoding = UnicodeCharacterTokenizer(
+            name="unicode_character_tokenizer_config_gen",
+            lowercase=False,
+            sequence_length=8,
+            errors="ignore",
+            replacement_char=0,
+            input_encoding="UTF-16",
+            output_encoding="UTF-16",
+        )
+        exp_config_different_encoding = {
+            "dtype": "int32",
+            "errors": "ignore",
+            "lowercase": False,
+            "name": "unicode_character_tokenizer_config_gen",
+            "normalization_form": None,
+            "replacement_char": 0,
+            "sequence_length": 8,
+            "input_encoding": "UTF-16",
+            "output_encoding": "UTF-16",
+            "trainable": True,
+        }
+        self.assertEqual(
+            tokenize_different_encoding.get_config(),
+            exp_config_different_encoding,
+        )
+
+    def test_saving(self):
+        input_data = tf.constant(["ninjas and samurais", "time travel"])
+
+        tokenizer = UnicodeCharacterTokenizer(
+            name="unicode_character_tokenizer_config_gen",
+            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),
+        )