Repository: keras-team/keras-nlp
Pull request #646 (2023-01-10 18:50:57)

Add FNet Preprocessor
https://colab.research.google.com/drive/1wd7ApfnTwS_xqa62Isdx3CGIZ95ixwxC?usp=sharing


Patch:
diff --git a/examples/README.md b/examples/README.md
index 7acc255c..e7426169 100644
--- a/examples/README.md
+++ b/examples/README.md
@@ -1,60 +1,5 @@
-# KerasNLP Example Models
+# KerasNLP Examples
 
-This examples directory contains complete recipes for training popular model
-architectures using KerasNLP. These are not part of the library itself, but
-rather serve to demonstrate how to use the library for common tasks, while
-simultaneously providing a mechanism to rigorously test library components.
-
-This directory is complementary to the NLP examples on
-[keras.io](https://keras.io/examples/). If you want to contribute a KerasNLP
-example and you're not sure if it should live on keras.io or in this directory,
-here's how they differ:
-
-- If an example can fit in < 300 lines of code and run in a Colab,
-  put it on keras.io.
-- If an example is too big for a single script or has high compute requirements,
-  add it here.
-
-In general, we will have a fairly high bar for what models we support in this
-directory. They should be widely used, practical models for solving standard
-problems in NLP.
-
-## Anatomy of an example
-
-Given a model named `modelname`, which involves both pretraining and finetuning
-on a downstream task, the contents of the `modelname` directory should be as
-follows:
-
-```shell
-modelname
-├── README.md
-├── __init__.py
-├── modelname_config.py
-├── modelname_model.py
-├── modelname_preprocess.py
-├── modelname_train.py
-└── modelname_finetune_X.py
-```
-
-- `README.md`: The README should contain complete instructions for downloading
-  data and training a model from scratch.
-- `__init__.py`: Empty (it's for imports).
-- `modelname_config.py`: This file should contain most of the configuration for
-  the model architecture, learning rate, etc, using simple Python constants. We
-  would like to avoid complex configuration setups (json, yaml, etc).
-- `modelname_preprocess.py`: If necessary. Standalone script to preprocess
-  inputs. If possible, prefer doing preprocessing dynamically with tf.data
-  inside the training and finetuning scripts.
-- `modelname_model.py`: This file should contain the actual `keras.Model` and
-  any custom layers needed for the example. Use KerasNLP components where ever
-  possible.
-- `modelname_train.py`: This file should be a runnable training script for
-  pretraining. If possible, this script should preprocess data dynamically
-  during training using `tf.data` and KerasNLP components (e.g. tokenizers).
-- `modelname_finetune_X.py`: Optional. There can be any number of these files,
-  for each task `X` we would like to support for finetuning. The file should be
-  a runnable training script which loads and finetunes a pretrained model.
-
-## Instructions for running on Google Cloud
-
-TODO(https://github.com/keras-team/keras-nlp/issues/178)
+The `examples/` directly contains scripts built on top of the library that do not fit well into
+the colab format used on [keras.io](https://keras.io/examples/). This includes recipes for
+pre-training models and evaluating models on benchmarks such as GLUE.
diff --git a/examples/bert/bert_finetune_glue.py b/examples/bert/bert_finetune_glue.py
deleted file mode 100644
index d1f76ad8..00000000
--- a/examples/bert/bert_finetune_glue.py
+++ /dev/null
@@ -1,205 +0,0 @@
-# Copyright 2022 The KerasNLP Authors
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     https://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Run finetuning on a GLUE task."""
-
-import tempfile
-
-import keras_tuner
-import tensorflow as tf
-import tensorflow_datasets as tfds
-from absl import app
-from absl import flags
-from tensorflow import keras
-
-import keras_nlp
-from examples.bert.bert_config import FINETUNING_CONFIG
-from examples.bert.bert_config import PREPROCESSING_CONFIG
-
-FLAGS = flags.FLAGS
-
-flags.DEFINE_string(
-    "vocab_file",
-    None,
-    "The vocabulary file for tokenization.",
-)
-
-flags.DEFINE_string(
-    "saved_model_input",
-    None,
-    "The directory to load the pretrained model.",
-)
-
-flags.DEFINE_string(
-    "saved_model_output",
-    None,
-    "The directory to save the finetuned model.",
-)
-
-flags.DEFINE_string(
-    "task_name",
-    "mrpc",
-    "The name of the GLUE task to finetune on.",
-)
-
-flags.DEFINE_bool(
-    "do_lower_case",
-    True,
-    "Whether to lower case the input text.",
-)
-
-flags.DEFINE_bool(
-    "do_evaluation",
-    True,
-    "Whether to run evaluation on test data.",
-)
-
-
-def load_data(task_name):
-    if task_name in ("cola", "sst2"):
-        feature_names = ("sentence",)
-    elif task_name in ("mrpc", "stsb", "rte", "wnli"):
-        feature_names = ("sentence1", "sentence2")
-    elif task_name in ("mnli", "mnli_matched", "mnli_mismatched"):
-        feature_names = ("premise", "hypothesis")
-    elif task_name in "qnli":
-        feature_names = ("question", "sentence")
-    elif task_name in "qqp":
-        feature_names = ("question1", "question2")
-    else:
-        raise ValueError(f"Unknown task_name {task_name}.")
-
-    test_suffix = ""
-    if task_name in ("mnli", "mnli_matched"):
-        # For "mnli", just run default to "mnli_matched".
-        task_name = "mnli"
-        test_suffix = "_matched"
-    elif task_name in ("mnli_mismatched",):
-        task_name = "mnli"
-        test_suffix = "_mismatched"
-
-    def split_features(x):
-        features = tuple([x[name] for name in feature_names])
-        labels = x["label"]
-        return features, labels
-
-    train_ds, test_ds, validation_ds = tfds.load(
-        f"glue/{task_name}",
-        split=["train", "test" + test_suffix, "validation" + test_suffix],
-    )
-    train_ds = train_ds.map(split_features, num_parallel_calls=tf.data.AUTOTUNE)
-    test_ds = test_ds.map(split_features, num_parallel_calls=tf.data.AUTOTUNE)
-    validation_ds = validation_ds.map(
-        split_features, num_parallel_calls=tf.data.AUTOTUNE
-    )
-    return train_ds, test_ds, validation_ds
-
-
-class BertHyperModel(keras_tuner.HyperModel):
-    """Creates a hypermodel to help with the search space for finetuning."""
-
-    def build(self, hp):
-        model = keras.models.load_model(FLAGS.saved_model_input, compile=False)
-        finetuning_model = keras_nlp.models.BertClassifier(
-            backbone=model,
-            num_classes=3 if FLAGS.task_name in ("mnli", "ax") else 2,
-        )
-        finetuning_model.compile(
-            optimizer=keras.optimizers.Adam(
-                learning_rate=hp.Choice(
-                    "lr", FINETUNING_CONFIG["learning_rates"]
-                ),
-            ),
-            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
-            metrics=[keras.metrics.SparseCategoricalAccuracy()],
-        )
-        return finetuning_model
-
-
-def main(_):
-    print(f"Reading input model from {FLAGS.saved_model_input}")
-
-    tokenizer = keras_nlp.tokenizers.WordPieceTokenizer(
-        vocabulary=FLAGS.vocab_file,
-        lowercase=FLAGS.do_lower_case,
-    )
-    packer = keras_nlp.layers.MultiSegmentPacker(
-        sequence_length=PREPROCESSING_CONFIG["max_seq_length"],
-        start_value=tokenizer.token_to_id("[CLS]"),
-        end_value=tokenizer.token_to_id("[SEP]"),
-    )
-
-    def preprocess_data(inputs, labels):
-        inputs = [tokenizer(x) for x in inputs]
-        token_ids, segment_ids = packer(inputs)
-        return {
-            "token_ids": token_ids,
-            "segment_ids": segment_ids,
-            "padding_mask": tf.cast(token_ids != 0, "int32"),
-        }, labels
-
-    # Read and preprocess GLUE task data.
-    train_ds, test_ds, validation_ds = load_data(FLAGS.task_name)
-
-    batch_size = FINETUNING_CONFIG["batch_size"]
-    train_ds = train_ds.batch(batch_size).map(
-        preprocess_data, num_parallel_calls=tf.data.AUTOTUNE
-    )
-    validation_ds = validation_ds.batch(batch_size).map(
-        preprocess_data, num_parallel_calls=tf.data.AUTOTUNE
-    )
-    test_ds = test_ds.batch(batch_size).map(
-        preprocess_data, num_parallel_calls=tf.data.AUTOTUNE
-    )
-
-    # Create a hypermodel object for a RandomSearch.
-    hypermodel = BertHyperModel()
-
-    # Initialize the random search over the 4 learning rate parameters, for 4
-    # trials and 3 epochs for each trial.
-    tuner = keras_tuner.RandomSearch(
-        hypermodel=hypermodel,
-        objective=keras_tuner.Objective("val_loss", direction="min"),
-        max_trials=4,
-        overwrite=True,
-        project_name="hyperparameter_tuner_results",
-        directory=tempfile.mkdtemp(),
-    )
-
-    tuner.search(
-        train_ds,
-        epochs=FINETUNING_CONFIG["epochs"],
-        validation_data=validation_ds,
-    )
-
-    # Extract the best hyperparameters after the search.
-    best_hp = tuner.get_best_hyperparameters()[0]
-    finetuning_model = tuner.get_best_models()[0]
-
-    print(
-        f"The best hyperparameters found are:\nLearning Rate: {best_hp['lr']}"
-    )
-
-    if FLAGS.do_evaluation:
-        print("Evaluating on test set.")
-        finetuning_model.evaluate(test_ds)
-
-    if FLAGS.saved_model_output:
-        print(f"Saving to {FLAGS.saved_model_output}")
-        finetuning_model.save(FLAGS.saved_model_output)
-
-
-if __name__ == "__main__":
-    flags.mark_flag_as_required("vocab_file")
-    flags.mark_flag_as_required("saved_model_input")
-    app.run(main)
diff --git a/examples/bert/README.md b/examples/bert_pretraining/README.md
similarity index 85%
rename from examples/bert/README.md
rename to examples/bert_pretraining/README.md
index 23098408..361a629a 100644
--- a/examples/bert/README.md
+++ b/examples/bert_pretraining/README.md
@@ -25,20 +25,16 @@ python3 examples/tools/split_sentences.py \
     --input_files $OUTPUT_DIR/wiki_example_data.txt \
     --output_directory $OUTPUT_DIR/sentence-split-data
 # Preprocess input for pretraining.
-python3 examples/bert/bert_preprocess.py \
+python3 examples/bert_pretraining/bert_create_pretraining_data.py \
     --input_files $OUTPUT_DIR/sentence-split-data/ \
     --vocab_file $OUTPUT_DIR/bert_vocab_uncased.txt \
     --output_file $OUTPUT_DIR/pretraining-data/pretraining.tfrecord
 # Run pretraining for 100 train steps only.
-python3 examples/bert/bert_train.py \
+python3 examples/bert_pretraining/bert_pretrain.py \
     --input_directory $OUTPUT_DIR/pretraining-data/ \
     --vocab_file $OUTPUT_DIR/bert_vocab_uncased.txt \
     --saved_model_output $OUTPUT_DIR/model/ \
     --num_train_steps 100
-# Run finetuning.
-python3 examples/bert/bert_finetune_glue.py \
-    --saved_model_input $OUTPUT_DIR/model/ \
-    --vocab_file $OUTPUT_DIR/bert_vocab_uncased.txt
 ```
 
 ## Installing dependencies
@@ -124,7 +120,7 @@ You can also use `examples/tools/train_word_piece_vocab.py` to train your own.
 
 ### Tokenize, mask, and combine sentences into training examples
 
-The `bert_preprocess.py` script will take in a set of sentence split files, and
+The ` bert_create_pretraining_data.py` script will take in a set of sentence split files, and
 set up training examples for the next sentence prediction and masked word tasks.
 
 The output of the script will be TFRecord files with a number of fields per
@@ -160,7 +156,7 @@ with the following:
 ```shell
 for file in path/to/sentence-split-data/*; do
     output="path/to/pretraining-data/$(basename -- "$file" .txt).tfrecord"
-    python3 examples/bert/bert_preprocess.py \
+    python3 examples/bert_pretraining/bert_create_pretraining_data.py \
         --input_files ${file} \
         --vocab_file bert_vocab_uncased.txt \
         --output_file ${output}
@@ -175,7 +171,7 @@ on an 8 core machine.
 NUM_JOBS=5
 for file in path/to/sentence-split-data/*; do
     output="path/to/pretraining-data/$(basename -- "$file" .txt).tfrecord"
-    echo python3 examples/bert/bert_preprocess.py \
+    echo python3 examples/bert_pretraining/bert_create_pretraining_data.py \
         --input_files ${file} \
         --vocab_file bert_vocab_uncased.txt \
         --output_file ${output}
@@ -190,35 +186,15 @@ python3 -c "from examples.utils.data_utils import preview_tfrecord; preview_tfre
 
 ### Running BERT pretraining
 
-After preprocessing, we can run pretraining with the `bert_train.py`
+After preprocessing, we can run pretraining with the `bert_pretrain.py`
 script. This will train a model and save it to the `--saved_model_output`
 directory. If you are willing to train from data stored on google cloud storage bucket (GCS), you can do it by setting the file path to
 the URL of GCS bucket. For example, `--input_directory=gs://your-bucket-name/you-data-path`. You can also save models directly to GCS by the same approach.
 
 ```shell
-python3 examples/bert/bert_train.py \
+python3 examples/bert_pretraining/bert_pretrain.py \
     --input_directory path/to/data/ \
     --vocab_file path/to/bert_vocab_uncased.txt \
     --model_size tiny \
     --saved_model_output path/to/model/
 ```
-
-## Evaluating BERT with GLUE
-
-After pretraining, we can evaluate the performance of a BERT model with the
-General Language Understanding Evaluation (GLUE) benchmark. This will
-finetune the model and running classification for a number of downstream tasks.
-
-The `bert_finetune_glue.py` script downloads the GLUE data for a specific
-tasks, reloads the pretraining model with appropriate finetuning heads, and runs
-training for a few epochs to finetune the model.
-
-```shell
-python3 examples/bert/bert_finetune_glue.py \
-    --saved_model_input path/to/model/ \
-    --vocab_file path/to/bert_vocab_uncased.txt \
-    --task_name mrpc
-```
-
-The script could be easily adapted to any other text classification finetuning
-tasks, where inputs can be any number of raw text sentences per sample.
diff --git a/examples/bert/__init__.py b/examples/bert_pretraining/__init__.py
similarity index 100%
rename from examples/bert/__init__.py
rename to examples/bert_pretraining/__init__.py
diff --git a/examples/bert/bert_config.py b/examples/bert_pretraining/bert_config.py
similarity index 95%
rename from examples/bert/bert_config.py
rename to examples/bert_pretraining/bert_config.py
index 951ee152..31a7f9eb 100644
--- a/examples/bert/bert_config.py
+++ b/examples/bert_pretraining/bert_config.py
@@ -77,9 +77,3 @@ TRAINING_CONFIG = {
     # Percentage of training steps used for learning rate warmup.
     "warmup_percentage": 0.1,
 }
-
-FINETUNING_CONFIG = {
-    "batch_size": 32,
-    "epochs": 3,
-    "learning_rates": [5e-5, 4e-5, 3e-5, 2e-5],
-}
diff --git a/examples/bert/bert_preprocess.py b/examples/bert_pretraining/bert_create_pretraining_data.py
similarity index 99%
rename from examples/bert/bert_preprocess.py
rename to examples/bert_pretraining/bert_create_pretraining_data.py
index c580bf8e..4d5c4c88 100644
--- a/examples/bert/bert_preprocess.py
+++ b/examples/bert_pretraining/bert_create_pretraining_data.py
@@ -44,7 +44,7 @@ import tensorflow_text as tf_text
 from absl import app
 from absl import flags
 
-from examples.bert.bert_config import PREPROCESSING_CONFIG
+from examples.bert_pretraining.bert_config import PREPROCESSING_CONFIG
 from examples.utils.scripting_utils import list_filenames_for_arg
 
 # Tokenization will happen with tensorflow and can easily OOM a GPU.
diff --git a/examples/bert/bert_train.py b/examples/bert_pretraining/bert_pretrain.py
similarity index 98%
rename from examples/bert/bert_train.py
rename to examples/bert_pretraining/bert_pretrain.py
index 09b08999..9bd0da49 100644
--- a/examples/bert/bert_train.py
+++ b/examples/bert_pretraining/bert_pretrain.py
@@ -22,9 +22,9 @@ from absl import logging
 from tensorflow import keras
 
 import keras_nlp
-from examples.bert.bert_config import MODEL_CONFIGS
-from examples.bert.bert_config import PREPROCESSING_CONFIG
-from examples.bert.bert_config import TRAINING_CONFIG
+from examples.bert_pretraining.bert_config import MODEL_CONFIGS
+from examples.bert_pretraining.bert_config import PREPROCESSING_CONFIG
+from examples.bert_pretraining.bert_config import TRAINING_CONFIG
 
 FLAGS = flags.FLAGS
 
diff --git a/keras_nlp/models/__init__.py b/keras_nlp/models/__init__.py
index 21d83940..ba473478 100644
--- a/keras_nlp/models/__init__.py
+++ b/keras_nlp/models/__init__.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from keras_nlp.models.albert.albert_backbone import AlbertBackbone
 from keras_nlp.models.albert.albert_preprocessor import AlbertPreprocessor
 from keras_nlp.models.albert.albert_tokenizer import AlbertTokenizer
 from keras_nlp.models.bert.bert_backbone import BertBackbone
@@ -28,6 +29,8 @@ from keras_nlp.models.distil_bert.distil_bert_preprocessor import (
 from keras_nlp.models.distil_bert.distil_bert_tokenizer import (
     DistilBertTokenizer,
 )
+from keras_nlp.models.f_net.f_net_preprocessor import FNetPreprocessor
+from keras_nlp.models.f_net.f_net_tokenizer import FNetTokenizer
 from keras_nlp.models.roberta.roberta_backbone import RobertaBackbone
 from keras_nlp.models.roberta.roberta_classifier import RobertaClassifier
 from keras_nlp.models.roberta.roberta_preprocessor import RobertaPreprocessor
diff --git a/keras_nlp/models/albert/albert_backbone.py b/keras_nlp/models/albert/albert_backbone.py
new file mode 100644
index 00000000..45017991
--- /dev/null
+++ b/keras_nlp/models/albert/albert_backbone.py
@@ -0,0 +1,266 @@
+# 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.
+
+"""ALBERT backbone model."""
+
+import tensorflow as tf
+from tensorflow import keras
+
+from keras_nlp.layers.position_embedding import PositionEmbedding
+from keras_nlp.layers.transformer_encoder import TransformerEncoder
+from keras_nlp.models.backbone import Backbone
+
+
+def albert_kernel_initializer(stddev=0.02):
+    return keras.initializers.TruncatedNormal(stddev=stddev)
+
+
+@keras.utils.register_keras_serializable(package="keras_nlp")
+class AlbertBackbone(Backbone):
+    """ALBERT encoder network.
+
+    This class implements a bi-directional Transformer-based encoder as
+    described in
+    ["ALBERT: A Lite BERT for Self-supervised Learning of Language Representations"](https://arxiv.org/abs/1909.11942).
+    ALBERT is a more efficient variant of BERT, and uses parameter reduction
+    techniques such as cross-layer parameter sharing and factorized embedding
+    parameterization. This model class includes the embedding lookups and
+    transformer layers, but not the masked language model or sentence order
+    prediction heads.
+
+    The default constructor gives a fully customizable, randomly initialized
+    ALBERT encoder with any number of layers, heads, and embedding dimensions.
+    To load preset architectures and weights, use the `from_preset` constructor.
+
+    Disclaimer: Pre-trained models are provided on an "as is" basis, without
+    warranties or conditions of any kind.
+
+    Args:
+        vocabulary_size: int. The size of the token vocabulary.
+        num_layers: int, must be divisible by `num_groups`. The number of
+            "virtual" layers, i.e., the total number of times the input sequence
+            will be fed through the groups in one forward pass. The input will
+            be routed to the correct group based on the layer index.
+        num_heads: int. The number of attention heads for each transformer.
+            The hidden size must be divisible by the number of attention heads.
+        embedding_dim: int. The size of the embeddings.
+        hidden_dim: int. The size of the transformer encoding and pooler layers.
+        intermediate_dim: int. The output dimension of the first Dense layer in
+            a two-layer feedforward network for each transformer.
+        num_groups: int. Number of groups, with each group having
+            `num_inner_repetitions` number of `TransformerEncoder` layers.
+        num_inner_repetitions: int. Number of `TransformerEncoder` layers per
+            group.
+        dropout: float. Dropout probability for the Transformer encoder.
+        max_sequence_length: int. The maximum sequence length that this encoder
+            can consume. If None, `max_sequence_length` uses the value from
+            sequence length. This determines the variable shape for positional
+            embeddings.
+        num_segments: int. The number of types that the 'segment_ids' input can
+            take.
+
+    Examples:
+    ```python
+    input_data = {
+        "token_ids": tf.ones(shape=(1, 12), dtype=tf.int64),
+        "segment_ids": tf.constant(
+            [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
+        ),
+        "padding_mask": tf.constant(
+            [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
+        ),
+    }
+
+    # Randomly initialized ALBERT encoder
+    model = keras_nlp.models.AlbertBackbone(
+        vocabulary_size=30000,
+        num_layers=12,
+        num_heads=12,
+        num_groups=1,
+        num_inner_repetitions=1,
+        embedding_dim=128,
+        hidden_dim=768,
+        intermediate_dim=3072,
+        max_sequence_length=12,
+    )
+    output = model(input_data)
+    ```
+    """
+
+    def __init__(
+        self,
+        vocabulary_size,
+        num_layers,
+        num_heads,
+        embedding_dim,
+        hidden_dim,
+        intermediate_dim,
+        num_groups=1,
+        num_inner_repetitions=1,
+        dropout=0.0,
+        max_sequence_length=512,
+        num_segments=2,
+        **kwargs,
+    ):
+
+        if num_layers % num_groups != 0:
+            raise ValueError(
+                "`num_layers` must be divisible by `num_groups`. Received: "
+                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 = keras.layers.Embedding(
+            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(
+            initializer=albert_kernel_initializer(),
+            sequence_length=max_sequence_length,
+            name="position_embedding",
+        )(token_embedding)
+        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(
+            name="embeddings_layer_norm",
+            axis=-1,
+            epsilon=1e-12,
+            dtype=tf.float32,
+        )(x)
+        x = keras.layers.Dropout(
+            dropout,
+            name="embeddings_dropout",
+        )(x)
+
+        # Project the embedding to `hidden_dim`.
+        x = 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(
+                    num_heads=num_heads,
+                    intermediate_dim=intermediate_dim,
+                    activation=lambda x: keras.activations.gelu(
+                        x, approximate=True
+                    ),
+                    dropout=dropout,
+                    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(
+            hidden_dim,
+            kernel_initializer=albert_kernel_initializer(),
+            activation="tanh",
+            name="pooled_dense",
+        )(x[:, cls_token_index, :])
+
+        # Instantiate using Functional API Model constructor
+        super().__init__(
+            inputs={
+                "token_ids": token_id_input,
+                "segment_ids": segment_id_input,
+                "padding_mask": padding_mask,
+            },
+            outputs={
+                "sequence_output": sequence_output,
+                "pooled_output": pooled_output,
+            },
+            **kwargs,
+        )
+        # All references to `self` below this line
+        self.vocabulary_size = vocabulary_size
+        self.num_layers = num_layers
+        self.num_heads = num_heads
+        self.num_groups = num_groups
+        self.num_inner_repetitions = num_inner_repetitions
+        self.embedding_dim = embedding_dim
+        self.hidden_dim = hidden_dim
+        self.intermediate_dim = intermediate_dim
+        self.dropout = dropout
+        self.max_sequence_length = max_sequence_length
+        self.num_segments = num_segments
+        self.cls_token_index = cls_token_index
+
+    def get_config(self):
+        return {
+            "vocabulary_size": self.vocabulary_size,
+            "num_layers": self.num_layers,
+            "num_heads": self.num_heads,
+            "num_groups": self.num_groups,
+            "num_inner_repetitions": self.num_inner_repetitions,
+            "embedding_dim": self.embedding_dim,
+            "hidden_dim": self.hidden_dim,
+            "intermediate_dim": self.intermediate_dim,
+            "dropout": self.dropout,
+            "max_sequence_length": self.max_sequence_length,
+            "num_segments": self.num_segments,
+            "name": self.name,
+            "trainable": self.trainable,
+        }
diff --git a/keras_nlp/models/bert/bert_backbone.py b/keras_nlp/models/bert/bert_backbone.py
index 8339d0f8..6fb8a25e 100644
--- a/keras_nlp/models/bert/bert_backbone.py
+++ b/keras_nlp/models/bert/bert_backbone.py
@@ -139,7 +139,7 @@ class BertBackbone(Backbone):
             name="segment_embedding",
         )(segment_id_input)
 
-        # Sum, normailze and apply dropout to embeddings.
+        # Sum, normalize and apply dropout to embeddings.
         x = keras.layers.Add()(
             (token_embedding, position_embedding, segment_embedding)
         )
diff --git a/keras_nlp/models/f_net/__init__.py b/keras_nlp/models/f_net/__init__.py
new file mode 100644
index 00000000..6e4df4e7
--- /dev/null
+++ b/keras_nlp/models/f_net/__init__.py
@@ -0,0 +1,13 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
diff --git a/keras_nlp/models/f_net/f_net_preprocessor.py b/keras_nlp/models/f_net/f_net_preprocessor.py
new file mode 100644
index 00000000..b36ba353
--- /dev/null
+++ b/keras_nlp/models/f_net/f_net_preprocessor.py
@@ -0,0 +1,179 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""FNet preprocessor layer."""
+
+from tensorflow import keras
+
+from keras_nlp.layers.multi_segment_packer import MultiSegmentPacker
+from keras_nlp.models.f_net.f_net_tokenizer import FNetTokenizer
+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.utils.register_keras_serializable(package="keras_nlp")
+class FNetPreprocessor(Preprocessor):
+    """An FNet preprocessing layer which tokenizes and packs inputs.
+
+    This preprocessing layer will do three things:
+
+     - Tokenize any number of input segments using the `tokenizer`.
+     - Pack the inputs together using a `keras_nlp.layers.MultiSegmentPacker`.
+       with the appropriate `"[CLS]"`, `"[SEP]"` and `"<pad>"` tokens.
+     - Construct a dictionary with keys `"token_ids"`, and `"segment_ids"`  that
+       can be passed directly to `keras_nlp.models.FNetBackbone`.
+
+    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.
+
+    Special care should be taken when using `tf.data` to map over an unlabeled
+    tuple of string segments. `tf.data.Dataset.map` will unpack this tuple
+    directly into the call arguments of this layer, rather than forward all
+    argument to `x`. To handle this case, it is recommended to  explicitly call
+    the layer, e.g. `ds.map(lambda seg1, seg2: preprocessor(x=(seg1, seg2)))`.
+
+    Args:
+        tokenizer: A `keras_nlp.models.FNetTokenizer` instance.
+        sequence_length: The length of the packed inputs.
+        truncate: string. The algorithm to truncate a list of batched segments
+            to fit within `sequence_length`. The value can be either
+            `round_robin` or `waterfall`:
+                - `"round_robin"`: Available space is assigned one token at a
+                    time in a round-robin fashion to the inputs that still need
+                    some, until the limit is reached.
+                - `"waterfall"`: The allocation of the budget is done using a
+                    "waterfall" algorithm that allocates quota in a
+                    left-to-right manner and fills up the buckets until we run
+                    out of budget. It supports an arbitrary number of segments.
+
+    Examples:
+    ```python
+    tokenizer = keras_nlp.models.FNetTokenizer(proto="model.spm")
+    preprocessor = keras_nlp.models.FNetPreprocessor(
+        tokenizer=tokenizer,
+        sequence_length=10,
+    )
+
+    # Tokenize and pack a single sentence.
+    sentence = tf.constant("The quick brown fox jumped.")
+    preprocessor(sentence)
+    # Same output.
+    preprocessor("The quick brown fox jumped.")
+
+    # Tokenize and a batch of single sentences.
+    sentences = tf.constant(
+        ["The quick brown fox jumped.", "Call me Ishmael."]
+    )
+    preprocessor(sentences)
+    # Same output.
+    preprocessor(
+        ["The quick brown fox jumped.", "Call me Ishmael."]
+    )
+
+    # Tokenize and pack a sentence pair.
+    first_sentence = tf.constant("The quick brown fox jumped.")
+    second_sentence = tf.constant("The fox tripped.")
+    preprocessor((first_sentence, second_sentence))
+
+    # Map a dataset to preprocess a single sentence.
+    features = tf.constant(
+        ["The quick brown fox jumped.", "Call me Ishmael."]
+    )
+    labels = tf.constant([0, 1])
+    ds = tf.data.Dataset.from_tensor_slices((features, labels))
+    ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
+
+    # Map a dataset to preprocess sentence pairs.
+    first_sentences = tf.constant(
+        ["The quick brown fox jumped.", "Call me Ishmael."]
+    )
+    second_sentences = tf.constant(
+        ["The fox tripped.", "Oh look, a whale."]
+    )
+    labels = tf.constant([1, 1])
+    ds = tf.data.Dataset.from_tensor_slices(
+        (
+            (first_sentences, second_sentences), labels
+        )
+    )
+    ds = ds.map(preprocessor, num_parallel_calls=tf.data.AUTOTUNE)
+
+    # Map a dataset to preprocess unlabeled sentence pairs.
+    first_sentences = tf.constant(
+        ["The quick brown fox jumped.", "Call me Ishmael."]
+    )
+    second_sentences = tf.constant(
+        ["The fox tripped.", "Oh look, a whale."]
+    )
+    ds = tf.data.Dataset.from_tensor_slices((first_sentences, second_sentences))
+    # Watch out for tf.data's default unpacking of tuples here!
+    # Best to invoke the `preprocessor` directly in this case.
+    ds = ds.map(
+        lambda s1, s2: preprocessor(x=(s1, s2)),
+        num_parallel_calls=tf.data.AUTOTUNE,
+    )
+    ```
+    """
+
+    def __init__(
+        self,
+        tokenizer,
+        sequence_length=512,
+        truncate="round_robin",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self._tokenizer = tokenizer
+        self.packer = MultiSegmentPacker(
+            start_value=self.tokenizer.cls_token_id,
+            end_value=self.tokenizer.sep_token_id,
+            pad_value=self.tokenizer.pad_token_id,
+            truncate=truncate,
+            sequence_length=sequence_length,
+        )
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "sequence_length": self.packer.sequence_length,
+                "truncate": self.packer.truncate,
+            }
+        )
+        return config
+
+    def call(self, x, y=None, sample_weight=None):
+        x = convert_inputs_to_list_of_tensor_segments(x)
+        x = [self.tokenizer(segment) for segment in x]
+        token_ids, segment_ids = self.packer(x)
+        x = {
+            "token_ids": token_ids,
+            "segment_ids": segment_ids,
+        }
+        return pack_x_y_sample_weight(x, y, sample_weight)
+
+    @classproperty
+    def tokenizer_cls(cls):
+        return FNetTokenizer
diff --git a/keras_nlp/models/f_net/f_net_tokenizer.py b/keras_nlp/models/f_net/f_net_tokenizer.py
new file mode 100644
index 00000000..c0d2c29c
--- /dev/null
+++ b/keras_nlp/models/f_net/f_net_tokenizer.py
@@ -0,0 +1,95 @@
+# 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.
+
+"""FNet tokenizer."""
+
+
+from tensorflow import keras
+
+from keras_nlp.tokenizers.sentence_piece_tokenizer import SentencePieceTokenizer
+from keras_nlp.utils.python_utils import classproperty
+
+
+@keras.utils.register_keras_serializable(package="keras_nlp")
+class FNetTokenizer(SentencePieceTokenizer):
+    """FNet tokenizer layer based on SentencePiece.
+
+    This tokenizer class will tokenize raw strings into integer sequences and
+    is based on `keras_nlp.tokenizers.SentencePieceTokenizer`. Unlike the
+    underlying tokenizer, it will check for all special tokens needed by
+    FNet models and provides a `from_preset()` method to automatically
+    download a matching vocabulary for a FNet preset.
+
+    This tokenizer does not provide truncation or padding of inputs. It can be
+    combined with a `keras_nlp.models.FNetPreprocessor` layer for input
+    packing.
+
+    If input is a batch of strings (rank > 0), the layer will output a
+    `tf.RaggedTensor` where the last dimension of the output is ragged.
+
+    If input is a scalar string (rank == 0), the layer will output a dense
+    `tf.Tensor` with static shape `[None]`.
+
+    Args:
+        proto: Either a `string` path to a SentencePiece proto file, or a
+            `bytes` object with a serialized SentencePiece proto. See the
+            [SentencePiece repository](https://github.com/google/sentencepiece)
+            for more details on the format.
+
+    Examples:
+
+    ```python
+    tokenizer = keras_nlp.models.FNetTokenizer(proto="model.spm")
+
+    # Batched inputs.
+    tokenizer(["the quick brown fox", "the earth is round"])
+
+    # Unbatched inputs.
+    tokenizer("the quick brown fox")
+
+    # Detokenization.
+    tokenizer.detokenize(tf.constant([[[2, 14, 2231, 886, 2385, 3]]))
+    ```
+    """
+
+    def __init__(self, proto, **kwargs):
+        super().__init__(proto=proto, **kwargs)
+
+        # Check for necessary special tokens.
+        cls_token = "[CLS]"
+        sep_token = "[SEP]"
+        pad_token = "<pad>"
+        for token in [cls_token, sep_token, pad_token]:
+            if token not in self.get_vocabulary():
+                raise ValueError(
+                    f"Cannot find token `'{token}'` in the provided "
+                    f"`vocabulary`. Please provide `'{token}'` in your "
+                    "`vocabulary` or use a pretrained `vocabulary` name."
+                )
+
+        self.cls_token_id = self.token_to_id(cls_token)
+        self.sep_token_id = self.token_to_id(sep_token)
+        self.pad_token_id = self.token_to_id(pad_token)
+
+    @classproperty
+    def presets(cls):
+        return {}
+
+    @classmethod
+    def from_preset(
+        cls,
+        preset,
+        **kwargs,
+    ):
+        raise NotImplementedError
diff --git a/tools/count_preset_params.py b/tools/count_preset_params.py
new file mode 100644
index 00000000..3a676ca1
--- /dev/null
+++ b/tools/count_preset_params.py
@@ -0,0 +1,55 @@
+# 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.
+"""
+Small utility script to count parameters in our preset checkpoints.
+
+Usage:
+python tools/count_preset_params.py
+python tools/count_preset_params.py --model BertBackbone
+python tools/count_preset_params.py --preset bert_base_multi
+"""
+
+import inspect
+
+from absl import app
+from absl import flags
+from keras.utils.layer_utils import count_params
+from tensorflow import keras
+
+import keras_nlp
+
+FLAGS = flags.FLAGS
+flags.DEFINE_string("model", None, "The name of a model, e.g. BertBackbone.")
+flags.DEFINE_string(
+    "preset", None, "The name of a preset, e.g. bert_base_multi."
+)
+
+
+def main(_):
+    for name, symbol in keras_nlp.models.__dict__.items():
+        if inspect.isclass(symbol) and issubclass(symbol, keras.Model):
+            if FLAGS.model and name != FLAGS.model:
+                continue
+            if not hasattr(symbol, "from_preset"):
+                continue
+            for preset in symbol.presets:
+                if FLAGS.preset and preset != FLAGS.preset:
+                    continue
+                model = symbol.from_preset(preset)
+                params = count_params(model.weights)
+                print(f"{name} {preset} {params}")
+
+
+if __name__ == "__main__":
+    app.run(main)


Test patch:
diff --git a/keras_nlp/models/albert/albert_backbone_test.py b/keras_nlp/models/albert/albert_backbone_test.py
new file mode 100644
index 00000000..d9f73886
--- /dev/null
+++ b/keras_nlp/models/albert/albert_backbone_test.py
@@ -0,0 +1,120 @@
+# 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.
+"""Test for ALBERT backbone model."""
+
+import os
+
+import tensorflow as tf
+from absl.testing import parameterized
+from tensorflow import keras
+
+from keras_nlp.models.albert.albert_backbone import AlbertBackbone
+
+
+class AlbertBackboneTest(tf.test.TestCase, parameterized.TestCase):
+    def setUp(self):
+        self.model = AlbertBackbone(
+            vocabulary_size=1000,
+            num_layers=2,
+            num_heads=2,
+            num_groups=1,
+            num_inner_repetitions=1,
+            embedding_dim=16,
+            hidden_dim=64,
+            intermediate_dim=128,
+            max_sequence_length=128,
+        )
+        self.batch_size = 8
+        self.input_batch = {
+            "token_ids": tf.ones(
+                (self.batch_size, self.model.max_sequence_length), dtype="int32"
+            ),
+            "segment_ids": tf.ones(
+                (self.batch_size, self.model.max_sequence_length), dtype="int32"
+            ),
+            "padding_mask": tf.ones(
+                (self.batch_size, self.model.max_sequence_length), dtype="int32"
+            ),
+        }
+
+        self.input_dataset = tf.data.Dataset.from_tensor_slices(
+            self.input_batch
+        ).batch(2)
+
+    def test_valid_call_albert(self):
+        self.model(self.input_batch)
+
+        # Check default name passed through
+        self.assertRegexpMatches(self.model.name, "albert_backbone")
+
+    def test_variable_sequence_length_call_albert(self):
+        for seq_length in (25, 50, 75):
+            input_data = {
+                "token_ids": tf.ones(
+                    (self.batch_size, seq_length), dtype="int32"
+                ),
+                "segment_ids": tf.ones(
+                    (self.batch_size, seq_length), dtype="int32"
+                ),
+                "padding_mask": tf.ones(
+                    (self.batch_size, seq_length), dtype="int32"
+                ),
+            }
+            self.model(input_data)
+
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_compile(self, jit_compile):
+        self.model.compile(jit_compile=jit_compile)
+        self.model.predict(self.input_batch)
+
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_compile_batched_ds(self, jit_compile):
+        self.model.compile(jit_compile=jit_compile)
+        self.model.predict(self.input_dataset)
+
+    def test_error_for_invalid_num_groups(self):
+        with self.assertRaises(ValueError):
+            self.model = AlbertBackbone(
+                vocabulary_size=1000,
+                num_layers=3,
+                num_heads=2,
+                num_groups=2,
+                num_inner_repetitions=1,
+                embedding_dim=16,
+                hidden_dim=64,
+                intermediate_dim=128,
+            )
+
+    @parameterized.named_parameters(
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
+    )
+    def test_saved_model(self, save_format, filename):
+        model_output = self.model(self.input_batch)
+        save_path = os.path.join(self.get_temp_dir(), filename)
+        self.model.save(save_path, save_format=save_format)
+        restored_model = keras.models.load_model(save_path)
+
+        # Check we got the real object back.
+        self.assertIsInstance(restored_model, AlbertBackbone)
+
+        # Check that output matches.
+        restored_output = restored_model(self.input_batch)
+        self.assertAllClose(
+            model_output["pooled_output"], restored_output["pooled_output"]
+        )
diff --git a/keras_nlp/models/f_net/f_net_preprocessor_test.py b/keras_nlp/models/f_net/f_net_preprocessor_test.py
new file mode 100644
index 00000000..9b254cf2
--- /dev/null
+++ b/keras_nlp/models/f_net/f_net_preprocessor_test.py
@@ -0,0 +1,156 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for FNet preprocessor layer."""
+
+import io
+import os
+
+import sentencepiece
+import tensorflow as tf
+from absl.testing import parameterized
+from tensorflow import keras
+
+from keras_nlp.models.f_net.f_net_preprocessor import FNetPreprocessor
+from keras_nlp.models.f_net.f_net_tokenizer import FNetTokenizer
+
+
+class FNetPreprocessorTest(tf.test.TestCase, parameterized.TestCase):
+    def setUp(self):
+        bytes_io = io.BytesIO()
+        vocab_data = tf.data.Dataset.from_tensor_slices(
+            ["the quick brown fox", "the earth is round"]
+        )
+        sentencepiece.SentencePieceTrainer.train(
+            sentence_iterator=vocab_data.as_numpy_iterator(),
+            model_writer=bytes_io,
+            vocab_size=10,
+            model_type="WORD",
+            pad_id=3,
+            unk_id=0,
+            bos_id=4,
+            eos_id=5,
+            pad_piece="<pad>",
+            unk_piece="<unk>",
+            bos_piece="[CLS]",
+            eos_piece="[SEP]",
+        )
+        self.proto = bytes_io.getvalue()
+
+        self.preprocessor = FNetPreprocessor(
+            tokenizer=FNetTokenizer(proto=self.proto),
+            sequence_length=12,
+        )
+
+    def test_tokenize_strings(self):
+        input_data = "the quick brown fox"
+        output = self.preprocessor(input_data)
+        self.assertAllEqual(
+            output["token_ids"], [4, 1, 9, 2, 7, 5, 3, 3, 3, 3, 3, 3]
+        )
+        self.assertAllEqual(
+            output["segment_ids"], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
+        )
+
+    def test_tokenize_list_of_strings(self):
+        # We should handle a list of strings as as batch.
+        input_data = ["the quick brown fox"] * 4
+        output = self.preprocessor(input_data)
+        self.assertAllEqual(
+            output["token_ids"],
+            [[4, 1, 9, 2, 7, 5, 3, 3, 3, 3, 3, 3]] * 4,
+        )
+        self.assertAllEqual(
+            output["segment_ids"], [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]] * 4
+        )
+
+    def test_tokenize_labeled_batch(self):
+        x = tf.constant(["the quick brown fox"] * 4)
+        y = tf.constant([1] * 4)
+        sw = tf.constant([1.0] * 4)
+        x_out, y_out, sw_out = self.preprocessor(x, y, sw)
+        self.assertAllEqual(
+            x_out["token_ids"],
+            [[4, 1, 9, 2, 7, 5, 3, 3, 3, 3, 3, 3]] * 4,
+        )
+        self.assertAllEqual(
+            x_out["segment_ids"], [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]] * 4
+        )
+        self.assertAllEqual(y_out, y)
+        self.assertAllEqual(sw_out, sw)
+
+    def test_tokenize_labeled_dataset(self):
+        x = tf.constant(["the quick brown fox"] * 4)
+        y = tf.constant([1] * 4)
+        sw = tf.constant([1.0] * 4)
+        ds = tf.data.Dataset.from_tensor_slices((x, y, sw))
+        ds = ds.map(self.preprocessor)
+        x_out, y_out, sw_out = ds.batch(4).take(1).get_single_element()
+        self.assertAllEqual(
+            x_out["token_ids"],
+            [[4, 1, 9, 2, 7, 5, 3, 3, 3, 3, 3, 3]] * 4,
+        )
+        self.assertAllEqual(
+            x_out["segment_ids"], [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]] * 4
+        )
+        self.assertAllEqual(y_out, y)
+        self.assertAllEqual(sw_out, sw)
+
+    def test_tokenize_multiple_sentences(self):
+        sentence_one = tf.constant("the quick brown fox")
+        sentence_two = tf.constant("the earth")
+        output = self.preprocessor((sentence_one, sentence_two))
+        self.assertAllEqual(
+            output["token_ids"],
+            [4, 1, 9, 2, 7, 5, 1, 6, 5, 3, 3, 3],
+        )
+        self.assertAllEqual(
+            output["segment_ids"], [0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0]
+        )
+
+    def test_tokenize_multiple_batched_sentences(self):
+        sentence_one = tf.constant(["the quick brown fox"] * 4)
+        sentence_two = tf.constant(["the earth"] * 4)
+        # The first tuple or list is always interpreted as an enumeration of
+        # separate sequences to concatenate.
+        output = self.preprocessor((sentence_one, sentence_two))
+        self.assertAllEqual(
+            output["token_ids"],
+            [[4, 1, 9, 2, 7, 5, 1, 6, 5, 3, 3, 3]] * 4,
+        )
+        self.assertAllEqual(
+            output["segment_ids"], [[0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0]] * 4
+        )
+
+    def test_errors_for_2d_list_input(self):
+        ambiguous_input = [["one", "two"], ["three", "four"]]
+        with self.assertRaises(ValueError):
+            self.preprocessor(ambiguous_input)
+
+    @parameterized.named_parameters(
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
+    )
+    def test_saved_model(self, save_format, filename):
+        input_data = tf.constant(["the quick brown fox"])
+        inputs = keras.Input(dtype="string", shape=())
+        outputs = self.preprocessor(inputs)
+        model = keras.Model(inputs, outputs)
+        path = os.path.join(self.get_temp_dir(), filename)
+        model.save(path, save_format=save_format)
+        restored_model = keras.models.load_model(path)
+        self.assertAllEqual(
+            model(input_data)["token_ids"],
+            restored_model(input_data)["token_ids"],
+        )
diff --git a/keras_nlp/models/f_net/f_net_tokenizer_test.py b/keras_nlp/models/f_net/f_net_tokenizer_test.py
new file mode 100644
index 00000000..dfeb9c23
--- /dev/null
+++ b/keras_nlp/models/f_net/f_net_tokenizer_test.py
@@ -0,0 +1,89 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for FNet tokenizer."""
+
+import io
+import os
+
+import sentencepiece
+import tensorflow as tf
+from absl.testing import parameterized
+from tensorflow import keras
+
+from keras_nlp.models.f_net.f_net_tokenizer import FNetTokenizer
+
+
+class FNetTokenizerTest(tf.test.TestCase, parameterized.TestCase):
+    def setUp(self):
+        bytes_io = io.BytesIO()
+        vocab_data = tf.data.Dataset.from_tensor_slices(
+            ["the quick brown fox", "the earth is round"]
+        )
+        sentencepiece.SentencePieceTrainer.train(
+            sentence_iterator=vocab_data.as_numpy_iterator(),
+            model_writer=bytes_io,
+            vocab_size=10,
+            model_type="WORD",
+            pad_id=3,
+            unk_id=0,
+            bos_id=4,
+            eos_id=5,
+            pad_piece="<pad>",
+            unk_piece="<unk>",
+            bos_piece="[CLS]",
+            eos_piece="[SEP]",
+        )
+        self.proto = bytes_io.getvalue()
+
+        self.tokenizer = FNetTokenizer(proto=self.proto)
+
+    def test_tokenize(self):
+        input_data = "the quick brown fox"
+        output = self.tokenizer(input_data)
+        self.assertAllEqual(output, [1, 9, 2, 7])
+
+    def test_tokenize_batch(self):
+        input_data = tf.constant(["the quick brown fox", "the earth is round"])
+        output = self.tokenizer(input_data)
+        self.assertAllEqual(output, [[1, 9, 2, 7], [1, 6, 8, 0]])
+
+    def test_detokenize(self):
+        input_data = tf.constant([[1, 9, 2, 7]])
+        output = self.tokenizer.detokenize(input_data)
+        self.assertEqual(output, tf.constant(["the quick brown fox"]))
+
+    def test_vocabulary_size(self):
+        tokenizer = FNetTokenizer(proto=self.proto)
+        self.assertEqual(tokenizer.vocabulary_size(), 10)
+
+    @parameterized.named_parameters(
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
+    )
+    def test_saved_model(self, save_format, filename):
+        input_data = tf.constant(["the quick brown fox"])
+
+        inputs = keras.Input(dtype="string", shape=())
+        outputs = self.tokenizer(inputs)
+        model = keras.Model(inputs, outputs)
+
+        path = os.path.join(self.get_temp_dir(), filename)
+        model.save(path, save_format=save_format)
+
+        restored_model = keras.models.load_model(path)
+        self.assertAllEqual(
+            model(input_data),
+            restored_model(input_data),
+        )