Fine-Tune Smaller Transformer Models: Text Classification | by Ida Silfverskiöld | May, 2024

This script creates clickbait titles and stores it in a new CSV in your root folder. You need to change the prompt template later to produce an equal amount of titles that are factual.

As I’m using a generative text model, Phi-3, some outputs won’t be usable, but that is to be expected. It will take some time for this to run, so go do something else with your time.

Once you’re finished you can store your finished CSV file with the clickbait and factual tiles in your Google Drive. Remember to set the text and label as fields, where the text is the title and the label is whether it is clickbait or factual.

Since I’ve already prepared the dataset we’ll use, please see this script to upload your custom dataset to HuggingFace.

Looking through the dataset, you’ll see that most clickbait articles that have been generated by Phi-3 has an exclamation mark at the end of it. This is something you want to make sure doesn’t happen, so it’s important to check the work of the LLM generating the data.

Remember that the script I provided you with splits your data into a training, test and validation set. I would recommend to have at least a training and test set for training the model.

If you’ve got your dataset sorted, we can go ahead and fine-tune the model.

Dataset & Model

If you haven’t opened up the cook book, do so here. The first part of this is deciding on your dataset and then your pre-trained model.

from datasets import load_dataset, DatasetDict
dataset = load_dataset("ilsilfverskiold/clickbait_titles_synthetic_data")
dataset

model_name = "albert/albert-base-v2"
your_path = "classify-clickbait"

I wen’t through the different models under the introduction section, where ALBERT and DistillBERT are smaller models and BERT and RoBERTa are larger.

For this case, as it’s not overly complex, I will go for ALBERT. I’m sure BERT can do better, but ALBERT is ten times smaller. RoBERTa is too big and may produce some overfitting with this dataset.

Remember, if you’re working with a different language then look for a base model that has been trained on a corpus from at least a similar language.

If you’re working with nordic languages I can recommend KB/bert-base-swedish-cased that I used to create a model for the IPTC newscodes categories.

Prepare The Dataset

Now we need to do a few things for this to work well.

We first convert our labels to a standardized numerical format that the trainer will understand.

from sklearn.preprocessing import LabelEncoder
label_encoder = LabelEncoder()
label_encoder.fit(dataset['train']['label'])
def encode_labels(example):
return {'encoded_label': label_encoder.transform([example['label']])[0]}
for split in dataset:
dataset[split] = dataset[split].map(encode_labels, batched=False)

Then we need to map the numerical representations back to the actual label names. This is so we can get the actual label names rather than the numerical reps when we do inference with the model.

from transformers import AutoConfig
unique_labels = sorted(list(set(dataset['train']['label'])))
id2label = {i: label for i, label in enumerate(unique_labels)}
label2id = {label: i for i, label in enumerate(unique_labels)}
config = AutoConfig.from_pretrained(model_name)
config.id2label = id2label
config.label2id = label2id
# Verify the correct labels
print("ID to Label Mapping:", config.id2label)
print("Label to ID Mapping:", config.label2id)

After this we’re ready to fetch the pre-trained model and it’s tokenizer. We use the config we set up with the labels when we import the model.

from transformers import AlbertForSequenceClassification, AlbertTokenizer
tokenizer = AlbertTokenizer.from_pretrained(model_name)
model = AlbertForSequenceClassification.from_pretrained(model_name, config=config)

If you’re using a different model such as BERT or RoBERTa, you can use AutoTokenizer and AutoModelForSequenceClassification which will automatically select the correct classes for your specified model.

This next function filters for invalid content and then makes sure the text data is properly tokenized and labeled, preparing the dataset for training.

def filter_invalid_content(example):
return isinstance(example['text'], str)
dataset = dataset.filter(filter_invalid_content, batched=False)
def encode_data(batch):
tokenized_inputs = tokenizer(batch["text"], padding=True, truncation=True, max_length=256)
tokenized_inputs["labels"] = batch["encoded_label"]
return tokenized_inputs
dataset_encoded = dataset.map(encode_data, batched=True)
dataset_encoded

dataset_encoded.set_format(type='torch', columns=['input_ids', 'attention_mask', 'labels'])

We also need to fetch a data collator to handle padding for our inputs.

from transformers import DataCollatorWithPadding
data_collator = DataCollatorWithPadding(tokenizer)

Evaluation Metrics

It’s not required for you to set up any evaluation metrics, such as accuracy, precision, recall or f1. However, you do need at least accuracy to understand how the model is performing.

Accuracy measures the amount of predictions the model got right across all categories. Precision measures how often predictions for a specific category are correct. Recall tells us how well the model can identify all instances within a specific category. The F1 Score is the weighted average of Precision and Recall.

I won’t go into detail on these metrics, but there are many others that write about this. For this case, I’m more interested in how it performs on new real data rather than synthetic data. So, what I look out for are metrics that are too good, indicating that it has overfitted.

We do though set up a function that let us look at the accuracy for each label rather than as an average. This is much more relevant when you have many labels, rather than just two.

from sklearn.preprocessing import LabelEncoder
from sklearn.metrics import accuracy_score, confusion_matrix
import numpy as np
label_encoder = LabelEncoder()
label_encoder.fit(unique_labels)
def per_label_accuracy(y_true, y_pred, labels):
cm = confusion_matrix(y_true, y_pred, labels=labels)
correct_predictions = cm.diagonal()
label_totals = cm.sum(axis=1)
per_label_acc = np.divide(correct_predictions, label_totals, out=np.zeros_like(correct_predictions, dtype=float), where=label_totals != 0)
return dict(zip(labels, per_label_acc))

We also set up the general compute metrics function. I am using all of these metrics here because this is general template I have for any text classifier, but you may decide which ones you want.

from sklearn.metrics import accuracy_score, recall_score, precision_score, f1_score
def compute_metrics(pred):
labels = pred.label_ids
preds = pred.predictions.argmax(-1)
decoded_labels = label_encoder.inverse_transform(labels)
decoded_preds = label_encoder.inverse_transform(preds)
precision = precision_score(decoded_labels, decoded_preds, average='weighted')
recall = recall_score(decoded_labels, decoded_preds, average='weighted')
f1 = f1_score(decoded_labels, decoded_preds, average='weighted')
acc = accuracy_score(decoded_labels, decoded_preds)
labels_list = list(label_encoder.classes_)
per_label_acc = per_label_accuracy(decoded_labels, decoded_preds, labels_list)
per_label_acc_metrics = {}
for label, accuracy in per_label_acc.items():
label_key = f"accuracy_label_{label}"
per_label_acc_metrics[label_key] = accuracy
return {
'accuracy': acc,
'f1': f1,
'precision': precision,
'recall': recall,
**per_label_acc_metrics
}

Once you’re decently satisfied, we can move on to setting up the training arguments and the trainer.

Training the Model

Next up we set up our training arguments. Here you can tweak the epochs, batch size and learning rate.

from transformers import Trainer, TrainingArguments
training_args = TrainingArguments(
output_dir=your_path,
num_train_epochs=3,
warmup_steps=500,
per_device_train_batch_size=16,
per_device_eval_batch_size=16,
weight_decay=0.01,
logging_steps=10,
evaluation_strategy='steps',
eval_steps=100,
learning_rate=2e-5,
save_steps=1000,
gradient_accumulation_steps=2
)

I chose to go with a learning rate and epochs based on the paper “How to Fine-Tune BERT for Text Classification?” but decreased the batch size.

Now we can go ahead and set up the trainer, with everything we’ve prepared, and run it.

trainer = Trainer(
model=model,
args=training_args,
train_dataset=dataset_encoded['train'],
eval_dataset=dataset_encoded['test'],
compute_metrics=compute_metrics,
tokenizer=tokenizer,
data_collator=data_collator,
)
trainer.train()

When in training, you need to look out for overfitting. As both the training and evaluation datasets are synthetic, the typical signs of overfitting might be unclear.

Keep an eye on the accuracy and loss for both the training and evaluation datasets. I.e. very low training and validation loss, along with too stellar evaluation metrics could be a sign of overfitting.

But remember binary classes with less complex tasks usually perform well.

You’ll see my results for one run I made below.

As you can see from the training metrics, they are a bit too good. Validation loss is fluctuating as well. This could be a very bad sign so you have to make sure to test the model on real data once it has finished training.

If you’re training a model with several classes, perhaps even with a skewed dataset, don’t worry if the average evaluation metrics aren’t great. Look at the metrics per label.

Evaluating the Model

Once it’s done training, you can run the final evaluation metrics, save the model and then save the state. This will build the metrics for when you push it to the hub for your model page.

trainer.evaluate()
trainer.save_model(your_path)
trainer.save_state()

Now you can run the HuggingFace pipeline in your notebook to test it.

from transformers import pipeline
pipe = pipeline('text-classification', model=your_path)

example_titles = [
"grab an example title",
"grab another example title",
"and another xample title"
]
for title in example_titles:
result = pipe(title)
print(f"Title: {title}")
print(f"Output: {result[0]['label']}")

Mine did fine on test data, however it missed a few clickbait articles that I personally found to be clickbait. For a production use case, it’s better to build a more diverse dataset (especially with synthetic data) so it can perform well on new real data.

Nevertheless, if you’re not satisfied, then you go back to the dataset, redo it or try with a different model.

If you are wondering, I have indeed gotten stellar results on some runs and less-than-stellar results on other runs with the same data, the same training parameters, and the same seed.

Testing the Model

Before you push the model, you can also test the model against other alternatives.

I asked GPT-3.5 to tell me which titles it thought was clickbait and factual, and it did really well which is to be expected, it is more than 1000x larger than Albert.

We can also compare some titles to what a fine-tuned FastText says versus the fine-tuned transformer encoder model.

Using FastText is very simple and computationally efficient, but it treats words in isolation and lacks deep contextual understanding.

Therefore, FastText doesn’t capture the context and nuances of language as well as a model that is transformer based.

Push to the Hub

If you’re satisfied with your model, you can push it to the HuggingFace hub to store it there.

You simply login with a write token you can find in your HuggingFace account under Settings.

!huggingface-cli login

And then push it.

tokenizer.push_to_hub("username/classify-clickbait")
trainer.push_to_hub("username/classify-clickbait")

Push the tokenizer just in case, especially if you are working with a version of Albert.

Now you can use it directly from there, mine you’ll find there.

Optimization Techniques

If you want to use a larger model like BERT, you can apply different techniques so you can distill it further after fine-tuning. I didn’t find it that much more successful than just using ALBERT, at least for this case.

BERT on its own though performed much better in general. Although I really like RoBERTa for most cases, it was prone to overfit on this specific dataset either because it was too small, not good enough or too artificial.

For every case you’ll have to estimate how much performance you can sacrifice for efficiency and eventually you learn which models do well in what situation.

Ending Notes

Would the model have performed better if we had used real data? It’s possible, but the accuracy may be lower unless the dataset is meticulously sorted.

This is hard work.

Using synthetic data can get the job done very quickly so you get up a prototype to work with. Synthetic data is much cleaner to work with.

You are also free to work with the larger open source LLMs, so it doesn’t break any rules for people that can’t access high quality data without breaching protocol.

I did not put down time and effort into building this dataset, but in all cases you should make sure you have varied data the model can learn from.