Unofficial Implementation of MiniLM-v2

This is an unofficial implementation of MiniLM-v2 using Huggingface Transformers, PyTorch and PyTorch Lightning. Currently, only BERT architectures are supported.

Data

To distill a model, you first need to pretokenize your data. This repository expects a gzipped jsonl file containing lists of tokenized sentences. See test/test_text.jsonl.gz for an example. You can use the minilm_v2/data.py script to pretokenize your data.

$ python minilm_v2/data.py -h
usage: data.py [-h] --text_files TEXT_FILES [TEXT_FILES ...] --save_path SAVE_PATH --tokenizer_name TOKENIZER_NAME
               [--batch_size BATCH_SIZE] [--num_lines]

Tokenizes text files and writes the token ids to a pickle file. Each line in a text file is considered as a new sample.
Optionally splits the input by sentences.

options:
  -h, --help            show this help message and exit
  --text_files TEXT_FILES [TEXT_FILES ...]
                        Text file paths to tokenize
  --save_path SAVE_PATH
                        Path to save the pickled token ids to
  --tokenizer_name TOKENIZER_NAME
                        Name of or path to local huggingface tokenizer
  --batch_size BATCH_SIZE
                        Number of lines to tokenize at once.
  --num_lines           Count number of lines before tokenizing. Useful for estimating ETA.

Distilling a Model

Next, you need to write a huggingface config.json file specifying the model architecture you want to distill to. See configs/lh-h384/config.json for an example. Then, modifiy the fit.yaml configuration file. It contains all model, data, training and logging parameters. Note that the student path should be the path of the directory containing the config.json file. See the Lightning CLI for details more information on all options. To start training, run python minilm_v2/fit.py fit.yaml.

NOTE: At the end of training, the model is saved in the directory of the architecture config.json file.

References

https://github.com/microsoft/unilm/tree/master/minilm https://arxiv.org/abs/2012.15828