Repository: keras-team/keras-nlp

Issue:
GLUE Finetuning script
There are several questions about the GLUE finetuning script:
1. GLUE finetuning is basically backbone model + classification head. Shall we make this classification head a standalone API, maybe named as `GlueClassifier`? 
2. If YES to 1, can we assume the `GlueClassifier` takes in two args about backbone model? 1) backbone model instance 2) a string identifier tells how to extract the classifier input from the backbone model's output, e.g., "pooled_output" for BERT.
3. For the script itself, do we want to have distribution strategy support, like multi-gpu or TPU? 
4. Do we want to let users be able to run the script with their model in a single command? This is tough because we need to pipe the backbone model instance to the script. 

My thoughts on these questions:
1. I would like to have `GlueClassifier` instead of each backbone model using its own classification head for GLUE.
2. yes.
3. Yes.
4. I would reply on users to copy the code and run, and we can have explicit message on which part should be customized to reflect the custom model.

Here is a draft colab: https://colab.research.google.com/gist/chenmoneygithub/9c6dfd911759ecdf8edf110054619378/glue-finetuning.ipynb


Discussion:
It appears `GlueClassifier` is essentially the same as `BertClassifier`. Why couldn't you reuse `BertClassifier`?
@jbischof This is my question 1) (should've been more clear) - do we want a reusable `GlueClassifier` or let users create their own classifier for each model? 
Forgive me for not following but if you replaced your custom classifier with a `BertClassifier` I'm pretty sure it would run OOB! Creating a glue script is helpful but I don't see why we would need extra modeling code to accomplish this.
Yeah, I think trying to write single classification head for all models will not generalize well. Your current implementation in the colab involves indexing into a dictionary output with to a "pooled" dense representation of some sort.

 - This will not work for RoBERTa, which has no pooled output by default. In fact, we don't even output a dictionary! The classifier for RoBERTa applies [two dense layers](https://github.com/keras-team/keras-nlp/blob/master/keras_nlp/models/roberta/roberta_tasks.py#L70-L80) to index 0 of the sequence embedding.
 - This would also not work if you wanted to simple average pool over an entire sequence.
 - For things like t5, an output string label is auto-regressively sampled from a language model output, and then converted back to an integer label.

We might also want to do some thinking about the different ways this could be run. The BERT paper describes a learning rate hyperparameter search, so there is no single learning rate for the real BERT GLUE score. And for large models these tasks can take many hours, so you might want to distribute runs across many workers.

I think all of this is to say maybe we should start with a script, get it working across a good cross section of models, and then take a look if there is any "library code" to pull out of it.
@jbischof It's about reusability. This GLUE finetuning script is expected to be able to easily adapt to custom models. If we don't provide a classification head, then users need to do. 

@mattdangerw Yea, relying on a string identifier is not quite stable. I am kinda unsure about how this classification part should be provided to our users. I want it to be inline with the script because it gives a clear clue what the finetuning part does, so users can adapt to their models. 

Agree we should start simple. I am thinking about using BERT, RoBERTa and GPT to create the story, and we can evaluate from there where we can improve. 
@chenmoneygithub following the design pattern of the library so far, we should expect each model to provide their own classification layer because, as @mattdangerw detailed nicely, each model may have a unique implementation. We will provide implementations for all models in our library, e.g., `BertClassifier`, `RobertaClassifier` etc. 

I think the input to this script should be a classification layer and not a raw encoder. Tell the user which args are expected and they can write a wrapper if necessary.
Just got back! Resuming this discussion:

@jbischof I am hoping users can use this script on their custom models. It looks trivial that they just need to implement a simple classification head, but without explicitly told so, it could take users (especially our basic users) a while to figure out this fact, that's why I want to have the classification head explicitly coded up in the script.

@jbischof @mattdangerw I think it is safe to assume that the users' model could always output a vector representation for the input sequence, so my proposal is to explicitly document that the `ClassificationHead` should take the vector representation as input. For the shipped code, I will just showcase with our pretrained BERT, WDYT?
To reflect our latest discussion - we will delete the classification head code from the script. 
Forgive my back and forth please, but putting in more thoughts i still find the usefulness of having a explicit ClassificationHead users can use. So this is my new plan: use a flag `use_default_classification_head`, setting which to `True` causes the eval to use the default classification head, otherwise it is using the custom head such like `BertClassifier`. I am going to make a PR and we can discuss from there. 

Resolving patch:
diff --git a/examples/glue_benchmark/glue.py b/examples/glue_benchmark/glue.py
new file mode 100644
--- /dev/null
+++ b/examples/glue_benchmark/glue.py
@@ -0,0 +1,236 @@
+# 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.
+import csv
+import os
+
+import numpy as np
+import tensorflow as tf
+import tensorflow_datasets as tfds
+from absl import app
+from absl import flags
+from tensorflow import keras
+
+import keras_nlp
+
+FLAGS = flags.FLAGS
+
+flags.DEFINE_string(
+    "task_name",
+    "mrpc",
+    "The name of the GLUE task to finetune on.",
+)
+
+flags.DEFINE_integer(
+    "batch_size",
+    32,
+    "Batch size of data.",
+)
+
+flags.DEFINE_integer(
+    "epochs",
+    2,
+    "Number of epochs to run finetuning.",
+)
+
+flags.DEFINE_float(
+    "learning_rate",
+    5e-5,
+    "Learning rate",
+)
+
+flags.DEFINE_string(
+    "submission_directory",
+    None,
+    "The directory to save the glue submission file.",
+)
+
+
+flags.DEFINE_string(
+    "finetuning_model_save_path",
+    None,
+    "The path to save the finetuning model. If None, the model is not saved.",
+)
+
+
+def load_data(task_name):
+    """
+    Load GLUE dataset.
+
+    Load GLUE dataset, and convert the dictionary format to (features, label),
+    where features is a tuple of all input sentences.
+    """
+    if task_name in ("cola", "sst2"):
+        feature_names = ("sentence",)
+    elif task_name in ("mrpc", "stsb", "rte", "wnli"):
+        feature_names = ("sentence1", "sentence2")
+    elif task_name in ("mnli", "mnli_matched", "mnli_mismatched", "ax"):
+        feature_names = ("premise", "hypothesis")
+    elif task_name in "qnli":
+        feature_names = ("question", "sentence")
+    elif task_name in "qqp":
+        feature_names = ("question1", "question2")
+    else:
+        raise ValueError(f"Unknown task_name {task_name}.")
+
+    test_suffix = ""
+    if task_name in ("mnli", "mnli_matched"):
+        # For "mnli", just run default to "mnli_matched".
+        task_name = "mnli"
+        test_suffix = "_matched"
+    elif task_name in ("mnli_mismatched"):
+        task_name = "mnli"
+        test_suffix = "_mismatched"
+
+    def split_features(x):
+        # GLUE comes with dictonary data, we convert it to a uniform format
+        # (features, label), where features is a tuple consisting of all
+        # features.
+        features = tuple([x[name] for name in feature_names])
+        label = x["label"]
+        return (features, label)
+
+    if task_name == "ax":
+        # AX is trained and evaluated on MNLI, and has its own test split.
+        train_ds, validation_ds = tfds.load(
+            "glue/mnli",
+            split=["train", "validation_matched"],
+        )
+        test_ds = tfds.load(
+            "glue/ax",
+            split="test",
+        )
+    else:
+        train_ds, test_ds, validation_ds = tfds.load(
+            f"glue/{task_name}",
+            split=["train", "test" + test_suffix, "validation" + test_suffix],
+        )
+    train_ds = train_ds.map(split_features, num_parallel_calls=tf.data.AUTOTUNE)
+    test_ds = test_ds.map(split_features, num_parallel_calls=tf.data.AUTOTUNE)
+    validation_ds = validation_ds.map(
+        split_features, num_parallel_calls=tf.data.AUTOTUNE
+    )
+    return train_ds, test_ds, validation_ds
+
+
+def preprocess_data(preprocess_fn, dataset):
+    """Run `proprocess_fn` on input dataset then batch & prefetch."""
+    return (
+        dataset.map(preprocess_fn)
+        .batch(FLAGS.batch_size)
+        .prefetch(tf.data.AUTOTUNE)
+    )
+
+
+def generate_submission_files(finetuning_model, test_ds):
+    """Generate GLUE leaderboard submission files."""
+    filenames = {
+        "cola": "CoLA.tsv",
+        "sst2": "SST-2.tsv",
+        "mrpc": "MRPC.tsv",
+        "qqp": "QQP.tsv",
+        "stsb": "STS-B.tsv",
+        "mnli_matched": "MNLI-m.tsv",
+        "mnli_mismatched": "MNLI-mm.tsv",
+        "qnli": "QNLI.tsv",
+        "rte": "RTE.tsv",
+        "wnli": "WNLI.tsv",
+        "ax": "AX.tsv",
+    }
+
+    labelnames = {
+        "mnli_matched": ["entailment", "neutral", "contradiction"],
+        "mnli_mismatched": ["entailment", "neutral", "contradiction"],
+        "ax": ["entailment", "neutral", "contradiction"],
+        "qnli": ["entailment", "not_entailment"],
+        "rte": ["entailment", "not_entailment"],
+    }
+    if not os.path.exists(FLAGS.submission_directory):
+        os.makedirs(FLAGS.submission_directory)
+    filename = FLAGS.submission_directory + "/" + filenames[FLAGS.task_name]
+
+    labelname = labelnames.get(FLAGS.task_name)
+    with tf.io.gfile.GFile(filename, "w") as f:
+        # GLUE requires a format of index + tab + prediction.
+        writer = csv.writer(f, delimiter="\t")
+        # Write the required headline for GLUE.
+        writer.writerow(["index", "prediction"])
+        predictions = finetuning_model.predict(test_ds)
+        predictions = np.argmax(predictions, -1)
+        for idx, pred in enumerate(predictions):
+            if labelname:
+                pred_value = labelname[int(pred)]
+            else:
+                pred_value = pred
+            writer.writerow([idx, pred_value])
+
+
+def main(_):
+    train_ds, test_ds, val_ds = load_data(FLAGS.task_name)
+
+    # ----- Custom code block starts -----
+    bert_model = keras_nlp.models.Bert.from_preset("bert_tiny_uncased_en")
+    bert_preprocessor = keras_nlp.models.BertPreprocessor.from_preset(
+        "bert_tiny_uncased_en"
+    )
+
+    # Users should change this function to implement the preprocessing required
+    # by the model.
+    def preprocess_fn(feature, label):
+        return bert_preprocessor(feature), label
+
+    # ----- Custom code block ends -----
+
+    train_ds = preprocess_data(preprocess_fn, train_ds)
+    val_ds = preprocess_data(preprocess_fn, val_ds)
+    test_ds = preprocess_data(preprocess_fn, test_ds)
+
+    loss = keras.losses.SparseCategoricalCrossentropy(from_logits=True)
+    metrics = [keras.metrics.SparseCategoricalAccuracy()]
+    if FLAGS.task_name == "stsb":
+        num_classes = 1
+        loss = keras.losses.MeanSquaredError()
+        metrics = [keras.losses.MeanSquaredError()]
+    elif FLAGS.task_name in ("mnli", "mnli_mismatched", "mnli_matched", "ax"):
+        num_classes = 3
+    else:
+        num_classes = 2
+
+    # ----- Custom code block starts -----
+    # Users should change this `BertClassifier` to your own classifier.
+    # Commonly the classifier is simply your model + several dense layers,
+    # please refer to "Make the Finetuning Model" section in README for
+    # detailed instructions.
+    finetuning_model = keras_nlp.models.BertClassifier(
+        backbone=bert_model,
+        num_classes=num_classes,
+    )
+    # ----- Custom code block ends -----
+
+    finetuning_model.compile(
+        optimizer=tf.keras.optimizers.Adam(FLAGS.learning_rate),
+        loss=loss,
+        metrics=metrics,
+    )
+
+    finetuning_model.fit(train_ds, validation_data=val_ds, epochs=FLAGS.epochs)
+
+    if FLAGS.submission_directory:
+        generate_submission_files(finetuning_model, test_ds)
+
+    if FLAGS.finetuning_model_save_path:
+        finetuning_model.save(FLAGS.finetuning_model_save_path)
+
+
+if __name__ == "__main__":
+    app.run(main)


Test patch: