Repository: keras-team/keras-nlp

Issue:
Add a utility to detokenize as lists of python strings
Going from tf.Tensor strings back to python string is confusing. We should add a utility to do that to the base tokenizer class.

We can discuss the exact name of the PR, suggest `.detokenize_to_strings()` as a starting point.

Usages:
```
>>> keras_nlp.layers.ByteTokenizer().detokenize_to_strings([104, 101, 108, 108, 111])
'hello'
>>> keras_nlp.layers.ByteTokenizer().detokenize_to_strings([[104, 101, 108, 108, 111]])
['hello']
```



Discussion:
aflah02: Thanks for the review! I'll get to work on this
aflah02: @mattdangerw 
I have some questions regarding testing the function as you had described [here](https://github.com/keras-team/keras-nlp/pull/119#issuecomment-1100578869)
Essentially we need to test whether it outputs the right string over byte strings right when used with something like the Unicode Character Tokenizer . For those tests I don't get how a simple subclass can help us test this out?
Won't it be better to test it with the Unicode Character Tokenizer Class and other classes directly?
mattdangerw: The test cases should live on the base class, as we are writing unit tests for functionality on the base class. But the tokenizer base class is not a usable thing standalone, it needs to be subclassed. So in tokenizer_test.py you could define a quick helper. For example

```
class PassThroughTokenizer(Tokenizer):
    __test__ = False  # for pytest

    def tokenize(self, inputs):
        return inputs

    def detokenize(self, inputs):
        return inputs
```

You could use that to fully test the logic of your helper function (by calling detokenize_to_string with raggeds, dense, scalar, etc.).
aflah02: @mattdangerw Sorry if I missed the message but I can't seem to find any confirmation for the naming of the method. Has it been decided?

Patch:
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index fa08c85f..460af4d4 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -47,6 +47,21 @@ request gets approved by the reviewer.
 
 Once the pull request is approved, a team member will take care of merging.
 
+## Developing on Windows
+
+For Windows development, we recommend using WSL (Windows Subsystem for Linux),
+so you can run the shell scripts in this repository. We will not support
+Windows Shell/PowerShell. You can refer
+[to these instructions](https://docs.microsoft.com/en-us/windows/wsl/install)
+for WSL installation.
+
+Note that if you are using Windows Subsystem for Linux (WSL), make sure you 
+clone the repo with Linux style LF line endings and change the default setting
+for line separator in your Text Editor before running the format
+or lint scripts. This is automatically done if you clone using git inside WSL.
+If there is conflict due to the line endings you might see an error
+like - `: invalid option`.
+
 ## Setup environment
 
 Python 3.7 or later is required.
@@ -112,12 +127,6 @@ If after running these the CI flow is still failing, try updating `flake8`,
 `isort` and `black`. This can be done by running `pip install --upgrade black`,
 `pip install --upgrade flake8`, and `pip install --upgrade isort`.
 
-Note that if you are using Windows Subsystem for Linux (WSL), make sure you 
-clone the repo with Linux style LF line endings before running the format
-or lint scripts. This is automatically done if you clone using git inside WSL.
-If there is conflict due to the line endings you might see an error
-like - `: invalid option`
-
 ## Community Guidelines
 
 This project follows 
diff --git a/examples/bert/bert_model.py b/examples/bert/bert_model.py
index 4569376a..99df0440 100644
--- a/examples/bert/bert_model.py
+++ b/examples/bert/bert_model.py
@@ -13,7 +13,11 @@
 # limitations under the License.
 """Bert model and layer implementations.
 
-We should work to replace this with components from the keras-nlp library.
+These components come from the tensorflow official model repository for BERT:
+https://github.com/tensorflow/models/tree/master/official/nlp/modeling
+
+This is to get us into a testable state. We should work to replace all of these
+components with components from the keras-nlp library.
 """
 
 import tensorflow as tf
@@ -484,7 +488,7 @@ class BertModel(keras.Model):
 
         self._position_embedding_layer = keras_nlp.layers.PositionEmbedding(
             initializer=initializer,
-            max_length=max_sequence_length,
+            sequence_length=max_sequence_length,
             name="position_embedding",
         )
 
diff --git a/examples/bert/create_pretraining_data.py b/examples/bert/create_pretraining_data.py
index 859503ea..7bf07d37 100644
--- a/examples/bert/create_pretraining_data.py
+++ b/examples/bert/create_pretraining_data.py
@@ -24,6 +24,9 @@ By default, it will duplicate the input data 10 times with different masks and
 sentence pairs, as will the original paper. So a 20gb source of wikipedia and
 bookscorpus will result in a 400gb dataset.
 
+This script is adapted from the original BERT respository:
+https://github.com/google-research/bert/blob/master/create_pretraining_data.py
+
 Usage:
 python create_pretraining_data.py \
     --input_files ~/datasets/bert-sentence-split-data/shard_0.txt \
diff --git a/keras_nlp/__init__.py b/keras_nlp/__init__.py
index 06afe693..edbe34c3 100644
--- a/keras_nlp/__init__.py
+++ b/keras_nlp/__init__.py
@@ -15,5 +15,6 @@
 from keras_nlp import layers
 from keras_nlp import metrics
 from keras_nlp import tokenizers
+from keras_nlp import utils
 
-__version__ = "0.1.1"
+__version__ = "0.2.0-dev.1"
diff --git a/keras_nlp/layers/__init__.py b/keras_nlp/layers/__init__.py
index 4d1f5ddf..ec1ea1ad 100644
--- a/keras_nlp/layers/__init__.py
+++ b/keras_nlp/layers/__init__.py
@@ -13,8 +13,9 @@
 # limitations under the License.
 
 from keras_nlp.layers.fnet_encoder import FNetEncoder
+from keras_nlp.layers.mlm_head import MLMHead
+from keras_nlp.layers.mlm_mask_generator import MLMMaskGenerator
 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.token_and_position_embedding import (
     TokenAndPositionEmbedding,
diff --git a/keras_nlp/layers/mlm_head.py b/keras_nlp/layers/mlm_head.py
new file mode 100644
index 00000000..09644205
--- /dev/null
+++ b/keras_nlp/layers/mlm_head.py
@@ -0,0 +1,213 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Masked Language Model (MLM) head."""
+
+import tensorflow as tf
+from tensorflow import keras
+
+
+class MLMHead(keras.layers.Layer):
+    """Masked Language Model (MLM) head.
+
+    This layer takes two inputs:
+     - `inputs`: which should be a tensor of encoded tokens with shape
+            `(batch_size, sequence_length, encoding_dim)`.
+     - `mask_positions`: which should be a tensor of integer positions to
+            predict with shape `(batch_size, masks_per_sequence)`.
+
+    The token encodings should usually be the last output of an encoder model,
+    and mask positions should be the interger positions you would like to
+    predict for the MLM task.
+
+    The layer will first gather the token encodings at the mask positions. These
+    gathered tokens will be passed through a dense layer the same size as
+    encoding dimension, then transformed to predictions the same size as the
+    input vocabulary. This layer will produce a single output with shape
+    `(batch_size, masks_per_sequence, vocabulary_size)`, which can be used to
+    compute an MLM loss function.
+
+    This layer is often be paired with `keras_nlp.layers.MLMMaskGenerator`,
+    which will help prepare inputs for the MLM task.
+
+    Args:
+        vocabulary_size: The total size of the vocabulary for predictions.
+        embedding_weights: Optional. The weights of the word embedding used
+            to transform input token ids. The transpose of this weight matrix
+            will be used to project a token embedding vector to a prediction
+            over all input words, as described in [1].
+        intermediate_activation: The activation function of inner dense layer.
+        activation: The activation function for the outputs of the layer.
+            Usually either `None` (return logits), or `"softmax"`
+            (return probabilities).
+        layer_norm_epsilon: float, defaults to 1e-5. The epsilon value in layer
+            normalization components.
+        kernel_initializer: string or tf.keras.initializers initializer,
+            defaults to "glorot_uniform". The kernel initializer for
+            the dense and multiheaded attention layers.
+        bias_initializer: string or tf.keras.initializers initializer,
+            defaults to "zeros". The bias initializer for
+            the dense and multiheaded attention layers.
+        name: string, defaults to None. The name of the layer.
+        **kwargs: other keyword arguments.
+
+    Examples:
+
+    ```python
+    batch_size = 32
+    vocab_size = 100
+    encoding_size = 32
+    seq_length = 50
+    mask_length = 10
+
+    # Generate a random encoding.
+    encoded_tokens = tf.random.normal([batch_size, seq_length, encoding_size])
+    # Generate random positions and labels
+    mask_positions = tf.random.uniform(
+        [batch_size, mask_length], maxval=seq_length, dtype="int32"
+    )
+    mask_ids = tf.random.uniform(
+        [batch_size, mask_length], maxval=vocab_size, dtype="int32"
+    )
+
+    # Predict an output word for each masked input token.
+    mask_preds = keras_nlp.layers.MLMHead(
+        vocabulary_size=vocab_size,
+        activation="softmax",
+    )(encoded_tokens, mask_positions=mask_positions)
+    # Calculate a loss.
+    keras.losses.sparse_categorical_crossentropy(mask_ids, mask_preds)
+    ```
+
+    References:
+        [1] [Press and Wolf, 2016](https://arxiv.org/abs/1608.05859)
+    """
+
+    def __init__(
+        self,
+        vocabulary_size=None,
+        embedding_weights=None,
+        intermediate_activation="relu",
+        activation=None,
+        layer_norm_epsilon=1e-05,
+        kernel_initializer="glorot_uniform",
+        bias_initializer="zeros",
+        name=None,
+        **kwargs,
+    ):
+        super().__init__(name=name, **kwargs)
+
+        self.vocabulary_size = vocabulary_size
+        self.embedding_weights = embedding_weights
+        self.intermediate_activation = keras.activations.get(
+            intermediate_activation
+        )
+        self.activation = keras.activations.get(activation)
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.kernel_initializer = keras.initializers.get(kernel_initializer)
+        self.bias_initializer = keras.initializers.get(bias_initializer)
+        self._built = False
+
+        if vocabulary_size is None and embedding_weights is None:
+            raise ValueError(
+                "One of `vocabulary_size` or `embedding_weights` must be set. "
+                "Received: `vocabulary_size=None`, `embedding_weights=None`"
+            )
+
+        if embedding_weights is not None:
+            shape = embedding_weights.shape
+            if vocabulary_size is not None and vocabulary_size != shape[0]:
+                raise ValueError(
+                    "`vocabulary_size` should match the first dimension of the "
+                    "shape of `embedding_weights`. Received: "
+                    f"`vocabulary_size={vocabulary_size}`, "
+                    f"`embedding_weights.shape={shape}`"
+                )
+            self.vocabulary_size = shape[0]
+
+    def _build(self, input_shape):
+        # Create layers based on input shape.
+        self._built = True
+        feature_size = input_shape[-1]
+
+        self._dense = keras.layers.Dense(
+            feature_size,
+            activation=self.intermediate_activation,
+            kernel_initializer=self.kernel_initializer,
+            bias_initializer=self.bias_initializer,
+        )
+        self._layer_norm = tf.keras.layers.LayerNormalization(
+            epsilon=self.layer_norm_epsilon,
+        )
+        if self.embedding_weights is None:
+            self._kernel = self.add_weight(
+                name="output_kernel",
+                shape=[feature_size, self.vocabulary_size],
+                initializer=self.kernel_initializer,
+                dtype=self.dtype,
+            )
+        self._bias = self.add_weight(
+            name="output_bias",
+            shape=[self.vocabulary_size],
+            initializer=self.bias_initializer,
+            dtype=self.dtype,
+        )
+
+    def call(self, inputs, mask_positions):
+        if not self._built:
+            self._build(inputs.shape)
+
+        # Gather the encoded tokens at the masked indices.
+        x = tf.gather(inputs, mask_positions, axis=1, batch_dims=1)
+
+        # Apply a trainable linear transformation and a layer norm.
+        x = self._dense(x)
+        x = self._layer_norm(x)
+
+        # Transform encodings to vocabulary_size predictions.
+        if self.embedding_weights is None:
+            outputs = tf.matmul(x, self._kernel)
+        else:
+            outputs = tf.matmul(
+                x,
+                tf.cast(self.embedding_weights, self.compute_dtype),
+                transpose_b=True,
+            )
+        outputs = outputs + self._bias
+
+        # Apply a final activation.
+        if self.activation is not None:
+            outputs = self.activation(outputs)
+
+        return outputs
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "vocabulary_size": self.vocabulary_size,
+                "intermediate_activation": keras.activations.serialize(
+                    self.intermediate_activation
+                ),
+                "activation": keras.activations.serialize(self.activation),
+                "layer_norm_epsilon": self.layer_norm_epsilon,
+                "kernel_initializer": keras.initializers.serialize(
+                    self.kernel_initializer
+                ),
+                "bias_initializer": keras.initializers.serialize(
+                    self.bias_initializer
+                ),
+            }
+        )
+        return config
diff --git a/keras_nlp/layers/preprocessing/mlm_mask_generator.py b/keras_nlp/layers/mlm_mask_generator.py
similarity index 98%
rename from keras_nlp/layers/preprocessing/mlm_mask_generator.py
rename to keras_nlp/layers/mlm_mask_generator.py
index 3c2bbd69..85b207a4 100644
--- a/keras_nlp/layers/preprocessing/mlm_mask_generator.py
+++ b/keras_nlp/layers/mlm_mask_generator.py
@@ -73,13 +73,13 @@ class MLMMaskGenerator(keras.layers.Layer):
     Examples:
 
     Basic usage.
-    >>> masker = keras_nlp.layers.preprocessing.MLMMaskGenerator( \
+    >>> masker = keras_nlp.layers.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( \
+    >>> masker = keras_nlp.layers.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]]))
diff --git a/keras_nlp/layers/position_embedding.py b/keras_nlp/layers/position_embedding.py
index 5d3b01f8..9bf8a615 100644
--- a/keras_nlp/layers/position_embedding.py
+++ b/keras_nlp/layers/position_embedding.py
@@ -32,7 +32,7 @@ class PositionEmbedding(keras.layers.Layer):
     corresponds to the sequence, that is, the penultimate dimension.
 
     Args:
-        max_length: The maximum length of the dynamic sequence.
+        sequence_length: The maximum length of the dynamic sequence.
         initializer: The initializer to use for the embedding weights. Defaults
             to "glorot_uniform".
         seq_axis: The axis of the input tensor where we add the embeddings.
@@ -43,7 +43,7 @@ class PositionEmbedding(keras.layers.Layer):
         input_dim=vocab_size, output_dim=embed_dim
     )
     position_embeddings = keras_nlp.layers.PositionEmbedding(
-        max_length=max_length
+        sequence_length=sequence_length
     )
 
     embedded_tokens = token_embeddings(inputs)
@@ -58,23 +58,23 @@ class PositionEmbedding(keras.layers.Layer):
 
     def __init__(
         self,
-        max_length,
+        sequence_length,
         initializer="glorot_uniform",
         **kwargs,
     ):
         super().__init__(**kwargs)
-        if max_length is None:
+        if sequence_length is None:
             raise ValueError(
-                "`max_length` must be an Integer, received `None`."
+                "`sequence_length` must be an Integer, received `None`."
             )
-        self.max_length = int(max_length)
+        self.sequence_length = int(sequence_length)
         self.initializer = keras.initializers.get(initializer)
 
     def get_config(self):
         config = super().get_config()
         config.update(
             {
-                "max_length": self.max_length,
+                "sequence_length": self.sequence_length,
                 "initializer": keras.initializers.serialize(self.initializer),
             }
         )
@@ -84,7 +84,7 @@ class PositionEmbedding(keras.layers.Layer):
         feature_size = input_shape[-1]
         self.position_embeddings = self.add_weight(
             "embeddings",
-            shape=[self.max_length, feature_size],
+            shape=[self.sequence_length, feature_size],
             initializer=self.initializer,
             trainable=True,
         )
diff --git a/keras_nlp/layers/preprocessing/__init__.py b/keras_nlp/layers/preprocessing/__init__.py
deleted file mode 100644
index c8059499..00000000
--- a/keras_nlp/layers/preprocessing/__init__.py
+++ /dev/null
@@ -1,15 +0,0 @@
-# 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/token_and_position_embedding.py b/keras_nlp/layers/token_and_position_embedding.py
index 035ddd02..03c5f8c6 100644
--- a/keras_nlp/layers/token_and_position_embedding.py
+++ b/keras_nlp/layers/token_and_position_embedding.py
@@ -28,7 +28,7 @@ class TokenAndPositionEmbedding(keras.layers.Layer):
     Args:
         vocabulary_size: The size of the vocabulary (should be no larger
             than 999)
-        max_length: The maximum length of input sequence
+        sequence_length: The maximum length of input sequence
         embedding_dim: The output dimension of the embedding layer
         embeddings_initializer: The initializer to use for the Embedding
             Layers
@@ -49,7 +49,7 @@ class TokenAndPositionEmbedding(keras.layers.Layer):
     inputs = tf.keras.Input(shape=(seq_length,))
     embedding_layer = keras_nlp.layers.TokenAndPositionEmbedding(
         vocabulary_size=vocab_size,
-        max_length=max_length,
+        sequence_length=seq_length,
         embedding_dim=embed_dim,
     )
     outputs = embedding_layer(inputs)
@@ -59,7 +59,7 @@ class TokenAndPositionEmbedding(keras.layers.Layer):
     def __init__(
         self,
         vocabulary_size,
-        max_length,
+        sequence_length,
         embedding_dim,
         embeddings_initializer="glorot_uniform",
         mask_zero=False,
@@ -70,16 +70,16 @@ class TokenAndPositionEmbedding(keras.layers.Layer):
             raise ValueError(
                 "`vocabulary_size` must be an Integer, received `None`."
             )
-        if max_length is None:
+        if sequence_length is None:
             raise ValueError(
-                "`max_length` must be an Integer, received `None`."
+                "`sequence_length` must be an Integer, received `None`."
             )
         if embedding_dim is None:
             raise ValueError(
                 "`embedding_dim` must be an Integer, received `None`."
             )
         self.vocabulary_size = int(vocabulary_size)
-        self.max_length = int(max_length)
+        self.sequence_length = int(sequence_length)
         self.embedding_dim = int(embedding_dim)
         self.token_embedding = keras.layers.Embedding(
             vocabulary_size,
@@ -88,7 +88,7 @@ class TokenAndPositionEmbedding(keras.layers.Layer):
             mask_zero=mask_zero,
         )
         self.position_embedding = keras_nlp.layers.PositionEmbedding(
-            max_length=max_length,
+            sequence_length=sequence_length,
             initializer=embeddings_initializer,
         )
         self.supports_masking = self.token_embedding.supports_masking
@@ -98,7 +98,7 @@ class TokenAndPositionEmbedding(keras.layers.Layer):
         config.update(
             {
                 "vocabulary_size": self.vocabulary_size,
-                "max_length": self.max_length,
+                "sequence_length": self.sequence_length,
                 "embedding_dim": self.embedding_dim,
                 "embeddings_initializer": keras.initializers.serialize(
                     self.token_embedding.embeddings_initializer
diff --git a/keras_nlp/tokenizers/unicode_character_tokenizer.py b/keras_nlp/tokenizers/unicode_character_tokenizer.py
index ade085e3..3d21e9ca 100644
--- a/keras_nlp/tokenizers/unicode_character_tokenizer.py
+++ b/keras_nlp/tokenizers/unicode_character_tokenizer.py
@@ -175,9 +175,9 @@ class UnicodeCharacterTokenizer(tokenizer.Tokenizer):
             kwargs["dtype"] = tf.int32
         else:
             dtype = tf.dtypes.as_dtype(kwargs["dtype"])
-            if not dtype.is_integer and dtype != tf.string:
+            if not dtype.is_integer:
                 raise ValueError(
-                    "Output dtype must be an integer type of a string. "
+                    "Output dtype must be an integer type. "
                     f"Received: dtype={dtype}"
                 )
 
@@ -251,6 +251,7 @@ class UnicodeCharacterTokenizer(tokenizer.Tokenizer):
             replacement_char=self.replacement_char,
             input_encoding=self.input_encoding,
         )
+        tokens = tf.cast(tokens, self.compute_dtype)
 
         if self.sequence_length:
             output_shape = tokens.shape.as_list()
diff --git a/keras_nlp/tokenizers/word_piece_tokenizer.py b/keras_nlp/tokenizers/word_piece_tokenizer.py
index 35f2bf81..cc1766ec 100644
--- a/keras_nlp/tokenizers/word_piece_tokenizer.py
+++ b/keras_nlp/tokenizers/word_piece_tokenizer.py
@@ -80,7 +80,7 @@ class WordPieceTokenizer(tokenizer.Tokenizer):
 
     If a more custom pre-tokenization step is desired, the layer can be
     configured to apply only the strict WordPiece algorithm by passing
-    `lowercase=False`, `strip_accents=False` and `split_pattern=None`. In
+    `lowercase=False`, `strip_accents=False` and `split=False`. In
     this case, inputs should be pre-split string tensors or ragged tensors.
 
     By default, the layer will output a `tf.RaggedTensor` where the last
@@ -101,10 +101,11 @@ class WordPieceTokenizer(tokenizer.Tokenizer):
             tokenization.
         strip_accents: If true, all accent marks will be removed from text
             before tokenization.
-        split_pattern: A regex pattern to match delimiters to split, or None
-            indicating that the input is pre-split and no splitting should be
-            performed. By default, all whitespace and punctuation marks will
-            be split on.
+        split: If true, input will be split according to `split_pattern`
+            and `keep_pattern`. If false, input should be split before calling
+            the layer.
+        split_pattern: A regex pattern to match delimiters to split. By default,
+            all whitespace and punctuation marks will be split on.
         keep_pattern: A regex pattern of delimiters contained in the
             `split_pattern` of delimeters that should be kept as independent
             tokens. By default, all punctuation marks will be kept as tokens.
@@ -167,8 +168,9 @@ class WordPieceTokenizer(tokenizer.Tokenizer):
         sequence_length: int = None,
         lowercase: bool = True,
         strip_accents: bool = True,
-        split_pattern: str = WHITESPACE_AND_PUNCTUATION_REGEX,
-        keep_pattern: str = PUNCTUATION_REGEX,
+        split: bool = True,
+        split_pattern: str = None,
+        keep_pattern: str = None,
         suffix_indicator: str = "##",
         oov_token: str = "[UNK]",
         **kwargs,
@@ -198,15 +200,33 @@ class WordPieceTokenizer(tokenizer.Tokenizer):
                 "Vocabulary must be an file path or list of terms. "
                 f"Received: vocabulary={vocabulary}"
             )
+        if oov_token is None:
+            raise ValueError("`oov_token` cannot be None.")
+
+        if split_pattern is None:
+            split_pattern = WHITESPACE_AND_PUNCTUATION_REGEX
+
+        if keep_pattern is None:
+            keep_pattern = PUNCTUATION_REGEX
 
         self.sequence_length = sequence_length
         self.lowercase = lowercase
         self.strip_accents = strip_accents
+        self.split = split
         self.split_pattern = split_pattern
         self.keep_pattern = keep_pattern
         self.suffix_indicator = suffix_indicator
         self.oov_token = oov_token
 
+        if oov_token not in self.vocabulary:
+            raise RuntimeError(
+                f'Cannot find `oov_token="{self.oov_token}"` in the '
+                "vocabulary.\n"
+                "You can either update the vocabulary to include "
+                f'`"{self.oov_token}"`, or pass a different value for '
+                "the `oov_token` argument when creating the tokenizer."
+            )
+
         self._fast_word_piece = tf_text.FastWordpieceTokenizer(
             vocab=self.vocabulary,
             token_out_type=self.compute_dtype,
@@ -246,6 +266,7 @@ class WordPieceTokenizer(tokenizer.Tokenizer):
                 "sequence_length": self.sequence_length,
                 "lowercase": self.lowercase,
                 "strip_accents": self.strip_accents,
+                "split": self.split,
                 "split_pattern": self.split_pattern,
                 "keep_pattern": self.keep_pattern,
                 "suffix_indicator": self.suffix_indicator,
@@ -269,7 +290,7 @@ class WordPieceTokenizer(tokenizer.Tokenizer):
             inputs = tf_text.normalize_utf8(inputs, "NFD")
             # Remove the accent marks.
             inputs = tf.strings.regex_replace(inputs, r"\p{Mn}", "")
-        if self.split_pattern:
+        if self.split:
             inputs = tf_text.regex_split(
                 inputs,
                 delim_regex_pattern=self.split_pattern,
diff --git a/keras_nlp/utils/tensor_utils.py b/keras_nlp/utils/tensor_utils.py
new file mode 100644
index 00000000..26fc815f
--- /dev/null
+++ b/keras_nlp/utils/tensor_utils.py
@@ -0,0 +1,47 @@
+# 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
+
+
+def _decode_strings_to_utf8(inputs):
+    """Recursively decodes to list of strings with 'utf-8' encoding."""
+    if isinstance(inputs, bytes):
+        # Handles the case when the input is a scalar string.
+        return inputs.decode("utf-8")
+    else:
+        # Recursively iterate when input is a list.
+        return [_decode_strings_to_utf8(x) for x in inputs]
+
+
+def tensor_to_string_list(inputs):
+    """Detokenize and convert tensor to nested lists of python strings.
+
+    This is a convenience method which converts each byte string to a python
+    string.
+
+    Args:
+        inputs: Input tensor, or dict/list/tuple of input tensors.
+        *args: Additional positional arguments.
+        **kwargs: Additional keyword arguments.
+    """
+    if not isinstance(inputs, (tf.RaggedTensor, tf.Tensor)):
+        inputs = tf.convert_to_tensor(inputs)
+    if isinstance(inputs, tf.RaggedTensor):
+        list_outputs = inputs.to_list()
+    elif isinstance(inputs, tf.Tensor):
+        list_outputs = inputs.numpy()
+        if inputs.shape.rank != 0:
+            list_outputs = list_outputs.tolist()
+    return _decode_strings_to_utf8(list_outputs)
diff --git a/setup.py b/setup.py
index e05c744e..bb458e4c 100644
--- a/setup.py
+++ b/setup.py
@@ -29,7 +29,7 @@ setup(
     ),
     long_description=README,
     long_description_content_type="text/markdown",
-    version="0.1.1",
+    version="0.2.0-dev.1",
     url="https://github.com/keras-team/keras-nlp",
     author="Keras team",
     author_email="keras-nlp@google.com",


Test patch:
diff --git a/keras_nlp/layers/mlm_head_test.py b/keras_nlp/layers/mlm_head_test.py
new file mode 100644
index 00000000..0271d915
--- /dev/null
+++ b/keras_nlp/layers/mlm_head_test.py
@@ -0,0 +1,179 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for Transformer Encoder."""
+
+import os
+
+import tensorflow as tf
+from tensorflow import keras
+
+from keras_nlp.layers import mlm_head
+
+
+class MLMHeadTest(tf.test.TestCase):
+    def test_valid_call(self):
+        head = mlm_head.MLMHead(
+            vocabulary_size=100,
+            activation="softmax",
+        )
+        encoded_tokens = keras.Input(shape=(10, 16))
+        positions = keras.Input(shape=(5,), dtype="int32")
+        outputs = head(encoded_tokens, mask_positions=positions)
+        model = keras.Model((encoded_tokens, positions), outputs)
+
+        token_data = tf.random.uniform(shape=(4, 10, 16))
+        position_data = tf.random.uniform(
+            shape=(4, 5), maxval=10, dtype="int32"
+        )
+        model((token_data, position_data))
+
+    def test_valid_call_with_embedding_weights(self):
+        embedding = keras.layers.Embedding(100, 16)
+        embedding.build((4, 10))
+        head = mlm_head.MLMHead(
+            vocabulary_size=100,
+            embedding_weights=embedding.embeddings,
+            activation="softmax",
+        )
+        encoded_tokens = keras.Input(shape=(10, 16))
+        positions = keras.Input(shape=(5,), dtype="int32")
+        outputs = head(encoded_tokens, mask_positions=positions)
+        model = keras.Model((encoded_tokens, positions), outputs)
+        token_data = tf.random.uniform(shape=(4, 10, 16))
+        position_data = tf.random.uniform(
+            shape=(4, 5), maxval=10, dtype="int32"
+        )
+        model((token_data, position_data))
+
+    def test_get_config_and_from_config(self):
+        head = mlm_head.MLMHead(
+            vocabulary_size=100,
+            kernel_initializer="HeNormal",
+            bias_initializer="Zeros",
+            activation="softmax",
+        )
+
+        config = head.get_config()
+
+        expected_params = {
+            "vocabulary_size": 100,
+            "kernel_initializer": keras.initializers.serialize(
+                keras.initializers.HeNormal()
+            ),
+            "bias_initializer": keras.initializers.serialize(
+                keras.initializers.Zeros()
+            ),
+            "activation": keras.activations.serialize(
+                keras.activations.softmax
+            ),
+        }
+
+        self.assertEqual(config, {**config, **expected_params})
+
+        restored = mlm_head.MLMHead.from_config(config)
+        restored_config = restored.get_config()
+
+        self.assertEqual(
+            restored_config, {**restored_config, **expected_params}
+        )
+        self.assertEqual(restored_config, config)
+
+    def test_value_error_when_neither_embedding_or_vocab_size_set(self):
+        with self.assertRaises(ValueError):
+            mlm_head.MLMHead()
+
+    def test_value_error_when_vocab_size_mismatch(self):
+        embedding = keras.layers.Embedding(100, 16)
+        embedding.build((4, 10))
+        with self.assertRaises(ValueError):
+            mlm_head.MLMHead(
+                vocabulary_size=101,
+                embedding_weights=embedding.embeddings,
+            )
+
+    def test_one_train_step(self):
+        head = mlm_head.MLMHead(
+            vocabulary_size=100,
+        )
+        encoded_tokens = keras.Input(shape=(10, 16))
+        positions = keras.Input(shape=(5,), dtype="int32")
+        outputs = head(encoded_tokens, mask_positions=positions)
+        model = keras.Model((encoded_tokens, positions), outputs)
+
+        token_data = tf.random.uniform(shape=(4, 10, 16))
+        position_data = tf.random.uniform(
+            shape=(4, 5), maxval=10, dtype="int32"
+        )
+        label_data = tf.random.uniform(shape=(4, 5), maxval=100, dtype="int32")
+
+        loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(
+            from_logits=False
+        )
+        optimizer = keras.optimizers.Adam()
+        with tf.GradientTape() as tape:
+            model((token_data, position_data))
+            pred = model((token_data, position_data))
+            loss = loss_fn(label_data, pred)
+        grad = tape.gradient(loss, model.trainable_variables)
+        self.assertGreater(len(grad), 1)
+        optimizer.apply_gradients(zip(grad, model.trainable_variables))
+
+    def test_checkpointing(self):
+        head1 = mlm_head.MLMHead(
+            vocabulary_size=100,
+            activation="softmax",
+        )
+        head2 = mlm_head.MLMHead(
+            vocabulary_size=100,
+            activation="softmax",
+        )
+        token_data = tf.random.uniform(shape=(4, 10, 16))
+        position_data = tf.random.uniform(
+            shape=(4, 5), maxval=10, dtype="int32"
+        )
+        # The weights of head1 and head2 are different.
+        head1_output = head1(token_data, mask_positions=position_data)
+        head2_output = head2(token_data, mask_positions=position_data)
+        self.assertNotAllClose(head1_output, head2_output)
+
+        checkpoint = tf.train.Checkpoint(head1)
+        checkpoint2 = tf.train.Checkpoint(head2)
+        save_path = checkpoint.save(self.get_temp_dir())
+        checkpoint2.restore(save_path)
+
+        head1_output = head1(token_data, mask_positions=position_data)
+        head2_output = head2(token_data, mask_positions=position_data)
+        self.assertAllClose(head1_output, head2_output)
+
+    def test_saving_model(self):
+        head = mlm_head.MLMHead(
+            vocabulary_size=100,
+            activation="softmax",
+        )
+        encoded_tokens = keras.Input(shape=(10, 16))
+        positions = keras.Input(shape=(5,), dtype="int32")
+        outputs = head(encoded_tokens, mask_positions=positions)
+        model = keras.Model((encoded_tokens, positions), outputs)
+
+        token_data = tf.random.uniform(shape=(4, 10, 16))
+        position_data = tf.random.uniform(
+            shape=(4, 5), maxval=10, dtype="int32"
+        )
+        model_output = model((token_data, position_data))
+        save_path = os.path.join(self.get_temp_dir(), "model")
+        model.save(save_path)
+        restored = keras.models.load_model(save_path)
+
+        restored_output = restored((token_data, position_data))
+        self.assertAllClose(model_output, restored_output)
diff --git a/keras_nlp/layers/preprocessing/mlm_mask_generator_test.py b/keras_nlp/layers/mlm_mask_generator_test.py
similarity index 99%
rename from keras_nlp/layers/preprocessing/mlm_mask_generator_test.py
rename to keras_nlp/layers/mlm_mask_generator_test.py
index 492a271c..342fa6f0 100644
--- a/keras_nlp/layers/preprocessing/mlm_mask_generator_test.py
+++ b/keras_nlp/layers/mlm_mask_generator_test.py
@@ -14,7 +14,7 @@
 
 import tensorflow as tf
 
-from keras_nlp.layers.preprocessing.mlm_mask_generator import MLMMaskGenerator
+from keras_nlp.layers.mlm_mask_generator import MLMMaskGenerator
 
 
 class MLMMaskGeneratorTest(tf.test.TestCase):
diff --git a/keras_nlp/layers/position_embedding_test.py b/keras_nlp/layers/position_embedding_test.py
index cd39f4fc..7d2741be 100644
--- a/keras_nlp/layers/position_embedding_test.py
+++ b/keras_nlp/layers/position_embedding_test.py
@@ -34,7 +34,7 @@ class PositionEmbeddingLayerTest(tf.test.TestCase):
         sequence_length = 21
         feature_size = 30
         test_layer = position_embedding.PositionEmbedding(
-            max_length=sequence_length
+            sequence_length=sequence_length
         )
         input_tensor = tf.keras.Input(shape=(sequence_length, feature_size))
         output_tensor = test_layer(input_tensor)
@@ -51,7 +51,7 @@ class PositionEmbeddingLayerTest(tf.test.TestCase):
         sequence_length = 21
         feature_size = 30
         test_layer = position_embedding.PositionEmbedding(
-            max_length=sequence_length
+            sequence_length=sequence_length
         )
         input_tensor = tf.keras.Input(
             shape=(feature_size, sequence_length, feature_size)
@@ -75,7 +75,7 @@ class PositionEmbeddingLayerTest(tf.test.TestCase):
         sequence_length = 21
         feature_size = 30
         test_layer = position_embedding.PositionEmbedding(
-            max_length=sequence_length, dtype="float16"
+            sequence_length=sequence_length, dtype="float16"
         )
         input_tensor = tf.keras.Input(shape=(sequence_length, feature_size))
         output_tensor = test_layer(input_tensor)
@@ -91,7 +91,7 @@ class PositionEmbeddingLayerTest(tf.test.TestCase):
         max_sequence_length = 40
         feature_size = 30
         test_layer = position_embedding.PositionEmbedding(
-            max_length=max_sequence_length
+            sequence_length=max_sequence_length
         )
         # Create a 3-dimensional input (the first dimension is implicit).
         input_tensor = tf.keras.Input(shape=(None, feature_size))
@@ -107,7 +107,7 @@ class PositionEmbeddingLayerTest(tf.test.TestCase):
         max_sequence_length = 60
         feature_size = 30
         test_layer = position_embedding.PositionEmbedding(
-            max_length=max_sequence_length
+            sequence_length=max_sequence_length
         )
         # Create a 4-dimensional input (the first dimension is implicit).
         input_tensor = tf.keras.Input(shape=(None, None, feature_size))
@@ -123,7 +123,7 @@ class PositionEmbeddingLayerTest(tf.test.TestCase):
         max_sequence_length = 40
         feature_size = 30
         test_layer = position_embedding.PositionEmbedding(
-            max_length=max_sequence_length
+            sequence_length=max_sequence_length
         )
         # Create a 3-dimensional input (the first dimension is implicit).
         input_tensor = tf.keras.Input(shape=(None, feature_size))
@@ -148,7 +148,7 @@ class PositionEmbeddingLayerTest(tf.test.TestCase):
         max_sequence_length = 4
         feature_size = 3
         test_layer = position_embedding.PositionEmbedding(
-            max_length=max_sequence_length,
+            sequence_length=max_sequence_length,
             initializer=custom_init,
         )
         inputs = tf.keras.Input(shape=(max_sequence_length, feature_size))
@@ -172,7 +172,7 @@ class PositionEmbeddingLayerTest(tf.test.TestCase):
         max_sequence_length = 4
         feature_size = 2
         test_layer = position_embedding.PositionEmbedding(
-            max_length=max_sequence_length,
+            sequence_length=max_sequence_length,
             initializer=custom_init,
         )
         # Create a 3-dimensional ragged input (the first dimension is implicit).
@@ -209,7 +209,7 @@ class PositionEmbeddingLayerTest(tf.test.TestCase):
         max_sequence_length = 4
         feature_size = 2
         test_layer = position_embedding.PositionEmbedding(
-            max_length=max_sequence_length,
+            sequence_length=max_sequence_length,
             initializer=custom_init,
         )
         # Create a 4-dimensional ragged input (the first dimension is implicit).
@@ -254,7 +254,7 @@ class PositionEmbeddingLayerTest(tf.test.TestCase):
         max_sequence_length = 4
         feature_size = 3
         test_layer = position_embedding.PositionEmbedding(
-            max_length=max_sequence_length
+            sequence_length=max_sequence_length
         )
         inputs = tf.keras.Input(shape=(max_sequence_length, feature_size))
         outputs = test_layer(inputs)
@@ -283,11 +283,11 @@ class PositionEmbeddingLayerTest(tf.test.TestCase):
     def test_get_config_and_from_config(self):
         max_sequence_length = 40
         test_layer = position_embedding.PositionEmbedding(
-            max_length=max_sequence_length
+            sequence_length=max_sequence_length
         )
         config = test_layer.get_config()
         expected_config_subset = {
-            "max_length": max_sequence_length,
+            "sequence_length": max_sequence_length,
             "initializer": {
                 "class_name": "GlorotUniform",
                 "config": {"seed": None},
@@ -306,7 +306,7 @@ class PositionEmbeddingLayerTest(tf.test.TestCase):
         max_sequence_length = 4
         feature_size = 6
         test_layer = position_embedding.PositionEmbedding(
-            max_length=max_sequence_length
+            sequence_length=max_sequence_length
         )
         inputs = tf.keras.Input(shape=(max_sequence_length, feature_size))
         outputs = test_layer(inputs)
diff --git a/keras_nlp/layers/token_and_position_embedding_test.py b/keras_nlp/layers/token_and_position_embedding_test.py
index a9916ab9..cfaadef6 100644
--- a/keras_nlp/layers/token_and_position_embedding_test.py
+++ b/keras_nlp/layers/token_and_position_embedding_test.py
@@ -32,7 +32,7 @@ class TokenAndPositionEmbeddingTest(tf.test.TestCase):
     def test_get_config_and_from_config(self):
         token_and_position_embed = TokenAndPositionEmbedding(
             vocabulary_size=5,
-            max_length=10,
+            sequence_length=10,
             embedding_dim=32,
         )
 
@@ -62,7 +62,7 @@ class TokenAndPositionEmbeddingTest(tf.test.TestCase):
         embedding_dim = 3
         test_layer = TokenAndPositionEmbedding(
             vocabulary_size=vocabulary_size,
-            max_length=sequence_length,
+            sequence_length=sequence_length,
             embedding_dim=embedding_dim,
             embeddings_initializer=custom_embed_init,
         )
@@ -106,7 +106,7 @@ class TokenAndPositionEmbeddingTest(tf.test.TestCase):
         embedding_dim = 3
         test_layer = TokenAndPositionEmbedding(
             vocabulary_size=vocabulary_size,
-            max_length=sequence_length,
+            sequence_length=sequence_length,
             embedding_dim=embedding_dim,
             embeddings_initializer=custom_embed_init,
         )
@@ -124,7 +124,7 @@ class TokenAndPositionEmbeddingTest(tf.test.TestCase):
     def test_mask_propagation(self):
         test_layer = TokenAndPositionEmbedding(
             vocabulary_size=5,
-            max_length=4,
+            sequence_length=4,
             embedding_dim=3,
             mask_zero=True,
         )
@@ -139,7 +139,7 @@ class TokenAndPositionEmbeddingTest(tf.test.TestCase):
         embedding_dim = 3
         test_layer = TokenAndPositionEmbedding(
             vocabulary_size=vocabulary_size,
-            max_length=sequence_length,
+            sequence_length=sequence_length,
             embedding_dim=embedding_dim,
         )
         inputs = keras.Input(shape=(sequence_length,))
diff --git a/keras_nlp/tokenizers/word_piece_tokenizer_test.py b/keras_nlp/tokenizers/word_piece_tokenizer_test.py
index 9a54f857..fc3e6237 100644
--- a/keras_nlp/tokenizers/word_piece_tokenizer_test.py
+++ b/keras_nlp/tokenizers/word_piece_tokenizer_test.py
@@ -135,9 +135,7 @@ class WordPieceTokenizerTest(tf.test.TestCase):
     def test_no_spliting(self):
         input_data = ["t o k e n", "m i s s i n g", "t o k e n"]
         vocab_data = ["[UNK]", "t o k e n"]
-        tokenizer = WordPieceTokenizer(
-            vocabulary=vocab_data, split_pattern=None
-        )
+        tokenizer = WordPieceTokenizer(vocabulary=vocab_data, split=False)
         call_output = tokenizer(input_data)
         self.assertAllEqual(call_output, [1, 0, 1])
 
@@ -148,7 +146,7 @@ class WordPieceTokenizerTest(tf.test.TestCase):
             vocabulary=vocab_data,
             lowercase=False,
             strip_accents=False,
-            split_pattern=None,
+            split=False,
         )
         call_output = tokenizer(input_data)
         self.assertAllEqual(call_output, [1, 2, 3, 4, 5, 6])
@@ -223,3 +221,25 @@ class WordPieceTokenizerTest(tf.test.TestCase):
             model(input_data),
             restored_model(input_data),
         )
+
+    def test_no_oov_token_in_vocabulary(self):
+        vocab_data = ["qu", "@@ick", "br", "@@OWN", "fox"]
+        with self.assertRaises(RuntimeError):
+            WordPieceTokenizer(
+                vocabulary=vocab_data,
+            )
+
+        vocab_data = ["@UNK@", "qu", "@@ick", "br", "@@OWN", "fox"]
+        with self.assertRaises(RuntimeError):
+            WordPieceTokenizer(
+                vocabulary=vocab_data,
+            )
+
+        vocab_data = ["UNK", "qu", "@@ick", "br", "@@OWN", "fox"]
+        with self.assertRaises(RuntimeError):
+            WordPieceTokenizer(
+                vocabulary=vocab_data,
+            )
+
+        with self.assertRaises(ValueError):
+            WordPieceTokenizer(vocabulary=vocab_data, oov_token=None)
diff --git a/keras_nlp/utils/tensor_utils_test.py b/keras_nlp/utils/tensor_utils_test.py
new file mode 100644
index 00000000..2c9103f3
--- /dev/null
+++ b/keras_nlp/utils/tensor_utils_test.py
@@ -0,0 +1,33 @@
+# 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 tensor_utils import tensor_to_string_list
+
+
+class TensorToStringListTest(tf.test.TestCase):
+    def test_detokenize_to_strings_for_ragged(self):
+        input_data = tf.ragged.constant([["▀▁▂▃", "samurai"]])
+        detokenize_output = tensor_to_string_list(input_data)
+        self.assertAllEqual(detokenize_output, [["▀▁▂▃", "samurai"]])
+
+    def test_detokenize_to_strings_for_dense(self):
+        input_data = tf.constant([["▀▁▂▃", "samurai"]])
+        detokenize_output = tensor_to_string_list(input_data)
+        self.assertAllEqual(detokenize_output, [["▀▁▂▃", "samurai"]])
+
+    def test_detokenize_to_strings_for_scalar(self):
+        input_data = tf.constant("▀▁▂▃")
+        detokenize_output = tensor_to_string_list(input_data)
+        self.assertEqual(detokenize_output, "▀▁▂▃")