training.py

# -*- coding: utf-8 -*-
"""Training

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1mqsFeFVF6wesNc3s-h1GE-xuTWYRymTP
"""

!pip install transformers

import torch
from transformers import GPT2Tokenizer, GPT2LMHeadModel
import numpy as np
import spacy

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)

from google.colab import drive
drive.mount('/content/drive')

file_path = '/content/drive/My Drive/490G1 Final Project/tweets.csv'
import pandas as pd
dataframe = pd.read_csv(file_path)
dataframe
dataframe = dataframe[['author', 'content']]

import re
def remove_line(text):
    return re.sub(r'http\S+', '', text)

dataframe['content'] = dataframe['content'].apply(lambda x: remove_line(x))
name_set = set(dataframe['author'].tolist())

tokenizer(dataframe['content'].tolist()[12])

print(name_set)

for (idx, row) in dataframe.iterrows():
    author = row.author
    content = row.content.strip()

from torch.utils.data import Dataset
from torch.utils.data import Dataset, DataLoader
import os
import json
import csv

# only trump this time
class TrumpDataset(Dataset):
    def __init__(self, file_path = '/content/drive/My Drive/490G1 Final Project/trump_tweet.txt'):
        super().__init__()

        # Trump part
        # self.tweet_list = []
        # self.end_of_text_token = "<|endoftext|>"
        # self.person_list = ['realDonaldTrump']
        # with open(file_path) as tweets:
        #     for line in tweets.readlines():
        #         line = line.strip('\n')
        #         self.tweet_list.append(f"realDonaldTrump:{line}{self.end_of_text_token}")
        
        for (idx, row) in dataframe.iterrows():
            author = row.author
            content = row.content.strip()
            if content == '':
                continue
            else:
                self.person_list.append(str(author))
                self.tweet_list.append(f"{author}:{content}{self.end_of_text_token}")

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

    def __getitem__(self, item):
        return self.tweet_list[item]
    
    def get_person(self):
        return self.person_list

from torch.utils.data import Dataset
from torch.utils.data import Dataset, DataLoader
import os
import json
import csv

# only trump this time
class TwitterDataset(Dataset):
    def __init__(self, author_name):
        super().__init__()
        self.tweet_list = []
        self.end_of_text_token = "<|endoftext|>"
        
        for (idx, row) in dataframe.iterrows():
            author = row.author
            content = row.content.strip()
            if content == '' and author != author_name:
                continue
            else:
                self.tweet_list.append(f"{author}:{content}{self.end_of_text_token}")

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

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

dataset = TrumpDataset()
trump_loader = DataLoader(dataset, batch_size=8, shuffle=True)

BATCH_SIZE = 16
EPOCHS = 5
LEARNING_RATE = 3e-5
WARMUP_STEPS = 2000
TRAINING_STEPS = -1
MAX_SEQ_LEN = 400
from transformers import AdamW, get_linear_schedule_with_warmup

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

from tqdm import tqdm
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=TRAINING_STEPS)
proc_seq_count = 0
sum_loss = 0.0
batch_count = 0

tmp_trump_tens = None
models_folder = "/content/drive/My Drive/490G1 Final Project/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,tweet in tqdm(enumerate(trump_loader)):

        trump_tens = torch.tensor(tokenizer.encode(tweet[0])).unsqueeze(0).to(device)
        #Skip sample from dataset if it is longer than MAX_SEQ_LEN

        if trump_tens.size()[1] > MAX_SEQ_LEN:
            continue
        
        #The first trump sequence in the sequence
        if not torch.is_tensor(tmp_trump_tens):
            tmp_trump_tens = trump_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_trump_tens.size()[1] + trump_tens.size()[1] > MAX_SEQ_LEN:
                work_trump_tens = tmp_trump_tens
                tmp_trump_tens = trump_tens
            else:
                #Add the joke to sequence, continue and try to add more
                tmp_trump_tens = torch.cat([tmp_trump_tens, trump_tens[:,1:]], dim=1)
                continue
        ################## Sequence ready, process it trough the model ##################
            
        outputs = model(work_trump_tens, labels=work_trump_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 == 1:
            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_trump_{epoch}.pt"))

from tqdm import tqdm

EPOCHS = 15

for id, name in enumerate(name_set):
    dataset = TwitterDataset(name)
    data_loader = DataLoader(dataset, batch_size=8, shuffle=True)
    model = GPT2LMHeadModel.from_pretrained('gpt2-medium')
    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=TRAINING_STEPS)
    proc_seq_count = 0
    sum_loss = 0.0
    batch_count = 0

    tmp_trump_tens = None
    models_folder = "/content/drive/My Drive/490G1 Final Project/trained_models"
    if not os.path.exists(models_folder):
        os.mkdir(models_folder)

    for epoch in range(EPOCHS):
        
        print(f"EPOCH {epoch} for {id}{name} started" + '=' * 30)
        
        for idx,tweet in tqdm(enumerate(data_loader)):

            trump_tens = torch.tensor(tokenizer.encode(tweet[0])).unsqueeze(0).to(device)
            #Skip sample from dataset if it is longer than MAX_SEQ_LEN

            if trump_tens.size()[1] > MAX_SEQ_LEN:
                continue
            
            #The first trump sequence in the sequence
            if not torch.is_tensor(tmp_trump_tens):
                tmp_trump_tens = trump_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_trump_tens.size()[1] + trump_tens.size()[1] > MAX_SEQ_LEN:
                    work_trump_tens = tmp_trump_tens
                    tmp_trump_tens = trump_tens
                else:
                    #Add the joke to sequence, continue and try to add more
                    tmp_trump_tens = torch.cat([tmp_trump_tens, trump_tens[:,1:]], dim=1)
                    continue
            ################## Sequence ready, process it trough the model ##################
                
            outputs = model(work_trump_tens, labels=work_trump_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 == 1:
                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_{name}_final.pt"))

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)

MODEL_EPOCH = 4

person_list = list(set(dataset.get_person()))

models_folder = "/content/drive/My Drive/490G1 Final Project/trained_models"

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

trump_output_file_path = f'generated_trump_tweet_{MODEL_EPOCH}.txt'

output_file_path = f'generated_tweet_{MODEL_EPOCH}.txt'

model.eval()
if os.path.exists(trump_output_file_path):
    os.remove(trump_output_file_path)
    
tweet_num = 0
with torch.no_grad():
    
    for person in person_list:
   
        for tweet_idx in range(10):
        
            tweet_finished = False

            cur_ids = torch.tensor(tokenizer.encode(f'{person}:')).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|>'):
                    tweet_finished = True
                    break

            
            if tweet_finished:
                
                tweet_num = tweet_num + 1
                
                output_list = list(cur_ids.squeeze().to('cpu').numpy())
                output_text = tokenizer.decode(output_list)
                #print(f"{tweet_num}:{output_text} \n")

                with open(trump_output_file_path, 'a') as f:
                    print(f"{tweet_num}:{output_text} \n")
                    f.write(f"{tweet_num}:{output_text} \n")