<jupyter_start><jupyter_text>This notebook is based on @mf1024's "Teach GPT-2 a sense of humour"I've graciously borrowed much of his logic for fine tuning my training loop Check out his notebook here https://github.com/mf1024/Transformers/blob/master/Teaching%20GPT-2%20a%20sense%20of%20humor.ipynb<jupyter_code>!pip install transformers
import torch
from torch.utils.data import Dataset
from torch.utils.data import Dataset, DataLoader
from transformers import GPT2Tokenizer, GPT2LMHeadModel, AdamW, get_linear_schedule_with_warmup
import numpy as np
import os
import random
output_dir = "./storage/models/"

device = 'cpu'
if torch.cuda.is_available():
    device = 'cuda'
tokenizer = GPT2Tokenizer.from_pretrained(output_dir)
model = GPT2LMHeadModel.from_pretrained(output_dir)
model = model.to(device)
FILE_PATH = os.path.join("storage","data", "film_text.txt")
from language_modelling import ScriptData
dataset = ScriptData(tokenizer= tokenizer, file_path= FILE_PATH )
script_loader = DataLoader(dataset,batch_size=4,shuffle=True)
BATCH_SIZE = 1
EPOCHS = 1
LEARNING_RATE = 0.00002
WARMUP_STEPS = 10000

model = model.to(device)
model.train()
optimizer = AdamW(model.parameters(), lr=LEARNING_RATE)
scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=WARMUP_STEPS, num_training_steps=-1)
script_count = 0
sum_loss = 0.0
batch_count = 0
for epoch in range(EPOCHS):
    print(f"EPOCH {epoch} started" + '=' * 30)
    for idx,script in enumerate(script_loader):
        outputs = model(script.to(device), labels=script.to(device))
        
        loss, logits = outputs[:2]                        
        loss.backward()
        sum_loss = sum_loss + loss.detach().data
                       
        script_count = script_count + 1
        if script_count == BATCH_SIZE:
            script_count = 0    
            batch_count += 1
            optimizer.step()
            scheduler.step() 
            optimizer.zero_grad()
            model.zero_grad()
            
        if batch_count == 200:
            model.eval()
            print(f"sum loss {sum_loss}")
            sample_outputs = model.generate(
                                    bos_token_id=random.randint(1,30000),
                                    do_sample=True,   
                                    top_k=50, 
                                    max_length = 1000,
                                    top_p=0.95, 
                                    num_return_sequences=1
                                )

            print("Output:\n" + 100 * '-')
            for i, sample_output in enumerate(sample_outputs):
                  print("{}: {}".format(i, tokenizer.decode(sample_output, skip_special_tokens=True)))
            
            batch_count = 0
            sum_loss = 0.0
            model.train()
from transformers import WEIGHTS_NAME, CONFIG_NAME
output_model_file = os.path.join(output_dir, WEIGHTS_NAME)
output_config_file = os.path.join(output_dir, CONFIG_NAME)

torch.save(model.state_dict(), output_model_file)
model.config.to_json_file(output_config_file)
tokenizer.save_vocabulary(output_dir)
model = GPT2LMHeadModel.from_pretrained(output_dir)
tokenizer = GPT2Tokenizer.from_pretrained(output_dir)
input_ids = tokenizer.encode('         He kisses her softly and takes out his gun.         ', return_tensors='pt')
model.eval()
sample_outputs = model.generate(
                        input_ids= input_ids,
                        num_beams= 5,
                        max_length = 1000,
                        top_p=0.85, 
                        num_return_sequences=3
                    )

print("Output:\n" + 100 * '-')
for i, sample_output in enumerate(sample_outputs):
      print("{}: {}".format(i, tokenizer.decode(sample_output, skip_special_tokens=True)))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.0942161737764983, 'fraction_numerical': 0.023030620256477362, 'mean_word_length': 2.415549597855228, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '74349', 'license': 'MIT', 'n_tokens_mistral': 1231, 'n_tokens_neox': 1136, 'n_words': 231, 'path': 'script_buddy/script_generation.ipynb', 'repo_name': 'sndrskth/script_buddy_v2', 'size': '3821'}