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
Please use pip install {library name} in order to install the libraries below if they are not installed. "transformers" is the Huggingface library.
importreimportpandasaspdimportnumpyasnpimporttorchfromtorch.utils.dataimportDataset,DataLoaderfromtransformersimportAutoTokenizer,AutoModelWithLMHeadimporttorch.optimasoptimWe set the device to enable GPU processing.
device=torch.device('cuda:0'iftorch.cuda.is_available()else'cpu')devicemovies_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_dfWe 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_namesAs 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)defremove_year(name):returnre.sub("\([0-9]+\)","",name).strip()movie_list=[remove_year(name)fornameinmovie_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]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())fornameinmovie_list])/len(movie_list)avg_lengthSince 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=10Before 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")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: ")extra_length=len(tokenizer.encode("movie: "))We create a simple dataset that extends the PyTorch Dataset class:
classMovieDataset(Dataset):def__init__(self,tokenizer,init_token,movie_titles,max_len):self.max_len=max_lenself.tokenizer=tokenizerself.eos=self.tokenizer.eos_tokenself.eos_id=self.tokenizer.eos_token_idself.movies=movie_titlesself.result=[]formovieinself.movies:# Encode the text using tokenizer.encode(). We ass EOS at the endtokenized=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 resultself.result.append(torch.tensor(padded))def__len__(self):returnlen(self.result)def__getitem__(self,item):returnself.result[item]defpad_truncate(self,name):name_length=len(name)-extra_lengthifname_length<self.max_len:difference=self.max_len-name_lengthresult=name+[self.eos_id]*differenceelifname_length>self.max_len:result=name[:self.max_len+2]+[self.eos_id]else:result=namereturnresultThen, 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:
deftrain(model,optimizer,dl,epochs):forepochinrange(epochs):foridx,batchinenumerate(dl):withtorch.set_grad_enabled(True):optimizer.zero_grad()batch=batch.to(device)output=model(batch,labels=batch)loss=output[0]loss.backward()optimizer.step()ifidx%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)The loss decreased consistently, which means that the model was learning.
Movie Name Generation
In order to verify, we generate 20 movie names that are not existent in the movie list.
The generation methodology is as follows:
- The task designator is initially fed into the model
- A choice from the top-k choices is selected. A common question is why not use the highest ranked choice always. The simple answer is that introducing randomness helps the model create different outputs. There are several sampling methods in the literature, such as top-k and nucleus sampling. Im this example, we use top-k, where k = 9. K is a hyperparameter that improves the performance with tweaking. Feel free to play around with it to see the effects.
- The choice is added to the sequence and the current sequence is fed to the model.
- Repeat steps 2 and 3 until either max_len is achieved or the EOS token is generated.
deftopk(probs,n=9):# The scores are initially softmaxed to convert to probabilitiesprobs=torch.softmax(probs,dim=-1)# PyTorch has its own topk method, which we use heretokensProb,topIx=torch.topk(probs,k=n)# The new selection pool (9 choices) is normalizedtokensProb=tokensProb/torch.sum(tokensProb)# Send to CPU for numpy handlingtokensProb=tokensProb.cpu().detach().numpy()# Make a random choice from the pool based on the new prob distributionchoice=np.random.choice(n,1,p=tokensProb)tokenId=topIx[choice][0]returnint(tokenId)defmodel_infer(model,tokenizer,init_token,max_length=10):# Preprocess the init token (task designator)init_id=tokenizer.encode(init_token)result=init_idinit_input=torch.tensor(init_id).unsqueeze(0).to(device)withtorch.set_grad_enabled(False):# Feed the init token to the modeloutput=model(init_input)# Flatten the logits at the final time steplogits=output.logits[0,-1]# Make a top-k choice and append to the resultresult.append(topk(logits))# For max_length times:foriinrange(max_length):# Feed the current sequence to the model and make a choiceinput=torch.tensor(result).unsqueeze(0).to(device)output=model(input)logits=output.logits[0,-1]res_id=topk(logits)# If the chosen token is EOS, return the resultifres_id==tokenizer.eos_token_id:returntokenizer.decode(result)else:# Append to the sequence result.append(res_id)# IF no EOS is generated, return after the max_lenreturntokenizer.decode(result)Generating 20 unique movie names:
results=set()whilelen(results)<20:name=model_infer(model,tokenizer,"movie:").replace("movie: ","").strip()ifnamenotinmovie_list:results.add(name)print(name)As shown, the movie names look realistic, meaning that the model learned how to generate movie names correctly.
PyTorch makes it very easy to save the model:
torch.save(model.state_dict(),"movie_gpt.pth")And, if you need to load the model in the future for quick inference without having to train:
model.load_state_dict(torch.load("movie_gpt.pth"))In this tutorial, we learnt how to fine-tune the Huggingface GPT model to perform movie name generation. The same methodology can be applied to any language model available on https://huggingface.co/models