Repository: keras-team/keras-nlp

Issue:
Add checkpoints for fine-tuned models
We have numerous checkpoints for text encoders but there's a lot of value in offering ready-to-go fine tuned models as well.

Thoughts:
* Let's start with fine-tuning on SST using a `BertBase` backbone to minimize the amount of novel code.
* Add a `weights` argument to `BertClassifier` and `classifier_checkpoints` struct to `bert_tasks.py`
* Each `classifier_checkpoints` entry will need to specify `num_classes`.
   * This is a baby step towards configs but I cannot imagine writing config-in-code class like `BertClassifierWith4Classes()`. Other tasks might need several params in the checkpoint struct.
* We need to choose if we want a default for `weights` or `backbone`. To me `backbone` is more natural than giving everyone sentiment analysis by default.
* Add example usage in the `BertClassifier` docstring.
* While there are plenty of checkpoints to convert around the web, I wonder if we should train our own. This would be a good e2e test of our current UX and code and good material for an eventual sentiment analysis demo.


Discussion:
This all sounds good to me!

SST as the first offering sounds great.

> Each classifier_checkpoints entry will need to specify num_classes.

This sounds good to me to! I think the "Keras style" worry was more specifically about config objects exposed in API symbols, not about tracking metadata for checkpoints internally.

It seems like for now at least, we have the pattern where the `weights`/`vocabulary` arguments will imply a value for other args. Backbone `weights` sets `vocabulary_size`, classifier `weights` sets `num_classes`, preprocessing `vocabulary` sets `lowercase`. I know that is probably not your ideal state, but at least we are consistent in the "mental model" we are giving to users. And can continue to discuss a way to avoid or mitigate the presence of arg-arg interactions in our signatures.

> I wonder if we should train our own.

sgtm! The most natural place for this seems like `bert/examples` for this I think? This is explicitly not library code, that will be used to generate checkpoints used in library code.

Some planning might be in order here broadly for the examples code. I.e.
 - We should probably rewrite the bert pretraining and glue fine tuning scripts so that checkpoints are the defacto currency passed between scripts.
 - We should probably make a standalone bert classification fine tuning script for uses like training our sst model, but should still figure out how that interacts with the glue fine tuning offering. Do we have two separate scripts--one that does sst inside glue and one that fine tunes for classification on a dataset separately? Or is the glue script a meta script that invoked our classifier script multiple times.
> Each classifier_checkpoints entry will need to specify num_classes.

Just to clarify - do we want to have this information in the saved file, or a config? My preference is to go with a config, e.g.:
```
"sst": {
  "num_classes": 2
  "weights_url": .....
}
```

> We need to choose if we want a default for weights or backbone. 

We may also go with no defaults? We can rely on users to choose their desired arch. 

> I wonder if we should train our own.

I vote for finetuning on our own side rather than load checkpoints. Doing finetuning lets us use our pretrained model offering as a real user, so we can find out the conflicts. 


Resolving patch:
diff --git a/keras_nlp/models/bert/bert_presets.py b/keras_nlp/models/bert/bert_presets.py
--- a/keras_nlp/models/bert/bert_presets.py
+++ b/keras_nlp/models/bert/bert_presets.py
@@ -220,3 +220,34 @@
         "vocabulary_hash": "bb6ca9b42e790e5cd986bbb16444d0e0",
     },
 }
+
+classifier_presets = {
+    "bert_tiny_uncased_en_sst2": {
+        "config": {
+            "backbone": {
+                "class_name": "keras_nlp>Bert",
+                "config": {
+                    "vocabulary_size": 30522,
+                    "hidden_dim": 128,
+                    "intermediate_dim": 512,
+                    "num_layers": 2,
+                    "num_heads": 2,
+                    "max_sequence_length": 512,
+                    "num_segments": 2,
+                    "dropout": 0.1,
+                },
+            },
+            "num_classes": 2,
+        },
+        "preprocessor_config": {
+            "lowercase": True,
+        },
+        "description": (
+            "bert_tiny_uncased_en backbone fine-tuned on the glue/sst2 dataset."
+        ),
+        "weights_url": "https://storage.googleapis.com/keras-nlp/models/bert_tiny_uncased_en_sst2/model.h5",
+        "weights_hash": "1f9c2d59f9e229e08f3fbd44239cfb0b",
+        "vocabulary_url": "https://storage.googleapis.com/keras-nlp/models/bert_tiny_uncased_en_sst2/vocab.txt",
+        "vocabulary_hash": "64800d5d8528ce344256daf115d4965e",
+    }
+}
diff --git a/keras_nlp/models/bert/bert_tasks.py b/keras_nlp/models/bert/bert_tasks.py
--- a/keras_nlp/models/bert/bert_tasks.py
+++ b/keras_nlp/models/bert/bert_tasks.py
@@ -14,15 +14,21 @@
 """BERT task specific models and heads."""
 
 import copy
+import os
 
 from tensorflow import keras
 
 from keras_nlp.models.bert.bert_models import Bert
 from keras_nlp.models.bert.bert_models import bert_kernel_initializer
 from keras_nlp.models.bert.bert_presets import backbone_presets
+from keras_nlp.models.bert.bert_presets import classifier_presets
 from keras_nlp.utils.python_utils import classproperty
 from keras_nlp.utils.python_utils import format_docstring
 
+CLASSIFIER_PRESET_NAMES = ", ".join(
+    list(backbone_presets) + list(classifier_presets)
+)
+
 
 @keras.utils.register_keras_serializable(package="keras_nlp")
 class BertClassifier(keras.Model):
@@ -110,11 +116,15 @@ def from_config(cls, config):
         return cls(**config)
 
     @classproperty
-    def presets(cls):
+    def backbone_presets(cls):
         return copy.deepcopy(backbone_presets)
 
+    @classproperty
+    def presets(cls):
+        return copy.deepcopy({**backbone_presets, **classifier_presets})
+
     @classmethod
-    @format_docstring(names=", ".join(backbone_presets))
+    @format_docstring(names=CLASSIFIER_PRESET_NAMES)
     def from_preset(
         cls,
         preset,
@@ -158,15 +168,31 @@ def from_preset(
         output = classifier(input_data)
         ```
         """
-        # Check if preset is backbone-only model
-        if preset in Bert.presets:
-            backbone = Bert.from_preset(preset, load_weights)
-            return cls(backbone, **kwargs)
-
-        # Otherwise must be one of class presets
-        # Currently no classifier-level presets, so must throw.
         if preset not in cls.presets:
             raise ValueError(
                 "`preset` must be one of "
                 f"""{", ".join(cls.presets)}. Received: {preset}."""
             )
+
+        # Check if preset is backbone-only model
+        if preset in cls.backbone_presets:
+            backbone = Bert.from_preset(preset, load_weights)
+            return cls(backbone, **kwargs)
+
+        # Otherwise must be one of class presets
+        metadata = cls.presets[preset]
+        config = metadata["config"]
+        model = cls.from_config({**config, **kwargs})
+
+        if not load_weights:
+            return model
+
+        weights = keras.utils.get_file(
+            "model.h5",
+            metadata["weights_url"],
+            cache_subdir=os.path.join("models", preset),
+            file_hash=metadata["weights_hash"],
+        )
+
+        model.load_weights(weights)
+        return model


Test patch:
diff --git a/keras_nlp/models/bert/bert_presets_test.py b/keras_nlp/models/bert/bert_presets_test.py
--- a/keras_nlp/models/bert/bert_presets_test.py
+++ b/keras_nlp/models/bert/bert_presets_test.py
@@ -184,7 +184,7 @@ def test_load_bert(self, load_weights):
     def test_load_bert_classifier(self, load_weights):
         for preset in BertClassifier.presets:
             classifier = BertClassifier.from_preset(
-                preset, num_classes=4, load_weights=load_weights
+                preset, load_weights=load_weights
             )
             input_data = {
                 "token_ids": tf.random.uniform(