Repository: keras-team/keras-nlp

Issue:
Add pretrained checkpoints for Bert
# Problem
We exposed an API for fully configurable and fixed Bert architectures in PR #288. Most users will not have the resources to do pretraining on their own, however. We should offer checkpoints for `models.BertBase` similar to those used in the original paper.

# Implementation

## Weights

Model checkpoints from the original paper repo ([gh](https://github.com/google-research/bert)) or similar source can be converted to be compatible with our implementation in PR #288. This code will be a one-off task and not added to the repo.

## Format

`keras-cv` is currently using `h5`, but that format is deprecated. An obvious alternative is a `tf.checkpoint` using `model.save_weights()`. The resulting `data` and `index` files will have to be zipped and uploaded to our cloud storage bucket.

## API

There some prior art in `keras-cv` and our pretrained vocab to use a string to identify a model of interest which is then checked against a dictionary (e.g., [weights.py](https://github.com/keras-team/keras-cv/blob/master/keras_cv/models/weights.py#L58)). However, this requires the user to dive into the code to find possible names and a lot of input validation and error handling in the implementation.

I propose to create an `enum` for each `Model` class (e.g., `models.BertBase`) with possible checkpoints that will expose these options as part of the API and work with users' autocomplete engines:
```
class BertBaseCkpt(Enum):
   uncased = 0
   cased = 1
   zh = 2
```

Under the hood this `enum` would then be the key for a dictionary of metadata for each model including checkpoint location, vocabulary location, and a description string explaining where the model came from and what data was used to train it. We could also write a method to print the description for each/all `enum` value(s) to again protect the user from having to read through utility functions.

```
base_models = {
   BertBaseCkpt.uncased: {
      "ckpt": ...
      "description": ...
      "vocab": ...
   }
}
```

Model checkpoints will be hosted in our Google Cloud storage bucket and loaded behind the scenes given an enum value. The user should never have to interact with this dictionary.

Here is the enhanced API:
```
def BertBase(ckpt=None, name=None, trainable=True):
    """Bi-directional Transformer-based encoder network (Bert) using "Base"
    architecture.

    This network implements a bi-directional Transformer-based encoder as
    described in ["BERT: Pre-training of Deep Bidirectional Transformers for
    Language Understanding"](https://arxiv.org/abs/1810.04805). It includes the
    embedding lookups and transformer layers, but not the masked language model
    or classification task networks.

    Args:
        ckpt: BertBaseCkpt. A checkpoint of pretrained weights to load into the model. 
             If `None`, the model will be randomly initialized.
        name: String, optional. Name of the model.
        trainable: Boolean, optional. If the model's variables should be
            trainable.
```

## Vocabulary
Note that this proposal does not cover the vocabulary file, which will be incorporated in our preprocessing work. One potential proposal is to host the `vocab.txt` file used to train the model in the same folder to be loaded if the user does not supply their own. The rationale is that except for very advanced use cases the vocabulary must match that used in training, and requiring the user to specify their own will expose our Cloud storage implementation to the user for little gain. However, this is beyond the scope of the proposal.


Discussion:
Thanks for putting up the nice design!

Some thoughts from my side:
1. Enum is something not commonly used in Keras code. My take is in general we prefer to use string identifier, @fchollet Tagging Francois for thoughts.
2. The arg `ckpt` is vague to our users, also it's not the actual checkpoint path. HuggingFace use a static method `from_pretrained`, which has a clear meaning "it is loading from pretrained model, and the pretrained model is XXX". 
3. checkpoint format is more preferrable than h5 saving format, as the later is technically deprecated.

Additionally, I would like to include the design of tokenizer saving/loading inside this discussion. A pretrained model should always come with its tokenizer, so we may not want to decouple the design. HuggingFace decouples the tokenizer from model, but uses the same string identifer, e.g., "bert-base-uncased" to glue those together. It seems we want to have tokenizer exist as a member of model, we should get the agreement. 
Thanks for the proposal! Some quick thoughts, will dig in more next week.

 - Using tf checkpoints sgtm. But probably good to check with KerasCV folk here. Will they be using tf checkpoints instead of h5 for newer weights they are hosting?
 - Arg name should probably still be `weights`, unless we have a good reason to change. We want a consistent API experience across keras as much as possible.
 - Also unsure about the enum type. Probably makes sense to follow the prior art in Keras unless we have a strong reason.
 - I wouldn't totally through away the checkpoint conversion code. We should try to keep all these in a tools or scripts directory, outside of `keras_nlp/' (the dir with library code). We should be able to reproduce the conversion as needed in the future, using checked in code (though not shipped code in our library).
Thanks for the comments! Talking to @LukeWood the model loading experience for `keras-cv` is not yet polished. The point of this proposal was to explore the design space so we can optimize the user experience across both packages. Maybe `keras-cv` will want to use our new method. Rejecting anything new out of hand means we cannot innovate.

This proposal was partly inspired by the `lang` handling in PR #287, where we want to prompt the user about which languages are supported and do input validation on the result. While a tiny change, the `enum` would solve this issue naturally. 

Resolving patch:
diff --git a/keras_nlp/models/bert.py b/keras_nlp/models/bert.py
--- a/keras_nlp/models/bert.py
+++ b/keras_nlp/models/bert.py
@@ -25,6 +25,27 @@ def _bert_kernel_initializer(stddev=0.02):
     return keras.initializers.TruncatedNormal(stddev=stddev)
 
 
+# Pretrained models
+BASE_PATH = "https://storage.googleapis.com/keras-nlp/models/"
+
+checkpoints = {
+    "bert_base": {
+        "uncased_en": {
+            "md5": "9b2b2139f221988759ac9cdd17050b31",
+            "description": "Base size of Bert where all input is lowercased. "
+            "Trained on English wikipedia + books corpora.",
+            "vocabulary_size": 30522,
+        },
+        "cased_en": {
+            "md5": "f94a6cb012e18f4fb8ec92abb91864e9",
+            "description": "Base size of Bert where case is maintained. "
+            "Trained on English wikipedia + books corpora.",
+            "vocabulary_size": 28996,
+        },
+    }
+}
+
+
 class BertCustom(keras.Model):
     """Bi-directional Transformer-based encoder network.
 
@@ -61,25 +82,29 @@ class BertCustom(keras.Model):
     Example usage:
     ```python
     # Randomly initialized Bert encoder
-    encoder = keras_nlp.models.BertCustom(
+    model = keras_nlp.models.BertCustom(
         vocabulary_size=30522,
         num_layers=12,
         num_heads=12,
         hidden_dim=768,
         intermediate_dim=3072,
-        max_sequence_length=12
+        max_sequence_length=12,
+        name="encoder",
     )
 
-     # Call encoder on the inputs.
+    # Call encoder on the inputs
     input_data = {
         "input_ids": tf.random.uniform(
-            shape=(1, 12), dtype=tf.int64, maxval=30522),
+            shape=(1, 12), dtype=tf.int64, maxval=model.vocabulary_size
+        ),
         "segment_ids": tf.constant(
-            [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)),
+            [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
+        ),
         "input_mask": tf.constant(
-            [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)),
+            [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
+        ),
     }
-    output = encoder(input_data)
+    output = model(input_data)
     ```
     """
 
@@ -156,7 +181,7 @@ def __init__(
                 ),
                 dropout=dropout,
                 kernel_initializer=_bert_kernel_initializer(),
-                name=f"""transformer_layer_{i}""",
+                name=f"transformer_layer_{i}",
             )(x, padding_mask=input_mask)
 
         # Construct the two Bert outputs. The pooled output is a dense layer on
@@ -225,10 +250,9 @@ class BertClassifier(keras.Model):
             trainable.
 
     Example usage:
-    ```
-    python
+    ```python
     # Randomly initialized Bert encoder
-    encoder = keras_nlp.models.BertCustom(
+    model = keras_nlp.models.BertCustom(
         vocabulary_size=30522,
         num_layers=12,
         num_heads=12,
@@ -237,16 +261,19 @@ class BertClassifier(keras.Model):
         max_sequence_length=12
     )
 
-     # Call classifier on the inputs.
+    # Call classifier on the inputs.
     input_data = {
         "input_ids": tf.random.uniform(
-            shape=(1, 12), dtype=tf.int64, maxval=30522),
+            shape=(1, 12), dtype=tf.int64, maxval=30522
+        ),
         "segment_ids": tf.constant(
-            [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)),
+            [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
+        ),
         "input_mask": tf.constant(
-            [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)),
+            [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
+        ),
     }
-    classifier = keras_nlp.models.BertClassifier(encoder, 4)
+    classifier = bert.BertClassifier(model, 4, name="classifier")
     logits = classifier(input_data)
     ```
     """
@@ -274,9 +301,8 @@ def __init__(
         self.num_classes = num_classes
 
 
-def BertBase(name=None, trainable=True):
-    """Bi-directional Transformer-based encoder network (Bert) using "Base"
-    architecture.
+MODEL_DOCSTRING = """Bi-directional Transformer-based encoder network (Bert)
+    using "{type}" architecture.
 
     This network implements a bi-directional Transformer-based encoder as
     described in ["BERT: Pre-training of Deep Bidirectional Transformers for
@@ -285,6 +311,12 @@ def BertBase(name=None, trainable=True):
     or classification task networks.
 
     Args:
+        weights: String, optional. Name of pretrained model to load weights.
+            Should be one of {names}.
+            If None, model is randomly initialized. Either `weights` or
+            `vocabulary_size` must be specified, but not both.
+        vocabulary_size: Int, optional. The size of the token vocabulary. Either
+            `weights` or `vocabularly_size` must be specified, but not both.
         name: String, optional. Name of the model.
         trainable: Boolean, optional. If the model's variables should be
             trainable.
@@ -292,23 +324,49 @@ def BertBase(name=None, trainable=True):
     Example usage:
     ```python
     # Randomly initialized BertBase encoder
-    encoder = keras_nlp.models.BertBase()
+    model = keras_nlp.models.BertBase(vocabulary_size=10000)
 
-     # Call encoder on the inputs.
-    input_data = {
+    # Call encoder on the inputs.
+    input_data = {{
         "input_ids": tf.random.uniform(
-            shape=(1, 512), dtype=tf.int64, maxval=encoder.vocabulary_size),
-        "segment_ids": tf.constant(
-            [0] * 200 + [1] * 312, shape=(1, 512)),
-        "input_mask": tf.constant(
-            [1] * 512, shape=(1, 512)),
-    }
-    output = encoder(input_data)
+            shape=(1, 512), dtype=tf.int64, maxval=model.vocabulary_size
+        ),
+        "segment_ids": tf.constant([0] * 200 + [1] * 312, shape=(1, 512)),
+        "input_mask": tf.constant([1] * 512, shape=(1, 512)),
+    }}
+    output = model(input_data)
+
+    # Load a pretrained model
+    model = keras_nlp.models.BertBase(weights="uncased_en")
+    # Call encoder on the inputs.
+    output = model(input_data)
     ```
-    """
+"""
+
+
+def BertBase(weights=None, vocabulary_size=None, name=None, trainable=True):
+
+    if (vocabulary_size is None and weights is None) or (
+        vocabulary_size and weights
+    ):
+        raise ValueError(
+            "One of `vocabulary_size` or `weights` must be specified "
+            "(but not both). "
+            f"Received: weights={weights}, "
+            f"vocabulary_size={vocabulary_size}"
+        )
+
+    if weights:
+        if weights not in checkpoints["bert_base"]:
+            raise ValueError(
+                "`weights` must be one of "
+                f"""{", ".join(checkpoints["bert_base"])}. """
+                f"Received: {weights}"
+            )
+        vocabulary_size = checkpoints["bert_base"][weights]["vocabulary_size"]
 
     model = BertCustom(
-        vocabulary_size=30522,
+        vocabulary_size=vocabulary_size,
         num_layers=12,
         num_heads=12,
         hidden_dim=768,
@@ -319,6 +377,26 @@ def BertBase(name=None, trainable=True):
         trainable=trainable,
     )
 
-    # TODO(jbischof): add some documentation or magic to load our checkpoints
-    # TODO(jbischof): attach the tokenizer
+    # TODO(jbischof): consider changing format from `h5` to
+    # `tf.train.Checkpoint` once
+    # https://github.com/keras-team/keras/issues/16946 is resolved
+    if weights:
+        filepath = keras.utils.get_file(
+            "model.h5",
+            BASE_PATH + "bert_base_" + weights + "/model.h5",
+            cache_subdir="models/bert_base/" + weights + "/",
+            file_hash=checkpoints["bert_base"][weights]["md5"],
+        )
+        model.load_weights(filepath)
+
+    # TODO(jbischof): attach the tokenizer or create separate tokenizer class
     return model
+
+
+setattr(
+    BertBase,
+    "__doc__",
+    MODEL_DOCSTRING.format(
+        type="Base", names=", ".join(checkpoints["bert_base"])
+    ),
+)


Test patch:
diff --git a/keras_nlp/models/bert_test.py b/keras_nlp/models/bert_test.py
--- a/keras_nlp/models/bert_test.py
+++ b/keras_nlp/models/bert_test.py
@@ -34,7 +34,7 @@ def test_valid_call_bert(self):
         )
         input_data = {
             "input_ids": tf.random.uniform(
-                shape=(1, 12), dtype=tf.int64, maxval=30522
+                shape=(1, 12), dtype=tf.int64, maxval=model.vocabulary_size
             ),
             "segment_ids": tf.constant(
                 [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
@@ -88,7 +88,7 @@ def test_valid_call_classifier(self):
         classifier(input_data)
 
     def test_valid_call_bert_base(self):
-        model = bert.BertBase(name="encoder")
+        model = bert.BertBase(vocabulary_size=10000, name="encoder")
         input_data = {
             "input_ids": tf.random.uniform(
                 shape=(1, 512), dtype=tf.int64, maxval=model.vocabulary_size
@@ -98,6 +98,26 @@ def test_valid_call_bert_base(self):
         }
         model(input_data)
 
+    def test_bert_base_vocab_error(self):
+        # Need `vocabulary_size` or `weights`
+        with self.assertRaises(ValueError):
+            bert.BertBase(name="encoder")
+
+        # Only one of `vocabulary_size` or `weights`
+        with self.assertRaises(ValueError):
+            bert.BertBase(
+                weights="bert_base_uncased",
+                vocabulary_size=10000,
+                name="encoder",
+            )
+
+        # Not a checkpoint name
+        with self.assertRaises(ValueError):
+            bert.BertBase(
+                weights="bert_base_clowntown",
+                name="encoder",
+            )
+
     def test_saving_model(self):
         model = bert.BertCustom(
             vocabulary_size=30522,
@@ -110,7 +130,7 @@ def test_saving_model(self):
         )
         input_data = {
             "input_ids": tf.random.uniform(
-                shape=(1, 12), dtype=tf.int64, maxval=30522
+                shape=(1, 12), dtype=tf.int64, maxval=model.vocabulary_size
             ),
             "segment_ids": tf.constant(
                 [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
diff --git a/network_tests/test_load_ckpts.py b/network_tests/test_load_ckpts.py
new file mode 100644
--- /dev/null
+++ b/network_tests/test_load_ckpts.py
@@ -0,0 +1,42 @@
+# 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 loading pretrained model checkpoints."""
+
+import tensorflow as tf
+
+import keras_nlp
+
+
+class BertCkptTest(tf.test.TestCase):
+    def test_load_bert_base_uncased(self):
+        model = keras_nlp.models.BertBase(weights="uncased_en")
+        input_data = {
+            "input_ids": tf.random.uniform(
+                shape=(1, 512), dtype=tf.int64, maxval=model.vocabulary_size
+            ),
+            "segment_ids": tf.constant([0] * 200 + [1] * 312, shape=(1, 512)),
+            "input_mask": tf.constant([1] * 512, shape=(1, 512)),
+        }
+        model(input_data)
+
+    def test_load_bert_base_cased(self):
+        model = keras_nlp.models.BertBase(weights="cased_en")
+        input_data = {
+            "input_ids": tf.random.uniform(
+                shape=(1, 512), dtype=tf.int64, maxval=model.vocabulary_size
+            ),
+            "segment_ids": tf.constant([0] * 200 + [1] * 312, shape=(1, 512)),
+            "input_mask": tf.constant([1] * 512, shape=(1, 512)),
+        }
+        model(input_data)