2022 intel corp s mayank kumar raunak javier turek nicole backage for reproducible training args seed a seed for reproducible training computes perplexity of the transformer model on data in testdata args model pretrained gpt2 model testdata data on which perplexity calculation is required contextlen the maximum total input sequence length after tokenization sequences longer than this will be truncated sequences shorter will be padded returns perplexity on input test data pad load original gpt2 and save off for quicker loading args modelname gpt2 returns gpt2 model reset the model to the original pretrained gpt2 weights after each iteration args origmodel original pretrained gpt2 model imported from transformers library device cpugpu maxsteps number of training steps returns original pretrained gpt2 model lmoptimizer adam optimizer with decoupled weight decay lmscheduler linear scheduler with the appropriate schedule save the perplexity differences to filename args contexts example on which the perplexity is calculated realperps perplexity after backpropagating on the selected context pastperps perplexity of model before training on the context filename file to store perplexity differences returns file with perplexity differences save the perplexity differences to filename collect individual igf values from pretrained transformer model maxsteps samples of training data to train secondary model args model pretrained gpt2 model origperp perplexity of original pretrained gpt2 model contextlen the maximum total input sequence length after tokenization sequences longer than this will be truncated sequences shorter will be padded traindata data to train model objectiveset contexts used to create x igx pairs which is the training data for secondary learner maxsteps to calculate training epochs of model device gpucpu filename to store intermediate perplexity differences recopymodel reset the model to the original pretrained gpt2 weights after each iteration returns file stored intermediate perplexity differences in intermediate stages initialize variables to record relevant information initialize the transformer model compute perplexity of initial transformer model for comparison do lm backprop compute perplexity after backpropagating on the selected context periodically save the stored x igx pairs reset the pretrained model to the original pretrained gpt2 weights after each iteration generate objective set and training set args contextlen the maximum total input sequence length after tokenization sequences longer than this will be truncated sequences shorter will be padded file tokenized data split into training set and objective set number size of objective dataset minlen minimum length of a context in objective set trim if true truncate the context if it exceeds context length returns generated objective set and training data generate objective set and training set designate the first number 100 articles that are long enough to be used as our objective set rest that are long enough are training data for secondary learner train the secondary learner igfmodel args secondarylearner secondary learner traindataset data to train secondary learner maxepochs number of epochs to train secondary learner batchsize batch size of training data of secondary learner evalfreq secondary model evaluation can be triggered at evalfreq igfmodelpath path to store trained secondary learner returns trained secondary learner we will use the first 512 pairs from our dataset as a test set for our secondary learner and the rest to train secondary learner model set up todo in original code this is written as number of actual batches seen not number of items seen but other places it is number of items instead improve consistency changed this to epochs for clarity iterate through batches until we ve used maxsteps batches model trains fairly quickly so we won t wait for a full epoch eval is triggered at evalfreq and end of epochs compute performance of the secondary learner after this batch our secondary learner we use a simple convolutional network as our secondary learner args model pretrained gpt2 model embeddings are from the pretrained model forward pass through the secondary learner args context context input to the secondary learner returns tensor after squeeze operation load the secondary learner args statepath path to save secondary learner model pretrained gpt2 returns secondary learner 2022 intel corp s mayank kumar raunak javier turek nicole backage for reproducible training args seed a seed for reproducible training computes perplexity of the transformer model on data in test_data args model pre trained gpt2 model test_data data on which perplexity calculation is required context_len the maximum total input sequence length after tokenization sequences longer than this will be truncated sequences shorter will be padded returns perplexity on input test data pad load original gpt2 and save off for quicker loading args model_name gpt 2 returns gpt 2 model reset the model to the original pretrained gpt 2 weights after each iteration args orig_model original pretrained gpt 2 model imported from transformers library device cpu gpu max_steps number of training steps returns original pretrained gpt 2 model lm_optimizer adam optimizer with decoupled weight decay lm_scheduler linear scheduler with the appropriate schedule save the perplexity differences to filename args contexts example on which the perplexity is calculated real_perps perplexity after back propagating on the selected context past_perps perplexity of model before training on the context filename file to store perplexity differences returns file with perplexity differences save the perplexity differences to filename collect individual igf values from pre trained transformer model max_steps samples of training data to train secondary model args model pre trained gpt2 model orig_perp perplexity of original pretrained gpt 2 model context_len the maximum total input sequence length after tokenization sequences longer than this will be truncated sequences shorter will be padded train_data data to train model objective_set contexts used to create x ig x pairs which is the training data for secondary learner max_steps to calculate training epochs of model device gpu cpu filename to store intermediate perplexity differences recopy_model reset the model to the original pretrained gpt 2 weights after each iteration returns file stored intermediate perplexity differences in intermediate stages initialize variables to record relevant information initialize the transformer model compute perplexity of initial transformer model for comparison do lm backprop update learning rate schedule compute perplexity after back propagating on the selected context periodically save the stored x ig x pairs reset the pretrained model to the original pretrained gpt 2 weights after each iteration generate objective set and training set args context_len the maximum total input sequence length after tokenization sequences longer than this will be truncated sequences shorter will be padded file tokenized data split into training set and objective set number size of objective dataset min_len minimum length of a context in objective set trim if true truncate the context if it exceeds context length returns generated objective set and training data generate objective set and training set designate the first number 100 articles that are long enough to be used as our objective set rest that are long enough are training data for secondary learner train the secondary learner igf_model args secondary_learner secondary learner train_dataset data to train secondary learner max_epochs number of epochs to train secondary learner batch_size batch size of training data of secondary learner eval_freq secondary model evaluation can be triggered at eval_freq igf_model_path path to store trained secondary learner returns trained secondary learner we will use the first 512 pairs from our dataset as a test set for our secondary learner and the rest to train secondary learner model set up todo in original code this is written as number of actual batches seen not number of items seen but other places it is number of items instead improve consistency changed this to epochs for clarity iterate through batches until we ve used max_steps batches model trains fairly quickly so we won t wait for a full epoch eval is triggered at eval_freq and end of epochs compute performance of the secondary learner after this batch our secondary learner we use a simple convolutional network as our secondary learner args model pre trained gpt2 model embeddings are from the pretrained model forward pass through the secondary learner args context context input to the secondary learner returns tensor after squeeze operation load the secondary learner args state_path path to save secondary learner model pretrained gpt 2 returns secondary learner this calls __init__
import copy
import logging
import random

import joblib
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from tqdm import tqdm

from transformers import AdamW, GPT2LMHeadModel, get_linear_schedule_with_warmup


logger = logging.getLogger(__name__)


def set_seed(seed):
    
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


def compute_perplexity(model, test_data, context_len):
    

    model.eval()
    device = next(model.parameters()).device
    eval_batch_size = 1
    context = torch.zeros((eval_batch_size, context_len), dtype=torch.long, device=device)
    eval_dataloader = DataLoader(test_data, shuffle=False, batch_size=eval_batch_size)
    eval_loss = torch.zeros(1, device=device)
    nb_eval_examples = 0
    for batch in eval_dataloader:
        batch.to(device)
        
        context.zero_()
        for i in range(eval_batch_size):
            context[i, :] = batch[i]
        outputs = model(context, labels=context)
        eval_loss += outputs[0].sum().item()
        nb_eval_examples += batch.size(0)
    eval_loss = eval_loss / nb_eval_examples
    perplexity = torch.exp(eval_loss)
    model.train()
    return perplexity


def load_gpt2(model_name="gpt2"):
    

    model = GPT2LMHeadModel.from_pretrained(model_name, output_hidden_states=True)
    torch.save(model.state_dict(), model_name + "local.pt")
    return model


def recopy_gpt2(orig_model, device, max_steps):
    
    model = copy.deepcopy(orig_model)
    model.to(device)

    no_decay = ["bias", "LayerNorm.weight"]
    optimizer_grouped_parameters = [
        {
            "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
            "weight_decay": 0.0,
        },
        {"params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], "weight_decay": 0.0},
    ]
    lm_optimizer = AdamW(optimizer_grouped_parameters, lr=5e-5, eps=1e-8)
    lm_scheduler = get_linear_schedule_with_warmup(lm_optimizer, 0, max_steps)
    torch.cuda.empty_cache()
    return model, lm_optimizer, lm_scheduler


def intermittent_save(contexts, real_perps, past_perps, filename):
    
    
    avg = np.array(real_perps).mean()
    std = np.array(real_perps).std()
    perp_diff = (real_perps - avg) / std
    data_final = list(zip(contexts, perp_diff, past_perps))
    joblib.dump(data_final, filename)


def collect_objective_set(
    model,
    orig_perp,
    context_len,
    train_data,
    objective_set,
    max_steps,
    device,
    filename="dev.jbl",
    recopy_model=recopy_gpt2,
):
    

    
    contexts = []
    real_perps = []
    past_perps = []

    
    orig_model = copy.deepcopy(model)
    orig_model.to(device="cpu")
    torch.cuda.empty_cache()

    
    model.train()
    model, lm_optimizer, lm_scheduler = recopy_model(orig_model, device, max_steps)

    for step in tqdm(range(max_steps)):
        context = torch.zeros((1, context_len), dtype=torch.long, device=device)
        story = random.choice(train_data)
        start = random.randint(0, len(story[0]) - context_len - 1)
        context[0, :] = story[0][start : start + context_len]
        lm_optimizer.zero_grad()
        outputs = model(context, labels=context)
        lm_loss = outputs[0]
        past_perp = compute_perplexity(model, context, context_len)
        model.train()
        lm_loss.backward()
        
        torch.nn.utils.clip_grad_norm_(model.parameters(), 3.0)
        lm_optimizer.step()
        lm_scheduler.step()  

        
        real_perp = compute_perplexity(model, objective_set, context_len)

        
        if step % 1000 == 0 and step > 1:
            intermittent_save(contexts, real_perps, past_perps, filename)

        
        model, lm_optimizer, lm_scheduler = recopy_model(orig_model, device, max_steps)

        past_perps.append(past_perp.item())
        real_perps.append(orig_perp - real_perp.item())
        contexts.append(np.array(context.cpu()))

    intermittent_save(contexts, real_perps, past_perps, filename)


def generate_datasets(
    context_len, file="data/tokenized_stories_train_wikitext103.jbl", number=100, min_len=1026, trim=True
):
    
    
    
    
    

    data = joblib.load(file)
    print("data loaded")
    objective_set = []
    if trim:
        for i, example in enumerate(data):
            if len(example[0]) > min_len:
                start = random.randint(0, len(example[0]) - context_len - 1)
                objective_set.append(example[0, start : start + context_len])
            if len(objective_set) >= number:
                break
        train_data = []
        for j in range(i + 1, len(data)):
            if len(data[j][0]) > min_len:
                train_data.append(data[j])
    else:
        objective_set = data[0:number]
        train_data = data[number:]

    joblib.dump(objective_set, "objective_set.jbl")
    print("objective set saved")
    return train_data, objective_set


def train_secondary_learner(
    secondary_learner, train_dataset, max_epochs, batch_size, eval_freq=50, igf_model_path="secondary_learner.pt"
):
    
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    
    
    test_dataset = train_dataset[:512]
    train_dataset = train_dataset[512:]
    train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=batch_size)
    test_dataloader = DataLoader(test_dataset, shuffle=False, batch_size=batch_size)

    
    loss = nn.MSELoss()
    test_loss = nn.MSELoss(reduction="sum")
    secondary_learner.to(device)
    q_optimizer = torch.optim.Adam(secondary_learner.parameters(), lr=0.00001)
    secondary_learner.train()

    
    
    
    best_test_loss = float("inf")
    
    for epoch in range(int(max_epochs)):
        tr_q_loss = 0.0
        secondary_learner.train()
        for step, batch in enumerate(train_dataloader):
            context = batch[0].to(device)
            real_q = batch[1].to(device)
            predicted_q = secondary_learner(context)
            q_optimizer.zero_grad()
            q_loss = loss(predicted_q, real_q.float())
            q_loss.backward()
            q_optimizer.step()
            tr_q_loss += q_loss.item()

            
            
            if (step % eval_freq == 0 and step > 0) or ((step + 1) == len(train_dataloader)):
                tr_loss = tr_q_loss / (step + 1)

                secondary_learner.eval()
                q_loss2 = 0.0
                sum_q2 = 0.0
                predicted = []
                actual = []
                
                for step2, batch2 in enumerate(test_dataloader):
                    features2 = batch2[0].to(device)
                    real_q2 = batch2[1].to(device)
                    predicted_q2 = secondary_learner(features2)
                    q_loss2 += test_loss(predicted_q2, real_q2).item()
                    sum_q2 += torch.sum(predicted_q2).item()
                    for ei, i in enumerate(predicted_q2.cpu().detach().numpy()):
                        predicted.append(i.item())
                    for ei, i in enumerate(real_q2.cpu().detach().numpy()):
                        actual.append(i.item())

                q_loss2 /= len(test_dataset)
                print(
                    "Epoch: ",
                    epoch,
                    "step: ",
                    step,
                    "Avg. q:",
                    sum_q2 / len(test_dataset),
                    "Train Loss: ",
                    tr_loss,
                    "Test Loss: ",
                    q_loss2,
                )
                if q_loss2 < best_test_loss:
                    joblib.dump((predicted, actual), "pred_vs_actual.jbl")
                    torch.save(secondary_learner.state_dict(), igf_model_path)
                    best_test_loss = q_loss2

            secondary_learner.train()
    return secondary_learner


class SecondaryLearner(nn.Module):
    

    def __init__(self, model):
        
        
        super(SecondaryLearner, self).__init__()
        self.embeddings = model.transformer.wte
        self.embeddings.weight = copy.deepcopy(model.transformer.wte.weight)
        self.conv = nn.Conv1d(self.embeddings.weight.size(1), 256, 3, padding=1)
        self.fc = nn.Sequential(nn.Linear(256, 32), nn.Dropout(p=0.1), nn.Linear(32, 32), nn.Linear(32, 1))

    def forward(self, context):
        
        pooled = torch.max(self.conv(self.embeddings(context).squeeze(1).transpose(1, 2)), 2)[0]
        qs = self.fc(pooled)
        return qs.squeeze(1)

    @classmethod
    def from_pretrained(cls, state_path, model):
        

        secondary_learner = cls(model)  
        state_dict = torch.load(state_path)
        secondary_learner.load_state_dict(state_dict)
        secondary_learner.embeddings = model.transformer.wte
        secondary_learner.embeddings.weight = copy.deepcopy(model.transformer.wte.weight)
        return secondary_learner