Repository: keras-team/keras-nlp

Issue:
Add bert tiny, small, medium network and pretrained weights 
The official bert repo lists a number of smaller variants and provides checkpoints.
https://github.com/google-research/bert

We should provide these as options in our bert model offering. Once we do this, we can also delete the config object for our bert example code, and use the randomly initialized architectures directly for our pretraining example:
https://github.com/keras-team/keras-nlp/blob/master/examples/bert/bert_config.py#L15


Discussion:
A couple notes:

I'm proposing we ditch "bert mini" as an offering. That seems like it would avoid some clutter, while still giving us good coverage of sizes of bert.
 - tiny: 2 layer
 - small: 4 layer
 - medium: 8 layer
 - base: 12 layer
 - large: 24 layer

These model will only be available in english, uncased for now. We may want to improve our error messages so that someone trying to instantiate `keras_nlp.models.BertSmall(weights="zh")` is told to look for other sizes of bert that do support that option.

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
@@ -16,6 +16,9 @@
 from keras_nlp.models.bert import BertClassifier
 from keras_nlp.models.bert import BertCustom
 from keras_nlp.models.bert import BertLarge
+from keras_nlp.models.bert import BertMedium
+from keras_nlp.models.bert import BertSmall
+from keras_nlp.models.bert import BertTiny
 from keras_nlp.models.roberta import RobertaBase
 from keras_nlp.models.roberta import RobertaClassifier
 from keras_nlp.models.roberta import RobertaCustom
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
@@ -68,6 +68,30 @@ def _handle_weights_and_vocab_size(bert_variant, weights, vocabulary_size):
 BASE_PATH = "https://storage.googleapis.com/keras-nlp/models/"
 
 checkpoints = {
+    "bert_tiny": {
+        "uncased_en": {
+            "md5": "c2b29fcbf8f814a0812e4ab89ef5c068",
+            "description": "Tiny size of BERT where all input is lowercased. "
+            "Trained on English Wikipedia + BooksCorpus.",
+            "vocabulary_size": 30522,
+        }
+    },
+    "bert_small": {
+        "uncased_en": {
+            "md5": "08632c9479b034f342ba2c2b7afba5f7",
+            "description": "Small size of BERT where all input is lowercased. "
+            "Trained on English Wikipedia + BooksCorpus.",
+            "vocabulary_size": 30522,
+        }
+    },
+    "bert_medium": {
+        "uncased_en": {
+            "md5": "bb990e1184ec6b6185450c73833cd661",
+            "description": "Medium size of BERT where all input is lowercased. "
+            "Trained on English Wikipedia + BooksCorpus.",
+            "vocabulary_size": 30522,
+        }
+    },
     "bert_base": {
         "uncased_en": {
             "md5": "9b2b2139f221988759ac9cdd17050b31",
@@ -409,6 +433,77 @@ def __init__(
 """
 
 
+def BertTiny(weights=None, vocabulary_size=None, name=None, trainable=True):
+    weights, vocabulary_size = _handle_weights_and_vocab_size(
+        "bert_tiny", weights, vocabulary_size
+    )
+
+    model = BertCustom(
+        vocabulary_size=vocabulary_size,
+        num_layers=2,
+        num_heads=2,
+        hidden_dim=128,
+        intermediate_dim=512,
+        dropout=0.1,
+        max_sequence_length=512,
+        name=name,
+        trainable=trainable,
+    )
+
+    if weights:
+        model.load_weights(weights)
+
+    # TODO(jbischof): attach the tokenizer or create separate tokenizer class.
+    # This comment applies to other variants as well.
+    return model
+
+
+def BertSmall(weights=None, vocabulary_size=None, name=None, trainable=True):
+    weights, vocabulary_size = _handle_weights_and_vocab_size(
+        "bert_small", weights, vocabulary_size
+    )
+
+    model = BertCustom(
+        vocabulary_size=vocabulary_size,
+        num_layers=4,
+        num_heads=8,
+        hidden_dim=512,
+        intermediate_dim=2048,
+        dropout=0.1,
+        max_sequence_length=512,
+        name=name,
+        trainable=trainable,
+    )
+
+    if weights:
+        model.load_weights(weights)
+
+    return model
+
+
+def BertMedium(weights=None, vocabulary_size=None, name=None, trainable=True):
+    weights, vocabulary_size = _handle_weights_and_vocab_size(
+        "bert_medium", weights, vocabulary_size
+    )
+
+    model = BertCustom(
+        vocabulary_size=vocabulary_size,
+        num_layers=8,
+        num_heads=8,
+        hidden_dim=512,
+        intermediate_dim=2048,
+        dropout=0.1,
+        max_sequence_length=512,
+        name=name,
+        trainable=trainable,
+    )
+
+    if weights:
+        model.load_weights(weights)
+
+    return model
+
+
 def BertBase(weights=None, vocabulary_size=None, name=None, trainable=True):
     weights, vocabulary_size = _handle_weights_and_vocab_size(
         "bert_base", weights, vocabulary_size
@@ -429,8 +524,6 @@ def BertBase(weights=None, vocabulary_size=None, name=None, trainable=True):
     if weights is not None:
         model.load_weights(weights)
 
-    # TODO(jbischof): attach the tokenizer or create separate tokenizer class.
-    # Applicable for other BERT variants as well.
     return model
 
 
@@ -457,6 +550,30 @@ def BertLarge(weights=None, vocabulary_size=None, name=None, trainable=True):
     return model
 
 
+setattr(
+    BertTiny,
+    "__doc__",
+    MODEL_DOCSTRING.format(
+        type="Tiny", names=", ".join(checkpoints["bert_tiny"])
+    ),
+)
+
+setattr(
+    BertSmall,
+    "__doc__",
+    MODEL_DOCSTRING.format(
+        type="Small", names=", ".join(checkpoints["bert_small"])
+    ),
+)
+
+setattr(
+    BertMedium,
+    "__doc__",
+    MODEL_DOCSTRING.format(
+        type="Medium", names=", ".join(checkpoints["bert_medium"])
+    ),
+)
+
 setattr(
     BertBase,
     "__doc__",


Test patch:
diff --git a/keras_nlp/network_tests/test_load_ckpts.py b/keras_nlp/network_tests/test_load_ckpts.py
--- a/keras_nlp/network_tests/test_load_ckpts.py
+++ b/keras_nlp/network_tests/test_load_ckpts.py
@@ -23,6 +23,9 @@
 @pytest.mark.slow
 class BertCkptTest(tf.test.TestCase, parameterized.TestCase):
     @parameterized.named_parameters(
+        ("tiny_uncased_en", keras_nlp.models.BertTiny, "uncased_en"),
+        ("small_uncased_en", keras_nlp.models.BertSmall, "uncased_en"),
+        ("medium_uncased_en", keras_nlp.models.BertMedium, "uncased_en"),
         ("base_uncased_en", keras_nlp.models.BertBase, "uncased_en"),
         ("base_cased_en", keras_nlp.models.BertBase, "cased_en"),
         ("base_zh", keras_nlp.models.BertBase, "zh"),