How to Fine-Tune GPT-2 for Text Generation | by François St-Amant | Towards Data Science

Natural Language Generation (NLG) has made incredible strides in recent years. In early 2019, OpenAI released GPT-2, a huge pretrained model (1.5B parameters) capable of generating text of human-like quality.
Generative Pretrained Transformer 2 (GPT-2) is, like the name says, based on the Transformer. It therefore uses the attention mechanism, which means it learns to focus on previous words that are most relevant to the context in order to predict the next word (for more on this, go here).
The goal of this article is to show you how you can fine-tune GPT-2 to generate context relevant text, based on the data you provide to it.
As an example, I will be generating song lyrics. The idea is to use the already trained model, fine-tune it to our specific data and then, based on what the model observes, generate what should follow in any given song.
GPT-2 on it’s own can generate decent quality text. However, if you want it to do even better for a specific context, you need to fine-tune it on your specific data. In my case, since I want to generate song lyrics, I will be using the following Kaggle dataset, which contains a total of 12,500 popular rock songs lyrics, all in English.
Let’s begin by importing the necessary librairies and preparing the data. I recommend using Google Colab for this project, as the access to a GPU will make things much faster for you.
As can be seen in lines 26 and 34–35, I created a small test set where I removed the last 20 words of every song. This will allow me to compare the generated text with the actual one to see how well the model performs.
In order to use GPT-2 on our data, we still need to do a few things. We need to tokenize the data, which is the process of converting a sequence of characters into tokens, i.e. separating a sentence into words.
We also need to ensure that every song respects a maximum of 1024 tokens.
The SongLyrics class will do just that for us during the training, for every song in our original dataframe.
We can now import the pretrained GPT-2 model, as well as the tokenizer. Also, like I mentionned earlier, GPT-2 is HUGE. It is likely that if you try to use it on your computer, you will be getting a bunch of CUDA Out of Memory errors.
An alternative that can be used is to accumulate the gradients.
The idea is simply that before calling for optimization to perform a step of gradient descent, it will sum the gradients of several operations. Then, it will divide that total by the number of accumulated steps, in order to get an average loss over the training sample. That means much fewer calculations.
And now, finally, we can create the training function that will use all our song lyrics to fine-tune GPT-2 so that it can predict quality verses in the future.
Feel free to play around with the various hyperparameters (batch size, learning rate, epochs, optimizer).
Then, finally, we can train the model.
model = train(dataset, model, tokenizer)
Using torch.save and torch.load, you could also save your trained model for future use.
The time has come to use our brand new fine-tuned model to generate lyrics. With the use of the following two functions, we can generate lyrics for all songs in our test dataset. Remember, I had removed the last 20 words for every song. Our model will now, for a given song, look at the lyrics he has and come up with what should be the end of the songs.
The generate function prepares the generation, while text_generation actually does it, for the entire test dataframe.
In line 6 is where we specify what should be the maximum length of a generation. I left it at 30 but that is because punctuation counts, and I will later remove the last couple of words, to ensure the generation finishes at the end of a sentence.
Two other hyperparameters are worth mentioning:
Temperature (line 8). It is used to scale the probabilities of a given word being generated. Therefore, a high temperature forces the model to make more original predictions while a smaller one keeps the model from going off topic.
Top p filtering (line 7). The model will sort the word probabilities in descending order. Then, it will sum those probabilities up to p while dropping the other words. This means the model only keeps themost relevant word probabilities, but does not only keep the best one, as more than one word can be appropriate given a sequence.
In the following code, I simply clean the generated text, ensure that it finishes at the end of a sentence (not in the middle of it) and store it in a new column, in the test dataset.
There are many ways to evaluate the quality of generated text. The most popular metric is called BLEU. The algorithm outputs a score between 0 and 1, depending how similar a generated text is to reality. A score of 1 indicates that every word that was generated is present in the real text.
Here is the code to evaluate BLEU score for the generated lyrics.
We obtain an average BLEU score of 0.685, which is pretty good. In comparison, the BLEU score for the GPT-2 model without any fine-tuning was of 0.288.
However, BLEU has it’s limits. It was originally created for machine translation and only looks at the vocabulary used to determine the quality of a generated text. That is a problem for us. Indeed, it is possible to generate high quality verses that use entirely different words than reality.
That is why I will do a subjective evaluation of the performance of the model. To do that, I created a small web interface (using Dash). The code is available on my Github repository.
The way the interface works is that you provide the app with some input words. Then, the model will use that to predict what the next couple of verses should be. Here are a few example outcomes.
In red is what the GPT-2 model predicted, given the input sequence in black. You see that it has managed to generate verses that make sense and that respect the context of what came prior! Also, it generates sentences of similar length, which is extremely important when it comes to keeping the rythmn of a song. In that regard, punctuation in the input text is absolutely essential when generating lyrics.
As the article shows, by fine-tuning GPT-2 to specific data, it is possible to generate context relevant text fairly easily.
For lyrics generation, the model can generate lyrics that respect both the context and the desired length of a sentence. Of course, improvements could be made to the model. For instance we could force it to generate verses that rhyme, something often necessary when writing song lyrics.
Thanks a lot for reading, I hope I could help!
The repository with all the code and models can be found right here: https://github.com/francoisstamant/lyrics-generation-with-GPT2