Move Bert to models folder

examples/bert/bert_config.py

@@ -12,60 +12,49 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# TODO(jbischof): remove in favor of BertBase, BertSmall, etc
 MODEL_CONFIGS = {
     "tiny": {
         "num_layers": 2,
-        "hidden_size": 128,
+        "hidden_dim": 128,
         "dropout": 0.1,
-        "num_attention_heads": 2,
-        "inner_size": 512,
-        "inner_activation": "gelu",
-        "initializer_range": 0.02,
+        "num_heads": 2,
+        "intermediate_dim": 512,
     },
     "mini": {
         "num_layers": 4,
-        "hidden_size": 256,
+        "hidden_dim": 256,
         "dropout": 0.1,
-        "num_attention_heads": 4,
-        "inner_size": 1024,
-        "inner_activation": "gelu",
-        "initializer_range": 0.02,
+        "num_heads": 4,
+        "intermediate_dim": 1024,
     },
     "small": {
         "num_layers": 4,
-        "hidden_size": 512,
+        "hidden_dim": 512,
         "dropout": 0.1,
-        "num_attention_heads": 8,
-        "inner_size": 2048,
-        "inner_activation": "gelu",
-        "initializer_range": 0.02,
+        "num_heads": 8,
+        "intermediate_dim": 2048,
     },
     "medium": {
         "num_layers": 8,
-        "hidden_size": 512,
+        "hidden_dim": 512,
         "dropout": 0.1,
-        "num_attention_heads": 8,
-        "inner_size": 2048,
-        "inner_activation": "gelu",
-        "initializer_range": 0.02,
+        "num_heads": 8,
+        "intermediate_dim": 2048,
     },
     "base": {
         "num_layers": 12,
-        "hidden_size": 768,
+        "hidden_dim": 768,
         "dropout": 0.1,
-        "num_attention_heads": 12,
-        "inner_size": 3072,
-        "inner_activation": "gelu",
-        "initializer_range": 0.02,
+        "num_heads": 12,
+        "intermediate_dim": 3072,
     },
     "large": {
         "num_layers": 24,
-        "hidden_size": 1024,
+        "hidden_dim": 1024,
         "dropout": 0.1,
-        "num_attention_heads": 16,
-        "inner_size": 4096,
-        "inner_activation": "gelu",
-        "initializer_range": 0.02,
+        "num_heads": 16,
+        "intermediate_dim": 4096,
     },
 }
 

examples/bert/bert_finetune_glue.py

@@ -24,17 +24,10 @@
 
 import keras_nlp
 from examples.bert.bert_config import FINETUNING_CONFIG
-from examples.bert.bert_config import MODEL_CONFIGS
 from examples.bert.bert_config import PREPROCESSING_CONFIG
 
 FLAGS = flags.FLAGS
 
-flags.DEFINE_string(
-    "model_size",
-    "tiny",
-    "One of: tiny, mini, small, medium, base, or large.",
-)
-
 flags.DEFINE_string(
     "vocab_file",
     None,
@@ -109,39 +102,14 @@ def to_tf_dataset(split):
     return train_ds, test_ds, validation_ds
 
 
-class BertClassificationFinetuner(keras.Model):
-    """Adds a classification head to a pre-trained BERT model for finetuning"""
-
-    def __init__(self, bert_model, num_classes, initializer, **kwargs):
-        super().__init__(**kwargs)
-        self.bert_model = bert_model
-        self._logit_layer = keras.layers.Dense(
-            num_classes,
-            kernel_initializer=initializer,
-            name="logits",
-        )
-
-    def call(self, inputs):
-        # Ignore the sequence output, use the pooled output.
-        _, pooled_output = self.bert_model(inputs)
-        return self._logit_layer(pooled_output)
-
-
 class BertHyperModel(keras_tuner.HyperModel):
     """Creates a hypermodel to help with the search space for finetuning."""
 
-    def __init__(self, model_config):
-        self.model_config = model_config
-
     def build(self, hp):
         model = keras.models.load_model(FLAGS.saved_model_input, compile=False)
-        model_config = self.model_config
-        finetuning_model = BertClassificationFinetuner(
-            bert_model=model,
+        finetuning_model = keras_nlp.models.BertClassifier(
+            base_model=model,
             num_classes=3 if FLAGS.task_name in ("mnli", "ax") else 2,
-            initializer=keras.initializers.TruncatedNormal(
-                stddev=model_config["initializer_range"]
-            ),
         )
         finetuning_model.compile(
             optimizer=keras.optimizers.Adam(
@@ -168,8 +136,6 @@ def main(_):
         end_value=tokenizer.token_to_id("[SEP]"),
     )
 
-    model_config = MODEL_CONFIGS[FLAGS.model_size]
-
     def preprocess_data(inputs, labels):
         inputs = [tokenizer(x) for x in inputs]
         token_ids, segment_ids = packer(inputs)
@@ -194,7 +160,7 @@ def preprocess_data(inputs, labels):
     )
 
     # Create a hypermodel object for a RandomSearch.
-    hypermodel = BertHyperModel(model_config)
+    hypermodel = BertHyperModel()
 
     # Initialize the random search over the 4 learning rate parameters, for 4
     # trials and 3 epochs for each trial.

examples/bert/bert_preprocess.py

@@ -375,6 +375,9 @@ def create_masked_lm_predictions(
 ):
     """Creates the predictions for the masked LM objective."""
 
+    # TODO(jbischof): replace with keras_nlp.layers.MLMMaskGenerator
+    # (Issue #166)
+
     cand_indexes = []
     for (i, token) in enumerate(tokens):
         if token == "[CLS]" or token == "[SEP]":