Tune a text classification model with Hugging Face Transformers

This notebook trains a SMS spam classifier with "distillibert-base-uncased" as the base model on a single GPU machine using the 🤗 Transformers library.

It first downloads a small dataset, copies it to DBFS, then converts it to a Spark DataFrame. Preprocessing up to tokenization is done on Spark. While DBFS is used as a convenience to access the datasets directly as local files on the driver, you can modify it to avoid use of DBFS.

Text tokenization of the SMS messages is done in transformers in the model's default tokenizer in order to have consistency in tokenization with the base model. The notebook uses the Trainer utility in the transformers library to fine-tune the model. The notebook wraps the tokenizer and trained model in a Transformers pipeline and logs the pipeline as an MLflow model. This make it easy to directly apply the pipeline as a UDF on Spark DataFrame string columns.

Cluster setup

For this notebook, Databricks recommends a single GPU cluster, such as a g4dn.xlarge on AWS or Standard_NC4as_T4_v3 on Azure. You can create a single machine cluster using the personal compute policy or by choosing "Single Node" when creating a cluster. This notebook works with Databricks Runtime ML GPU version 11.1 or greater. Databricks Runtime ML GPU versions 9.1 through 10.4 can be used by replacing the following command with %pip install --upgrade transformers datasets evaluate.

The transformers library is installed by default on Databricks Runtime ML. This notebook also requires 🤗 Datasets and 🤗 Evaluate, which you can install using %pip.

Set up any parameters for the notebook.

• The base model DistilBERT base model (uncased) is a great foundational model that is smaller and faster than BERT base model (uncased), but still provides similar behavior. This notebook fine tunes this base model.
• The tutorial_path sets the path in DBFS that the notebook uses to write the sample dataset. It is deleted by the last command in this notebook.

Data download and loading

Start by downloading the dataset and load it into a Spark DataFrame. The SMS Spam Collection Dataset is available from the UCI Machine Learning Repository.

Unzip the downloaded archive.

Copy the file to DBFS.

Load the dataset into a DataFrame. The file is tab separated and does not contain a header, so we specify the separator using sep and specify the column names explicitly.

Data preparation

The datasets passed into the transformers trainer for text classification need to have integer labels [0, 1].

Collect the labels and generate a mapping from labels to IDs and vice versa. transformers models need these mappings to correctly translate the integer values into the human readable labels.

Replace the string labels with the IDs in the DataFrame.

Hugging Face datasets supports loading from Spark DataFrames using datasets.Dataset.from_spark. See the Hugging Face documentation to learn more about the from_spark() method.

Dataset.from_spark caches the dataset. In this example, the model is trained on the driver, and the cached data is parallelized using Spark, so cache_dir must be accessible to the driver and to all the workers. You can use the Databricks File System (DBFS) root(AWS| Azure |GCP) or mount point (AWS | Azure | GCP).

By using DBFS, you can reference "local" paths when creating the transformers compatible datasets used for model training.

Tokenize and shuffle the datasets for training. Since the Trainer does not need the untokenized text columns for training, the notebook removes them from the dataset. This isn't necessary, but not removing the column results in a warning during training. In this step, datasets also caches the transformed datasets on local disk for fast subsequent loading during model training.

Model training

For model training, this notebook largely uses default behavior. However, you can use the full range of metrics and parameters available to the Trainer to adjust your model training behavior.

Create the evaluation metric to log. Loss is also logged, but adding other metrics such as accuracy can make modeling performance easier to understand.

Construct default training arguments. This is where you would set many of your training parameters, such as the learning rate. Refer to transformers documentation for the full range of arguments you can set.

Create the model to train from the base model, specifying the label mappings and the number of classes.

Using a data collator batches input in training and evaluation datasets. Using the DataCollatorWithPadding with defaults gives good baseline performance for text classification.

Construct the trainer object with the model, arguments, datasets, collator, and metrics created above.

Construct the MLflow wrapper class to store the model as a pipeline. When loading the pipeline, this model uses the GPU if CUDA is available. This model hardwires the batchsize to use with the transformers pipeline. You'll want to set this with the hardware you will use for inference in mind.

Train the model, logging metrics and results to MLflow. This task is very easy for BERT-based models. Don't be surprised is the evaluation accuracy is 1 or close to 1.

Batch inference

Load the model as a UDF using MLflow and use it for batch scoring.

Cleanup

Remove the files placed in DBFS.