<jupyter_start><jupyter_text>Fine-tuning GPT-2 on a jokes dataset in PyTorchThis notebook was adapted from this blog post - [Fine-tuning large Transformer models on a single GPU in PyTorch - Teaching GPT-2 a sense of humor](https://mf1024.github.io/2019/11/12/Fun-With-GPT-2/). Full credits to the blog post author :)Here I demonstrate how to fine-tune a pre-trained GPT-2 model on a jokes dataset. Let's see if the model can learn to crack some jokes!For this experiment, I will use a pre-trained GPT-2 medium-sized model from the huggingface [transformers repository](https://github.com/huggingface/transformers).<jupyter_code>import torch
from transformers import GPT2Tokenizer, GPT2LMHeadModel
import numpy as np

import logging
logging.getLogger().setLevel(logging.CRITICAL)

import warnings
warnings.filterwarnings('ignore')

device = 'cpu'
if torch.cuda.is_available():
    device = 'cuda'
tokenizer = GPT2Tokenizer.from_pretrained('gpt2-medium')
model = GPT2LMHeadModel.from_pretrained('gpt2-medium')
model = model.to(device)
def choose_from_top(probs, n=5):
    ind = np.argpartition(probs, -n)[-n:]
    top_prob = probs[ind]
    top_prob = top_prob / np.sum(top_prob) # Normalize
    choice = np.random.choice(n, 1, p = top_prob)
    token_id = ind[choice][0]
    return int(token_id)<jupyter_output><empty_output><jupyter_text>PyTorch Dataset module for Short jokes datasetFor fine-tuning the GPT2 model, I will use this [Short Jokes dataset](https://www.kaggle.com/abhinavmoudgil95/short-jokes) published on Kaggle. After each joke, I add "" which is recognized by the GPT2 model as and end of text marker. The marker will allow me to concatenate many jokes in a single input sequence. WARNING : Jokes dataset was made from scrapping web collections. Many of them can be offensive and inappropriate<jupyter_code>from torch.utils.data import Dataset
from torch.utils.data import Dataset, DataLoader
import os
import json
import csv

class JokesDataset(Dataset):
    def __init__(self, jokes_dataset_path = 'jokes/'):
        super().__init__()

        short_jokes_path = os.path.join(jokes_dataset_path, 'shortjokes.csv')

        self.joke_list = []
        self.end_of_text_token = "<|endoftext|>"
        
        with open(short_jokes_path) as csv_file:
            csv_reader = csv.reader(csv_file, delimiter=',')
            
            x = 0
            for row in csv_reader:
                joke_str = f"JOKE:{row[1]}{self.end_of_text_token}"
                self.joke_list.append(joke_str)
        
    def __len__(self):
        return len(self.joke_list)

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

dataset = JokesDataset()
joke_loader = DataLoader(dataset, batch_size=1, shuffle=True)<jupyter_output><empty_output><jupyter_text>HyperparametersI tested many(more than 5) hyperparameter sets till I found one that works the best. I mostly tuned ***BATCH_SIZE*** (in this case, it's the number of forward-backward passes between each optimization step), ***EOPOCHS***, and ***LEARNING_RATE***.For a parameter value starting point for fine-tuning, I inspired from [this](https://github.com/huggingface/transformers/blob/master/examples/question-answering/run_squad.py) and [this](https://github.com/huggingface/transformers/blob/master/examples/text-classification/run_glue.py) huggingface fine-tuning code.<jupyter_code>BATCH_SIZE = 16
EPOCHS = 5
LEARNING_RATE = 3e-5
WARMUP_STEPS = 5000
MAX_SEQ_LEN = 400
from transformers import AdamW, get_linear_schedule_with_warmup

device = 'cpu'
if torch.cuda.is_available():
    device = 'cuda'<jupyter_output><empty_output><jupyter_text>Model trainingI will train the model and save the model weights after each epoch and then I will try to generate jokes with each version of the weight to see which performs the best.<jupyter_code>model = model.to(device)
model.train()
optimizer = AdamW(model.parameters(), lr=LEARNING_RATE)
scheduler = get_linear_schedule_with_warmup(optimizer, WARMUP_STEPS, 1)
proc_seq_count = 0
sum_loss = 0.0
batch_count = 0

tmp_jokes_tens = None
models_folder = "trained_models"
if not os.path.exists(models_folder):
    os.mkdir(models_folder)

for epoch in range(EPOCHS):
    
    print(f"EPOCH {epoch} started" + '=' * 30)
    
    for idx,joke in enumerate(joke_loader):
        
        #################### "Fit as many joke sequences into MAX_SEQ_LEN sequence as possible" logic start ####
        joke_tens = torch.tensor(tokenizer.encode(joke[0])).unsqueeze(0).to(device)
        #Skip sample from dataset if it is longer than MAX_SEQ_LEN
        if joke_tens.size()[1] > MAX_SEQ_LEN:
            continue
        
        #The first joke sequence in the sequence
        if not torch.is_tensor(tmp_jokes_tens):
            tmp_jokes_tens = joke_tens
            continue
        else:
            #The next joke does not fit in so we process the sequence and leave the last joke 
            #as the start for next sequence 
            if tmp_jokes_tens.size()[1] + joke_tens.size()[1] > MAX_SEQ_LEN:
                work_jokes_tens = tmp_jokes_tens
                tmp_jokes_tens = joke_tens
            else:
                #Add the joke to sequence, continue and try to add more
                tmp_jokes_tens = torch.cat([tmp_jokes_tens, joke_tens[:,1:]], dim=1)
                continue
        ################## Sequence ready, process it trough the model ##################
            
        outputs = model(work_jokes_tens, labels=work_jokes_tens)
        loss, logits = outputs[:2]                        
        loss.backward()
        sum_loss = sum_loss + loss.detach().data
                       
        proc_seq_count = proc_seq_count + 1
        if proc_seq_count == BATCH_SIZE:
            proc_seq_count = 0    
            batch_count += 1
            optimizer.step()
            scheduler.step() 
            optimizer.zero_grad()
            model.zero_grad()

        if batch_count == 100:
            print(f"sum loss {sum_loss}")
            batch_count = 0
            sum_loss = 0.0
    
    # Store the model after each epoch to compare the performance of them
    torch.save(model.state_dict(), os.path.join(models_folder, f"gpt2_medium_joker_{epoch}.pt"))<jupyter_output>EPOCH 0 started==============================
sum loss 5468.6435546875
sum loss 5481.8779296875
sum loss 5478.85693359375
sum loss 5465.9462890625
sum loss 5463.58447265625
sum loss 5456.69140625
sum loss 5451.001953125
sum loss 5449.91455078125
sum loss 5438.00732421875
EPOCH 1 started==============================
sum loss 5433.62255859375
sum loss 5411.2529296875
sum loss 5373.556640625
sum loss 5393.3291015625
sum loss 5374.06005859375
sum loss 5362.5751953125
sum loss 5338.7001953125
sum loss 5356.826171875
sum loss 5323.91455078125
sum loss 5330.11279296875
EPOCH 2 started==============================
sum loss 5297.46435546875
sum loss 5278.43896484375
sum loss 5263.87890625
sum loss 5263.6982421875
sum loss 5258.423828125
sum loss 5253.38232421875
sum loss 5244.58447265625
sum loss 5231.24755859375
sum loss 5212.47705078125
sum loss 5205.92626953125
EPOCH 3 started==============================
sum loss 5175.8896484375
sum loss 5161.28955078125
sum loss 5159.814453125
sum loss [...]<jupyter_text>Generating the jokes<jupyter_code>MODEL_EPOCH = 3

models_folder = "trained_models"

model_path = os.path.join(models_folder, f"gpt2_medium_joker_{MODEL_EPOCH}.pt")
model.load_state_dict(torch.load(model_path))

jokes_output_file_path = f'generated_{MODEL_EPOCH}.jokes'
model.eval()
if os.path.exists(jokes_output_file_path):
    os.remove(jokes_output_file_path)
    
joke_num = 0
with torch.no_grad():
   
        for joke_idx in range(1000):
        
            joke_finished = False

            cur_ids = torch.tensor(tokenizer.encode("JOKE:")).unsqueeze(0).to(device)

            for i in range(100):
                outputs = model(cur_ids, labels=cur_ids)
                loss, logits = outputs[:2]
                softmax_logits = torch.softmax(logits[0,-1], dim=0) #Take the first(from only one in this case) batch and the last predicted embedding
                if i < 3:
                    n = 20
                else:
                    n = 3
                next_token_id = choose_from_top(softmax_logits.to('cpu').numpy(), n=n) #Randomly(from the topN probability distribution) select the next word
                cur_ids = torch.cat([cur_ids, torch.ones((1,1)).long().to(device) * next_token_id], dim = 1) # Add the last word to the running sequence

                if next_token_id in tokenizer.encode('<|endoftext|>'):
                    joke_finished = True
                    break

            
            if joke_finished:
                
                joke_num = joke_num + 1
                
                output_list = list(cur_ids.squeeze().to('cpu').numpy())
                output_text = tokenizer.decode(output_list)

                with open(jokes_output_file_path, 'a') as f:
                    f.write(f"{output_text} \n\n")<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.10561748151418165, 'fraction_numerical': 0.06103079130338815, 'mean_word_length': 3.1023087369850613, 'pattern_counts': {'":': 0, '<': 23, '<?xml version=': 0, '>': 24, 'https://': 5, 'lorem ipsum': 0, 'www.': 1, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '535972', 'license': 'MIT', 'n_tokens_mistral': 3279, 'n_tokens_neox': 2863, 'n_words': 759, 'path': 'Training/Pytorch - GPT2 fine tuning and jokes generation/Pytorch_Fine-tuning GPT2_medium_for_jokes_generations.ipynb', 'repo_name': 'baaastijn/notebooks', 'size': '9061'}