Repository: keras-team/keras-nlp
Pull request #122 (2022-04-16 05:47:29)

Add ROUGE Metric
Resolves https://github.com/keras-team/keras-nlp/issues/67

@mattdangerw, @chenmoneygithub, this PR is now ready for review :)

mattdangerw: Thank you! This looks awesome!

Before diving into review could you share a colab showing the expected end to end use case for this with some string translations and references?

You should be able to add these lines to the top of the colab.

```
!git clone --banch your-branch-name your-remote-url
!cd keras-nlp && pip install . -q
```

It looks like we would need people to tokenize themselves before calling into this metric, which I think is ok (and maybe even preferable). But we should be clear what our expected end to end flow is. See as reference...

https://colab.sandbox.google.com/github/huggingface/notebooks/blob/master/course/videos/rouge_metric.ipynb

Maybe use `TextVectorization` for now to show how this could be done at a word level?

It would also be helpful to compare this argument list and docstring to the ones we expect for the Rouge-N variant. Could you show what the full arg list and usage for Rouge-N could look like? I know we are still figure out how to implement, so that part could be a markdown codeblock, not actually runnable.
abheesht17: Sure thing! Will share a notebook ASAP :)
abheesht17: @mattdangerw, here is the notebook: https://colab.research.google.com/drive/1X8ppbzcsBE1wKTB-XnfKhLRCPT6EHQTH?usp=sharing :)
mattdangerw: @abheesht17 thanks for the colab! That is super helpful. Thoughts...

I am still somewhat torn on whether we might want to somehow accept strings. The hugging face implementation is certainly a little more usable on the face of it, as is the [rouge-score package](https://github.com/google-research/google-research/tree/master/rouge#using-pip) it is based on. Though it seems like they will need to bake in a lot of assumptions (including language stemming!).

This brings up a lot of questions. Curious your thoughts here.
 - If you wanted to report rouge score, say for a paper, would you need to recreate the tokenization and stemming exactly like the package I linked, for comparability with other models? 
 - How could you do that from our package?
 - Do people regularly report rouge with other tokenizers?
 - How does people handle other languages with rogue generally (especially ones without whitespace splitting for tokens)?

Second, more minor point, why do you need to expand_dims on the first metric, could we remove that requirement or is that baked into keras.metrics somehow?
chenmoneygithub: @abheesht17 Hi, we had some discussions around this, and also reached out to Google research team for their insights. Briefly, we will have ROUGE metric working at string inputs, and by default provide a standard tokenizer so that different works would report ROUGE score based on the same calculation mechanism. 

To make our work compatible with existing ROUGE scores reporting, we can depend this metric on [this rouge_score package](https://github.com/google-research/google-research/tree/master/rouge). We still want to deliver Keras metrics so that ROUGE can be easily integrated into training flow. To do that, one possible way is to use [`tf.py_function()`](https://www.tensorflow.org/api_docs/python/tf/py_function) to wrap the rouge function call from existing `rouge_score` package. 

We can leave tokenizer customization a TODO until we know how to make our tokenizers compatible with `rouge_score` package. 

Sorry about the inconvenience! 
abheesht17: Hello, @chenmoneygithub! Thank you! Sorry for not replying earlier to @mattdangerw's questions.

So, essentially, the gist is:

1. We will still keep ROUGE-L as a subclass of keras.metrics.Metric. 
2. We will use [Google Research's ROUGE package](https://github.com/google-research/google-research/tree/master/rouge).
3. We will take string inputs. 
4. For future work, we can figure out how to make our tokenisers compatible with package's tokenisers and give an option to the user to pass an arg for that.

I had a brief look at the tokenisers provided by the package. Their default tokeniser merely returns a list of string tokens:

```
>>> df = DefaultTokenizer()
>>> df.tokenize("hello, this is fun")
['hello', 'this', 'is', 'fun']
```

If you see [this file](https://github.com/google-research/google-research/blob/master/rouge/tokenizers.py#L18-L20), we need to do two things to implement a custom tokeniser:
1. The tokeniser should be a subclass of the `Tokenizer` class present in the ROUGE package.
2. It should have a `tokenize()` method which returns a list of string tokens.

I don't think (1) is very important; after all, in the internal ROUGE package implementation, they will just call the tokenize method. See here: https://github.com/google-research/google-research/blob/master/rouge/rouge_scorer.py#L73 and https://github.com/google-research/google-research/blob/master/rouge/rouge_scorer.py#L125.

Currently, our tokenisers return a RaggedTensor on passing a string:

```
String output.
>>> vocab = ["[UNK]", "the", "qu", "##ick", "br", "##own", "fox", "."]
>>> inputs = "The quick brown fox."
>>> tokenizer = keras_nlp.tokenizers.WordPieceTokenizer(
...     vocabulary=vocab, dtype="string")
>>> tokenizer(inputs)
<tf.RaggedTensor [[b'the', b'qu', b'##ick', b'br', b'##own', b'fox', b'.']]>
```

So, I guess in the tokenizer class in KerasNLP, we can simply add an option, `return_list`, and if this is True, we can convert the RaggedTensor output to a list. I'll try this out.
abheesht17: Small update:

The functionality for `split_summaries` and providing a custom tokeniser was only recently added to the rouge-score package. The authors have not released a version with the above two functionalities. So, I'm just going for a basic implementation right now.

See these two commits:
1. https://github.com/google-research/google-research/commit/61ce9f0ca76025dac5b671c0631e443a9975a8a3 (custom tokeniser) 
2. https://github.com/google-research/google-research/commit/ed3e2bc5c85713225b3a0e75f7404d642274089f (split summaries)
abheesht17: @mattdangerw, @chenmoneygithub, how do I convert a string Tensor to a Python string (graph ops)? The rouge-score package does not accept string tensors (I tried passing a tensor with a single string).

```
>>> scorer.score(tf.constant("hello"), tf.constant("bye"))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "C:\Users\abheesht\anaconda3\envs\keras_nlp\lib\site-packages\rouge_score\rouge_scorer.py", line 88, in score
    target_tokens = tokenize.tokenize(target, self._stemmer)
  File "C:\Users\abheesht\anaconda3\envs\keras_nlp\lib\site-packages\rouge_score\tokenize.py", line 42, in tokenize
    text = text.lower()
  File "C:\Users\abheesht\anaconda3\envs\keras_nlp\lib\site-packages\tensorflow\python\framework\ops.py", line 513, in __getattr__
    self.__getattribute__(name)
AttributeError: 'tensorflow.python.framework.ops.EagerTensor' object has no attribute 'lower'
```

I tried a bunch of things, none of them seem to work. `str(tensor_var)` works in Eager mode, but won't work in graph mode. `tf.strings.as_string` did not work either.

```
>>> tf.strings.as_string(tf.constant("weifwf"))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "C:\Users\abheesht\anaconda3\envs\keras_nlp\lib\site-packages\tensorflow\python\ops\gen_string_ops.py", line 74, in as_string
    _ops.raise_from_not_ok_status(e, name)
  File "C:\Users\abheesht\anaconda3\envs\keras_nlp\lib\site-packages\tensorflow\python\framework\ops.py", line 7186, in raise_from_not_ok_status     
    raise core._status_to_exception(e) from None  # pylint: disable=protected-access
tensorflow.python.framework.errors_impl.InvalidArgumentError: Value for attr 'T' of string is not in the list of allowed values: float, double, int32, uint8, int16, int8, int64, bfloat16, uint16, half, uint32, uint64, complex64, complex128, bool, variant
        ; NodeDef: {{node AsString}}; Op<name=AsString; signature=input:T -> output:string; attr=T:type,allowed=[DT_FLOAT, DT_DOUBLE, DT_INT32, DT_UINT8, DT_INT16, DT_INT8, DT_INT64, DT_BFLOAT16, DT_UINT16, DT_HALF, DT_UINT32, DT_UINT64, DT_COMPLEX64, DT_COMPLEX128, DT_BOOL, DT_VARIANT]; attr=precision:int,default=-1; attr=scientific:bool,default=false; attr=shortest:bool,default=false; attr=width:int,default=-1; attr=fill:string,default=""> [Op:AsString]
```
mattdangerw: @abheesht17 re string tensors to string lists, we actually have a PR demonstrating this! https://github.com/keras-team/keras-nlp/pull/124

It is definitely more tricky than it should be. This is an issue at the tensorflow level.

For the tokenization question in general, it would be awesome if you could use any keras_nlp tokenizer (with any configuration), for rougue. So be able to handle int output (would the rouge_score package accept integer tokens?), as well as stripping padding tokens if present. But agreed with Chen above that this is a power user, nice to have feature.

The 95% use case will be using the package with it build it tokenization, to get comparable Rouge* score to other projects.
mattdangerw: Oh also missed the part re graph mode. Using py_function https://www.tensorflow.org/api_docs/python/tf/py_function to wrap your computation, should give you eager tensors, even if it is being called in a function. You could also do numpy_function, which will give numpy inputs directly.
abheesht17: @mattdangerw, thank you for the link! Will go through it.


The rouge-score package API for computing rouge score does not accept a list of tokens. It only accepts raw Python strings.
```
>>> from rouge_score import rouge_scorer
>>> scorer = rouge_scorer.RougeScorer(['rouge1', 'rougeL'], use_stemmer=True)
>>> scores = scorer.score('The quick brown fox jumps over the lazy dog',
...                       'The quick brown dog jumps on the log.')
>>> scores
{'rouge1': Score(precision=0.75, recall=0.6666666666666666, fmeasure=0.7058823529411765), 'rougeL': Score(precision=0.625, recall=0.5555555555555556, fmeasure=0.5882352941176471)}
```

```
>>> scores = scorer.score(['The', 'quick'], ['quick'])                         
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "C:\Users\abheesht\anaconda3\envs\keras_nlp\lib\site-packages\rouge_score\rouge_scorer.py", line 88, in score
    target_tokens = tokenize.tokenize(target, self._stemmer)
  File "C:\Users\abheesht\anaconda3\envs\keras_nlp\lib\site-packages\rouge_score\tokenize.py", line 42, in tokenize
    text = text.lower()
AttributeError: 'list' object has no attribute 'lower'
```

So, in order to use a tokeniser, we need to provide it as an arg to the rouge scorer function (https://github.com/google-research/google-research/blob/master/rouge/rouge_scorer.py#L73). However, unfortunately, this is not yet a part of a PyPi release: https://github.com/keras-team/keras-nlp/pull/122#issuecomment-1109956255. So, I am holding off on this for now.
chenmoneygithub: @abheesht17 We can stick to its basic functionality for now, and doing customized tokenizer support in a following PR.
mattdangerw: @abheesht17 one other note that might be helpful for this. I'm working on adding some additional notes on how we want to think about new dependencies here:
https://github.com/keras-team/keras-nlp/pull/158/files

I actually added a note there with this metric specifically as an example. Basically, we should not add this as a core dependency of keras-nlp, but rather import the library with a try/catch and give instructions on how to install deps. Something like...

```python
try:
    import rouge_score
except ImportError:
    pass


class RougeL(keras.metrics.Metric):
    def __init__(self):
        if rouge_score is None:
            raise ImportError(
                'RougeL metrics requires the rouge_score package. '
                '`pip install rouge-score`.')
```
abheesht17: @mattdangerw, @chenmoneygithub, I've made the required changes. Apologies once again for the delay!

Still confused about the graph ops stuff though. Since we are using tf.py_function, we can't really pass an object of this class to `model.compile`.

Patch:
diff --git a/API_DESIGN_GUIDE.md b/API_DESIGN_GUIDE.md
index 151463e4..e1536ff1 100644
--- a/API_DESIGN_GUIDE.md
+++ b/API_DESIGN_GUIDE.md
@@ -39,16 +39,20 @@ and add installation instructions for the specific symbol:
 try:
     import rouge_score
 except ImportError:
-    pass
+    rouge_score = None
 
-class RougeL(keras.metrics.Metric):
+class Rouge(keras.metrics.Metric):
     def __init__(self):
         if rouge_score is None:
             raise ImportError(
-                'RougeL metrics requires the rouge_score package. '
-                '`pip install rouge-score`.')
+                "ROUGE metric requires the `rouge_score` package."
+                "Please install it with `pip install rouge_score`."
+            )
 ```
 
+Additionally, to ensure that unit tests don't fail, please add the corresponding
+library to the `extras_require["tests"]` list in `setup.py`.
+
 ## Keep computation inside TensorFlow graph
 
 Our layers, metrics, and tokenizers should be fast and efficient, which means
diff --git a/README.md b/README.md
index 590c6ec5..c52da04c 100644
--- a/README.md
+++ b/README.md
@@ -22,9 +22,14 @@ We are a new and growing project, and welcome [contributions](CONTRIBUTING.md).
 
 ## Quick Links
 
+### For everyone
+
 - [Home Page](https://keras.io/keras_nlp)
 - [Developer Guides](https://keras.io/guides/keras_nlp)
 - [API Reference](https://keras.io/api/keras_nlp)
+
+### For contributors
+
 - [Contributing Guide](CONTRIBUTING.md)
 - [Roadmap](ROADMAP.md)
 - [Style Guide](STYLE_GUIDE.md)
@@ -71,7 +76,7 @@ outputs = keras.layers.Dense(1, activation="sigmoid")(outputs)
 model = keras.Model(inputs, outputs)
 
 # Run a single batch of gradient descent.
-model.compile(loss="binary_crossentropy", jit_compile=True)
+model.compile(optimizer="adam", loss="binary_crossentropy", jit_compile=True)
 model.train_on_batch(x, y)
 ```
 
@@ -100,4 +105,4 @@ Thank you to all of our wonderful contributors!
 
 <a href="https://github.com/keras-team/keras-nlp/graphs/contributors">
   <img src="https://contrib.rocks/image?repo=keras-team/keras-nlp" />
-</a>
\ No newline at end of file
+</a>
diff --git a/ROADMAP.md b/ROADMAP.md
index 49ed455d..4d92b2fe 100644
--- a/ROADMAP.md
+++ b/ROADMAP.md
@@ -16,10 +16,10 @@
   part of the library itself, but rather a way to vet our components and show
   best practices.
 
-- **A community of NLP practioners.** KerasNLP is a actively growing project,
+- **A community of NLP practitioners.** KerasNLP is an actively growing project,
   and we welcome contributors on all fronts of our development. We hope that our
   guides and examples can be both a valuable resource to experienced
-  practitioners, and a accessible entry point to newcomers to the field.
+  practitioners and an accessible entry point to newcomers to the field.
 
 ## What KerasNLP is not
 
@@ -49,14 +49,14 @@
   easy to train and use within the library. However, we need to support other
   types of models, such as n-gram or word2vec approaches that might run more
   easily on limited hardware. We will always want the most practical tool for
-  the task, regardless of architecture.
+  the task, regardless of the architecture.
 
 ## Focus areas for 2022
 
 At this point in our development cycle, we are primarily interested in providing
 building blocks for a short list of "greatest hits" NLP models (such as BERT,
 GPT-2, word2vec). Given a popular model architecture (e.g. a
-sequence-to-sequence transformer like T5) and a end-to-end task (e.g.
+sequence-to-sequence transformer like T5) and an end-to-end task (e.g.
 summarization), we should have a clear code example in mind and a list of
 components to use.
 
@@ -74,11 +74,11 @@ The major tasks within this effort:
   like to add support for byte-level BPE tokenization (the RoBERTa and GPT
   tokenizer) within the Tensorflow graph.
 - Pre-trained sub-word tokenizers for any language. Training a tokenizer can
-  add a lot of friction to a project, particularly when working working in a
+  add a lot of friction to a project, particularly when working in a
   language where examples are less readily available. We would like to support
   a pre-trained tokenization offering that allows a user to choose a tokenizer,
   language, and vocabulary size and then download an off the shelf vocabulary.
-- A standardized way to training tokenizer vocabularies. As another way to
+- A standardized way to train tokenizer vocabularies. As another way to
   reduce the friction of training a tokenizer, we should offer a standardized
   experience for training new vocabularies.
 
@@ -104,7 +104,7 @@ As we move forward with KerasNLP as a whole, we expect development for new
 components (say, a new attention mechanism) to happen in tandem with an
 example demonstrating the component in an end-to-end architecture.
 
-By the end of 2022, we should have a actively growing collection of examples
+By the end of 2022, we should have an actively growing collection of examples
 models, with a standardized set of training scripts, that match expected
 performance as reported in publications.
 
@@ -125,7 +125,7 @@ the end of 2022.
 We are looking to develop simple, modular components that make it easy to build
 end-to-end workflows for each of these tasks.
 
-Currently projects in this area include:
+Currently, projects in this area include:
 
 - Utilities for generating sequences of text using greedy or beam search.
 - Metrics for evaluating the quality of generated sequences, such a ROUGE and
diff --git a/examples/README.md b/examples/README.md
new file mode 100644
index 00000000..7acc255c
--- /dev/null
+++ b/examples/README.md
@@ -0,0 +1,60 @@
+# KerasNLP Example Models
+
+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)
diff --git a/examples/bert/README.md b/examples/bert/README.md
index 7913ea68..93ccbafc 100644
--- a/examples/bert/README.md
+++ b/examples/bert/README.md
@@ -1,14 +1,14 @@
-# BERT with keras-nlp
+# BERT with KerasNLP
 
-This example will show how to train a Bidirectional Encoder
-Representations from Transformers (BERT) model end-to-end using the keras-nlp
+This example demonstrates how to train a Bidirectional Encoder
+Representations from Transformers (BERT) model end-to-end using the KerasNLP
 library. This README contains instructions on how to run pretraining directly
-from raw data, followed by fine tuning and evaluation on the GLUE dataset.
+from raw data, followed by finetuning and evaluation on the GLUE dataset.
 
 ## Quickly test out the code
 
 To exercise the code in this directory by training a tiny BERT model, you can
-run the following commands from the base of the keras-nlp repository. This can
+run the following commands from the base directory of the repository. This can
 be useful to validate any code changes, but note that a useful BERT model would
 need to be trained for much longer on a much larger dataset.
 
@@ -18,47 +18,42 @@ DATA_URL=https://storage.googleapis.com/tensorflow/keras-nlp/examples/bert
 
 # Create a virtual env and install dependencies.
 mkdir $OUTPUT_DIR
-python3 -m venv $OUTPUT_DIR/env
-source $OUTPUT_DIR/env/bin/activate
+python3 -m venv $OUTPUT_DIR/env && source $OUTPUT_DIR/env/bin/activate
 pip install -e ".[tests,examples]"
 
 # Download example data.
 wget ${DATA_URL}/bert_vocab_uncased.txt -O $OUTPUT_DIR/bert_vocab_uncased.txt
 wget ${DATA_URL}/wiki_example_data.txt -O $OUTPUT_DIR/wiki_example_data.txt
 
-# Run preprocessing.
-python3 examples/bert/create_sentence_split_data.py \
+# Parse input data and split into sentences.
+python3 examples/tools/split_sentences.py \
     --input_files $OUTPUT_DIR/wiki_example_data.txt \
-    --output_directory $OUTPUT_DIR/sentence-split-data --num_shards 1
-python3 examples/bert/create_pretraining_data.py \
+    --output_directory $OUTPUT_DIR/sentence-split-data
+# Preprocess input for pretraining.
+python3 examples/bert/bert_preprocess.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.
-python3 examples/bert/run_pretraining.py \
+# Run pretraining for 100 train steps only.
+python3 examples/bert/bert_train.py \
     --input_files $OUTPUT_DIR/pretraining-data/ \
     --vocab_file $OUTPUT_DIR/bert_vocab_uncased.txt \
-    --bert_config_file examples/bert/configs/bert_tiny.json \
-    --num_warmup_steps 20 \
-    --num_train_steps 200 \
-    --saved_model_output $OUTPUT_DIR/model/
-
+    --saved_model_output $OUTPUT_DIR/model/ \
+    --num_train_steps 100
 # Run finetuning.
-python3 examples/bert/run_glue_finetuning.py \
+python3 examples/bert/bert_finetune_glue.py \
     --saved_model_input $OUTPUT_DIR/model/ \
-    --vocab_file $OUTPUT_DIR/bert_vocab_uncased.txt \
-    --bert_config_file examples/bert/configs/bert_tiny.json
+    --vocab_file $OUTPUT_DIR/bert_vocab_uncased.txt
 ```
 
 ## Installing dependencies
 
-Pip dependencies for all keras-nlp examples are listed in `setup.py`. To install
-both the keras-nlp library from source and all other dependencies required to
-run the example, run the below command. You may want to install to a self
-contained environment (e.g. a container or a virtualenv).
+Pip dependencies for all KerasNLP examples are listed in `setup.py`. The
+following command will create a virtual environment, install all dependencies,
+and install KerasNLP from source.
 
 ```shell
+python3 -m venv path/to/venv && source path/to/venv/bin/activate
 pip install -e ".[examples]"
 ```
 
@@ -66,7 +61,7 @@ pip install -e ".[examples]"
 
 Training a BERT model happens in two stages. First, the model is "pretrained" on
 a large corpus of input text. This is computationally expensive. After
-pretraining, the model can be "fine tuned" on a downstream task with much
+pretraining, the model can be "finetuned" on a downstream task with a much
 smaller amount of labeled data.
 
 ### Downloading pretraining data
@@ -75,7 +70,8 @@ The GLUE pretraining data (Wikipedia + BooksCorpus) is fairly large. The raw
 input data takes roughly ~20GB of space, and after preprocessing, the full
 corpus will take ~400GB.
 
-The latest wikipedia dump can be downloaded [at this link](https://dumps.wikimedia.org/enwiki/latest/enwiki-latest-pages-articles.xml.bz2),
+The latest wikipedia dump can be downloaded
+[at this link](https://dumps.wikimedia.org/enwiki/latest/enwiki-latest-pages-articles.xml.bz2),
 or via command line:
 
 ```shell
@@ -84,14 +80,14 @@ curl -O https://dumps.wikimedia.org/enwiki/latest/enwiki-latest-pages-articles.x
 The dump can be extracted with the `wikiextractor` tool.
 
 ```shell
-python -m wikiextractor.WikiExtractor enwiki-latest-pages-articles.xml.bz2
+python3 -m wikiextractor.WikiExtractor enwiki-latest-pages-articles.xml.bz2
 ```
 
 BooksCorpus is no longer hosted by
-[it's creators](https://yknzhu.wixsite.com/mbweb), but you can find instructions
-for downloading or reproducing the corpus in this
-[repository](https://github.com/soskek/bookcorpus). We suggest the pre-made file
-downloads listed at the top of the README. Alternatively, you can forgo it 
+[its creators](https://yknzhu.wixsite.com/mbweb), but you can find instructions
+for downloading or reproducing the corpus in
+[this repository](https://github.com/soskek/bookcorpus). We suggest the pre-made file
+downloads listed at the top of the README. Alternatively, you can forgo it
 entirely and pretrain solely on wikipedia.
 
 Preparing the pretraining data will happen in two stages. First, raw text needs
@@ -101,31 +97,27 @@ next sentence predictions.
 
 ### Splitting raw text into sentences
 
-The `create_sentence_split_data.py` will process raw input files and split them
-into output files where each line contains a sentence, and a blank line marks
-the start of a new document.
-
-The script supports two types of inputs files. Plain text files, where each
-individual file is assumed to be an entire document, and wikipedia dump files
-in the format outputted by the wikiextractor tool (each document is enclosed in
-`<doc>` tags).
+Next, use `examples/tools/split_sentences.py` to process raw input files and
+split them into output files where each line contains a sentence, and a blank
+line marks the start of a new document. We need this for the next-sentence
+prediction task used by BERT.
 
-For example, if wikipedia files are located in `~/datasets/wikipedia` and
+For example, if Wikipedia files are located in `~/datasets/wikipedia` and
 bookscorpus in `~/datasets/bookscorpus`, the following command will output
 sentence split documents to a configurable number of output file shards:
 
 ```shell
-python examples/bert/create_sentence_split_data.py \
+python3 examples/tools/split_sentences.py \
     --input_files ~/datasets/wikipedia,~/datasets/bookscorpus \
     --output_directory ~/datasets/sentence-split-data
 ```
 
 ### Computing a WordPiece vocabulary
 
-The `create_vocabulary.py` script allows you to compute your own WordPiece
-vocabulary for use with BERT. In most cases however, it is desirable to use the
-standard BERT vocabularies from the original models. You can download the
-English uncased vocabulary
+The easiest and best approach when training BERT is to use the official
+vocabularies from the original project, which have become somewhat standard.
+
+You can download the English uncased vocabulary
 [here](https://storage.googleapis.com/tensorflow/keras-nlp/examples/bert/bert_vocab_uncased.txt),
 or in your terminal run:
 
@@ -133,11 +125,12 @@ or in your terminal run:
 curl -O https://storage.googleapis.com/tensorflow/keras-nlp/examples/bert/bert_vocab_uncased.txt
 ```
 
+You can also use `examples/tools/train_word_piece_vocab.py` to train your own.
+
 ### Tokenize, mask, and combine sentences into training examples
 
-The `create_pretraining_data.py` scrip will take in a set of sentence split
-files, and set up training examples for the next sentence prediction and masked
-word tasks.
+The `bert_preprocess.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
 example. Below shows a complete output example with the addition of a string
@@ -172,21 +165,22 @@ with the following:
 ```shell
 for file in path/to/sentence-split-data/*; do
     output="path/to/pretraining-data/$(basename -- "$file" .txt).tfrecord"
-    python examples/bert/create_pretraining_data.py \
+    python3 examples/bert/bert_preprocess.py \
         --input_files ${file} \
         --vocab_file bert_vocab_uncased.txt \
         --output_file ${output}
 done
 ```
 
-If memory is available, this could be further sped up by running this script
-multiple times in parallel:
+If enough memory is available, this could be further sped up by running this script
+multiple times in parallel. The following will take 3-4 hours on the entire dataset
+on an 8 core machine.
 
 ```shell
 NUM_JOBS=5
 for file in path/to/sentence-split-data/*; do
     output="path/to/pretraining-data/$(basename -- "$file" .txt).tfrecord"
-    echo python examples/bert/create_pretraining_data.py \
+    echo python3 examples/bert/bert_preprocess.py \
         --input_files ${file} \
         --vocab_file bert_vocab_uncased.txt \
         --output_file ${output}
@@ -196,20 +190,19 @@ done | parallel -j ${NUM_JOBS}
 To preview a sample of generated data files, you can run the command below:
 
 ```shell
-python -c "from keras_nlp.utils.tensor_utils import preview_tfrecord; preview_tfrecord('/path/to/tfrecord_file')"
+python3 -c "from examples.utils.data_utils import preview_tfrecord; preview_tfrecord('path/to/tfrecord_file')"
 ```
 
 ### Running BERT pretraining
 
-After preprocessing, we can run pretraining with the `run_pretraining.py`
+After preprocessing, we can run pretraining with the `bert_train.py`
 script. This will train a model and save it to the `--saved_model_output`
 directory.
 
 ```shell
-python3 examples/bert/run_pretraining.py \
+python3 examples/bert/bert_train.py \
     --input_files path/to/data/ \
     --vocab_file path/to/bert_vocab_uncased.txt \
-    --bert_config_file examples/bert/configs/bert_tiny.json \
     --saved_model_output path/to/model/
 ```
 
@@ -217,18 +210,17 @@ python3 examples/bert/run_pretraining.py \
 
 After pretraining, we can evaluate the performance of a BERT model with the
 General Language Understanding Evaluation (GLUE) benchmark. This will
-fine tune the model and running classification for a number of downstream tasks.
+finetune the model and running classification for a number of downstream tasks.
 
-The `run_glue_finetuning.py` script downloads the GLUE data for a specific
+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/run_glue_finetuning.py \
+python3 examples/bert/bert_finetune_glue.py \
     --saved_model_input path/to/model/ \
-    --vocab_file path/to/bert_vocab_uncased.txt \
-    --bert_config_file examples/bert/configs/bert_tiny.json \
+    --vocab_file path/to/bert_vocab_uncased.txt
 ```
 
-The script could be easily adapted to any other text classification fine-tuning
+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/bert_config.py b/examples/bert/bert_config.py
new file mode 100644
index 00000000..80334ae6
--- /dev/null
+++ b/examples/bert/bert_config.py
@@ -0,0 +1,102 @@
+# 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.
+
+MODEL_CONFIGS = {
+    "tiny": {
+        "num_layers": 2,
+        "hidden_size": 128,
+        "hidden_dropout": 0.1,
+        "num_attention_heads": 2,
+        "attention_dropout": 0.1,
+        "inner_size": 512,
+        "inner_activation": "gelu",
+        "initializer_range": 0.02,
+    },
+    "mini": {
+        "num_layers": 4,
+        "hidden_size": 256,
+        "hidden_dropout": 0.1,
+        "num_attention_heads": 4,
+        "attention_dropout": 0.1,
+        "inner_size": 1024,
+        "inner_activation": "gelu",
+        "initializer_range": 0.02,
+    },
+    "small": {
+        "num_layers": 4,
+        "hidden_size": 512,
+        "hidden_dropout": 0.1,
+        "num_attention_heads": 8,
+        "attention_dropout": 0.1,
+        "inner_size": 2048,
+        "inner_activation": "gelu",
+        "initializer_range": 0.02,
+    },
+    "medium": {
+        "num_layers": 8,
+        "hidden_size": 512,
+        "hidden_dropout": 0.1,
+        "num_attention_heads": 8,
+        "attention_dropout": 0.1,
+        "inner_size": 2048,
+        "inner_activation": "gelu",
+        "initializer_range": 0.02,
+    },
+    "base": {
+        "num_layers": 12,
+        "hidden_size": 768,
+        "hidden_dropout": 0.1,
+        "num_attention_heads": 12,
+        "attention_dropout": 0.1,
+        "inner_size": 3072,
+        "inner_activation": "gelu",
+        "initializer_range": 0.02,
+    },
+    "large": {
+        "num_layers": 24,
+        "hidden_size": 1024,
+        "hidden_dropout": 0.1,
+        "num_attention_heads": 16,
+        "attention_dropout": 0.1,
+        "inner_size": 4096,
+        "inner_activation": "gelu",
+        "initializer_range": 0.02,
+    },
+}
+
+# Currently we have the same set of training parameters for all configurations.
+# We should see if we need to split this for different architecture sizes.
+
+PREPROCESSING_CONFIG = {
+    "max_seq_length": 512,
+    "max_predictions_per_seq": 76,
+    "dupe_factor": 10,
+    "masked_lm_prob": 0.15,
+    "short_seq_prob": 0.1,
+}
+
+TRAINING_CONFIG = {
+    "batch_size": 256,
+    "epochs": 10,
+    "learning_rate": 1e-4,
+    "num_train_steps": 1_000_000,
+    # 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/run_glue_finetuning.py b/examples/bert/bert_finetune_glue.py
similarity index 82%
rename from examples/bert/run_glue_finetuning.py
rename to examples/bert/bert_finetune_glue.py
index d15ecbbb..598a09fb 100644
--- a/examples/bert/run_glue_finetuning.py
+++ b/examples/bert/bert_finetune_glue.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 """Run finetuning on a GLUE task."""
 
-import json
+import tempfile
 
 import datasets
 import keras_tuner
@@ -23,54 +23,54 @@ from absl import app
 from absl import flags
 from tensorflow import keras
 
+from examples.bert.bert_config import FINETUNING_CONFIG
+from examples.bert.bert_config import MODEL_CONFIGS
+from examples.bert.bert_config import PREPROCESSING_CONFIG
+
 FLAGS = flags.FLAGS
 
 flags.DEFINE_string(
-    "bert_config_file",
-    None,
-    "The json config file for the bert model parameters.",
+    "model_size",
+    "tiny",
+    "One of: tiny, mini, small, medium, base, or large.",
 )
 
 flags.DEFINE_string(
     "vocab_file",
     None,
-    "The vocabulary file that the BERT model was trained on.",
+    "The vocabulary file for tokenization.",
 )
 
 flags.DEFINE_string(
     "saved_model_input",
     None,
-    "The directory containing the input pretrained model to finetune.",
+    "The directory to load the pretrained model.",
 )
 
 flags.DEFINE_string(
-    "saved_model_output", None, "The directory to save the finetuned model in."
+    "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."
+    "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. Should be True for uncased "
-    "models and False for cased models.",
+    "Whether to lower case the input text.",
 )
 
 flags.DEFINE_bool(
     "do_evaluation",
     True,
-    "Whether to run evaluation on the validation set for a given task.",
+    "Whether to run evaluation on test data.",
 )
 
-flags.DEFINE_integer("batch_size", 32, "The batch size.")
-
-flags.DEFINE_integer("epochs", 3, "The number of training epochs.")
-
-flags.DEFINE_integer("max_seq_length", 128, "Maximum sequence length.")
-
 
 def pack_inputs(
     inputs,
@@ -173,23 +173,25 @@ class BertClassificationFinetuner(keras.Model):
 class BertHyperModel(keras_tuner.HyperModel):
     """Creates a hypermodel to help with the search space for finetuning."""
 
-    def __init__(self, bert_config):
-        self.bert_config = bert_config
+    def __init__(self, model_config):
+        self.model_config = model_config
 
     def build(self, hp):
         model = keras.models.load_model(FLAGS.saved_model_input, compile=False)
-        bert_config = self.bert_config
+        model_config = self.model_config
         finetuning_model = BertClassificationFinetuner(
             bert_model=model,
-            hidden_size=bert_config["hidden_size"],
+            hidden_size=model_config["hidden_size"],
             num_classes=3 if FLAGS.task_name in ("mnli", "ax") else 2,
             initializer=keras.initializers.TruncatedNormal(
-                stddev=bert_config["initializer_range"]
+                stddev=model_config["initializer_range"]
             ),
         )
         finetuning_model.compile(
             optimizer=keras.optimizers.Adam(
-                learning_rate=hp.Choice("lr", [5e-5, 4e-5, 3e-5, 2e-5])
+                learning_rate=hp.Choice(
+                    "lr", FINETUNING_CONFIG["learning_rates"]
+                ),
             ),
             loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
             metrics=[keras.metrics.SparseCategoricalAccuracy()],
@@ -213,14 +215,13 @@ def main(_):
     end_id = vocab.index("[SEP]")
     pad_id = vocab.index("[PAD]")
 
-    with open(FLAGS.bert_config_file, "r") as bert_config_file:
-        bert_config = json.loads(bert_config_file.read())
+    model_config = MODEL_CONFIGS[FLAGS.model_size]
 
     def preprocess_data(inputs, labels):
         inputs = [tokenizer.tokenize(x).merge_dims(1, -1) for x in inputs]
         inputs = pack_inputs(
             inputs,
-            FLAGS.max_seq_length,
+            PREPROCESSING_CONFIG["max_seq_length"],
             start_of_sequence_id=start_id,
             end_of_segment_id=end_id,
             padding_id=pad_id,
@@ -230,18 +231,19 @@ def main(_):
     # Read and preprocess GLUE task data.
     train_ds, test_ds, validation_ds = load_data(FLAGS.task_name)
 
-    train_ds = train_ds.batch(FLAGS.batch_size).map(
+    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(FLAGS.batch_size).map(
+    validation_ds = validation_ds.batch(batch_size).map(
         preprocess_data, num_parallel_calls=tf.data.AUTOTUNE
     )
-    test_ds = test_ds.batch(FLAGS.batch_size).map(
+    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(bert_config)
+    hypermodel = BertHyperModel(model_config)
 
     # Initialize the random search over the 4 learning rate parameters, for 4
     # trials and 3 epochs for each trial.
@@ -251,9 +253,14 @@ def main(_):
         max_trials=4,
         overwrite=True,
         project_name="hyperparameter_tuner_results",
+        directory=tempfile.mkdtemp(),
     )
 
-    tuner.search(train_ds, epochs=FLAGS.epochs, validation_data=validation_ds)
+    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]
@@ -276,6 +283,5 @@ def main(_):
 
 if __name__ == "__main__":
     flags.mark_flag_as_required("vocab_file")
-    flags.mark_flag_as_required("bert_config_file")
     flags.mark_flag_as_required("saved_model_input")
     app.run(main)
diff --git a/examples/bert/create_pretraining_data.py b/examples/bert/bert_preprocess.py
similarity index 93%
rename from examples/bert/create_pretraining_data.py
rename to examples/bert/bert_preprocess.py
index 7bf07d37..36b6bc0e 100644
--- a/examples/bert/create_pretraining_data.py
+++ b/examples/bert/bert_preprocess.py
@@ -44,7 +44,8 @@ import tensorflow_text as tf_text
 from absl import app
 from absl import flags
 
-from examples.bert.bert_utils import list_filenames_for_arg
+from examples.bert.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.
 # Restrict the script to run CPU as GPU will not offer speedup here anyway.
@@ -55,7 +56,7 @@ FLAGS = flags.FLAGS
 flags.DEFINE_string(
     "input_files",
     None,
-    "Comma seperated list of directories, files, or globs.",
+    "Comma seperated list of directories, globs or files.",
 )
 
 flags.DEFINE_string(
@@ -67,39 +68,19 @@ flags.DEFINE_string(
 flags.DEFINE_string(
     "vocab_file",
     None,
-    "The vocabulary file that the BERT model was trained on.",
+    "The vocabulary file for tokenization.",
 )
 
 flags.DEFINE_bool(
     "do_lower_case",
     True,
-    "Whether to lower case the input text. Should be True for uncased "
-    "models and False for cased models.",
+    "Whether to lower case the input text.",
 )
 
-flags.DEFINE_integer("max_seq_length", 128, "Maximum sequence length.")
-
-flags.DEFINE_integer(
-    "max_predictions_per_seq",
-    20,
-    "Maximum number of masked LM predictions per sequence.",
-)
-
-flags.DEFINE_integer("random_seed", 12345, "Random seed for data generation.")
-
 flags.DEFINE_integer(
-    "dupe_factor",
-    10,
-    "Number of times to duplicate the input data (with different masks).",
-)
-
-flags.DEFINE_float("masked_lm_prob", 0.15, "Masked LM probability.")
-
-flags.DEFINE_float(
-    "short_seq_prob",
-    0.1,
-    "Probability of creating sequences which are shorter than the "
-    "maximum length.",
+    "random_seed",
+    12345,
+    "Random seed for data generation.",
 )
 
 
@@ -500,11 +481,11 @@ def main(_):
         input_filenames,
         tokenizer,
         vocab,
-        FLAGS.max_seq_length,
-        FLAGS.dupe_factor,
-        FLAGS.short_seq_prob,
-        FLAGS.masked_lm_prob,
-        FLAGS.max_predictions_per_seq,
+        PREPROCESSING_CONFIG["max_seq_length"],
+        PREPROCESSING_CONFIG["dupe_factor"],
+        PREPROCESSING_CONFIG["short_seq_prob"],
+        PREPROCESSING_CONFIG["masked_lm_prob"],
+        PREPROCESSING_CONFIG["max_predictions_per_seq"],
         rng,
     )
 
@@ -515,8 +496,8 @@ def main(_):
     write_instance_to_example_files(
         instances,
         vocab,
-        FLAGS.max_seq_length,
-        FLAGS.max_predictions_per_seq,
+        PREPROCESSING_CONFIG["max_seq_length"],
+        PREPROCESSING_CONFIG["max_predictions_per_seq"],
         FLAGS.output_file,
     )
 
diff --git a/examples/bert/run_pretraining.py b/examples/bert/bert_train.py
similarity index 77%
rename from examples/bert/run_pretraining.py
rename to examples/bert/bert_train.py
index 66768b02..77aca591 100644
--- a/examples/bert/run_pretraining.py
+++ b/examples/bert/bert_train.py
@@ -12,7 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import json
+import os
+import shutil
 import sys
 
 import tensorflow as tf
@@ -20,58 +21,54 @@ from absl import app
 from absl import flags
 from tensorflow import keras
 
+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.bert_model import BertModel
-from examples.bert.bert_utils import list_filenames_for_arg
+from examples.utils.scripting_utils import list_filenames_for_arg
 
 FLAGS = flags.FLAGS
 
 flags.DEFINE_string(
     "input_files",
     None,
-    "Comma seperated list of directories, files, or globs for input data.",
+    "Comma seperated list of directories, globs or files.",
 )
 
 flags.DEFINE_string(
-    "saved_model_output", None, "Output directory to save the model to."
-)
-
-flags.DEFINE_string(
-    "bert_config_file",
+    "saved_model_output",
     None,
-    "The json config file for the bert model parameters.",
+    "Output directory to save the model to.",
 )
 
 flags.DEFINE_string(
-    "vocab_file",
+    "checkpoint_save_directory",
     None,
-    "The vocabulary file that the BERT model was trained on.",
+    "Output directory to save checkpoints to.",
 )
 
-flags.DEFINE_integer("epochs", 10, "The number of training epochs.")
-
-flags.DEFINE_integer("batch_size", 256, "The training batch size.")
-
-flags.DEFINE_float("learning_rate", 1e-4, "The initial learning rate for Adam.")
-
-flags.DEFINE_integer("max_seq_length", 128, "Maximum sequence length.")
+flags.DEFINE_bool(
+    "skip_restore",
+    False,
+    "Skip restoring from checkpoint if True",
+)
 
-flags.DEFINE_integer(
-    "max_predictions_per_seq",
-    20,
-    "Maximum number of masked LM predictions per sequence.",
+flags.DEFINE_string(
+    "model_size",
+    "tiny",
+    "One of: tiny, mini, small, medium, base, or large.",
 )
 
-flags.DEFINE_integer(
-    "num_warmup_steps",
-    10000,
-    "The number of warmup steps during which the learning rate will increase "
-    "till a threshold.",
+flags.DEFINE_string(
+    "vocab_file",
+    None,
+    "The vocabulary file for tokenization.",
 )
 
 flags.DEFINE_integer(
     "num_train_steps",
-    1000000,
-    "The total fixed number of steps till which the model will train.",
+    None,
+    "Override the pre-configured number of train steps..",
 )
 
 
@@ -342,11 +339,18 @@ class LinearDecayWithWarmup(keras.optimizers.schedules.LearningRateSchedule):
             ),
         )
 
+    def get_config(self):
+        return {
+            "learning_rate": self.learning_rate,
+            "num_warmup_steps": self.warmup_steps,
+            "num_train_steps": self.train_steps,
+        }
+
 
 def decode_record(record):
     """Decodes a record to a TensorFlow example."""
-    seq_length = FLAGS.max_seq_length
-    lm_length = FLAGS.max_predictions_per_seq
+    seq_length = PREPROCESSING_CONFIG["max_seq_length"]
+    lm_length = PREPROCESSING_CONFIG["max_predictions_per_seq"]
     name_to_features = {
         "input_ids": tf.io.FixedLenFeature([seq_length], tf.int64),
         "input_mask": tf.io.FixedLenFeature([seq_length], tf.int64),
@@ -379,8 +383,17 @@ def main(_):
         for line in vocab_file:
             vocab.append(line.strip())
 
-    with open(FLAGS.bert_config_file, "r") as bert_config_file:
-        bert_config = json.loads(bert_config_file.read())
+    model_config = MODEL_CONFIGS[FLAGS.model_size]
+
+    if tf.config.list_logical_devices("TPU"):
+        # Connect to TPU and create TPU strategy.
+        resolver = tf.distribute.cluster_resolver.TPUClusterResolver.connect(
+            tpu="local"
+        )
+        strategy = tf.distribute.TPUStrategy(resolver)
+    else:
+        # Use default strategy if not using TPU.
+        strategy = tf.distribute.get_strategy()
 
     # Decode and batch data.
     dataset = tf.data.TFRecordDataset(input_filenames)
@@ -388,33 +401,65 @@ def main(_):
         lambda record: decode_record(record),
         num_parallel_calls=tf.data.experimental.AUTOTUNE,
     )
-    dataset = dataset.batch(FLAGS.batch_size, drop_remainder=True)
+    dataset = dataset.batch(TRAINING_CONFIG["batch_size"], drop_remainder=True)
     dataset = dataset.repeat()
 
-    # Create a BERT model the input config.
-    model = BertModel(
-        vocab_size=len(vocab),
-        **bert_config,
-    )
-    # Make sure model has been called.
-    model(model.inputs)
-    model.summary()
-
-    learning_rate_schedule = LinearDecayWithWarmup(
-        learning_rate=FLAGS.learning_rate,
-        num_warmup_steps=FLAGS.num_warmup_steps,
-        num_train_steps=FLAGS.num_train_steps,
-    )
+    with strategy.scope():
+        # Create a BERT model the input config.
+        model = BertModel(
+            vocab_size=len(vocab),
+            **model_config,
+        )
+        # Make sure model has been called.
+        model(model.inputs)
+        model.summary()
+
+        # Allow overriding train steps from the command line for quick testing.
+        if FLAGS.num_train_steps is not None:
+            num_train_steps = FLAGS.num_train_steps
+        else:
+            num_train_steps = TRAINING_CONFIG["num_train_steps"]
+        num_warmup_steps = int(
+            num_train_steps * TRAINING_CONFIG["warmup_percentage"]
+        )
+        learning_rate_schedule = LinearDecayWithWarmup(
+            learning_rate=TRAINING_CONFIG["learning_rate"],
+            num_warmup_steps=num_warmup_steps,
+            num_train_steps=num_train_steps,
+        )
+        optimizer = keras.optimizers.Adam(learning_rate=learning_rate_schedule)
+
+        pretraining_model = BertPretrainer(model)
+        pretraining_model.compile(
+            optimizer=optimizer,
+        )
+
+    epochs = TRAINING_CONFIG["epochs"]
+    steps_per_epoch = num_train_steps // epochs
+
+    callbacks = []
+    if FLAGS.checkpoint_save_directory is not None:
+        if os.path.exists(FLAGS.checkpoint_save_directory):
+            if not os.path.isdir(FLAGS.checkpoint_save_directory):
+                raise ValueError(
+                    "`checkpoint_save_directory` should be a directory, but "
+                    f"{FLAGS.checkpoint_save_directory} is not a directory."
+                    " Please set `checkpoint_save_directory` as a directory."
+                )
+
+            elif FLAGS.skip_restore:
+                # Clear up the directory if users want to skip restoring.
+                shutil.rmtree(FLAGS.checkpoint_save_directory)
+        checkpoint_path = FLAGS.checkpoint_save_directory + "/checkpoint"
+        callbacks.append(
+            tf.keras.callbacks.BackupAndRestore(backup_dir=checkpoint_path)
+        )
 
-    # Wrap with pretraining heads and call fit.
-    pretraining_model = BertPretrainer(model)
-    pretraining_model.compile(
-        optimizer=keras.optimizers.Adam(learning_rate=learning_rate_schedule)
-    )
-    # TODO(mattdangerw): Add TPU strategy support.
-    steps_per_epoch = FLAGS.num_train_steps // FLAGS.epochs
     pretraining_model.fit(
-        dataset, epochs=FLAGS.epochs, steps_per_epoch=steps_per_epoch
+        dataset,
+        epochs=epochs,
+        steps_per_epoch=steps_per_epoch,
+        callbacks=callbacks,
     )
 
     print(f"Saving to {FLAGS.saved_model_output}")
@@ -424,6 +469,5 @@ def main(_):
 if __name__ == "__main__":
     flags.mark_flag_as_required("input_files")
     flags.mark_flag_as_required("vocab_file")
-    flags.mark_flag_as_required("bert_config_file")
     flags.mark_flag_as_required("saved_model_output")
     app.run(main)
diff --git a/examples/bert/configs/bert_base.json b/examples/bert/configs/bert_base.json
deleted file mode 100644
index 80eb2621..00000000
--- a/examples/bert/configs/bert_base.json
+++ /dev/null
@@ -1,11 +0,0 @@
-{
-  "num_layers": 12,
-  "hidden_size": 768,
-  "hidden_dropout": 0.1,
-  "num_attention_heads": 12,
-  "attention_dropout": 0.1,
-  "inner_size": 3072,
-  "inner_activation": "gelu",
-  "initializer_range": 0.02,
-  "max_sequence_length": 512
-}
diff --git a/examples/bert/configs/bert_large.json b/examples/bert/configs/bert_large.json
deleted file mode 100644
index 1638feb0..00000000
--- a/examples/bert/configs/bert_large.json
+++ /dev/null
@@ -1,11 +0,0 @@
-{
-  "num_layers": 24,
-  "hidden_size": 1024,
-  "hidden_dropout": 0.1,
-  "num_attention_heads": 16,
-  "attention_dropout": 0.1,
-  "inner_size": 4096,
-  "inner_activation": "gelu",
-  "initializer_range": 0.02,
-  "max_sequence_length": 512
-}
diff --git a/examples/bert/configs/bert_medium.json b/examples/bert/configs/bert_medium.json
deleted file mode 100644
index af12f066..00000000
--- a/examples/bert/configs/bert_medium.json
+++ /dev/null
@@ -1,11 +0,0 @@
-{
-  "num_layers": 8,
-  "hidden_size": 512,
-  "hidden_dropout": 0.1,
-  "num_attention_heads": 8,
-  "attention_dropout": 0.1,
-  "inner_size": 2048,
-  "inner_activation": "gelu",
-  "initializer_range": 0.02,
-  "max_sequence_length": 512
-}
diff --git a/examples/bert/configs/bert_mini.json b/examples/bert/configs/bert_mini.json
deleted file mode 100644
index 8fd4fa7e..00000000
--- a/examples/bert/configs/bert_mini.json
+++ /dev/null
@@ -1,11 +0,0 @@
-{
-  "num_layers": 4,
-  "hidden_size": 256,
-  "hidden_dropout": 0.1,
-  "num_attention_heads": 4,
-  "attention_dropout": 0.1,
-  "inner_size": 1024,
-  "inner_activation": "gelu",
-  "initializer_range": 0.02,
-  "max_sequence_length": 512
-}
diff --git a/examples/bert/configs/bert_small.json b/examples/bert/configs/bert_small.json
deleted file mode 100644
index 1d744e87..00000000
--- a/examples/bert/configs/bert_small.json
+++ /dev/null
@@ -1,11 +0,0 @@
-{
-  "num_layers": 4,
-  "hidden_size": 512,
-  "hidden_dropout": 0.1,
-  "num_attention_heads": 8,
-  "attention_dropout": 0.1,
-  "inner_size": 2048,
-  "inner_activation": "gelu",
-  "initializer_range": 0.02,
-  "max_sequence_length": 512
-}
diff --git a/examples/bert/configs/bert_tiny.json b/examples/bert/configs/bert_tiny.json
deleted file mode 100644
index 3521a015..00000000
--- a/examples/bert/configs/bert_tiny.json
+++ /dev/null
@@ -1,11 +0,0 @@
-{
-  "num_layers": 2,
-  "hidden_size": 128,
-  "hidden_dropout": 0.1,
-  "num_attention_heads": 2,
-  "attention_dropout": 0.1,
-  "inner_size": 512,
-  "inner_activation": "gelu",
-  "initializer_range": 0.02,
-  "max_sequence_length": 512
-}
diff --git a/examples/tools/README.md b/examples/tools/README.md
new file mode 100644
index 00000000..efe525e9
--- /dev/null
+++ b/examples/tools/README.md
@@ -0,0 +1,37 @@
+# KerasNLP Modeling Tools
+
+This directory contains runnable scripts that are not specific to a specific
+model architecture, but still useful for end-to-end workflows.
+
+## split_sentences.py
+
+The `split_sentences.py` script will process raw input files and split them into
+output files where each line contains a sentence, and a blank line marks the
+start of a new document. This is useful for tasks like next sentence prediction
+where the boundaries between sentences are needed for training.
+
+The script supports two types of inputs files. Plain text files, where each
+individual file is assumed to be an entire document, and wikipedia dump files
+in the format outputted by the wikiextractor tool (each document is enclosed in
+`<doc>` tags).
+
+Example usage:
+
+```shell
+python examples/tools/split_sentences.py \
+    --input_files ~/datasets/wikipedia,~/datasets/bookscorpus \
+    --output_directory ~/datasets/sentence-split-data
+```
+
+### train_word_piece_vocabulary.py
+
+The `train_word_piece_vocabulary.py` script allows you to compute your own
+WordPiece vocabulary.
+
+Example usage:
+
+```shell
+python examples/tools/train_word_piece_vocabulary.py \
+    --input_files ~/datasets/my-raw-dataset/ \
+    --output_file vocab.txt
+```
diff --git a/examples/tools/__init__.py b/examples/tools/__init__.py
new file mode 100644
index 00000000..6e4df4e7
--- /dev/null
+++ b/examples/tools/__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/examples/bert/create_sentence_split_data.py b/examples/tools/split_sentences.py
similarity index 97%
rename from examples/bert/create_sentence_split_data.py
rename to examples/tools/split_sentences.py
index e11818a3..373ab6a4 100644
--- a/examples/bert/create_sentence_split_data.py
+++ b/examples/tools/split_sentences.py
@@ -22,7 +22,7 @@ This script will run muliprocessed, and the number of concurrent process and
 output file shards can be controlled with `--num_jobs` and `--num_shards`.
 
 Usage:
-python create_sentence_split_data.py \
+python examples/tools/create_sentence_split_data.py \
     --input_files ~/datasets/wikipedia,~/datasets/bookscorpus \
     --output_directory ~/datasets/bert-sentence-split-data
 """
@@ -38,7 +38,7 @@ from absl import app
 from absl import flags
 from tensorflow import keras
 
-from examples.bert.bert_utils import list_filenames_for_arg
+from examples.utils.scripting_utils import list_filenames_for_arg
 
 FLAGS = flags.FLAGS
 
diff --git a/examples/bert/create_vocabulary.py b/examples/tools/train_word_piece_vocab.py
similarity index 96%
rename from examples/bert/create_vocabulary.py
rename to examples/tools/train_word_piece_vocab.py
index 34529f6a..d5cbff16 100644
--- a/examples/bert/create_vocabulary.py
+++ b/examples/tools/train_word_piece_vocab.py
@@ -16,7 +16,7 @@
 This script will create wordpiece vocabularies suitable for pretraining BERT.
 
 Usage:
-python create_vocabulary.py \
+python examples/tools/train_word_piece_vocabulary.py \
     --input_files ~/datasets/bert-sentence-split-data/ \
     --output_file vocab.txt
 """
@@ -29,7 +29,7 @@ from absl import app
 from absl import flags
 from tensorflow_text.tools.wordpiece_vocab import bert_vocab_from_dataset
 
-from examples.bert.bert_utils import list_filenames_for_arg
+from examples.utils.scripting_utils import list_filenames_for_arg
 
 FLAGS = flags.FLAGS
 
diff --git a/examples/utils/__init__.py b/examples/utils/__init__.py
new file mode 100644
index 00000000..6e4df4e7
--- /dev/null
+++ b/examples/utils/__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/examples/utils/data_utils.py b/examples/utils/data_utils.py
new file mode 100644
index 00000000..b9f50af6
--- /dev/null
+++ b/examples/utils/data_utils.py
@@ -0,0 +1,30 @@
+# 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.
+"""Utility functions for data handling."""
+
+import os
+
+import tensorflow as tf
+from google import protobuf
+
+
+def preview_tfrecord(filepath):
+    """Pretty prints a single record from a tfrecord file."""
+    dataset = tf.data.TFRecordDataset(os.path.expanduser(filepath))
+    example = tf.train.Example()
+    example.ParseFromString(next(iter(dataset)).numpy())
+    formatted = protobuf.text_format.MessageToString(
+        example, use_short_repeated_primitives=True
+    )
+    print(formatted)
diff --git a/examples/bert/bert_utils.py b/examples/utils/scripting_utils.py
similarity index 95%
rename from examples/bert/bert_utils.py
rename to examples/utils/scripting_utils.py
index 83eb6ed7..586a5c12 100644
--- a/examples/bert/bert_utils.py
+++ b/examples/utils/scripting_utils.py
@@ -11,7 +11,7 @@
 # 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.
-"""Utility files for BERT scripts."""
+"""Utility functions for writing training scripts."""
 
 import glob
 import os
diff --git a/keras_nlp/layers/__init__.py b/keras_nlp/layers/__init__.py
index ec1ea1ad..a4cc2508 100644
--- a/keras_nlp/layers/__init__.py
+++ b/keras_nlp/layers/__init__.py
@@ -15,8 +15,10 @@
 from keras_nlp.layers.fnet_encoder import FNetEncoder
 from keras_nlp.layers.mlm_head import MLMHead
 from keras_nlp.layers.mlm_mask_generator import MLMMaskGenerator
+from keras_nlp.layers.multi_segment_packer import MultiSegmentPacker
 from keras_nlp.layers.position_embedding import PositionEmbedding
 from keras_nlp.layers.sine_position_encoding import SinePositionEncoding
+from keras_nlp.layers.start_end_packer import StartEndPacker
 from keras_nlp.layers.token_and_position_embedding import (
     TokenAndPositionEmbedding,
 )
diff --git a/keras_nlp/layers/multi_segment_packer.py b/keras_nlp/layers/multi_segment_packer.py
new file mode 100644
index 00000000..a6cb001d
--- /dev/null
+++ b/keras_nlp/layers/multi_segment_packer.py
@@ -0,0 +1,223 @@
+# 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.
+
+"""BERT token packing layer."""
+
+import tensorflow as tf
+import tensorflow_text as tf_text
+from tensorflow import keras
+
+
+class MultiSegmentPacker(keras.layers.Layer):
+    """Packs multiple sequences into a single fixed width model input.
+
+    This layer packs multiple input sequences into a single fixed width sequence
+    containing start and end delimeters, forming an dense input suitable for a
+    classification task for BERT and BERT-like models.
+
+    Takes as input a list or tuple of token segments. The layer will process
+    inputs as follows:
+     - Truncate all input segments to fit within `sequence_length` according to
+       the `truncator` strategy.
+     - Concatenate all input segments, adding a single `start_value` at the
+       start of the entire sequence, and multiple `end_value`s at the end of
+       each segment.
+     - Pad the resulting sequence to `sequence_length` using `pad_tokens`.
+     - Calculate a separate tensor of "segment ids", with integer type and the
+       same shape as the packed token output, where each integer index of the
+       segment the token originated from. The segment id of the `start_value`
+       is always 0, and the segment id of each `end_value` is the segment that
+       precedes it.
+
+    If inputs are batched, inputs should be `tf.RaggedTensor`s with shape
+    `[batch_size, None]` and will be packed and converted to a dense tensor with
+    shape `[batch_size, sequence_length]`.
+
+    If inputs are unbatched, inputs should be dense rank-1 tensors of any shape,
+    and will be packed to shape `[sequence_length]`.
+
+    Args:
+        sequence_length: The desired output length.
+        start_value: The id or token that is to be placed at the start of each
+            sequence (called "[CLS]" for BERT). The dtype must mach the dtype of
+            the input tensors to the layer.
+        end_value: The id or token that is to be placed at the end of each
+            input segment (called "[SEP]" for BERT). The dtype much mach the
+            dtype of the input tensors to the layer.
+        pad_value: The id or token that is to be placed into the unused
+            positions after the last segment in the sequence
+            (called "[PAD]" for BERT).
+        truncator: The algorithm to truncate a list of batched segments to fit a
+            per-example length limit. 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 support arbitrary number of segments.
+
+    Returns:
+        A tuple with two elements. The first is the dense, packed token
+        sequence. The second is an integer tensor of the same shape, containing
+        the segment ids.
+
+    Examples:
+
+    *Pack a single input for classification.*
+    >>> seq1 = tf.constant([1, 2, 3, 4])
+    >>> packer = keras_nlp.layers.MultiSegmentPacker(
+    ...     8, start_value=101, end_value=102)
+    >>> packer(seq1)
+    (<tf.Tensor: shape=(8,), dtype=int32,
+        numpy=array([101, 1, 2, 3, 4, 102, 0, 0], dtype=int32)>,
+     <tf.Tensor: shape=(8,), dtype=int32,
+        numpy=array([0, 0, 0, 0, 0, 0, 0, 0], dtype=int32)>)
+
+    *Pack multiple inputs for classification.*
+    >>> seq1 = tf.constant([1, 2, 3, 4])
+    >>> seq2 = tf.constant([11, 12, 13, 14])
+    >>> packer = keras_nlp.layers.MultiSegmentPacker(
+    ...     8, start_value=101, end_value=102)
+    >>> packer((seq1, seq2))
+    (<tf.Tensor: shape=(8,), dtype=int32,
+        numpy=array([101,   1,   2,   3, 102,  11,  12, 102], dtype=int32)>,
+     <tf.Tensor: shape=(8,), dtype=int32,
+        numpy=array([0, 0, 0, 0, 0, 1, 1, 1], dtype=int32)>)
+
+    Reference:
+        [Devlin et al., 2018](https://arxiv.org/abs/1810.04805).
+    """
+
+    def __init__(
+        self,
+        sequence_length,
+        start_value,
+        end_value,
+        pad_value=None,
+        truncator="round_robin",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.sequence_length = sequence_length
+        if truncator not in ("round_robin", "waterfall"):
+            raise ValueError(
+                "Only 'round_robin' and 'waterfall' algorithms are "
+                "supported. Received %s" % truncator
+            )
+        self.truncator = truncator
+        self.start_value = start_value
+        self.end_value = end_value
+        self.pad_value = pad_value
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "sequence_length": self.sequence_length,
+                "start_value": self.start_value,
+                "end_value": self.end_value,
+                "pad_value": self.pad_value,
+                "truncator": self.truncator,
+            }
+        )
+        return config
+
+    def _sanitize_inputs(self, inputs):
+        """Force inputs to a list of rank 2 ragged tensors."""
+        # Sanitize inputs.
+        if not isinstance(inputs, (list, tuple)):
+            inputs = [inputs]
+        if not inputs:
+            raise ValueError("At least one input is required for packing")
+        input_ranks = [x.shape.rank for x in inputs]
+        if not all(0 < rank < 3 for rank in input_ranks):
+            raise ValueError(
+                "All inputs for packing must have rank 1 or 2. "
+                f"Received input ranks: {input_ranks}"
+            )
+        if None in input_ranks or len(set(input_ranks)) > 1:
+            raise ValueError(
+                "All inputs for packing must have the same rank. "
+                f"Received input ranks: {input_ranks}"
+            )
+        return inputs
+
+    def _trim_inputs(self, inputs):
+        """Trim inputs to desired length."""
+        num_special_tokens = len(inputs) + 1
+        if self.truncator == "round_robin":
+            return tf_text.RoundRobinTrimmer(
+                self.sequence_length - num_special_tokens
+            ).trim(inputs)
+        elif self.truncator == "waterfall":
+            return tf_text.WaterfallTrimmer(
+                self.sequence_length - num_special_tokens
+            ).trim(inputs)
+        else:
+            raise ValueError("Unsupported truncator: %s" % self.truncator)
+
+    def _combine_inputs(self, segments):
+        """Combine inputs with start and end values added."""
+        dtype = segments[0].dtype
+        start_value = tf.convert_to_tensor(self.start_value, dtype=dtype)
+        end_value = tf.convert_to_tensor(self.end_value, dtype=dtype)
+
+        start_column = tf.tile([start_value], [segments[0].nrows()])
+        start_column = tf.expand_dims(start_column, 1)
+        end_column = tf.tile([end_value], [segments[0].nrows()])
+        end_column = tf.expand_dims(end_column, 1)
+        ones_column = tf.ones_like(start_column, dtype=tf.int32)
+
+        segments_to_combine = [start_column]
+        segment_ids_to_combine = [ones_column * 0]
+        for i, seg in enumerate(segments):
+            # Combine all segments adding end tokens.
+            segments_to_combine.append(seg)
+            segments_to_combine.append(end_column)
+
+            # Combine segment ids accounting for end tokens.
+            segment_ids_to_combine.append(tf.ones_like(seg, dtype=tf.int32) * i)
+            segment_ids_to_combine.append(ones_column * i)
+
+        token_ids = tf.concat(segments_to_combine, 1)
+        segment_ids = tf.concat(segment_ids_to_combine, 1)
+        return token_ids, segment_ids
+
+    def call(self, inputs):
+        inputs = self._sanitize_inputs(inputs)
+
+        # If rank 1, add a batch dim and convert to ragged.
+        rank_1 = inputs[0].shape.rank == 1
+        if rank_1:
+            inputs = [tf.expand_dims(x, 0) for x in inputs]
+            inputs = [tf.RaggedTensor.from_tensor(x) for x in inputs]
+
+        segments = self._trim_inputs(inputs)
+        token_ids, segment_ids = self._combine_inputs(segments)
+
+        # Pad to dense tensor output.
+        shape = tf.cast([-1, self.sequence_length], "int64")
+        token_ids = token_ids.to_tensor(
+            shape=shape, default_value=self.pad_value
+        )
+        segment_ids = segment_ids.to_tensor(shape=shape)
+
+        # Remove the batch dim if added.
+        if rank_1:
+            token_ids = tf.squeeze(token_ids, 0)
+            segment_ids = tf.squeeze(segment_ids, 0)
+
+        return (token_ids, segment_ids)
diff --git a/keras_nlp/layers/start_end_packer.py b/keras_nlp/layers/start_end_packer.py
new file mode 100644
index 00000000..704f4892
--- /dev/null
+++ b/keras_nlp/layers/start_end_packer.py
@@ -0,0 +1,166 @@
+# 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.
+
+"""Start End Packer implementation based on `keras.layers.Layer`."""
+
+import tensorflow as tf
+from tensorflow import keras
+
+
+class StartEndPacker(keras.layers.Layer):
+    """Adds start and end tokens to a sequence and pads to a fixed length.
+
+    This layer is useful when tokenizing inputs for tasks like translation,
+    where each sequence should include a start and end marker. It should
+    be called after tokenization. The layer will first trim inputs to fit, then
+    add start/end tokens, and finally pad, if necessary, to `sequence_length`.
+
+    If input is batched, input should be a `tf.RaggedTensor` with shape
+    `[batch_size, None]` and will be packed and converted to a dense tensor with
+    shape `[batch_size, sequence_length]`.
+    If input is unbatched, input should be a dense rank-1 tensor of any shape,
+    and will be packed to shape `[sequence_length]`.
+
+    Args:
+        sequence_length: int. The desired output length.
+        start_value: int/str. The ID or token that is to be placed at the start
+            of each sequence. The dtype must match the dtype of the input
+            tensors to the layer. If None, no start value will be added.
+        end_value: int/str. The ID or token that is to be placed at the end of
+            each input segment. The dtype must match the dtype of the input
+            tensors to the layer. If None, no end value will be added.
+        pad_value: int/str. The ID or token that is to be placed into the
+            unused positions after the last segment in the sequence. If None,
+            0 or "" will be added depending on the dtype of the input tensor.
+
+    Examples:
+
+    Unbatched input (int).
+    >>> input_data = tf.constant([5, 6, 7])
+    >>> start_end_packer = keras_nlp.layers.StartEndPacker(
+    ...     sequence_length=7, start_value=1, end_value=2,
+    ... )
+    >>> start_end_packer(input_data)
+    <tf.Tensor: shape=(7,), dtype=int32, numpy=
+    array([1, 5, 6, 7, 2, 0, 0], dtype=int32)>
+
+    Batched input (int).
+    >>> input_data = tf.ragged.constant([[5, 6, 7], [8, 9, 10, 11, 12, 13, 14]])
+    >>> start_end_packer = keras_nlp.layers.StartEndPacker(
+    ...     sequence_length=6, start_value=1, end_value=2,
+    ... )
+    >>> start_end_packer(input_data)
+    <tf.Tensor: shape=(2, 6), dtype=int32, numpy=
+    array([[ 1,  5,  6,  7,  2,  0],
+           [ 1,  8,  9, 10, 11,  2]], dtype=int32)>
+
+    Unbatched input (str).
+    >>> input_data = tf.constant(["this", "is", "fun"])
+    >>> start_end_packer = keras_nlp.layers.StartEndPacker(
+    ...     sequence_length=6, start_value="
", end_value="
",
+    ...     pad_value="<pad>"
+    ... )
+    >>> start_end_packer(input_data)
+    <tf.Tensor: shape=(6,), dtype=string, numpy=
+    array([b'
', b'this', b'is', b'fun', b'
', b'<pad>'], dtype=object)>
+
+    Batched input (str).
+    >>> input_data = tf.ragged.constant([["this", "is", "fun"], ["awesome"]])
+    >>> start_end_packer = keras_nlp.layers.StartEndPacker(
+    ...     sequence_length=6, start_value="
", end_value="
",
+    ...     pad_value="<pad>"
+    ... )
+    >>> start_end_packer(input_data)
+    <tf.Tensor: shape=(2, 6), dtype=string, numpy=
+    array([[b'
', b'this', b'is', b'fun', b'
', b'<pad>'],
+           [b'
', b'awesome', b'
', b'<pad>', b'<pad>', b'<pad>']],
+          dtype=object)>
+    """
+
+    def __init__(
+        self,
+        sequence_length,
+        start_value=None,
+        end_value=None,
+        pad_value=None,
+        name=None,
+        **kwargs,
+    ):
+        super().__init__(name=name, **kwargs)
+
+        self.sequence_length = sequence_length
+        self.start_value = start_value
+        self.end_value = end_value
+        self.pad_value = pad_value
+
+    def call(self, inputs):
+        if not isinstance(inputs, (tf.Tensor, tf.RaggedTensor)):
+            inputs = tf.convert_to_tensor(inputs)
+
+        input_is_dense = isinstance(inputs, tf.Tensor)
+        input_is_ragged = isinstance(inputs, tf.RaggedTensor)
+
+        if input_is_dense:
+            if inputs.shape.rank != 1:
+                raise ValueError(
+                    "Input must either be dense with rank 1, or ragged with "
+                    "rank 2. Received dense input with "
+                    f"rank={inputs.shape.rank}"
+                )
+
+            # Add a new axis at the beginning and convert to ragged tensor.
+            inputs = tf.RaggedTensor.from_tensor(tf.expand_dims(inputs, axis=0))
+        elif input_is_ragged:
+            if inputs.shape.rank != 2:
+                raise ValueError(
+                    "Input must either be dense with rank 1, or ragged with "
+                    "rank 2. Received ragged input with "
+                    f"rank={inputs.shape.rank}"
+                )
+
+        batch_size = tf.shape(inputs)[0]
+
+        # Concatenate start and end tokens.
+        if self.start_value is not None:
+            start_token_id_tensor = tf.fill((batch_size, 1), self.start_value)
+            inputs = tf.concat([start_token_id_tensor, inputs], axis=-1)
+        if self.end_value is not None:
+            end_token_id_tensor = tf.fill((batch_size, 1), self.end_value)
+
+            # Trim to leave room for end token.
+            inputs = inputs[..., : self.sequence_length - 1]
+            inputs = tf.concat([inputs, end_token_id_tensor], axis=-1)
+
+        # Pad to desired length.
+        inputs = inputs.to_tensor(
+            default_value=self.pad_value,
+            shape=(batch_size, self.sequence_length),
+        )
+
+        if input_is_dense:
+            inputs = tf.squeeze(inputs, axis=0)
+
+        return inputs
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "sequence_length": self.sequence_length,
+                "start_value": self.start_value,
+                "end_value": self.end_value,
+                "pad_value": self.pad_value,
+            }
+        )
+        return config
diff --git a/keras_nlp/layers/token_and_position_embedding.py b/keras_nlp/layers/token_and_position_embedding.py
index e2607ee6..cd21d40c 100644
--- a/keras_nlp/layers/token_and_position_embedding.py
+++ b/keras_nlp/layers/token_and_position_embedding.py
@@ -80,15 +80,22 @@ class TokenAndPositionEmbedding(keras.layers.Layer):
         self.vocabulary_size = int(vocabulary_size)
         self.sequence_length = int(sequence_length)
         self.embedding_dim = int(embedding_dim)
+        self.embeddings_initializer = keras.initializers.get(
+            embeddings_initializer
+        )
         self.token_embedding = keras.layers.Embedding(
             vocabulary_size,
             embedding_dim,
-            embeddings_initializer=embeddings_initializer,
+            embeddings_initializer=self.embeddings_initializer,
             mask_zero=mask_zero,
+            name="token_embedding"
+            + str(keras.backend.get_uid("token_embedding")),
         )
         self.position_embedding = keras_nlp.layers.PositionEmbedding(
             sequence_length=sequence_length,
-            initializer=embeddings_initializer,
+            initializer=self.embeddings_initializer,
+            name="position_embedding"
+            + str(keras.backend.get_uid("position_embedding")),
         )
         self.supports_masking = self.token_embedding.supports_masking
 
@@ -100,7 +107,7 @@ class TokenAndPositionEmbedding(keras.layers.Layer):
                 "sequence_length": self.sequence_length,
                 "embedding_dim": self.embedding_dim,
                 "embeddings_initializer": keras.initializers.serialize(
-                    self.token_embedding.embeddings_initializer
+                    self.embeddings_initializer
                 ),
                 "mask_zero": self.token_embedding.mask_zero,
             },
diff --git a/keras_nlp/layers/transformer_decoder.py b/keras_nlp/layers/transformer_decoder.py
index f8cdd105..c71d7283 100644
--- a/keras_nlp/layers/transformer_decoder.py
+++ b/keras_nlp/layers/transformer_decoder.py
@@ -30,12 +30,23 @@ class TransformerDecoder(keras.layers.Layer):
     paper [Attention is All You Need](https://arxiv.org/abs/1706.03762). Users
     can instantiate multiple instances of this class to stack up a decoder.
 
+    This layer will always apply a causal mask to the decoder attention layer.
     This layer will correctly compute an attention mask from an implicit
     Keras padding mask (for example, by passing `mask_zero=True` to a
     `keras.layers.Embedding` layer). See the Masking and Padding
     [guide](https://keras.io/guides/understanding_masking_and_padding/)
     for more details.
 
+    This layer can be called with either one or two inputs. The number of inputs
+    must be consistent across all calls. The options are as follows:
+        `layer(decoder_sequence)`: no cross-attention will be built into the
+            decoder block. This is useful when building a "decoder-only"
+            transformer such as GPT-2.
+        `layer(decoder_sequence, encoder_sequence)`: cross-attention will be
+            built into the decoder block. This is useful when building an
+            "encoder-decoder" transformer, such as the original transformer
+            model described in Attention is All You Need.
+
     Args:
         intermediate_dim: int, the hidden size of feedforward network.
         num_heads: int, the number of heads in MultiHeadAttention.
@@ -102,7 +113,7 @@ class TransformerDecoder(keras.layers.Layer):
         self._built = False
         self.supports_masking = True
 
-    def _build(self, input_shape):
+    def _build(self, input_shape, include_cross_attention):
         # Create layers based on input shape.
         self._built = True
         feature_size = input_shape[-1]
@@ -115,29 +126,37 @@ class TransformerDecoder(keras.layers.Layer):
             kernel_initializer=self.kernel_initializer,
             bias_initializer=self.bias_initializer,
         )
-        self._encoder_decoder_attention_layer = keras.layers.MultiHeadAttention(
-            num_heads=self.num_heads,
-            key_dim=self._attention_head_size,
-            value_dim=feature_size,
-            dropout=self.dropout,
-            kernel_initializer=self.kernel_initializer,
-            bias_initializer=self.bias_initializer,
-        )
 
         self._decoder_attention_layernorm = keras.layers.LayerNormalization(
             epsilon=self.layer_norm_epsilon,
         )
-        self._enc_dec_attention_layernorm = keras.layers.LayerNormalization(
-            epsilon=self.layer_norm_epsilon,
-        )
+
+        self._cross_attention_layer = None
+        if include_cross_attention:
+            # Create layers for cross attention.
+            self._cross_attention_layer = keras.layers.MultiHeadAttention(
+                num_heads=self.num_heads,
+                key_dim=self._attention_head_size,
+                value_dim=feature_size,
+                dropout=self.dropout,
+                kernel_initializer=self.kernel_initializer,
+                bias_initializer=self.bias_initializer,
+            )
+
+            self._cross_attention_layernorm = keras.layers.LayerNormalization(
+                epsilon=self.layer_norm_epsilon,
+            )
+
+            self._cross_attention_dropout = keras.layers.Dropout(
+                rate=self.dropout,
+            )
+
         self._feedforward_layernorm = keras.layers.LayerNormalization(
             epsilon=self.layer_norm_epsilon,
         )
 
         self._self_attention_dropout = keras.layers.Dropout(rate=self.dropout)
-        self._enc_dec_attentiondropout = keras.layers.Dropout(
-            rate=self.dropout,
-        )
+
         # First dense layer in the feedforward network, which maps input
         # feauture size to dimension `self.intermediate_dim`.
         self._intermediate_dense = keras.layers.Dense(
@@ -166,18 +185,20 @@ class TransformerDecoder(keras.layers.Layer):
     def call(
         self,
         decoder_sequence,
-        encoder_sequence,
+        encoder_sequence=None,
         decoder_padding_mask=None,
         decoder_attention_mask=None,
         encoder_padding_mask=None,
         encoder_attention_mask=None,
-        use_causal_mask=False,
     ):
         """Forward pass of the TransformerDecoder.
 
         Args:
             decoder_sequence: a Tensor. The decoder input sequence.
-            encoder_sequence: a Tensor. The decoder input sequence.
+            encoder_sequence: a Tensor. The encoder input sequence. For decoder
+                only models (like GPT2), this should be left None. Once the
+                model is called once without an encoder_sequence, you cannot
+                call it again with encoder_sequence.
             decoder_padding_mask: a boolean Tensor, the padding mask of decoder
                 sequence, must of shape [batch_size, decoder_sequence_length].
             decoder_attention_mask: a boolean Tensor. Customized decoder
@@ -188,29 +209,45 @@ class TransformerDecoder(keras.layers.Layer):
             encoder_attention_mask: a boolean Tensor. Customized encoder
                 sequence mask, must of shape
                 [batch_size, encoder_sequence_length, encoder_sequence_length].
-            use_causal_mask: bool, defaults to False. If true, causal mask
-                (masking out future input) is applied on the decoder sequence.
-
         Returns:
             A Tensor of the same shape as the `decoder_sequence`.
         """
+        has_encoder_sequence = encoder_sequence is not None
         if not self._built:
-            self._build(decoder_sequence.shape)
-        encoder_mask = merge_padding_and_attention_mask(
-            encoder_sequence, encoder_padding_mask, encoder_attention_mask
-        )
+            self._build(decoder_sequence.shape, has_encoder_sequence)
+
+        is_cross_attention = self._cross_attention_layer is not None
+        if not is_cross_attention and has_encoder_sequence:
+            raise ValueError(
+                "The number of call arguments to "
+                "`keras_nlp.layers.TransformerDecoder` should not change. "
+                "Use `layer(decoder_sequence, encoder_sequence)` to "
+                "build a layer with cross attention, or "
+                "`layer(decoder_sequence)` to build a layer without. "
+                "This layer has been built without cross attention, but "
+                "you are trying to call it with encoder_sequence."
+            )
+        elif is_cross_attention and not has_encoder_sequence:
+            raise ValueError(
+                "The number of call arguments to "
+                "`keras_nlp.layers.TransformerDecoder` should not change. "
+                "Use `layer(decoder_sequence, encoder_sequence)` to "
+                "build a layer with cross attention, or "
+                "`layer(decoder_sequence)` to build a layer without. "
+                "This layer has been built with cross attention, but "
+                "you did not provide encoder_sequence."
+            )
         decoder_mask = merge_padding_and_attention_mask(
             decoder_sequence, decoder_padding_mask, decoder_attention_mask
         )
-        if use_causal_mask:
-            causal_mask = tf.cast(
-                compute_causal_mask(decoder_sequence),
-                dtype=tf.int32,
-            )
-            if decoder_mask is None:
-                decoder_mask = causal_mask
-            else:
-                decoder_mask = tf.minimum(decoder_mask, causal_mask)
+        causal_mask = tf.cast(
+            compute_causal_mask(decoder_sequence),
+            dtype=tf.int32,
+        )
+        if decoder_mask is None:
+            decoder_mask = causal_mask
+        else:
+            decoder_mask = tf.minimum(decoder_mask, causal_mask)
 
         # Decoder input self-attention.
         self_attended = self._self_attention_layer(
@@ -220,28 +257,34 @@ class TransformerDecoder(keras.layers.Layer):
             attention_mask=decoder_mask,
         )
         self_attended = self._self_attention_dropout(self_attended)
-        self_attended = self._add_and_norm(
+        attention_output = self._add_and_norm(
             self_attended, decoder_sequence, self._decoder_attention_layernorm
         )
-        # Encoder-decoder attention.
-        encoder_decoder_attended = self._encoder_decoder_attention_layer(
-            query=self_attended,
-            value=encoder_sequence,
-            key=encoder_sequence,
-            attention_mask=encoder_mask,
-        )
-        encoder_decoder_attended = self._enc_dec_attentiondropout(
-            encoder_decoder_attended,
-        )
-        encoder_decoder_attended = self._add_and_norm(
-            encoder_decoder_attended,
-            self_attended,
-            self._enc_dec_attention_layernorm,
-        )
+
+        if self._cross_attention_layer is not None:
+            encoder_mask = merge_padding_and_attention_mask(
+                encoder_sequence, encoder_padding_mask, encoder_attention_mask
+            )
+            # Cross attention.
+            cross_attended = self._cross_attention_layer(
+                query=attention_output,
+                value=encoder_sequence,
+                key=encoder_sequence,
+                attention_mask=encoder_mask,
+            )
+            cross_attended = self._cross_attention_dropout(
+                cross_attended,
+            )
+            attention_output = self._add_and_norm(
+                cross_attended,
+                attention_output,
+                self._cross_attention_layernorm,
+            )
+
         # Feedforward.
-        feed_forward_output = self._feed_forward(encoder_decoder_attended)
+        feed_forward_output = self._feed_forward(attention_output)
         return self._add_and_norm(
-            encoder_decoder_attended,
+            attention_output,
             feed_forward_output,
             self._feedforward_layernorm,
         )
diff --git a/keras_nlp/metrics/__init__.py b/keras_nlp/metrics/__init__.py
index 7152a970..55ade6dc 100644
--- a/keras_nlp/metrics/__init__.py
+++ b/keras_nlp/metrics/__init__.py
@@ -13,3 +13,5 @@
 # limitations under the License.
 
 from keras_nlp.metrics.perplexity import Perplexity
+from keras_nlp.metrics.rouge_l import RougeL
+from keras_nlp.metrics.rouge_n import RougeN
diff --git a/keras_nlp/metrics/rouge_base.py b/keras_nlp/metrics/rouge_base.py
new file mode 100644
index 00000000..22d4adf3
--- /dev/null
+++ b/keras_nlp/metrics/rouge_base.py
@@ -0,0 +1,223 @@
+# 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.
+
+"""ROUGE metric implementation based on `keras.metrics.Metric`."""
+
+
+import types
+
+import tensorflow as tf
+from tensorflow import keras
+
+from keras_nlp.utils.tensor_utils import tensor_to_string_list
+
+try:
+    import rouge_score
+    from rouge_score import rouge_scorer
+except ImportError:
+    rouge_score = None
+
+
+class RougeBase(keras.metrics.Metric):
+    """ROUGE metric.
+
+    This class implements two variants of the ROUGE metric - ROUGE-N,
+    and ROUGE-L.
+
+    Note on input shapes:
+    For `y_true` and `y_pred`, this class supports scalar values and batch
+    inputs of shapes `()`, `(batch_size,)` and `(batch_size, 1)`.
+
+    Args:
+        variant: string. One of "rougeN", "rougeL". Defaults to
+            "rouge2". For "rougeN", N lies in the range [1, 9].
+        use_stemmer: bool. Whether Porter Stemmer should be used to strip word
+            suffixes to improve matching. Defaults to False.
+        dtype: string or tf.dtypes.Dtype. Precision of metric computation. If
+            not specified, it defaults to tf.float32.
+        name: string. Name of the metric instance.
+        **kwargs: Other keyword arguments.
+    """
+
+    def __init__(
+        self,
+        variant="rouge2",
+        use_stemmer=False,
+        dtype=None,
+        name="rouge",
+        **kwargs,
+    ):
+        super().__init__(name=name, dtype=dtype, **kwargs)
+
+        if rouge_score is None:
+            raise ImportError(
+                "ROUGE metric requires the `rouge_score` package. "
+                "Please install it with `pip install rouge-score`."
+            )
+
+        if not tf.as_dtype(self.dtype).is_floating:
+            raise ValueError(
+                "`dtype` must be a floating point type. "
+                f"Received: dtype={dtype}"
+            )
+
+        if variant not in tuple(
+            ("rouge" + str(order) for order in range(1, 10))
+        ) + ("rougeL",):
+            raise ValueError(
+                "Invalid variant of ROUGE. Should be one of: rougeN, rougeL, "
+                "with N ranging from 1 to 9. Received: "
+                f"variant={variant}"
+            )
+
+        self.variant = variant
+        self.use_stemmer = use_stemmer
+
+        # To-do: Add split_summaries and tokenizer options after the maintainers
+        # of rouge_scorer have released a new version.
+        self._rouge_scorer = rouge_scorer.RougeScorer(
+            rouge_types=[self.variant],
+            use_stemmer=use_stemmer,
+        )
+
+        self._rouge_precision = self.add_weight(
+            name="rouge_precision",
+            initializer="zeros",
+            dtype=self.dtype,
+        )
+        self._rouge_recall = self.add_weight(
+            name="rouge_recall",
+            initializer="zeros",
+            dtype=self.dtype,
+        )
+        self._rouge_f1_score = self.add_weight(
+            name="rouge_f1_score",
+            initializer="zeros",
+            dtype=self.dtype,
+        )
+
+        self._number_of_samples = self.add_weight(
+            name="number_of_samples", initializer="zeros", dtype=self.dtype
+        )
+
+    def __new__(cls, *args, **kwargs):
+        # Temporary workaround for Keras bug with dictionary return types.
+        # Wraps `result()` with a python dictionary that also supports variable
+        # assignment. We have to do this with __new__ because the base metric
+        # class wraps the `results()` method.
+        # TODO: Remove this snippet of code once the Keras bug is fixed.
+        obj = super().__new__(cls)
+
+        class MetricDict(dict):
+            """A dictionary that supports variable assignment."""
+
+            pass
+
+        def wrap_result(result_fn):
+            return tf.__internal__.decorator.make_decorator(
+                result_fn, lambda obj, *args: MetricDict(result_fn(*args))
+            )
+
+        obj.result = types.MethodType(wrap_result(obj.result), obj)
+        return obj
+
+    def update_state(self, y_true, y_pred, sample_weight=None):
+        # Three possible shapes for y_true and y_pred: Python string,
+        # [batch_size] and [batch_size, 1]. In the latter two cases, we have
+        # strings in the tensor/list.
+
+        def validate_and_fix_rank(inputs, tensor_name):
+            if not isinstance(inputs, tf.Tensor):
+                inputs = tf.convert_to_tensor(inputs)
+
+            if inputs.shape.rank == 0:
+                return inputs[tf.newaxis]
+            elif inputs.shape.rank == 1:
+                return inputs
+            elif inputs.shape.rank == 2:
+                if inputs.shape[1] != 1:
+                    raise ValueError(
+                        f"{tensor_name} must be of shape `[batch_size, 1]`. "
+                        f"Found shape: {inputs.shape}"
+                    )
+                else:
+                    return tf.squeeze(inputs, axis=1)
+            else:
+                raise ValueError(
+                    f"{tensor_name} must be of rank 0 (scalar input), 1 or 2. "
+                    f"Found rank: {inputs.shape.rank}"
+                )
+
+        y_true = validate_and_fix_rank(y_true, "y_true")
+        y_pred = validate_and_fix_rank(y_pred, "y_pred")
+
+        batch_size = tf.shape(y_true)[0]
+
+        def calculate_rouge_score(reference, hypothesis):
+            reference = tensor_to_string_list(reference)
+            hypothesis = tensor_to_string_list(hypothesis)
+            score = self._rouge_scorer.score(reference, hypothesis)[
+                self.variant
+            ]
+            return tf.cast(
+                tf.constant([score.precision, score.recall, score.fmeasure]),
+                dtype=self.dtype,
+            )
+
+        for batch_idx in range(batch_size):
+            score = tf.py_function(
+                func=calculate_rouge_score,
+                inp=[y_true[batch_idx], y_pred[batch_idx]],
+                Tout=self.dtype,
+            )
+            self._rouge_precision.assign_add(score[0])
+            self._rouge_recall.assign_add(score[1])
+            self._rouge_f1_score.assign_add(score[2])
+
+        self._number_of_samples.assign_add(
+            tf.cast(batch_size, dtype=self.dtype)
+        )
+
+    def result(self):
+        if self._number_of_samples == 0:
+            return {
+                "precision": 0.0,
+                "recall": 0.0,
+                "f1_score": 0.0,
+            }
+
+        rouge_precision = self._rouge_precision / self._number_of_samples
+        rouge_recall = self._rouge_recall / self._number_of_samples
+        rouge_f1_score = self._rouge_f1_score / self._number_of_samples
+        return {
+            "precision": rouge_precision,
+            "recall": rouge_recall,
+            "f1_score": rouge_f1_score,
+        }
+
+    def reset_state(self):
+        self._rouge_precision.assign(0.0)
+        self._rouge_recall.assign(0.0)
+        self._rouge_f1_score.assign(0.0)
+        self._number_of_samples.assign(0.0)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "variant": self.variant,
+                "use_stemmer": self.use_stemmer,
+            }
+        )
+        return config
diff --git a/keras_nlp/metrics/rouge_l.py b/keras_nlp/metrics/rouge_l.py
new file mode 100644
index 00000000..f6969a85
--- /dev/null
+++ b/keras_nlp/metrics/rouge_l.py
@@ -0,0 +1,129 @@
+# 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.
+
+"""ROUGE-L metric implementation based on `keras.metrics.Metric`."""
+
+
+from keras_nlp.metrics.rouge_base import RougeBase
+
+
+class RougeL(RougeBase):
+    """ROUGE-L metric.
+
+    This class implements the ROUGE-L variant of the ROUGE metric. The ROUGE-L
+    metric is traditionally used for evaluating summarisation systems.
+    Succinctly put, ROUGE-L is a score based on the length of the longest
+    common subsequence present in the reference text and the hypothesis text.
+
+    Note on input shapes:
+    For `y_true` and `y_pred`, this class supports scalar values and batch
+    inputs of shapes `()`, `(batch_size,)` and `(batch_size, 1)`.
+
+    Args:
+        use_stemmer: bool. Whether Porter Stemmer should be used to strip word
+            suffixes to improve matching. Defaults to False.
+        dtype: string or tf.dtypes.Dtype. Precision of metric computation. If
+               not specified, it defaults to tf.float32.
+        name: string. Name of the metric instance.
+        **kwargs: Other keyword arguments.
+
+    Examples:
+
+    1. Various Input Types.
+    1.1. Python string.
+    >>> rouge_l = keras_nlp.metrics.RougeL()
+    >>> y_true = "the tiny little cat was found under the big funny bed"
+    >>> y_pred = "the cat was under the bed"
+    >>> rouge_l(y_true, y_pred)["f1_score"]
+    <tf.Tensor: shape=(), dtype=float32, numpy=0.7058824>
+
+    1.2. rank 1 inputs.
+    a. Python list.
+    >>> rouge_l = keras_nlp.metrics.RougeL()
+    >>> y_true = [
+    ...     "the tiny little cat was found under the big funny bed",
+    ...     "i really love contributing to KerasNLP",
+    ... ]
+    >>> y_pred = [
+    ...     "the cat was under the bed",
+    ...     "i love contributing to KerasNLP",
+    ... ]
+    >>> rouge_l(y_true, y_pred)["f1_score"]
+    <tf.Tensor: shape=(), dtype=float32, numpy=0.80748665>
+
+    b. Tensor
+    >>> rouge_l = keras_nlp.metrics.RougeL()
+    >>> y_true = tf.constant(
+    ...     [
+    ...         "the tiny little cat was found under the big funny bed",
+    ...         "i really love contributing to KerasNLP",
+    ...     ]
+    ... )
+    >>> y_pred = tf.constant(
+    ...     [
+    ...         "the cat was under the bed",
+    ...         "i love contributing to KerasNLP",
+    ...     ]
+    ... )
+    >>> rouge_l(y_true, y_pred)["f1_score"]
+    <tf.Tensor: shape=(), dtype=float32, numpy=0.80748665>
+
+    1.3. rank 2 inputs.
+    >>> rouge_l = keras_nlp.metrics.RougeL()
+    >>> y_true = tf.constant(
+    ...     [
+    ...         ["the tiny little cat was found under the big funny bed"],
+    ...         ["i really love contributing to KerasNLP"],
+    ...     ]
+    ... )
+    >>> y_pred = tf.constant(
+    ...     [
+    ...         ["the cat was under the bed"],
+    ...         ["i love contributing to KerasNLP"],
+    ...     ]
+    ... )
+    >>> rouge_l(y_true, y_pred)["f1_score"]
+    <tf.Tensor: shape=(), dtype=float32, numpy=0.80748665>
+
+    3. Pass the metric to `model.compile()`.
+    >>> inputs = keras.Input(shape=(), dtype='string')
+    >>> outputs = tf.strings.lower(inputs)
+    >>> model = keras.Model(inputs, outputs)
+    >>> model.compile(metrics=[keras_nlp.metrics.RougeL()])
+    >>> x = tf.constant(["HELLO THIS IS FUN"])
+    >>> y = tf.constant(["hello this is awesome"])
+    >>> metric_dict = model.evaluate(x, y, return_dict=True)
+    >>> metric_dict["f1_score"]
+     0.75
+    """
+
+    def __init__(
+        self,
+        use_stemmer=False,
+        dtype=None,
+        name="rouge-l",
+        **kwargs,
+    ):
+        super().__init__(
+            variant="rougeL",
+            use_stemmer=use_stemmer,
+            dtype=dtype,
+            name=name,
+            **kwargs,
+        )
+
+    def get_config(self):
+        config = super().get_config()
+        del config["variant"]
+        return config
diff --git a/keras_nlp/metrics/rouge_n.py b/keras_nlp/metrics/rouge_n.py
new file mode 100644
index 00000000..4bfe532e
--- /dev/null
+++ b/keras_nlp/metrics/rouge_n.py
@@ -0,0 +1,163 @@
+# 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.
+
+"""ROUGE-N metric implementation based on `keras.metrics.Metric`."""
+
+
+from keras_nlp.metrics.rouge_base import RougeBase
+
+
+class RougeN(RougeBase):
+    """ROUGE-N metric.
+
+    This class implements the ROUGE-N variant of the ROUGE metric. The ROUGE-N
+    metric is traditionally used for evaluating summarisation systems.
+    Succinctly put, ROUGE-N is a score based on the number of matching n-grams
+    between the reference text and the hypothesis text.
+
+    Note on input shapes:
+    For `y_true` and `y_pred`, this class supports scalar values and batch
+    inputs of shapes `()`, `(batch_size,)` and `(batch_size, 1)`.
+
+    Args:
+        order: The order of n-grams which are to be matched. It should lie in
+            range [1, 9]. Defaults to 2.
+        use_stemmer: bool. Whether Porter Stemmer should be used to strip word
+            suffixes to improve matching. Defaults to False.
+        dtype: string or tf.dtypes.Dtype. Precision of metric computation. If
+               not specified, it defaults to tf.float32.
+        name: string. Name of the metric instance.
+        **kwargs: Other keyword arguments.
+
+    Examples:
+
+    1. Various Input Types.
+    1.1. Python string.
+    >>> rouge_n = keras_nlp.metrics.RougeN(order=2)
+    >>> y_true = "the tiny little cat was found under the big funny bed"
+    >>> y_pred = "the cat was under the bed"
+    >>> rouge_n(y_true, y_pred)["f1_score"]
+    <tf.Tensor: shape=(), dtype=float32, numpy=0.26666668>
+
+    1.2. rank 1 inputs.
+    a. Python list.
+    >>> rouge_n = keras_nlp.metrics.RougeN(order=2)
+    >>> y_true = [
+    ...     "the tiny little cat was found under the big funny bed",
+    ...     "i really love contributing to KerasNLP",
+    ... ]
+    >>> y_pred = [
+    ...     "the cat was under the bed",
+    ...     "i love contributing to KerasNLP",
+    ... ]
+    >>> rouge_n(y_true, y_pred)["f1_score"]
+    <tf.Tensor: shape=(), dtype=float32, numpy=0.4666667>
+
+    b. Tensor.
+    >>> rouge_n = keras_nlp.metrics.RougeN(order=2)
+    >>> y_true = tf.constant(
+    ...     [
+    ...         "the tiny little cat was found under the big funny bed",
+    ...         "i really love contributing to KerasNLP",
+    ...     ]
+    ... )
+    >>> y_pred = tf.constant(
+    ...     [
+    ...         "the cat was under the bed",
+    ...         "i love contributing to KerasNLP",
+    ...     ]
+    ... )
+    >>> rouge_n(y_true, y_pred)["f1_score"]
+    <tf.Tensor: shape=(), dtype=float32, numpy=0.4666667>
+
+    1.3. rank 2 inputs.
+    >>> rouge_n = keras_nlp.metrics.RougeN(order=2)
+    >>> y_true = tf.constant(
+    ...     [
+    ...         ["the tiny little cat was found under the big funny bed"],
+    ...         ["i really love contributing to KerasNLP"],
+    ...     ]
+    ... )
+    >>> y_pred = tf.constant(
+    ...     [
+    ...         ["the cat was under the bed"],
+    ...         ["i love contributing to KerasNLP"],
+    ...     ]
+    ... )
+    >>> rouge_n(y_true, y_pred)["f1_score"]
+    <tf.Tensor: shape=(), dtype=float32, numpy=0.4666667>
+
+    2. Consider trigrams for calculating ROUGE-N.
+    >>> rouge_n = keras_nlp.metrics.RougeN(order=3)
+    >>> y_true = tf.constant(
+    ...     [
+    ...         "the tiny little cat was found under the big funny bed",
+    ...         "i really love contributing to KerasNLP",
+    ...     ]
+    ... )
+    >>> y_pred = tf.constant(
+    ...     [
+    ...         "the cat was under the bed",
+    ...         "i love contributing to KerasNLP",
+    ...     ]
+    ... )
+    >>> rouge_n(y_true, y_pred)["f1_score"]
+    <tf.Tensor: shape=(), dtype=float32, numpy=0.2857143>
+
+    3. Pass the metric to `model.compile()`.
+    >>> inputs = keras.Input(shape=(), dtype='string')
+    >>> outputs = tf.strings.lower(inputs)
+    >>> model = keras.Model(inputs, outputs)
+    >>> model.compile(metrics=[keras_nlp.metrics.RougeN()])
+    >>> x = tf.constant(["HELLO THIS IS FUN"])
+    >>> y = tf.constant(["hello this is awesome"])
+    >>> metric_dict = model.evaluate(x, y, return_dict=True)
+    >>> metric_dict["f1_score"]
+    0.6666666865348816
+    """
+
+    def __init__(
+        self,
+        order=2,
+        use_stemmer=False,
+        dtype=None,
+        name="rouge-n",
+        **kwargs,
+    ):
+        if order not in range(1, 10):
+            raise ValueError(
+                "Invalid `order` value. Should lie in the range [1, 9]."
+                f"Received order={order}"
+            )
+
+        super().__init__(
+            variant=f"rouge{order}",
+            use_stemmer=use_stemmer,
+            dtype=dtype,
+            name=name,
+            **kwargs,
+        )
+
+        self.order = order
+
+    def get_config(self):
+        config = super().get_config()
+        del config["variant"]
+
+        config.update(
+            {
+                "order": self.order,
+            }
+        )
+        return config
diff --git a/keras_nlp/tokenizers/byte_tokenizer.py b/keras_nlp/tokenizers/byte_tokenizer.py
index fcbad637..eb0057a8 100644
--- a/keras_nlp/tokenizers/byte_tokenizer.py
+++ b/keras_nlp/tokenizers/byte_tokenizer.py
@@ -14,9 +14,6 @@
 
 """Byte Tokenizer."""
 
-from typing import Any
-from typing import Dict
-
 import numpy as np
 import tensorflow as tf
 import tensorflow_text as tf_text
@@ -248,7 +245,7 @@ class ByteTokenizer(tokenizer.Tokenizer):
         )
         return decoded
 
-    def get_config(self) -> Dict[str, Any]:
+    def get_config(self):
         config = super().get_config()
         config.update(
             {
diff --git a/keras_nlp/tokenizers/unicode_character_tokenizer.py b/keras_nlp/tokenizers/unicode_character_tokenizer.py
index 17462537..c1955764 100644
--- a/keras_nlp/tokenizers/unicode_character_tokenizer.py
+++ b/keras_nlp/tokenizers/unicode_character_tokenizer.py
@@ -12,9 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any
-from typing import Dict
-
 import tensorflow as tf
 import tensorflow_text as tf_text
 
@@ -247,7 +244,7 @@ class UnicodeCharacterTokenizer(tokenizer.Tokenizer):
         self.output_encoding = output_encoding
         self.vocabulary_size = vocabulary_size
 
-    def get_config(self) -> Dict[str, Any]:
+    def get_config(self):
         config = super().get_config()
         config.update(
             {
diff --git a/keras_nlp/tokenizers/word_piece_tokenizer.py b/keras_nlp/tokenizers/word_piece_tokenizer.py
index 79ec5057..9b108213 100644
--- a/keras_nlp/tokenizers/word_piece_tokenizer.py
+++ b/keras_nlp/tokenizers/word_piece_tokenizer.py
@@ -12,11 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any
-from typing import Dict
 from typing import Iterable
 from typing import List
-from typing import Union
 
 import tensorflow as tf
 import tensorflow_text as tf_text
@@ -102,7 +99,7 @@ class WordPieceTokenizer(tokenizer.Tokenizer):
     be either an integer or string type.
 
     Args:
-        vocabulary: A list of strings or a string string filename path. If
+        vocabulary: A list of strings or a string filename path. If
             passing a list, each element of the list should be a single word
             piece token string. If passing a filename, the file should be a
             plain text file containing a single word piece token per line.
@@ -175,7 +172,7 @@ class WordPieceTokenizer(tokenizer.Tokenizer):
 
     def __init__(
         self,
-        vocabulary: Union[Iterable[str], str] = None,
+        vocabulary=None,
         sequence_length: int = None,
         lowercase: bool = True,
         strip_accents: bool = True,
@@ -262,11 +259,11 @@ class WordPieceTokenizer(tokenizer.Tokenizer):
     def token_to_id(self, token: str) -> int:
         """Convert a string token to an integer id."""
         # This will be slow, but keep memory usage down compared to building a
-        # dict. Assuming the main use case is looking up a few special tokens
+        # . Assuming the main use case is looking up a few special tokens
         # early in the vocab, this should be fine.
         return self.vocabulary.index(token)
 
-    def get_config(self) -> Dict[str, Any]:
+    def get_config(self):
         config = super().get_config()
         config.update(
             {
diff --git a/keras_nlp/utils/tensor_utils.py b/keras_nlp/utils/tensor_utils.py
index 4f00b2b1..26fc815f 100644
--- a/keras_nlp/utils/tensor_utils.py
+++ b/keras_nlp/utils/tensor_utils.py
@@ -12,10 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
-
 import tensorflow as tf
-from google import protobuf
 
 
 def _decode_strings_to_utf8(inputs):
@@ -48,14 +45,3 @@ def tensor_to_string_list(inputs):
         if inputs.shape.rank != 0:
             list_outputs = list_outputs.tolist()
     return _decode_strings_to_utf8(list_outputs)
-
-
-def preview_tfrecord(filepath):
-    """Pretty prints a single record from a tfrecord file."""
-    dataset = tf.data.TFRecordDataset(os.path.expanduser(filepath))
-    example = tf.train.Example()
-    example.ParseFromString(next(iter(dataset)).numpy())
-    formatted = protobuf.text_format.MessageToString(
-        example, use_short_repeated_primitives=True
-    )
-    print(formatted)
diff --git a/setup.cfg b/setup.cfg
index 2fe96a95..67f100fb 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -6,6 +6,7 @@ filterwarnings =
     ignore::RuntimeWarning
     ignore::PendingDeprecationWarning
     ignore::FutureWarning
+    ignore::UserWarning
 
 addopts=-vv
 
diff --git a/setup.py b/setup.py
index 371a3e08..fac287ae 100644
--- a/setup.py
+++ b/setup.py
@@ -48,6 +48,7 @@ setup(
             "isort",
             "pytest",
             "pytest-cov",
+            "rouge-score",
         ],
         "examples": [
             "datasets",  # For GLUE in BERT example.


Test patch:
diff --git a/keras_nlp/layers/multi_segment_packer_test.py b/keras_nlp/layers/multi_segment_packer_test.py
new file mode 100644
index 00000000..a03b20cd
--- /dev/null
+++ b/keras_nlp/layers/multi_segment_packer_test.py
@@ -0,0 +1,178 @@
+# 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 Transformer Decoder."""
+
+import tensorflow as tf
+from tensorflow import keras
+
+from keras_nlp.layers.multi_segment_packer import MultiSegmentPacker
+
+
+class MultiSegmentPackerTest(tf.test.TestCase):
+    def test_trim_single_input_ints(self):
+        input_data = tf.range(3, 10)
+        packer = MultiSegmentPacker(8, start_value=1, end_value=2)
+        output = packer(input_data)
+        self.assertAllEqual(
+            output, ([1, 3, 4, 5, 6, 7, 8, 2], [0, 0, 0, 0, 0, 0, 0, 0])
+        )
+
+    def test_trim_single_input_strings(self):
+        input_data = tf.constant(["a", "b", "c", "d"])
+        packer = MultiSegmentPacker(5, start_value="[CLS]", end_value="[SEP]")
+        output = packer(input_data)
+        self.assertAllEqual(
+            output, (["[CLS]", "a", "b", "c", "[SEP]"], [0, 0, 0, 0, 0])
+        )
+
+    def test_trim_multiple_inputs_round_robin(self):
+        seq1 = tf.constant(["a", "b", "c"])
+        seq2 = tf.constant(["x", "y", "z"])
+        packer = MultiSegmentPacker(
+            7, start_value="[CLS]", end_value="[SEP]", truncator="round_robin"
+        )
+        output = packer([seq1, seq2])
+        self.assertAllEqual(
+            output,
+            (
+                ["[CLS]", "a", "b", "[SEP]", "x", "y", "[SEP]"],
+                [0, 0, 0, 0, 1, 1, 1],
+            ),
+        )
+
+    def test_trim_multiple_inputs_waterfall(self):
+        seq1 = tf.constant(["a", "b", "c"])
+        seq2 = tf.constant(["x", "y", "z"])
+        packer = MultiSegmentPacker(
+            7, start_value="[CLS]", end_value="[SEP]", truncator="waterfall"
+        )
+        output = packer([seq1, seq2])
+        self.assertAllEqual(
+            output,
+            (
+                ["[CLS]", "a", "b", "c", "[SEP]", "x", "[SEP]"],
+                [0, 0, 0, 0, 0, 1, 1],
+            ),
+        )
+
+    def test_trim_batched_inputs_round_robin(self):
+        seq1 = tf.ragged.constant([["a", "b", "c"], ["a", "b", "c"]])
+        seq2 = tf.ragged.constant([["x", "y", "z"], ["x", "y", "z"]])
+        packer = MultiSegmentPacker(
+            7, start_value="[CLS]", end_value="[SEP]", truncator="round_robin"
+        )
+        output = packer([seq1, seq2])
+        self.assertAllEqual(
+            output,
+            (
+                [
+                    ["[CLS]", "a", "b", "[SEP]", "x", "y", "[SEP]"],
+                    ["[CLS]", "a", "b", "[SEP]", "x", "y", "[SEP]"],
+                ],
+                [
+                    [0, 0, 0, 0, 1, 1, 1],
+                    [0, 0, 0, 0, 1, 1, 1],
+                ],
+            ),
+        )
+
+    def test_trim_batched_inputs_waterfall(self):
+        seq1 = tf.ragged.constant([["a", "b", "c"], ["a", "b"]])
+        seq2 = tf.ragged.constant([["x", "y", "z"], ["x", "y", "z"]])
+        packer = MultiSegmentPacker(
+            7, start_value="[CLS]", end_value="[SEP]", truncator="waterfall"
+        )
+        output = packer([seq1, seq2])
+        self.assertAllEqual(
+            output,
+            (
+                [
+                    ["[CLS]", "a", "b", "c", "[SEP]", "x", "[SEP]"],
+                    ["[CLS]", "a", "b", "[SEP]", "x", "y", "[SEP]"],
+                ],
+                [
+                    [0, 0, 0, 0, 0, 1, 1],
+                    [0, 0, 0, 0, 1, 1, 1],
+                ],
+            ),
+        )
+
+    def test_pad_inputs(self):
+        seq1 = tf.constant(["a"])
+        seq2 = tf.constant(["x"])
+        packer = MultiSegmentPacker(
+            6, start_value="[CLS]", end_value="[SEP]", pad_value="[PAD]"
+        )
+        output = packer([seq1, seq2])
+        self.assertAllEqual(
+            output,
+            (
+                ["[CLS]", "a", "[SEP]", "x", "[SEP]", "[PAD]"],
+                [0, 0, 0, 1, 1, 0],
+            ),
+        )
+
+    def test_pad_batched_inputs(self):
+        seq1 = tf.ragged.constant([["a"], ["a"]])
+        seq2 = tf.ragged.constant([["x"], ["x", "y"]])
+        packer = MultiSegmentPacker(
+            7, start_value="[CLS]", end_value="[SEP]", pad_value="[PAD]"
+        )
+        output = packer([seq1, seq2])
+        self.assertAllEqual(
+            output,
+            (
+                [
+                    ["[CLS]", "a", "[SEP]", "x", "[SEP]", "[PAD]", "[PAD]"],
+                    ["[CLS]", "a", "[SEP]", "x", "y", "[SEP]", "[PAD]"],
+                ],
+                [
+                    [0, 0, 0, 1, 1, 0, 0],
+                    [0, 0, 0, 1, 1, 1, 0],
+                ],
+            ),
+        )
+
+    def test_config(self):
+        seq1 = tf.ragged.constant([["a", "b", "c"], ["a", "b"]])
+        seq2 = tf.ragged.constant([["x", "y", "z"], ["x", "y", "z"]])
+        original_packer = MultiSegmentPacker(
+            7, start_value="[CLS]", end_value="[SEP]", truncator="waterfall"
+        )
+        cloned_packer = MultiSegmentPacker.from_config(
+            original_packer.get_config()
+        )
+        self.assertAllEqual(
+            original_packer([seq1, seq2]),
+            cloned_packer([seq1, seq2]),
+        )
+
+    def test_saving(self):
+        seq1 = tf.ragged.constant([["a", "b", "c"], ["a", "b"]])
+        seq2 = tf.ragged.constant([["x", "y", "z"], ["x", "y", "z"]])
+        packer = MultiSegmentPacker(
+            7, start_value="[CLS]", end_value="[SEP]", truncator="waterfall"
+        )
+        inputs = (
+            keras.Input(dtype="string", ragged=True, shape=(None,)),
+            keras.Input(dtype="string", ragged=True, shape=(None,)),
+        )
+        outputs = packer(inputs)
+        model = keras.Model(inputs, outputs)
+        model.save(self.get_temp_dir())
+        restored_model = keras.models.load_model(self.get_temp_dir())
+        self.assertAllEqual(
+            model((seq1, seq2)),
+            restored_model((seq1, seq2)),
+        )
diff --git a/keras_nlp/layers/start_end_packer_test.py b/keras_nlp/layers/start_end_packer_test.py
new file mode 100644
index 00000000..b36bc3b1
--- /dev/null
+++ b/keras_nlp/layers/start_end_packer_test.py
@@ -0,0 +1,151 @@
+# 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 Start End Packer layer."""
+
+
+import tensorflow as tf
+from tensorflow import keras
+
+from keras_nlp.layers.start_end_packer import StartEndPacker
+
+
+class StartEndPackerTest(tf.test.TestCase):
+    def test_dense_input(self):
+        input_data = tf.constant([5, 6, 7])
+        start_end_packer = StartEndPacker(sequence_length=5)
+        output = start_end_packer(input_data)
+        expected_output = [5, 6, 7, 0, 0]
+        self.assertAllEqual(output, expected_output)
+
+    def test_dense_input_error(self):
+        input_data = tf.constant([[5, 6, 7]])
+        start_end_packer = StartEndPacker(sequence_length=5)
+        with self.assertRaisesRegex(
+            ValueError,
+            "Input must either be dense with rank 1, or ragged with rank 2. "
+            "Received dense input with rank=2",
+        ):
+            start_end_packer(input_data)
+
+    def test_ragged_input(self):
+        input_data = tf.ragged.constant([[5, 6, 7], [8, 9, 10, 11]])
+        start_end_packer = StartEndPacker(sequence_length=5)
+        output = start_end_packer(input_data)
+        expected_output = [[5, 6, 7, 0, 0], [8, 9, 10, 11, 0]]
+        self.assertAllEqual(output, expected_output)
+
+    def test_ragged_input_error(self):
+        input_data = tf.ragged.constant([[[5, 6, 7], [8, 9, 10, 11]]])
+        start_end_packer = StartEndPacker(sequence_length=5)
+        with self.assertRaisesRegex(
+            ValueError,
+            "Input must either be dense with rank 1, or ragged with rank 2. "
+            "Received ragged input with "
+            "rank=3",
+        ):
+            start_end_packer(input_data)
+
+    def test_start_end_token(self):
+        input_data = tf.ragged.constant([[5, 6, 7], [8, 9, 10, 11]])
+        start_end_packer = StartEndPacker(
+            sequence_length=6, start_value=1, end_value=2
+        )
+        output = start_end_packer(input_data)
+        expected_output = [[1, 5, 6, 7, 2, 0], [1, 8, 9, 10, 11, 2]]
+        self.assertAllEqual(output, expected_output)
+
+    def test_start_end_padding_value(self):
+        input_data = tf.ragged.constant([[5, 6, 7], [8, 9, 10, 11]])
+        start_end_packer = StartEndPacker(
+            sequence_length=7, start_value=1, end_value=2, pad_value=3
+        )
+        output = start_end_packer(input_data)
+        expected_output = [[1, 5, 6, 7, 2, 3, 3], [1, 8, 9, 10, 11, 2, 3]]
+        self.assertAllEqual(output, expected_output)
+
+    def test_end_token_value_during_truncation(self):
+        input_data = tf.ragged.constant([[5, 6], [8, 9, 10, 11, 12, 13]])
+        start_end_packer = StartEndPacker(
+            sequence_length=5, start_value=1, end_value=2, pad_value=0
+        )
+        output = start_end_packer(input_data)
+        expected_output = [[1, 5, 6, 2, 0], [1, 8, 9, 10, 2]]
+        self.assertAllEqual(output, expected_output)
+
+    def test_string_input(self):
+        input_data = tf.ragged.constant(
+            [["KerasNLP", "is", "awesome"], ["amazing"]]
+        )
+        start_end_packer = StartEndPacker(
+            sequence_length=5,
+            start_value="[START]",
+            end_value="[END]",
+            pad_value="[PAD]",
+        )
+        output = start_end_packer(input_data)
+        expected_output = [
+            ["[START]", "KerasNLP", "is", "awesome", "[END]"],
+            ["[START]", "amazing", "[END]", "[PAD]", "[PAD]"],
+        ]
+        self.assertAllEqual(output, expected_output)
+
+    def test_functional_model(self):
+        input_data = tf.ragged.constant([[5, 6, 7], [8, 9, 10, 11]])
+        start_end_packer = StartEndPacker(
+            sequence_length=7, start_value=1, end_value=2, pad_value=3
+        )
+
+        inputs = keras.Input(dtype=tf.int32, shape=())
+        outputs = start_end_packer(inputs)
+        model = keras.Model(inputs, outputs)
+        model_output = model(input_data)
+
+        expected_output = [[1, 5, 6, 7, 2, 3, 3], [1, 8, 9, 10, 11, 2, 3]]
+        self.assertAllEqual(model_output, expected_output)
+
+    def test_batch(self):
+        start_end_packer = StartEndPacker(
+            sequence_length=7, start_value=1, end_value=2, pad_value=3
+        )
+
+        ds = tf.data.Dataset.from_tensor_slices(
+            tf.ragged.constant([[5, 6, 7], [8, 9, 10, 11]])
+        )
+        ds = ds.batch(2).map(start_end_packer)
+        output = ds.take(1).get_single_element()
+
+        exp_output = [[1, 5, 6, 7, 2, 3, 3], [1, 8, 9, 10, 11, 2, 3]]
+
+        for i in range(output.shape[0]):
+            self.assertAllEqual(output[i], exp_output[i])
+
+    def test_get_config(self):
+        start_end_packer = StartEndPacker(
+            sequence_length=512,
+            start_value=10,
+            end_value=20,
+            pad_value=100,
+            name="start_end_packer_test",
+        )
+
+        config = start_end_packer.get_config()
+        expected_config_subset = {
+            "sequence_length": 512,
+            "start_value": 10,
+            "end_value": 20,
+            "pad_value": 100,
+        }
+
+        self.assertEqual(config, {**config, **expected_config_subset})
diff --git a/keras_nlp/layers/transformer_decoder_test.py b/keras_nlp/layers/transformer_decoder_test.py
index 43460892..80a18e12 100644
--- a/keras_nlp/layers/transformer_decoder_test.py
+++ b/keras_nlp/layers/transformer_decoder_test.py
@@ -29,32 +29,52 @@ class TransformerDecoderTest(tf.test.TestCase):
             intermediate_dim=4,
             num_heads=2,
         )
-        output = decoder(encoder_input, decoder_input)
+        output = decoder(decoder_input, encoder_input)
         model = keras.Model(
             inputs=[decoder_input, encoder_input],
             outputs=output,
         )
-
         encoder_sequence = tf.random.uniform(shape=[2, 4, 6])
         decoder_sequence = tf.random.uniform(shape=[2, 4, 6])
-        model([encoder_sequence, decoder_sequence])
+        model([decoder_sequence, encoder_sequence])
 
-    def test_valid_call_with_mask(self):
+    def test_valid_call_without_cross_attention(self):
+        decoder_input = keras.Input(shape=[4, 6])
         decoder = transformer_decoder.TransformerDecoder(
             intermediate_dim=4,
             num_heads=2,
         )
-        encoder_sequence = tf.random.uniform(shape=[2, 4, 6])
+        output = decoder(decoder_input)
+        model = keras.Model(
+            inputs=decoder_input,
+            outputs=output,
+        )
         decoder_sequence = tf.random.uniform(shape=[2, 4, 6])
-        encoder_padding_mask = encoder_sequence[:, :, 0] > 0.5
-        decoder_padding_mask = decoder_sequence[:, :, 0] > 0.5
-        decoder(
-            decoder_sequence,
-            encoder_sequence,
-            decoder_padding_mask=decoder_padding_mask,
-            encoder_padding_mask=encoder_padding_mask,
-            use_causal_mask=True,
+        model(decoder_sequence)
+
+    def test_invalid_calls(self):
+        encoder_input = keras.Input(shape=[4, 6])
+        decoder_input = keras.Input(shape=[4, 6])
+
+        # with cross-attention.
+        decoder = transformer_decoder.TransformerDecoder(
+            intermediate_dim=4,
+            num_heads=2,
+        )
+        decoder(decoder_input, encoder_input)
+        # should raise ValueError if encoder_input is not provided
+        with self.assertRaises(ValueError):
+            decoder(decoder_input)
+
+        # without cross-attention.
+        decoder = transformer_decoder.TransformerDecoder(
+            intermediate_dim=4,
+            num_heads=2,
         )
+        decoder(decoder_input)
+        # should raise ValueError if encoder_input is provided
+        with self.assertRaises(ValueError):
+            decoder(decoder_input, encoder_input)
 
     def test_get_config_and_from_config(self):
         decoder = transformer_decoder.TransformerDecoder(
@@ -65,7 +85,6 @@ class TransformerDecoderTest(tf.test.TestCase):
         )
 
         config = decoder.get_config()
-
         expected_config_subset = {
             "intermediate_dim": 4,
             "num_heads": 2,
@@ -79,14 +98,11 @@ class TransformerDecoderTest(tf.test.TestCase):
                 keras.initializers.Zeros()
             ),
         }
-
         self.assertEqual(config, {**config, **expected_config_subset})
         self.assertEqual(config, {**config, **expected_config_subset})
-
         restored_decoder = transformer_decoder.TransformerDecoder.from_config(
             config,
         )
-
         self.assertEqual(
             restored_decoder.get_config(), {**config, **expected_config_subset}
         )
@@ -100,7 +116,7 @@ class TransformerDecoderTest(tf.test.TestCase):
                 kernel_initializer="Invalid",
             )
 
-    def test_one_training_step_of_transformer_encoder(self):
+    def test_one_training_step_of_transformer_with_cross_attention(self):
         class MyModel(keras.Model):
             def __init__(self):
                 super(MyModel, self).__init__()
@@ -128,12 +144,39 @@ class TransformerDecoderTest(tf.test.TestCase):
         self.assertGreater(len(grad), 1)
         optimizer.apply_gradients(zip(grad, model.trainable_variables))
 
+    def test_one_training_step_of_transformer_without_cross_attention(self):
+        class MyModel(keras.Model):
+            def __init__(self):
+                super(MyModel, self).__init__()
+                self._decoder = transformer_decoder.TransformerDecoder(
+                    intermediate_dim=4,
+                    num_heads=2,
+                )
+                self._dense = keras.layers.Dense(1, activation="sigmoid")
+
+            def call(self, decoder_input):
+                x = self._decoder(decoder_input)
+                return self._dense(x)
+
+        model = MyModel()
+
+        decoder_sequence = tf.random.uniform(shape=[2, 4, 6])
+        label = tf.cast(decoder_sequence[:, :, 0] >= 0.5, dtype=tf.int32)
+
+        loss_fn = keras.losses.BinaryCrossentropy(from_logits=False)
+        optimizer = keras.optimizers.Adam()
+        with tf.GradientTape() as tape:
+            pred = model(decoder_sequence)
+            loss = loss_fn(label, pred)
+        grad = tape.gradient(loss, model.trainable_variables)
+        self.assertGreater(len(grad), 1)
+        optimizer.apply_gradients(zip(grad, model.trainable_variables))
+
     def test_checkpointing_transformer_decoder(self):
         decoder1 = transformer_decoder.TransformerDecoder(
             intermediate_dim=4,
             num_heads=2,
         )
-
         decoder2 = transformer_decoder.TransformerDecoder(
             intermediate_dim=4,
             num_heads=2,
@@ -156,10 +199,38 @@ class TransformerDecoderTest(tf.test.TestCase):
         checkpoint2.restore(save_path)
 
         decoder1_output = decoder1(decoder_sequence, encoder_sequence)
-        decoder2_output = decoder2(
-            decoder_sequence,
-            encoder_sequence,
+        decoder2_output = decoder2(decoder_sequence, encoder_sequence)
+        self.assertAllClose(decoder1_output, decoder2_output)
+
+    def test_checkpointing_transformer_decoder_without_cross_attention(self):
+        decoder1 = transformer_decoder.TransformerDecoder(
+            intermediate_dim=4,
+            num_heads=2,
         )
+
+        decoder2 = transformer_decoder.TransformerDecoder(
+            intermediate_dim=4,
+            num_heads=2,
+        )
+
+        decoder_sequence = tf.random.uniform(shape=[2, 4, 6])
+        decoder1(decoder_sequence)
+        decoder2(decoder_sequence)
+        # The weights of decoder1 and decoder2 are different.
+        self.assertFalse(
+            all(
+                decoder1._output_dense.trainable_variables[0][0]
+                == decoder2._output_dense.trainable_variables[0][0]
+            )
+        )
+        checkpoint = tf.train.Checkpoint(decoder1)
+        checkpoint2 = tf.train.Checkpoint(decoder2)
+        temp_dir = self.get_temp_dir()
+        save_path = checkpoint.save(temp_dir)
+        checkpoint2.restore(save_path)
+
+        decoder1_output = decoder1(decoder_sequence)
+        decoder2_output = decoder2(decoder_sequence)
         self.assertAllClose(decoder1_output, decoder2_output)
 
     def test_mask_propagation(self):
@@ -174,6 +245,17 @@ class TransformerDecoderTest(tf.test.TestCase):
         outputs = decoder(decoder_sequence, encoder_sequence)
         self.assertAllEqual(outputs._keras_mask, mask)
 
+    def test_mask_propagation_without_cross_attention(self):
+        decoder = transformer_decoder.TransformerDecoder(
+            intermediate_dim=4,
+            num_heads=2,
+        )
+        decoder_sequence = tf.random.uniform(shape=[1, 4, 6])
+        mask = tf.constant([[True, True, False, False]])
+        decoder_sequence._keras_mask = mask
+        outputs = decoder(decoder_sequence)
+        self.assertAllEqual(outputs._keras_mask, mask)
+
     def test_save_model(self):
         encoder_input = keras.Input(shape=[4, 6])
         decoder_input = keras.Input(shape=[4, 6])
@@ -188,11 +270,32 @@ class TransformerDecoderTest(tf.test.TestCase):
         )
         encoder_sequence = tf.random.uniform(shape=[2, 4, 6])
         decoder_sequence = tf.random.uniform(shape=[2, 4, 6])
-        model([encoder_sequence, decoder_sequence])
+        model([decoder_sequence, encoder_sequence])
         path = os.path.join(self.get_temp_dir(), "model")
         model.save(path)
-        loaded_model = keras.models.load_model(path)
 
+        loaded_model = keras.models.load_model(path)
         model_output = model([decoder_sequence, encoder_sequence])
         loaded_model_output = loaded_model([decoder_sequence, encoder_sequence])
         self.assertAllClose(model_output, loaded_model_output)
+
+    def test_save_model_without_cross_attention(self):
+        decoder_input = keras.Input(shape=[4, 6])
+        decoder = transformer_decoder.TransformerDecoder(
+            intermediate_dim=4,
+            num_heads=2,
+        )
+        output = decoder(decoder_input)
+        model = keras.Model(
+            inputs=decoder_input,
+            outputs=output,
+        )
+        decoder_sequence = tf.random.uniform(shape=[2, 4, 6])
+        model(decoder_sequence)
+        path = os.path.join(self.get_temp_dir(), "model")
+        model.save(path)
+        loaded_model = keras.models.load_model(path)
+
+        model_output = model(decoder_sequence)
+        loaded_model_output = loaded_model(decoder_sequence)
+        self.assertAllClose(model_output, loaded_model_output)
diff --git a/keras_nlp/metrics/rouge_l_test.py b/keras_nlp/metrics/rouge_l_test.py
new file mode 100644
index 00000000..d130e121
--- /dev/null
+++ b/keras_nlp/metrics/rouge_l_test.py
@@ -0,0 +1,228 @@
+# 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 RougeL."""
+
+import tensorflow as tf
+from tensorflow import keras
+
+from keras_nlp.metrics import RougeL
+
+
+class RougeLTest(tf.test.TestCase):
+    def setUp(self):
+        super().setUp()
+
+        def assertDictAlmostEqual(d1, d2, delta=1e-3, typecast_to_numpy=True):
+            for key, val in d1.items():
+                if typecast_to_numpy:
+                    val = val.numpy()
+                self.assertAlmostEqual(val, d2[key], delta=delta)
+
+        def assertDictAllValuesNotEqual(d1, d2):
+            for key, val in d1.items():
+                self.assertNotEqual(val, d2[key])
+
+        self.assertDictAlmostEqual = assertDictAlmostEqual
+        self.assertDictAllValuesNotEqual = assertDictAllValuesNotEqual
+
+    def test_initialization(self):
+        rouge = RougeL()
+        result = rouge.result()
+
+        self.assertDictEqual(
+            result, {"precision": 0.0, "recall": 0.0, "f1_score": 0.0}
+        )
+
+    def test_string_input(self):
+        rouge = RougeL(use_stemmer=False)
+        y_true = "the tiny little cat was found under the big funny bed"
+        y_pred = "the cat was under the bed"
+
+        rouge_val = rouge(y_true, y_pred)
+        self.assertDictAlmostEqual(
+            rouge_val, {"precision": 1.0, "recall": 0.545, "f1_score": 0.706}
+        )
+
+    def test_string_list_input(self):
+        rouge = RougeL(use_stemmer=False)
+        y_true = [
+            "the tiny little cat was found under the big funny bed",
+            "i really love contributing to KerasNLP",
+        ]
+        y_pred = [
+            "the cat was under the bed",
+            "i love contributing to KerasNLP",
+        ]
+
+        rouge_val = rouge(y_true, y_pred)
+        self.assertDictAlmostEqual(
+            rouge_val, {"precision": 1.0, "recall": 0.689, "f1_score": 0.807}
+        )
+
+    def test_tensor_input(self):
+        rouge = RougeL(use_stemmer=False)
+        y_true = tf.constant(
+            [
+                "the tiny little cat was found under the big funny bed",
+                "i really love contributing to KerasNLP",
+            ]
+        )
+        y_pred = tf.constant(
+            ["the cat was under the bed", "i love contributing to KerasNLP"]
+        )
+
+        rouge_val = rouge(y_true, y_pred)
+        self.assertDictAlmostEqual(
+            rouge_val, {"precision": 1.0, "recall": 0.689, "f1_score": 0.807}
+        )
+
+    def test_rank_2_input(self):
+        rouge = RougeL(use_stemmer=False)
+        y_true = tf.constant(
+            [
+                ["the tiny little cat was found under the big funny bed"],
+                ["i really love contributing to KerasNLP"],
+            ]
+        )
+        y_pred = tf.constant(
+            [["the cat was under the bed"], ["i love contributing to KerasNLP"]]
+        )
+
+        rouge_val = rouge(y_true, y_pred)
+        self.assertDictAlmostEqual(
+            rouge_val, {"precision": 1.0, "recall": 0.689, "f1_score": 0.807}
+        )
+
+    def test_model_compile(self):
+        inputs = keras.Input(shape=(), dtype="string")
+        outputs = tf.strings.lower(inputs)
+        model = keras.Model(inputs, outputs)
+
+        model.compile(metrics=[RougeL()])
+
+        x = tf.constant(["HELLO THIS IS FUN"])
+        y = tf.constant(["hello this is awesome"])
+
+        output = model.evaluate(x, y, return_dict=True)
+        del output["loss"]
+        self.assertDictAlmostEqual(
+            output,
+            {"precision": 0.75, "recall": 0.75, "f1_score": 0.75},
+            typecast_to_numpy=False,
+        )
+
+    def test_reset_state(self):
+        rouge = RougeL()
+        y_true = tf.constant(
+            ["hey, this is great fun", "i love contributing to KerasNLP"]
+        )
+        y_pred = tf.constant(
+            [
+                "great fun indeed",
+                "KerasNLP is awesome, i love contributing to it",
+            ]
+        )
+
+        rouge.update_state(y_true, y_pred)
+        rouge_val = rouge.result()
+        self.assertDictAllValuesNotEqual(
+            rouge_val, {"precision": 0.0, "recall": 0.0, "f1_score": 0.0}
+        )
+
+        rouge.reset_state()
+        rouge_val = rouge.result()
+        self.assertDictEqual(
+            rouge_val, {"precision": 0.0, "recall": 0.0, "f1_score": 0.0}
+        )
+
+    def test_update_state(self):
+        rouge = RougeL()
+        y_true_1 = tf.constant(
+            [
+                "the tiny little cat was found under the big funny bed",
+                "i really love contributing to KerasNLP",
+            ]
+        )
+        y_pred_1 = tf.constant(
+            ["the cat was under the bed", "i love contributing to KerasNLP"]
+        )
+
+        rouge.update_state(y_true_1, y_pred_1)
+        rouge_val = rouge.result()
+        self.assertDictAlmostEqual(
+            rouge_val, {"precision": 1.0, "recall": 0.689, "f1_score": 0.807}
+        )
+
+        y_true_2 = tf.constant(["what is your favourite show"])
+        y_pred_2 = tf.constant(["my favourite show is silicon valley"])
+
+        rouge.update_state(y_true_2, y_pred_2)
+        rouge_val = rouge.result()
+        self.assertDictAlmostEqual(
+            rouge_val, {"precision": 0.778, "recall": 0.593, "f1_score": 0.66}
+        )
+
+    def test_merge_state(self):
+        rouge_1 = RougeL()
+        rouge_2 = RougeL()
+
+        y_true_1 = tf.constant(
+            [
+                "the tiny little cat was found under the big funny bed",
+                "i really love contributing to KerasNLP",
+            ]
+        )
+        y_pred_1 = tf.constant(
+            ["the cat was under the bed", "i love contributing to KerasNLP"]
+        )
+
+        y_true_2 = tf.constant(["what is your favourite show"])
+        y_pred_2 = tf.constant(["my favourite show is silicon valley"])
+
+        y_true_3 = tf.constant(["lorem ipsum dolor sit amet"])
+        y_pred_3 = tf.constant(["lorem ipsum is simply dummy text"])
+
+        rouge_1.update_state(y_true_1, y_pred_1)
+        rouge_1.update_state(y_true_2, y_pred_2)
+        rouge_val = rouge_1.result()
+        self.assertDictAlmostEqual(
+            rouge_val, {"precision": 0.778, "recall": 0.593, "f1_score": 0.66}
+        )
+
+        rouge_2.update_state(y_true_3, y_pred_3)
+        rouge_val = rouge_2.result()
+        self.assertDictAlmostEqual(
+            rouge_val, {"precision": 0.333, "recall": 0.4, "f1_score": 0.364}
+        )
+
+        merged_rouge = RougeL()
+        merged_rouge.merge_state([rouge_1, rouge_2])
+        rouge_val = merged_rouge.result()
+        self.assertDictAlmostEqual(
+            rouge_val, {"precision": 0.667, "recall": 0.545, "f1_score": 0.586}
+        )
+
+    def test_get_config(self):
+        rouge = RougeL(
+            use_stemmer=True,
+            dtype=tf.float32,
+            name="rouge_l_test",
+        )
+
+        config = rouge.get_config()
+        expected_config_subset = {
+            "use_stemmer": True,
+        }
+        self.assertEqual(config, {**config, **expected_config_subset})
diff --git a/keras_nlp/metrics/rouge_n_test.py b/keras_nlp/metrics/rouge_n_test.py
new file mode 100644
index 00000000..2183afe3
--- /dev/null
+++ b/keras_nlp/metrics/rouge_n_test.py
@@ -0,0 +1,254 @@
+# 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 RougeN."""
+
+import tensorflow as tf
+from tensorflow import keras
+
+from keras_nlp.metrics import RougeN
+
+
+class RougeNTest(tf.test.TestCase):
+    def setUp(self):
+        super().setUp()
+
+        def assertDictAlmostEqual(d1, d2, delta=1e-3, typecast_to_numpy=True):
+            for key, val in d1.items():
+                if typecast_to_numpy:
+                    val = val.numpy()
+                self.assertAlmostEqual(val, d2[key], delta=delta)
+
+        def assertDictAllValuesNotEqual(d1, d2):
+            for key, val in d1.items():
+                self.assertNotEqual(val, d2[key])
+
+        self.assertDictAlmostEqual = assertDictAlmostEqual
+        self.assertDictAllValuesNotEqual = assertDictAllValuesNotEqual
+
+    def test_initialization(self):
+        rouge = RougeN()
+        result = rouge.result()
+
+        self.assertDictEqual(
+            result, {"precision": 0.0, "recall": 0.0, "f1_score": 0.0}
+        )
+
+    def test_string_input(self):
+        rouge = RougeN(order=2, use_stemmer=False)
+        y_true = "the tiny little cat was found under the big funny bed"
+        y_pred = "the cat was under the bed"
+
+        rouge_val = rouge(y_true, y_pred)
+        self.assertDictAlmostEqual(
+            rouge_val, {"precision": 0.4, "recall": 0.2, "f1_score": 0.267}
+        )
+
+    def test_string_list_input(self):
+        rouge = RougeN(order=2, use_stemmer=False)
+        y_true = [
+            "the tiny little cat was found under the big funny bed",
+            "i really love contributing to KerasNLP",
+        ]
+        y_pred = [
+            "the cat was under the bed",
+            "i love contributing to KerasNLP",
+        ]
+
+        rouge_val = rouge(y_true, y_pred)
+        self.assertDictAlmostEqual(
+            rouge_val, {"precision": 0.575, "recall": 0.4, "f1_score": 0.467}
+        )
+
+    def test_tensor_input(self):
+        rouge = RougeN(order=2, use_stemmer=False)
+        y_true = tf.constant(
+            [
+                "the tiny little cat was found under the big funny bed",
+                "i really love contributing to KerasNLP",
+            ]
+        )
+        y_pred = tf.constant(
+            ["the cat was under the bed", "i love contributing to KerasNLP"]
+        )
+
+        rouge_val = rouge(y_true, y_pred)
+        self.assertDictAlmostEqual(
+            rouge_val, {"precision": 0.575, "recall": 0.4, "f1_score": 0.467}
+        )
+
+    def test_rank_2_input(self):
+        rouge = RougeN(order=2, use_stemmer=False)
+        y_true = tf.constant(
+            [
+                ["the tiny little cat was found under the big funny bed"],
+                ["i really love contributing to KerasNLP"],
+            ]
+        )
+        y_pred = tf.constant(
+            [["the cat was under the bed"], ["i love contributing to KerasNLP"]]
+        )
+
+        rouge_val = rouge(y_true, y_pred)
+        self.assertDictAlmostEqual(
+            rouge_val, {"precision": 0.575, "recall": 0.4, "f1_score": 0.467}
+        )
+
+    def test_model_compile(self):
+        inputs = keras.Input(shape=(), dtype="string")
+        outputs = tf.strings.lower(inputs)
+        model = keras.Model(inputs, outputs)
+
+        model.compile(metrics=[RougeN()])
+
+        x = tf.constant(["HELLO THIS IS FUN"])
+        y = tf.constant(["hello this is awesome"])
+
+        output = model.evaluate(x, y, return_dict=True)
+        del output["loss"]
+        self.assertDictAlmostEqual(
+            output,
+            {"precision": 0.667, "recall": 0.667, "f1_score": 0.667},
+            typecast_to_numpy=False,
+        )
+
+    def test_incorrect_order(self):
+        with self.assertRaises(ValueError):
+            _ = RougeN(order=10)
+
+    def test_different_order(self):
+        rouge = RougeN(order=3, use_stemmer=False)
+        y_true = tf.constant(
+            [
+                "the tiny little cat was found under the big funny bed",
+                "i really love contributing to KerasNLP",
+            ]
+        )
+        y_pred = tf.constant(
+            ["the cat was under the bed", "i love contributing to KerasNLP"]
+        )
+
+        rouge_val = rouge(y_true, y_pred)
+        self.assertDictAlmostEqual(
+            rouge_val,
+            {"precision": 0.333, "recall": 0.25, "f1_score": 0.286},
+            typecast_to_numpy=False,
+        )
+
+    def test_reset_state(self):
+        rouge = RougeN()
+        y_true = tf.constant(
+            ["hey, this is great fun", "i love contributing to KerasNLP"]
+        )
+        y_pred = tf.constant(
+            [
+                "great fun indeed",
+                "KerasNLP is awesome, i love contributing to it",
+            ]
+        )
+
+        rouge.update_state(y_true, y_pred)
+        rouge_val = rouge.result()
+        self.assertDictAllValuesNotEqual(
+            rouge_val, {"precision": 0.0, "recall": 0.0, "f1_score": 0.0}
+        )
+
+        rouge.reset_state()
+        rouge_val = rouge.result()
+        self.assertDictEqual(
+            rouge_val, {"precision": 0.0, "recall": 0.0, "f1_score": 0.0}
+        )
+
+    def test_update_state(self):
+        rouge = RougeN()
+        y_true_1 = tf.constant(
+            [
+                "the tiny little cat was found under the big funny bed",
+                "i really love contributing to KerasNLP",
+            ]
+        )
+        y_pred_1 = tf.constant(
+            ["the cat was under the bed", "i love contributing to KerasNLP"]
+        )
+
+        rouge.update_state(y_true_1, y_pred_1)
+        rouge_val = rouge.result()
+        self.assertDictAlmostEqual(
+            rouge_val, {"precision": 0.575, "recall": 0.4, "f1_score": 0.467}
+        )
+
+        y_true_2 = tf.constant(["what is your favourite show"])
+        y_pred_2 = tf.constant(["my favourite show is silicon valley"])
+
+        rouge.update_state(y_true_2, y_pred_2)
+        rouge_val = rouge.result()
+        self.assertDictAlmostEqual(
+            rouge_val, {"precision": 0.45, "recall": 0.35, "f1_score": 0.385}
+        )
+
+    def test_merge_state(self):
+        rouge_1 = RougeN()
+        rouge_2 = RougeN()
+
+        y_true_1 = tf.constant(
+            [
+                "the tiny little cat was found under the big funny bed",
+                "i really love contributing to KerasNLP",
+            ]
+        )
+        y_pred_1 = tf.constant(
+            ["the cat was under the bed", "i love contributing to KerasNLP"]
+        )
+
+        y_true_2 = tf.constant(["what is your favourite show"])
+        y_pred_2 = tf.constant(["my favourite show is silicon valley"])
+
+        y_true_3 = tf.constant(["lorem ipsum dolor sit amet"])
+        y_pred_3 = tf.constant(["lorem ipsum is simply dummy text"])
+
+        rouge_1.update_state(y_true_1, y_pred_1)
+        rouge_1.update_state(y_true_2, y_pred_2)
+        rouge_val = rouge_1.result()
+        self.assertDictAlmostEqual(
+            rouge_val, {"precision": 0.45, "recall": 0.35, "f1_score": 0.385}
+        )
+
+        rouge_2.update_state(y_true_3, y_pred_3)
+        rouge_val = rouge_2.result()
+        self.assertDictAlmostEqual(
+            rouge_val, {"precision": 0.2, "recall": 0.25, "f1_score": 0.222}
+        )
+
+        merged_rouge = RougeN()
+        merged_rouge.merge_state([rouge_1, rouge_2])
+        rouge_val = merged_rouge.result()
+        self.assertDictAlmostEqual(
+            rouge_val, {"precision": 0.388, "recall": 0.325, "f1_score": 0.344}
+        )
+
+    def test_get_config(self):
+        rouge = RougeN(
+            order=5,
+            use_stemmer=True,
+            dtype=tf.float32,
+            name="rouge_n_test",
+        )
+
+        config = rouge.get_config()
+        expected_config_subset = {
+            "order": 5,
+            "use_stemmer": True,
+        }
+
+        self.assertEqual(config, {**config, **expected_config_subset})
diff --git a/keras_nlp/tokenizers/byte_tokenizer_test.py b/keras_nlp/tokenizers/byte_tokenizer_test.py
index 80b86bc7..752bb4bd 100644
--- a/keras_nlp/tokenizers/byte_tokenizer_test.py
+++ b/keras_nlp/tokenizers/byte_tokenizer_test.py
@@ -130,7 +130,6 @@ class ByteTokenizerTest(tf.test.TestCase):
             [104, 97, 104, 97],
         ]
         for i in range(output.shape[0]):
-            print(output[i])
             self.assertAllEqual(output[i], exp_output[i])
 
     def test_tokenize_first_batch_second_with_sequence_length(self):
@@ -150,7 +149,6 @@ class ByteTokenizerTest(tf.test.TestCase):
             [104, 97, 104, 97, 0, 0, 0, 0, 0, 0],
         ]
         for i in range(output.shape[0]):
-            print(output[i])
             self.assertAllEqual(output[i], exp_output[i])
 
     def test_batch_first_tokenize_second(self):
@@ -169,7 +167,6 @@ class ByteTokenizerTest(tf.test.TestCase):
             [104, 97, 104, 97],
         ]
         for i in range(output.shape[0]):
-            print(output[i])
             self.assertAllEqual(output[i], exp_output[i])
 
     def test_batch_first_tokenize_second_with_sequence_length(self):
@@ -188,7 +185,6 @@ class ByteTokenizerTest(tf.test.TestCase):
             [104, 97, 104, 97, 0, 0, 0, 0, 0, 0],
         ]
         for i in range(output.shape[0]):
-            print(output[i])
             self.assertAllEqual(output[i], exp_output[i])
 
     def test_functional_model(self):