

TITLE
Script for distilling zero-shot classifier to more efficient student

COMMENTS
2

REACTIONS
+1: 0
-1: 0
laugh: 0
hooray: 0
heart: 1
rocket: 0
eyes: 0

BODY
This PR introduces a script that provides a way to improve the speed and memory performance of a zero-shot classifier by training a more efficient student model from the zero-shot teacher's predictions over an unlabeled dataset.

For a given sequence, the zero-shot classification pipeline requires each possible label to be fed through the large NLI model separately. This requirement slows results considerably, particularly for tasks with a large number of classes `K`.

Given (1) an unlabeled corpus and (2) a set of candidate class names, this script allows a user to train a standard classification head with `K` output dimensions. The script generates a softmax distribution for the provided data & class names, and a student classifier is then fine-tuned on these proxy labels. The resulting student model can be used for classifying novel text instances over these `K` classes with an order-of-magnitude boost in inference speed in addition to decreased memory usage.

A teacher NLI model can be distilled to a student model by running `distill_classifier.py` like so:

```
python distill_classifier.py \
--data_file unlabeled_data.txt \
--class_names_file class_names.txt \
--output_dir ./distilled_model
```

A number of other args are provided as well, such as `--teacher_name_or_path` and `--student_name_or_path` for specifying the pre-trained student & teacher models to be used (by default `roberta-large-mnli` and `distillbert-base-uncased`) and `--hypothesis_template` for customizing the [hypothesis template](https://huggingface.co/transformers/main_classes/pipelines.html#transformers.ZeroShotClassificationPipeline.__call__) used by the teacher zero-shot model. The training is implemented via `Trainer`, so any [`TrainingArguments`](https://huggingface.co/transformers/main_classes/trainer.html#trainingarguments) can be specified as well.

The resulting model can then be used trivially in a text classification pipeline or in any other way:

```python
model = AutoModelForSequenceClassification.from_pretrained("./distilled_model")
tokenizer = AutoTokenizer.from_pretrained("./distilled_model")
distilled_classifier = TextClassificationPipeline(model=model, tokenizer=tokenizer)
```

See the included [README.md](https://github.com/joeddav/transformers/blob/zero-shot-distillation/examples/research_projects/zero-shot-distillation/README.md) for more details and examples.

Soon I'll introduce a similar script for self-training an NLI model, boosting the model's performance after training on only unlabeled data, which model can then be subsequently distilled with this script like any NLI model.

**Update**: I also just added a link to a working [colab notebook demo](https://colab.research.google.com/drive/1mjBjd0cR8G57ZpsnFCS3ngGyo5nCa9ya?usp=sharing).