Repository: keras-team/keras-nlp
Pull request #592 (2022-12-21 02:13:54)

GPT2 Text Generation APIs
You can view the demo via this [link](https://colab.sandbox.google.com/drive/183PgXP09Fydl3tKN4bTKIql2pkx-kO1X?pli=1&usp=drive_open)


Patch:
diff --git a/keras_nlp/__init__.py b/keras_nlp/__init__.py
index 050a189d..5717f8bc 100644
--- a/keras_nlp/__init__.py
+++ b/keras_nlp/__init__.py
@@ -12,6 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# sentencepiece is segfaulting on tf-nightly if tensorflow is imported first.
+# This is a temporary fix to restore our nightly testing to green, while we look
+# for a longer term solution.
+try:
+    import sentencepiece
+except ImportError:
+    pass
+
 from keras_nlp import layers
 from keras_nlp import metrics
 from keras_nlp import models
diff --git a/keras_nlp/models/__init__.py b/keras_nlp/models/__init__.py
index 6d0d142d..66acc027 100644
--- a/keras_nlp/models/__init__.py
+++ b/keras_nlp/models/__init__.py
@@ -33,6 +33,13 @@ from keras_nlp.models.distil_bert.distil_bert_tokenizer import (
 from keras_nlp.models.f_net.f_net_backbone import FNetBackbone
 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.gpt2.gpt2_backbone import GPT2Backbone
+from keras_nlp.models.gpt2.gpt2_causal_lm import GPT2CausalLM
+from keras_nlp.models.gpt2.gpt2_causal_lm_preprocessor import (
+    GPT2CausalLMPreprocessor,
+)
+from keras_nlp.models.gpt2.gpt2_preprocessor import GPT2Preprocessor
+from keras_nlp.models.gpt2.gpt2_tokenizer import GPT2Tokenizer
 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/gpt2/gpt2_causal_lm.py b/keras_nlp/models/gpt2/gpt2_causal_lm.py
new file mode 100644
index 00000000..b1078503
--- /dev/null
+++ b/keras_nlp/models/gpt2/gpt2_causal_lm.py
@@ -0,0 +1,239 @@
+# 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.
+"""GPT2 Causal LM (Language Model)."""
+
+import copy
+
+import tensorflow as tf
+from tensorflow import keras
+
+import keras_nlp
+from keras_nlp.models.gpt2.gpt2_backbone import GPT2Backbone
+from keras_nlp.models.gpt2.gpt2_causal_lm_preprocessor import (
+    GPT2CausalLMPreprocessor,
+)
+from keras_nlp.models.gpt2.gpt2_presets import backbone_presets
+from keras_nlp.models.task import Task
+from keras_nlp.utils.python_utils import classproperty
+
+
+@keras.utils.register_keras_serializable(package="keras_nlp")
+class GPT2CausalLM(Task):
+    """An end-to-end GPT2 model for causal langauge modeling.
+
+    A causal language model (LM) predicts the next token based on previous
+    tokens the next token based on previous tokens, which is the way GPT2 gets
+    pretrained. You can finetune `GPT2CausalLM` to generate text similar to
+    the custom dataset. `GPT2CausalLM` also has a method `generate()`, which
+    generates text based on given prompt.
+
+    This model can optionally be configured with a `preprocessor` layer, in
+    which case it will automatically apply preprocessing to raw inputs during
+    `fit()`, `predict()`, and `evaluate()`. This is done by default when
+    creating the model with `from_preset()`.
+
+    Disclaimer: Pre-trained models are provided on an "as is" basis, without
+    warranties or conditions of any kind. The underlying model is provided by a
+    third party and subject to a separate license, available
+    [here](https://github.com/openai/gpt-2).
+
+    Args:
+        backbone: A `keras_nlp.models.GPT2Backbone` instance.
+        preprocessor: A `keras_nlp.models.GPT2CausalLMPreprocessor` or `None`.
+            If `None`, this model will not apply preprocessing, and inputs
+            should be preprocessed before calling the model.
+
+    Examples:
+
+    Use `generate()` method to do text generation.
+    ```python
+    gpt2_lm = keras_nlp.models.GPT2CausalLM.from_preset("gpt2_base_en")
+    gpt2_lm.generate("I want to say", max_length=30)
+
+    # Generate with batched prompts.
+    gpt2_lm.generate(["This is a", "Where are you"], max_length=30)
+    ```
+
+    Use a custom sampler for text generation.
+    ```python
+    gpt2_lm = keras_nlp.models.GPT2CausalLM.from_preset("gpt2_base_en")
+
+    # Use string identifier to set sampler.
+    gpt2_lm.generate("I want to say", max_length=30, sampler="top_p")
+
+    # Construct a sampler instance.
+    sampler = keras_nlp.samplers.BeamSampler(num_beams=2)
+    gpt2_lm.generate("I want to say", max_length=30, sampler=sampler)
+    ```
+
+    Map raw string to languages model logit predictions.
+    ```python
+    gpt2_lm = keras_nlp.models.GPT2CausalLM.from_preset("gpt2_base_en")
+    gpt2_lm.predict(["You know this is just a test string"])
+    ```
+
+    Load a pretrained GPT2 and fit on a string dataset.
+    ```python
+    features = [
+        "I don't listen to music while coding.",
+        "But I watch youtube while coding!",
+    ]
+    ds = tf.data.Dataset.from_tensor_slices(features)
+
+    # Create a `GPT2CausalLM` and fit your data.
+    gpt2_lm = keras_nlp.models.GPT2CausalLM.from_preset(
+        "gpt2_base_en",
+    )
+    gpt2_lm.compile(
+        loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+    )
+    gpt2_lm.fit(ds, batch_size=2)
+    ```
+
+    Load a pretrained `GPT2CausalLM` with custom preprocessor, and predict on
+    string inputs.
+    ```python
+    # Use a shorter sequence length.
+    preprocessor = keras_nlp.models.GPT2CausalLMPreprocessor.from_preset(
+        "gpt2_base_en",
+        sequence_length=128,
+    )
+
+    # Create a `GPT2CausalLM`, using pretrained GPT2 and custom preprocessor.
+    gpt2_lm = keras_nlp.models.GPT2CausalLM.from_preset(
+        "gpt2_base_en",
+        preprocessor=preprocessor,
+    )
+    gpt2_lm.predict(["You know this is still a test string"])
+    ```
+
+    Fit your preprocessed data with randomly initialized GPT2. This is useful
+    when you want to do data preprocessing inside `tf.data` pipeline.
+    ```python
+    # Define preprocessed input.
+    features = {
+        "token_ids": tf.constant(
+            [[1, 2, 3, 4, 0, 0]] * 2, shape=(2, 6)
+        ),
+        "padding_mask": tf.constant(
+            [[1, 1, 1, 1, 0, 0]] * 2, shape=(2, 6)
+        ),
+    }
+    labels = tf.constant(
+        [[2, 3, 4, 0, 0, 0]] * 2, shape=(2, 6)
+    )
+    sample_weight = tf.constant(
+        [[1, 1, 1, 0, 0, 0]] * 2, shape=(2, 6)
+    )
+
+    # Randomly initialize a GPT2 backbone.
+    backbone = keras_nlp.models.GPT2Backbone(
+        vocabulary_size=50257,
+        num_layers=2,
+        num_heads=2,
+        hidden_dim=128,
+        intermediate_dim=256,
+        max_sequence_length=128,
+    )
+    # Create a `GPT2CausalLM` without preprocessor and fit the data.
+    gpt2_lm = keras_nlp.models.GPT2CausalLM(backbone, preprocessor=None)
+    gpt2_lm.compile(
+        loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+    )
+    gpt2_lm.fit(
+        x=features,
+        y=labels,
+        sample_weight=sample_weight,
+        batch_size=2,
+    )
+    ```
+
+    """
+
+    def __init__(self, backbone, preprocessor=None, **kwargs):
+        inputs = backbone.input
+        x = backbone(inputs)
+        # Use token embedding weights to project from the token representation
+        # to vocabulary logits.
+        outputs = tf.matmul(
+            x,
+            backbone.token_embedding.embeddings,
+            transpose_b=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
+
+    @classproperty
+    def presets(cls):
+        return copy.deepcopy(backbone_presets)
+
+    @classproperty
+    def backbone_cls(cls):
+        return GPT2Backbone
+
+    @classproperty
+    def preprocessor_cls(cls):
+        return GPT2CausalLMPreprocessor
+
+    def _get_token_probability(self, prompt, mask):
+        model_inputs = {
+            "token_ids": prompt,
+            "padding_mask": mask,
+        }
+        return self(model_inputs)
+
+    def generate(
+        self,
+        prompt,
+        max_length,
+        sampler="top_k",
+    ):
+        """Generate text.
+
+        This method generates text based on given `prompt`. Generation will
+        continue until `max_length` is met, and all tokens generated after
+        `end_token` will be truncated. The sampling approach used can be
+        controlled via the sampler argument.
+
+        Args:
+            prompt: a string, string Tensor or string RaggedTensor. The prompt
+                text for generation.
+            max_length: int. The max length of generated sequence.
+            sampler: a string or `keras_nlp.samplers.Sampler` instance. The
+                sampler to be used for text generation.
+        """
+        end_token_id = self.preprocessor.tokenizer.end_token_id
+
+        sampler = keras_nlp.samplers.get(sampler)
+        if hasattr(self, "jit_compile"):
+            # `jit_compile` is a public property as of tf 2.12. hasattr is for
+            # backward compat.
+            sampler.jit_compile = self.jit_compile
+        sampler.run_eagerly = self.run_eagerly
+        generated = sampler(
+            self.preprocessor.tokenizer(prompt),
+            self._get_token_probability,
+            max_length=max_length,
+            end_token_id=end_token_id,
+        )
+        return self.preprocessor.tokenizer.detokenize(generated)
diff --git a/keras_nlp/models/gpt2/gpt2_causal_lm_preprocessor.py b/keras_nlp/models/gpt2/gpt2_causal_lm_preprocessor.py
new file mode 100644
index 00000000..803d111c
--- /dev/null
+++ b/keras_nlp/models/gpt2/gpt2_causal_lm_preprocessor.py
@@ -0,0 +1,96 @@
+# 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.
+
+"""GPT2 Causal LM preprocessor layer."""
+
+from absl import logging
+from tensorflow import keras
+
+from keras_nlp.models.gpt2.gpt2_preprocessor import GPT2Preprocessor
+from keras_nlp.utils.keras_utils import pack_x_y_sample_weight
+
+
+@keras.utils.register_keras_serializable(package="keras_nlp")
+class GPT2CausalLMPreprocessor(GPT2Preprocessor):
+    """GPT2 Causal LM preprocessor.
+
+    This preprocessor is majorly used as the preprocesor for `GPT2CausalLM`.
+    This class subclasses `keras_nlp.models.GPT2Preprocessor` and keeps most of
+    its functionality. The only change is `GPT2CausalLMPreprocessor` sets
+    `y` (label) and `sample_weights` field by shifting the input sequence one
+    step towards left, and drop the last token as it does not have a successor,
+    e.g., if the tokenized input is `[1, 2, 3, 0, 0]` with
+    `padding_mask = [1, 1, 1, 0, 0]`, then after preprocessing, we
+    will have `x = [1, 2, 3, 0]` and `y = [2, 3, 0, 0]`, with
+    `padding_mask = [1, 1, 1, 0]` and `sample_weights = [1, 1, 0, 0]`.
+
+    Args:
+        tokenizer: A `keras_nlp.models.GPT2Tokenizer` instance.
+        sequence_length: The length of the packed inputs.
+
+    Examples:
+    ```python
+    # Load the preprocessor from a preset.
+    preprocessor = keras_nlp.models.GPT2CausalLMPreprocessor.from_preset(
+        "gpt2_base_en"
+    )
+
+    # 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):
+        if y is not None or sample_weight is not None:
+            logging.warning(
+                "`GPT2CausalLMPreprocessor` 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."
+            )
+
+        x = super().call(x)
+        token_ids, padding_mask = x["token_ids"], x["padding_mask"]
+        # 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 = token_ids[..., 1:]
+        sample_weight = padding_mask[..., 1:]
+        return pack_x_y_sample_weight(x, y, sample_weight)
diff --git a/keras_nlp/models/gpt2/gpt2_preprocessor.py b/keras_nlp/models/gpt2/gpt2_preprocessor.py
new file mode 100644
index 00000000..5f6090a2
--- /dev/null
+++ b/keras_nlp/models/gpt2/gpt2_preprocessor.py
@@ -0,0 +1,172 @@
+# 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.
+
+"""GPT2 preprocessor layer."""
+
+import copy
+
+import tensorflow as tf
+from tensorflow import keras
+
+from keras_nlp.models.gpt2.gpt2_presets import backbone_presets
+from keras_nlp.models.gpt2.gpt2_tokenizer import GPT2Tokenizer
+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 GPT2Preprocessor(Preprocessor):
+    """GPT2 preprocessing layer which tokenizes and packs inputs.
+
+    This preprocessing layer will do 2 things:
+
+    - Tokenize the input using the `tokenizer`.
+    - Construct a dictionary with keys `"token_ids"`, `"padding_mask"`, that can
+        be passed directly to a `keras_nlp.models.GPT2Backbone`.
+
+    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.
+
+    `GPT2Preprocessor` forces the input to have only one segment, as GPT2 is
+    mainly used for generation tasks. For tasks having multi-segment inputs
+    like "glue/mnli", please use a model designed for classification purposes
+    such as BERT or RoBERTa.
+
+    Args:
+        tokenizer: A `keras_nlp.models.GPT2Tokenizer` instance.
+        sequence_length: The length of the packed inputs.
+
+    Examples:
+    ```python
+    # Load the preprocessor from a preset.
+    preprocessor = keras_nlp.models.GPT2Preprocessor.from_preset("gpt2_base_en")
+
+    # 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", "gi gi gi gi"])
+    preprocessor(sentences)
+    # Same output.
+    preprocessor(["taco tuesday", "gi gi gi gi"])
+
+    # 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)
+
+    # Alternatively, you can create a preprocessor from your own vocabulary.
+    # The usage is exactly the same as above.
+    vocab = {
+        "
": 0,
+        "<pad>": 1,
+        "
": 2,
+        "Ġafter": 5,
+        "noon": 6,
+        "Ġsun": 7,
+    }
+    merges = ["Ġ a", "Ġ s", "Ġ n", "e r", "n o", "o n", "Ġs u", "Ġa f", "no on"]
+    merges += ["Ġsu n", "Ġaf t", "Ġaft er"]
+
+    tokenizer = keras_nlp.models.GPT2Tokenizer(
+        vocabulary=vocab,
+        merges=merges,
+    )
+    preprocessor = keras_nlp.models.GPT2Preprocessor(
+        tokenizer=tokenizer,
+        sequence_length=20,
+    )
+    ```
+    """
+
+    def __init__(
+        self,
+        tokenizer,
+        sequence_length,
+        **kwargs,
+    ):
+
+        super().__init__(**kwargs)
+
+        self._tokenizer = tokenizer
+        self.sequence_length = sequence_length
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "sequence_length": self.sequence_length,
+            }
+        )
+        return config
+
+    def call(self, x, y=None, sample_weight=None):
+        x = convert_inputs_to_list_of_tensor_segments(x)
+        if len(x) > 1:
+            raise ValueError(
+                "GPT2 requires each input feature to contain only "
+                f"one segment, but received {len(x)}. If you are using GPT2 "
+                "for a multi-segment classification task, please refer to "
+                "classification models like BERT or RoBERTa."
+            )
+        token_ids = self.tokenizer(x[0])
+        input_is_1d = len(token_ids.shape) == 1
+        if input_is_1d:
+            token_ids = tf.RaggedTensor.from_tensor([token_ids])
+        mask = tf.ones_like(token_ids, dtype=tf.bool)
+        mask = mask.to_tensor(shape=(None, self.sequence_length))
+        token_ids = token_ids.to_tensor(shape=(None, self.sequence_length))
+        if input_is_1d:
+            # If the input is a single string, we let the output be a 1D tensor.
+            token_ids = tf.squeeze(token_ids, axis=0)
+            mask = tf.squeeze(mask, axis=0)
+        x = {
+            "token_ids": token_ids,
+            "padding_mask": mask,
+        }
+
+        return pack_x_y_sample_weight(x, y, sample_weight)
+
+    @classproperty
+    def presets(cls):
+        return copy.deepcopy(backbone_presets)
+
+    @classproperty
+    def tokenizer_cls(cls):
+        return GPT2Tokenizer
diff --git a/keras_nlp/samplers/sampler.py b/keras_nlp/samplers/sampler.py
index e11a13cb..dde87a66 100644
--- a/keras_nlp/samplers/sampler.py
+++ b/keras_nlp/samplers/sampler.py
@@ -57,6 +57,12 @@ class Sampler:
 
     The inputs and outputs of Sampler class are both token ids.
 
+    Subclassers should always implement the `get_next_token()` method, which
+    gets the next token based on probability distribution over vocab tokens.
+    Please check available subclass samplers for examples. If you need more
+    control over the sampling process, please implement `sample()` method
+    instead, see `keras_nlp.samplers.BeamSampler` for examples.
+
     Examples:
 
     Basic usage:
@@ -83,7 +89,7 @@ class Sampler:
 
     prompt = tf.fill((8, 1), 1)
 
-    sampler = keras_nlp.samplers.Greedy()
+    sampler = keras_nlp.samplers.GreedySampler()
     # Print the generated sequence (token ids).
     print(sampler(prompt, token_probability_fn, max_length=10, end_token_id=2))
     ```
@@ -114,7 +120,7 @@ class Sampler:
         return model(inputs)
 
     prompt = tokenizer("the quick brown fox")
-    sampler = keras_nlp.samplers.Greedy()
+    sampler = keras_nlp.samplers.GreedySampler()
     generated = sampler(
         prompt,
         token_probability_fn,
diff --git a/requirements-common.txt b/requirements-common.txt
index 0d4f7f89..07899151 100644
--- a/requirements-common.txt
+++ b/requirements-common.txt
@@ -3,7 +3,7 @@ astor
 numpy~=1.23.2  # Numpy 1.24 breaks tests on ragged tensors
 packaging
 black>=22
-black[juypter]
+black[jupyter]
 flake8
 isort
 pytest


Test patch:
diff --git a/keras_nlp/models/gpt2/gpt2_causal_lm_preprocessor_test.py b/keras_nlp/models/gpt2/gpt2_causal_lm_preprocessor_test.py
new file mode 100644
index 00000000..7ad2cce9
--- /dev/null
+++ b/keras_nlp/models/gpt2/gpt2_causal_lm_preprocessor_test.py
@@ -0,0 +1,127 @@
+# 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.
+
+"""Tests for GPT2 causal LM preprocessor layer."""
+
+import os
+
+import tensorflow as tf
+from absl.testing import parameterized
+from tensorflow import keras
+
+from keras_nlp.models.gpt2.gpt2_causal_lm_preprocessor import (
+    GPT2CausalLMPreprocessor,
+)
+from keras_nlp.models.gpt2.gpt2_tokenizer import GPT2Tokenizer
+
+
+class GPT2CausalLMPreprocessorTest(tf.test.TestCase, parameterized.TestCase):
+    def setUp(self):
+        vocab = {
+            "
": 0,
+            "!": 1,
+            "air": 2,
+            "Ġair": 3,
+            "plane": 4,
+            "Ġat": 5,
+            "port": 6,
+            "Ġkoh": 7,
+            "li": 8,
+            "Ġis": 9,
+            "Ġthe": 10,
+            "Ġbest": 11,
+        }
+
+        merges = ["Ġ a", "Ġ t", "Ġ k", "Ġ i", "Ġ b", "a i", "p l", "n e"]
+        merges += ["Ġa t", "p o", "r t", "o h", "l i", "Ġi s", "Ġb e", "s t"]
+        merges += ["Ġt h", "ai r", "pl a", "Ġk oh", "Ġth e", "Ġbe st", "po rt"]
+        merges += ["Ġai r", "Ġa i", "pla ne"]
+
+        self.preprocessor = GPT2CausalLMPreprocessor(
+            tokenizer=GPT2Tokenizer(
+                vocabulary=vocab,
+                merges=merges,
+            ),
+            sequence_length=8,
+        )
+
+    def test_tokenize_strings(self):
+        input_data = "airplane at airport"
+
+        x, y, sw = self.preprocessor(input_data)
+        self.assertAllEqual(x["token_ids"], [2, 4, 5, 3, 6, 0, 0])
+        self.assertAllEqual(x["padding_mask"], [1, 1, 1, 1, 1, 0, 0])
+        self.assertAllEqual(y, [4, 5, 3, 6, 0, 0, 0])
+        self.assertAllEqual(sw, [1, 1, 1, 1, 0, 0, 0])
+
+    def test_tokenize_list_of_strings(self):
+        input_data = ["airplane at airport"] * 4
+
+        x, y, sw = self.preprocessor(input_data)
+        self.assertAllEqual(
+            x["token_ids"],
+            [[2, 4, 5, 3, 6, 0, 0]] * 4,
+        )
+        self.assertAllEqual(x["padding_mask"], [[1, 1, 1, 1, 1, 0, 0]] * 4)
+        self.assertAllEqual(y, [[4, 5, 3, 6, 0, 0, 0]] * 4)
+        self.assertAllEqual(sw, [[1, 1, 1, 1, 0, 0, 0]] * 4)
+
+    def test_tokenize_labeled_batch(self):
+        x = tf.constant(["airplane at airport"] * 4)
+        y = tf.constant([1] * 4)
+        sw = tf.constant([1.0] * 4)
+        x, y, sw = self.preprocessor(x, y, sw)
+        self.assertAllEqual(
+            x["token_ids"],
+            [[2, 4, 5, 3, 6, 0, 0]] * 4,
+        )
+        self.assertAllEqual(x["padding_mask"], [[1, 1, 1, 1, 1, 0, 0]] * 4)
+        self.assertAllEqual(y, [[4, 5, 3, 6, 0, 0, 0]] * 4)
+        self.assertAllEqual(sw, [[1, 1, 1, 1, 0, 0, 0]] * 4)
+
+    def test_tokenize_labeled_dataset(self):
+        x = tf.constant(["airplane at airport"] * 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, y, sw = ds.batch(4).take(1).get_single_element()
+
+        self.assertAllEqual(
+            x["token_ids"],
+            [[2, 4, 5, 3, 6, 0, 0]] * 4,
+        )
+        self.assertAllEqual(x["padding_mask"], [[1, 1, 1, 1, 1, 0, 0]] * 4)
+        self.assertAllEqual(y, [[4, 5, 3, 6, 0, 0, 0]] * 4)
+        self.assertAllEqual(sw, [[1, 1, 1, 1, 0, 0, 0]] * 4)
+
+    @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(["airplane at airport"])
+
+        inputs = keras.Input(dtype="string", shape=())
+        outputs, y, sw = 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/gpt2/gpt2_causal_lm_test.py b/keras_nlp/models/gpt2/gpt2_causal_lm_test.py
new file mode 100644
index 00000000..8baf2ebb
--- /dev/null
+++ b/keras_nlp/models/gpt2/gpt2_causal_lm_test.py
@@ -0,0 +1,160 @@
+# 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.
+"""Tests for GPT2 causal LM model."""
+
+import os
+
+import tensorflow as tf
+from absl.testing import parameterized
+from tensorflow import keras
+
+from keras_nlp.models.gpt2.gpt2_backbone import GPT2Backbone
+from keras_nlp.models.gpt2.gpt2_causal_lm import GPT2CausalLM
+from keras_nlp.models.gpt2.gpt2_causal_lm_preprocessor import (
+    GPT2CausalLMPreprocessor,
+)
+from keras_nlp.models.gpt2.gpt2_tokenizer import GPT2Tokenizer
+
+
+class GPT2CausalLMTest(tf.test.TestCase, parameterized.TestCase):
+    def setUp(self):
+        backbone = GPT2Backbone(
+            vocabulary_size=1000,
+            num_layers=2,
+            num_heads=2,
+            hidden_dim=64,
+            intermediate_dim=128,
+            max_sequence_length=128,
+        )
+        vocab = {
+            "
": 0,
+            "!": 1,
+            "air": 2,
+            "Ġair": 3,
+            "plane": 4,
+            "Ġat": 5,
+            "port": 6,
+            "Ġkoh": 7,
+            "li": 8,
+            "Ġis": 9,
+            "Ġthe": 10,
+            "Ġbest": 11,
+        }
+
+        merges = ["Ġ a", "Ġ t", "Ġ k", "Ġ i", "Ġ b", "a i", "p l", "n e"]
+        merges += ["Ġa t", "p o", "r t", "o h", "l i", "Ġi s", "Ġb e", "s t"]
+        merges += ["Ġt h", "ai r", "pl a", "Ġk oh", "Ġth e", "Ġbe st", "po rt"]
+        merges += ["Ġai r", "Ġa i", "pla ne"]
+
+        self.preprocessor = GPT2CausalLMPreprocessor(
+            GPT2Tokenizer(vocabulary=vocab, merges=merges),
+            sequence_length=8,
+        )
+        self.causal_lm = GPT2CausalLM(
+            backbone,
+            preprocessor=self.preprocessor,
+        )
+        self.causal_lm_no_preprocessing = GPT2CausalLM(
+            backbone,
+            preprocessor=None,
+        )
+
+        self.raw_batch = tf.constant(
+            [
+                " airplane at airport",
+                " the airplane is the best",
+                " the best airport",
+                " kohli is the best",
+            ]
+        )
+        self.preprocessed_batch = self.preprocessor(self.raw_batch)[0]
+        self.raw_dataset = tf.data.Dataset.from_tensor_slices(
+            self.raw_batch
+        ).batch(2)
+        self.preprocessed_dataset = self.raw_dataset.map(self.preprocessor)
+
+    def test_valid_call_causal_lm(self):
+        self.causal_lm(self.preprocessed_batch)
+
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_gpt2_causal_lm_predict(self, jit_compile):
+        self.causal_lm.compile(jit_compile=jit_compile)
+        self.causal_lm.predict(self.raw_batch)
+
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_gpt2_causal_lm_predict_no_preprocessing(self, jit_compile):
+        self.causal_lm_no_preprocessing.compile(jit_compile=jit_compile)
+        self.causal_lm_no_preprocessing.predict(self.preprocessed_batch)
+
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_gpt2_causal_lm_fit(self, jit_compile):
+        self.causal_lm.compile(
+            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+            jit_compile=jit_compile,
+        )
+        self.causal_lm.fit(self.raw_dataset)
+
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_gpt2_causal_lm_fit_no_preprocessing(self, jit_compile):
+        self.causal_lm_no_preprocessing.compile(
+            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+            jit_compile=jit_compile,
+        )
+        self.causal_lm_no_preprocessing.fit(self.preprocessed_dataset)
+
+    @parameterized.named_parameters(
+        ("jit_compile_false", False), ("jit_compile_true", True)
+    )
+    def test_gpt2_causal_lm_generate(self, jit_compile):
+        self.causal_lm_no_preprocessing.compile(jit_compile=jit_compile)
+        self.causal_lm.generate(
+            self.raw_batch,
+            max_length=10,
+        )
+
+        # String input
+        prompt = " airplane"
+        generated = self.causal_lm.generate(
+            prompt,
+            max_length=10,
+        )
+        generated = generated.numpy().decode("utf-8")
+        self.assertTrue(prompt in generated)
+
+    @parameterized.named_parameters(
+        ("tf_format", "tf", "model"),
+        ("keras_format", "keras_v3", "model.keras"),
+    )
+    def test_saving_model(self, save_format, filename):
+        keras.utils.set_random_seed(42)
+        model_output = self.causal_lm.predict(self.raw_batch)
+        save_path = os.path.join(self.get_temp_dir(), filename)
+        self.causal_lm.save(save_path, save_format)
+        restored_model = keras.models.load_model(save_path)
+
+        # Check we got the real object back.
+        self.assertIsInstance(restored_model, GPT2CausalLM)
+
+        # Check that output matches.
+        keras.utils.set_random_seed(42)
+        restored_output = restored_model.predict(self.raw_batch)
+        self.assertAllClose(model_output, restored_output)
diff --git a/keras_nlp/models/gpt2/gpt2_preprocessor_test.py b/keras_nlp/models/gpt2/gpt2_preprocessor_test.py
new file mode 100644
index 00000000..d06f8f9f
--- /dev/null
+++ b/keras_nlp/models/gpt2/gpt2_preprocessor_test.py
@@ -0,0 +1,122 @@
+# 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.
+
+"""Tests for GPT2 preprocessor layer."""
+
+import os
+
+import tensorflow as tf
+from absl.testing import parameterized
+from tensorflow import keras
+
+from keras_nlp.models.gpt2.gpt2_preprocessor import GPT2Preprocessor
+from keras_nlp.models.gpt2.gpt2_tokenizer import GPT2Tokenizer
+
+
+class GPT2PreprocessorTest(tf.test.TestCase, parameterized.TestCase):
+    def setUp(self):
+        vocab = {
+            "
": 0,
+            "!": 1,
+            "air": 2,
+            "Ġair": 3,
+            "plane": 4,
+            "Ġat": 5,
+            "port": 6,
+            "Ġkoh": 7,
+            "li": 8,
+            "Ġis": 9,
+            "Ġthe": 10,
+            "Ġbest": 11,
+        }
+
+        merges = ["Ġ a", "Ġ t", "Ġ k", "Ġ i", "Ġ b", "a i", "p l", "n e"]
+        merges += ["Ġa t", "p o", "r t", "o h", "l i", "Ġi s", "Ġb e", "s t"]
+        merges += ["Ġt h", "ai r", "pl a", "Ġk oh", "Ġth e", "Ġbe st", "po rt"]
+        merges += ["Ġai r", "Ġa i", "pla ne"]
+
+        self.preprocessor = GPT2Preprocessor(
+            tokenizer=GPT2Tokenizer(
+                vocabulary=vocab,
+                merges=merges,
+            ),
+            sequence_length=8,
+        )
+
+    def test_tokenize_strings(self):
+        input_data = "airplane at airport"
+
+        output = self.preprocessor(input_data)
+        self.assertAllEqual(output["token_ids"], [2, 4, 5, 3, 6, 0, 0, 0])
+        self.assertAllEqual(output["padding_mask"], [1, 1, 1, 1, 1, 0, 0, 0])
+
+    def test_tokenize_list_of_strings(self):
+        input_data = ["airplane at airport"] * 4
+
+        output = self.preprocessor(input_data)
+        self.assertAllEqual(
+            output["token_ids"],
+            [[2, 4, 5, 3, 6, 0, 0, 0]] * 4,
+        )
+
+        self.assertAllEqual(
+            output["padding_mask"], [[1, 1, 1, 1, 1, 0, 0, 0]] * 4
+        )
+
+    def test_tokenize_labeled_batch(self):
+        x = tf.constant(["airplane at airport"] * 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"], [[2, 4, 5, 3, 6, 0, 0, 0]] * 4)
+        self.assertAllEqual(
+            x_out["padding_mask"], [[1, 1, 1, 1, 1, 0, 0, 0]] * 4
+        )
+        self.assertAllEqual(y_out, y)
+        self.assertAllEqual(sw_out, sw)
+
+    def test_tokenize_labeled_dataset(self):
+        x = tf.constant(["airplane at airport"] * 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"], [[2, 4, 5, 3, 6, 0, 0, 0]] * 4)
+        self.assertAllEqual(
+            x_out["padding_mask"], [[1, 1, 1, 1, 1, 0, 0, 0]] * 4
+        )
+        self.assertAllEqual(y_out, y)
+        self.assertAllEqual(sw_out, sw)
+
+    @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(["airplane at airport"])
+
+        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"],
+        )