Fine tuning BERT for text classification tasks

With simpletransformers

We first use the simpletransformers package which is more beginner-friendly. The basic steps for finetuning a classifier using simpletrasnformers are:

• Initialize a model based on a specific architechture (BERT, DistilBERT, etc)
• Train the model with train_model()
• Evaluate the model with eval_model()
• Make predictions on (unlabelled) data with predict()

from simpletransformers.classification import ClassificationModel, ClassificationArgs
import logging

logging.basicConfig(level=logging.INFO)
transformers_logger = logging.getLogger("transformers")
transformers_logger.setLevel(logging.WARNING)

We need to preprocess the data first before we start the finetuning process. In this step, we split the dataset into train and test sets to have a fully held-out test set that can be used to evaluate our classifier.

We can also create a validation that is used during the fine tuning process for hyperparameter tuning, but that is not mandatory.

from sklearn.model_selection import train_test_split

train_df, test_df = train_test_split(sexism_data, stratify=sexism_data['sexist'], test_size=0.2)

We now convert the dataframes into a format that can be read by simpletransformers. This is a dataframe with the columns ‘text’ and ‘labels’. The ‘labels’ column should be numerical, so we use one-hot encoding to transform our boolean sexist labels to numerical ones.

from sklearn.preprocessing import LabelEncoder
le = LabelEncoder()
le.fit(train_df['sexist'])
train_df['labels'] = le.transform(train_df['sexist'])
test_df['labels'] = le.transform(test_df['sexist'])

# to see which number was mapped to which class:
list(le.inverse_transform([0,1]))

[np.False_, np.True_]

So, 0 is non-sexist and 1 is sexist. We now have the appropriate data structure.

The next step is setting the training parameters and loading the classification model, in this case, DistilBERT (Sanh et al., 2019), a lightweight model that can be trained relatively quickly compared to other transformer variants like BERT and RoBERTa.

For training parameters, we have many to choose from such as the learning rate, whether we want to stop early or not, where we should save the model, and more. You can find all of them here.

As a minimal setup, we will just set the number of epochs, i.e., the number of passes the model does over the full training set. For recent transformer models, epochs are usually set to 2 or 3, after which overfitting may happen.

use_cuda is a parameter that signals whether the GPU should be used or not. It will be set based on our check earlier.

# Optional model configuration
model_args = ClassificationArgs(num_train_epochs=3, overwrite_output_dir=True)

# Create a ClassificationModel
model = ClassificationModel(
    "distilbert", "distilbert-base-uncased", args=model_args, use_cuda=gpu_avail,
)

# we set some additional parameters when using a GPU
if gpu_avail:
    model_args.use_multiprocessing=False
    model_args.use_multiprocessing_for_evaluation=False

We are now finally ready to begin training! This might take a while, especially when we’re not using a GPU.

# Train the model
model.train_model(train_df)

Epoch:   0%|          | 0/3 [00:00<?, ?it/s]

Running Epoch 1 of 3:   0%|          | 0/20 [00:00<?, ?it/s]

Running Epoch 2 of 3:   0%|          | 0/20 [00:00<?, ?it/s]

Running Epoch 3 of 3:   0%|          | 0/20 [00:00<?, ?it/s]

(60, 0.5564038594563802)

After training our model, we can use it to make predictions for unlabeled datapoints to classify whether they are sexist or not.

sexist_tweet = "A woman will never be truly fulfilled in life if she doesn’t have a committed long-term relationship with a man"
predictions, raw_outputs = model.predict([sexist_tweet])
le.inverse_transform(predictions)

  0%|          | 0/1 [00:00<?, ?it/s]

array([ True])

nonsexist_tweet = "International Women's Day (IWD) is a holiday celebrated annually on March 8 as a focal point in the women's rights movement."
predictions, raw_outputs = model.predict([nonsexist_tweet])
le.inverse_transform(predictions)

  0%|          | 0/1 [00:00<?, ?it/s]

array([ True])

We can also use the held-out test set to quantitatively evaluate our model.

# Evaluate the model
result, model_outputs, wrong_predictions = model.eval_model(test_df)
result

Running Evaluation:   0%|          | 0/1 [00:00<?, ?it/s]

{'mcc': np.float64(0.6713171133426189),
 'accuracy': 0.825,
 'f1_score': 0.8444444444444444,
 'tp': np.int64(19),
 'tn': np.int64(14),
 'fp': np.int64(6),
 'fn': np.int64(1),
 'auroc': np.float64(0.9299999999999999),
 'auprc': np.float64(0.9545054047259929),
 'eval_loss': 0.4366455078125}

# you can also use sklearn's neat classification report to get more metrics
from sklearn.metrics import classification_report

preds, _ = model.predict(list(test_df['text'].values))
# preds = le.inverse_transform(preds)

print(classification_report(test_df['labels'], preds))

  0%|          | 0/1 [00:00<?, ?it/s]

              precision    recall  f1-score   support

           0       0.93      0.70      0.80        20
           1       0.76      0.95      0.84        20

    accuracy                           0.82        40
   macro avg       0.85      0.82      0.82        40
weighted avg       0.85      0.82      0.82        40

Optional: HuggingFace transformers

We now repeat the same process with the HuggingFace transformers Python library. Additionally, we also use the accelerate library, which helps make our code more efficient. We will again use DistilBERT.

from transformers import DistilBertTokenizerFast, DistilBertForSequenceClassification
from transformers import Trainer, TrainingArguments

We will set some of the configurations, including whether to use a GPU or not.

model_name = 'distilbert-base-uncased'
if gpu_avail:
    device_name = 'cuda'
else:
    device_name = 'cpu'

# This is the maximum number of tokens in any document; the rest will be truncated.
max_length = 512

# This is the name of the directory where we'll save our model. You can name it whatever you want.
cached_model_directory_name = 'output_hf'

We will reuse the train-test splits we created for simpletransformers, but change the data structure slightly.

train_texts = train_df['text'].values
train_labels = train_df['labels'].values

test_texts = test_df['text'].values
test_labels = test_df['labels'].values

Compared to simpletransformers, we get a closer look at what happens ‘under the hood’ with huggingface. We will see the transformation of the text better — each tweet will be truncated if they’re more than 512 tokens or padded if they’re fewer than 512 tokens.

The tokens will be separated into “word pieces” using the transformers tokenizers (‘DistilBertTokenizerFast’ in this case to match the DistiBERT model). And some special tokens will also be added such as CLS (start token of every tweet) and SEP (separator between each sentence {not tweet}):

tokenizer = DistilBertTokenizerFast.from_pretrained(model_name)

We now encode our texts using the tokenizer.

from datasets import Dataset

train_df = Dataset.from_pandas(train_df)
test_df = Dataset.from_pandas(test_df)

def tokenize_function(examples):
    return tokenizer(examples["text"], padding="max_length", truncation=True)


tokenized_train_df = train_df.map(tokenize_function, batched=True)
tokenized_test_df = test_df.map(tokenize_function, batched=True)

Map:   0%|          | 0/160 [00:00<?, ? examples/s]

Map:   0%|          | 0/40 [00:00<?, ? examples/s]

We now load the DistilBERT model and specify that it should use the GPU.

model = DistilBertForSequenceClassification.from_pretrained(model_name, num_labels=len(le.classes_)).to()

As we did with simpletransformers, we now set the training parameters, i.e., the number of epochs.

import accelerate

training_args = TrainingArguments(
    num_train_epochs=3,              # total number of training epochs
    output_dir='./results',          # output directory
    report_to='none'
)

from sklearn.metrics import accuracy_score
def compute_metrics(pred):
    labels = pred.label_ids
    preds = pred.predictions.argmax(-1)
    acc = accuracy_score(labels, preds)
    return {
        'accuracy': acc,
  }

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_train_df,
    compute_metrics=compute_metrics
)

trainer.train()

<div>
  &#10;      <progress value='60' max='60' style='width:300px; height:20px; vertical-align: middle;'></progress>
  [60/60 02:58, Epoch 3/3]
</div>
&#10;

TrainOutput(global_step=60, training_loss=0.39763174057006834, metrics={'train_runtime': 181.7385, 'train_samples_per_second': 2.641, 'train_steps_per_second': 0.33, 'total_flos': 63584351354880.0, 'train_loss': 0.39763174057006834, 'epoch': 3.0})

trainer.save_model(cached_model_directory_name)

# trainer = DistilBertForSequenceClassification.from_pretrained(cached_model_directory_name)

predicted_results = trainer.predict(tokenized_test_df)

predicted_labels = predicted_results.predictions.argmax(-1) # Get the highest probability prediction
predicted_labels = predicted_labels.flatten().tolist()      # Flatten the predictions into a 1D list
predicted_labels[0:5]

[1, 1, 1, 1, 1]

print(classification_report(tokenized_test_df['labels'],
                            predicted_labels))

              precision    recall  f1-score   support

           0       0.94      0.75      0.83        20
           1       0.79      0.95      0.86        20

    accuracy                           0.85        40
   macro avg       0.86      0.85      0.85        40
weighted avg       0.86      0.85      0.85        40