Repository: keras-team/keras-nlp

Issue:
Add a standalone tokenizer for Distilbert
Following on with https://github.com/keras-team/keras-nlp/pull/449, we would like to expose a standalone, pure string -> int sequence tokenizer for our models. This will allow for any possible packing and padding, de-tokenization, etc.

This will be a new class wrapping the `WordPieceTokenizer`.

A few steps:

 - Add the `DistilBertTokenizer` class subclassing `WordPieceTokenizer` following the example in https://github.com/keras-team/keras-nlp/pull/449
 - Have `DistilBertPreprocessor` take in a `tokenizer` as an argument, and serialize it properly in get_config/from_config.
 - Add a documented property to access the tokenizer from the preprocessor layer.
 - Track the special token ids on the tokenizer class, and use these ids from the preprocessor like `self.tokenizer.cls_token_id`.
 - Expose `DistilBertTokenizer` in `models/__init__.py`.
 - Add documentation and tests!


Discussion:
Working on this.

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
@@ -20,6 +20,9 @@
 from keras_nlp.models.distilbert.distilbert_preprocessing import (
     DistilBertPreprocessor,
 )
+from keras_nlp.models.distilbert.distilbert_preprocessing import (
+    DistilBertTokenizer,
+)
 from keras_nlp.models.distilbert.distilbert_tasks import DistilBertClassifier
 from keras_nlp.models.gpt2.gpt2_models import GPT2
 from keras_nlp.models.roberta.roberta_models import Roberta
diff --git a/keras_nlp/models/distilbert/distilbert_preprocessing.py b/keras_nlp/models/distilbert/distilbert_preprocessing.py
--- a/keras_nlp/models/distilbert/distilbert_preprocessing.py
+++ b/keras_nlp/models/distilbert/distilbert_preprocessing.py
@@ -17,39 +17,126 @@
 
 from keras_nlp.layers.multi_segment_packer import MultiSegmentPacker
 from keras_nlp.tokenizers.word_piece_tokenizer import WordPieceTokenizer
+from keras_nlp.utils.python_utils import classproperty
+
+
+@keras.utils.register_keras_serializable(package="keras_nlp")
+class DistilBertTokenizer(WordPieceTokenizer):
+    """A DistilBERT tokenizer using WordPiece subword segmentation.
+
+    This tokenizer class will tokenize raw strings into integer sequences and
+    is based on `keras_nlp.tokenizers.WordPieceTokenizer`. Unlike the
+    underlying tokenizer, it will check for all special tokens needed by DistilBERT
+    models and provides a `from_preset()` method to automatically download
+    a matching vocabulary for a DistilBERT preset.
+
+    This tokenizer does not provide truncation or padding of inputs. It can be
+    combined with a `keras_nlp.models.DistilBertPreprocessor` layer for input packing.
+
+    If input is a batch of strings (rank > 0), the layer will output a
+    `tf.RaggedTensor` where the last dimension of the output is ragged.
+
+    If input is a scalar string (rank == 0), the layer will output a dense
+    `tf.Tensor` with static shape `[None]`.
+
+    Args:
+        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.
+        lowercase: If true, the input text will be first lowered before
+            tokenization.
+
+    Examples:
+
+    Batched input.
+    >>> vocab = ["[UNK]", "[CLS]", "[SEP]", "[PAD]"]
+    >>> vocab += ["The", "qu", "##ick", "brown", "fox", "."]
+    >>> inputs = ["The quick brown fox.", "The fox."]
+    >>> tokenizer = keras_nlp.models.DistilBertTokenizer(vocabulary=vocab)
+    >>> tokenizer(inputs)
+    <tf.RaggedTensor [[4, 5, 6, 7, 8, 9], [4, 8, 9]]>
+
+    Unbatched input.
+    >>> vocab = ["[UNK]", "[CLS]", "[SEP]", "[PAD]"]
+    >>> vocab += ["The", "qu", "##ick", "brown", "fox", "."]
+    >>> inputs = "The fox."
+    >>> tokenizer = keras_nlp.models.DistilBertTokenizer(vocabulary=vocab)
+    >>> tokenizer(inputs)
+    <tf.Tensor: shape=(3,), dtype=int32, numpy=array([4, 8, 9], dtype=int32)>
+
+    Detokenization.
+    >>> vocab = ["[UNK]", "[CLS]", "[SEP]", "[PAD]"]
+    >>> vocab += ["The", "qu", "##ick", "brown", "fox", "."]
+    >>> inputs = "The quick brown fox."
+    >>> tokenizer = keras_nlp.models.DistilBertTokenizer(vocabulary=vocab)
+    >>> tokenizer.detokenize(tokenizer.tokenize(inputs)).numpy().decode('utf-8')
+    'The quick brown fox .'
+    """
+
+    def __init__(
+        self,
+        vocabulary,
+        lowercase=False,
+        **kwargs,
+    ):
+        super().__init__(
+            vocabulary=vocabulary,
+            lowercase=lowercase,
+            **kwargs,
+        )
+
+        # Check for necessary special tokens.
+        cls_token = "[CLS]"
+        sep_token = "[SEP]"
+        pad_token = "[PAD]"
+        for token in [cls_token, pad_token, sep_token]:
+            if token not in self.get_vocabulary():
+                raise ValueError(
+                    f"Cannot find token `'{token}'` in the provided "
+                    f"`vocabulary`. Please provide `'{token}'` in your "
+                    "`vocabulary` or use a pretrained `vocabulary` name."
+                )
+
+        self.cls_token_id = self.token_to_id(cls_token)
+        self.sep_token_id = self.token_to_id(sep_token)
+        self.pad_token_id = self.token_to_id(pad_token)
+
+    @classproperty
+    def presets(cls):
+        raise NotImplementedError
+
+    @classmethod
+    def from_preset(
+        cls,
+        preset,
+        **kwargs,
+    ):
+        raise NotImplementedError
 
 
 @keras.utils.register_keras_serializable(package="keras_nlp")
 class DistilBertPreprocessor(keras.layers.Layer):
-    """DistilBERT preprocessing layer.
+    """A DistilBERT preprocessing layer which tokenizes and packs inputs.
 
     This preprocessing layer will do three things:
 
-     - Tokenize any number of inputs using a
-       `keras_nlp.tokenizers.WordPieceTokenizer`.
+     - Tokenize any number of inputs using the `tokenizer`.
      - Pack the inputs together using a `keras_nlp.layers.MultiSegmentPacker`.
        with the appropriate `"[CLS]"`, `"[SEP]"` and `"[PAD]"` tokens.
-     - Construct a dictionary of with keys `"token_ids"`, `"padding_mask"`,
+     - Construct a dictionary of with keys `"token_ids"` and `"padding_mask"`,
        that can be passed directly to a DistilBERT 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 WordPiece tokenizer can be accessed via the `tokenizer` property on this
-    layer, and can be used directly for custom packing on inputs.
-
     Args:
-        vocabulary: A list of vocabulary terms or a vocabulary filename.
-        lowercase: If `True`, input will be lowercase before tokenization. If
-            `vocabulary` is set to a pretrained vocabulary, this parameter will
-            be inferred.
-        sequence_length: The length of the packed inputs. Only used if
-            `pack_inputs` is True.
-        truncate: The algorithm to truncate a list of batched segments to fit
-            within `sequence_length`. Only used if
-            `pack_inputs` is True. The value can be either `round_robin` or
-            `waterfall`:
+        tokenizer: A `keras_nlp.models.DistilBertTokenizer` instance.
+        sequence_length: The length of the packed inputs.
+        truncate: string. The algorithm to truncate a list of batched segments
+            to fit within `sequence_length`. The value can be either
+            `round_robin` or `waterfall`:
                 - `"round_robin"`: Available space is assigned one token at a
                     time in a round-robin fashion to the inputs that still need
                     some, until the limit is reached.
@@ -60,9 +147,13 @@ class DistilBertPreprocessor(keras.layers.Layer):
 
     Examples:
     ```python
-    preprocessor = keras_nlp.models.DistilBertPreprocessor(
-        vocabulary="./vocab.txt",
-    )
+    vocab = ["[PAD]", "[UNK]", "[CLS]", "[SEP]"]
+    vocab += ["The", "qu", "##ick", "br", "##own", "fox", "tripped"]
+    vocab += ["Call", "me", "Ish", "##mael", "."]
+    vocab += ["Oh", "look", "a", "whale"]
+    vocab += ["I", "forgot", "my", "home", "##work"]
+    tokenizer = keras_nlp.models.DistilBertTokenizer(vocabulary=vocab)
+    preprocessor = keras_nlp.models.DistilBertPreprocessor(tokenizer)
 
     # Tokenize and pack a single sentence directly.
     preprocessor("The quick brown fox jumped.")
@@ -84,7 +175,10 @@ class DistilBertPreprocessor(keras.layers.Layer):
     second_sentences = ["The fox tripped.", "Oh look, a whale."]
     labels = [1, 1]
     ds = tf.data.Dataset.from_tensor_slices(
-        ((first_sentences, second_sentences), labels))
+        (
+            (first_sentences, second_sentences), labels
+        )
+    )
     ds = ds.map(
         lambda x, y: (preprocessor(x), y),
         num_parallel_calls=tf.data.AUTOTUNE,
@@ -94,56 +188,43 @@ class DistilBertPreprocessor(keras.layers.Layer):
 
     def __init__(
         self,
-        vocabulary="uncased_en",
-        lowercase=False,
+        tokenizer,
         sequence_length=512,
         truncate="round_robin",
         **kwargs,
     ):
         super().__init__(**kwargs)
-
-        self.tokenizer = WordPieceTokenizer(
-            vocabulary=vocabulary,
-            lowercase=lowercase,
-        )
-
-        # Check for necessary special tokens.
-        cls_token = "[CLS]"
-        sep_token = "[SEP]"
-        pad_token = "[PAD]"
-        for token in [cls_token, pad_token, sep_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._tokenizer = tokenizer
         self.packer = MultiSegmentPacker(
-            start_value=self.tokenizer.token_to_id(cls_token),
-            end_value=self.tokenizer.token_to_id(sep_token),
-            pad_value=self.pad_token_id,
+            start_value=self.tokenizer.cls_token_id,
+            end_value=self.tokenizer.sep_token_id,
+            pad_value=self.tokenizer.pad_token_id,
             truncate=truncate,
             sequence_length=sequence_length,
         )
 
-    def vocabulary_size(self):
-        """Returns the vocabulary size of the tokenizer."""
-        return self.tokenizer.vocabulary_size()
+    @property
+    def tokenizer(self):
+        """The `keras_nlp.models.DistilBertTokenizer` used to tokenize strings."""
+        return self._tokenizer
 
     def get_config(self):
         config = super().get_config()
         config.update(
             {
-                "vocabulary": self.tokenizer.vocabulary,
-                "lowercase": self.tokenizer.lowercase,
+                "tokenizer": keras.layers.serialize(self.tokenizer),
                 "sequence_length": self.packer.sequence_length,
                 "truncate": self.packer.truncate,
             }
         )
         return config
 
+    @classmethod
+    def from_config(cls, config):
+        if "tokenizer" in config:
+            config["tokenizer"] = keras.layers.deserialize(config["tokenizer"])
+        return cls(**config)
+
     def call(self, inputs):
         if not isinstance(inputs, (list, tuple)):
             inputs = [inputs]
@@ -152,9 +233,13 @@ def call(self, inputs):
         token_ids, _ = self.packer(inputs)
         return {
             "token_ids": token_ids,
-            "padding_mask": token_ids != self.pad_token_id,
+            "padding_mask": token_ids != self.tokenizer.pad_token_id,
         }
 
+    @classproperty
+    def presets(cls):
+        raise NotImplementedError
+
     @classmethod
     def from_preset(
         cls,


Test patch:
diff --git a/keras_nlp/models/distilbert/distilbert_preprocessing_test.py b/keras_nlp/models/distilbert/distilbert_preprocessing_test.py
--- a/keras_nlp/models/distilbert/distilbert_preprocessing_test.py
+++ b/keras_nlp/models/distilbert/distilbert_preprocessing_test.py
@@ -13,14 +13,73 @@
 # limitations under the License.
 """Tests for DistilBERT preprocessing layers."""
 
+import os
+
 import tensorflow as tf
+from absl.testing import parameterized
+from tensorflow import keras
 
 from keras_nlp.models.distilbert.distilbert_preprocessing import (
     DistilBertPreprocessor,
 )
+from keras_nlp.models.distilbert.distilbert_preprocessing import (
+    DistilBertTokenizer,
+)
+
+
+class DistilBertTokenizerTest(tf.test.TestCase, parameterized.TestCase):
+    def setUp(self):
+        self.vocab = ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"]
+        self.vocab += ["THE", "QUICK", "BROWN", "FOX"]
+        self.vocab += ["the", "quick", "brown", "fox"]
+
+    def test_tokenize(self):
+        input_data = "THE QUICK BROWN FOX."
+        tokenizer = DistilBertTokenizer(vocabulary=self.vocab)
+        output = tokenizer(input_data)
+        self.assertAllEqual(output, [5, 6, 7, 8, 1])
+
+    def test_tokenize_batch(self):
+        input_data = tf.constant(["THE QUICK BROWN FOX.", "THE FOX."])
+        tokenizer = DistilBertTokenizer(vocabulary=self.vocab)
+        output = tokenizer(input_data)
+        self.assertAllEqual(output, [[5, 6, 7, 8, 1], [5, 8, 1]])
+
+    def test_lowercase(self):
+        input_data = "THE QUICK BROWN FOX."
+        tokenizer = DistilBertTokenizer(vocabulary=self.vocab, lowercase=True)
+        output = tokenizer(input_data)
+        self.assertAllEqual(output, [9, 10, 11, 12, 1])
+
+    def test_detokenize(self):
+        input_tokens = [[5, 6, 7, 8]]
+        tokenizer = DistilBertTokenizer(vocabulary=self.vocab)
+        output = tokenizer.detokenize(input_tokens)
+        self.assertAllEqual(output, ["THE QUICK BROWN FOX"])
+
+    def test_vocabulary_size(self):
+        tokenizer = DistilBertTokenizer(vocabulary=self.vocab)
+        self.assertEqual(tokenizer.vocabulary_size(), 13)
+
+    @parameterized.named_parameters(
+        ("save_format_tf", "tf"), ("save_format_h5", "h5")
+    )
+    def test_saving_model(self, save_format):
+        input_data = tf.constant(["THE QUICK BROWN FOX."])
+        tokenizer = DistilBertTokenizer(vocabulary=self.vocab)
+        inputs = keras.Input(dtype="string", shape=())
+        outputs = tokenizer(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),
+            restored_model(input_data),
+        )
 
 
-class DistilBertPreprocessorTest(tf.test.TestCase):
+class DistilBertPreprocessorTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
         self.vocab = ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"]
         self.vocab += ["THE", "QUICK", "BROWN", "FOX"]
@@ -29,29 +88,90 @@ def setUp(self):
     def test_tokenize(self):
         input_data = ["THE QUICK BROWN FOX."]
         preprocessor = DistilBertPreprocessor(
-            vocabulary=self.vocab,
+            DistilBertTokenizer(vocabulary=self.vocab),
             sequence_length=8,
         )
         output = preprocessor(input_data)
         self.assertAllEqual(output["token_ids"], [2, 5, 6, 7, 8, 1, 3, 0])
         self.assertAllEqual(output["padding_mask"], [1, 1, 1, 1, 1, 1, 1, 0])
 
-    def test_lowercase(self):
-        input_data = ["THE QUICK BROWN FOX."]
+    def test_tokenize_batch(self):
+        input_data = tf.constant(
+            [
+                "THE QUICK BROWN FOX.",
+                "THE QUICK BROWN FOX.",
+                "THE QUICK BROWN FOX.",
+                "THE QUICK BROWN FOX.",
+            ]
+        )
         preprocessor = DistilBertPreprocessor(
-            vocabulary=self.vocab,
+            DistilBertTokenizer(vocabulary=self.vocab),
             sequence_length=8,
-            lowercase=True,
         )
         output = preprocessor(input_data)
-        self.assertAllEqual(output["token_ids"], [2, 9, 10, 11, 12, 1, 3, 0])
+        self.assertAllEqual(output["token_ids"], [[2, 5, 6, 7, 8, 1, 3, 0]] * 4)
+        self.assertAllEqual(
+            output["padding_mask"], [[1, 1, 1, 1, 1, 1, 1, 0]] * 4
+        )
 
-    def test_detokenize(self):
-        input_data = [[5, 6, 7, 8]]
-        preprocessor = DistilBertPreprocessor(vocabulary=self.vocab)
-        output = preprocessor.tokenizer.detokenize(input_data)
-        self.assertAllEqual(output, ["THE QUICK BROWN FOX"])
+    def test_tokenize_multiple_sentences(self):
+        sentence_one = "THE QUICK"
+        sentence_two = "BROWN FOX."
+        preprocessor = DistilBertPreprocessor(
+            DistilBertTokenizer(vocabulary=self.vocab),
+            sequence_length=8,
+        )
+        # The first tuple or list is always interpreted as an enumeration of
+        # separate sequences to concatenate.
+        output = preprocessor((sentence_one, sentence_two))
+        self.assertAllEqual(output["token_ids"], [2, 5, 6, 3, 7, 8, 1, 3])
+        self.assertAllEqual(output["padding_mask"], [1, 1, 1, 1, 1, 1, 1, 1])
 
-    def test_vocabulary_size(self):
-        preprocessor = DistilBertPreprocessor(vocabulary=self.vocab)
-        self.assertEqual(preprocessor.vocabulary_size(), 13)
+    def test_tokenize_multiple_batched_sentences(self):
+        sentence_one = tf.constant(
+            [
+                "THE QUICK",
+                "THE QUICK",
+                "THE QUICK",
+                "THE QUICK",
+            ]
+        )
+        sentence_two = tf.constant(
+            [
+                "BROWN FOX.",
+                "BROWN FOX.",
+                "BROWN FOX.",
+                "BROWN FOX.",
+            ]
+        )
+        preprocessor = DistilBertPreprocessor(
+            DistilBertTokenizer(vocabulary=self.vocab),
+            sequence_length=8,
+        )
+        # The first tuple or list is always interpreted as an enumeration of
+        # separate sequences to concatenate.
+        output = preprocessor((sentence_one, sentence_two))
+        self.assertAllEqual(output["token_ids"], [[2, 5, 6, 3, 7, 8, 1, 3]] * 4)
+        self.assertAllEqual(
+            output["padding_mask"], [[1, 1, 1, 1, 1, 1, 1, 1]] * 4
+        )
+
+    @parameterized.named_parameters(
+        ("save_format_tf", "tf"), ("save_format_h5", "h5")
+    )
+    def test_saving_model(self, save_format):
+        input_data = tf.constant(["THE QUICK BROWN FOX."])
+        preprocessor = DistilBertPreprocessor(
+            DistilBertTokenizer(vocabulary=self.vocab),
+            sequence_length=8,
+        )
+        inputs = keras.Input(dtype="string", shape=())
+        outputs = 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"],
+        )