bert_base_zh, bert_base_multi_cased: Add BERT Base Variants

keras_nlp/models/bert.py

@@ -33,15 +33,26 @@ def _bert_kernel_initializer(stddev=0.02):
         "uncased_en": {
             "md5": "9b2b2139f221988759ac9cdd17050b31",
             "description": "Base size of Bert where all input is lowercased. "
-            "Trained on English wikipedia + books corpora.",
+            "Trained on English Wikipedia + BooksCorpus.",
             "vocabulary_size": 30522,
         },
         "cased_en": {
             "md5": "f94a6cb012e18f4fb8ec92abb91864e9",
             "description": "Base size of Bert where case is maintained. "
-            "Trained on English wikipedia + books corpora.",
+            "Trained on English Wikipedia + BooksCorpus.",
             "vocabulary_size": 28996,
         },
+        "zh": {
+            "md5": "79afa421e386076e62ab42dad555ab0c",
+            "description": "Base size of BERT. Trained on Chinese Wikipedia.",
+            "vocabulary_size": 21128,
+        },
+        "multi_cased": {
+            "md5": "b0631cec0a1f2513c6cfd75ba29c33aa",
+            "description": "Base size of BERT. Trained on Wikipedias of "
+            "104 languages.",
+            "vocabulary_size": 119547,
+        },
     }
 }
 

network_tests/test_load_ckpts.py

@@ -14,24 +14,20 @@
 """Tests for loading pretrained model checkpoints."""
 
 import tensorflow as tf
+from absl.testing import parameterized
 
 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 = {
-            "token_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)),
-            "padding_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")
+class BertCkptTest(tf.test.TestCase, parameterized.TestCase):
+    @parameterized.named_parameters(
+        ("uncased_en", "uncased_en"),
+        ("cased_en", "cased_en"),
+        ("zh", "zh"),
+        ("multi_cased", "multi_cased"),
+    )
+    def test_load_bert_base(self, weights):
+        model = keras_nlp.models.BertBase(weights=weights)
         input_data = {
             "token_ids": tf.random.uniform(
                 shape=(1, 512), dtype=tf.int64, maxval=model.vocabulary_size