Repository: keras-team/keras-nlp
Pull request #1421 (2024-02-02 18:14:12)

Add layer attributes for all functional models
This proposes a new style for backbone and tasks:

```python
# === Layers ===
self.token_embedding = ...
self.transformer_layers = []
for i in range(num_layers):
    self.transformer_layers.append(...)

# === Functional Model ===
inputs = keras.Input(...)
x = self.token_embedding(inputs)
for layer in self.transformer_layers:
    x = layer(x)
outputs = x
super.__init__(inputs, outputs)

# === Config ===
self.num_layers = num_layers
```

The main goal is to allow more readable manipulations of the model, e.g. a custom LoRA routine:

```python
backbone = GPT2Backbone(...)
for layer in self.transformer_layers:
    # Use different lora rank for different matrices.
    layer.self_attention.query_projection.enable_lora(16)
    layer.self_attention.key_projection.enable_lora(4)
```

mattdangerw: Draft PR just for feedback for now. I've only touched less than half of the models. Tests will fail.
mattdangerw: I also removed unused features from our pipeline model class, just to keep our tasks as clean as possible.

Patch:
diff --git a/keras_nlp/models/albert/albert_backbone.py b/keras_nlp/models/albert/albert_backbone.py
index 414bb97e..99a7cc97 100644
--- a/keras_nlp/models/albert/albert_backbone.py
+++ b/keras_nlp/models/albert/albert_backbone.py
@@ -118,67 +118,47 @@ class AlbertBackbone(Backbone):
                 f"`num_layers={num_layers}` and `num_groups={num_groups}`."
             )
 
-        # Index of classification token in the vocabulary
-        cls_token_index = 0
-        # Inputs
-        token_id_input = keras.Input(
-            shape=(None,), dtype="int32", name="token_ids"
-        )
-        segment_id_input = keras.Input(
-            shape=(None,), dtype="int32", name="segment_ids"
-        )
-        padding_mask = keras.Input(
-            shape=(None,), dtype="int32", name="padding_mask"
-        )
-
-        # Embed tokens, positions, and segment ids.
-        token_embedding_layer = ReversibleEmbedding(
+        # === Layers ===
+        self.token_embedding = ReversibleEmbedding(
             input_dim=vocabulary_size,
             output_dim=embedding_dim,
             embeddings_initializer=albert_kernel_initializer(),
             name="token_embedding",
         )
-        token_embedding = token_embedding_layer(token_id_input)
-        position_embedding = PositionEmbedding(
+        self.position_embedding = PositionEmbedding(
             initializer=albert_kernel_initializer(),
             sequence_length=max_sequence_length,
             name="position_embedding",
-        )(token_embedding)
-        segment_embedding = keras.layers.Embedding(
+        )
+        self.segment_embedding = keras.layers.Embedding(
             input_dim=num_segments,
             output_dim=embedding_dim,
             embeddings_initializer=albert_kernel_initializer(),
             name="segment_embedding",
-        )(segment_id_input)
-
-        # Sum, normalize and apply dropout to embeddings.
-        x = keras.layers.Add()(
-            (token_embedding, position_embedding, segment_embedding)
         )
-        x = keras.layers.LayerNormalization(
+        self.embeddings_add = keras.layers.Add(
+            name="embeddings_add",
+        )
+        self.embeddings_layer_norm = keras.layers.LayerNormalization(
             name="embeddings_layer_norm",
             axis=-1,
             epsilon=1e-12,
             dtype="float32",
-        )(x)
-        x = keras.layers.Dropout(
+        )
+        self.embeddings_dropout = keras.layers.Dropout(
             dropout,
             name="embeddings_dropout",
-        )(x)
-
-        # Project the embedding to `hidden_dim`.
-        x = keras.layers.Dense(
+        )
+        self.embeddings_projection = keras.layers.Dense(
             hidden_dim,
             kernel_initializer=albert_kernel_initializer(),
             name="embedding_projection",
-        )(x)
-
-        def get_group_layer(group_idx):
-            """Defines a group `num_inner_repetitions` transformer layers and
-            returns the callable.
-            """
-            transformer_layers = [
-                TransformerEncoder(
+        )
+        self.transformer_layers = []
+        for group_idx in range(num_groups):
+            inner_layers = []
+            for inner_idx in range(num_inner_repetitions):
+                layer = TransformerEncoder(
                     num_heads=num_heads,
                     intermediate_dim=intermediate_dim,
                     activation=gelu_approximate,
@@ -187,43 +167,51 @@ class AlbertBackbone(Backbone):
                     kernel_initializer=albert_kernel_initializer(),
                     name=f"group_{group_idx}_inner_layer_{inner_idx}",
                 )
-                for inner_idx in range(num_inner_repetitions)
-            ]
-
-            def call(x, padding_mask):
-                for transformer_layer in transformer_layers:
-                    x = transformer_layer(x, padding_mask=padding_mask)
-                return x
-
-            return call
-
-        num_calls_per_group = num_layers // num_groups
-        for group_idx in range(num_groups):
-            # Define the group. A group in ALBERT terminology is any number of
-            # repeated attention and FFN blocks.
-            group_layer = get_group_layer(group_idx)
-
-            # Assume num_layers = 8, num_groups = 4. Then, the order of group
-            # calls will be 0, 0, 1, 1, 2, 2, 3, 3.
-            for call in range(num_calls_per_group):
-                x = group_layer(x, padding_mask=padding_mask)
-
-        # Construct the two ALBERT outputs. The pooled output is a dense layer on
-        # top of the [CLS] token.
-        sequence_output = x
-        pooled_output = keras.layers.Dense(
+                inner_layers.append(layer)
+            self.transformer_layers.append(inner_layers)
+        self.pooled_dense = keras.layers.Dense(
             hidden_dim,
             kernel_initializer=albert_kernel_initializer(),
             activation="tanh",
             name="pooled_dense",
-        )(x[:, cls_token_index, :])
+        )
 
-        # Instantiate using Functional API Model constructor
+        # === Functional Model ===
+        # Inputs
+        token_id_input = keras.Input(
+            shape=(None,), dtype="int32", name="token_ids"
+        )
+        segment_id_input = keras.Input(
+            shape=(None,), dtype="int32", name="segment_ids"
+        )
+        padding_mask_input = keras.Input(
+            shape=(None,), dtype="int32", name="padding_mask"
+        )
+        # Embed tokens, positions, and segment ids.
+        tokens = self.token_embedding(token_id_input)
+        positions = self.position_embedding(tokens)
+        segments = self.segment_embedding(segment_id_input)
+        # Sum, normalize and apply dropout to embeddings.
+        x = self.embeddings_add((tokens, positions, segments))
+        x = self.embeddings_layer_norm(x)
+        x = self.embeddings_dropout(x)
+        x = self.embeddings_projection(x)
+        # Call transformer layers with repeated groups.
+        num_calls_per_group = num_layers // num_groups
+        for group in self.transformer_layers:
+            for _ in range(num_calls_per_group):
+                for transformer_layer in group:
+                    x = transformer_layer(x, padding_mask=padding_mask_input)
+        # Construct the two ALBERT outputs. The pooled output is a dense layer
+        # on top of the [CLS] token.
+        sequence_output = x
+        cls_token_index = 0
+        pooled_output = self.pooled_dense(x[:, cls_token_index, :])
         super().__init__(
             inputs={
                 "token_ids": token_id_input,
                 "segment_ids": segment_id_input,
-                "padding_mask": padding_mask,
+                "padding_mask": padding_mask_input,
             },
             outputs={
                 "sequence_output": sequence_output,
@@ -231,7 +219,8 @@ class AlbertBackbone(Backbone):
             },
             **kwargs,
         )
-        # All references to `self` below this line
+
+        # === Config ===
         self.vocabulary_size = vocabulary_size
         self.num_layers = num_layers
         self.num_heads = num_heads
@@ -244,7 +233,6 @@ class AlbertBackbone(Backbone):
         self.max_sequence_length = max_sequence_length
         self.num_segments = num_segments
         self.cls_token_index = cls_token_index
-        self.token_embedding = token_embedding_layer
 
     def get_config(self):
         config = super().get_config()
diff --git a/keras_nlp/models/albert/albert_classifier.py b/keras_nlp/models/albert/albert_classifier.py
index b0ed7bca..e8621cf6 100644
--- a/keras_nlp/models/albert/albert_classifier.py
+++ b/keras_nlp/models/albert/albert_classifier.py
@@ -155,30 +155,37 @@ class AlbertClassifier(Task):
         dropout=0.1,
         **kwargs,
     ):
-        inputs = backbone.input
-        pooled = backbone(inputs)["pooled_output"]
-        pooled = keras.layers.Dropout(dropout)(pooled)
-        outputs = keras.layers.Dense(
+        # === Layers ===
+        self.backbone = backbone
+        self.preprocessor = preprocessor
+        self.output_dense = keras.layers.Dense(
             num_classes,
             kernel_initializer=albert_kernel_initializer(),
             activation=activation,
             name="logits",
-        )(pooled)
-        # Instantiate using Functional API Model constructor
+        )
+        self.output_dropout = keras.layers.Dropout(
+            dropout,
+            name="output_dropout",
+        )
+
+        # === Functional Model ===
+        inputs = backbone.input
+        pooled = backbone(inputs)["pooled_output"]
+        pooled = self.output_dropout(pooled)
+        outputs = self.output_dense(pooled)
         super().__init__(
             inputs=inputs,
             outputs=outputs,
-            include_preprocessing=preprocessor is not None,
             **kwargs,
         )
-        # All references to `self` below this line
-        self._backbone = backbone
-        self._preprocessor = preprocessor
+
+        # === Config ===
         self.num_classes = num_classes
         self.activation = keras.activations.get(activation)
         self.dropout = dropout
 
-        # Default compilation
+        # === Default compilation ===
         logit_output = self.activation == keras.activations.linear
         self.compile(
             loss=keras.losses.SparseCategoricalCrossentropy(
diff --git a/keras_nlp/models/albert/albert_masked_lm.py b/keras_nlp/models/albert/albert_masked_lm.py
index e95af7c2..139be0a7 100644
--- a/keras_nlp/models/albert/albert_masked_lm.py
+++ b/keras_nlp/models/albert/albert_masked_lm.py
@@ -97,32 +97,35 @@ class AlbertMaskedLM(Task):
     """
 
     def __init__(self, backbone, preprocessor=None, **kwargs):
+        # === Layers ===
+        self.backbone = backbone
+        self.preprocessor = preprocessor
+        self.masked_lm_head = MaskedLMHead(
+            vocabulary_size=backbone.vocabulary_size,
+            token_embedding=backbone.token_embedding,
+            intermediate_activation=gelu_approximate,
+            kernel_initializer=albert_kernel_initializer(),
+            name="mlm_head",
+        )
+
+        # === Functional Model ===
         inputs = {
             **backbone.input,
             "mask_positions": keras.Input(
                 shape=(None,), dtype="int32", name="mask_positions"
             ),
         }
-
         backbone_outputs = backbone(backbone.input)
-        outputs = MaskedLMHead(
-            vocabulary_size=backbone.vocabulary_size,
-            token_embedding=backbone.token_embedding,
-            intermediate_activation=gelu_approximate,
-            kernel_initializer=albert_kernel_initializer(),
-            name="mlm_head",
-        )(backbone_outputs["sequence_output"], inputs["mask_positions"])
-
+        outputs = self.masked_lm_head(
+            backbone_outputs["sequence_output"], inputs["mask_positions"]
+        )
         super().__init__(
             inputs=inputs,
             outputs=outputs,
-            include_preprocessing=preprocessor is not None,
-            **kwargs
+            **kwargs,
         )
 
-        self.backbone = backbone
-        self.preprocessor = preprocessor
-
+        # === Default compilation ===
         self.compile(
             loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
             optimizer=keras.optimizers.Adam(5e-5),
diff --git a/keras_nlp/models/backbone.py b/keras_nlp/models/backbone.py
index 69da5659..9e22be3f 100644
--- a/keras_nlp/models/backbone.py
+++ b/keras_nlp/models/backbone.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from keras_nlp.backend import config
 from keras_nlp.backend import keras
 from keras_nlp.utils.preset_utils import check_preset_class
 from keras_nlp.utils.preset_utils import load_from_preset
@@ -23,24 +24,40 @@ from keras_nlp.utils.python_utils import format_docstring
 class Backbone(keras.Model):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self._token_embedding = None
         self._functional_layer_ids = set(
             id(layer) for layer in self._flatten_layers()
         )
+        self._initialized = True
 
     def __dir__(self):
-        # Temporary fixes for weight saving. This mimics the following PR for
+        if config.keras_3():
+            return super().__dir__()
+
+        # Temporary fixes for Keras 2 saving. This mimics the following PR for
         # older version of Keras: https://github.com/keras-team/keras/pull/18982
         def filter_fn(attr):
-            if attr == "_layer_checkpoint_dependencies":
+            if attr in [
+                "_layer_checkpoint_dependencies",
+                "transformer_layers",
+                "encoder_transformer_layers",
+                "decoder_transformer_layers",
+            ]:
                 return False
             return id(getattr(self, attr)) not in self._functional_layer_ids
 
         return filter(filter_fn, super().__dir__())
 
     def __setattr__(self, name, value):
-        # Work around torch setattr for properties.
-        if name in ["token_embedding"]:
+        # Work around setattr issues for Keras 2 and Keras 3 torch backend.
+        # Since all our state is covered by functional model we can route
+        # around custom setattr calls.
+        is_property = isinstance(getattr(type(self), name, None), property)
+        is_unitialized = not hasattr(self, "_initialized")
+        is_torch = config.backend() == "torch"
+        is_keras_2 = not config.keras_3()
+        if is_torch and (is_property or is_unitialized):
+            return object.__setattr__(self, name, value)
+        if is_keras_2 and is_unitialized:
             return object.__setattr__(self, name, value)
         return super().__setattr__(name, value)
 
@@ -48,18 +65,13 @@ class Backbone(keras.Model):
     def token_embedding(self):
         """A `keras.layers.Embedding` instance for embedding token ids.
 
-        This layer integer token ids to the hidden dim of the model.
+        This layer embeds integer token ids to the hidden dim of the model.
         """
         return self._token_embedding
 
     @token_embedding.setter
     def token_embedding(self, value):
-        # Workaround tf.keras h5 checkpoint loading, which is sensitive to layer
-        # count mismatches and does not deduplicate layers. This could go away
-        # if we update our checkpoints to the newer `.weights.h5` format.
-        self._setattr_tracking = False
         self._token_embedding = value
-        self._setattr_tracking = True
 
     def get_config(self):
         # Don't chain to super here. The default `get_config()` for functional
diff --git a/keras_nlp/models/bart/bart_backbone.py b/keras_nlp/models/bart/bart_backbone.py
index 2679b84a..fdb8e5df 100644
--- a/keras_nlp/models/bart/bart_backbone.py
+++ b/keras_nlp/models/bart/bart_backbone.py
@@ -102,59 +102,34 @@ class BartBackbone(Backbone):
         max_sequence_length=1024,
         **kwargs,
     ):
-        # Encoder inputs
-        encoder_token_id_input = keras.Input(
-            shape=(None,), dtype="int32", name="encoder_token_ids"
-        )
-        encoder_padding_mask = keras.Input(
-            shape=(None,), dtype="int32", name="encoder_padding_mask"
-        )
-
-        # Decoder inputs.
-        decoder_token_id_input = keras.Input(
-            shape=(None,), dtype="int32", name="decoder_token_ids"
-        )
-        decoder_padding_mask = keras.Input(
-            shape=(None,), dtype="int32", name="decoder_padding_mask"
-        )
-
-        # Token embedding layer. This layer is shared by encoder and decoder.
-        token_embedding_layer = ReversibleEmbedding(
+        # === Layers ===
+        self.token_embedding = ReversibleEmbedding(
             input_dim=vocabulary_size,
             output_dim=hidden_dim,
             embeddings_initializer=bart_kernel_initializer(),
             name="token_embedding",
         )
-
-        # ===== Encoder =====
-
-        # Embed tokens and positions.
-        token_embedding = token_embedding_layer(encoder_token_id_input)
-        # Position embedding parameters are not shared by encode and decoder.
-        position_embedding = PositionEmbedding(
+        self.encoder_position_embedding = PositionEmbedding(
             initializer=bart_kernel_initializer(),
             sequence_length=max_sequence_length,
             name="encoder_position_embedding",
-        )(token_embedding)
-
-        # Sum, normalize and apply dropout to embeddings.
-        x = keras.layers.Add(name="encoder_embeddings_add")(
-            (token_embedding, position_embedding)
         )
-        x = keras.layers.LayerNormalization(
+        self.encoder_embeddings_add = keras.layers.Add(
+            name="encoder_embeddings_add",
+        )
+        self.encoder_embeddings_layer_norm = keras.layers.LayerNormalization(
             name="encoder_embeddings_layer_norm",
             axis=-1,
             epsilon=1e-5,
             dtype="float32",
-        )(x)
-        x = keras.layers.Dropout(
+        )
+        self.encoder_embeddings_dropout = keras.layers.Dropout(
             dropout,
             name="encoder_embeddings_dropout",
-        )(x)
-
-        # Apply successive transformer encoder blocks.
+        )
+        self.encoder_transformer_layers = []
         for i in range(num_layers):
-            x = TransformerEncoder(
+            layer = TransformerEncoder(
                 num_heads=num_heads,
                 intermediate_dim=intermediate_dim,
                 activation=keras.activations.gelu,
@@ -162,39 +137,29 @@ class BartBackbone(Backbone):
                 layer_norm_epsilon=1e-5,
                 kernel_initializer=bart_kernel_initializer(),
                 name=f"transformer_encoder_layer_{i}",
-            )(x, padding_mask=encoder_padding_mask)
-
-        encoder_output = x
-
-        # ===== Decoder =====
-
-        # Embed tokens and positions.
-        token_embedding = token_embedding_layer(decoder_token_id_input)
-        # Position embedding parameters are not shared by encode and decoder.
-        position_embedding = PositionEmbedding(
+            )
+            self.encoder_transformer_layers.append(layer)
+        self.decoder_position_embedding = PositionEmbedding(
             initializer=bart_kernel_initializer(),
             sequence_length=max_sequence_length,
             name="decoder_position_embedding",
-        )(token_embedding)
-
-        # Sum, normalize and apply dropout to embeddings.
-        x = keras.layers.Add(name="decoder_embeddings_add")(
-            (token_embedding, position_embedding)
         )
-        x = keras.layers.LayerNormalization(
+        self.decoder_embeddings_add = keras.layers.Add(
+            name="decoder_embeddings_add",
+        )
+        self.decoder_embeddings_layer_norm = keras.layers.LayerNormalization(
             name="decoder_embeddings_layer_norm",
             axis=-1,
             epsilon=1e-5,
             dtype="float32",
-        )(x)
-        x = keras.layers.Dropout(
+        )
+        self.decoder_embeddings_dropout = keras.layers.Dropout(
             dropout,
             name="decoder_embeddings_dropout",
-        )(x)
-
-        # Apply successive transformer decoder blocks.
+        )
+        self.decoder_transformer_layers = []
         for i in range(num_layers):
-            transformer_decoder_layer = TransformerDecoder(
+            layer = TransformerDecoder(
                 intermediate_dim=intermediate_dim,
                 num_heads=num_heads,
                 dropout=dropout,
@@ -203,22 +168,51 @@ class BartBackbone(Backbone):
                 kernel_initializer=bart_kernel_initializer(),
                 name=f"transformer_decoder_layer_{i}",
             )
-            x = transformer_decoder_layer(
+            self.decoder_transformer_layers.append(layer)
+
+        # === Functional Model ===
+        encoder_token_id_input = keras.Input(
+            shape=(None,), dtype="int32", name="encoder_token_ids"
+        )
+        encoder_padding_mask_input = keras.Input(
+            shape=(None,), dtype="int32", name="encoder_padding_mask"
+        )
+        decoder_token_id_input = keras.Input(
+            shape=(None,), dtype="int32", name="decoder_token_ids"
+        )
+        decoder_padding_mask_input = keras.Input(
+            shape=(None,), dtype="int32", name="decoder_padding_mask"
+        )
+        # Encoder.
+        tokens = self.token_embedding(encoder_token_id_input)
+        positions = self.encoder_position_embedding(tokens)
+        x = self.encoder_embeddings_add((tokens, positions))
+        x = self.encoder_embeddings_layer_norm(x)
+        x = self.encoder_embeddings_dropout(x)
+        for transformer_layer in self.encoder_transformer_layers:
+            x = transformer_layer(x, padding_mask=encoder_padding_mask_input)
+        encoder_output = x
+        # Decoder.
+        tokens = self.token_embedding(decoder_token_id_input)
+        positions = self.decoder_position_embedding(tokens)
+        x = self.decoder_embeddings_add((tokens, positions))
+        x = self.decoder_embeddings_layer_norm(x)
+        x = self.decoder_embeddings_dropout(x)
+        for transformer_layer in self.decoder_transformer_layers:
+            x = transformer_layer(
                 decoder_sequence=x,
                 encoder_sequence=encoder_output,
-                decoder_padding_mask=decoder_padding_mask,
-                encoder_padding_mask=encoder_padding_mask,
+                decoder_padding_mask=decoder_padding_mask_input,
+                encoder_padding_mask=encoder_padding_mask_input,
             )
-
         decoder_output = x
-
         # Instantiate using Functional API Model constructor
         super().__init__(
             inputs={
                 "encoder_token_ids": encoder_token_id_input,
-                "encoder_padding_mask": encoder_padding_mask,
+                "encoder_padding_mask": encoder_padding_mask_input,
                 "decoder_token_ids": decoder_token_id_input,
-                "decoder_padding_mask": decoder_padding_mask,
+                "decoder_padding_mask": decoder_padding_mask_input,
             },
             outputs={
                 "encoder_sequence_output": encoder_output,
@@ -227,7 +221,7 @@ class BartBackbone(Backbone):
             **kwargs,
         )
 
-        # All references to `self` below this line
+        # === Config ===
         self.vocabulary_size = vocabulary_size
         self.num_layers = num_layers
         self.num_heads = num_heads
@@ -235,7 +229,6 @@ class BartBackbone(Backbone):
         self.intermediate_dim = intermediate_dim
         self.dropout = dropout
         self.max_sequence_length = max_sequence_length
-        self.token_embedding = token_embedding_layer
 
     def get_config(self):
         config = super().get_config()
diff --git a/keras_nlp/models/bart/bart_seq_2_seq_lm.py b/keras_nlp/models/bart/bart_seq_2_seq_lm.py
index 2131519c..c17eafdb 100644
--- a/keras_nlp/models/bart/bart_seq_2_seq_lm.py
+++ b/keras_nlp/models/bart/bart_seq_2_seq_lm.py
@@ -185,24 +185,21 @@ class BartSeq2SeqLM(GenerativeTask):
         preprocessor=None,
         **kwargs,
     ):
+        # === Layers ===
+        self.backbone = backbone
+        self.preprocessor = preprocessor
+
+        # === Functional Model ===
         inputs = backbone.input
         hidden_states = backbone(inputs)["decoder_sequence_output"]
         outputs = backbone.token_embedding(hidden_states, reverse=True)
-
-        # Instantiate using Functional API Model constructor.
         super().__init__(
             inputs=inputs,
             outputs=outputs,
-            include_preprocessing=preprocessor is not None,
             **kwargs,
         )
 
-        self.backbone = backbone
-        self.preprocessor = preprocessor
-        self.generate_function = None
-        self._sampler = None
-
-        # Default compilation
+        # === Default compilation ===
         self.compile(
             loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
             optimizer=keras.optimizers.Adam(2e-5),
@@ -280,33 +277,28 @@ class BartSeq2SeqLM(GenerativeTask):
             cross-attention layer.
         """
         # Embedding layers.
-        token_embedding = self.backbone.get_layer("token_embedding")(
-            decoder_token_ids
+        tokens = self.backbone.token_embedding(decoder_token_ids)
+        positions = self.backbone.decoder_position_embedding(
+            tokens,
+            start_index=self_attention_cache_update_index,
         )
-        position_embedding = self.backbone.get_layer(
-            "decoder_position_embedding"
-        )(token_embedding, start_index=self_attention_cache_update_index)
-
         # Sum, normalize and apply dropout to embeddings.
-        x = self.backbone.get_layer("decoder_embeddings_add")(
-            (token_embedding, position_embedding)
-        )
-        x = self.backbone.get_layer("decoder_embeddings_layer_norm")(x)
-        x = self.backbone.get_layer("decoder_embeddings_dropout")(x)
+        x = self.backbone.decoder_embeddings_add((tokens, positions))
+        x = self.backbone.decoder_embeddings_layer_norm(x)
+        x = self.backbone.decoder_embeddings_dropout(x)
 
         # Every decoder layer has a separate cache for the self-attention layer
         # and the cross-attention layer. We update all of them separately.
         self_attention_caches = []
         cross_attention_caches = []
-        for i in range(self.backbone.num_layers):
+        for i, layer in enumerate(self.backbone.decoder_transformer_layers):
             current_self_attention_cache = self_attention_cache[:, i, ...]
             current_cross_attention_cache = cross_attention_cache[:, i, ...]
-
             (
                 x,
                 next_self_attention_cache,
                 next_cross_attention_cache,
-            ) = self.backbone.get_layer(f"transformer_decoder_layer_{i}")(
+            ) = layer(
                 decoder_sequence=x,
                 encoder_sequence=encoder_hidden_states,
                 encoder_padding_mask=encoder_padding_mask,
@@ -315,7 +307,6 @@ class BartSeq2SeqLM(GenerativeTask):
                 cross_attention_cache=current_cross_attention_cache,
                 cross_attention_cache_update_index=cross_attention_cache_update_index,
             )
-
             if self_attention_cache_update_index is not None:
                 self_attention_caches.append(next_self_attention_cache)
             if cross_attention_cache_update_index is not None:
@@ -337,26 +328,13 @@ class BartSeq2SeqLM(GenerativeTask):
 
     def call_encoder(self, token_ids, padding_mask):
         """Does a forward pass on the encoder and returns the encoder output."""
-
-        # Embedding layers.
-        token_embedding = self.backbone.get_layer("token_embedding")(token_ids)
-        position_embedding = self.backbone.get_layer(
-            "encoder_position_embedding"
-        )(token_embedding)
-
-        # Sum, normalize and apply dropout to embeddings.
-        x = self.backbone.get_layer("encoder_embeddings_add")(
-            (token_embedding, position_embedding)
-        )
-        x = self.backbone.get_layer("encoder_embeddings_layer_norm")(x)
-        x = self.backbone.get_layer("encoder_embeddings_dropout")(x)
-
-        # Transformer encoder layers.
-        for i in range(self.backbone.num_layers):
-            x = self.backbone.get_layer(f"transformer_encoder_layer_{i}")(
-                x, padding_mask=padding_mask
-            )
-
+        tokens = self.backbone.token_embedding(token_ids)
+        positions = self.backbone.encoder_position_embedding(tokens)
+        x = self.backbone.decoder_embeddings_add((tokens, positions))
+        x = self.backbone.encoder_embeddings_layer_norm(x)
+        x = self.backbone.encoder_embeddings_dropout(x)
+        for transformer_layer in self.backbone.encoder_transformer_layers:
+            x = transformer_layer(x, padding_mask=padding_mask)
         return x
 
     def _initialize_cache(self, encoder_token_ids, decoder_token_ids):
diff --git a/keras_nlp/models/bert/bert_backbone.py b/keras_nlp/models/bert/bert_backbone.py
index 174b0f0e..f511de36 100644
--- a/keras_nlp/models/bert/bert_backbone.py
+++ b/keras_nlp/models/bert/bert_backbone.py
@@ -99,57 +99,40 @@ class BertBackbone(Backbone):
         num_segments=2,
         **kwargs,
     ):
-        # Index of classification token in the vocabulary
-        cls_token_index = 0
-        # Inputs
-        token_id_input = keras.Input(
-            shape=(None,), dtype="int32", name="token_ids"
-        )
-        segment_id_input = keras.Input(
-            shape=(None,), dtype="int32", name="segment_ids"
-        )
-        padding_mask = keras.Input(
-            shape=(None,), dtype="int32", name="padding_mask"
-        )
-
-        # Embed tokens, positions, and segment ids.
-        token_embedding_layer = ReversibleEmbedding(
+        # === Layers ===
+        self.token_embedding = ReversibleEmbedding(
             input_dim=vocabulary_size,
             output_dim=hidden_dim,
             embeddings_initializer=bert_kernel_initializer(),
             name="token_embedding",
         )
-        token_embedding = token_embedding_layer(token_id_input)
-        position_embedding = PositionEmbedding(
+        self.position_embedding = PositionEmbedding(
             initializer=bert_kernel_initializer(),
             sequence_length=max_sequence_length,
             name="position_embedding",
-        )(token_embedding)
-        segment_embedding = keras.layers.Embedding(
+        )
+        self.segment_embedding = keras.layers.Embedding(
             input_dim=num_segments,
             output_dim=hidden_dim,
             embeddings_initializer=bert_kernel_initializer(),
             name="segment_embedding",
-        )(segment_id_input)
-
-        # Sum, normalize and apply dropout to embeddings.
-        x = keras.layers.Add()(
-            (token_embedding, position_embedding, segment_embedding)
         )
-        x = keras.layers.LayerNormalization(
+        self.embeddings_add = keras.layers.Add(
+            name="embeddings_add",
+        )
+        self.embeddings_layer_norm = keras.layers.LayerNormalization(
             name="embeddings_layer_norm",
             axis=-1,
             epsilon=1e-12,
             dtype="float32",
-        )(x)
-        x = keras.layers.Dropout(
+        )
+        self.embeddings_dropout = keras.layers.Dropout(
             dropout,
             name="embeddings_dropout",
-        )(x)
-
-        # Apply successive transformer encoder blocks.
+        )
+        self.transformer_layers = []
         for i in range(num_layers):
-            x = TransformerEncoder(
+            layer = TransformerEncoder(
                 num_heads=num_heads,
                 intermediate_dim=intermediate_dim,
                 activation=gelu_approximate,
@@ -157,24 +140,45 @@ class BertBackbone(Backbone):
                 layer_norm_epsilon=1e-12,
                 kernel_initializer=bert_kernel_initializer(),
                 name=f"transformer_layer_{i}",
-            )(x, padding_mask=padding_mask)
-
-        # Construct the two BERT outputs. The pooled output is a dense layer on
-        # top of the [CLS] token.
-        sequence_output = x
-        pooled_output = keras.layers.Dense(
+            )
+            self.transformer_layers.append(layer)
+        self.pooled_dense = keras.layers.Dense(
             hidden_dim,
             kernel_initializer=bert_kernel_initializer(),
             activation="tanh",
             name="pooled_dense",
-        )(x[:, cls_token_index, :])
+        )
 
-        # Instantiate using Functional API Model constructor
+        # === Functional Model ===
+        token_id_input = keras.Input(
+            shape=(None,), dtype="int32", name="token_ids"
+        )
+        segment_id_input = keras.Input(
+            shape=(None,), dtype="int32", name="segment_ids"
+        )
+        padding_mask_input = keras.Input(
+            shape=(None,), dtype="int32", name="padding_mask"
+        )
+        # Embed tokens, positions, and segment ids.
+        tokens = self.token_embedding(token_id_input)
+        positions = self.position_embedding(tokens)
+        segments = self.segment_embedding(segment_id_input)
+        # Sum, normalize and apply dropout to embeddings.
+        x = self.embeddings_add((tokens, positions, segments))
+        x = self.embeddings_layer_norm(x)
+        x = self.embeddings_dropout(x)
+        for transformer_layer in self.transformer_layers:
+            x = transformer_layer(x, padding_mask=padding_mask_input)
+        # Construct the two BERT outputs. The pooled output is a dense layer on
+        # top of the [CLS] token.
+        sequence_output = x
+        cls_token_index = 0
+        pooled_output = self.pooled_dense(x[:, cls_token_index, :])
         super().__init__(
             inputs={
                 "token_ids": token_id_input,
                 "segment_ids": segment_id_input,
-                "padding_mask": padding_mask,
+                "padding_mask": padding_mask_input,
             },
             outputs={
                 "sequence_output": sequence_output,
@@ -183,7 +187,7 @@ class BertBackbone(Backbone):
             **kwargs,
         )
 
-        # All references to `self` below this line
+        # === Config ===
         self.vocabulary_size = vocabulary_size
         self.num_layers = num_layers
         self.num_heads = num_heads
@@ -193,7 +197,6 @@ class BertBackbone(Backbone):
         self.max_sequence_length = max_sequence_length
         self.num_segments = num_segments
         self.cls_token_index = cls_token_index
-        self.token_embedding = token_embedding_layer
 
     def get_config(self):
         config = super().get_config()
diff --git a/keras_nlp/models/bert/bert_classifier.py b/keras_nlp/models/bert/bert_classifier.py
index 2a9aa548..3ddb90a9 100644
--- a/keras_nlp/models/bert/bert_classifier.py
+++ b/keras_nlp/models/bert/bert_classifier.py
@@ -140,30 +140,37 @@ class BertClassifier(Task):
         dropout=0.1,
         **kwargs,
     ):
-        inputs = backbone.input
-        pooled = backbone(inputs)["pooled_output"]
-        pooled = keras.layers.Dropout(dropout)(pooled)
-        outputs = keras.layers.Dense(
+        # === Layers ===
+        self.backbone = backbone
+        self.preprocessor = preprocessor
+        self.output_dropout = keras.layers.Dropout(
+            dropout,
+            name="classifier_dropout",
+        )
+        self.output_dense = keras.layers.Dense(
             num_classes,
             kernel_initializer=bert_kernel_initializer(),
             activation=activation,
             name="logits",
-        )(pooled)
-        # Instantiate using Functional API Model constructor
+        )
+
+        # === Functional Model ===
+        inputs = backbone.input
+        pooled = backbone(inputs)["pooled_output"]
+        pooled = self.output_dropout(pooled)
+        outputs = self.output_dense(pooled)
         super().__init__(
             inputs=inputs,
             outputs=outputs,
-            include_preprocessing=preprocessor is not None,
             **kwargs,
         )
-        # All references to `self` below this line
-        self.backbone = backbone
-        self.preprocessor = preprocessor
+
+        # === Config ===
         self.num_classes = num_classes
         self.activation = keras.activations.get(activation)
         self.dropout = dropout
 
-        # Default compilation
+        # === Default compilation ===
         logit_output = self.activation == keras.activations.linear
         self.compile(
             loss=keras.losses.SparseCategoricalCrossentropy(
diff --git a/keras_nlp/models/bert/bert_masked_lm.py b/keras_nlp/models/bert/bert_masked_lm.py
index d4c12d10..555c562f 100644
--- a/keras_nlp/models/bert/bert_masked_lm.py
+++ b/keras_nlp/models/bert/bert_masked_lm.py
@@ -101,6 +101,18 @@ class BertMaskedLM(Task):
         preprocessor=None,
         **kwargs,
     ):
+        # === Layers ===
+        self.backbone = backbone
+        self.preprocessor = preprocessor
+        self.masked_lm_head = MaskedLMHead(
+            vocabulary_size=backbone.vocabulary_size,
+            token_embedding=backbone.token_embedding,
+            intermediate_activation="gelu",
+            kernel_initializer=bert_kernel_initializer(),
+            name="mlm_head",
+        )
+
+        # === Functional Model ===
         inputs = {
             **backbone.input,
             "mask_positions": keras.Input(
@@ -108,22 +120,16 @@ class BertMaskedLM(Task):
             ),
         }
         backbone_outputs = backbone(backbone.input)
-        outputs = MaskedLMHead(
-            vocabulary_size=backbone.vocabulary_size,
-            token_embedding=backbone.token_embedding,
-            intermediate_activation="gelu",
-            kernel_initializer=bert_kernel_initializer(),
-            name="mlm_head",
-        )(backbone_outputs["sequence_output"], inputs["mask_positions"])
-
-        # Instantiate using Functional API Model constructor
+        outputs = self.masked_lm_head(
+            backbone_outputs["sequence_output"], inputs["mask_positions"]
+        )
         super().__init__(
             inputs=inputs,
             outputs=outputs,
-            include_preprocessing=preprocessor is not None,
             **kwargs,
         )
-        # All references to `self` below this line
+
+        # === Default compilation ===
         self.backbone = backbone
         self.preprocessor = preprocessor
         self.compile(
diff --git a/keras_nlp/models/bloom/bloom_backbone.py b/keras_nlp/models/bloom/bloom_backbone.py
index 2fd18ed7..4f12bacd 100644
--- a/keras_nlp/models/bloom/bloom_backbone.py
+++ b/keras_nlp/models/bloom/bloom_backbone.py
@@ -28,7 +28,7 @@ def _bloom_kernel_initializer(stddev=0.02):
 
 @keras_nlp_export("keras_nlp.models.BloomBackbone")
 class BloomBackbone(Backbone):
-    """A Bloom decoder network.
+    """A BLOOM decoder network.
 
     This network implements a Transformer-based decoder network, BigScience
     Language Open-science Open-access Multilingual (BLOOM), as descriped in
@@ -92,49 +92,58 @@ class BloomBackbone(Backbone):
         intermediate_dim,
         dropout=0.0,
         layer_norm_epsilon=1e-5,
-        max_sequence_length=512,
+        max_sequence_length=2048,
         **kwargs,
     ):
-        token_ids = keras.Input(shape=(None,), dtype="int32", name="token_ids")
-        padding_mask = keras.Input(
-            shape=(None,), dtype="int32", name="padding_mask"
-        )
-
-        # Embed tokens
-        token_embedding_layer = ReversibleEmbedding(
+        # === Layers ===
+        self.token_embedding = ReversibleEmbedding(
             input_dim=vocabulary_size,
             output_dim=hidden_dim,
             embeddings_initializer=_bloom_kernel_initializer(stddev=0.02),
             tie_weights=False,
             name="token_embedding",
         )
-        token_embedding = token_embedding_layer(token_ids)
-
-        x = keras.layers.LayerNormalization(
-            epsilon=layer_norm_epsilon, name="token_embedding_layernorm"
-        )(token_embedding)
-
+        self.embeddings_layer_norm = keras.layers.LayerNormalization(
+            epsilon=layer_norm_epsilon,
+            name="token_embedding_layernorm",
+        )
+        self.transformer_layers = []
         for i in range(num_layers):
-            x = BloomDecoder(
+            layer = BloomDecoder(
                 num_heads=num_heads,
                 intermediate_dim=intermediate_dim,
                 dropout=dropout,
                 layer_norm_epsilon=layer_norm_epsilon,
                 name=f"transformer_layer_{i}",
-            )(x, decoder_padding_mask=padding_mask)
-
-        sequence_output = keras.layers.LayerNormalization(
-            epsilon=layer_norm_epsilon, name="final_layernorm"
-        )(x)
+            )
+            self.transformer_layers.append(layer)
+        self.layer_norm = keras.layers.LayerNormalization(
+            epsilon=layer_norm_epsilon,
+            name="final_layernorm",
+        )
 
+        # === Functional Model ===
+        token_id_input = keras.Input(
+            shape=(None,), dtype="int32", name="token_ids"
+        )
+        padding_mask_input = keras.Input(
+            shape=(None,), dtype="int32", name="padding_mask"
+        )
+        x = self.token_embedding(token_id_input)
+        x = self.embeddings_layer_norm(x)
+        for transformer_layer in self.transformer_layers:
+            x = transformer_layer(x, decoder_padding_mask=padding_mask_input)
+        sequence_output = self.layer_norm(x)
         super().__init__(
             inputs={
-                "token_ids": token_ids,
-                "padding_mask": padding_mask,
+                "token_ids": token_id_input,
+                "padding_mask": padding_mask_input,
             },
             outputs=sequence_output,
             **kwargs,
         )
+
+        # === Config ===
         self.vocabulary_size = vocabulary_size
         self.num_layers = num_layers
         self.num_heads = num_heads
@@ -143,7 +152,6 @@ class BloomBackbone(Backbone):
         self.dropout = dropout
         self.layer_norm_epsilon = layer_norm_epsilon
         self.max_sequence_length = max_sequence_length
-        self.token_embedding = token_embedding_layer
 
     def get_config(self):
         config = super().get_config()
diff --git a/keras_nlp/models/bloom/bloom_causal_lm_preprocessor.py b/keras_nlp/models/bloom/bloom_causal_lm_preprocessor.py
new file mode 100644
index 00000000..01f3c88d
--- /dev/null
+++ b/keras_nlp/models/bloom/bloom_causal_lm_preprocessor.py
@@ -0,0 +1,182 @@
+# Copyright 2023 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 absl import logging
+
+from keras_nlp.api_export import keras_nlp_export
+from keras_nlp.backend import ops
+from keras_nlp.models.bloom.bloom_preprocessor import BloomPreprocessor
+from keras_nlp.utils.keras_utils import (
+    convert_inputs_to_list_of_tensor_segments,
+)
+from keras_nlp.utils.keras_utils import pack_x_y_sample_weight
+
+
+@keras_nlp_export("keras_nlp.models.BloomCausalLMPreprocessor")
+class BloomCausalLMPreprocessor(BloomPreprocessor):
+    """BLOOM Causal LM preprocessor.
+
+    This preprocessing layer is meant for use with
+    `keras_nlp.models.BloomCausalLM`. By default, it will take in batches of
+    strings, and return outputs in a `(x, y, sample_weight)` format, where the
+    `y` label is the next token id in the `x` sequence.
+
+    For use with generation, the layer also exposes two methods
+    `generate_preprocess()` and `generate_postprocess()`. When this preprocessor
+    is attached to a `keras_nlp.models.BloomCausalLM` instance, these methods
+    will be called implicitly in `generate()`. They can also be called
+    standalone (e.g. to precompute preprocessing inputs for generation in a
+    separate process).
+
+    Args:
+        tokenizer: A `keras_nlp.models.BloomTokenizer` instance.
+        sequence_length: The length of the packed inputs.
+        add_start_token: If `True`, the preprocessor will prepend the tokenizer
+            start token to each input sequence.
+        add_end_token: If `True`, the preprocessor will append the tokenizer
+            end token to each input sequence.
+
+    Call arguments:
+        x: A string, `tf.Tensor` or list of python strings.
+        y: Label data. Should always be `None` as the layer generates labels.
+        sample_weight: Label weights. Should always be `None` as the layer
+            generates label weights.
+        sequence_length: Pass to override the configured `sequence_length` of
+            the layer.
+
+    Examples:
+    ```python
+    # Load the preprocessor from a preset.
+    preprocessor = keras_nlp.models.BloomCausalLMPreprocessor.from_preset(
+        "bloom_560m_multi"
+    )
+
+    # Tokenize and pack a single sentence.
+    sentence = tf.constant("League of legends")
+    preprocessor(sentence)
+    # Same output.
+    preprocessor("League of legends")
+
+    # Tokenize a batch of sentences.
+    sentences = tf.constant(["Taco tuesday", "Fish taco please!"])
+    preprocessor(sentences)
+    # Same output.
+    preprocessor(["Taco tuesday", "Fish taco please!"])
+
+    # Map a dataset to preprocess a single sentence.
+    features = tf.constant(
+        [
+            "Avatar 2 is amazing!",
+            "Well, I am not sure.",
+        ]
+    )
+    labels = tf.constant([1, 0])
+    ds = tf.data.Dataset.from_tensor_slices((features, labels))
+    ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
+
+    # Map a dataset to preprocess unlabled sentences.
+    ds = tf.data.Dataset.from_tensor_slices(features)
+    ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
+    ```
+    """
+
+    def call(
+        self,
+        x,
+        y=None,
+        sample_weight=None,
+        sequence_length=None,
+    ):
+        if y is not None or sample_weight is not None:
+            logging.warning(
+                "`BloomCausalLMPreprocessor` generates `y` and `sample_weight` "
+                "based on your input data, but your data already contains `y` "
+                "or `sample_weight`. Your `y` and `sample_weight` will be "
+                "ignored."
+            )
+        sequence_length = sequence_length or self.sequence_length
+
+        x = convert_inputs_to_list_of_tensor_segments(x)[0]
+        x = self.tokenizer(x)
+        # Pad with one extra token to account for the truncation below.
+        token_ids, padding_mask = self.packer(
+            x,
+            sequence_length=sequence_length + 1,
+            add_start_value=self.add_start_token,
+            add_end_value=self.add_end_token,
+        )
+        # The last token does not have a next token, so we truncate it out.
+        x = {
+            "token_ids": token_ids[..., :-1],
+            "padding_mask": padding_mask[..., :-1],
+        }
+        # Target `y` will be the next token.
+        y, sample_weight = token_ids[..., 1:], padding_mask[..., 1:]
+        return pack_x_y_sample_weight(x, y, sample_weight)
+
+    def generate_preprocess(
+        self,
+        x,
+        sequence_length=None,
+    ):
+        """Covert strings to integer token input for generation.
+
+        Similar to calling the layer for training, this method takes in strings
+        or tensor strings, tokenizes and packs the input, and computes a padding
+        mask masking all inputs not filled in with a padded value.
+
+        Unlike calling the layer for training, this method does not compute
+        labels and will never append a `tokenizer.end_token_id` to the end of
+        the sequence (as generation is expected to continue at the end of the
+        inputted prompt).
+        """
+        if not self.built:
+            self.build(None)
+
+        x = convert_inputs_to_list_of_tensor_segments(x)[0]
+        x = self.tokenizer(x)
+        token_ids, padding_mask = self.packer(
+            x, sequence_length=sequence_length, add_end_value=False
+        )
+        return {
+            "token_ids": token_ids,
+            "padding_mask": padding_mask,
+        }
+
+    def generate_postprocess(
+        self,
+        x,
+    ):
+        """Covert integer token output to strings for generation.
+
+        This method reverses `generate_preprocess()`, by first removing all
+        padding and start/end tokens, and then converting the integer sequence
+        back to a string.
+        """
+        if not self.built:
+            self.build(None)
+
+        token_ids, padding_mask = x["token_ids"], x["padding_mask"]
+        token_ids = ops.convert_to_numpy(token_ids)
+        padding_mask = ops.convert_to_numpy(padding_mask)
+        # Strip any special tokens during detokenization (e.g. the start and
+        # end markers). In the future we could make this configurable.
+        padding_mask = (
+            padding_mask
+            & (token_ids != self.tokenizer.eos_token_id)
+            & (token_ids != self.tokenizer.bos_token_id)
+        )
+        token_ids = tf.ragged.boolean_mask(token_ids, padding_mask)
+        return self.tokenizer.detokenize(token_ids)
diff --git a/keras_nlp/models/bloom/bloom_preprocessor.py b/keras_nlp/models/bloom/bloom_preprocessor.py
new file mode 100644
index 00000000..d57ccd24
--- /dev/null
+++ b/keras_nlp/models/bloom/bloom_preprocessor.py
@@ -0,0 +1,182 @@
+# 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 copy
+
+from keras_nlp.api_export import keras_nlp_export
+from keras_nlp.layers.preprocessing.start_end_packer import StartEndPacker
+from keras_nlp.models.bloom.bloom_presets import backbone_presets
+from keras_nlp.models.bloom.bloom_tokenizer import BloomTokenizer
+from keras_nlp.models.preprocessor import Preprocessor
+from keras_nlp.utils.keras_utils import (
+    convert_inputs_to_list_of_tensor_segments,
+)
+from keras_nlp.utils.keras_utils import pack_x_y_sample_weight
+from keras_nlp.utils.python_utils import classproperty
+
+
+@keras_nlp_export("keras_nlp.models.BloomPreprocessor")
+class BloomPreprocessor(Preprocessor):
+    """BLOOM preprocessing layer which tokenizes and packs inputs.
+
+    This preprocessing layer will do 2 things:
+
+    - Tokenize the inputs using the `tokenizer`.
+    - Construct a dictionary with keys `"token_ids"`, `"padding_mask"`, that can
+        be passed directly to a `keras_nlp.models.BloomBackbone`.
+
+    This layer can be used directly with `tf.data.Dataset.map` to preprocess
+    string data in the `(x, y, sample_weight)` format used by
+    `keras.Model.fit`.
+
+    The call method of this layer accepts three arguments, `x`, `y`, and
+    `sample_weight`. `x` can be a python string or tensor representing a single
+    segment, a list of python strings representing a batch of single segments,
+    or a list of tensors representing multiple segments to be packed together.
+    `y` and `sample_weight` are both optional, can have any format, and will be
+    passed through unaltered.
+
+    Args:
+        tokenizer: A `keras_nlp.models.BloomTokenizer` instance.
+        sequence_length: The length of the packed inputs.
+        add_start_token: If `True`, the preprocessor will prepend the tokenizer
+            start token to each input sequence.
+        add_end_token: If `True`, the preprocessor will append the tokenizer
+            end token to each input sequence.
+
+    Call arguments:
+        x: A string, `tf.Tensor` or list of python strings.
+        y: Any label data. Will be passed through unaltered.
+        sample_weight: Any label weight data. Will be passed through unaltered.
+        sequence_length: Pass to override the configured `sequence_length` of
+            the layer.
+
+    Examples:
+
+    Directly calling the layer on data.
+    ```python
+    preprocessor = keras_nlp.models.BloomPreprocessor.from_preset(
+        "bloom_560m_multi"
+    )
+    # Tokenize and pack a single sentence.
+    preprocessor("The quick brown fox jumped.")
+
+    # Tokenize a batch of single sentences.
+    preprocessor(["The quick brown fox jumped.", "Call me Ishmael."])
+
+    # Custom vocabulary.
+    features = ["a quick fox.", "a fox quick."]
+    vocab = {"<pad>": 0, "
":1, "
":2, "a": 3, "Ġquick": 4, "Ġfox": 5}
+    merges = ["Ġ q", "u i", "c k", "ui ck", "Ġq uick"]
+    merges += ["Ġ f", "o x", "Ġf ox"]
+    tokenizer = keras_nlp.models.BloomTokenizer(
+        vocabulary=vocab,
+        merges=merges,
+    )
+    preprocessor = keras_nlp.models.BloomPreprocessor(tokenizer=tokenizer)
+    preprocessor("The quick brown fox jumped.")
+    ```
+
+    Mapping with `tf.data.Dataset`.
+    ```python
+    preprocessor = keras_nlp.models.BloomPreprocessor.from_preset(
+        "bloom_560m_multi"
+    )
+
+    text = tf.constant(["The quick brown fox jumped.", "Call me Ishmael."])
+    label = tf.constant([1, 1])
+
+    # Map labeled single sentences.
+    ds = tf.data.Dataset.from_tensor_slices((text, label))
+    ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
+
+    # Map unlabeled single sentences.
+    ds = tf.data.Dataset.from_tensor_slices(text)
+    ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
+    ```
+    """
+
+    def __init__(
+        self,
+        tokenizer,
+        sequence_length=2048,
+        add_start_token=True,
+        add_end_token=True,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.tokenizer = tokenizer
+        self.sequence_length = sequence_length
+        self.add_start_token = add_start_token
+        self.add_end_token = add_end_token
+
+    def build(self, input_shape):
+        # Defer packer creation to `build()` so that we can be sure tokenizer
+        # assets have loaded when restoring a saved model.
+        self.packer = StartEndPacker(
+            start_value=self.tokenizer.bos_token_id,
+            end_value=self.tokenizer.eos_token_id,
+            pad_value=self.tokenizer.pad_token_id,
+            sequence_length=self.sequence_length,
+            return_padding_mask=True,
+        )
+        self.built = True
+
+    def call(
+        self,
+        x,
+        y=None,
+        sample_weight=None,
+        sequence_length=None,
+    ):
+        x = convert_inputs_to_list_of_tensor_segments(x)
+        if len(x) != 1:
+            raise ValueError(
+                "BLOOM requires each input feature to contain only "
+                f"one segment, but received {len(x)}. If you are using BLOOM "
+                "for a multi-segment classification task, please refer to "
+                "classification models like BERT or RoBERTa."
+            )
+        sequence_length = sequence_length or self.sequence_length
+        token_ids, padding_mask = self.packer(
+            self.tokenizer(x[0]),
+            sequence_length=sequence_length,
+            add_start_value=self.add_start_token,
+            add_end_value=self.add_end_token,
+        )
+        x = {
+            "token_ids": token_ids,
+            "padding_mask": padding_mask,
+        }
+        return pack_x_y_sample_weight(x, y, sample_weight)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "sequence_length": self.sequence_length,
+                "add_start_token": self.add_start_token,
+                "add_end_token": self.add_end_token,
+            }
+        )
+        return config
+
+    @classproperty
+    def presets(cls):
+        return copy.deepcopy(backbone_presets)
+
+    @classproperty
+    def tokenizer_cls(cls):
+        return BloomTokenizer
diff --git a/keras_nlp/models/deberta_v3/deberta_v3_backbone.py b/keras_nlp/models/deberta_v3/deberta_v3_backbone.py
index aa5077ec..87531aef 100644
--- a/keras_nlp/models/deberta_v3/deberta_v3_backbone.py
+++ b/keras_nlp/models/deberta_v3/deberta_v3_backbone.py
@@ -108,46 +108,32 @@ class DebertaV3Backbone(Backbone):
         bucket_size=256,
         **kwargs,
     ):
-        # Inputs
-        token_id_input = keras.Input(
-            shape=(None,), dtype="int32", name="token_ids"
-        )
-        padding_mask = keras.Input(
-            shape=(None,), dtype="int32", name="padding_mask"
-        )
-
-        # Embed tokens.
-        token_embedding_layer = ReversibleEmbedding(
+        # === Layers ===
+        self.token_embedding = ReversibleEmbedding(
             input_dim=vocabulary_size,
             output_dim=hidden_dim,
             embeddings_initializer=deberta_kernel_initializer(),
             name="token_embedding",
         )
-        x = token_embedding_layer(token_id_input)
-
-        # Normalize and apply dropout to embeddings.
-        x = keras.layers.LayerNormalization(
+        self.embeddings_layer_norm = keras.layers.LayerNormalization(
             epsilon=1e-7,
             dtype="float32",
             name="embeddings_layer_norm",
-        )(x)
-        x = keras.layers.Dropout(
+        )
+        self.embeddings_dropout = keras.layers.Dropout(
             dropout,
             name="embeddings_dropout",
-        )(x)
-
-        # Relative embedding layer.
-        rel_embeddings = RelativeEmbedding(
+        )
+        self.relative_embeddings = RelativeEmbedding(
             hidden_dim=hidden_dim,
             bucket_size=bucket_size,
             layer_norm_epsilon=1e-7,
             kernel_initializer=deberta_kernel_initializer(),
             name="rel_embedding",
-        )(x)
-
-        # Apply successive DeBERTa encoder blocks.
+        )
+        self.transformer_layers = []
         for i in range(num_layers):
-            x = DisentangledAttentionEncoder(
+            layer = DisentangledAttentionEncoder(
                 num_heads=num_heads,
                 intermediate_dim=intermediate_dim,
                 max_position_embeddings=max_sequence_length,
@@ -157,22 +143,36 @@ class DebertaV3Backbone(Backbone):
                 layer_norm_epsilon=1e-7,
                 kernel_initializer=deberta_kernel_initializer(),
                 name=f"disentangled_attention_encoder_layer_{i}",
-            )(
+            )
+            self.transformer_layers.append(layer)
+
+        # === Functional Model ===
+        token_id_input = keras.Input(
+            shape=(None,), dtype="int32", name="token_ids"
+        )
+        padding_mask_input = keras.Input(
+            shape=(None,), dtype="int32", name="padding_mask"
+        )
+        x = self.token_embedding(token_id_input)
+        x = self.embeddings_layer_norm(x)
+        x = self.embeddings_dropout(x)
+        rel_embeddings = self.relative_embeddings(x)
+        for transformer_layer in self.transformer_layers:
+            x = transformer_layer(
                 x,
                 rel_embeddings=rel_embeddings,
-                padding_mask=padding_mask,
+                padding_mask=padding_mask_input,
             )
-
-        # Instantiate using Functional API Model constructor
         super().__init__(
             inputs={
                 "token_ids": token_id_input,
-                "padding_mask": padding_mask,
+                "padding_mask": padding_mask_input,
             },
             outputs=x,
             **kwargs,
         )
-        # All references to `self` below this line
+
+        # === Config ===
         self.vocabulary_size = vocabulary_size
         self.num_layers = num_layers
         self.num_heads = num_heads
@@ -182,7 +182,6 @@ class DebertaV3Backbone(Backbone):
         self.max_sequence_length = max_sequence_length
         self.bucket_size = bucket_size
         self.start_token_index = 0
-        self.token_embedding = token_embedding_layer
 
     def get_config(self):
         config = super().get_config()
diff --git a/keras_nlp/models/deberta_v3/deberta_v3_classifier.py b/keras_nlp/models/deberta_v3/deberta_v3_classifier.py
index b0312206..f5249cb3 100644
--- a/keras_nlp/models/deberta_v3/deberta_v3_classifier.py
+++ b/keras_nlp/models/deberta_v3/deberta_v3_classifier.py
@@ -163,32 +163,44 @@ class DebertaV3Classifier(Task):
         dropout=0.0,
         **kwargs,
     ):
-        inputs = backbone.input
+        # === Layers ===
+        self.backbone = backbone
+        self.preprocessor = preprocessor
+        self.pooled_dropout = keras.layers.Dropout(
+            dropout,
+            name="pooled_dropout",
+        )
         hidden_dim = hidden_dim or backbone.hidden_dim
-
-        x = backbone(inputs)[:, backbone.start_token_index, :]
-        x = keras.layers.Dropout(dropout, name="pooled_dropout")(x)
-        x = keras.layers.Dense(
+        self.pooled_dense = keras.layers.Dense(
             hidden_dim,
             activation=keras.activations.gelu,
             name="pooled_dense",
-        )(x)
-        x = keras.layers.Dropout(backbone.dropout, name="classifier_dropout")(x)
-        outputs = keras.layers.Dense(
+        )
+        self.output_dropout = keras.layers.Dropout(
+            backbone.dropout,
+            name="classifier_dropout",
+        )
+        self.output_dense = keras.layers.Dense(
             num_classes,
             kernel_initializer=deberta_kernel_initializer(),
             activation=activation,
             name="logits",
-        )(x)
+        )
 
-        # Instantiate using Functional API Model constructor
+        # === Functional Model ===
+        inputs = backbone.input
+        x = backbone(inputs)[:, backbone.start_token_index, :]
+        x = self.pooled_dropout(x)
+        x = self.pooled_dense(x)
+        x = self.output_dropout(x)
+        outputs = self.output_dense(x)
         super().__init__(
             inputs=inputs,
             outputs=outputs,
-            include_preprocessing=preprocessor is not None,
             **kwargs,
         )
-        # All references to `self` below this line
+
+        # === Config ===
         self.backbone = backbone
         self.preprocessor = preprocessor
         self.num_classes = num_classes
@@ -196,7 +208,7 @@ class DebertaV3Classifier(Task):
         self.hidden_dim = hidden_dim
         self.dropout = dropout
 
-        # Default compilation
+        # === Default compilation ===
         logit_output = self.activation == keras.activations.linear
         self.compile(
             loss=keras.losses.SparseCategoricalCrossentropy(
diff --git a/keras_nlp/models/deberta_v3/deberta_v3_masked_lm.py b/keras_nlp/models/deberta_v3/deberta_v3_masked_lm.py
index bf6a850a..fadb4c0e 100644
--- a/keras_nlp/models/deberta_v3/deberta_v3_masked_lm.py
+++ b/keras_nlp/models/deberta_v3/deberta_v3_masked_lm.py
@@ -104,32 +104,33 @@ class DebertaV3MaskedLM(Task):
         preprocessor=None,
         **kwargs,
     ):
-        inputs = {
-            **backbone.input,
-            "mask_positions": keras.Input(
-                shape=(None,), dtype="int32", name="mask_positions"
-            ),
-        }
-        backbone_outputs = backbone(backbone.input)
-        outputs = MaskedLMHead(
+        # === Layers ===
+        self.backbone = backbone
+        self.preprocessor = preprocessor
+        self.masked_lm_head = MaskedLMHead(
             vocabulary_size=backbone.vocabulary_size,
             token_embedding=backbone.token_embedding,
             intermediate_activation=keras.activations.gelu,
             kernel_initializer=deberta_kernel_initializer(),
             name="mlm_head",
-        )(backbone_outputs, inputs["mask_positions"])
+        )
 
-        # Instantiate using Functional API Model constructor
+        # === Functional Model ===
+        inputs = {
+            **backbone.input,
+            "mask_positions": keras.Input(
+                shape=(None,), dtype="int32", name="mask_positions"
+            ),
+        }
+        x = backbone(backbone.input)
+        outputs = self.masked_lm_head(x, inputs["mask_positions"])
         super().__init__(
             inputs=inputs,
             outputs=outputs,
-            include_preprocessing=preprocessor is not None,
             **kwargs,
         )
-        # All references to `self` below this line
-        self.backbone = backbone
-        self.preprocessor = preprocessor
 
+        # === Default compilation ===
         self.compile(
             loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
             optimizer=keras.optimizers.Adam(5e-5),
diff --git a/keras_nlp/models/distil_bert/distil_bert_backbone.py b/keras_nlp/models/distil_bert/distil_bert_backbone.py
index a3634215..f97e24f5 100644
--- a/keras_nlp/models/distil_bert/distil_bert_backbone.py
+++ b/keras_nlp/models/distil_bert/distil_bert_backbone.py
@@ -100,39 +100,29 @@ class DistilBertBackbone(Backbone):
         max_sequence_length=512,
         **kwargs,
     ):
-        # Inputs
-        token_id_input = keras.Input(
-            shape=(None,), dtype="int32", name="token_ids"
-        )
-        padding_mask = keras.Input(
-            shape=(None,), dtype="int32", name="padding_mask"
-        )
-
-        # Embed tokens and positions.
-        embedding_layer = TokenAndPositionEmbedding(
+        # === Layers ===
+        self.embeddings = TokenAndPositionEmbedding(
             vocabulary_size=vocabulary_size,
             sequence_length=max_sequence_length,
             embedding_dim=hidden_dim,
             embeddings_initializer=distilbert_kernel_initializer(),
             name="token_and_position_embedding",
         )
-        x = embedding_layer(token_id_input)
-
-        # Normalize and apply dropout to embeddings.
-        x = keras.layers.LayerNormalization(
+        # Keep the token_embedding property for consistency across models.
+        self.token_embedding = self.embeddings.token_embedding
+        self.embeddings_layer_norm = keras.layers.LayerNormalization(
             axis=-1,
             epsilon=1e-12,
             dtype="float32",
             name="embeddings_layer_norm",
-        )(x)
-        x = keras.layers.Dropout(
+        )
+        self.embeddings_dropout = keras.layers.Dropout(
             dropout,
             name="embeddings_dropout",
-        )(x)
-
-        # Apply successive transformer encoder blocks.
+        )
+        self.transformer_layers = []
         for i in range(num_layers):
-            x = TransformerEncoder(
+            layer = TransformerEncoder(
                 num_heads=num_heads,
                 intermediate_dim=intermediate_dim,
                 activation="gelu",
@@ -140,18 +130,31 @@ class DistilBertBackbone(Backbone):
                 layer_norm_epsilon=1e-12,
                 kernel_initializer=distilbert_kernel_initializer(),
                 name=f"transformer_layer_{i}",
-            )(x, padding_mask=padding_mask)
+            )
+            self.transformer_layers.append(layer)
 
-        # Instantiate using Functional API Model constructor
+        # === Functional Model ===
+        token_id_input = keras.Input(
+            shape=(None,), dtype="int32", name="token_ids"
+        )
+        padding_mask_input = keras.Input(
+            shape=(None,), dtype="int32", name="padding_mask"
+        )
+        x = self.embeddings(token_id_input)
+        x = self.embeddings_layer_norm(x)
+        x = self.embeddings_dropout(x)
+        for transformer_layer in self.transformer_layers:
+            x = transformer_layer(x, padding_mask=padding_mask_input)
         super().__init__(
             inputs={
                 "token_ids": token_id_input,
-                "padding_mask": padding_mask,
+                "padding_mask": padding_mask_input,
             },
             outputs=x,
             **kwargs,
         )
-        # All references to `self` below this line
+
+        # === Config ===
         self.vocabulary_size = vocabulary_size
         self.num_layers = num_layers
         self.num_heads = num_heads
@@ -160,7 +163,6 @@ class DistilBertBackbone(Backbone):
         self.dropout = dropout
         self.max_sequence_length = max_sequence_length
         self.cls_token_index = 0
-        self.token_embedding = embedding_layer.token_embedding
 
     def get_config(self):
         config = super().get_config()
diff --git a/keras_nlp/models/distil_bert/distil_bert_classifier.py b/keras_nlp/models/distil_bert/distil_bert_classifier.py
index 42de1cee..cf0db978 100644
--- a/keras_nlp/models/distil_bert/distil_bert_classifier.py
+++ b/keras_nlp/models/distil_bert/distil_bert_classifier.py
@@ -150,39 +150,46 @@ class DistilBertClassifier(Task):
         dropout=0.2,
         **kwargs,
     ):
-        inputs = backbone.input
+        # === Layers ===
+        self.backbone = backbone
+        self.preprocessor = preprocessor
         hidden_dim = hidden_dim or backbone.hidden_dim
-
-        x = backbone(inputs)[:, backbone.cls_token_index, :]
-        x = keras.layers.Dense(
+        self.pooled_dense = keras.layers.Dense(
             hidden_dim,
             activation="relu",
             kernel_initializer=distilbert_kernel_initializer(),
             name="pooled_dense",
-        )(x)
-        x = keras.layers.Dropout(dropout, name="classifier_dropout")(x)
-        outputs = keras.layers.Dense(
+        )
+        self.output_dropout = keras.layers.Dropout(
+            dropout,
+            name="output_dropout",
+        )
+        self.output_dense = keras.layers.Dense(
             num_classes,
             kernel_initializer=distilbert_kernel_initializer(),
             activation=activation,
             name="logits",
-        )(x)
+        )
 
-        # Instantiate using Functional API Model constructor
+        # === Functional Model ===
+        inputs = backbone.input
+        x = backbone(inputs)[:, backbone.cls_token_index, :]
+        x = self.pooled_dense(x)
+        x = self.output_dropout(x)
+        outputs = self.output_dense(x)
         super().__init__(
             inputs=inputs,
             outputs=outputs,
-            include_preprocessing=preprocessor is not None,
             **kwargs,
         )
-        # All references to `self` below this line
-        self.backbone = backbone
-        self.preprocessor = preprocessor
+
+        # === Config ===
         self.num_classes = num_classes
         self.activation = keras.activations.get(activation)
         self.hidden_dim = hidden_dim
         self.dropout = dropout
 
+        # === Default compilation ===
         logit_output = self.activation == keras.activations.linear
         self.compile(
             loss=keras.losses.SparseCategoricalCrossentropy(
diff --git a/keras_nlp/models/distil_bert/distil_bert_masked_lm.py b/keras_nlp/models/distil_bert/distil_bert_masked_lm.py
index 71cb117d..9be43f8a 100644
--- a/keras_nlp/models/distil_bert/distil_bert_masked_lm.py
+++ b/keras_nlp/models/distil_bert/distil_bert_masked_lm.py
@@ -104,6 +104,18 @@ class DistilBertMaskedLM(Task):
         preprocessor=None,
         **kwargs,
     ):
+        # === Layers ===
+        self.backbone = backbone
+        self.preprocessor = preprocessor
+        self.masked_lm_head = MaskedLMHead(
+            vocabulary_size=backbone.vocabulary_size,
+            token_embedding=backbone.token_embedding,
+            intermediate_activation="gelu",
+            kernel_initializer=distilbert_kernel_initializer(),
+            name="mlm_head",
+        )
+
+        # === Functional Model ===
         inputs = {
             **backbone.input,
             "mask_positions": keras.Input(
@@ -111,25 +123,16 @@ class DistilBertMaskedLM(Task):
             ),
         }
         backbone_outputs = backbone(backbone.input)
-        outputs = MaskedLMHead(
-            vocabulary_size=backbone.vocabulary_size,
-            token_embedding=backbone.token_embedding,
-            intermediate_activation="gelu",
-            kernel_initializer=distilbert_kernel_initializer(),
-            name="mlm_head",
-        )(backbone_outputs, inputs["mask_positions"])
-
-        # Instantiate using Functional API Model constructor
+        outputs = self.masked_lm_head(
+            backbone_outputs, inputs["mask_positions"]
+        )
         super().__init__(
             inputs=inputs,
             outputs=outputs,
-            include_preprocessing=preprocessor is not None,
             **kwargs,
         )
-        # All references to `self` below this line
-        self.backbone = backbone
-        self.preprocessor = preprocessor
 
+        # === Default compilation ===
         self.compile(
             loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
             optimizer=keras.optimizers.Adam(5e-5),
diff --git a/keras_nlp/models/electra/electra_backbone.py b/keras_nlp/models/electra/electra_backbone.py
index 66d1db8c..2e2a0197 100644
--- a/keras_nlp/models/electra/electra_backbone.py
+++ b/keras_nlp/models/electra/electra_backbone.py
@@ -94,65 +94,44 @@ class ElectraBackbone(Backbone):
         num_segments=2,
         **kwargs,
     ):
-        # Index of classification token in the vocabulary
-        cls_token_index = 0
-        # Inputs
-        token_id_input = keras.Input(
-            shape=(None,), dtype="int32", name="token_ids"
-        )
-        segment_id_input = keras.Input(
-            shape=(None,), dtype="int32", name="segment_ids"
-        )
-        padding_mask = keras.Input(
-            shape=(None,), dtype="int32", name="padding_mask"
-        )
-
-        # Embed tokens, positions, and segment ids.
-        token_embedding_layer = ReversibleEmbedding(
+        # === Layers ===
+        self.token_embedding = ReversibleEmbedding(
             input_dim=vocab_size,
             output_dim=embedding_dim,
             embeddings_initializer=electra_kernel_initializer(),
             name="token_embedding",
         )
-        token_embedding = token_embedding_layer(token_id_input)
-        position_embedding = PositionEmbedding(
+        self.position_embedding = PositionEmbedding(
             initializer=electra_kernel_initializer(),
             sequence_length=max_sequence_length,
             name="position_embedding",
-        )(token_embedding)
-        segment_embedding = keras.layers.Embedding(
+        )
+        self.segment_embedding = keras.layers.Embedding(
             input_dim=num_segments,
             output_dim=embedding_dim,
             embeddings_initializer=electra_kernel_initializer(),
             name="segment_embedding",
-        )(segment_id_input)
-
-        # Add all embeddings together.
-        x = keras.layers.Add()(
-            (token_embedding, position_embedding, segment_embedding),
         )
-        # Layer normalization
-        x = keras.layers.LayerNormalization(
+        self.embeddings_add = keras.layers.Add()
+        self.embeddings_layer_norm = keras.layers.LayerNormalization(
             name="embeddings_layer_norm",
             axis=-1,
             epsilon=1e-12,
             dtype="float32",
-        )(x)
-        # Dropout
-        x = keras.layers.Dropout(
+        )
+        self.embeddings_dropout = keras.layers.Dropout(
             dropout,
             name="embeddings_dropout",
-        )(x)
+        )
         if hidden_dim != embedding_dim:
-            x = keras.layers.Dense(
+            self.embeddings_projection = keras.layers.Dense(
                 hidden_dim,
                 kernel_initializer=electra_kernel_initializer(),
                 name="embeddings_projection",
-            )(x)
-
-        # Apply successive transformer encoder blocks.
+            )
+        self.transformer_layers = []
         for i in range(num_layers):
-            x = TransformerEncoder(
+            layer = TransformerEncoder(
                 num_heads=num_heads,
                 intermediate_dim=intermediate_dim,
                 activation=gelu_approximate,
@@ -160,24 +139,49 @@ class ElectraBackbone(Backbone):
                 layer_norm_epsilon=1e-12,
                 kernel_initializer=electra_kernel_initializer(),
                 name=f"transformer_layer_{i}",
-            )(x, padding_mask=padding_mask)
-
-        sequence_output = x
-        # Construct the two ELECTRA outputs. The pooled output is a dense layer on
-        # top of the [CLS] token.
-        pooled_output = keras.layers.Dense(
+            )
+            self.transformer_layers.append(layer)
+        self.pooled_dense = keras.layers.Dense(
             hidden_dim,
             kernel_initializer=electra_kernel_initializer(),
             activation="tanh",
             name="pooled_dense",
-        )(x[:, cls_token_index, :])
+        )
 
-        # Instantiate using Functional API Model constructor
+        # === Functional Model ===
+        token_id_input = keras.Input(
+            shape=(None,), dtype="int32", name="token_ids"
+        )
+        segment_id_input = keras.Input(
+            shape=(None,), dtype="int32", name="segment_ids"
+        )
+        padding_mask_input = keras.Input(
+            shape=(None,), dtype="int32", name="padding_mask"
+        )
+        # Embed tokens, positions, and segment ids.
+        tokens = self.token_embedding(token_id_input)
+        positions = self.position_embedding(tokens)
+        segments = self.segment_embedding(segment_id_input)
+        # Add all embeddings together.
+        x = self.embeddings_add((tokens, positions, segments))
+        x = self.embeddings_layer_norm(x)
+        x = self.embeddings_dropout(x)
+        if hidden_dim != embedding_dim:
+            x = self.embeddings_projection(x)
+        # Apply successive transformer encoder blocks.
+        for transformer_layer in self.transformer_layers:
+            x = transformer_layer(x, padding_mask=padding_mask_input)
+        # Index of classification token in the vocabulary
+        cls_token_index = 0
+        sequence_output = x
+        # Construct the two ELECTRA outputs. The pooled output is a dense layer on
+        # top of the [CLS] token.
+        pooled_output = self.pooled_dense(x[:, cls_token_index, :])
         super().__init__(
             inputs={
                 "token_ids": token_id_input,
                 "segment_ids": segment_id_input,
-                "padding_mask": padding_mask,
+                "padding_mask": padding_mask_input,
             },
             outputs={
                 "sequence_output": sequence_output,
@@ -186,7 +190,7 @@ class ElectraBackbone(Backbone):
             **kwargs,
         )
 
-        # All references to self below this line
+        # === Config ===
         self.vocab_size = vocab_size
         self.num_layers = num_layers
         self.num_heads = num_heads
@@ -197,7 +201,6 @@ class ElectraBackbone(Backbone):
         self.max_sequence_length = max_sequence_length
         self.num_segments = num_segments
         self.cls_token_index = cls_token_index
-        self.token_embedding = token_embedding_layer
 
     def get_config(self):
         config = super().get_config()
diff --git a/keras_nlp/models/f_net/f_net_backbone.py b/keras_nlp/models/f_net/f_net_backbone.py
index ac4d290b..9103a10d 100644
--- a/keras_nlp/models/f_net/f_net_backbone.py
+++ b/keras_nlp/models/f_net/f_net_backbone.py
@@ -101,61 +101,44 @@ class FNetBackbone(Backbone):
         num_segments=4,
         **kwargs,
     ):
-        # Index of classification token in the vocabulary
-        cls_token_index = 0
-        # Inputs
-        token_id_input = keras.Input(
-            shape=(None,), dtype="int32", name="token_ids"
-        )
-        segment_id_input = keras.Input(
-            shape=(None,), dtype="int32", name="segment_ids"
-        )
-
-        # Embed tokens, positions, and segment ids.
-        token_embedding_layer = ReversibleEmbedding(
+        # === Layers ===
+        self.token_embedding = ReversibleEmbedding(
             input_dim=vocabulary_size,
             output_dim=hidden_dim,
             embeddings_initializer=f_net_kernel_initializer(),
             name="token_embedding",
         )
-        token_embedding = token_embedding_layer(token_id_input)
-        position_embedding = PositionEmbedding(
+        self.position_embedding = PositionEmbedding(
             initializer=f_net_kernel_initializer(),
             sequence_length=max_sequence_length,
             name="position_embedding",
-        )(token_embedding)
-        segment_embedding = keras.layers.Embedding(
+        )
+        self.segment_embedding = keras.layers.Embedding(
             input_dim=num_segments,
             output_dim=hidden_dim,
             embeddings_initializer=f_net_kernel_initializer(),
             name="segment_embedding",
-        )(segment_id_input)
-
-        # Sum, normalize and apply dropout to embeddings.
-        x = keras.layers.Add()(
-            (token_embedding, position_embedding, segment_embedding)
         )
-        x = keras.layers.LayerNormalization(
+        self.embeddings_add = keras.layers.Add()
+        self.embeddings_layer_norm = keras.layers.LayerNormalization(
             name="embeddings_layer_norm",
             axis=-1,
             epsilon=1e-12,
             dtype="float32",
-        )(x)
-
-        x = keras.layers.Dense(
+        )
+        self.embedding_projection = keras.layers.Dense(
             hidden_dim,
             kernel_initializer=f_net_kernel_initializer(),
             bias_initializer=f_net_bias_initializer(),
             name="embedding_projection",
-        )(x)
-        x = keras.layers.Dropout(
+        )
+        self.embeddings_dropout = keras.layers.Dropout(
             dropout,
             name="embeddings_dropout",
-        )(x)
-
-        # Apply successive FNet encoder blocks.
+        )
+        self.transformer_layers = []
         for i in range(num_layers):
-            x = FNetEncoder(
+            layer = FNetEncoder(
                 intermediate_dim=intermediate_dim,
                 activation=gelu_approximate,
                 dropout=dropout,
@@ -163,19 +146,41 @@ class FNetBackbone(Backbone):
                 kernel_initializer=f_net_kernel_initializer(),
                 bias_initializer=f_net_bias_initializer(),
                 name=f"f_net_layer_{i}",
-            )(x)
-
-        # Construct the two FNet outputs. The pooled output is a dense layer on
-        # top of the [CLS] token.
-        sequence_output = x
-        pooled_output = keras.layers.Dense(
+            )
+            self.transformer_layers.append(layer)
+        self.pooled_dense = keras.layers.Dense(
             hidden_dim,
             kernel_initializer=f_net_kernel_initializer(),
             bias_initializer=f_net_bias_initializer(),
             activation="tanh",
             name="pooled_dense",
-        )(x[:, cls_token_index, :])
+        )
 
+        # === Functional Model ===
+        token_id_input = keras.Input(
+            shape=(None,), dtype="int32", name="token_ids"
+        )
+        segment_id_input = keras.Input(
+            shape=(None,), dtype="int32", name="segment_ids"
+        )
+        # Embed tokens, positions, and segment ids.
+        tokens = self.token_embedding(token_id_input)
+        positions = self.position_embedding(tokens)
+        segments = self.segment_embedding(segment_id_input)
+        # Sum, normalize and apply dropout to embeddings.
+        x = self.embeddings_add((tokens, positions, segments))
+        x = self.embeddings_layer_norm(x)
+        x = self.embedding_projection(x)
+        x = self.embeddings_dropout(x)
+        # Apply successive FNet encoder blocks.
+        for transformer_layer in self.transformer_layers:
+            x = transformer_layer(x)
+        # Index of classification token in the vocabulary
+        cls_token_index = 0
+        # Construct the two FNet outputs. The pooled output is a dense layer on
+        # top of the [CLS] token.
+        sequence_output = x
+        pooled_output = self.pooled_dense(x[:, cls_token_index, :])
         # Instantiate using Functional API Model constructor
         super().__init__(
             inputs={
@@ -189,7 +194,7 @@ class FNetBackbone(Backbone):
             **kwargs,
         )
 
-        # All references to `self` below this line
+        # === Config ===
         self.vocabulary_size = vocabulary_size
         self.num_layers = num_layers
         self.hidden_dim = hidden_dim
@@ -198,7 +203,6 @@ class FNetBackbone(Backbone):
         self.max_sequence_length = max_sequence_length
         self.num_segments = num_segments
         self.cls_token_index = cls_token_index
-        self.token_embedding = token_embedding_layer
 
     def get_config(self):
         config = super().get_config()
diff --git a/keras_nlp/models/f_net/f_net_classifier.py b/keras_nlp/models/f_net/f_net_classifier.py
index f6485485..f4ee31d1 100644
--- a/keras_nlp/models/f_net/f_net_classifier.py
+++ b/keras_nlp/models/f_net/f_net_classifier.py
@@ -109,29 +109,37 @@ class FNetClassifier(Task):
         dropout=0.1,
         **kwargs,
     ):
-        inputs = backbone.input
-        pooled = backbone(inputs)["pooled_output"]
-        pooled = keras.layers.Dropout(dropout)(pooled)
-        outputs = keras.layers.Dense(
+        # === Layers ===
+        self.backbone = backbone
+        self.preprocessor = preprocessor
+        self.output_dropout = keras.layers.Dropout(
+            dropout,
+            name="output_dropout",
+        )
+        self.output_dense = keras.layers.Dense(
             num_classes,
             kernel_initializer=f_net_kernel_initializer(),
             activation=activation,
             name="logits",
-        )(pooled)
-        # Instantiate using Functional API Model constructor
+        )
+
+        # === Functional Model ===
+        inputs = backbone.input
+        pooled = backbone(inputs)["pooled_output"]
+        pooled = self.output_dropout(pooled)
+        outputs = self.output_dense(pooled)
         super().__init__(
             inputs=inputs,
             outputs=outputs,
-            include_preprocessing=preprocessor is not None,
             **kwargs,
         )
-        # All references to `self` below this line
-        self.backbone = backbone
-        self.preprocessor = preprocessor
+
+        # === Config ===
         self.num_classes = num_classes
         self.activation = keras.activations.get(activation)
         self.dropout = dropout
 
+        # === Default compilation ===
         logit_output = self.activation == keras.activations.linear
         self.compile(
             loss=keras.losses.SparseCategoricalCrossentropy(
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 d7048cd5..c0eb231d 100644
--- a/keras_nlp/models/f_net/f_net_masked_lm.py
+++ b/keras_nlp/models/f_net/f_net_masked_lm.py
@@ -101,6 +101,18 @@ class FNetMaskedLM(Task):
         preprocessor=None,
         **kwargs,
     ):
+        # === Layers ===
+        self.backbone = backbone
+        self.preprocessor = preprocessor
+        self.masked_lm_head = MaskedLMHead(
+            vocabulary_size=backbone.vocabulary_size,
+            token_embedding=backbone.token_embedding,
+            intermediate_activation="gelu",
+            kernel_initializer=f_net_kernel_initializer(),
+            name="mlm_head",
+        )
+
+        # === Functional Model ===
         inputs = {
             **backbone.input,
             "mask_positions": keras.Input(
@@ -108,24 +120,16 @@ class FNetMaskedLM(Task):
             ),
         }
         backbone_outputs = backbone(backbone.input)
-        outputs = MaskedLMHead(
-            vocabulary_size=backbone.vocabulary_size,
-            token_embedding=backbone.token_embedding,
-            intermediate_activation="gelu",
-            kernel_initializer=f_net_kernel_initializer(),
-            name="mlm_head",
-        )(backbone_outputs["sequence_output"], inputs["mask_positions"])
-
-        # Instantiate using Functional API Model constructor
+        outputs = self.masked_lm_head(
+            backbone_outputs["sequence_output"], inputs["mask_positions"]
+        )
         super().__init__(
             inputs=inputs,
             outputs=outputs,
-            include_preprocessing=preprocessor is not None,
             **kwargs,
         )
-        # All references to `self` below this line
-        self.backbone = backbone
-        self.preprocessor = preprocessor
+
+        # === Default compilation ===
         self.compile(
             loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
             optimizer=keras.optimizers.Adam(5e-5),
diff --git a/keras_nlp/models/gpt2/gpt2_backbone.py b/keras_nlp/models/gpt2/gpt2_backbone.py
index 89c23f71..d3f4a415 100644
--- a/keras_nlp/models/gpt2/gpt2_backbone.py
+++ b/keras_nlp/models/gpt2/gpt2_backbone.py
@@ -97,68 +97,73 @@ class GPT2Backbone(Backbone):
         max_sequence_length=1024,
         **kwargs,
     ):
-        # Inputs
-        token_ids = keras.Input(shape=(None,), dtype="int32", name="token_ids")
-        padding_mask = keras.Input(
-            shape=(None,), dtype="int32", name="padding_mask"
-        )
-
-        # Embed tokens, positions.
-        token_embedding_layer = ReversibleEmbedding(
+        # === Layers ===
+        self.token_embedding = ReversibleEmbedding(
             input_dim=vocabulary_size,
             output_dim=hidden_dim,
             embeddings_initializer=_gpt_2_kernel_initializer(stddev=0.01),
             name="token_embedding",
         )
-        token_embedding = token_embedding_layer(token_ids)
-
-        # Can't use `TokenAndPositionEmbedding` layer here because of different
-        # initializers.
-        position_embedding = PositionEmbedding(
+        self.position_embedding = PositionEmbedding(
             initializer=_gpt_2_kernel_initializer(stddev=0.02),
             sequence_length=max_sequence_length,
             name="position_embedding",
-        )(token_embedding)
-
-        # Sum and apply dropout to embeddings.
-        x = keras.layers.Add(name="embeddings_add")(
-            (token_embedding, position_embedding)
         )
-        x = keras.layers.Dropout(
+        self.embeddings_add = keras.layers.Add(
+            name="embeddings_add",
+        )
+        self.embeddings_dropout = keras.layers.Dropout(
             dropout,
             name="embeddings_dropout",
-        )(x)
-
-        # Apply successive transformer decoder blocks.
+        )
+        self.transformer_layers = []
         for i in range(num_layers):
-            x = TransformerDecoder(
-                intermediate_dim=intermediate_dim,
-                num_heads=num_heads,
-                dropout=dropout,
-                layer_norm_epsilon=1e-05,
-                activation=gelu_approximate,
-                kernel_initializer=_gpt_2_kernel_initializer(stddev=0.02),
-                normalize_first=True,
-                name=f"transformer_layer_{i}",
-            )(x, decoder_padding_mask=padding_mask)
-
-        sequence_output = keras.layers.LayerNormalization(
+            self.transformer_layers.append(
+                TransformerDecoder(
+                    intermediate_dim=intermediate_dim,
+                    num_heads=num_heads,
+                    dropout=dropout,
+                    layer_norm_epsilon=1e-05,
+                    activation=gelu_approximate,
+                    kernel_initializer=_gpt_2_kernel_initializer(stddev=0.02),
+                    normalize_first=True,
+                    name=f"transformer_layer_{i}",
+                )
+            )
+        self.layer_norm = keras.layers.LayerNormalization(
             name="layer_norm",
             axis=-1,
             epsilon=1e-05,
             dtype="float32",
-        )(x)
+        )
 
-        # Instantiate using Functional API Model constructor
+        # === Functional Model ===
+        token_id_input = keras.Input(
+            shape=(None,), dtype="int32", name="token_ids"
+        )
+        padding_mask_input = keras.Input(
+            shape=(None,), dtype="int32", name="padding_mask"
+        )
+        # Embed inputs.
+        tokens = self.token_embedding(token_id_input)
+        positions = self.position_embedding(tokens)
+        x = self.embeddings_add((tokens, positions))
+        x = self.embeddings_dropout(x)
+        # Apply transformer layers.
+        for transformer_layer in self.transformer_layers:
+            x = transformer_layer(x, decoder_padding_mask=padding_mask_input)
+        sequence_output = self.layer_norm(x)
+        # Instantiate using the Functional constructor.
         super().__init__(
             inputs={
-                "token_ids": token_ids,
-                "padding_mask": padding_mask,
+                "token_ids": token_id_input,
+                "padding_mask": padding_mask_input,
             },
             outputs=sequence_output,
             **kwargs,
         )
-        # All references to `self` below this line
+
+        # === Config ===
         self.vocabulary_size = vocabulary_size
         self.num_layers = num_layers
         self.num_heads = num_heads
@@ -166,7 +171,6 @@ class GPT2Backbone(Backbone):
         self.intermediate_dim = intermediate_dim
         self.dropout = dropout
         self.max_sequence_length = max_sequence_length
-        self.token_embedding = token_embedding_layer
 
     def get_config(self):
         config = super().get_config()
diff --git a/keras_nlp/models/gpt2/gpt2_causal_lm.py b/keras_nlp/models/gpt2/gpt2_causal_lm.py
index 44eebd0a..e154c88b 100644
--- a/keras_nlp/models/gpt2/gpt2_causal_lm.py
+++ b/keras_nlp/models/gpt2/gpt2_causal_lm.py
@@ -155,23 +155,21 @@ class GPT2CausalLM(GenerativeTask):
         preprocessor=None,
         **kwargs,
     ):
+        # === Layers ===
+        self.backbone = backbone
+        self.preprocessor = preprocessor
+
+        # === Functional Model ===
         inputs = backbone.input
         hidden_states = backbone(inputs)
         outputs = backbone.token_embedding(hidden_states, reverse=True)
-
-        # Instantiate using Functional API Model constructor.
         super().__init__(
             inputs=inputs,
             outputs=outputs,
-            include_preprocessing=preprocessor is not None,
             **kwargs,
         )
-        self.backbone = backbone
-        self.preprocessor = preprocessor
-        self.generate_function = None
-        self._sampler = None
 
-        # Default compilation
+        # === Default compilation ===
         self.compile(
             loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
             optimizer=keras.optimizers.Adam(2e-5),
@@ -216,27 +214,25 @@ class GPT2CausalLM(GenerativeTask):
             the final hidden representation of the input tokens, and `cache` is
             the decoding cache.
         """
-        token_embedding = self.backbone.get_layer("token_embedding")(token_ids)
-        position_embedding = self.backbone.get_layer("position_embedding")(
-            token_embedding, start_index=cache_update_index
-        )
-        x = self.backbone.get_layer("embeddings_add")(
-            (token_embedding, position_embedding)
+        tokens = self.backbone.token_embedding(token_ids)
+        positions = self.backbone.position_embedding(
+            tokens, start_index=cache_update_index
         )
-        x = self.backbone.get_layer("embeddings_dropout")(x)
+        x = self.backbone.embeddings_add((tokens, positions))
+        x = self.backbone.embeddings_dropout(x)
         # Each decoder layer has a cache; we update them separately.
         caches = []
-        for i in range(self.backbone.num_layers):
+        for i, transformer_layer in enumerate(self.backbone.transformer_layers):
             current_cache = cache[:, i, ...]
-            x, next_cache = self.backbone.get_layer(f"transformer_layer_{i}")(
+            x, next_cache = transformer_layer(
                 x,
                 self_attention_cache=current_cache,
                 self_attention_cache_update_index=cache_update_index,
             )
             caches.append(next_cache)
         cache = ops.stack(caches, axis=1)
-        hidden_states = x = self.backbone.get_layer("layer_norm")(x)
-        logits = self.backbone.get_layer("token_embedding")(x, reverse=True)
+        hidden_states = x = self.backbone.layer_norm(x)
+        logits = self.backbone.token_embedding(x, reverse=True)
         return logits, hidden_states, cache
 
     def _build_cache(self, token_ids):
diff --git a/keras_nlp/models/gpt_neo_x/gpt_neo_x_backbone.py b/keras_nlp/models/gpt_neo_x/gpt_neo_x_backbone.py
index 6804331a..5bbc11af 100644
--- a/keras_nlp/models/gpt_neo_x/gpt_neo_x_backbone.py
+++ b/keras_nlp/models/gpt_neo_x/gpt_neo_x_backbone.py
@@ -77,29 +77,20 @@ class GPTNeoXBackbone(Backbone):
         max_sequence_length=512,
         **kwargs,
     ):
-        # Inputs
-        token_ids = keras.Input(shape=(None,), dtype="int32", name="token_ids")
-        padding_mask = keras.Input(
-            shape=(None,), dtype="int32", name="padding_mask"
-        )
-
-        # Embed tokens
-        token_embedding_layer = ReversibleEmbedding(
+        # === Layers ===
+        self.token_embedding = ReversibleEmbedding(
             input_dim=vocabulary_size,
             output_dim=hidden_dim,
             embeddings_initializer=_gpt_neo_x_kernel_initializer(stddev=0.01),
             name="token_embedding",
         )
-        token_embedding = token_embedding_layer(token_ids)
-
-        x = keras.layers.Dropout(
+        self.embeddings_dropout = keras.layers.Dropout(
             dropout,
             name="embeddings_dropout",
-        )(token_embedding)
-
-        # Apply successive transformer decoder blocks.
+        )
+        self.transformer_layers = []
         for i in range(num_layers):
-            x = GPTNeoXDecoder(
+            layer = GPTNeoXDecoder(
                 intermediate_dim=intermediate_dim,
                 num_heads=num_heads,
                 dropout=dropout,
@@ -110,25 +101,38 @@ class GPTNeoXBackbone(Backbone):
                 activation=gelu_approximate,
                 kernel_initializer=_gpt_neo_x_kernel_initializer(stddev=0.02),
                 name=f"transformer_layer_{i}",
-            )(x, decoder_padding_mask=padding_mask)
-
-        sequence_output = keras.layers.LayerNormalization(
+            )
+            self.transformer_layers.append(layer)
+        self.layer_norm = keras.layers.LayerNormalization(
             name="layer_norm",
             axis=-1,
             epsilon=layer_norm_epsilon,
             dtype="float32",
-        )(x)
+        )
 
-        # Instantiate using Functional API Model constructor
+        # === Functional Model ===
+        token_id_input = keras.Input(
+            shape=(None,), dtype="int32", name="token_ids"
+        )
+        padding_mask_input = keras.Input(
+            shape=(None,), dtype="int32", name="padding_mask"
+        )
+        # Embed tokens.
+        x = self.token_embedding(token_id_input)
+        x = self.embeddings_dropout(x)
+        for transformer_layer in self.transformer_layers:
+            x = transformer_layer(x, decoder_padding_mask=padding_mask_input)
+        sequence_output = self.layer_norm(x)
         super().__init__(
             inputs={
-                "token_ids": token_ids,
-                "padding_mask": padding_mask,
+                "token_ids": token_id_input,
+                "padding_mask": padding_mask_input,
             },
             outputs=sequence_output,
             **kwargs,
         )
-        # All references to `self` below this line
+
+        # === Config ===
         self.vocabulary_size = vocabulary_size
         self.num_layers = num_layers
         self.num_heads = num_heads
@@ -139,7 +143,6 @@ class GPTNeoXBackbone(Backbone):
         self.rotary_max_wavelength = rotary_max_wavelength
         self.max_sequence_length = max_sequence_length
         self.layer_norm_epsilon = layer_norm_epsilon
-        self.token_embedding = token_embedding_layer
 
     def get_config(self):
         config = super().get_config()
diff --git a/keras_nlp/models/gpt_neo_x/gpt_neo_x_causal_lm.py b/keras_nlp/models/gpt_neo_x/gpt_neo_x_causal_lm.py
index 0f813470..bef32017 100644
--- a/keras_nlp/models/gpt_neo_x/gpt_neo_x_causal_lm.py
+++ b/keras_nlp/models/gpt_neo_x/gpt_neo_x_causal_lm.py
@@ -52,23 +52,21 @@ class GPTNeoXCausalLM(GenerativeTask):
         preprocessor=None,
         **kwargs,
     ):
+        # === Layers ===
+        self.backbone = backbone
+        self.preprocessor = preprocessor
+
+        # === Functional Model ===
         inputs = backbone.input
         hidden_states = backbone(inputs)
         outputs = backbone.token_embedding(hidden_states, reverse=True)
-
-        # Instantiate using Functional API Model constructor.
         super().__init__(
             inputs=inputs,
             outputs=outputs,
-            include_preprocessing=preprocessor is not None,
             **kwargs,
         )
-        self.backbone = backbone
-        self.preprocessor = preprocessor
-        self.generate_function = None
-        self._sampler = None
 
-        # Default compilation
+        # === Default compilation ===
         self.compile(
             loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
             optimizer=keras.optimizers.Adam(2e-5),
@@ -109,20 +107,20 @@ class GPTNeoXCausalLM(GenerativeTask):
             the final hidden representation of the input tokens, and `cache` is
             the decoding cache.
         """
-        token_embedding = self.backbone.get_layer("token_embedding")(token_ids)
-        x = self.backbone.get_layer("embeddings_dropout")(token_embedding)
+        token_embedding = self.backbone.token_embedding(token_ids)
+        x = self.backbone.embeddings_dropout(token_embedding)
         # Each decoder layer has a cache; we update them separately.
         caches = []
-        for i in range(self.backbone.num_layers):
+        for i, transformer_layer in enumerate(self.backbone.transformer_layers):
             current_cache = cache[:, i, ...]
-            x, next_cache = self.backbone.get_layer(f"transformer_layer_{i}")(
+            x, next_cache = transformer_layer(
                 x,
                 self_attention_cache=current_cache,
                 self_attention_cache_update_index=cache_update_index,
             )
             caches.append(next_cache)
         cache = ops.stack(caches, axis=1)
-        x = self.backbone.get_layer("layer_norm")(x)
+        x = self.backbone.layer_norm(x)
         hidden_states = x
         logits = self.backbone.token_embedding(hidden_states, reverse=True)
         return logits, hidden_states, cache
diff --git a/keras_nlp/models/llama/llama_backbone.py b/keras_nlp/models/llama/llama_backbone.py
index 63438544..46cfdc37 100644
--- a/keras_nlp/models/llama/llama_backbone.py
+++ b/keras_nlp/models/llama/llama_backbone.py
@@ -75,26 +75,17 @@ class LlamaBackbone(Backbone):
         max_sequence_length=4096,
         **kwargs,
     ):
-        # Inputs
-        token_ids = keras.Input(shape=(None,), dtype="int32", name="token_ids")
-        padding_mask = keras.Input(
-            shape=(None,), dtype="int32", name="padding_mask"
-        )
-
-        # Embed tokens
-        token_embedding = ReversibleEmbedding(
+        # === Layers ===
+        self.token_embedding = ReversibleEmbedding(
             input_dim=vocabulary_size,
             output_dim=hidden_dim,
             embeddings_initializer=_llama_kernel_initializer(stddev=0.01),
             tie_weights=False,
             name="token_embedding",
-        )(token_ids)
-
-        x = token_embedding
-
-        # Apply successive transformer decoder blocks.
+        )
+        self.transformer_layers = []
         for i in range(num_layers):
-            x = LlamaDecoder(
+            layer = LlamaDecoder(
                 intermediate_dim=intermediate_dim,
                 num_query_heads=num_query_heads,
                 num_key_value_heads=num_key_value_heads,
@@ -105,23 +96,34 @@ class LlamaBackbone(Backbone):
                 activation=ops.silu,
                 kernel_initializer=_llama_kernel_initializer(stddev=0.02),
                 name=f"transformer_layer_{i}",
-            )(x, decoder_padding_mask=padding_mask)
-
-        sequence_output = LlamaLayerNorm(
+            )
+            self.transformer_layers.append(layer)
+        self.layer_norm = LlamaLayerNorm(
             name="layer_norm",
             epsilon=layer_norm_epsilon,
-        )(x)
+        )
 
-        # Instantiate using Functional API Model constructor
+        # === Functional Model ===
+        token_id_input = keras.Input(
+            shape=(None,), dtype="int32", name="token_ids"
+        )
+        padding_mask_input = keras.Input(
+            shape=(None,), dtype="int32", name="padding_mask"
+        )
+        x = self.token_embedding(token_id_input)
+        for transformer_layer in self.transformer_layers:
+            x = transformer_layer(x, decoder_padding_mask=padding_mask_input)
+        sequence_output = self.layer_norm(x)
         super().__init__(
             inputs={
-                "token_ids": token_ids,
-                "padding_mask": padding_mask,
+                "token_ids": token_id_input,
+                "padding_mask": padding_mask_input,
             },
             outputs=sequence_output,
             **kwargs,
         )
-        # All references to `self` below this line
+
+        # === Config ===
         self.vocabulary_size = vocabulary_size
         self.num_layers = num_layers
         self.num_query_heads = num_query_heads
@@ -150,7 +152,3 @@ class LlamaBackbone(Backbone):
             }
         )
         return config
-
-    @property
-    def token_embedding(self):
-        return self.get_layer("token_embedding")
diff --git a/keras_nlp/models/mistral/mistral_backbone.py b/keras_nlp/models/mistral/mistral_backbone.py
index 42cec8b2..107e5699 100644
--- a/keras_nlp/models/mistral/mistral_backbone.py
+++ b/keras_nlp/models/mistral/mistral_backbone.py
@@ -109,17 +109,9 @@ class MistralBackbone(Backbone):
         dropout=0,
         **kwargs,
     ):
-        # Get the dtype
+        # === Layers ===
         dtype = kwargs.pop("dtype", keras.backend.floatx())
-
-        # Inputs
-        token_ids = keras.Input(shape=(None,), dtype="int32", name="token_ids")
-        padding_mask = keras.Input(
-            shape=(None,), dtype="int32", name="padding_mask"
-        )
-
-        # Embed Tokens
-        token_embedding_layer = ReversibleEmbedding(
+        self.token_embedding = ReversibleEmbedding(
             input_dim=vocabulary_size,
             output_dim=hidden_dim,
             tie_weights=False,
@@ -127,11 +119,9 @@ class MistralBackbone(Backbone):
             dtype=dtype,
             name="token_embedding",
         )
-        x = token_embedding_layer(token_ids)
-
-        # Apply successive transformer decoder blocks
+        self.transformer_layers = []
         for i in range(num_layers):
-            x = MistralTransformerDecoder(
+            layer = MistralTransformerDecoder(
                 intermediate_dim=intermediate_dim,
                 num_query_heads=num_query_heads,
                 num_key_value_heads=num_key_value_heads,
@@ -144,25 +134,35 @@ class MistralBackbone(Backbone):
                 dropout=dropout,
                 dtype=dtype,
                 name=f"transformer_layer_{i}",
-            )(x, decoder_padding_mask=padding_mask)
-
-        sequence_output = MistralLayerNormalization(
-            name="sequence_output_layernorm",
+            )
+            self.transformer_layers.append(layer)
+        self.layer_norm = MistralLayerNormalization(
             epsilon=layer_norm_epsilon,
             dtype=dtype,
-        )(x)
+            name="sequence_output_layernorm",
+        )
 
-        # Instantiate using Functional API Model constructor
+        # === Functional Model ===
+        token_id_input = keras.Input(
+            shape=(None,), dtype="int32", name="token_ids"
+        )
+        padding_mask_input = keras.Input(
+            shape=(None,), dtype="int32", name="padding_mask"
+        )
+        x = self.token_embedding(token_id_input)
+        for transformer_layer in self.transformer_layers:
+            x = transformer_layer(x, decoder_padding_mask=padding_mask_input)
+        sequence_output = self.layer_norm(x)
         super().__init__(
             inputs={
-                "token_ids": token_ids,
-                "padding_mask": padding_mask,
+                "token_ids": token_id_input,
+                "padding_mask": padding_mask_input,
             },
             outputs=sequence_output,
             **kwargs,
         )
 
-        # All references to `self` below this line
+        # === Config ===
         self.vocabulary_size = vocabulary_size
         self.num_layers = num_layers
         self.num_query_heads = num_query_heads
@@ -174,7 +174,6 @@ class MistralBackbone(Backbone):
         self.sliding_window = sliding_window
         self.layer_norm_epsilon = layer_norm_epsilon
         self.dropout = dropout
-        self.token_embedding = token_embedding_layer
 
     def get_config(self):
         config = super().get_config()
diff --git a/keras_nlp/models/opt/opt_backbone.py b/keras_nlp/models/opt/opt_backbone.py
index ff1495ba..d04f4b57 100644
--- a/keras_nlp/models/opt/opt_backbone.py
+++ b/keras_nlp/models/opt/opt_backbone.py
@@ -93,25 +93,18 @@ class OPTBackbone(Backbone):
         max_sequence_length=2048,
         **kwargs,
     ):
-        # Decoder inputs.
-        token_ids = keras.Input(shape=(None,), dtype="int32", name="token_ids")
-        padding_mask = keras.Input(
-            shape=(None,), dtype="int32", name="padding_mask"
-        )
-
-        # Embed tokens and positions.
-        embedding_layer = TokenAndPositionEmbedding(
+        # === Layers ===
+        self.embeddings = TokenAndPositionEmbedding(
             vocabulary_size=vocabulary_size,
             sequence_length=max_sequence_length,
             embedding_dim=hidden_dim,
             embeddings_initializer=opt_kernel_initializer(),
             name="embeddings",
         )
-        x = embedding_layer(token_ids)
-
-        # Apply successive transformer decoder blocks.
+        self.token_embedding = self.embeddings.token_embedding
+        self.transformer_layers = []
         for i in range(num_layers):
-            x = TransformerDecoder(
+            layer = TransformerDecoder(
                 intermediate_dim=intermediate_dim,
                 num_heads=num_heads,
                 dropout=dropout,
@@ -120,27 +113,36 @@ class OPTBackbone(Backbone):
                 normalize_first=True,
                 kernel_initializer=opt_kernel_initializer(),
                 name=f"transformer_layer_{i}",
-            )(x, decoder_padding_mask=padding_mask)
-
-        # Add a final layer norm.
-        x = keras.layers.LayerNormalization(
-            name="layer_norm",
+            )
+            self.transformer_layers.append(layer)
+        self.layer_norm = keras.layers.LayerNormalization(
             axis=-1,
             epsilon=1e-5,
             dtype="float32",
-        )(x)
+            name="layer_norm",
+        )
 
-        # Instantiate using Functional API Model constructor
+        # === Functional Model ===
+        token_id_input = keras.Input(
+            shape=(None,), dtype="int32", name="token_ids"
+        )
+        padding_mask_input = keras.Input(
+            shape=(None,), dtype="int32", name="padding_mask"
+        )
+        x = self.embeddings(token_id_input)
+        for transformer_layer in self.transformer_layers:
+            x = transformer_layer(x, decoder_padding_mask=padding_mask_input)
+        x = self.layer_norm(x)
         super().__init__(
             inputs={
-                "token_ids": token_ids,
-                "padding_mask": padding_mask,
+                "token_ids": token_id_input,
+                "padding_mask": padding_mask_input,
             },
             outputs=x,
             **kwargs,
         )
 
-        # All references to `self` below this line
+        # === Config ===
         self.vocabulary_size = vocabulary_size
         self.num_layers = num_layers
         self.num_heads = num_heads
@@ -148,7 +150,6 @@ class OPTBackbone(Backbone):
         self.intermediate_dim = intermediate_dim
         self.dropout = dropout
         self.max_sequence_length = max_sequence_length
-        self.token_embedding = embedding_layer.token_embedding
 
     def get_config(self):
         return {
diff --git a/keras_nlp/models/opt/opt_causal_lm.py b/keras_nlp/models/opt/opt_causal_lm.py
index 6197a87f..9715bc6b 100644
--- a/keras_nlp/models/opt/opt_causal_lm.py
+++ b/keras_nlp/models/opt/opt_causal_lm.py
@@ -155,23 +155,21 @@ class OPTCausalLM(GenerativeTask):
         preprocessor=None,
         **kwargs,
     ):
+        # === Layers ===
+        self.backbone = backbone
+        self.preprocessor = preprocessor
+
+        # === Functional Model ===
         inputs = backbone.input
         hidden_states = backbone(inputs)
         outputs = backbone.token_embedding(hidden_states, reverse=True)
-
-        # Instantiate using Functional API Model constructor.
         super().__init__(
             inputs=inputs,
             outputs=outputs,
-            include_preprocessing=preprocessor is not None,
             **kwargs,
         )
-        self.backbone = backbone
-        self.preprocessor = preprocessor
-        self.generate_function = None
-        self._sampler = None
 
-        # Default compilation
+        # === Default compilation ===
         self.compile(
             loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
             optimizer=keras.optimizers.Adam(2e-5),
@@ -216,21 +214,19 @@ class OPTCausalLM(GenerativeTask):
             the final hidden representation of the input tokens, and `cache` is
             the decoding cache.
         """
-        x = self.backbone.get_layer("embeddings")(
-            token_ids, start_index=cache_update_index
-        )
+        x = self.backbone.embeddings(token_ids, start_index=cache_update_index)
         # Each decoder layer has a cache; we update them separately.
         caches = []
-        for i in range(self.backbone.num_layers):
+        for i, transformer_layer in enumerate(self.backbone.transformer_layers):
             current_cache = cache[:, i, ...]
-            x, next_cache = self.backbone.get_layer(f"transformer_layer_{i}")(
+            x, next_cache = transformer_layer(
                 x,
                 self_attention_cache=current_cache,
                 self_attention_cache_update_index=cache_update_index,
             )
             caches.append(next_cache)
         cache = ops.stack(caches, axis=1)
-        x = self.backbone.get_layer("layer_norm")(x)
+        x = self.backbone.layer_norm(x)
         hidden_states = x
         logits = self.backbone.token_embedding(hidden_states, reverse=True)
         return logits, hidden_states, cache
diff --git a/keras_nlp/models/roberta/roberta_backbone.py b/keras_nlp/models/roberta/roberta_backbone.py
index 8495b5cb..614104d8 100644
--- a/keras_nlp/models/roberta/roberta_backbone.py
+++ b/keras_nlp/models/roberta/roberta_backbone.py
@@ -98,39 +98,28 @@ class RobertaBackbone(Backbone):
         max_sequence_length=512,
         **kwargs,
     ):
-        # Inputs
-        token_id_input = keras.Input(
-            shape=(None,), dtype="int32", name="token_ids"
-        )
-        padding_mask = keras.Input(
-            shape=(None,), dtype="int32", name="padding_mask"
-        )
-
-        # Embed tokens and positions.
-        embedding_layer = TokenAndPositionEmbedding(
+        # === Layers ===
+        self.embeddings = TokenAndPositionEmbedding(
             vocabulary_size=vocabulary_size,
             sequence_length=max_sequence_length,
             embedding_dim=hidden_dim,
             embeddings_initializer=roberta_kernel_initializer(),
             name="embeddings",
         )
-        embedding = embedding_layer(token_id_input)
-
-        # Sum, normalize and apply dropout to embeddings.
-        x = keras.layers.LayerNormalization(
-            name="embeddings_layer_norm",
+        self.token_embedding = self.embeddings.token_embedding
+        self.embeddings_layer_norm = keras.layers.LayerNormalization(
             axis=-1,
             epsilon=1e-5,  # Original paper uses this epsilon value
             dtype="float32",
-        )(embedding)
-        x = keras.layers.Dropout(
+            name="embeddings_layer_norm",
+        )
+        self.embeddings_dropout = keras.layers.Dropout(
             dropout,
             name="embeddings_dropout",
-        )(x)
-
-        # Apply successive transformer encoder blocks.
+        )
+        self.transformer_layers = []
         for i in range(num_layers):
-            x = TransformerEncoder(
+            layer = TransformerEncoder(
                 num_heads=num_heads,
                 intermediate_dim=intermediate_dim,
                 activation="gelu",
@@ -138,18 +127,31 @@ class RobertaBackbone(Backbone):
                 layer_norm_epsilon=1e-5,
                 kernel_initializer=roberta_kernel_initializer(),
                 name=f"transformer_layer_{i}",
-            )(x, padding_mask=padding_mask)
+            )
+            self.transformer_layers.append(layer)
 
-        # Instantiate using Functional API Model constructor
+        # === Functional Model ===
+        token_id_input = keras.Input(
+            shape=(None,), dtype="int32", name="token_ids"
+        )
+        padding_mask_input = keras.Input(
+            shape=(None,), dtype="int32", name="padding_mask"
+        )
+        x = self.embeddings(token_id_input)
+        x = self.embeddings_layer_norm(x)
+        x = self.embeddings_dropout(x)
+        for transformer_layer in self.transformer_layers:
+            x = transformer_layer(x, padding_mask=padding_mask_input)
         super().__init__(
             inputs={
                 "token_ids": token_id_input,
-                "padding_mask": padding_mask,
+                "padding_mask": padding_mask_input,
             },
             outputs=x,
             **kwargs,
         )
-        # All references to `self` below this line
+
+        # === Config ===
         self.vocabulary_size = vocabulary_size
         self.num_layers = num_layers
         self.num_heads = num_heads
@@ -158,7 +160,6 @@ class RobertaBackbone(Backbone):
         self.dropout = dropout
         self.max_sequence_length = max_sequence_length
         self.start_token_index = 0
-        self.token_embedding = embedding_layer.token_embedding
 
     def get_config(self):
         config = super().get_config()
diff --git a/keras_nlp/models/roberta/roberta_classifier.py b/keras_nlp/models/roberta/roberta_classifier.py
index 9098d954..e3d7666f 100644
--- a/keras_nlp/models/roberta/roberta_classifier.py
+++ b/keras_nlp/models/roberta/roberta_classifier.py
@@ -144,38 +144,50 @@ class RobertaClassifier(Task):
         dropout=0.0,
         **kwargs,
     ):
-        inputs = backbone.input
+        # === Layers ===
+        self.backbone = backbone
+        self.preprocessor = preprocessor
+        self.pooled_dropout = keras.layers.Dropout(
+            dropout,
+            name="pooled_dropout",
+        )
         hidden_dim = hidden_dim or backbone.hidden_dim
-
-        x = backbone(inputs)[:, backbone.start_token_index, :]
-        x = keras.layers.Dropout(dropout, name="pooled_dropout")(x)
-        x = keras.layers.Dense(
-            hidden_dim, activation="tanh", name="pooled_dense"
-        )(x)
-        x = keras.layers.Dropout(dropout, name="classifier_dropout")(x)
-        outputs = keras.layers.Dense(
+        self.pooled_dense = keras.layers.Dense(
+            hidden_dim,
+            activation="tanh",
+            name="pooled_dense",
+        )
+        self.output_dropout = keras.layers.Dropout(
+            dropout,
+            name="output_dropout",
+        )
+        self.output_dense = keras.layers.Dense(
             num_classes,
             kernel_initializer=roberta_kernel_initializer(),
             activation=activation,
             name="logits",
-        )(x)
+        )
 
-        # Instantiate using Functional API Model constructor
+        # === Functional Model ===
+        inputs = backbone.input
+        x = backbone(inputs)[:, backbone.start_token_index, :]
+        x = self.pooled_dropout(x)
+        x = self.pooled_dense(x)
+        x = self.output_dropout(x)
+        outputs = self.output_dense(x)
         super().__init__(
             inputs=inputs,
             outputs=outputs,
-            include_preprocessing=preprocessor is not None,
             **kwargs,
         )
-        # All references to `self` below this line
-        self.backbone = backbone
-        self.preprocessor = preprocessor
+
+        # === Config ===
         self.num_classes = num_classes
         self.activation = keras.activations.get(activation)
         self.hidden_dim = hidden_dim
         self.dropout = dropout
 
-        # Default compilation
+        # === Default compilation ===
         logit_output = self.activation == keras.activations.linear
         self.compile(
             loss=keras.losses.SparseCategoricalCrossentropy(
diff --git a/keras_nlp/models/roberta/roberta_masked_lm.py b/keras_nlp/models/roberta/roberta_masked_lm.py
index 1517f259..0e62f4cf 100644
--- a/keras_nlp/models/roberta/roberta_masked_lm.py
+++ b/keras_nlp/models/roberta/roberta_masked_lm.py
@@ -103,6 +103,18 @@ class RobertaMaskedLM(Task):
         preprocessor=None,
         **kwargs,
     ):
+        # === Layers ===
+        self.backbone = backbone
+        self.preprocessor = preprocessor
+        self.masked_lm_head = MaskedLMHead(
+            vocabulary_size=backbone.vocabulary_size,
+            token_embedding=backbone.token_embedding,
+            intermediate_activation="gelu",
+            kernel_initializer=roberta_kernel_initializer(),
+            name="mlm_head",
+        )
+
+        # === Functional Model ===
         inputs = {
             **backbone.input,
             "mask_positions": keras.Input(
@@ -110,25 +122,16 @@ class RobertaMaskedLM(Task):
             ),
         }
         backbone_outputs = backbone(backbone.input)
-        outputs = MaskedLMHead(
-            vocabulary_size=backbone.vocabulary_size,
-            token_embedding=backbone.token_embedding,
-            intermediate_activation="gelu",
-            kernel_initializer=roberta_kernel_initializer(),
-            name="mlm_head",
-        )(backbone_outputs, inputs["mask_positions"])
-
-        # Instantiate using Functional API Model constructor
+        outputs = self.masked_lm_head(
+            backbone_outputs, inputs["mask_positions"]
+        )
         super().__init__(
             inputs=inputs,
             outputs=outputs,
-            include_preprocessing=preprocessor is not None,
             **kwargs,
         )
-        # All references to `self` below this line
-        self.backbone = backbone
-        self.preprocessor = preprocessor
 
+        # === Default compilation ===
         self.compile(
             loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
             optimizer=keras.optimizers.Adam(5e-5),
diff --git a/keras_nlp/models/t5/t5_backbone.py b/keras_nlp/models/t5/t5_backbone.py
index 6e76094d..3fdf69ff 100644
--- a/keras_nlp/models/t5/t5_backbone.py
+++ b/keras_nlp/models/t5/t5_backbone.py
@@ -85,45 +85,21 @@ class T5Backbone(Backbone):
         tie_embedding_weights=True,
         **kwargs,
     ):
-        # Encoder inputs
-        encoder_token_ids = keras.Input(
-            shape=(None,), dtype="int32", name="encoder_token_ids"
-        )
-        encoder_padding_mask = keras.Input(
-            shape=(None,), dtype="int32", name="encoder_padding_mask"
-        )
-
-        # Decoder inputs.
-        decoder_token_ids = keras.Input(
-            shape=(None,), dtype="int32", name="decoder_token_ids"
-        )
-        decoder_padding_mask = keras.Input(
-            shape=(None,), dtype="int32", name="decoder_padding_mask"
-        )
-
         # Token embedding layer. This layer is shared by encoder and decoder.
-        token_embedding_layer = ReversibleEmbedding(
+        self.token_embedding = ReversibleEmbedding(
             input_dim=vocabulary_size,
             output_dim=hidden_dim,
             tie_weights=tie_embedding_weights,
             embeddings_initializer=keras.initializers.TruncatedNormal(1.0),
             name="token_embedding",
         )
-
-        # ===== Encoder =====
-
-        # Embed tokens.
-        token_embedding = token_embedding_layer(encoder_token_ids)
-        x = keras.layers.Dropout(
+        self.encoder_embedding_dropout = keras.layers.Dropout(
             dropout,
             name="encoder_embedding_dropout",
-        )(token_embedding)
-
-        encoder_attention_mask = encoder_padding_mask[:, None, :]
-
-        position_bias = None
+        )
+        self.encoder_transformer_layers = []
         for i in range(num_layers):
-            output = T5TransformerLayer(
+            layer = T5TransformerLayer(
                 is_decoder=False,
                 hidden_dim=hidden_dim,
                 intermediate_dim=intermediate_dim,
@@ -135,39 +111,23 @@ class T5Backbone(Backbone):
                 use_gated_activation=use_gated_activation,
                 use_relative_attention_bias=bool(i == 0),
                 name=f"transformer_encoder_layer_{i}",
-            )(
-                x,
-                attention_mask=encoder_attention_mask,
-                position_bias=position_bias,
-                use_causal_mask=False,
             )
-            if isinstance(output, tuple):
-                x, position_bias = output
-
-        x = T5LayerNorm(
+            self.encoder_transformer_layers.append(layer)
+        self.encoder_layer_norm = T5LayerNorm(
             epsilon=layer_norm_epsilon,
             name="encoder_output_layer_norm",
-        )(x)
-        x = keras.layers.Dropout(
+        )
+        self.encoder_dropout = keras.layers.Dropout(
             dropout,
             name="encoder_output_dropout",
-        )(x)
-        encoder_output = x
-
-        # ===== Decoder =====
-
-        # Embed tokens.
-        token_embedding = token_embedding_layer(decoder_token_ids)
-        x = keras.layers.Dropout(
+        )
+        self.decoder_embedding_dropout = keras.layers.Dropout(
             dropout,
             name="decoder_embedding_dropout",
-        )(token_embedding)
-
-        decoder_attention_mask = decoder_padding_mask[:, None, :]
-
-        position_bias = None
+        )
+        self.decoder_transformer_layers = []
         for i in range(num_layers):
-            output = T5TransformerLayer(
+            layer = T5TransformerLayer(
                 is_decoder=True,
                 hidden_dim=hidden_dim,
                 intermediate_dim=intermediate_dim,
@@ -179,7 +139,54 @@ class T5Backbone(Backbone):
                 use_gated_activation=use_gated_activation,
                 use_relative_attention_bias=bool(i == 0),
                 name=f"transformer_decoder_layer_{i}",
-            )(
+            )
+            self.decoder_transformer_layers.append(layer)
+        self.decoder_layer_norm = T5LayerNorm(
+            epsilon=layer_norm_epsilon,
+            name="decoder_output_layer_norm",
+        )
+        self.decoder_dropout = keras.layers.Dropout(
+            dropout,
+            name="decoder_output_dropout",
+        )
+
+        # === Functional Model ===
+        encoder_token_id_input = keras.Input(
+            shape=(None,), dtype="int32", name="encoder_token_ids"
+        )
+        encoder_padding_mask_input = keras.Input(
+            shape=(None,), dtype="int32", name="encoder_padding_mask"
+        )
+        decoder_token_id_input = keras.Input(
+            shape=(None,), dtype="int32", name="decoder_token_ids"
+        )
+        decoder_padding_mask_input = keras.Input(
+            shape=(None,), dtype="int32", name="decoder_padding_mask"
+        )
+        # Encoder.
+        x = self.token_embedding(encoder_token_id_input)
+        x = self.encoder_embedding_dropout(x)
+        encoder_attention_mask = encoder_padding_mask_input[:, None, :]
+        position_bias = None
+        for transformer_layer in self.encoder_transformer_layers:
+            output = transformer_layer(
+                x,
+                attention_mask=encoder_attention_mask,
+                position_bias=position_bias,
+                use_causal_mask=False,
+            )
+            if isinstance(output, tuple):
+                x, position_bias = output
+        x = self.encoder_layer_norm(x)
+        x = self.encoder_dropout(x)
+        encoder_output = x
+        # Decoder.
+        x = self.token_embedding(decoder_token_id_input)
+        x = self.decoder_embedding_dropout(x)
+        decoder_attention_mask = decoder_padding_mask_input[:, None, :]
+        position_bias = None
+        for transformer_layer in self.decoder_transformer_layers:
+            output = transformer_layer(
                 x,
                 attention_mask=decoder_attention_mask,
                 position_bias=position_bias,
@@ -189,23 +196,15 @@ class T5Backbone(Backbone):
             )
             if isinstance(output, tuple):
                 x, position_bias = output
-
-        x = T5LayerNorm(
-            epsilon=layer_norm_epsilon,
-            name="decoder_output_layer_norm",
-        )(x)
-        x = keras.layers.Dropout(
-            dropout,
-            name="decoder_output_dropout",
-        )(x)
+        x = self.decoder_layer_norm(x)
+        x = self.decoder_dropout(x)
         decoder_output = x
-
         super().__init__(
             {
-                "encoder_token_ids": encoder_token_ids,
-                "encoder_padding_mask": encoder_padding_mask,
-                "decoder_token_ids": decoder_token_ids,
-                "decoder_padding_mask": decoder_padding_mask,
+                "encoder_token_ids": encoder_token_id_input,
+                "encoder_padding_mask": encoder_padding_mask_input,
+                "decoder_token_ids": decoder_token_id_input,
+                "decoder_padding_mask": decoder_padding_mask_input,
             },
             outputs={
                 "encoder_sequence_output": encoder_output,
@@ -213,7 +212,8 @@ class T5Backbone(Backbone):
             },
             **kwargs,
         )
-        # All references to `self` below this line
+
+        # === Config ===
         self.vocabulary_size = vocabulary_size
         self.hidden_dim = hidden_dim
         self.intermediate_dim = intermediate_dim
@@ -225,7 +225,6 @@ class T5Backbone(Backbone):
         self.use_gated_activation = use_gated_activation
         self.layer_norm_epsilon = layer_norm_epsilon
         self.tie_embedding_weights = tie_embedding_weights
-        self.token_embedding = token_embedding_layer
 
     def get_config(self):
         config = super().get_config()
diff --git a/keras_nlp/models/task.py b/keras_nlp/models/task.py
index ee28e3a9..1fe8d0b7 100644
--- a/keras_nlp/models/task.py
+++ b/keras_nlp/models/task.py
@@ -31,18 +31,25 @@ class Task(PipelineModel):
     """Base class for Task models."""
 
     def __init__(self, *args, **kwargs):
-        self._backbone = None
-        self._preprocessor = None
         super().__init__(*args, **kwargs)
         self._functional_layer_ids = set(
             id(layer) for layer in self._flatten_layers()
         )
+        self._initialized = True
 
     def __dir__(self):
-        # Temporary fixes for weight saving. This mimics the following PR for
+        if config.keras_3():
+            return super().__dir__()
+
+        # Temporary fixes for Keras 2 saving. This mimics the following PR for
         # older version of Keras: https://github.com/keras-team/keras/pull/18982
         def filter_fn(attr):
-            if attr == "_layer_checkpoint_dependencies":
+            if attr in [
+                "_layer_checkpoint_dependencies",
+                "transformer_layers",
+                "encoder_transformer_layers",
+                "decoder_transformer_layers",
+            ]:
                 return False
             return id(getattr(self, attr)) not in self._functional_layer_ids
 
@@ -99,17 +106,28 @@ class Task(PipelineModel):
         super().compile(optimizer=optimizer, loss=loss, **kwargs)
 
     def preprocess_samples(self, x, y=None, sample_weight=None):
-        return self.preprocessor(x, y=y, sample_weight=sample_weight)
+        if self.preprocessor is not None:
+            return self.preprocessor(x, y=y, sample_weight=sample_weight)
+        else:
+            return super().preprocess_samples(x, y, sample_weight)
 
     def __setattr__(self, name, value):
-        # Work around torch setattr for properties.
-        if name in ["backbone", "preprocessor"]:
+        # Work around setattr issues for Keras 2 and Keras 3 torch backend.
+        # Since all our state is covered by functional model we can route
+        # around custom setattr calls.
+        is_property = isinstance(getattr(type(self), name, None), property)
+        is_unitialized = not hasattr(self, "_initialized")
+        is_torch = config.backend() == "torch"
+        is_keras_2 = not config.keras_3()
+        if is_torch and (is_property or is_unitialized):
+            return object.__setattr__(self, name, value)
+        if is_keras_2 and is_unitialized:
             return object.__setattr__(self, name, value)
         return super().__setattr__(name, value)
 
     @property
     def backbone(self):
-        """A `keras.Model` instance providing the backbone submodel."""
+        """A `keras.Model` instance providing the backbone sub-model."""
         return self._backbone
 
     @backbone.setter
@@ -123,7 +141,6 @@ class Task(PipelineModel):
 
     @preprocessor.setter
     def preprocessor(self, value):
-        self.include_preprocessing = value is not None
         self._preprocessor = value
 
     def get_config(self):
diff --git a/keras_nlp/models/whisper/whisper_backbone.py b/keras_nlp/models/whisper/whisper_backbone.py
index 32cfab21..2e842190 100644
--- a/keras_nlp/models/whisper/whisper_backbone.py
+++ b/keras_nlp/models/whisper/whisper_backbone.py
@@ -116,75 +116,40 @@ class WhisperBackbone(Backbone):
     ):
         assert_tf_backend(self.__class__.__name__)
 
-        # Encoder inputs. Note that the encoder does not have a padding mask:
-        # https://github.com/openai/whisper/blob/v20230124/whisper/model.py#L132.
-        encoder_feature_input = keras.Input(
-            shape=(None, num_mels), dtype="float32", name="encoder_features"
-        )
-
-        # Decoder inputs.
-        decoder_token_id_input = keras.Input(
-            shape=(None,), dtype="int32", name="decoder_token_ids"
-        )
-        decoder_padding_mask = keras.Input(
-            shape=(None,), dtype="int32", name="decoder_padding_mask"
-        )
-
-        # ====== Encoder ======
-
-        # Embed the input features. This consists of two 1D convolutional
-        # layers.
-        # For the first layer, we use `padding="same"` since that corresponds to
-        # a padding size of 1.
-        encoder_conv_layer_1 = keras.layers.Conv1D(
+        # === Layers ===
+        self.encoder_conv_layer_1 = keras.layers.Conv1D(
             filters=hidden_dim,
             kernel_size=3,
             strides=1,
             padding="same",
             name="encoder_token_embedding_conv_layer_1",
         )
-        embedded_features = keras.activations.gelu(
-            encoder_conv_layer_1(encoder_feature_input),
-            approximate=False,
-        )
-
-        # For the second conv. layer, we cannot use `padding="same"` since
-        # that corresponds to a padding size of 1.5 (since stride is 2). Hence,
-        # we will manually pad the input.
-        embedded_features = Padder()(embedded_features)
-        encoder_conv_layer_2 = keras.layers.Conv1D(
+        self.encoder_conv_layer_2 = keras.layers.Conv1D(
             filters=hidden_dim,
             kernel_size=3,
             strides=2,
             padding="valid",
             name="encoder_token_embedding_conv_layer_2",
         )
-        embedded_features = keras.activations.gelu(
-            encoder_conv_layer_2(embedded_features),
-            approximate=False,
+        self.encoder_padder = Padder(
+            name="encoder_padder",
         )
-
-        # The position embedding layer for the encoder is a sinusoidal embedding
-        # layer: https://github.com/openai/whisper/blob/v20230124/whisper/model.py#L137.
-        # Hence, we set it to be non-trainable.
-        # TODO: We can use `keras_nlp.layers.SinePositionEncoding` layer.
-        position_embedding = PositionEmbedding(
+        self.encoder_position_embedding = PositionEmbedding(
             initializer=whisper_kernel_initializer(),
             sequence_length=max_encoder_sequence_length // 2,
             name="encoder_position_embedding",
             trainable=False,
-        )(embedded_features)
-
-        # Sum and apply dropout to embeddings.
-        x = keras.layers.Add()((embedded_features, position_embedding))
-        x = keras.layers.Dropout(
+        )
+        self.encoder_embeddings_add = keras.layers.Add(
+            name="encoder_embeddings_add",
+        )
+        self.encoder_embeddings_dropout = keras.layers.Dropout(
             dropout,
             name="encoder_embeddings_dropout",
-        )(x)
-
-        # Apply successive transformer encoder blocks.
+        )
+        self.encoder_transformer_layers = []
         for i in range(num_layers):
-            x = WhisperEncoder(
+            layer = WhisperEncoder(
                 num_heads=num_heads,
                 intermediate_dim=intermediate_dim,
                 activation=keras.activations.gelu,
@@ -193,37 +158,29 @@ class WhisperBackbone(Backbone):
                 kernel_initializer=whisper_kernel_initializer(),
                 normalize_first=True,
                 name=f"transformer_encoder_layer_{i}",
-            )(x)
-
-        x = keras.layers.LayerNormalization(
+            )
+            self.encoder_transformer_layers.append(layer)
+        self.encoder_layer_norm = keras.layers.LayerNormalization(
             name="encoder_layer_norm",
             axis=-1,
             epsilon=1e-5,
             dtype="float32",
-        )(x)
-        encoder_output = x
-
-        # ====== Decoder ======
-
-        # Embed tokens and positions.
-        embedding_layer = TokenAndPositionEmbedding(
+        )
+        self.decoder_embeddings = TokenAndPositionEmbedding(
             vocabulary_size=vocabulary_size,
             sequence_length=max_decoder_sequence_length,
             embedding_dim=hidden_dim,
             embeddings_initializer=whisper_kernel_initializer(),
             name="decoder_token_and_position_embedding",
         )
-        x = embedding_layer(decoder_token_id_input)
-
-        # Apply dropout to embeddings.
-        x = keras.layers.Dropout(
+        self.token_embedding = self.decoder_embeddings.token_embedding
+        self.decoder_embeddings_dropout = keras.layers.Dropout(
             dropout,
             name="decoder_embeddings_dropout",
-        )(x)
-
-        # Apply successive transformer decoder blocks.
+        )
+        self.decoder_transformer_layers = []
         for i in range(num_layers):
-            transformer_decoder_layer = WhisperDecoder(
+            layer = WhisperDecoder(
                 intermediate_dim=intermediate_dim,
                 num_heads=num_heads,
                 dropout=dropout,
@@ -233,26 +190,70 @@ class WhisperBackbone(Backbone):
                 normalize_first=True,
                 name=f"transformer_decoder_layer_{i}",
             )
-            x = transformer_decoder_layer(
-                decoder_sequence=x,
-                encoder_sequence=encoder_output,
-                decoder_padding_mask=decoder_padding_mask,
-            )
-
-        x = keras.layers.LayerNormalization(
+            self.decoder_transformer_layers.append(layer)
+        self.decoder_layer_norm = keras.layers.LayerNormalization(
             name="decoder_layer_norm",
             axis=-1,
             epsilon=1e-5,
             dtype="float32",
-        )(x)
-        decoder_output = x
+        )
 
-        # Instantiate using Functional API Model constructor
+        # === Functional Model ===
+        # Note that the encoder does not have a padding mask:
+        # https://github.com/openai/whisper/blob/v20230124/whisper/model.py#L132.
+        encoder_feature_input = keras.Input(
+            shape=(None, num_mels), dtype="float32", name="encoder_features"
+        )
+        decoder_token_id_input = keras.Input(
+            shape=(None,), dtype="int32", name="decoder_token_ids"
+        )
+        decoder_padding_mask_input = keras.Input(
+            shape=(None,), dtype="int32", name="decoder_padding_mask"
+        )
+        # Encoder.
+        # Embed the input features. This consists of two 1D convolutional
+        # layers.
+        # For the first layer, we use `padding="same"` since that corresponds to
+        # a padding size of 1.
+        embedded_features = keras.activations.gelu(
+            self.encoder_conv_layer_1(encoder_feature_input),
+            approximate=False,
+        )
+        # For the second conv. layer, we cannot use `padding="same"` since
+        # that corresponds to a padding size of 1.5 (since stride is 2). Hence,
+        # we will manually pad the input.
+        embedded_features = Padder()(embedded_features)
+        embedded_features = keras.activations.gelu(
+            self.encoder_conv_layer_2(embedded_features),
+            approximate=False,
+        )
+        # The position embedding layer for the encoder is a sinusoidal embedding
+        # layer: https://github.com/openai/whisper/blob/v20230124/whisper/model.py#L137.
+        # Hence, we set it to be non-trainable.
+        # TODO: We can use `keras_nlp.layers.SinePositionEncoding` layer.
+        positions = self.encoder_position_embedding(embedded_features)
+        x = self.encoder_embeddings_add((embedded_features, positions))
+        x = self.encoder_embeddings_dropout(x)
+        for transformer_layer in self.encoder_transformer_layers:
+            x = transformer_layer(x)
+        x = self.encoder_layer_norm(x)
+        encoder_output = x
+        # Decoder.
+        x = self.decoder_embeddings(decoder_token_id_input)
+        x = self.decoder_embeddings_dropout(x)
+        for transformer_layer in self.decoder_transformer_layers:
+            x = transformer_layer(
+                decoder_sequence=x,
+                encoder_sequence=encoder_output,
+                decoder_padding_mask=decoder_padding_mask_input,
+            )
+        x = self.decoder_layer_norm(x)
+        decoder_output = x
         super().__init__(
             inputs={
                 "encoder_features": encoder_feature_input,
                 "decoder_token_ids": decoder_token_id_input,
-                "decoder_padding_mask": decoder_padding_mask,
+                "decoder_padding_mask": decoder_padding_mask_input,
             },
             outputs={
                 "encoder_sequence_output": encoder_output,
@@ -261,7 +262,7 @@ class WhisperBackbone(Backbone):
             **kwargs,
         )
 
-        # All references to `self` below this line
+        # === Config ===
         self.vocabulary_size = vocabulary_size
         self.num_layers = num_layers
         self.num_heads = num_heads
@@ -271,7 +272,6 @@ class WhisperBackbone(Backbone):
         self.dropout = dropout
         self.max_encoder_sequence_length = max_encoder_sequence_length
         self.max_decoder_sequence_length = max_decoder_sequence_length
-        self.token_embedding = embedding_layer
 
     def get_config(self):
         config = super().get_config()
diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_classifier.py b/keras_nlp/models/xlm_roberta/xlm_roberta_classifier.py
index 67a9dd5b..45d79eb3 100644
--- a/keras_nlp/models/xlm_roberta/xlm_roberta_classifier.py
+++ b/keras_nlp/models/xlm_roberta/xlm_roberta_classifier.py
@@ -157,30 +157,45 @@ class XLMRobertaClassifier(Task):
         dropout=0.0,
         **kwargs,
     ):
-        inputs = backbone.input
+        # === Layers ===
+        self.backbone = backbone
+        self.preprocessor = preprocessor
+        self.pooled_dropout = keras.layers.Dropout(
+            dropout,
+            name="pooled_dropout",
+        )
         hidden_dim = hidden_dim or backbone.hidden_dim
-
-        x = backbone(inputs)[:, backbone.start_token_index, :]
-        x = keras.layers.Dropout(dropout, name="pooled_dropout")(x)
-        x = keras.layers.Dense(
-            hidden_dim, activation="tanh", name="pooled_dense"
-        )(x)
-        x = keras.layers.Dropout(dropout, name="classifier_dropout")(x)
-        outputs = keras.layers.Dense(
+        self.pooled_dense = keras.layers.Dense(
+            hidden_dim,
+            activation="tanh",
+            name="pooled_dense",
+        )
+        self.output_dropout = keras.layers.Dropout(
+            dropout,
+            name="output_dropout",
+        )
+        self.output_dense = keras.layers.Dense(
             num_classes,
             kernel_initializer=roberta_kernel_initializer(),
             activation=activation,
             name="logits",
-        )(x)
+        )
 
+        # === Functional Model ===
+        inputs = backbone.input
+        x = backbone(inputs)[:, backbone.start_token_index, :]
+        x = self.pooled_dropout(x)
+        x = self.pooled_dense(x)
+        x = self.output_dropout(x)
+        outputs = self.output_dense(x)
         # Instantiate using Functional API Model constructor
         super().__init__(
             inputs=inputs,
             outputs=outputs,
-            include_preprocessing=preprocessor is not None,
             **kwargs,
         )
-        # All references to `self` below this line
+
+        # === Config ===
         self.backbone = backbone
         self.preprocessor = preprocessor
         self.num_classes = num_classes
@@ -188,6 +203,7 @@ class XLMRobertaClassifier(Task):
         self.hidden_dim = hidden_dim
         self.dropout = dropout
 
+        # === Default compilation ===
         logit_output = self.activation == keras.activations.linear
         self.compile(
             loss=keras.losses.SparseCategoricalCrossentropy(
@@ -198,9 +214,6 @@ class XLMRobertaClassifier(Task):
             jit_compile=True,
         )
 
-    def preprocess_samples(self, x, y=None, sample_weight=None):
-        return self.preprocessor(x, y=y, sample_weight=sample_weight)
-
     def get_config(self):
         config = super().get_config()
         config.update(
diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_masked_lm.py b/keras_nlp/models/xlm_roberta/xlm_roberta_masked_lm.py
index f0dfc85e..b29aa30d 100644
--- a/keras_nlp/models/xlm_roberta/xlm_roberta_masked_lm.py
+++ b/keras_nlp/models/xlm_roberta/xlm_roberta_masked_lm.py
@@ -106,6 +106,18 @@ class XLMRobertaMaskedLM(Task):
         preprocessor=None,
         **kwargs,
     ):
+        # === Layers ===
+        self.backbone = backbone
+        self.preprocessor = preprocessor
+        self.masked_lm_head = MaskedLMHead(
+            vocabulary_size=backbone.vocabulary_size,
+            token_embedding=backbone.token_embedding,
+            intermediate_activation="gelu",
+            kernel_initializer=roberta_kernel_initializer(),
+            name="mlm_head",
+        )
+
+        # === Functional Model ===
         inputs = {
             **backbone.input,
             "mask_positions": keras.Input(
@@ -113,25 +125,16 @@ class XLMRobertaMaskedLM(Task):
             ),
         }
         backbone_outputs = backbone(backbone.input)
-        outputs = MaskedLMHead(
-            vocabulary_size=backbone.vocabulary_size,
-            token_embedding=backbone.token_embedding,
-            intermediate_activation="gelu",
-            kernel_initializer=roberta_kernel_initializer(),
-            name="mlm_head",
-        )(backbone_outputs, inputs["mask_positions"])
-
-        # Instantiate using Functional API Model constructor.
+        outputs = self.masked_lm_head(
+            backbone_outputs, inputs["mask_positions"]
+        )
         super().__init__(
             inputs=inputs,
             outputs=outputs,
-            include_preprocessing=preprocessor is not None,
             **kwargs,
         )
-        # All references to `self` below this line
-        self.backbone = backbone
-        self.preprocessor = preprocessor
 
+        # === Default compilation ===
         self.compile(
             loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
             optimizer=keras.optimizers.Adam(5e-5),
diff --git a/keras_nlp/models/xlnet/xlnet_backbone.py b/keras_nlp/models/xlnet/xlnet_backbone.py
index 1d1b4d23..45a4b8f4 100644
--- a/keras_nlp/models/xlnet/xlnet_backbone.py
+++ b/keras_nlp/models/xlnet/xlnet_backbone.py
@@ -103,42 +103,25 @@ class XLNetBackbone(Backbone):
         bias_initializer="zeros",
         **kwargs,
     ):
-        # Inputs
-        token_id_input = keras.Input(
-            shape=(None,), dtype="int32", name="token_ids"
-        )
-        padding_mask = keras.Input(
-            shape=(None,), dtype="int32", name="padding_mask"
-        )
-        segment_ids = keras.Input(
-            shape=(None,), dtype="int32", name="segment_ids"
-        )
-
-        # Content and Query Embedding
-        word_emb, pos_emb = ContentAndQueryEmbedding(
+        # === Layers ===
+        self.content_query_embedding = ContentAndQueryEmbedding(
             vocabulary_size=vocabulary_size,
             hidden_dim=hidden_dim,
             dropout=dropout,
             name="content_query_embedding",
-        )(token_id_input=token_id_input)
-
-        # Apply XLNetAttentionMaskLayer and XLNetSegmentMatrixLayer Layers
-        # to get the processed attention masks and segment matrix.
-        attn_mask_content, attn_mask_query = XLNetAttentionMaskLayer(
+        )
+        self.attn_mask_layer = XLNetAttentionMaskLayer(
             hidden_dim=hidden_dim,
             kernel_initializer_range=kernel_initializer_range,
             name="encoder_block_attn_mask_layer",
-        )(padding_mask)
-        seg_mat = XLNetSegmentMatrixLayer(name="encoder_block_seg_mat_layer")(
-            segment_ids
         )
-
-        output_content = word_emb
-
-        # Encoders
+        self.seg_mat_layer = XLNetSegmentMatrixLayer(
+            name="encoder_block_seg_mat_layer",
+        )
         head_dim = hidden_dim // num_heads
+        self.transformer_layers = []
         for i in range(num_layers):
-            output_content, output_query = XLNetEncoder(
+            layer = XLNetEncoder(
                 num_heads=num_heads,
                 hidden_dim=hidden_dim,
                 head_dim=head_dim,
@@ -149,27 +132,52 @@ class XLNetBackbone(Backbone):
                 kernel_initializer_range=kernel_initializer_range,
                 bias_initializer=bias_initializer,
                 name=f"xlnet_encoder_{i}",
-            )(
+            )
+            self.transformer_layers.append(layer)
+        self.dropout = keras.layers.Dropout(
+            dropout,
+            name="dropout",
+        )
+
+        # === Functional Model ===
+        token_id_input = keras.Input(
+            shape=(None,), dtype="int32", name="token_ids"
+        )
+        padding_mask_input = keras.Input(
+            shape=(None,), dtype="int32", name="padding_mask"
+        )
+        segment_id_input = keras.Input(
+            shape=(None,), dtype="int32", name="segment_ids"
+        )
+        # Content and Query Embedding
+        word_emb, pos_emb = self.content_query_embedding(token_id_input)
+        # Apply XLNetAttentionMaskLayer and XLNetSegmentMatrixLayer Layers
+        # to get the processed attention masks and segment matrix.
+        attn_mask_content, attn_mask_query = self.attn_mask_layer(
+            padding_mask_input
+        )
+        seg_mat = self.seg_mat_layer(segment_id_input)
+        output_content = word_emb
+        for transformer_layer in self.transformer_layers:
+            output_content, output_query = transformer_layer(
                 output_content=output_content,
                 attn_mask_content=attn_mask_content,
                 attn_mask_query=attn_mask_query,
                 pos_emb=pos_emb,
                 seg_mat=seg_mat,
             )
-
-        output = keras.layers.Dropout(dropout)(output_content)
-
+        output = self.dropout(output_content)
         super().__init__(
             inputs={
                 "token_ids": token_id_input,
-                "padding_mask": padding_mask,
-                "segment_ids": segment_ids,
+                "padding_mask": padding_mask_input,
+                "segment_ids": segment_id_input,
             },
             outputs=output,
             **kwargs,
         )
 
-        # All references to `self` below this line
+        # === Config ===
         self.vocabulary_size = vocabulary_size
         self.num_layers = num_layers
         self.num_heads = num_heads
diff --git a/keras_nlp/utils/pipeline_model.py b/keras_nlp/utils/pipeline_model.py
index fa08aaf9..89a2f818 100644
--- a/keras_nlp/utils/pipeline_model.py
+++ b/keras_nlp/utils/pipeline_model.py
@@ -142,27 +142,15 @@ def _train_validation_split(arrays, validation_split):
 class PipelineModel(keras.Model):
     """A model which allows automatically applying preprocessing."""
 
-    def __init__(self, *args, include_preprocessing=True, **kwargs):
+    def __init__(self, *args, **kwargs):
         # Workaround for https://github.com/keras-team/keras/issues/17270
         # Reset any attempt to overwrite this classes base class to this class
         # can continue to be used for functional and non-functional models.
         PipelineModel.__bases__ = (keras.Model,)
         super().__init__(*args, **kwargs)
-        self.include_preprocessing = include_preprocessing
-
-    def preprocess_features(self, x):
-        """An overridable function which preprocesses features."""
-        return x
-
-    def preprocess_labels(self, y):
-        """An overridable function which preprocesses labels."""
-        return y
 
     def preprocess_samples(self, x, y=None, sample_weight=None):
         """An overridable function which preprocesses entire samples."""
-        x = self.preprocess_features(x)
-        if y is not None:
-            y = self.preprocess_labels(y)
         return pack_x_y_sample_weight(x, y, sample_weight)
 
     # ========================================================================
@@ -184,10 +172,9 @@ class PipelineModel(keras.Model):
             )
 
         x = _convert_inputs_to_dataset(x, y, sample_weight, batch_size)
-        if self.include_preprocessing:
-            x = x.map(
-                self.preprocess_samples, num_parallel_calls=tf.data.AUTOTUNE
-            ).prefetch(tf.data.AUTOTUNE)
+        x = x.map(
+            self.preprocess_samples, num_parallel_calls=tf.data.AUTOTUNE
+        ).prefetch(tf.data.AUTOTUNE)
 
         if validation_data is not None:
             if not isinstance(validation_data, tf.data.Dataset):
@@ -221,10 +208,9 @@ class PipelineModel(keras.Model):
         # needs preprocessing.
         kwargs.pop("_use_cached_eval_dataset", None)
         x = _convert_inputs_to_dataset(x, y, sample_weight, batch_size)
-        if self.include_preprocessing:
-            x = x.map(
-                self.preprocess_samples, num_parallel_calls=tf.data.AUTOTUNE
-            ).prefetch(tf.data.AUTOTUNE)
+        x = x.map(
+            self.preprocess_samples, num_parallel_calls=tf.data.AUTOTUNE
+        ).prefetch(tf.data.AUTOTUNE)
         return super().evaluate(
             x=x,
             y=None,
@@ -239,11 +225,9 @@ class PipelineModel(keras.Model):
         **kwargs,
     ):
         x = _convert_inputs_to_dataset(x, None, None, batch_size)
-        if self.include_preprocessing:
-            x = x.map(
-                self.preprocess_samples, num_parallel_calls=tf.data.AUTOTUNE
-            ).prefetch(tf.data.AUTOTUNE)
-
+        x = x.map(
+            self.preprocess_samples, num_parallel_calls=tf.data.AUTOTUNE
+        ).prefetch(tf.data.AUTOTUNE)
         return super().predict(
             x=x,
             batch_size=None,
@@ -257,14 +241,13 @@ class PipelineModel(keras.Model):
         sample_weight=None,
         **kwargs,
     ):
-        if self.include_preprocessing:
-            data = self.preprocess_samples(x, y, sample_weight)
-            x, y, sample_weight = keras.utils.unpack_x_y_sample_weight(data)
-            x = ops.convert_to_tensor(x)
-            if y is not None:
-                y = ops.convert_to_tensor(y)
-            if sample_weight is not None:
-                sample_weight = ops.convert_to_tensor(sample_weight)
+        data = self.preprocess_samples(x, y, sample_weight)
+        x, y, sample_weight = keras.utils.unpack_x_y_sample_weight(data)
+        x = ops.convert_to_tensor(x)
+        if y is not None:
+            y = ops.convert_to_tensor(y)
+        if sample_weight is not None:
+            sample_weight = ops.convert_to_tensor(sample_weight)
         return super().train_on_batch(
             x=x,
             y=y,
@@ -279,14 +262,13 @@ class PipelineModel(keras.Model):
         sample_weight=None,
         **kwargs,
     ):
-        if self.include_preprocessing:
-            data = self.preprocess_samples(x, y, sample_weight)
-            x, y, sample_weight = keras.utils.unpack_x_y_sample_weight(data)
-            x = ops.convert_to_tensor(x)
-            if y is not None:
-                y = ops.convert_to_tensor(y)
-            if sample_weight is not None:
-                sample_weight = ops.convert_to_tensor(sample_weight)
+        data = self.preprocess_samples(x, y, sample_weight)
+        x, y, sample_weight = keras.utils.unpack_x_y_sample_weight(data)
+        x = ops.convert_to_tensor(x)
+        if y is not None:
+            y = ops.convert_to_tensor(y)
+        if sample_weight is not None:
+            sample_weight = ops.convert_to_tensor(sample_weight)
         return super().test_on_batch(
             x=x,
             y=y,
@@ -299,10 +281,9 @@ class PipelineModel(keras.Model):
         x,
         **kwargs,
     ):
-        if self.include_preprocessing:
-            data = self.preprocess_samples(x)
-            x, _, _ = keras.utils.unpack_x_y_sample_weight(data)
-            x = ops.convert_to_tensor(x)
+        data = self.preprocess_samples(x)
+        x, _, _ = keras.utils.unpack_x_y_sample_weight(data)
+        x = ops.convert_to_tensor(x)
         return super().predict_on_batch(
             x=x,
             **kwargs,


Test patch:
diff --git a/keras_nlp/models/bloom/bloom_causal_lm_preprocessor_test.py b/keras_nlp/models/bloom/bloom_causal_lm_preprocessor_test.py
new file mode 100644
index 00000000..6caf8fdd
--- /dev/null
+++ b/keras_nlp/models/bloom/bloom_causal_lm_preprocessor_test.py
@@ -0,0 +1,94 @@
+# Copyright 2023 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 pytest
+
+from keras_nlp.models.bloom.bloom_causal_lm_preprocessor import (
+    BloomCausalLMPreprocessor,
+)
+from keras_nlp.models.bloom.bloom_tokenizer import BloomTokenizer
+from keras_nlp.tests.test_case import TestCase
+
+
+class BloomCausalLMPreprocessorTest(TestCase):
+    def setUp(self):
+        self.vocab = ["<pad>", "
", "
"]
+        self.vocab += ["!", "air", "Ġair", "plane", "Ġat", "port"]
+        self.vocab = dict([(token, i) for i, token in enumerate(self.vocab)])
+        self.merges = ["Ġ a", "Ġ t", "Ġ i", "Ġ b", "a i", "p l", "n e"]
+        self.merges += ["Ġa t", "p o", "r t", "Ġt h", "ai r", "pl a", "po rt"]
+        self.merges += ["Ġai r", "Ġa i", "pla ne"]
+        self.tokenizer = BloomTokenizer(
+            vocabulary=self.vocab,
+            merges=self.merges,
+        )
+        self.init_kwargs = {
+            "tokenizer": self.tokenizer,
+            "sequence_length": 8,
+        }
+        self.input_data = ["airplane at airport"]
+
+    def test_preprocessor_basics(self):
+        self.run_preprocessing_layer_test(
+            cls=BloomCausalLMPreprocessor,
+            init_kwargs=self.init_kwargs,
+            input_data=self.input_data,
+            expected_output=(
+                {
+                    "token_ids": [[1, 4, 6, 7, 5, 8, 2, 0]],
+                    "padding_mask": [[1, 1, 1, 1, 1, 1, 1, 0]],
+                },
+                [[4, 6, 7, 5, 8, 2, 0, 0]],  # Pass through labels.
+                [[1, 1, 1, 1, 1, 1, 0, 0]],  # Pass through sample_weights.
+            ),
+        )
+
+    def test_no_start_end_token(self):
+        input_data = ["airplane at airport"] * 4
+
+        preprocessor = BloomCausalLMPreprocessor(
+            **self.init_kwargs,
+            add_start_token=False,
+            add_end_token=False,
+        )
+        x, y, sw = preprocessor(input_data)
+        self.assertAllEqual(x["token_ids"], [[4, 6, 7, 5, 8, 0, 0, 0]] * 4)
+        self.assertAllEqual(x["padding_mask"], [[1, 1, 1, 1, 1, 0, 0, 0]] * 4)
+        self.assertAllEqual(y, [[6, 7, 5, 8, 0, 0, 0, 0]] * 4)
+        self.assertAllEqual(sw, [[1, 1, 1, 1, 0, 0, 0, 0]] * 4)
+
+    def test_generate_preprocess(self):
+        input_data = "airplane at airport"
+        preprocessor = BloomCausalLMPreprocessor(**self.init_kwargs)
+        x = preprocessor.generate_preprocess(input_data)
+        self.assertAllEqual(x["token_ids"], [1, 4, 6, 7, 5, 8, 0, 0])
+        self.assertAllEqual(x["padding_mask"], [1, 1, 1, 1, 1, 1, 0, 0])
+
+    def test_generate_postprocess(self):
+        input_data = {
+            "token_ids": [1, 4, 6, 7, 5, 8, 0, 0],
+            "padding_mask": [1, 1, 1, 1, 1, 1, 0, 0],
+        }
+        preprocessor = BloomCausalLMPreprocessor(**self.init_kwargs)
+        x = preprocessor.generate_postprocess(input_data)
+        self.assertAllEqual(x, "airplane at airport")
+
+    @pytest.mark.extra_large
+    def test_all_presets(self):
+        for preset in BloomCausalLMPreprocessor.presets:
+            self.run_preset_test(
+                cls=BloomCausalLMPreprocessor,
+                preset=preset,
+                input_data=self.input_data,
+            )
diff --git a/keras_nlp/models/bloom/bloom_preprocessor_test.py b/keras_nlp/models/bloom/bloom_preprocessor_test.py
new file mode 100644
index 00000000..bb800063
--- /dev/null
+++ b/keras_nlp/models/bloom/bloom_preprocessor_test.py
@@ -0,0 +1,80 @@
+# Copyright 2023 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 pytest
+
+from keras_nlp.models.bloom.bloom_preprocessor import BloomPreprocessor
+from keras_nlp.models.bloom.bloom_tokenizer import BloomTokenizer
+from keras_nlp.tests.test_case import TestCase
+
+
+class BloomPreprocessorTest(TestCase):
+    def setUp(self):
+        self.vocab = ["<pad>", "
", "
"]
+        self.vocab += ["!", "air", "Ġair", "plane", "Ġat", "port"]
+        self.vocab = dict([(token, i) for i, token in enumerate(self.vocab)])
+        self.merges = ["Ġ a", "Ġ t", "Ġ i", "Ġ b", "a i", "p l", "n e"]
+        self.merges += ["Ġa t", "p o", "r t", "Ġt h", "ai r", "pl a", "po rt"]
+        self.merges += ["Ġai r", "Ġa i", "pla ne"]
+        self.tokenizer = BloomTokenizer(
+            vocabulary=self.vocab,
+            merges=self.merges,
+        )
+        self.init_kwargs = {
+            "tokenizer": self.tokenizer,
+            "sequence_length": 8,
+        }
+        self.input_data = ["airplane at airport"]
+
+    def test_preprocessor_basics(self):
+        self.run_preprocessing_layer_test(
+            cls=BloomPreprocessor,
+            init_kwargs=self.init_kwargs,
+            input_data=self.input_data,
+            expected_output={
+                "token_ids": [[1, 4, 6, 7, 5, 8, 2, 0]],
+                "padding_mask": [[1, 1, 1, 1, 1, 1, 1, 0]],
+            },
+        )
+
+    def test_no_start_end_token(self):
+        input_data = ["airplane at airport"] * 4
+
+        preprocessor = BloomPreprocessor(
+            tokenizer=BloomTokenizer(
+                vocabulary=self.vocab,
+                merges=self.merges,
+            ),
+            sequence_length=8,
+            add_start_token=False,
+            add_end_token=False,
+        )
+        x = preprocessor(input_data)
+        self.assertAllEqual(x["token_ids"], [[4, 6, 7, 5, 8, 0, 0, 0]] * 4)
+        self.assertAllEqual(x["padding_mask"], [[1, 1, 1, 1, 1, 0, 0, 0]] * 4)
+
+    def test_sequence_length_override(self):
+        input_data = "airplane at airport"
+        preprocessor = BloomPreprocessor(**self.init_kwargs)
+        x = preprocessor(input_data, sequence_length=4)
+        self.assertAllEqual(x["token_ids"], [1, 4, 6, 2])
+
+    @pytest.mark.extra_large
+    def test_all_presets(self):
+        for preset in BloomPreprocessor.presets:
+            self.run_preset_test(
+                cls=BloomPreprocessor,
+                preset=preset,
+                input_data=self.input_data,
+            )
diff --git a/keras_nlp/models/task_test.py b/keras_nlp/models/task_test.py
index 09fe1b00..bf82e4fa 100644
--- a/keras_nlp/models/task_test.py
+++ b/keras_nlp/models/task_test.py
@@ -32,11 +32,11 @@ class SimplePreprocessor(Preprocessor):
 
 class SimpleTask(Task):
     def __init__(self, preprocessor=None, activation=None, **kwargs):
+        self.preprocessor = preprocessor
+        self.activation = keras.activations.get(activation)
         inputs = keras.Input(shape=(5,))
         outputs = keras.layers.Dense(5)(inputs)
         super().__init__(inputs, outputs, **kwargs)
-        self.preprocessor = preprocessor
-        self.activation = keras.activations.get(activation)
 
 
 class TestTask(TestCase):
diff --git a/keras_nlp/utils/pipeline_model_test.py b/keras_nlp/utils/pipeline_model_test.py
index 4c7c7f19..ae71f9f5 100644
--- a/keras_nlp/utils/pipeline_model_test.py
+++ b/keras_nlp/utils/pipeline_model_test.py
@@ -36,8 +36,9 @@ class FeaturePipeline(PipelineModel):
         super().__init__(**kwargs)
         self.dense = keras.layers.Dense(1)
 
-    def preprocess_features(self, x):
-        return tf.strings.to_number(x)
+    def preprocess_samples(self, x, y=None, sample_weight=None):
+        x = tf.strings.to_number(x)
+        return keras.utils.pack_x_y_sample_weight(x, y, sample_weight)
 
     def call(self, inputs):
         return self.dense(inputs)
@@ -48,8 +49,10 @@ class LabelPipeline(PipelineModel):
         super().__init__(**kwargs)
         self.dense = keras.layers.Dense(1)
 
-    def preprocess_labels(self, y):
-        return tf.strings.to_number(y)
+    def preprocess_samples(self, x, y=None, sample_weight=None):
+        if y is not None:
+            y = tf.strings.to_number(y)
+        return keras.utils.pack_x_y_sample_weight(x, y, sample_weight)
 
     def call(self, inputs):
         return self.dense(inputs)
@@ -63,8 +66,7 @@ class DataPipeline(PipelineModel):
         self.dense = keras.layers.Dense(1)
 
     def preprocess_samples(self, x, y=None, sample_weight=None):
-        x = tf.strings.to_number(x)
-        y = x
+        y = x = tf.strings.to_number(x)
         return keras.utils.pack_x_y_sample_weight(x, y, sample_weight)
 
     def call(self, inputs):
@@ -77,8 +79,9 @@ class FunctionalPipeline(PipelineModel):
         outputs = keras.layers.Dense(1)(inputs)
         super().__init__(inputs, outputs, **kwargs)
 
-    def preprocess_features(self, x):
-        return tf.strings.to_number(x)
+    def preprocess_samples(self, x, y=None, sample_weight=None):
+        x = tf.strings.to_number(x)
+        return keras.utils.pack_x_y_sample_weight(x, y, sample_weight)
 
     def get_config(self):
         return {}
@@ -167,19 +170,6 @@ class TestFeaturePreprocessingModel(TestCase):
         model.fit(x=x, y=y, batch_size=8)
         model.fit(tf.data.Dataset.from_tensor_slices((x, y)).batch(8))
 
-    def test_fit_no_preprocessing(self):
-        x = np.random.uniform(size=(100, 5))
-        y = np.random.uniform(size=(100, 1))
-        sw = np.random.uniform(size=(100, 1))
-        model = FeaturePipeline(include_preprocessing=False)
-        model.compile(loss="mse")
-        # With sample weight.
-        model.fit(x=x, y=y, sample_weight=sw, batch_size=8)
-        model.fit(tf.data.Dataset.from_tensor_slices((x, y, sw)).batch(8))
-        # Without sample weight.
-        model.fit(x=x, y=y, batch_size=8)
-        model.fit(tf.data.Dataset.from_tensor_slices((x, y)).batch(8))
-
     def test_evaluate_with_preprocessing(self):
         x = tf.strings.as_string(np.random.uniform(size=(100, 5)))
         y = np.random.uniform(size=(100, 1))
@@ -193,19 +183,6 @@ class TestFeaturePreprocessingModel(TestCase):
         model.evaluate(x=x, y=y, batch_size=8)
         model.evaluate(tf.data.Dataset.from_tensor_slices((x, y)).batch(8))
 
-    def test_evaluate_no_preprocessing(self):
-        x = np.random.uniform(size=(100, 5))
-        y = np.random.uniform(size=(100, 1))
-        sw = np.random.uniform(size=(100, 1))
-        model = FeaturePipeline(include_preprocessing=False)
-        model.compile(loss="mse")
-        # With sample weight.
-        model.evaluate(x=x, y=y, sample_weight=sw, batch_size=8)
-        model.evaluate(tf.data.Dataset.from_tensor_slices((x, y, sw)).batch(8))
-        # Without sample weight.
-        model.evaluate(x=x, y=y, batch_size=8)
-        model.evaluate(tf.data.Dataset.from_tensor_slices((x, y)).batch(8))
-
     def test_predict_with_preprocessing(self):
         x = tf.strings.as_string(np.random.uniform(size=(100, 5)))
         model = FeaturePipeline()
@@ -213,13 +190,6 @@ class TestFeaturePreprocessingModel(TestCase):
         model.predict(x=x, batch_size=8)
         model.predict(tf.data.Dataset.from_tensor_slices(x).batch(8))
 
-    def test_predict_no_preprocessing(self):
-        x = np.random.uniform(size=(100, 5))
-        model = FeaturePipeline(include_preprocessing=False)
-        model.compile(loss="mse")
-        model.predict(x=x, batch_size=8)
-        model.predict(tf.data.Dataset.from_tensor_slices(x).batch(8))
-
     def test_on_batch(self):
         x = tf.strings.as_string(np.random.uniform(size=(8, 5)))
         y = np.random.uniform(size=(8, 1))
@@ -234,19 +204,6 @@ class TestFeaturePreprocessingModel(TestCase):
         model.test_on_batch(x=x, y=y)
         model.predict_on_batch(x=x)
 
-    def test_on_batch_no_preprocessing(self):
-        x = np.random.uniform(size=(8, 5))
-        y = np.random.uniform(size=(8, 1))
-        sw = np.random.uniform(size=(8, 1))
-        model = FeaturePipeline(include_preprocessing=False)
-        model.compile(loss="mse")
-        # With sample weight.
-        model.train_on_batch(x=x, y=y, sample_weight=sw)
-        model.test_on_batch(x=x, y=y, sample_weight=sw)
-        # Without sample weight.
-        model.train_on_batch(x=x, y=y)
-        model.test_on_batch(x=x, y=y)
-
     def test_saved_model(self):
         model = FeaturePipeline()
         x = tf.strings.as_string(np.random.uniform(size=(8, 5)))
@@ -278,19 +235,6 @@ class TestLabelPreprocessingModel(TestCase):
         model.fit(x=x, y=y, batch_size=8)
         model.fit(tf.data.Dataset.from_tensor_slices((x, y)).batch(8))
 
-    def test_fit_no_preprocessing(self):
-        x = np.random.uniform(size=(100, 5))
-        y = np.random.uniform(size=(100, 1))
-        sw = np.random.uniform(size=(100, 1))
-        model = LabelPipeline(include_preprocessing=False)
-        model.compile(loss="mse")
-        # With sample weight.
-        model.fit(x=x, y=y, sample_weight=sw, batch_size=8)
-        model.fit(tf.data.Dataset.from_tensor_slices((x, y, sw)).batch(8))
-        # Without sample weight.
-        model.fit(x=x, y=y, batch_size=8)
-        model.fit(tf.data.Dataset.from_tensor_slices((x, y)).batch(8))
-
     def test_evaluate_with_preprocessing(self):
         x = np.random.uniform(size=(100, 5))
         y = tf.strings.as_string(np.random.uniform(size=(100, 1)))
@@ -304,19 +248,6 @@ class TestLabelPreprocessingModel(TestCase):
         model.evaluate(x=x, y=y, batch_size=8)
         model.evaluate(tf.data.Dataset.from_tensor_slices((x, y)).batch(8))
 
-    def test_evaluate_no_preprocessing(self):
-        x = np.random.uniform(size=(100, 5))
-        y = np.random.uniform(size=(100, 1))
-        sw = np.random.uniform(size=(100, 1))
-        model = LabelPipeline(include_preprocessing=False)
-        model.compile(loss="mse")
-        # With sample weight.
-        model.evaluate(x=x, y=y, sample_weight=sw, batch_size=8)
-        model.evaluate(tf.data.Dataset.from_tensor_slices((x, y, sw)).batch(8))
-        # Without sample weight.
-        model.evaluate(x=x, y=y, batch_size=8)
-        model.evaluate(tf.data.Dataset.from_tensor_slices((x, y)).batch(8))
-
     def test_predict_with_preprocessing(self):
         x = np.random.uniform(size=(100, 5))
         model = LabelPipeline()
@@ -338,20 +269,6 @@ class TestLabelPreprocessingModel(TestCase):
         model.test_on_batch(x=x, y=y)
         model.predict_on_batch(x=x)
 
-    def test_on_batch_no_preprocessing(self):
-        x = np.random.uniform(size=(8, 5))
-        y = np.random.uniform(size=(8, 1))
-        sw = np.random.uniform(size=(8, 1))
-        model = LabelPipeline(include_preprocessing=False)
-        model.compile(loss="mse")
-        # With sample weight.
-        model.train_on_batch(x=x, y=y, sample_weight=sw)
-        model.test_on_batch(x=x, y=y, sample_weight=sw)
-        # Without sample weight.
-        model.train_on_batch(x=x, y=y)
-        model.test_on_batch(x=x, y=y)
-        model.predict_on_batch(x=x)
-
     def test_saved_model(self):
         model = LabelPipeline()
         x = np.random.uniform(size=(8, 5))
@@ -377,14 +294,6 @@ class TestDataPreprocessingModel(TestCase):
         model.fit(x=data, batch_size=8)
         model.fit(tf.data.Dataset.from_tensor_slices(data).batch(8))
 
-    def test_fit_no_preprocessing(self):
-        x = np.random.uniform(size=(100, 1))
-        y = np.random.uniform(size=(100, 1))
-        model = DataPipeline(include_preprocessing=False)
-        model.compile(loss="mse")
-        model.fit(x=x, y=y, batch_size=8)
-        model.fit(tf.data.Dataset.from_tensor_slices((x, y)).batch(8))
-
     def test_evaluate_with_preprocessing(self):
         data = tf.strings.as_string(np.random.uniform(size=(100, 1)))
         model = DataPipeline()
@@ -392,14 +301,6 @@ class TestDataPreprocessingModel(TestCase):
         model.evaluate(x=data, batch_size=8)
         model.evaluate(tf.data.Dataset.from_tensor_slices(data).batch(8))
 
-    def test_evaluate_no_preprocessing(self):
-        x = np.random.uniform(size=(100, 1))
-        y = np.random.uniform(size=(100, 1))
-        model = DataPipeline(include_preprocessing=False)
-        model.compile(loss="mse")
-        model.evaluate(x=x, y=y, batch_size=8)
-        model.evaluate(tf.data.Dataset.from_tensor_slices((x, y)).batch(8))
-
     def test_predict_with_preprocessing(self):
         x = tf.strings.as_string(np.random.uniform(size=(100, 1)))
         model = DataPipeline()
@@ -407,13 +308,6 @@ class TestDataPreprocessingModel(TestCase):
         model.predict(x=x, batch_size=8)
         model.predict(tf.data.Dataset.from_tensor_slices(x).batch(8))
 
-    def test_predict_no_preprocessing(self):
-        x = np.random.uniform(size=(100, 1))
-        model = DataPipeline(include_preprocessing=False)
-        model.compile(loss="mse")
-        model.predict(x=x, batch_size=8)
-        model.predict(tf.data.Dataset.from_tensor_slices(x).batch(8))
-
     def test_on_batch(self):
         data = tf.strings.as_string(np.random.uniform(size=(8, 1)))
         model = DataPipeline()
@@ -426,20 +320,6 @@ class TestDataPreprocessingModel(TestCase):
         model.test_on_batch(x=data)
         model.predict_on_batch(x=data)
 
-    def test_on_batch_no_preprocessing(self):
-        x = np.random.uniform(size=(8, 1))
-        y = np.random.uniform(size=(8, 1))
-        sw = np.random.uniform(size=(8, 1))
-        model = DataPipeline(include_preprocessing=False)
-        model.compile(loss="mse")
-        # With sample weight.
-        model.train_on_batch(x=x, y=y, sample_weight=sw)
-        model.test_on_batch(x=x, y=y, sample_weight=sw)
-        # Without sample weight.
-        model.train_on_batch(x=x, y=y)
-        model.test_on_batch(x=x, y=y)
-        model.predict_on_batch(x=x)
-
     def test_saved_model(self):
         model = DataPipeline()
         data = tf.strings.as_string(np.random.uniform(size=(8, 1)))
@@ -472,19 +352,6 @@ class TestFunctional(TestCase):
         model.fit(x=x, y=y, batch_size=8)
         model.fit(tf.data.Dataset.from_tensor_slices((x, y)).batch(8))
 
-    def test_fit_no_preprocessing(self):
-        x = np.random.uniform(size=(100, 5))
-        y = np.random.uniform(size=(100, 1))
-        sw = np.random.uniform(size=(100, 1))
-        model = FunctionalPipeline(include_preprocessing=False)
-        model.compile(loss="mse")
-        # With sample weight.
-        model.fit(x=x, y=y, sample_weight=sw, batch_size=8)
-        model.fit(tf.data.Dataset.from_tensor_slices((x, y, sw)).batch(8))
-        # Without sample weight.
-        model.fit(x=x, y=y, batch_size=8)
-        model.fit(tf.data.Dataset.from_tensor_slices((x, y)).batch(8))
-
     def test_saved_model(self):
         model = FunctionalPipeline()
         x = tf.strings.as_string(np.random.uniform(size=(8, 5)))