Repository: keras-team/keras-nlp

Issue:
Add a RoBERTa preprocessing layer
Now that we have landed the `BytePairTokenizer`, we should be able to land the preprocessing for the RoBERTa model, modeled after the `BertPreprocessor` layer.


Discussion:
See these files as a template:

https://github.com/keras-team/keras-nlp/blob/master/keras_nlp/models/bert/bert_preprocessing.py
https://github.com/keras-team/keras-nlp/blob/master/keras_nlp/models/bert/bert_preprocessing_test.py

Since we are not yet distributing weights for RoBERTa, the `from_preset` static constructor can simple return a `NotImplementedError`.
Blocked on #417.
Working on this. Will stack the PR on top of #417.

Resolving patch:
diff --git a/keras_nlp/models/__init__.py b/keras_nlp/models/__init__.py
--- a/keras_nlp/models/__init__.py
+++ b/keras_nlp/models/__init__.py
@@ -21,6 +21,7 @@
 )
 from keras_nlp.models.gpt2.gpt2_models import Gpt2
 from keras_nlp.models.roberta.roberta_models import Roberta
+from keras_nlp.models.roberta.roberta_preprocessing import RobertaPreprocessor
 from keras_nlp.models.roberta.roberta_tasks import RobertaClassifier
 from keras_nlp.models.xlm_roberta.xlm_roberta_models import XLMRoberta
 from keras_nlp.models.xlm_roberta.xlm_roberta_preprocessing import (
diff --git a/keras_nlp/models/roberta/roberta_preprocessing.py b/keras_nlp/models/roberta/roberta_preprocessing.py
--- a/keras_nlp/models/roberta/roberta_preprocessing.py
+++ b/keras_nlp/models/roberta/roberta_preprocessing.py
@@ -12,10 +12,187 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+"""RoBERTa preprocessing layers."""
+
 import tensorflow as tf
 import tensorflow_text as tf_text
 from tensorflow import keras
 
+from keras_nlp.tokenizers.byte_pair_tokenizer import BytePairTokenizer
+
+
+@keras.utils.register_keras_serializable(package="keras_nlp")
+class RobertaPreprocessor(keras.layers.Layer):
+    """RoBERTa preprocessing layer.
+
+    This preprocessing layer will do three things:
+
+    - Tokenize any number of inputs using a
+      `keras_nlp.tokenizers.BytePairTokenizer`.
+    - Pack the inputs together with the appropriate `"
"`, `"
"` and
+      `"<pad>"` tokens, i.e., adding a single `"
"` at the start of the
+      entire sequence, `"

"` at the end of each segment, save the last
+      and a `"
"` at the end of the entire sequence.
+    - Construct a dictionary with keys `"token_ids"` and `"padding_mask"` that
+      can be passed directly to a RoBERTa model.
+
+    This layer will accept either a tuple of (possibly batched) inputs, or a
+    single input tensor. If a single tensor is passed, it will be packed
+    equivalently to a tuple with a single element.
+
+    The `BytePairTokenizer` can be accessed via the `tokenizer` property on
+    this layer, and can be used directly for custom packing on inputs.
+
+    Args:
+        vocabulary: string or dict, maps token to integer ids. If it is a
+            string, it should be the file path to a json file.
+        merges: string or list, contains the merge rule. If it is a string,
+            it should be the file path to merge rules. The merge rule file
+            should have one merge rule per line. Every merge rule contains
+            merge entities separated by a space. Please refer to this example:
+            https://storage.googleapis.com/keras-nlp/models/roberta_base/merges.txt.
+        sequence_length: The length of the packed inputs.
+        truncate: string. The algorithm to truncate a list of batched segments
+            to fit within `sequence_length`. The value can be either
+            `round_robin` or `waterfall`:
+                - `"round_robin"`: Available space is assigned one token at a
+                    time in a round-robin fashion to the inputs that still need
+                    some, until the limit is reached.
+                - `"waterfall"`: The allocation of the budget is done using a
+                    "waterfall" algorithm that allocates quota in a
+                    left-to-right manner and fills up the buckets until we run
+                    out of budget. It supports an arbitrary number of segments.
+
+    Examples:
+    ```python
+    vocab = {
+        "
": 0,
+        "<pad>": 1,
+        "
": 2,
+        "reful": 3,
+        "gent": 4,
+        "Ġafter": 5,
+        "noon": 6,
+        "Ġsun": 7,
+        "Ġbright": 8,
+        "Ġnight": 9,
+        "Ġmoon": 10,
+    }
+    merges = ["Ġ a", "Ġ m", "Ġ s", "Ġ b", "Ġ n", "r e", "f u", "g e", "n t"]
+    merges += ["e r", "n o", "o n", "i g", "h t"]
+    merges += ["Ġs u", "Ġa f", "Ġm o", "Ġb r","ge nt", "no on", "re fu", "ig ht"]
+    merges += ["Ġn ight", "Ġsu n", "Ġaf t", "Ġmo on", "Ġbr ight", "refu l", "Ġaft er"]
+
+    preprocessor = keras_nlp.models.RobertaPreprocessor(
+        vocabulary=vocab,
+        merges=merges,
+        sequence_length=20,
+    )
+
+    # Tokenize and pack a single sentence directly.
+    preprocessor(" afternoon sun")
+
+    # Tokenize and pack multiples sentences directly.
+    preprocessor((" afternoon sun", "refulgent sun"))
+
+    # Map a dataset to preprocess a single sentence at a time.
+    features = [" afternoon sun", "refulgent sun"]
+    labels = [0, 1]
+    ds = tf.data.Dataset.from_tensor_slices((features, labels))
+    ds = ds.map(
+        lambda x, y: (preprocessor(x), y),
+        num_parallel_calls=tf.data.AUTOTUNE,
+    )
+
+    # Map a dataset to preprocess multiple sentences.
+    first_sentences = [" afternoon sun", "refulgent sun"]
+    second_sentences = [" night moon", " bright moon"]
+    labels = [1, 1]
+    ds = tf.data.Dataset.from_tensor_slices(
+        (
+            (first_sentences, second_sentences), labels
+        )
+    )
+    ds = ds.map(
+        lambda x, y: (preprocessor(x), y),
+        num_parallel_calls=tf.data.AUTOTUNE,
+    )
+    ```
+    """
+
+    def __init__(
+        self,
+        vocabulary,
+        merges,
+        sequence_length=512,
+        truncate="round_robin",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.tokenizer = BytePairTokenizer(
+            vocabulary=vocabulary,
+            merges=merges,
+        )
+
+        # Check for necessary special tokens.
+        start_token = "
"
+        pad_token = "<pad>"
+        end_token = "
"
+        for token in [start_token, pad_token, end_token]:
+            if token not in self.tokenizer.get_vocabulary():
+                raise ValueError(
+                    f"Cannot find token `'{token}'` in the provided "
+                    f"`vocabulary`. Please provide `'{token}'` in your "
+                    "`vocabulary` or use a pretrained `vocabulary` name."
+                )
+
+        self.pad_token_id = self.tokenizer.token_to_id(pad_token)
+        self.packer = RobertaMultiSegmentPacker(
+            start_value=self.tokenizer.token_to_id(start_token),
+            end_value=self.tokenizer.token_to_id(end_token),
+            pad_value=self.pad_token_id,
+            truncate=truncate,
+            sequence_length=sequence_length,
+        )
+
+    def vocabulary_size(self):
+        """Returns the vocabulary size of the tokenizer."""
+        return self.tokenizer.vocabulary_size()
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "vocabulary": self.tokenizer.vocabulary,
+                "merges": self.tokenizer.merges,
+                "sequence_length": self.packer.sequence_length,
+                "truncate": self.packer.truncate,
+            }
+        )
+        return config
+
+    def call(self, inputs):
+        if not isinstance(inputs, (list, tuple)):
+            inputs = [inputs]
+
+        inputs = [self.tokenizer(x) for x in inputs]
+        token_ids = self.packer(inputs)
+        return {
+            "token_ids": token_ids,
+            "padding_mask": token_ids != self.pad_token_id,
+        }
+
+    @classmethod
+    def from_preset(
+        cls,
+        preset,
+        sequence_length=None,
+        truncate="round_robin",
+        **kwargs,
+    ):
+        raise NotImplementedError
+
 
 # TODO: This is a temporary, unexported layer until we find a way to make the
 # `MultiSegmentPacker` layer more generic.


Test patch:
diff --git a/keras_nlp/models/roberta/roberta_preprocessing_test.py b/keras_nlp/models/roberta/roberta_preprocessing_test.py
new file mode 100644
--- /dev/null
+++ b/keras_nlp/models/roberta/roberta_preprocessing_test.py
@@ -0,0 +1,152 @@
+# 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 RoBERTa preprocessing layers."""
+
+import os
+
+import tensorflow as tf
+from absl.testing import parameterized
+from tensorflow import keras
+
+from keras_nlp.models.roberta.roberta_preprocessing import RobertaPreprocessor
+
+
+class RobertaPreprocessorTest(tf.test.TestCase, parameterized.TestCase):
+    def setUp(self):
+        vocab = {
+            "
": 0,
+            "<pad>": 1,
+            "
": 2,
+            "Ġair": 3,
+            "plane": 4,
+            "Ġat": 5,
+            "port": 6,
+            "Ġkoh": 7,
+            "li": 8,
+            "Ġis": 9,
+            "Ġthe": 10,
+            "Ġbest": 11,
+        }
+
+        merges = ["Ġ a", "Ġ t", "Ġ k", "Ġ i", "Ġ b", "Ġa i", "p l", "n e"]
+        merges += ["Ġa t", "p o", "r t", "o h", "l i", "Ġi s", "Ġb e", "s t"]
+        merges += ["Ġt h", "Ġai r", "pl a", "Ġk oh", "Ġth e", "Ġbe st", "po rt"]
+        merges += ["pla ne"]
+
+        self.preprocessor = RobertaPreprocessor(
+            vocabulary=vocab,
+            merges=merges,
+            sequence_length=12,
+        )
+
+    def test_tokenize(self):
+        input_data = [" airplane at airport"]
+
+        output = self.preprocessor(input_data)
+        self.assertAllEqual(
+            output["token_ids"], [0, 3, 4, 5, 3, 6, 2, 1, 1, 1, 1, 1]
+        )
+        self.assertAllEqual(
+            output["padding_mask"], [1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0]
+        )
+
+    def test_tokenize_batch(self):
+        input_data = tf.constant(
+            [
+                " airplane at airport",
+                " airplane at airport",
+                " airplane at airport",
+                " airplane at airport",
+            ]
+        )
+
+        output = self.preprocessor(input_data)
+        self.assertAllEqual(
+            output["token_ids"],
+            [[0, 3, 4, 5, 3, 6, 2, 1, 1, 1, 1, 1]] * 4,
+        )
+
+        self.assertAllEqual(
+            output["padding_mask"], [[1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0]] * 4
+        )
+
+    def test_tokenize_multiple_sentences(self):
+        sentence_one = " airplane at airport"
+        sentence_two = " kohli is the best"
+
+        output = self.preprocessor((sentence_one, sentence_two))
+        print(output)
+        self.assertAllEqual(
+            output["token_ids"], [0, 3, 4, 5, 3, 2, 2, 7, 8, 9, 10, 2]
+        )
+        self.assertAllEqual(
+            output["padding_mask"], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
+        )
+
+    def test_tokenize_multiple_batched_sentences(self):
+        sentence_one = tf.constant(
+            [
+                " airplane at airport",
+                " airplane at airport",
+                " airplane at airport",
+                " airplane at airport",
+            ]
+        )
+        sentence_two = tf.constant(
+            [
+                " kohli is the best",
+                " kohli is the best",
+                " kohli is the best",
+                " kohli is the best",
+            ]
+        )
+
+        output = self.preprocessor((sentence_one, sentence_two))
+        self.assertAllEqual(
+            output["token_ids"],
+            [[0, 3, 4, 5, 3, 2, 2, 7, 8, 9, 10, 2]] * 4,
+        )
+        self.assertAllEqual(
+            output["padding_mask"], [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]] * 4
+        )
+
+    def test_detokenize(self):
+        input_tokens = [[0, 3, 4, 5, 3, 6, 2, 1, 1, 1, 1, 1]]
+        output = self.preprocessor.tokenizer.detokenize(input_tokens)
+        self.assertAllEqual(
+            output, ["
 airplane at airport
<pad><pad><pad><pad><pad>"]
+        )
+
+    def test_vocabulary_size(self):
+        self.assertEqual(self.preprocessor.vocabulary_size(), 12)
+
+    @parameterized.named_parameters(
+        ("save_format_tf", "tf"), ("save_format_h5", "h5")
+    )
+    def test_saving_model(self, save_format):
+        input_data = tf.constant([" airplane at airport"])
+
+        inputs = keras.Input(dtype="string", shape=())
+        outputs = self.preprocessor(inputs)
+        model = keras.Model(inputs, outputs)
+
+        path = os.path.join(self.get_temp_dir(), "model")
+        model.save(path, save_format=save_format)
+
+        restored_model = keras.models.load_model(path)
+        self.assertAllEqual(
+            model(input_data)["token_ids"],
+            restored_model(input_data)["token_ids"],
+        )
diff --git a/keras_nlp/models/xlm_roberta/xlm_roberta_preprocessing_test.py b/keras_nlp/models/xlm_roberta/xlm_roberta_preprocessing_test.py
--- a/keras_nlp/models/xlm_roberta/xlm_roberta_preprocessing_test.py
+++ b/keras_nlp/models/xlm_roberta/xlm_roberta_preprocessing_test.py
@@ -11,7 +11,8 @@
 # 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 BERT preprocessing layers."""
+
+"""Tests for XLM-RoBERTa preprocessing layers."""
 
 import io