Movie Name Generation Using GPT-2

Since its reveal in 2017 in the popular paper Attention Is All You Need (https://arxiv.org/abs/1706.03762), the Transformer quickly became the most popular model in NLP. The ability to process text in a non-sequential way (as opposed to RNNs) allowed for training of big models. The attention mechanism it introduced proved extremely useful in generalizing text.

Following the paper, several popular transformers surfaced, the most popular of which is GPT. GPT models are developed and trained by OpenAI, one of the leaders in AI research. The latest release of GPT is GPT-3, which has 175 billion parameters. The model was very advanced to the point where OpenAI chose not to open-source it. People can access it through an API after a signup process and a long queue.

However, GPT-2, their previous release is open-source and available on many deep learning frameworks.

In this excercise, we use Huggingface and PyTorch to fine-tune a GPT-2 model for movie name generation.

Overview:

Imports and Data Loading
Data Preprocessing
Setup and Training
Movie Name Generation
Model Saving and Loading

Imports and Data Loading

Please use pip install {library name} in order to install the libraries below if they are not installed. "transformers" is the Huggingface library.

import re
import pandas as pd
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
from transformers import AutoTokenizer, AutoModelWithLMHead
import torch.optim as optim

We set the device to enable GPU processing.

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
device

device(type='cuda', index=0)

Data Preprocessing

movies_file = "movies.csv"

Since the file is in CSV format, we use pandas.read_csv() to read the file

raw_df = pd.read_csv(movies_file)
raw_df

We can see that we have 9742 movie names in the title column. Since the other columns are not useful for us, we will only keep the title column.

movie_names = raw_df['title']
movie_names

0                                Toy Story (1995)
1                                  Jumanji (1995)
2                         Grumpier Old Men (1995)
3                        Waiting to Exhale (1995)
4              Father of the Bride Part II (1995)
                          ...                    
9737    Black Butler: Book of the Atlantic (2017)
9738                 No Game No Life: Zero (2017)
9739                                 Flint (2017)
9740          Bungo Stray Dogs: Dead Apple (2018)
9741          Andrew Dice Clay: Dice Rules (1991)
Name: title, Length: 9742, dtype: object

As seen, the movie names all end with the release year. While it may be interesting to keep the years in the names and let the model output years for generated movies, we can safely assume it does not help the model in understanding movie names.

We remove them with a simple regex expression:

movie_list = list(movie_names)

def remove_year(name):
    return re.sub("\([0-9]+\)", "", name).strip()

movie_list = [remove_year(name) for name in movie_list]

The final movie list looks ready for training. Notice that we do not need to tokenize or process the text any further since GPT2 comes with its own tokenizer that handles text in the approriate way.

movie_list[:5]

['Toy Story',
 'Jumanji',
 'Grumpier Old Men',
 'Waiting to Exhale',
 'Father of the Bride Part II']

However, we should still acquire a fixed length input. We use the average movie name length in words in order to place a safe max length.

avg_length = sum([len(name.split()) for name in movie_list])/len(movie_list)
avg_length

3.2991172243892426

Since the average movie name length in words is 3.3, we can assume that a max length of 10 will cover most of the instances.

max_length = 10

Setup and Training

Before creating the dataset, we download the model and the tokenizer. We need the tokenizer in order to tokenize the data.

tokenizer = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelWithLMHead.from_pretrained("gpt2")

/usr/local/lib/python3.7/dist-packages/transformers/models/auto/modeling_auto.py:698: FutureWarning: The class `AutoModelWithLMHead` is deprecated and will be removed in a future version. Please use `AutoModelForCausalLM` for causal language models, `AutoModelForMaskedLM` for masked language models and `AutoModelForSeq2SeqLM` for encoder-decoder models.
  FutureWarning,

We send the model to the device and initialize the optimizer.

model = model.to(device)

optimizer = optim.AdamW(model.parameters(), lr=3e-4)

According to the GPT-2 paper, to fine-tune the model, use a task designator.

For our purposes, the designator is simply "movie: ". This will be added to the beginning of every example.

To correctly pad and truncate the instances, we find the number of tokens used by this designator:

tokenizer.encode("movie: ")

[41364, 25, 220]

extra_length = len(tokenizer.encode("movie: "))

We create a simple dataset that extends the PyTorch Dataset class:

class MovieDataset(Dataset):  
    def __init__(self, tokenizer, init_token, movie_titles, max_len):
        self.max_len = max_len
        self.tokenizer = tokenizer
        self.eos = self.tokenizer.eos_token
        self.eos_id = self.tokenizer.eos_token_id
        self.movies = movie_titles
        self.result = []

        for movie in self.movies:
            # Encode the text using tokenizer.encode(). We ass EOS at the end
            tokenized = self.tokenizer.encode(init_token + movie + self.eos)
            
            # Padding/truncating the encoded sequence to max_len 
            padded = self.pad_truncate(tokenized)            

            # Creating a tensor and adding to the result
            self.result.append(torch.tensor(padded))

    def __len__(self):
        return len(self.result)


    def __getitem__(self, item):
        return self.result[item]

    def pad_truncate(self, name):
        name_length = len(name) - extra_length
        if name_length < self.max_len:
            difference = self.max_len - name_length
            result = name + [self.eos_id] * difference
        elif name_length > self.max_len:
            result = name[:self.max_len + 2]+[self.eos_id] 
        else:
            result = name
        return result

Then, we create the dataset:

dataset = MovieDataset(tokenizer, "movie: ", movie_list, max_length)

Using a batch_size of 32, we create the dataloader:

dataloader = DataLoader(dataset, batch_size=32, shuffle=True, drop_last=True)

GPT-2 is capable of several tasks, including summarization, generation, and translation. To train for generation, use the same as input as labels:

def train(model, optimizer, dl, epochs):    
    for epoch in range(epochs):
        for idx, batch in enumerate(dl):
             with torch.set_grad_enabled(True):
                optimizer.zero_grad()
                batch = batch.to(device)
                output = model(batch, labels=batch)
                loss = output[0]
                loss.backward()
                optimizer.step()
                if idx % 50 == 0:
                    print("loss: %f, %d"%(loss, idx))

When training a language model, it is easy to overfit the model. This is due to the fact that there is no clear evaluation metric. With most tasks, one can use cross-validation to guarantee not to overfit. For our purposes, we only use 2 epochs for training

train(model=model, optimizer=optimizer, dl=dataloader, epochs=2)

loss: 9.313371, 0
loss: 2.283597, 50
loss: 1.748692, 100
loss: 2.109853, 150
loss: 1.902950, 200
loss: 2.051265, 250
loss: 2.213011, 300
loss: 1.370941, 0
loss: 1.346577, 50
loss: 1.278894, 100
loss: 1.373716, 150
loss: 1.419072, 200
loss: 1.505586, 250
loss: 1.493220, 300

The loss decreased consistently, which means that the model was learning.