Repository: keras-team/keras-nlp

Issue:
Port DistilBert to follow Bert model changes
We had a few larger changes to the bert API:
https://github.com/keras-team/keras-nlp/pull/387
https://github.com/keras-team/keras-nlp/pull/390

As well as some minor cleanups:
https://github.com/keras-team/keras-nlp/pull/397
https://github.com/keras-team/keras-nlp/pull/398
https://github.com/keras-team/keras-nlp/pull/399

We should mirror these where applicable to the DistilBert model code. Since we do not yet host pretrained weight for this model, we can just raise a `NotImplementedError` in `from_preset` for now.


Discussion:

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
@@ -15,8 +15,7 @@
 from keras_nlp.models.bert.bert_models import Bert
 from keras_nlp.models.bert.bert_preprocessing import BertPreprocessor
 from keras_nlp.models.bert.bert_tasks import BertClassifier
-from keras_nlp.models.distilbert.distilbert_models import DistilBertBase
-from keras_nlp.models.distilbert.distilbert_models import DistilBertCustom
+from keras_nlp.models.distilbert.distilbert_models import DistilBert
 from keras_nlp.models.distilbert.distilbert_preprocessing import (
     DistilBertPreprocessor,
 )
diff --git a/keras_nlp/models/distilbert/distilbert_models.py b/keras_nlp/models/distilbert/distilbert_models.py
--- a/keras_nlp/models/distilbert/distilbert_models.py
+++ b/keras_nlp/models/distilbert/distilbert_models.py
@@ -29,8 +29,8 @@ def distilbert_kernel_initializer(stddev=0.02):
 
 
 @keras.utils.register_keras_serializable(package="keras_nlp")
-class DistilBertCustom(keras.Model):
-    """DistilBERT encoder network with custom hyperparmeters.
+class DistilBert(keras.Model):
+    """DistilBERT encoder network.
 
     This network implements a bi-directional Transformer-based encoder as
     described in ["DistilBERT, a distilled version of BERT: smaller, faster,
@@ -38,9 +38,10 @@ class DistilBertCustom(keras.Model):
     embedding lookups and transformer layers, but not the masked language model
     or classification task networks.
 
-    This class gives a fully customizable DistilBERT model with any number of layers,
-    heads, and embedding dimensions. For specific DistilBERT architectures defined in
-    the paper, see, for example, `keras_nlp.models.DistilBertBase`.
+    The default constructor gives a fully customizable, randomly initalized
+    DistilBERT encoder with any number of layers, heads, and embedding
+    dimensions. To load preset architectures and weights, use the `from_presets`
+    constructor.
 
     Args:
         vocabulary_size: int. The size of the token vocabulary.
@@ -55,32 +56,27 @@ class DistilBertCustom(keras.Model):
             can consume. If None, `max_sequence_length` uses the value from
             sequence length. This determines the variable shape for positional
             embeddings.
-        name: string, optional. Name of the model.
-        trainable: boolean, optional. If the model's variables should be
-            trainable.
 
     Examples:
     ```python
-    # Randomly initialized DistilBERT encoder
-    model = keras_nlp.models.DistilBertCustom(
-        vocabulary_size=30522,
-        num_layers=12,
-        num_heads=12,
-        hidden_dim=768,
-        intermediate_dim=3072,
-        max_sequence_length=12,
-        name="encoder",
-    )
-
-    # Call encoder on the inputs
     input_data = {
         "token_ids": tf.random.uniform(
-            shape=(1, 12), dtype=tf.int64, maxval=model.vocabulary_size
+            shape=(1, 12), dtype=tf.int64, maxval=30522
         ),
         "padding_mask": tf.constant(
             [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
         ),
     }
+
+    # Randomly initialized DistilBERT encoder
+    model = keras_nlp.models.DistilBert(
+        vocabulary_size=30522,
+        num_layers=6,
+        num_heads=12,
+        hidden_dim=768,
+        intermediate_dim=3072,
+        max_sequence_length=12,
+    )
     output = model(input_data)
     ```
     """
@@ -94,8 +90,7 @@ def __init__(
         intermediate_dim,
         dropout=0.1,
         max_sequence_length=512,
-        name=None,
-        trainable=True,
+        **kwargs,
     ):
 
         # Inputs
@@ -138,6 +133,10 @@ def __init__(
                 name=f"transformer_layer_{i}",
             )(x, padding_mask=padding_mask)
 
+        # Set default for `name` if none given
+        if "name" not in kwargs:
+            kwargs["name"] = "backbone"
+
         # Instantiate using Functional API Model constructor
         super().__init__(
             inputs={
@@ -145,8 +144,7 @@ def __init__(
                 "padding_mask": padding_mask,
             },
             outputs=x,
-            name=name,
-            trainable=trainable,
+            **kwargs,
         )
         # All references to `self` below this line
         self.vocabulary_size = vocabulary_size
@@ -174,64 +172,11 @@ def get_config(self):
     def from_config(cls, config):
         return cls(**config)
 
-
-MODEL_DOCSTRING = """DistilBert "{type}" architecture.
-
-This network implements a bi-directional Transformer-based encoder as
-described in ["DistilBERT, a distilled version of BERT: smaller, faster,
-cheaper and lighter"](https://arxiv.org/abs/1910.01108). It includes the
-embedding lookups and transformer layers, but not the masked language model
-or classification task networks.
-
-Args:
-    vocabulary_size: int, optional. The size of the token vocabulary.
-    name: string, optional. Name of the model.
-    trainable: boolean, optional. If the model's variables should be
-        trainable.
-
-Examples:
-```python
-# Randomly initialized DistilBert{type} encoder
-model = keras_nlp.models.DistilBert{type}(vocabulary_size=10000)
-
-# Call encoder on the inputs.
-input_data = {{
-    "token_ids": tf.random.uniform(
-        shape=(1, 512), dtype=tf.int64, maxval=model.vocabulary_size
-    ),
-    "padding_mask": tf.constant([1] * 512, shape=(1, 512)),
-}}
-output = model(input_data)
-```
-"""
-
-
-def DistilBertBase(vocabulary_size, name=None, trainable=True):
-
-    model = DistilBertCustom(
-        vocabulary_size=vocabulary_size,
-        num_layers=6,
-        num_heads=12,
-        hidden_dim=768,
-        intermediate_dim=3072,
-        dropout=0.1,
-        max_sequence_length=512,
-        name=name,
-        trainable=trainable,
-    )
-
-    return model
-
-
-def model_class_by_name(classname):
-    """Return model class given the class name."""
-    return {
-        "DistilBertBase": DistilBertBase,
-    }[classname]
-
-
-setattr(
-    DistilBertBase,
-    "__doc__",
-    MODEL_DOCSTRING.format(type="Base"),
-)
+    @classmethod
+    def from_preset(
+        cls,
+        preset,
+        load_weights=True,
+        **kwargs,
+    ):
+        raise NotImplementedError
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
@@ -18,78 +18,80 @@
 from keras_nlp.layers.multi_segment_packer import MultiSegmentPacker
 from keras_nlp.tokenizers.word_piece_tokenizer import WordPieceTokenizer
 
-PREPROCESSOR_DOCSTRING = """DistilBERT preprocessor with pretrained vocabularies.
-
-This preprocessing layer will do three things:
-
- - Tokenize any number of inputs using a
-   `keras_nlp.tokenizers.WordPieceTokenizer`.
- - 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"`, 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: One of a list of vocabulary terms, a vocabulary filename, or
-        the name of the pretrained vocabulary.
-    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`:
-            - `"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
-preprocessor = keras_nlp.models.DistilBertPreprocessor(vocabulary="uncased_en")
-
-# Tokenize and pack a single sentence directly.
-preprocessor("The quick brown fox jumped.")
-
-# Tokenize and pack a multiple sentence directly.
-preprocessor(("The quick brown fox jumped.", "Call me Ishmael."))
-
-# Map a dataset to preprocess a single sentence.
-features = ["The quick brown fox jumped.", "I forgot my homework."]
-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 a multiple sentences.
-first_sentences = ["The quick brown fox jumped.", "Call me Ishmael."]
-second_sentences = ["The fox tripped.", "Oh look, a whale."]
-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,
-)
-```
-"""
-
 
 @keras.utils.register_keras_serializable(package="keras_nlp")
 class DistilBertPreprocessor(keras.layers.Layer):
+    """DistilBERT preprocessing layer.
+
+    This preprocessing layer will do three things:
+
+     - Tokenize any number of inputs using a
+       `keras_nlp.tokenizers.WordPieceTokenizer`.
+     - 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"`,
+       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`:
+                - `"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
+    preprocessor = keras_nlp.models.DistilBertPreprocessor(
+        vocabulary="./vocab.txt",
+    )
+
+    # Tokenize and pack a single sentence directly.
+    preprocessor("The quick brown fox jumped.")
+
+    # Tokenize and pack a multiple sentence directly.
+    preprocessor(("The quick brown fox jumped.", "Call me Ishmael."))
+
+    # Map a dataset to preprocess a single sentence.
+    features = ["The quick brown fox jumped.", "I forgot my homework."]
+    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 a multiple sentences.
+    first_sentences = ["The quick brown fox jumped.", "Call me Ishmael."]
+    second_sentences = ["The fox tripped.", "Oh look, a whale."]
+    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="uncased_en",
@@ -153,9 +155,12 @@ def call(self, inputs):
             "padding_mask": token_ids != self.pad_token_id,
         }
 
-
-setattr(
-    DistilBertPreprocessor,
-    "__doc__",
-    PREPROCESSOR_DOCSTRING,
-)
+    @classmethod
+    def from_preset(
+        cls,
+        preset,
+        sequence_length=None,
+        truncate="round_robin",
+        **kwargs,
+    ):
+        raise NotImplementedError


Test patch:
diff --git a/keras_nlp/models/distilbert/distilbert_models_test.py b/keras_nlp/models/distilbert/distilbert_models_test.py
--- a/keras_nlp/models/distilbert/distilbert_models_test.py
+++ b/keras_nlp/models/distilbert/distilbert_models_test.py
@@ -19,13 +19,12 @@
 from absl.testing import parameterized
 from tensorflow import keras
 
-from keras_nlp.models.distilbert.distilbert_models import DistilBertBase
-from keras_nlp.models.distilbert.distilbert_models import DistilBertCustom
+from keras_nlp.models.distilbert.distilbert_models import DistilBert
 
 
 class DistilBertTest(tf.test.TestCase, parameterized.TestCase):
     def setUp(self):
-        self.model = DistilBertCustom(
+        self.model = DistilBert(
             vocabulary_size=1000,
             num_layers=2,
             num_heads=2,
@@ -63,33 +62,19 @@ def test_variable_sequence_length_call_distilbert(self):
             }
             self.model(input_data)
 
-    def test_valid_call_distilbert_base(self):
-        model = DistilBertBase(vocabulary_size=1000, name="encoder")
-        input_data = {
-            "token_ids": tf.ones(
-                (self.batch_size, self.model.max_sequence_length), dtype="int32"
-            ),
-            "padding_mask": tf.ones(
-                (self.batch_size, self.model.max_sequence_length), dtype="int32"
-            ),
-        }
-        model(input_data)
-
     @parameterized.named_parameters(
         ("jit_compile_false", False), ("jit_compile_true", True)
     )
     def test_distilbert_base_compile(self, jit_compile):
-        model = DistilBertBase(vocabulary_size=1000, name="encoder")
-        model.compile(jit_compile=jit_compile)
-        model.predict(self.input_batch)
+        self.model.compile(jit_compile=jit_compile)
+        self.model.predict(self.input_batch)
 
     @parameterized.named_parameters(
         ("jit_compile_false", False), ("jit_compile_true", True)
     )
     def test_distilbert_base_compile_batched_ds(self, jit_compile):
-        model = DistilBertBase(vocabulary_size=1000, name="encoder")
-        model.compile(jit_compile=jit_compile)
-        model.predict(self.input_dataset)
+        self.model.compile(jit_compile=jit_compile)
+        self.model.predict(self.input_dataset)
 
     @parameterized.named_parameters(
         ("save_format_tf", "tf"), ("save_format_h5", "h5")
@@ -101,7 +86,7 @@ def test_saving_model(self, save_format):
         restored_model = keras.models.load_model(save_path)
 
         # Check we got the real object back.
-        self.assertIsInstance(restored_model, DistilBertCustom)
+        self.assertIsInstance(restored_model, DistilBert)
 
         # Check that output matches.
         restored_output = restored_model(self.input_batch)