Repository: keras-team/keras-nlp
Pull request #785 (2023-02-27 08:08:32)

Fix Fenced Doc-String #782
Following changes were made in Doc-String Examples -  

- Add models to the list in docstring_test.py
- Import from the root of the library
- Fix typo in f_net_classifier.py
- add [MASK] in Vocaburay
- remove num_heads parameter 

The number of failed examples has gone down from 5 to 1, however I am getting 1 error in `File "/content/keras-nlp/keras-nlp/keras-nlp/keras_nlp/models/f_net/f_net_masked_lm.py", line 73, in keras_nlp.models.f_net.f_net_masked_lm.FNetMaskedLM`- 

`Node: 'f_net_masked_lm_1/f_net_backbone_3/segment_embedding/embedding_lookup'
    indices[0,3] = 4 is not in [0, 4)
    	 [[{{node f_net_masked_lm_1/f_net_backbone_3/segment_embedding/embedding_lookup}}]] [Op:__inference_train_function_458156]
`
Might need help from @abheesht17 , or any other developer who can assist.
Thanks!
Fixes keras-team/keras-nlp #782

abheesht17: Thanks, @atharvapurdue! The reason why it's throwing an error is because the default value of  `num_segments` is 4 for FNet. Hence, this throws an embedding lookup error: https://github.com/keras-team/keras-nlp/blob/master/keras_nlp/models/f_net/f_net_backbone.py#L126-L131. The best way to resolve it is to just change

```
        "segment_ids": tf.constant(
            [[1, 0, 0, 4, 0, 6, 7, 8]] * 2, shape=(2, 8)
        ),
```

to values less than 4. Just change it to something like  `[[0, 0, 0, 1, 1, 1, 0, 0]] * 2`.
atharvapurdue: Yes! @abheesht17. I thought so, but wasn't sure. I will get make the necessary changes. Thank you!
atharvapurdue: Done !  @abheesht17 , Thanks for your help!
You can see the doc_string_test passing all tests here - https://colab.research.google.com/drive/1xVlUGJIDSH44bF851HZ01yD24ocJmul5#scrollTo=8DErjytmUTv_ 



Patch:
diff --git a/README.md b/README.md
index 3523dc69..7d01c2cb 100644
--- a/README.md
+++ b/README.md
@@ -47,14 +47,14 @@ We are a new and growing project and welcome [contributions](CONTRIBUTING.md).
 To install the latest official release:
 
 ```
-pip install keras-nlp tensorflow --upgrade
+pip install keras-nlp --upgrade
 ```
 
 To install the latest unreleased changes to the library, we recommend using
 pip to install directly from the master branch on github:
 
 ```
-pip install git+https://github.com/keras-team/keras-nlp.git tensorflow --upgrade
+pip install git+https://github.com/keras-team/keras-nlp.git --upgrade
 ```
 
 ## Quickstart
@@ -64,7 +64,6 @@ Fine-tune BERT on a small sentiment analysis task using the
 
 ```python
 import keras_nlp
-from tensorflow import keras
 import tensorflow_datasets as tfds
 
 imdb_train, imdb_test = tfds.load(
@@ -73,23 +72,12 @@ imdb_train, imdb_test = tfds.load(
     as_supervised=True,
     batch_size=16,
 )
-classifier = keras_nlp.models.BertClassifier.from_preset(
-    "bert_base_en_uncased",
-)
-classifier.compile(
-    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
-    optimizer=keras.optimizers.Adam(5e-5),
-    metrics=keras.metrics.SparseCategoricalAccuracy(),
-    jit_compile=True,
-)
-classifier.fit(
-    imdb_train,
-    validation_data=imdb_test,
-    epochs=1,
-)
-
-# Predict a new example
-classifier.predict(["What an amazing movie, three hours of pure bliss!"])
+# Load a BERT model.
+classifier = keras_nlp.models.BertClassifier.from_preset("bert_base_en_uncased")
+# Fine-tune on IMDb movie reviews.
+classifier.fit(imdb_train, validation_data=imdb_test)
+# Predict two new examples.
+classifier.predict(["What an amazing movie!", "A total waste of my time."])
 ```
 
 For more in depth guides and examples, visit https://keras.io/keras_nlp/.
@@ -107,7 +95,7 @@ KerasNLP provides access to pre-trained models via the `keras_nlp.models` API.
 These pre-trained models are provided on an "as is" basis, without warranties
 or conditions of any kind. The following underlying models are provided by third
 parties, and subject to separate licenses:
-DistilBERT, RoBERTa, XLM-RoBERTa, DeBERTa, and GPT-2.
+BART, DeBERTa, DistilBERT, GPT-2, OPT, RoBERTa, and XLM-RoBERTa.
 
 ## Citing KerasNLP
 
diff --git a/keras_nlp/layers/cached_multi_head_attention.py b/keras_nlp/layers/cached_multi_head_attention.py
index 094d381e..a86bb93d 100644
--- a/keras_nlp/layers/cached_multi_head_attention.py
+++ b/keras_nlp/layers/cached_multi_head_attention.py
@@ -62,44 +62,24 @@ class CachedMultiHeadAttention(keras.layers.MultiHeadAttention):
         value,
         key=None,
         attention_mask=None,
-        use_causal_mask=False,
         cache=None,
-        cache_index=None,
+        cache_index=0,
     ):
         if not self._built_from_signature:
             self._build_from_signature(query=query, value=value, key=key)
         if key is None:
             key = value
 
-        attention_mask = self._compute_attention_mask(
-            query,
-            value,
-            key=key,
-            attention_mask=attention_mask,
-            use_causal_mask=use_causal_mask,
-        )
-
         query = self._query_dense(query)
         key = self._key_dense(key)
         value = self._value_dense(value)
 
         if cache is not None:
-            if cache_index is None:
-                seq_len = tf.shape(query)[1]
-                k_update_indices = [0, 0, 0, 0, 0]
-                v_update_indices = [0, 1, 0, 0, 0]
-                cache_index = seq_len - 1
-            else:
-                k_update_indices = [0, 0, cache_index, 0, 0]
-                v_update_indices = [0, 1, cache_index, 0, 0]
-            cache = dynamic_update_slice(
-                cache, key[:, tf.newaxis, ...], k_update_indices
-            )
-            cache = dynamic_update_slice(
-                cache, value[:, tf.newaxis, ...], v_update_indices
-            )
-            key = cache[:, 0, : cache_index + 1, :, :]
-            value = cache[:, 1, : cache_index + 1, :, :]
+            key_cache, value_cache = tf.unstack(cache, axis=1)
+            start = [0, cache_index, 0, 0]
+            key = dynamic_update_slice(key_cache, key, start)
+            value = dynamic_update_slice(value_cache, value, start)
+            cache = tf.stack((key, value), axis=1)
 
         query = tf.multiply(query, 1.0 / tf.math.sqrt(float(self._key_dim)))
         attention_scores = tf.einsum(self._dot_product_equation, key, query)
diff --git a/keras_nlp/layers/position_embedding.py b/keras_nlp/layers/position_embedding.py
index 755750c7..64b04173 100644
--- a/keras_nlp/layers/position_embedding.py
+++ b/keras_nlp/layers/position_embedding.py
@@ -17,8 +17,6 @@
 import tensorflow as tf
 from tensorflow import keras
 
-SEQUENCE_AXIS = -2
-
 
 @keras.utils.register_keras_serializable(package="keras_nlp")
 class PositionEmbedding(keras.layers.Layer):
@@ -103,11 +101,12 @@ class PositionEmbedding(keras.layers.Layer):
 
         super().build(input_shape)
 
-    def call(self, inputs):
+    def call(self, inputs, start_index=0):
         if isinstance(inputs, tf.RaggedTensor):
             bounding_shape = inputs.bounding_shape()
             position_embeddings = self._trim_and_broadcast_position_embeddings(
                 bounding_shape,
+                start_index,
             )
             # then apply row lengths to recreate the same ragged shape as inputs
             return tf.RaggedTensor.from_tensor(
@@ -117,12 +116,18 @@ class PositionEmbedding(keras.layers.Layer):
         else:
             return self._trim_and_broadcast_position_embeddings(
                 tf.shape(inputs),
+                start_index,
             )
 
-    def _trim_and_broadcast_position_embeddings(self, shape):
-        input_length = shape[SEQUENCE_AXIS]
+    def _trim_and_broadcast_position_embeddings(self, shape, start_index):
+        feature_length = shape[-1]
+        sequence_length = shape[-2]
         # trim to match the length of the input sequence, which might be less
         # than the sequence_length of the layer.
-        position_embeddings = self.position_embeddings[:input_length, :]
+        position_embeddings = tf.slice(
+            self.position_embeddings,
+            (start_index, 0),
+            (sequence_length, feature_length),
+        )
         # then broadcast to add the missing dimensions to match "shape"
         return tf.broadcast_to(position_embeddings, shape)
diff --git a/keras_nlp/layers/transformer_decoder.py b/keras_nlp/layers/transformer_decoder.py
index afbdb672..469c0a4f 100644
--- a/keras_nlp/layers/transformer_decoder.py
+++ b/keras_nlp/layers/transformer_decoder.py
@@ -212,7 +212,7 @@ class TransformerDecoder(keras.layers.Layer):
         encoder_padding_mask=None,
         encoder_attention_mask=None,
         cache=None,
-        cache_index=None,
+        cache_index=0,
     ):
         """Forward pass of the TransformerDecoder.
 
@@ -236,8 +236,7 @@ class TransformerDecoder(keras.layers.Layer):
                 tokens. `cache` is of shape [B, 2, max_seq_len, num_heads,
                 key_dims].
             cache_index: a int or int Tensor, the index of the current token
-                being processed. If `cache_index=None` while `cache` is set, it
-                means it's the first pass to build the cache.
+                being processed.
         Returns:
             Either a tuple of (outputs, cache) if a cache was passed, or a
             single value outputs if a cache was not passed.
@@ -272,7 +271,19 @@ class TransformerDecoder(keras.layers.Layer):
         x = decoder_sequence  # Intermediate result.
 
         # Compute self attention mask.
-        self_attention_mask = compute_causal_mask(decoder_sequence)
+        batch_size = tf.shape(decoder_sequence)[0]
+        input_length = output_length = tf.shape(decoder_sequence)[1]
+        # We need to handle a rectangular causal mask when doing cached
+        # decoding. For generative inference, `decoder_sequence` will
+        # generally be length 1, and `cache` will be the full generation length.
+        if cache is not None:
+            input_length = tf.shape(cache)[2]
+        self_attention_mask = compute_causal_mask(
+            batch_size,
+            input_length,
+            output_length,
+            cache_index,
+        )
         decoder_mask = merge_padding_and_attention_mask(
             decoder_sequence, decoder_padding_mask, decoder_attention_mask
         )
diff --git a/keras_nlp/layers/transformer_layer_utils.py b/keras_nlp/layers/transformer_layer_utils.py
index 78b83997..2ce24d8a 100644
--- a/keras_nlp/layers/transformer_layer_utils.py
+++ b/keras_nlp/layers/transformer_layer_utils.py
@@ -18,21 +18,28 @@ import tensorflow as tf
 from absl import logging
 
 
-def compute_causal_mask(inputs):
-    input_shape = tf.shape(inputs)
-    batch_size, sequence_length = input_shape[0], input_shape[1]
-    i = tf.range(sequence_length)[:, tf.newaxis]
-    j = tf.range(sequence_length)
-    mask = tf.cast(i >= j, dtype="int32")
-    mask = tf.reshape(mask, (1, input_shape[1], input_shape[1]))
-    mult = tf.concat(
-        [
-            tf.expand_dims(batch_size, -1),
-            tf.constant([1, 1], dtype=tf.int32),
-        ],
-        axis=0,
-    )
-    return tf.tile(mask, mult)
+def compute_causal_mask(batch_size, input_length, output_length, cache_index=0):
+    """Compute a causal attention mask for a transformer decoder.
+
+    Args:
+        batch_size: batch size for the mask.
+        input_length: the length of key/value tensors in the attention layer.
+        output_length: the length of query tensors in the attention layer.
+        cache_index: the current index for cached generation. If passed, the
+            query sequence will be considered to start at `cache_index` rather
+            than zero. For example, a causal mask with `output_length=1` and
+            `cache_index=5` would allow the query tensor to attend to the first
+            five positions of the key/value tensors.
+
+    Return:
+        A causal attention mask with shape
+        `(batch_size, output_length, input_length)` that can be passed to a
+        attention layer.
+    """
+    i = tf.range(output_length)[:, tf.newaxis] + cache_index
+    j = tf.range(input_length)
+    mask = tf.cast(i >= j, dtype="int32")[tf.newaxis, :, :]
+    return tf.broadcast_to(mask, (batch_size, output_length, input_length))
 
 
 def merge_padding_and_attention_mask(
@@ -40,7 +47,7 @@ def merge_padding_and_attention_mask(
     padding_mask,
     attention_mask,
 ):
-    """Merge padding mask with users' customized mask.
+    """Merge the padding mask with a customized attention mask.
 
     Args:
         inputs: the input sequence.
diff --git a/keras_nlp/models/albert/albert_masked_lm_preprocessor.py b/keras_nlp/models/albert/albert_masked_lm_preprocessor.py
index 1c874a0b..1f3f7348 100644
--- a/keras_nlp/models/albert/albert_masked_lm_preprocessor.py
+++ b/keras_nlp/models/albert/albert_masked_lm_preprocessor.py
@@ -118,9 +118,9 @@ class AlbertMaskedLMPreprocessor(AlbertPreprocessor):
 
     proto = bytes_io.getvalue()
 
-    tokenizer = AlbertTokenizer(proto=proto)
+    tokenizer = keras_nlp.models.AlbertTokenizer(proto=proto)
 
-    preprocessor = AlbertMaskedLMPreprocessor(
+    preprocessor = keras_nlp.models.AlbertMaskedLMPreprocessor(
         tokenizer=tokenizer
     )
 
diff --git a/keras_nlp/models/deberta_v3/deberta_v3_presets.py b/keras_nlp/models/deberta_v3/deberta_v3_presets.py
index b40d3879..6dd20e56 100644
--- a/keras_nlp/models/deberta_v3/deberta_v3_presets.py
+++ b/keras_nlp/models/deberta_v3/deberta_v3_presets.py
@@ -33,7 +33,7 @@ backbone_presets = {
             ),
             "params": 70682112,
             "official_name": "DeBERTaV3",
-            "path": "deberta",
+            "path": "deberta_v3",
         },
         "weights_url": "https://storage.googleapis.com/keras-nlp/models/deberta_v3_extra_small_en/v1/model.h5",
         "weights_hash": "d8e10327107e5c5e20b45548a5028619",
@@ -59,7 +59,7 @@ backbone_presets = {
             ),
             "params": 141304320,
             "official_name": "DeBERTaV3",
-            "path": "deberta",
+            "path": "deberta_v3",
         },
         "weights_url": "https://storage.googleapis.com/keras-nlp/models/deberta_v3_small_en/v1/model.h5",
         "weights_hash": "84118eb7c5a735f2061ecccaf71bb888",
@@ -85,7 +85,7 @@ backbone_presets = {
             ),
             "params": 183831552,
             "official_name": "DeBERTaV3",
-            "path": "deberta",
+            "path": "deberta_v3",
         },
         "weights_url": "https://storage.googleapis.com/keras-nlp/models/deberta_v3_base_en/v1/model.h5",
         "weights_hash": "cebce044aeed36aec9b94e3b8a255430",
@@ -111,7 +111,7 @@ backbone_presets = {
             ),
             "params": 434012160,
             "official_name": "DeBERTaV3",
-            "path": "deberta",
+            "path": "deberta_v3",
         },
         "weights_url": "https://storage.googleapis.com/keras-nlp/models/deberta_v3_large_en/v1/model.h5",
         "weights_hash": "bce7690f358a9e39304f8c0ebc71a745",
@@ -137,7 +137,7 @@ backbone_presets = {
             ),
             "params": 278218752,
             "official_name": "DeBERTaV3",
-            "path": "deberta",
+            "path": "deberta_v3",
         },
         "weights_url": "https://storage.googleapis.com/keras-nlp/models/deberta_v3_base_multi/v1/model.h5",
         "weights_hash": "26e5a824b26afd2ee336835bd337bbeb",
diff --git a/keras_nlp/models/distil_bert/distil_bert_masked_lm_preprocessor.py b/keras_nlp/models/distil_bert/distil_bert_masked_lm_preprocessor.py
index bf19585b..ff515ef3 100644
--- a/keras_nlp/models/distil_bert/distil_bert_masked_lm_preprocessor.py
+++ b/keras_nlp/models/distil_bert/distil_bert_masked_lm_preprocessor.py
@@ -67,7 +67,7 @@ class DistilBertMaskedLMPreprocessor(DistilBertPreprocessor):
 
     # Alternatively, you can create a preprocessor from your own vocabulary.
     # The usage is exactly the same as above.
-    vocab = ["[PAD]", "[UNK]", "[CLS]", "[SEP]"]
+    vocab = ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"]
     vocab += ["The", "qu", "##ick", "br", "##own", "fox", "tripped"]
     vocab += ["Call", "me", "Ish", "##mael", "."]
     vocab += ["Oh", "look", "a", "whale"]
diff --git a/keras_nlp/models/distil_bert/distil_bert_preprocessor.py b/keras_nlp/models/distil_bert/distil_bert_preprocessor.py
index 67a44e17..696c6ca7 100644
--- a/keras_nlp/models/distil_bert/distil_bert_preprocessor.py
+++ b/keras_nlp/models/distil_bert/distil_bert_preprocessor.py
@@ -139,7 +139,7 @@ class DistilBertPreprocessor(Preprocessor):
 
     # Alternatively, you can create a preprocessor from your own vocabulary.
     # The usage is exactly the same as above.
-    vocab = ["[PAD]", "[UNK]", "[CLS]", "[SEP]"]
+    vocab = ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"]
     vocab += ["The", "qu", "##ick", "br", "##own", "fox", "tripped"]
     vocab += ["Call", "me", "Ish", "##mael", "."]
     vocab += ["Oh", "look", "a", "whale"]
diff --git a/keras_nlp/models/f_net/f_net_classifier.py b/keras_nlp/models/f_net/f_net_classifier.py
index 6e43b76f..712ae7e8 100644
--- a/keras_nlp/models/f_net/f_net_classifier.py
+++ b/keras_nlp/models/f_net/f_net_classifier.py
@@ -67,18 +67,17 @@ class FNetClassifier(Task):
     }
     labels = [0, 3]
 
-    # Randomly initialize a Fnet backbone
+    # Randomly initialize a FNet backbone.
     backbone = keras_nlp.models.FNetBackbone(
         vocabulary_size=32000,
         num_layers=12,
-        num_heads=12,
         hidden_dim=768,
         intermediate_dim=3072,
         max_sequence_length=12,
     )
 
-    # Create a Fnet classifier and fit your data.
-    classifier = keras_nlp.models.FnetClassifier(
+    # Create a FNet classifier and fit your data.
+    classifier = keras_nlp.models.FNetClassifier(
         backbone,
         num_classes=4,
         preprocessor=None,
diff --git a/keras_nlp/models/f_net/f_net_masked_lm.py b/keras_nlp/models/f_net/f_net_masked_lm.py
index 86278f75..5b367397 100644
--- a/keras_nlp/models/f_net/f_net_masked_lm.py
+++ b/keras_nlp/models/f_net/f_net_masked_lm.py
@@ -76,7 +76,7 @@ class FNetMaskedLM(Task):
             [[1, 2, 0, 4, 0, 6, 7, 8]] * 2, shape=(2, 8)
         ),
         "segment_ids": tf.constant(
-            [[1, 0, 0, 4, 0, 6, 7, 8]] * 2, shape=(2, 8)
+            [[0, 0, 0, 1, 1, 1, 0, 0]] * 2, shape=(2, 8)
         ),
         "mask_positions": tf.constant([[2, 4]] * 2, shape=(2, 2))
     }
@@ -87,7 +87,6 @@ class FNetMaskedLM(Task):
     backbone = keras_nlp.models.FNetBackbone(
         vocabulary_size=50265,
         num_layers=12,
-        num_heads=12,
         hidden_dim=768,
         intermediate_dim=3072,
         max_sequence_length=12
diff --git a/keras_nlp/models/f_net/f_net_masked_lm_preprocessor.py b/keras_nlp/models/f_net/f_net_masked_lm_preprocessor.py
index d7c12f54..206e9144 100644
--- a/keras_nlp/models/f_net/f_net_masked_lm_preprocessor.py
+++ b/keras_nlp/models/f_net/f_net_masked_lm_preprocessor.py
@@ -112,8 +112,8 @@ class FNetMaskedLMPreprocessor(FNetPreprocessor):
         user_defined_symbols="[MASK]",
     )
     proto = bytes_io.getvalue()
-    tokenizer = FNetTokenizer(proto=proto)
-    preprocessor = FNetMaskedLMPreprocessor(tokenizer=tokenizer)
+    tokenizer = keras_nlp.models.FNetTokenizer(proto=proto)
+    preprocessor = keras_nlp.models.FNetMaskedLMPreprocessor(tokenizer=tokenizer)
     ```
     """
 
diff --git a/keras_nlp/models/f_net/f_net_presets.py b/keras_nlp/models/f_net/f_net_presets.py
index 3d55e084..63aa10b1 100644
--- a/keras_nlp/models/f_net/f_net_presets.py
+++ b/keras_nlp/models/f_net/f_net_presets.py
@@ -32,7 +32,7 @@ backbone_presets = {
             ),
             "params": 82861056,
             "official_name": "FNet",
-            "path": "fnet",
+            "path": "f_net",
         },
         "weights_url": "https://storage.googleapis.com/keras-nlp/models/f_net_base_en/v1/model.h5",
         "weights_hash": "35db90842b85a985a0e54c86c00746fe",
@@ -57,7 +57,7 @@ backbone_presets = {
             ),
             "params": 236945408,
             "official_name": "FNet",
-            "path": "fnet",
+            "path": "f_net",
         },
         "weights_url": "https://storage.googleapis.com/keras-nlp/models/f_net_large_en/v1/model.h5",
         "weights_hash": "7ae4a3faa67ff054f8cecffb5619f779",
diff --git a/keras_nlp/models/gpt2/gpt2_causal_lm.py b/keras_nlp/models/gpt2/gpt2_causal_lm.py
index 61526b85..63537fb6 100644
--- a/keras_nlp/models/gpt2/gpt2_causal_lm.py
+++ b/keras_nlp/models/gpt2/gpt2_causal_lm.py
@@ -17,7 +17,6 @@ import copy
 
 import tensorflow as tf
 from tensorflow import keras
-from tensorflow.compiler.tf2xla.python.xla import dynamic_update_slice
 
 import keras_nlp
 from keras_nlp.models.gpt2.gpt2_backbone import GPT2Backbone
@@ -201,7 +200,7 @@ class GPT2CausalLM(Task):
     def preprocessor_cls(cls):
         return GPT2CausalLMPreprocessor
 
-    def call_with_cache(self, token_ids, padding_mask, cache, cache_index=None):
+    def call_with_cache(self, token_ids, padding_mask, cache, cache_index):
         """Forward pass of `GPT2CausalLM` with cache.
 
         `call_with_cache` adds an additional forward pass for the model for
@@ -213,101 +212,49 @@ class GPT2CausalLM(Task):
             token_ids: a dense int Tensor, input token ids.
             padding_mask: a dense bool Tensor, input padding mask.
             cache: a dense float Tensor, the cache of key and value.
-            cache_index: int, or int Tensor, defaults to None. If set, it
-                represents the index of current inputs in the whole sequence.
+            cache_index: int, or int Tensor. The index of current inputs in the
+                whole sequence.
 
         Returns:
             A (logits, cache) tuple. Where the first output is the language
             model logits for the input token_ids and the second output is the
             cache.
         """
-        transformer_layers = []
-        for i in range(self.backbone.num_layers):
-            transformer_layers.append(
-                self.backbone.get_layer(f"transformer_layer_{i}")
-            )
-        layer_norm = self.backbone.get_layer("layer_norm")
-        token_embeddings = self.backbone.get_layer("token_embedding")
-        position_embeddings = self.backbone.get_layer("position_embedding")
-        embeddings_add = self.backbone.get_layer("embeddings_add")
-        embeddings_dropout = self.backbone.get_layer("embeddings_dropout")
-
-        token_embedding = token_embeddings(token_ids)
-        if cache_index is None:
-            position_embedding = position_embeddings(token_embedding)
-        else:
-            position_embedding = position_embeddings.position_embeddings[
-                cache_index, :
-            ]
-        x = embeddings_add((token_embedding, position_embedding))
-        x = embeddings_dropout(x)
-        # Each `TransformerDecoder` layer has a cache, we update separately.
+        token_embedding = self.backbone.get_layer("token_embedding")(token_ids)
+        position_embedding = self.backbone.get_layer("position_embedding")(
+            token_embedding, start_index=cache_index
+        )
+        x = self.backbone.get_layer("embeddings_add")(
+            (token_embedding, position_embedding)
+        )
+        x = self.backbone.get_layer("embeddings_dropout")(x)
+        # Each decoder layer has a cache; we update them separately.
         caches = tf.unstack(cache, axis=1)
-        for i, transformer_layer in enumerate(transformer_layers):
+        for i in range(self.backbone.num_layers):
             current_cache = caches[i]
-            x, current_cache = transformer_layer(
+            x, next_cache = self.backbone.get_layer(f"transformer_layer_{i}")(
                 x,
                 decoder_padding_mask=padding_mask,
                 cache=current_cache,
                 cache_index=cache_index,
             )
-            caches[i] = current_cache
+            caches[i] = next_cache
         cache = tf.stack(caches, axis=1)
-        x = layer_norm(x)
+        x = self.backbone.get_layer("layer_norm")(x)
         x = tf.matmul(
             x,
-            self.backbone.token_embedding.embeddings,
+            self.backbone.get_layer("token_embedding").embeddings,
             transpose_b=True,
         )
         return x, cache
 
-    def build_initial_cache(self, initial_inputs, max_length):
-        """Build initial cache based on the prompt.
-
-        This method should be called before the decoding loop to build the
-        initial cache. The cache is of shape [batch_size, `self.num_layers`, 2
-        max_length, `self.num_heads`, `self.hidden_dim // self.num_heads`].
-        The first dim represents it's a key or value in multi-head attention.
-
-        Args:
-            initial_inputs: a dense Tensor, the initial inputs to the decoding
-                loop.
-            max_length: int, the max length of the generated sequence.
-
-        Returns:
-            A (logits, cache) tuple. The first output is the language
-            model logits for the input token_ids and the second output is the
-            cache.
-        """
-        token_ids = initial_inputs["token_ids"]
-        padding_mask = initial_inputs["padding_mask"]
-
-        if max_length < self.backbone.max_sequence_length:
-            token_ids = token_ids[:, :max_length]
-            padding_mask = padding_mask[:, :max_length]
-
-        batch_size = tf.shape(token_ids)[0]
-        outputs = tf.zeros(
-            [batch_size, max_length, self.backbone.vocabulary_size]
-        )
-        cache = tf.zeros(
-            [
-                batch_size,
-                self.backbone.num_layers,
-                2,
-                max_length,
-                self.backbone.num_heads,
-                self.backbone.hidden_dim // self.backbone.num_heads,
-            ],
-        )
-
-        output, cache = self.call_with_cache(
-            token_ids,
-            padding_mask,
-            cache=cache,
-        )
-        outputs = dynamic_update_slice(outputs, output, [0, 0, 0])
-        return outputs, cache
+    def build_empty_cache(self, batch_size, max_length):
+        """Build an empty cache for use with `call_with_cache()`."""
+        num_layers = self.backbone.num_layers
+        num_heads = self.backbone.num_heads
+        head_dim = self.backbone.hidden_dim // self.backbone.num_heads
+        shape = [batch_size, num_layers, 2, max_length, num_heads, head_dim]
+        return tf.zeros(shape)
 
     def _get_token_probability(
         self,
@@ -316,23 +263,9 @@ class GPT2CausalLM(Task):
         cache=None,
         cache_index=None,
     ):
-        model_inputs = {
-            "token_ids": prompt,
-            "padding_mask": mask,
-        }
-        if cache_index is None and cache is None:
-            return self(model_inputs)
-        if cache_index is not None:
-            batch_size = tf.shape(prompt)[0]
-            prompt = tf.slice(prompt, [0, cache_index], [batch_size, 1])
-            mask = mask[:, : cache_index + 1]
-        output, cache = self.call_with_cache(
-            prompt,
-            mask,
-            cache,
-            cache_index,
-        )
-        return output, cache
+        batch_size = tf.shape(prompt)[0]
+        prompt = tf.slice(prompt, [0, cache_index], [batch_size, 1])
+        return self.call_with_cache(prompt, mask, cache, cache_index)
 
     def generate(
         self,
@@ -356,28 +289,38 @@ class GPT2CausalLM(Task):
         """
         if self.preprocessor is None:
             raise ValueError(
-                "`self.preprocessor` is None, please make sure `preprocessor` "
-                "is set before calling `generate`."
+                "`self.preprocessor` is `None`, please make sure "
+                "`preprocessor` is set before calling `generate`."
             )
-        end_token_id = self.preprocessor.tokenizer.end_token_id
-
         sampler = keras_nlp.samplers.get(sampler)
-        if hasattr(self, "jit_compile"):
-            # `jit_compile` is a public property as of tf 2.12. hasattr is for
-            # backward compat.
-            sampler.jit_compile = self.jit_compile
+        # `jit_compile` is a property of keras.Model after tf 2.12.
+        # Use `getattr()` for backwards compatibility.
+        sampler.jit_compile = getattr(self, "jit_compile", True)
         sampler.run_eagerly = self.run_eagerly
-        x, _, _ = self.preprocessor(prompt)
-        if len(x["token_ids"].shape) == 1:
-            x["token_ids"] = x["token_ids"][tf.newaxis, :]
-            x["padding_mask"] = x["padding_mask"][tf.newaxis, :]
-        _, cache = self.build_initial_cache(x, max_length)
-        next_token_probability = self._get_token_probability
+
+        # Tokenize.
+        prompt = self.preprocessor.tokenizer(prompt)
+
+        # Create and seed the cache before generation.
+        token_ids = prompt
+        if prompt.shape.rank == 1:
+            token_ids = tf.RaggedTensor.from_tensor(prompt[tf.newaxis, :])
+        token_ids = token_ids.to_tensor(shape=(None, max_length))
+        # Pass a padding mask of all ones when seeing the cache. The mask will
+        # not affect cached key/values for input tokens we care about.
+        padding_mask = tf.ones_like(token_ids, dtype=tf.bool)
+        batch_size = tf.shape(token_ids)[0]
+        cache = self.build_empty_cache(batch_size, max_length)
+        _, cache = self.call_with_cache(token_ids, padding_mask, cache, 0)
+
+        # Run generation.
         generated = sampler(
-            self.preprocessor.tokenizer(prompt),
-            next_token_probability,
+            prompt,
+            self._get_token_probability,
             max_length=max_length,
-            end_token_id=end_token_id,
+            end_token_id=self.preprocessor.tokenizer.end_token_id,
             cache=cache,
         )
+
+        # Detokenize.
         return self.preprocessor.tokenizer.detokenize(generated)
diff --git a/keras_nlp/samplers/sampler.py b/keras_nlp/samplers/sampler.py
index 97b5d067..50c58e56 100644
--- a/keras_nlp/samplers/sampler.py
+++ b/keras_nlp/samplers/sampler.py
@@ -74,7 +74,7 @@ class Sampler:
     VOCAB_SIZE = 10
 
     # Create a dummy model to predict the next token. Note that the output is
-    # random without training, here we jsut demo how `samplers` works.
+    # random without training, here we just demo how `samplers` works.
     model = keras.Sequential(
         [
             keras.Input(shape=[None]),
@@ -172,18 +172,6 @@ class Sampler:
         prompt = tf.ragged.boolean_mask(prompt, mask)
         return prompt, mask
 
-    def _validate_token_probability_fn(
-        self, token_probability_fn, prompt, mask
-    ):
-        """Helper method to validate `token_probability_fn` output."""
-        test_pred = token_probability_fn(prompt, mask=mask)
-        if len(test_pred.shape) != 3:
-            raise ValueError(
-                "Output of `token_probability_fn` is not a 3D tensor, "
-                "please provide a function with the output shape "
-                "[batch_size, sequence_length, vocab_size]."
-            )
-
     def _pad_prompt(self, prompt, max_length):
         """Pad prompt to `max_length`."""
         mask = tf.ones_like(prompt, dtype=tf.bool)
@@ -237,7 +225,6 @@ class Sampler:
         # current prompt.
         prompt, mask = self._pad_prompt(prompt, max_length)
         original_padding_mask = tf.identity(mask)
-        self._validate_token_probability_fn(token_probability_fn, prompt, mask)
 
         # Convert `sample` method to a `tf.function` if `self.run_eagerly=False`
         # , and turn on `jit_compile` accordingly.
@@ -359,7 +346,7 @@ class Sampler:
             return [current_index, prompt, mask, cache]
 
         if cache is None:
-            current_index, prompt, mask = tf.while_loop(
+            _, prompt, _ = tf.while_loop(
                 cond=lambda current_index, prompt, mask: tf.less(
                     current_index, max_length
                 ),
@@ -368,7 +355,7 @@ class Sampler:
             )
             return prompt
         # Run a while loop till `max_length` of tokens has been generated.
-        current_index, prompt, mask, cache = tf.while_loop(
+        _, prompt, _, _ = tf.while_loop(
             cond=lambda current_index, prompt, mask, cache: tf.less(
                 current_index, max_length
             ),


Test patch:
diff --git a/keras_nlp/layers/cached_multi_head_attention_test.py b/keras_nlp/layers/cached_multi_head_attention_test.py
index d2e21bc2..0fc20c6f 100644
--- a/keras_nlp/layers/cached_multi_head_attention_test.py
+++ b/keras_nlp/layers/cached_multi_head_attention_test.py
@@ -28,7 +28,11 @@ class CachedMultiHeadAttentionTest(tf.test.TestCase, parameterized.TestCase):
         x = tf.random.uniform(shape=[2, 2, 8])
         layer(query=x, value=x)
 
-    def test_cache_call_is_correct(self):
+    @parameterized.named_parameters(
+        ("graph", False),
+        ("eager", True),
+    )
+    def test_cache_call_is_correct(self, eager):
         batch_size = 2
         seq_len = 5
         num_heads = 2
@@ -37,43 +41,36 @@ class CachedMultiHeadAttentionTest(tf.test.TestCase, parameterized.TestCase):
         layer = CachedMultiHeadAttention(num_heads=num_heads, key_dim=key_dim)
         x = tf.random.uniform(shape=[batch_size, seq_len, num_heads * key_dim])
         cache = tf.zeros([batch_size, 2, seq_len, num_heads, key_dim])
-        # Build the intial cache.
-        intial_seq_len = 2
-        initial_inputs = x[:, :intial_seq_len, :]
+        # Use a causal mask.
+        mask = tf.linalg.band_part(tf.ones((seq_len, seq_len)), -1, 0)
         outputs = tf.zeros_like(x)
-        output, cache = layer(
-            query=initial_inputs,
-            value=initial_inputs,
-            use_causal_mask=True,
-            cache=cache,
-        )
-        # Update the outputs in place.
-        outputs = dynamic_update_slice(outputs, output, [0, 0, 0])
 
-        def call(i, cache, outputs):
-            def loop_body(i, cache, outputs):
+        def call(outputs, cache):
+            def loop_body(i, outputs, cache):
                 # Compute the rest tokens.
-                current_input = x[:, i : i + 1, :]
-                output, cache = layer(
-                    query=current_input,
-                    value=current_input,
-                    use_causal_mask=True,
+                next_input = x[:, i : i + 1, :]
+                next_mask = mask[i : i + 1, :]
+                next_output, cache = layer(
+                    query=next_input,
+                    value=next_input,
                     cache=cache,
                     cache_index=i,
+                    attention_mask=next_mask,
                 )
-                outputs = dynamic_update_slice(outputs, output, [0, i, 0])
-                return i + 1, cache, outputs
+                outputs = dynamic_update_slice(outputs, next_output, [0, i, 0])
+                return i + 1, outputs, cache
 
-            i, cache, cached_outputs = tf.while_loop(
-                cond=lambda i, cache, outputs: i < seq_len,
+            _, outputs, cache = tf.while_loop(
+                cond=lambda i, outputs, cache: i < seq_len,
                 body=loop_body,
-                loop_vars=[i, cache, outputs],
+                loop_vars=[0, outputs, cache],
             )
-            return cached_outputs
+            return outputs, cache
 
-        cached_outputs = call(intial_seq_len, cache, outputs)
-        graph_call = tf.function(call)
-        graph_cached_outputs = graph_call(intial_seq_len, cache, outputs)
-        normal_outputs, _ = layer(query=x, value=x, use_causal_mask=True)
-        self.assertAllClose(cached_outputs, normal_outputs)
-        self.assertAllClose(graph_cached_outputs, normal_outputs)
+        call = call if eager else tf.function(call)
+        output, cache = call(outputs, cache)
+
+        no_loop_outputs, _ = layer(x, x, attention_mask=mask)
+        _, no_loop_cache = layer(x, x, cache=cache, attention_mask=mask)
+        self.assertAllClose(output, no_loop_outputs)
+        self.assertAllClose(cache, no_loop_cache)
diff --git a/keras_nlp/layers/position_embedding_test.py b/keras_nlp/layers/position_embedding_test.py
index 8f6389e1..f627cfde 100644
--- a/keras_nlp/layers/position_embedding_test.py
+++ b/keras_nlp/layers/position_embedding_test.py
@@ -284,24 +284,12 @@ class PositionEmbeddingLayerTest(tf.test.TestCase, parameterized.TestCase):
     def test_get_config_and_from_config(self):
         max_sequence_length = 40
         test_layer = position_embedding.PositionEmbedding(
-            sequence_length=max_sequence_length
+            sequence_length=max_sequence_length,
+            initializer="zeros",
         )
         config = test_layer.get_config()
-        expected_config_subset = {
-            "sequence_length": max_sequence_length,
-            "initializer": {
-                "class_name": "GlorotUniform",
-                "config": {"seed": None},
-            },
-        }
-        self.assertEqual(config, {**config, **expected_config_subset})
-
-        restored_encoder = position_embedding.PositionEmbedding.from_config(
-            config,
-        )
-        self.assertEqual(
-            restored_encoder.get_config(), {**config, **expected_config_subset}
-        )
+        restored = position_embedding.PositionEmbedding.from_config(config)
+        self.assertEqual(restored.get_config(), config)
 
     @parameterized.named_parameters(
         ("tf_format", "tf", "model"),
diff --git a/keras_nlp/layers/transformer_decoder_test.py b/keras_nlp/layers/transformer_decoder_test.py
index 4896c23f..6b1929b2 100644
--- a/keras_nlp/layers/transformer_decoder_test.py
+++ b/keras_nlp/layers/transformer_decoder_test.py
@@ -266,56 +266,50 @@ class TransformerDecoderTest(tf.test.TestCase, parameterized.TestCase):
         outputs = decoder(decoder_sequence)
         self.assertAllEqual(outputs._keras_mask, mask)
 
-    def test_cached_decoding_is_correct(self):
+    @parameterized.named_parameters(
+        ("graph", False),
+        ("eager", True),
+    )
+    def test_cached_decoding_is_correct(self, eager):
         batch_size = 2
         seq_len = 5
         num_heads = 2
-        hidden_dim = 8
+        head_dim = 4
 
         layer = transformer_decoder.TransformerDecoder(
             intermediate_dim=4,
             num_heads=num_heads,
         )
-        x = tf.random.uniform(shape=[batch_size, seq_len, hidden_dim])
-        cache = tf.zeros(
-            [batch_size, 2, seq_len, num_heads, hidden_dim // num_heads]
-        )
-        # Build the intial cache.
-        initial_seq_len = 2
-        initial_inputs = x[:, :initial_seq_len, :]
+        x = tf.random.uniform(shape=[batch_size, seq_len, num_heads * head_dim])
+        cache = tf.zeros([batch_size, 2, seq_len, num_heads, head_dim])
         outputs = tf.zeros_like(x)
-        output, cache = layer(
-            decoder_sequence=initial_inputs,
-            cache=cache,
-        )
-        # Update the outputs in place.
-        outputs = dynamic_update_slice(outputs, output, [0, 0, 0])
 
-        def call(i, cache, outputs):
-            def loop_body(i, cache, outputs):
+        def call(outputs, cache):
+            def loop_body(i, outputs, cache):
                 # Compute the rest tokens.
-                current_input = x[:, i : i + 1, :]
-                output, cache = layer(
-                    decoder_sequence=current_input,
+                next_input = x[:, i : i + 1, :]
+                next_output, cache = layer(
+                    decoder_sequence=next_input,
                     cache=cache,
                     cache_index=i,
                 )
-                outputs = dynamic_update_slice(outputs, output, [0, i, 0])
-                return i + 1, cache, outputs
+                outputs = dynamic_update_slice(outputs, next_output, [0, i, 0])
+                return i + 1, outputs, cache
 
-            i, cache, cached_outputs = tf.while_loop(
-                cond=lambda i, cache, outputs: i < seq_len,
+            _, outputs, cache = tf.while_loop(
+                cond=lambda i, outputs, cache: i < seq_len,
                 body=loop_body,
-                loop_vars=[i, cache, outputs],
+                loop_vars=[0, outputs, cache],
             )
-            return cached_outputs
-
-        cached_outputs = call(initial_seq_len, cache, outputs)
-        graph_call = tf.function(call)
-        graph_cached_outputs = graph_call(initial_seq_len, cache, outputs)
-        normal_outputs = layer(decoder_sequence=x)
-        self.assertAllClose(cached_outputs, normal_outputs)
-        self.assertAllClose(graph_cached_outputs, normal_outputs)
+            return outputs, cache
+
+        call = call if eager else tf.function(call)
+        output, cache = call(outputs, cache)
+
+        no_loop_outputs = layer(x)
+        _, no_loop_cache = layer(x, cache=cache)
+        self.assertAllClose(output, no_loop_outputs)
+        self.assertAllClose(cache, no_loop_cache)
 
     @parameterized.named_parameters(
         ("tf_format", "tf", "model"),
diff --git a/keras_nlp/layers/transformer_layer_utils_test.py b/keras_nlp/layers/transformer_layer_utils_test.py
index 73ef7bcf..82ee1d93 100644
--- a/keras_nlp/layers/transformer_layer_utils_test.py
+++ b/keras_nlp/layers/transformer_layer_utils_test.py
@@ -19,8 +19,7 @@ import keras_nlp.layers.transformer_layer_utils as utils
 
 class TransformerEncoderTest(tf.test.TestCase):
     def test_compute_causal_mask(self):
-        inputs = tf.random.uniform(shape=[1, 2, 2])
-        mask = utils.compute_causal_mask(inputs)
+        mask = utils.compute_causal_mask(1, 2, 2)
         self.assertTrue((mask.numpy() == [[1, 0], [1, 1]]).all())
 
     def test_merge_padding_and_attention_mask(self):
diff --git a/keras_nlp/models/bert/bert_presets_test.py b/keras_nlp/models/bert/bert_presets_test.py
index b06c44dc..bd2fb435 100644
--- a/keras_nlp/models/bert/bert_presets_test.py
+++ b/keras_nlp/models/bert/bert_presets_test.py
@@ -174,7 +174,7 @@ class BertPresetFullTest(tf.test.TestCase, parameterized.TestCase):
 
     This every presets for BERT and is only run manually.
     Run with:
-    `pytest keras_nlp/models/bert_presets_test.py --run_extra_large`
+    `pytest keras_nlp/models/bert/bert_presets_test.py --run_extra_large`
     """
 
     @parameterized.named_parameters(
diff --git a/keras_nlp/tests/doc_tests/docstring_test.py b/keras_nlp/tests/doc_tests/docstring_test.py
index f88c9aa5..12d31e8c 100644
--- a/keras_nlp/tests/doc_tests/docstring_test.py
+++ b/keras_nlp/tests/doc_tests/docstring_test.py
@@ -108,12 +108,15 @@ def test_fenced_docstrings():
             "keras_nlp.models.backbone",
             "keras_nlp.models.preprocessor",
             "keras_nlp.models.task",
-            # Preprocessors and tokenizers which use `model.spm` (temporary).
-            "keras_nlp.models.albert.albert_preprocessor",
-            "keras_nlp.models.albert.albert_tokenizer",
             "keras_nlp.models.xlm_roberta.xlm_roberta_preprocessor",
             "keras_nlp.models.f_net.f_net_preprocessor",
             "keras_nlp.models.f_net.f_net_tokenizer",
+            # Preprocessors and tokenizers which use `model.spm` (temporary).
+            "keras_nlp.models.albert.albert_preprocessor",
+            "keras_nlp.models.albert.albert_tokenizer",
+            "keras_nlp.tokenizers.byte_pair_tokenizer",
+            "keras_nlp.tokenizers.sentence_piece_tokenizer",
+            "keras_nlp.tokenizers.word_piece_tokenizer",
         ]:
             continue