Fine-tuning DistilBERT on senator tweets

A guide to fine-tuning DistilBERT on the tweets of American Senators with snscrape, SQLite, and Transformers (PyTorch) on Google Colab.

Introduction

Tweets are short bits of text that can (sometimes) be packed with valuable data. In the case of United States Senators, official Twitter accounts contain their opinions on a wide range of political issues and are (generally) in good grammatical form. Kaggle’s Toxic Comment Classification Challenge famously demonstrated the power of the transformer in text classification on tweets. But can a transformer perform as well on a binary task trained on a much smaller dataset of more diplomatic veiled language?

To find out, I fine-tuned the DistilBERT transformer model on a custom dataset of all 2021 tweets from US Senators. The result is a powerful text classification model that can determine a senator’s political party based on a single tweet with 90.8% accuracy.

In this article, I take you through the data procurement and fine-tuning process, demonstrating how I:

• Scrape all official Twitter accounts of US Senators with snscrape, clean them with Preprocessor and store them in a local sqlite database.
• Transform tweets into a HuggingFace Dataset and fine-tune the DistilBERT base model with PyTorch.
• Evaluate the resulting classification model and push it to the Huggingface model hub.

Part 1: Creating the dataset

I decided to create my own dataset for this project simply because I couldn’t find any previously curated Twitter datasets that included tweets from all senators in the most recent session of Congress over a long enough period of time. This code lives in its own notebook here.

Here’s the outline for creating the dataset:

1. Scrape all tweets from 2021 from individual senator Twitter accounts using snscrape.
2. Preprocess the tweets with preprocessor and pandas (using chaining methods).
3. Store the tweets in a local sqlite database that I can query later on to train the model.

Scrape (historical) tweets

Tweepy is one of the most popular and easy-to-use packages for scraping Twitter. Unfortunately, it only returns tweets posted within the last seven days. This drawback, along with the fact that we only have 100 accounts to scrape from (50 states x 2 senators each), makes it impossible to train a model that will generalize and produce accurate results (trust me, I tried it).

We need to have a well-rounded sample of tweets that captures the senators’ opinions on a broad range of social, economic, and foreign policy issues. Furthermore, to make the most up-to-date model possible, we can only use tweets from the current session of Congress, which started on January 3, 2021.

First, let’s take a look at the list of senator usernames that I created based on data gathered from the UC San Diego Library:

Next, let’s practice just scraping a single senator’s profile:

This is a total mess! We need to translate the data from the API call to something interpretable.

Pre-process tweets

I created two simple functions to package the data up into a neat and tidy data frame.

This looks much better, but we’re still missing the full names and party affiliations for all the tweets. Next, we join the extracted tweet data with senators_usernames.csv, convert the timestamps to readable formats, and put everything in a data frame. To make the code more readable and efficient, we use modern Pandas methods, as detailed here by Tom Augspurger.

Store tweets

The code above shows my workflow for scraping, cleaning, and storing the tweets for just a single scraped Twitter profile. To gather data for all 100 senators of the 117th United States Congress, I place most of the code in a loop.

First, we must connect to the sqlite engine.

Now, we loop over all of the senator usernames and append the resulting tweet data to an empty list.

After scraping the data, we loop over the list of data frames and write each one to the SQLite database.

Store tweets

I chose to store these tweets in a SQLite database for a few reasons:

1. I want to update the database with new tweets periodically, to keep the model up-to-date.
2. I’d like to store the data in a central location with options to query subsets of it for future projects.
3. I may try using other tokenizers in the future.
4. VSCode has wonderful SQLite integration and sqlalchemy makes it easy to access with pandas.

Our dataset is now complete and can be accessed here!

Part 2: Fine-tune DistilBERT

We’ve created the dataset and now it’s time to prepare the dataset for fine-tuning on our downstream task.

Transformer models are pre-trained on immense text corpora for many days and to the tune of thousands of dollars. Luckily, the process of fine-tuning allows mere mortals to harness the power of these models, while only requiring small datasets of task-specific data (i.e. our tweets dataset). (If you’re not familiar with fine-tuning I suggest you start here.)

I chose to fine-tune DistilBERT, a “smaller, faster, cheaper, lighter” version of BERT because it has a history of performing well on text classification tasks and retains 97% of the language understanding capabilities of the original BERT model. This makes it ideal to train on Google Colab, using its free tier.

Here’s the outline for fine-tuning:

1. Set up Google Colab.
2. Do some last-minute pre-processing.
3. Tokenize the dataset.
4. Specify model configuration and hyperparameters.
5. Train!
6. Evaluate model

The full fine-tuning notebook can be found here.

Set up Google Colab

First, we install all the necessary packages. Make sure to specify datasets==1.18.3 and transformers[sentencepiece]==4.16.2 if you want to reproduce the code in this notebook.

Mount your Google Drive and define a path to your Colab data to ensure your model saves before the runtime session times out.

Last-minute pre-processing

First, we manually upload TWEETS.db to Google Drive, connect to the database, and read in the tweets.

We clean up the tweets using the preprocessor package and assign both Independent senators, Bernie Sanders and Angus S. King, Jr., to the Democrat class because they both caucus with Democrats.

We want to confirm that we only have our two target classes in the dataset and create a label mapping to later convert them to integers.

The DistilBERT tokenizer will only accept the training data if it has a labels column, so we’ll add another column with our mapped labels.

Next, we create a Dataset from the data frame and split it into a test and train set.

Then, we cast the labels column as a class-encoded column, which tells the dataset where to find the labels for the data and automatically formats it for training.

Finally, the Dataset is ready for tokenizing!

Tokenize dataset

We start by loading the Transformers AutoTokenizer for DistilBERT.

Next, we make a list of columns to remove from the dataset upon tokenization and tokenize the tweets in the text column.

In the tokenize() function, we set padding=True to pad each tweet to the maximum length tweet in the batch. We also set truncation=True and max_length=512 to truncate each tweet to the maximum sequence length for DistilBERT.

After tokenization is complete, we prepare the dataset to be fed to the model by setting its format to "torch" and creating data loads to correctly reshape the dataset.

Now that the encoded dataset contains all the necessary inputs for DistilBERT, we instantiate a DataLoader for each split of the dataset to feed it to the model.

Load the model

Before training the model, we need to load it from the HuggingFace Model Hub and specify its configuration and hyperparameters.

Besides providing the model checkpoint, the only other configuration we need to specify is num_labels. I programmatically define the number of labels by setting it equal to num_classes in the Dataset.

Define hyperparameters, optimizer, and learning rate scheduler

The learning_rate and num_epochs are the only hyperparameters we will define for this task.

• Epochs are the number of training iterations over the dataset. In the past, I’ve had success with num_epochs=5.
• Learning rate specifies how much the model parameters are adjusted after each training epoch. It’s typical to start with learning_rate=5e-5 and increase it in subsequent training sessions, if necessary.

Optimization is the process in which model parameters are adjusted after each training epoch (or batch) to iteratively reduce model error. The optimizer is the specific algorithm used in the process. For our task, we follow HuggingFace’s lead and use the AdamW() optimizer, to which we pass our selected learning_rate.

Next, we move the model to the GPU, since we are using Google Colab.

Train model

Finally, we are ready to train! I copied this code, and the following evaluation loop code, from Chapter 3 of HuggingFace’s Course. I highly recommend the course for both learning and as a reference.

Save model

I usually step away from the computer while my models are training, so I always execute a code cell that saves the model locally in case the notebook times out.

Because Google Drive has limited storage space and model files tend to be quite large, pushing the model to Huggingface’s Model Hub is a great option to both store the model and share it with others. There are several ways to upload a model to the Hub, but I found the easiest and quickest way was to use huggingface_hub within the training notebook.

We start by importing the module and installing git-lfs, an open-source Git extension for versioning large files.

Next, we log in to our Huggingface account (create an account here, if you haven’t already).

The last step is to configure your Git settings and push the model.

Our model is saved and can be access here!

Evaluate model

Now that we’ve saved the model, we need to assess its performance. The 🤗 datasets library offers several NLP performance metrics that we can easily import to evaluate the model.

Specifying metric = load_metric("glue", "mrpc") will instantiate a metric object from the HuggingFace metrics repository, that will calculate the accuracy and F1 score of the model.

If you need a refresher on the confusion matrix and binary performance matrix, I highly recommend checking out this article.

Both the accuracy and F1 score are looking good! Let’s inference the model with some more recent tweets that weren’t in the training or test dataset.

Here’s a tweet from Texas Senator Ted Cruz, a well-known Republican:

After sending the tweet through the tokenizer and feeding it to the model, we are given the logits.

Because the raw output of the model, called logits, are relatively uninterpretable, we apply a softmax function, which re-scales the values so that we get probabilities between [0, 1]. Because our model is binary, we get a probability for each class. Recall that during data preprocessing, we assigned the tweets to numeric classes: {'Republican': '0', 'Democrat': '1'}. This means that the first value in the resulting tensor is the probability that the tweet belongs to the Republican class. Below, you can see that the model predicts Senator Cruz’s tweet as Republican with a probability of 99.999%!

Now let’s try a tweet from Massachusetts Senator Elizabeth Warren, a Democrat:

The model assigns Senator Warren’s tweet to the Democrat class with a probability of 100%!

Conclusion

In this article, I illustrate how I created a custom dataset of US Senator tweets and used it to fine-tune an accurate DistilBERT text classification model that predicts the political party affiliation of a senator based on a tweet they have posted.

Here’s a brief summary of all the steps involved:

1. Scrape historical tweets from Twitter using snscrape.

2. Clean and store tweets in a local SQLite database.

3. Tokenize the dataset and fine-tune a DistilBERT model using 🤗 transformers and PyTorch in Google Colab.

4. Push the model to the 🤗 Model Hub.

5. Evaluate model performance on the test set and inference the model on two unseen tweets.

Next steps

Both the senator-tweets dataset and distilbert-political-tweets model have many potential applications. A model that predicts the “partisan-ness” of senators in real-time based off of their tweets could be an interesting tool for prospective voters. A semantic search or question answering application to determine the political positions that senators take on specific issues could be similarly enlightening for voters.

Final thoughts

I completed this project without any fancy API keys, EC2 instances, or S3 buckets. In fact, I only used two Google Colab (free edition) notebooks to create the dataset and fine-tune the model. In other words, this project is accessible! I hope this article inspires you to get out there and build something you’re interested in, something original, with all of the phenomenal open-source resources available to us during this “golden age” of NLP.

Project links

• Github repository 🐍
• Final version of fine-tuned model ⚙️
• Senator tweets dataset 💾
• Streamlit app (coming soon!) 🎈

If you’d like to stay up-to-date on the latest data science trends, technologies, and packages, consider becoming a Medium member. You’ll get unlimited access to articles and blogs like Towards Data Science and you’ll be supporting my writing. (I earn a small commission for each membership).

Want to connect?

• 📖 Follow me on Medium
• 💌 Subscribe to get an email whenever I publish
• 🖌️ Check out my generative AI blog
• 🔗 Take a look at my portfolio
• 👩‍🏫 I’m also a data science coach!

I’ve also written

References

[1] Kaggle, Toxic Comment Classification Challenge (2018), Kaggle

[2] S. Bollinger, Congressional Twitter Accounts (2022), UC San Diego Library

[3] T. Augspurger, Modern Pandas (Part 1) (2022), Tom Augspurger

[4] T. Peng, The Staggering Cost of Training SOTA Models (2019), Synced

[5] S. Cellat, Fine-Tuning Transformer-Based Language Models (2021), Y Meadows

[6] V. Sahn, Smaller, faster, cheaper, lighter: Introducing DistilBERT, a distilled version of BERT (2019), Huggingface

[7] V. Sahn, L. Debut, J. Chaumond and T. Wolf, DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter (2020), Huggingface

[8] Senate Demorats, Our Caucus (2022), Senate Democrats

[9] Huggingface, A full training (2022), Huggingface

[10] Huggingface, Metrics (2022), Huggingface

[11] J. Brownlee, What is a Confusion Matrix in Machine Learning (2020), Machine Learning Mastery

[12] T. Cruz, https://twitter.com/SenTedCruz/status/1501309601605734408 (2022), Twitter

[13] E. Warren, https://twitter.com/ewarren/status/1498410222683148291, (2022), Twitter

Resources

• Huggingface Course — Write your training loop in PyTorch (Article)
• Huggingface Course — A full training (Notebook)
• Huggingface Course — Sharing pretrained models (Article)
• Huggingface — Share a model

Want to connect?

Contact me on LinkedIn ✍️.