Tutorial: Fine-Tuning a Language Model on Text Files with IMDB

Introduction

In this tutorial we will be showing an end-to-end example of fine-tuning a Transformer language model on a custom dataset in text files format.

By the end of this you should be able to:

1. Build a dataset with the LanguageModelDatasets class, and their DataLoaders
2. Build a LanguageModelTuner quickly, find a good learning rate, and train with the One-Cycle Policy
3. Save that model away, to be used with deployment or other HuggingFace libraries
4. Apply inference using both the Tuner available function as well as with the EasyTextGenerator class within AdaptNLP

Installing the Library

This tutorial utilizies the latest AdaptNLP version, as well as parts of the fastai library. Please run the below code to install them:

!pip install adaptnlp -U

(or pip3)

Getting the Dataset

First we need a dataset. We will use the fastai library to download the full IMDB Movie Reviews dataset

URLs holds a namespace of many data endpoints, and untar_data is a function that can download and extract any data from a given URL.

Combining both, we can download the data:

If we look at what was downloaded, we will find a train and test folder:

In each are folders seperating each text file by class:

As a result, we can say the dataset follows the following format:

• train
  • class_a
    • text1.txt
    • text2.txt
    • ...
  • class_b
    • text1.txt
    • ...
• test (or valid)
  • class_a
    • text1.txt
    • ...
  • class_b
    • text1.txt
    • ...

Now that we have the dataset, and we know the format it is in, let's pick a viable model to train with

Picking a Model with the Hub

AdaptNLP has a HFModelHub class that allows you to communicate with the HuggingFace Hub and pick a model from it, as well as a namespace HF_TASKS class with a list of valid tasks we can search by.

Let's try and find one suitable for sequence classification.

First we need to import the class and generate an instance of it:

Next we can search for a model:

Let's look at a few:

These are models specifically tagged with the text-generation tag, so you may not see a few models you would expect such as bert_base_cased.

We'll use that first model, distilgpt2:

Now that we have picked a model, let's use the data API to prepare our data

Building TaskDatasets with LanguageModelDatasets

Each task has a high-level data wrapper around the TaskDatasets class. In our case this is the LanguageModelDatasets class:

There are multiple different constructors for the LanguageModelDatasets class, and you should never call the main constructor directly.

We will be using from_folders method:

Anything you would normally pass to the tokenizer call (such as max_length, padding) should go in tokenize_kwargs, and anything going to the AutoTokenizer.from_pretrained constructor should be passed to the auto_kwargs.

In our case we have a train_path and valid_path

Also, we will set a block_size of 128, and it is not a masked language model:

And finally turn it into some AdaptiveDataLoaders.

These are just fastai's DataLoaders class, but it overrides a few functions to have it work nicely with HuggingFace's Dataset class

Finally, let's view a batch of data with the show_batch function:

When training a language model, the input and output are made to be the exact same, so there isn't a shown noticable difference here.

Building Tuner

Next we need to build a compatible Tuner for our problem. These tuners contain good defaults for our problem space, including loss functions and metrics.

First let's import the LanguageModelTuner and view it's documentation

Next we'll pass in our DataLoaders, the name of our model, and the tokenizer:

By default we can see that it used CrossEntropyLoss as our loss function, and Perplexity as our metric

Finally we just need to train our model!

Fine-Tuning

And all that's left is to tune. There are only 4 or 5 functions you can call on our tuner currently, and this is by design to make it simplistic. In case you don't want to be boxed in however, if you pass in expose_fastai_api=True to our earlier call, it will expose the entirety of Learner to you, so you can call fit_one_cycle, lr_find, and everything else as Tuner uses fastai under the hood.

First, let's call lr_find, which uses fastai's Learning Rate Finder to help us pick a learning rate.

It recommends a learning rate of around 5e-5, so we will use that.

Let's look at the documentation for tune:

We can pass in a number of epochs, a learning rate, a strategy, and additional fastai callbacks to call.

Valid strategies live in the Strategy namespace class, and consist of:

• OneCycle (Also called the One-Cycle Policy)
• CosineAnnealing
• SGDR

In this tutorial we will train with the One-Cycle policy, as currently it is one of the best schedulers to use.

Saving Model

Now that we have a trained model, let's save those weights away.

Calling tuner.save will save both the model and the tokenizer in the same format as how HuggingFace does:

Performing Inference

There are two ways to get predictions, the first is with the .predict method in our tuner. This is great for if you just finished training and want to see how your model performs on some new data! The other method is with AdaptNLP's inference API, which we will show afterwards

In Tuner

First let's write a sentence to test with

And then predict with it:

With the Inference API

Next we will use the EasyTextGenerator class, which AdaptNLP offers:

We simply construct the class:

And call the tag_text method, passing in the sentence, the location of our saved model, and some names for our classes:

And we got the exact same output!