Repository: keras-team/keras-nlp

Issue:
Add a vocabulary_size argument to UnicodeCharacterTokenizer
We should add a `vocabulary_size` argument to the `UnicodeCharacterTokenizer` layer that clamps the output code point values to be in the range `[0, vocabulary_size)`.

We should also add a `vocabulary_size()` method as with other tokenizers that returns this value if it has been set, and return `None` otherwise.

Potential docstring:

```
        vocabulary_size: Set the vocabulary `vocabulary_size`,
            by clamping all codepoints to the range [0, vocabulary_size).
            Effectively this will make the `vocabulary_size - 1` id the
            the OOV token.
```

Discussion:
aflah02: @mattdangerw I think the PR is ready for review now. [Here's a demo](https://colab.research.google.com/gist/aflah02/5542e35a383940c5870996af0e55d4ef/vocablimit.ipynb)

Patch:
diff --git a/examples/bert/run_glue_finetuning.py b/examples/bert/run_glue_finetuning.py
index 0bc7ffdf..156bc601 100644
--- a/examples/bert/run_glue_finetuning.py
+++ b/examples/bert/run_glue_finetuning.py
@@ -16,6 +16,7 @@
 import json
 
 import datasets
+import keras_tuner
 import tensorflow as tf
 import tensorflow_text as tftext
 from absl import app
@@ -68,8 +69,6 @@ flags.DEFINE_integer("batch_size", 32, "The batch size.")
 
 flags.DEFINE_integer("epochs", 3, "The number of training epochs.")
 
-flags.DEFINE_float("learning_rate", 2e-5, "The initial learning rate for Adam.")
-
 flags.DEFINE_integer("max_seq_length", 128, "Maximum sequence length.")
 
 
@@ -167,9 +166,32 @@ class BertClassificationFinetuner(keras.Model):
         return self._logit_layer(outputs)
 
 
+class BertHyperModel(keras_tuner.HyperModel):
+    """Creates a hypermodel to help with the search space for finetuning."""
+
+    def __init__(self, bert_config):
+        self.bert_config = bert_config
+
+    def build(self, hp):
+        model = keras.models.load_model(FLAGS.saved_model_input, compile=False)
+        bert_config = self.bert_config
+        finetuning_model = BertClassificationFinetuner(
+            bert_model=model,
+            hidden_size=bert_config["hidden_size"],
+            num_classes=3 if FLAGS.task_name in ("mnli", "ax") else 2,
+        )
+        finetuning_model.compile(
+            optimizer=keras.optimizers.Adam(
+                learning_rate=hp.Choice("lr", [5e-5, 4e-5, 3e-5, 2e-5])
+            ),
+            loss="sparse_categorical_crossentropy",
+            metrics=["accuracy"],
+        )
+        return finetuning_model
+
+
 def main(_):
     print(f"Reading input model from {FLAGS.saved_model_input}")
-    model = keras.models.load_model(FLAGS.saved_model_input)
 
     vocab = []
     with open(FLAGS.vocab_file, "r") as vocab_file:
@@ -200,6 +222,7 @@ def main(_):
 
     # Read and preprocess GLUE task data.
     train_ds, test_ds, validation_ds = load_data(FLAGS.task_name)
+
     train_ds = train_ds.batch(FLAGS.batch_size).map(
         preprocess_data, num_parallel_calls=tf.data.AUTOTUNE
     )
@@ -210,18 +233,27 @@ def main(_):
         preprocess_data, num_parallel_calls=tf.data.AUTOTUNE
     )
 
-    finetuning_model = BertClassificationFinetuner(
-        bert_model=model,
-        hidden_size=bert_config["hidden_size"],
-        num_classes=3 if FLAGS.task_name in ("mnli", "ax") else 2,
-    )
-    finetuning_model.compile(
-        optimizer=keras.optimizers.Adam(learning_rate=FLAGS.learning_rate),
-        loss="sparse_categorical_crossentropy",
-        metrics=["accuracy"],
+    # Create a hypermodel object for a RandomSearch.
+    hypermodel = BertHyperModel(bert_config)
+
+    # Initialize the random search over the 4 learning rate parameters, for 4
+    # trials and 3 epochs for each trial.
+    tuner = keras_tuner.RandomSearch(
+        hypermodel=hypermodel,
+        objective=keras_tuner.Objective("val_loss", direction="min"),
+        max_trials=4,
+        overwrite=True,
+        project_name="hyperparameter_tuner_results",
     )
-    finetuning_model.fit(
-        train_ds, epochs=FLAGS.epochs, validation_data=validation_ds
+
+    tuner.search(train_ds, epochs=FLAGS.epochs, validation_data=validation_ds)
+
+    # Extract the best hyperparameters after the search.
+    best_hp = tuner.get_best_hyperparameters()[0]
+    finetuning_model = tuner.get_best_models()[0]
+
+    print(
+        f"The best hyperparameters found are:\nLearning Rate: {best_hp['lr']}"
     )
 
     if FLAGS.do_evaluation:
diff --git a/keras_nlp/__init__.py b/keras_nlp/__init__.py
index edbe34c3..2918cc6a 100644
--- a/keras_nlp/__init__.py
+++ b/keras_nlp/__init__.py
@@ -17,4 +17,4 @@ from keras_nlp import metrics
 from keras_nlp import tokenizers
 from keras_nlp import utils
 
-__version__ = "0.2.0-dev.1"
+__version__ = "0.2.0.dev2"
diff --git a/keras_nlp/layers/fnet_encoder.py b/keras_nlp/layers/fnet_encoder.py
index 47937392..90a3fc33 100644
--- a/keras_nlp/layers/fnet_encoder.py
+++ b/keras_nlp/layers/fnet_encoder.py
@@ -21,12 +21,11 @@ from tensorflow import keras
 class FNetEncoder(keras.layers.Layer):
     """FNet encoder.
 
-    This class follows the architecture of FNet encoder layer in paper
-    "FNet: Mixing Tokens with Fourier Transforms"
-    (https://arxiv.org/abs/2105.03824). Users can instantiate multiple instances
-    of this class to stack up the encoder.
+    This class follows the architecture of FNet encoder layer in the
+    [FNet paper](https://arxiv.org/abs/2105.03824). Users can instantiate
+    multiple instances of this class to stack up the encoder.
 
-    Note on padding: In the official FNet code, padding tokens are added to the
+    Note on masking: In the official FNet code, padding tokens are added to the
     the input. However, the padding masks are deleted, i.e., mixing of
     all tokens is done. This is because certain frequencies will be zeroed
     out if we apply padding masks in every encoder layer. Hence, we don't
@@ -36,14 +35,14 @@ class FNetEncoder(keras.layers.Layer):
         intermediate_dim: int. The hidden size of feedforward network.
         dropout: float, defaults to 0. The dropout value, applied in the
             feedforward network.
-        activation: string or `tf.keras.activations`, defaults to "relu". The
+        activation: string or `keras.activations`, defaults to "relu". The
             activation function of feedforward network.
         layer_norm_epsilon: float, defaults to 1e-5. The epsilon value in layer
             normalization components.
-        kernel_initializer: "string" or `tf.keras.initializers` initializer,
+        kernel_initializer: "string" or `keras.initializers` initializer,
             defaults to "glorot_uniform". The kernel initializer for the dense
             layers.
-        bias_initializer: "string" or `tf.keras.initializers` initializer,
+        bias_initializer: "string" or `keras.initializers` initializer,
             defaults to "zeros". The bias initializer for the dense layers.
         name: string, defaults to None. The name of the layer.
         **kwargs: other keyword arguments.
@@ -56,9 +55,9 @@ class FNetEncoder(keras.layers.Layer):
         intermediate_dim=64)
 
     # Create a simple model containing the encoder.
-    input = tf.keras.Input(shape=[4, 6])
+    input = keras.Input(shape=[10, 64])
     output = encoder(input)
-    model = tf.keras.Model(inputs=input, outputs=output)
+    model = keras.Model(inputs=input, outputs=output)
 
     # Call encoder on the inputs.
     input_data = tf.random.uniform(shape=[1, 10, 64])
@@ -66,7 +65,7 @@ class FNetEncoder(keras.layers.Layer):
     ```
 
     References:
-        [Lee-Thorp et al., 2021](https://arxiv.org/abs/2105.03824)
+     - [Lee-Thorp et al., 2021](https://arxiv.org/abs/2105.03824)
     """
 
     def __init__(
diff --git a/keras_nlp/layers/mlm_head.py b/keras_nlp/layers/mlm_head.py
index 09644205..53c1a7ff 100644
--- a/keras_nlp/layers/mlm_head.py
+++ b/keras_nlp/layers/mlm_head.py
@@ -22,10 +22,11 @@ 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)`.
+       `(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)`.
+       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
@@ -46,17 +47,18 @@ class MLMHead(keras.layers.Layer):
         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].
+            over all input words, as described
+            [here](https://arxiv.org/abs/1608.05859).
         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,
+        kernel_initializer: string or `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,
+        bias_initializer: string or `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.
@@ -91,7 +93,7 @@ class MLMHead(keras.layers.Layer):
     ```
 
     References:
-        [1] [Press and Wolf, 2016](https://arxiv.org/abs/1608.05859)
+     - [Press and Wolf, 2016](https://arxiv.org/abs/1608.05859)
     """
 
     def __init__(
@@ -147,7 +149,7 @@ class MLMHead(keras.layers.Layer):
             kernel_initializer=self.kernel_initializer,
             bias_initializer=self.bias_initializer,
         )
-        self._layer_norm = tf.keras.layers.LayerNormalization(
+        self._layer_norm = keras.layers.LayerNormalization(
             epsilon=self.layer_norm_epsilon,
         )
         if self.embedding_weights is None:
diff --git a/keras_nlp/layers/mlm_mask_generator.py b/keras_nlp/layers/mlm_mask_generator.py
index 85b207a4..6576dc50 100644
--- a/keras_nlp/layers/mlm_mask_generator.py
+++ b/keras_nlp/layers/mlm_mask_generator.py
@@ -73,15 +73,15 @@ class MLMMaskGenerator(keras.layers.Layer):
     Examples:
 
     Basic usage.
-    >>> masker = keras_nlp.layers.MLMMaskGenerator( \
-            vocabulary_size=10, mask_selection_rate=0.2, mask_token_id=0, \
-            mask_selection_length=5)
+    >>> 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.MLMMaskGenerator( \
-            vocabulary_size=10, mask_selection_rate=0.5, mask_token_id=0, \
-            mask_selection_length=5)
+    >>> 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 9bf8a615..2012f010 100644
--- a/keras_nlp/layers/position_embedding.py
+++ b/keras_nlp/layers/position_embedding.py
@@ -34,10 +34,10 @@ class PositionEmbedding(keras.layers.Layer):
     Args:
         sequence_length: The maximum length of the dynamic sequence.
         initializer: The initializer to use for the embedding weights. Defaults
-            to "glorot_uniform".
+            to `"glorot_uniform"`.
         seq_axis: The axis of the input tensor where we add the embeddings.
 
-    Example:
+    Examples:
     ```python
     token_embeddings = layers.Embedding(
         input_dim=vocab_size, output_dim=embed_dim
@@ -52,8 +52,7 @@ class PositionEmbedding(keras.layers.Layer):
     ```
 
     Reference:
-        [BERT: Pre-training of Deep Bidirectional Transformers for Language
-        Understanding](https://arxiv.org/abs/1810.04805).
+     - [Devlin et al., 2019](https://arxiv.org/abs/1810.04805)
     """
 
     def __init__(
diff --git a/keras_nlp/layers/sine_position_encoding.py b/keras_nlp/layers/sine_position_encoding.py
index 4112ca8a..0ce8ec3c 100644
--- a/keras_nlp/layers/sine_position_encoding.py
+++ b/keras_nlp/layers/sine_position_encoding.py
@@ -35,7 +35,7 @@ class SinePositionEncoding(keras.layers.Layer):
             curves, as described in Attention is All You Need. Defaults to
             10000.
 
-    Example:
+    Examples:
     ```python
     # create a simple embedding layer with sinusoidal positional encoding
     seq_len = 100
@@ -50,7 +50,7 @@ class SinePositionEncoding(keras.layers.Layer):
     ```
 
     References:
-      [Attention is All You Need](https://arxiv.org/abs/1706.03762)
+     - [Vaswani et al., 2017](https://arxiv.org/abs/1706.03762)
     """
 
     def __init__(
diff --git a/keras_nlp/layers/token_and_position_embedding.py b/keras_nlp/layers/token_and_position_embedding.py
index 03c5f8c6..c6b9fb2d 100644
--- a/keras_nlp/layers/token_and_position_embedding.py
+++ b/keras_nlp/layers/token_and_position_embedding.py
@@ -41,7 +41,7 @@ class TokenAndPositionEmbedding(keras.layers.Layer):
             used in the vocabulary
             (input_dim should equal size of vocabulary + 1).
 
-    Example:
+    Examples:
     ```python
     seq_length = 50
     vocab_size = 5000
diff --git a/keras_nlp/layers/transformer_decoder.py b/keras_nlp/layers/transformer_decoder.py
index 92d59c2e..f8cdd105 100644
--- a/keras_nlp/layers/transformer_decoder.py
+++ b/keras_nlp/layers/transformer_decoder.py
@@ -26,23 +26,29 @@ from keras_nlp.layers.transformer_layer_utils import (  # isort:skip
 class TransformerDecoder(keras.layers.Layer):
     """Transformer decoder.
 
-    This class follows the architecture of transformer decoder layer in paper
-    "Attention is All You Need" (https://arxiv.org/abs/1706.03762). Users can
-    instantiate multiple instances of this class to stack up the decoder.
+    This class follows the architecture of the transformer decoder layer in the
+    paper [Attention is All You Need](https://arxiv.org/abs/1706.03762). Users
+    can instantiate multiple instances of this class to stack up a decoder.
+
+    This layer will correctly compute an attention mask from an implicit
+    Keras padding mask (for example, by passing `mask_zero=True` to a
+    `keras.layers.Embedding` layer). See the Masking and Padding
+    [guide](https://keras.io/guides/understanding_masking_and_padding/)
+    for more details.
 
     Args:
         intermediate_dim: int, the hidden size of feedforward network.
         num_heads: int, the number of heads in MultiHeadAttention.
         dropout: float, defaults to 0. the dropout value, shared by
             MultiHeadAttention and feedforward network.
-        activation: string or tf.keras.activations, defaults to "relu". the
+        activation: string or `keras.activations`, defaults to "relu". the
             activation function of feedforward network.
         layer_norm_epsilon: float, defaults to 1e-5. The eps value in layer
             normalization components.
-        kernel_initializer: string or tf.keras.initializers initializer,
+        kernel_initializer: string or `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,
+        bias_initializer: string or `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.
@@ -55,21 +61,21 @@ class TransformerDecoder(keras.layers.Layer):
         intermediate_dim=64, num_heads=8)
 
     # Create a simple model containing the decoder.
-    decoder_input = tf.keras.Input(shape=[4, 6])
-    encoder_input = tf.keras.Input(shape=[4, 6])
+    decoder_input = keras.Input(shape=[10, 64])
+    encoder_input = keras.Input(shape=[10, 64])
     output = decoder(decoder_input, encoder_input)
-    model = tf.keras.Model(inputs=[decoder_input, encoder_input],
+    model = keras.Model(inputs=[decoder_input, encoder_input],
         outputs=output)
 
     # Call decoder on the inputs.
-    decoder_input_data = tf.random.uniform(shape=[1, 10, 64])
-    encoder_input_data = tf.random.uniform(shape=[1, 10, 64])
+    decoder_input_data = tf.random.uniform(shape=[2, 10, 64])
+    encoder_input_data = tf.random.uniform(shape=[2, 10, 64])
     decoder_output = model([decoder_input_data, encoder_input_data])
 
     ```
 
     References:
-        [Vaswani et al., 20XX](https://arxiv.org/abs/1706.03762)
+     - [Vaswani et al., 2017](https://arxiv.org/abs/1706.03762)
 
     """
 
diff --git a/keras_nlp/layers/transformer_encoder.py b/keras_nlp/layers/transformer_encoder.py
index 870228be..c67409d9 100644
--- a/keras_nlp/layers/transformer_encoder.py
+++ b/keras_nlp/layers/transformer_encoder.py
@@ -24,23 +24,30 @@ from keras_nlp.layers.transformer_layer_utils import (  # isort:skip
 class TransformerEncoder(keras.layers.Layer):
     """Transformer encoder.
 
-    This class follows the architecture of transformer encoder layer in paper
-    "Attention is All You Need" (https://arxiv.org/abs/1706.03762). Users can
-    instantiate multiple instances of this class to stack up the encoder.
+    This class follows the architecture of the transformer encoder layer in the
+    paper [Attention is All You Need](https://arxiv.org/abs/1706.03762). Users
+    can instantiate multiple instances of this class to stack up an encoder.
+
+    This layer will correctly compute an attention mask from an implicit
+    Keras padding mask (for example, by passing `mask_zero=True` to a
+    `keras.layers.Embedding` layer). See the Masking and Padding
+    [guide](https://keras.io/guides/understanding_masking_and_padding/)
+    for more details.
 
     Args:
         intermediate_dim: int, the hidden size of feedforward network.
-        num_heads: int, the number of heads in MultiHeadAttention.
+        num_heads: int, the number of heads in the
+            `keras.layers.MultiHeadAttention` layer.
         dropout: float, defaults to 0. the dropout value, shared by
-            MultiHeadAttention and feedforward network.
-        activation: string or `tf.keras.activations`, defaults to "relu". the
+            `keras.layers.MultiHeadAttention` and feedforward network.
+        activation: string or `keras.activations`, defaults to "relu". the
             activation function of feedforward network.
         layer_norm_epsilon: float, defaults to 1e-5. The epsilon value in layer
             normalization components.
-        kernel_initializer: string or tf.keras.initializers initializer,
+        kernel_initializer: string or `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,
+        bias_initializer: string or `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.
@@ -54,18 +61,18 @@ class TransformerEncoder(keras.layers.Layer):
         intermediate_dim=64, num_heads=8)
 
     # Create a simple model containing the encoder.
-    input = tf.keras.Input(shape=[4, 6])
+    input = keras.Input(shape=[10, 64])
     output = encoder(input)
-    model = tf.keras.Model(inputs=input, outputs=output)
+    model = keras.Model(inputs=input, outputs=output)
 
     # Call encoder on the inputs.
-    input_data = tf.random.uniform(shape=[1, 10, 64])
+    input_data = tf.random.uniform(shape=[2, 10, 64])
     output = model(input_data)
 
     ```
 
     References:
-        [Vaswani et al., 20XX](https://arxiv.org/abs/1706.03762)
+     - [Vaswani et al., 2017](https://arxiv.org/abs/1706.03762)
     """
 
     def __init__(
diff --git a/keras_nlp/tokenizers/byte_tokenizer.py b/keras_nlp/tokenizers/byte_tokenizer.py
index 1a106cd9..fcbad637 100644
--- a/keras_nlp/tokenizers/byte_tokenizer.py
+++ b/keras_nlp/tokenizers/byte_tokenizer.py
@@ -30,19 +30,25 @@ class ByteTokenizer(tokenizer.Tokenizer):
     This tokenizer is a vocabulary-free tokenizer which will tokenize text as
     as raw bytes from [0, 256).
 
+    Tokenizer outputs can either be padded and truncated with a
+    `sequence_length` argument, or left un-truncated. The exact output will
+    depend on the rank of the input tensors.
+
     If input is a batch of strings:
     By default, the layer will output a `tf.RaggedTensor` where the last
     dimension of the output is ragged. If `sequence_length` is set, the layer
     will output a dense `tf.Tensor` where all inputs have been padded or
-    truncated to `sequence_length`. The output dtype can be controlled via the
-    `dtype` argument, which should be an integer type
-    (tf.int16, tf.int32, etc.).
+    truncated to `sequence_length`.
 
     If input is a scalar string:
     There are two cases here. If `sequence_length` is set, the output will be
     a dense `tf.Tensor` of shape `[sequence_length]`. Otherwise, the output will
     be a dense `tf.Tensor` of shape `[None]`.
 
+    The output dtype can be controlled via the
+    `dtype` argument, which should be an integer type
+    (tf.int16, tf.int32, etc.).
+
     Args:
         lowercase: boolean. If True, the input text will be converted to
             lowercase before tokenization.
@@ -89,7 +95,7 @@ class ByteTokenizer(tokenizer.Tokenizer):
     <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.
+    Tokenize, then batch for ragged outputs.
     >>> tokenizer = keras_nlp.tokenizers.ByteTokenizer()
     >>> ds = tf.data.Dataset.from_tensor_slices(["hello", "fun"])
     >>> ds = ds.map(tokenizer)
@@ -97,14 +103,14 @@ class ByteTokenizer(tokenizer.Tokenizer):
     >>> ds.take(1).get_single_element()
     <tf.RaggedTensor [[104, 101, 108, 108, 111], [102, 117, 110]]>
 
-    Batch the inputs and then tokenize.
+    Batch, then tokenize for ragged outputs.
     >>> tokenizer = keras_nlp.tokenizers.ByteTokenizer()
     >>> ds = tf.data.Dataset.from_tensor_slices(["hello", "fun"])
     >>> ds = ds.batch(2).map(tokenizer)
     >>> ds.take(1).get_single_element()
     <tf.RaggedTensor [[104, 101, 108, 108, 111], [102, 117, 110]]>
 
-    Tokenize first, then batch the dataset up (`sequence_length` provided).
+    Tokenize, then batch for dense outputs (`sequence_length` provided).
     >>> tokenizer = keras_nlp.tokenizers.ByteTokenizer(sequence_length=5)
     >>> ds = tf.data.Dataset.from_tensor_slices(["hello", "fun"])
     >>> ds = ds.map(tokenizer)
@@ -114,7 +120,7 @@ class ByteTokenizer(tokenizer.Tokenizer):
     array([[104, 101, 108, 108, 111],
            [102, 117, 110,   0,   0]], dtype=int32)>
 
-    Batch the inputs and then tokenize (`sequence_length` provided).
+    Batch, then tokenize for dense outputs. (`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)
diff --git a/keras_nlp/tokenizers/tokenizer.py b/keras_nlp/tokenizers/tokenizer.py
index baccd457..03602c82 100644
--- a/keras_nlp/tokenizers/tokenizer.py
+++ b/keras_nlp/tokenizers/tokenizer.py
@@ -20,32 +20,42 @@ from tensorflow import keras
 class Tokenizer(keras.layers.Layer):
     """A base class for tokenizer layers.
 
-    The class is intended as a base class when implementing a tokenizer as a
-    `keras.layers.Layer`. It contains two new methods `tokenize()` and
-    `detokenize()`.
+    Tokenizers in the KerasNLP library should all subclass this layer.
+    The class provides two core methods `tokenize()` and `detokenize()` for
+    going from plain text to sequences and back. A tokenizer is a subclass of
+    `keras.layers.Layer` and can be combined into a `keras.Model`.
 
-    Subclassers should always implement the `tokenize` method, which will also
-    be the default when invoking the layer directly on inputs.
+    Subclassers should always implement the `tokenize()` method, which will also
+    be the default when calling the layer directly on inputs.
 
-    Subclassers should implement the `detokenize()` method if the tokenization
-    is reversible. Otherwise, this can be skipped.
+    Subclassers can optionally implement the `detokenize()` method if the
+    tokenization is reversible. Otherwise, this can be skipped.
 
     Subclassers should implement `get_vocabulary()`, `vocabulary_size()`,
     `token_to_id()` and `id_to_token()` if applicable. For some simple
-    "pre-tokenizers," such whitespace splitting, these methods may not apply,
-    and can be skipped.
+    "vocab free" tokenizers, such as a whitespace splitter show below, these
+    methods do not apply and can be skipped.
 
     Examples:
 
     ```python
     class WhitespaceSplitterTokenizer(keras_nlp.tokenizers.Tokenizer):
         def tokenize(self, inputs):
-            return tf.strings.split(inputs).to_tensor()
+            return tf.strings.split(inputs)
 
         def detokenize(self, inputs):
-            return tf.strings.reduce_join([inputs], separator=" ", axis=-1)
+            return tf.strings.reduce_join(inputs, separator=" ", axis=-1)
 
-    tokenizer = WhitespaceTokenizer()
+    tokenizer = WhitespaceSplitterTokenizer()
+
+    # Tokenize some inputs.
+    tokenizer.tokenize("This is a test")
+
+    # Shorthard for `tokenize()`.
+    tokenizer("This is a test")
+
+    # Detokenize some outputs.
+    tokenizer.detokenize(["This", "is", "a", "test"])
     ```
     """
 
diff --git a/keras_nlp/tokenizers/unicode_character_tokenizer.py b/keras_nlp/tokenizers/unicode_character_tokenizer.py
index 3d21e9ca..9682cb16 100644
--- a/keras_nlp/tokenizers/unicode_character_tokenizer.py
+++ b/keras_nlp/tokenizers/unicode_character_tokenizer.py
@@ -27,6 +27,24 @@ class UnicodeCharacterTokenizer(tokenizer.Tokenizer):
     This tokenizer is a vocabulary free tokenizer which tokenizes text as
     unicode characters codepoints.
 
+    Tokenizer outputs can either be padded and truncated with a
+    `sequence_length` argument, or left un-truncated. The exact output will
+    depend on the rank of the input tensors.
+
+    If input is a batch of strings (rank > 0):
+    By default, the layer will output a `tf.RaggedTensor` where the last
+    dimension of the output is ragged. If `sequence_length` is set, the layer
+    will output a dense `tf.Tensor` where all inputs have been padded or
+    truncated to `sequence_length`.
+
+    If input is a scalar string (rank == 0):
+    By default, the layer will output a dense `tf.Tensor` with static shape
+    `[None]`. If `sequence_length` is set, the output will be
+    a dense `tf.Tensor` of shape `[sequence_length]`.
+
+    The output dtype can be controlled via the `dtype` argument, which should be
+    an integer type (tf.int16, tf.int32, etc.).
+
     Args:
         lowercase: If true, the input text will be first lowered before
             tokenization.
@@ -45,6 +63,10 @@ class UnicodeCharacterTokenizer(tokenizer.Tokenizer):
             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".
+        vocabulary_size: Set the vocabulary `vocabulary_size`,
+            by clamping all codepoints to the range [0, vocabulary_size).
+            Effectively this will make the `vocabulary_size - 1` id the
+            the OOV value.
 
     Examples:
 
@@ -74,7 +96,7 @@ class UnicodeCharacterTokenizer(tokenizer.Tokenizer):
         [2346, 2369, 2360, 2381, 2340, 2325,    0,    0],
         [1705, 1578, 1575, 1576,    0,    0,    0,    0]], dtype=int32)>
 
-    Tokenize first, then batch the dataset.
+    Tokenize, then batch for ragged outputs.
     >>> inputs = ["Book", "पुस्तक", "کتاب"]
     >>> tokenizer = keras_nlp.tokenizers.UnicodeCharacterTokenizer()
     >>> ds = tf.data.Dataset.from_tensor_slices(inputs)
@@ -85,7 +107,7 @@ class UnicodeCharacterTokenizer(tokenizer.Tokenizer):
         [2346, 2369, 2360, 2381, 2340, 2325],
         [1705, 1578, 1575, 1576]]>
 
-    Batch the inputs and then tokenize.
+    Batch, then tokenize for ragged outputs.
     >>> inputs = ["Book", "पुस्तक", "کتاب"]
     >>> tokenizer = keras_nlp.tokenizers.UnicodeCharacterTokenizer()
     >>> ds = tf.data.Dataset.from_tensor_slices(inputs)
@@ -95,7 +117,7 @@ class UnicodeCharacterTokenizer(tokenizer.Tokenizer):
         [2346, 2369, 2360, 2381, 2340, 2325],
         [1705, 1578, 1575, 1576]]>
 
-    Tokenize first, then batch for dense outputs (`sequence_length` provided).
+    Tokenize, then batch for dense outputs (`sequence_length` provided).
     >>> inputs = ["Book", "पुस्तक", "کتاب"]
     >>> tokenizer = keras_nlp.tokenizers.UnicodeCharacterTokenizer(
     ...     sequence_length=5)
@@ -108,8 +130,7 @@ class UnicodeCharacterTokenizer(tokenizer.Tokenizer):
         [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).
+    Batch, then tokenize for dense outputs (`sequence_length` provided).
     >>> inputs = ["Book", "पुस्तक", "کتاب"]
     >>> tokenizer = keras_nlp.tokenizers.UnicodeCharacterTokenizer(
     ...     sequence_length=5)
@@ -121,7 +142,7 @@ class UnicodeCharacterTokenizer(tokenizer.Tokenizer):
         [2346, 2369, 2360, 2381, 2340],
         [1705, 1578, 1575, 1576,    0]], dtype=int32)>
 
-    Tokenization showcasing truncation of long sequences.
+    Tokenization with truncation.
     >>> inputs = ["I Like to Travel a Lot", "मैं किताबें पढ़ना पसंद करता हूं"]
     >>> tokenizer = keras_nlp.tokenizers.UnicodeCharacterTokenizer(
     ...     sequence_length=5)
@@ -130,13 +151,26 @@ class UnicodeCharacterTokenizer(tokenizer.Tokenizer):
         numpy=array([[ 105,   32,  108,  105,  107],
        [2350, 2376, 2306,   32, 2325]], dtype=int32)>
 
+    Tokenization with vocabulary_size.
+    >>> latin_ext_cutoff = 592
+    >>> tokenizer = keras_nlp.tokenizers.UnicodeCharacterTokenizer(
+    ...     vocabulary_size=latin_ext_cutoff)
+    >>> tokenizer("¿Cómo estás?")
+    <tf.Tensor: shape=(10,), dtype=int32,
+    numpy=array([191,  99, 243, 109, 111,  32, 101, 115, 116, 225, 115,  63],
+    dtype=int32)>
+    >>> tokenizer("आप कैसे हैं")
+    <tf.Tensor: shape=(11,), dtype=int32,
+    numpy=array([591, 591,  32, 591, 591, 591, 591,  32, 591, 591, 591],
+    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
+    Detokenization with padding.
     >>> tokenizer = keras_nlp.tokenizers.UnicodeCharacterTokenizer(
     ...     sequence_length=7)
     >>> dataset = tf.data.Dataset.from_tensor_slices(["a b c", "b c", "a"])
@@ -150,13 +184,11 @@ class UnicodeCharacterTokenizer(tokenizer.Tokenizer):
     <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)
+    ...     errors="replace", replacement_char=88)
     >>> tokenizer.detokenize(inputs).numpy().decode('utf-8')
-    'niKnja'
+    'niXnja'
     """
 
     def __init__(
@@ -168,6 +200,7 @@ class UnicodeCharacterTokenizer(tokenizer.Tokenizer):
         replacement_char: int = 65533,
         input_encoding: str = "UTF-8",
         output_encoding: str = "UTF-8",
+        vocabulary_size: int = None,
         **kwargs,
     ) -> None:
         # Check dtype and provide a default.
@@ -213,6 +246,7 @@ class UnicodeCharacterTokenizer(tokenizer.Tokenizer):
         self.replacement_char = replacement_char
         self.input_encoding = input_encoding
         self.output_encoding = output_encoding
+        self.vocabulary_size = vocabulary_size
 
     def get_config(self) -> Dict[str, Any]:
         config = super().get_config()
@@ -225,10 +259,16 @@ class UnicodeCharacterTokenizer(tokenizer.Tokenizer):
                 "replacement_char": self.replacement_char,
                 "input_encoding": self.input_encoding,
                 "output_encoding": self.output_encoding,
+                "vocabulary_size": self.vocabulary_size,
             }
         )
         return config
 
+    def vocabulary_size(self) -> int:
+        """Get the size of the tokenizer vocabulary. None implies no vocabulary
+        size was provided"""
+        return self.vocabulary_size
+
     def tokenize(self, inputs):
         if not isinstance(inputs, (tf.Tensor, tf.RaggedTensor)):
             inputs = tf.convert_to_tensor(inputs)
@@ -260,6 +300,12 @@ class UnicodeCharacterTokenizer(tokenizer.Tokenizer):
 
         if scalar_input:
             tokens = tf.squeeze(tokens, 0)
+
+        # Optionally clamps the output code point values to be in the
+        # range [0, vocabulary_size)
+        if self.vocabulary_size:
+            tokens = tf.clip_by_value(tokens, 0, self.vocabulary_size - 1)
+
         return tokens
 
     def detokenize(self, inputs):
diff --git a/keras_nlp/tokenizers/word_piece_tokenizer.py b/keras_nlp/tokenizers/word_piece_tokenizer.py
index cc1766ec..79ec5057 100644
--- a/keras_nlp/tokenizers/word_piece_tokenizer.py
+++ b/keras_nlp/tokenizers/word_piece_tokenizer.py
@@ -83,12 +83,23 @@ class WordPieceTokenizer(tokenizer.Tokenizer):
     `lowercase=False`, `strip_accents=False` and `split=False`. In
     this case, inputs should be pre-split string tensors or ragged tensors.
 
+    Tokenizer outputs can either be padded and truncated with a
+    `sequence_length` argument, or left un-truncated. The exact output will
+    depend on the rank of the input tensors.
+
+    If input is a batch of strings (rank > 0):
     By default, the layer will output a `tf.RaggedTensor` where the last
-    dimension of the output is ragged after whitespace splitting and sub-word
-    tokenizing. If `sequence_length` is set, the layer will output a dense
-    `tf.Tensor` where all inputs have been padded or truncated to
-    `sequence_length`. The output dtype can be controlled via the `dtype`
-    argument, which should be either an integer or string type.
+    dimension of the output is ragged. If `sequence_length` is set, the layer
+    will output a dense `tf.Tensor` where all inputs have been padded or
+    truncated to `sequence_length`.
+
+    If input is a scalar string (rank == 0):
+    By default, the layer will output a dense `tf.Tensor` with static shape
+    `[None]`. If `sequence_length` is set, the output will be
+    a dense `tf.Tensor` of shape `[sequence_length]`.
+
+    The output dtype can be controlled via the `dtype` argument, which should
+    be either an integer or string type.
 
     Args:
         vocabulary: A list of strings or a string string filename path. If
@@ -115,8 +126,8 @@ class WordPieceTokenizer(tokenizer.Tokenizer):
             must be included in the vocab.
 
     References:
-        - [Schuster and Nakajima, 2012](https://research.google/pubs/pub37842/)
-        - [Song et al., 2020](https://arxiv.org/abs/2012.15524)
+     - [Schuster and Nakajima, 2012](https://research.google/pubs/pub37842/)
+     - [Song et al., 2020](https://arxiv.org/abs/2012.15524)
 
     Examples:
 
diff --git a/setup.py b/setup.py
index bb458e4c..6ec18947 100644
--- a/setup.py
+++ b/setup.py
@@ -29,7 +29,7 @@ setup(
     ),
     long_description=README,
     long_description_content_type="text/markdown",
-    version="0.2.0-dev.1",
+    version="0.2.0.dev2",
     url="https://github.com/keras-team/keras-nlp",
     author="Keras team",
     author_email="keras-nlp@google.com",
@@ -53,6 +53,7 @@ setup(
             "datasets",  # For GLUE in BERT example.
             "nltk",
             "wikiextractor",
+            "keras-tuner",
         ],
     },
     classifiers=[


Test patch:
diff --git a/keras_nlp/tokenizers/unicode_character_tokenizer_test.py b/keras_nlp/tokenizers/unicode_character_tokenizer_test.py
index 34df6a50..334a8139 100644
--- a/keras_nlp/tokenizers/unicode_character_tokenizer_test.py
+++ b/keras_nlp/tokenizers/unicode_character_tokenizer_test.py
@@ -59,6 +59,46 @@ class UnicodeCharacterTokenizerTest(tf.test.TestCase):
             ],
         )
 
+    def test_tokenize_scalar_with_vocabulary_size(self):
+        input_data = "ninja"
+        tokenizer = UnicodeCharacterTokenizer(vocabulary_size=105)
+        call_output = tokenizer(input_data)
+        tokenize_output = tokenizer.tokenize(input_data)
+
+        self.assertAllEqual(call_output, [104, 104, 104, 104, 97])
+        self.assertAllEqual(tokenize_output, [104, 104, 104, 104, 97])
+
+    def test_tokenize_dense_with_vocabulary_size(self):
+        input_data = tf.constant(["ninja", "samurai", "▀▁▂▃"])
+        tokenizer = UnicodeCharacterTokenizer(
+            sequence_length=10, vocabulary_size=105
+        )
+        call_output = tokenizer(input_data)
+        self.assertIsInstance(call_output, tf.Tensor)
+        self.assertAllEqual(
+            call_output,
+            [
+                [104, 104, 104, 104, 97, 0, 0, 0, 0, 0],
+                [104, 97, 104, 104, 104, 97, 104, 0, 0, 0],
+                [104, 104, 104, 104, 0, 0, 0, 0, 0, 0],
+            ],
+        )
+
+    def test_tokenize_ragged_with_vocabulary_size(self):
+        input_data = tf.constant(["ninja", "samurai", "▀▁▂▃"])
+        tokenizer = UnicodeCharacterTokenizer(vocabulary_size=105)
+        call_output = tokenizer(input_data)
+        tokenize_output = tokenizer.tokenize(input_data)
+        self.assertIsInstance(call_output, tf.RaggedTensor)
+        exp_outputs = [
+            [104, 104, 104, 104, 97],
+            [104, 97, 104, 104, 104, 97, 104],
+            [104, 104, 104, 104],
+        ]
+        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_detokenize(self):
         input_data = tf.ragged.constant(
             [
@@ -232,6 +272,7 @@ class UnicodeCharacterTokenizerTest(tf.test.TestCase):
             sequence_length=11,
             normalization_form="NFC",
             errors="strict",
+            vocabulary_size=None,
         )
         cloned_tokenizer = UnicodeCharacterTokenizer.from_config(
             original_tokenizer.get_config()
@@ -255,6 +296,7 @@ class UnicodeCharacterTokenizerTest(tf.test.TestCase):
             normalization_form="NFC",
             errors="ignore",
             replacement_char=0,
+            vocabulary_size=100,
         )
         exp_config = {
             "dtype": "int32",
@@ -267,6 +309,7 @@ class UnicodeCharacterTokenizerTest(tf.test.TestCase):
             "input_encoding": "UTF-8",
             "output_encoding": "UTF-8",
             "trainable": True,
+            "vocabulary_size": 100,
         }
         self.assertEqual(tokenizer.get_config(), exp_config)
 
@@ -278,6 +321,7 @@ class UnicodeCharacterTokenizerTest(tf.test.TestCase):
             replacement_char=0,
             input_encoding="UTF-16",
             output_encoding="UTF-16",
+            vocabulary_size=None,
         )
         exp_config_different_encoding = {
             "dtype": "int32",
@@ -290,6 +334,7 @@ class UnicodeCharacterTokenizerTest(tf.test.TestCase):
             "input_encoding": "UTF-16",
             "output_encoding": "UTF-16",
             "trainable": True,
+            "vocabulary_size": None,
         }
         self.assertEqual(
             tokenize_different_encoding.get_config(),
@@ -305,6 +350,7 @@ class UnicodeCharacterTokenizerTest(tf.test.TestCase):
             sequence_length=20,
             normalization_form="NFKC",
             errors="replace",
+            vocabulary_size=None,
         )
         inputs = keras.Input(dtype="string", shape=())
         outputs = tokenizer(inputs)