2022 intel corp s mayank kumar raunak javier turek nicole beckage implementation of a new method for finetuning transformer models that we call information gain filtration igf on wikitext data set and compared the results with the standard finetuning method steps followed in the code 1 generate a objective dataset of pairs x igx igxinformativeness of context x our ig information gain model is learning to predict the informativeness of a particular context informativeness is the change in metric between the models accuracy on an objective set before and after seeing that context for casual language modeling the metric is perplexity 2 a secondary learner is trained to infer a function approximation for ig using the dataset created in 1 3 the learner created in 2 is used to inform the finetuning process and filter out low informative samples last a plot is generated to compare the performance of igf to standard finetuning without any filtering prerequisite libraries collecting n pairs for training the secondary learner args contextlen the maximum total input sequence length after tokenization sequences longer than this will be truncated sequences shorter will be padded maxsteps to calculate training epochs of secondary learner sizeobjectiveset size of objective data set used to create x igx pairs which is the training data for secondary learner minlen the minimum length of the article to be used as objective set trim if true truncate the context if it exceeds context length datafile tokenized data set split for training and evaluation of model igfdatafile file to store i igx paired data set to train secondary learner returns data stored in igfdatafile generates same data everytime generate traindata and objectiveset keeps model same across runs model lmoptimizer lmscheduler recopygpt2model device maxsteps store original model weights can we train on gpu load pretrained model collect igf pairs and save to file demo jbl clean up delete model and data we don t need anymore train the secondary learner args secondarylearnertraindata data set with x igx pairs to train secondary learner where igx measure of informativeness and x context secondarylearnermaxepochs number of epochs to train secondary learner secondarylearnerbatchsize batch size to train secondary learner evalfreq object secondary model evaluation can be triggered at evalfreq igfmodelpath path to store trained secondary learner returns trained secondary learner load pretrained model initialize secondary learner to use embedding weights of model train secondary learner finetune with igf if secondarylearner is not none else standard finetuning args model pretrained gpt2 model traindataset data set to train gpt2 model testdataset evaluate gpt2 model contextlen the maximum total input sequence length after tokenization sequences longer than this will be truncated sequences shorter will be padded maxsteps to calculate training epochs batchsize batch size to train gpt2 model threshold the threshold value used by secondary learner to filter the traindata and allow only informative data as input to the model recopymodel reset the model to the original pretrained gpt2 weights after each iteration secondarylearner selection of igf as finetuning method if not none evalinterval number of batches after which decay the selectivity of our secondary learner filter from 1 standard deviation above average to 1 below average finetunedmodelname name of the final finaltuned gpt2 model returns finetuned gpt2 model compute the performance of the transformer model at the beginning here we implement the simple nonconstant threshold for the predicted igx value we will decay the selectivity of our secondary learner filter from 1 standard deviation above average to 1 below average after 10 batches if we passed the filter add the context to the batch once the batch is filled with enough contexts backprop on the batch do lm backprop compute the performance of the transformer model at this batch break out of the loop after 60 batches save finetuned transformer model do some cleaning up so we can reinitialize for the next run of this function required parameters function calls collecting n pairs of context and information gainx igx for training the secondary learner load train data for secondary learner train secondary learner load pretrained gpt2 model generate train and test data to train and evaluate gpt2 model finetuning of the gpt2 model using igf information gain filtration 2022 intel corp s mayank kumar raunak javier turek nicole beckage implementation of a new method for fine tuning transformer models that we call information gain filtration igf on wikitext data set and compared the results with the standard fine tuning method steps followed in the code 1 generate a objective dataset of pairs x ig x ig x informativeness of context x our ig information gain model is learning to predict the informativeness of a particular context informativeness is the change in metric between the model s accuracy on an objective set before and after seeing that context for casual language modeling the metric is perplexity 2 a secondary learner is trained to infer a function approximation for ig using the dataset created in 1 3 the learner created in 2 is used to inform the fine tuning process and filter out low informative samples last a plot is generated to compare the performance of igf to standard fine tuning without any filtering prerequisite libraries collecting n pairs for training the secondary learner args context_len the maximum total input sequence length after tokenization sequences longer than this will be truncated sequences shorter will be padded max_steps to calculate training epochs of secondary learner size_objective_set size of objective data set used to create x ig x pairs which is the training data for secondary learner min_len the minimum length of the article to be used as objective set trim if true truncate the context if it exceeds context length data_file tokenized data set split for training and evaluation of model igf_data_file file to store i ig x paired data set to train secondary learner returns data stored in igf_data_file generates same data everytime generate train_data and objective_set keeps model same across runs model lm_optimizer lm_scheduler recopy_gpt2 model device max_steps store original model weights can we train on gpu load pretrained model collect igf pairs and save to file demo jbl clean up delete model and data we don t need anymore train the secondary learner args secondary_learner_train_data data set with x ig x pairs to train secondary learner where ig x measure of informativeness and x context secondary_learner_max_epochs number of epochs to train secondary learner secondary_learner_batch_size batch size to train secondary learner eval_freq object secondary model evaluation can be triggered at eval_freq igf_model_path path to store trained secondary learner returns trained secondary learner load pre trained model initialize secondary learner to use embedding weights of model train secondary learner fine tune with igf if secondary_learner is not none else standard fine tuning args model pre trained gpt 2 model train_dataset data set to train gpt 2 model test_dataset evaluate 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 max_steps to calculate training epochs batch_size batch size to train gpt 2 model threshold the threshold value used by secondary learner to filter the train_data and allow only informative data as input to the model recopy_model reset the model to the original pretrained gpt 2 weights after each iteration secondary_learner selection of igf as fine tuning method if not none eval_interval number of batches after which decay the selectivity of our secondary learner filter from 1 standard deviation above average to 1 below average fine tuned_model_name name of the final final tuned gpt 2 model returns fine tuned gpt 2 model compute the performance of the transformer model at the beginning here we implement the simple non constant threshold for the predicted ig x value we will decay the selectivity of our secondary learner filter from 1 standard deviation above average to 1 below average after 10 batches if we passed the filter add the context to the batch once the batch is filled with enough contexts backprop on the batch do lm backprop update learning rate schedule compute the performance of the transformer model at this batch break out of the loop after 60 batches save finetuned transformer model do some cleaning up so we can reinitialize for the next run of this function required parameters function calls collecting n pairs of context and information gain x ig x for training the secondary learner load train data for secondary learner train secondary learner load pretrained gpt2 model generate train and test data to train and evaluate gpt2 model fine tuning of the gpt2 model using igf information gain filtration
import argparse
import random

import joblib
import numpy as np
import torch
from igf.igf import (
    SecondaryLearner,
    collect_objective_set,
    compute_perplexity,
    generate_datasets,
    load_gpt2,
    recopy_gpt2,
    set_seed,
    train_secondary_learner,
)
from torch.utils.data import DataLoader, RandomSampler

from transformers import GPT2LMHeadModel


def generate_n_pairs(
    context_len=32,
    max_steps=10,
    size_objective_set=100,
    min_len=1026,
    trim=True,
    data_file="data/tokenized_stories_train_wikitext103.jbl",
    igf_data_file="igf_context_pairs.jbl",
):
    
    
    set_seed(3)
    
    train_data, objective_set = generate_datasets(
        context_len, data_file, number=size_objective_set, min_len=1026, trim=True
    )
    
    set_seed(4)
    
    
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    
    model = load_gpt2("gpt2").to(device)
    print("computing perplexity on objective set")
    orig_perp = compute_perplexity(model, objective_set, context_len).item()
    print("perplexity on objective set:", orig_perp)

    
    collect_objective_set(model, orig_perp, context_len, train_data, objective_set, max_steps, device, igf_data_file)

    
    del model, train_data, objective_set
    torch.cuda.empty_cache()


def training_secondary_learner(
    secondary_learner_train_data,
    secondary_learner_max_epochs=15,
    secondary_learner_batch_size=128,
    eval_freq=100,
    igf_model_path="igf_model.pt",
):
    

    set_seed(42)

    
    model = GPT2LMHeadModel.from_pretrained("gpt2")

    
    secondary_learner = SecondaryLearner(model)

    
    secondary_learner = train_secondary_learner(
        secondary_learner,
        secondary_learner_train_data,
        max_epochs=secondary_learner_max_epochs,
        batch_size=secondary_learner_batch_size,
        eval_freq=100,
        igf_model_path=igf_model_path,
    )

    del model, secondary_learner_train_data
    torch.cuda.empty_cache()

    return secondary_learner


def finetune(
    model,
    train_dataset,
    test_dataset,
    context_len=32,
    max_steps=1000,
    batch_size=16,
    threshold=1.0,
    recopy_model=recopy_gpt2,
    secondary_learner=None,
    eval_interval=10,
    finetuned_model_name="gpt2_finetuned.pt",
):
    

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    train_sampler = RandomSampler(train_dataset)
    train_dataloader = DataLoader(train_dataset, sampler=train_sampler)

    num_train_epochs = max_steps // (len(train_dataset)) + 1
    global_step = 0
    context = torch.zeros((1, context_len), dtype=torch.long, device=device)
    model, lm_optimizer, lm_scheduler = recopy_model(model, device, max_steps)

    model.train()
    if secondary_learner is not None:
        secondary_learner.to(device)
        secondary_learner.eval()
    contexts = []
    examples = 0

    observed_qs = []
    test_perps = []

    
    real_perp = compute_perplexity(model, test_dataset, context_len)
    test_perps.append(real_perp)
    print("Test perplexity, step", global_step, ":", real_perp)
    for epoch in range(int(num_train_epochs)):
        for step, example in enumerate(train_dataloader):
            torch.cuda.empty_cache()
            start = random.randint(0, example.size(2) - context_len - 1)
            context[0, :] = example[0, 0, start : start + context_len]
            lm_optimizer.zero_grad()
            outputs = model(context, labels=context)
            do_backprop = True

            if secondary_learner is not None:
                predicted_q = secondary_learner.forward(
                    torch.tensor(context, dtype=torch.long, device=device).unsqueeze(0)
                )[0].item()
                observed_qs.append(float(predicted_q))

                
                
                

                if global_step == 10:
                    threshold = -1
                if predicted_q < threshold:
                    do_backprop = False

            
            if do_backprop:
                contexts.append(np.array(context.cpu()))
                lm_loss = outputs[0]
                lm_loss.backward()
                examples += 1

            del outputs

            
            if examples == batch_size:
                torch.cuda.empty_cache()
                examples = 0
                
                torch.nn.utils.clip_grad_norm_(model.parameters(), 3.0)
                lm_optimizer.step()
                lm_scheduler.step()  
                global_step += 1
                
                if global_step % eval_interval == 0:
                    real_perp = compute_perplexity(model, test_dataset, context_len)
                    test_perps.append(real_perp)

                    print("Test perplexity, step", global_step, ":", real_perp)
            
            if max_steps > 0 and global_step > 60:
                break
        if max_steps > 0 and global_step > 60:
            break

    
    torch.save(model.state_dict(), finetuned_model_name)
    torch.cuda.empty_cache()
    
    del lm_optimizer
    del lm_scheduler
    return model


def main():
    parser = argparse.ArgumentParser(description="Fine-tune a transformer model with IGF on a language modeling task")

    
    parser.add_argument(
        "--data_dir",
        default=None,
        type=str,
        required=True,
        help="The input data dir. Should contain data files for WikiText.",
    )
    parser.add_argument(
        "--model_name_or_path",
        default=None,
        type=str,
        required=True,
        help="Path to pretrained model or model identifier from huggingface.co/models",
    )
    parser.add_argument(
        "--data_file",
        type=str,
        default=None,
        help=(
            "A jbl file containing tokenized data which can be split as objective dataset, "
            "train_dataset and test_dataset."
        ),
    )

    parser.add_argument(
        "--igf_data_file",
        type=str,
        default=None,
        help="A jbl file containing the context and information gain pairs to train secondary learner.",
    )

    parser.add_argument(
        "--output_dir",
        default=None,
        type=str,
        required=True,
        help="The output directory where the final fine-tuned model is stored.",
    )

    parser.add_argument(
        "--tokenizer_name",
        default=None,
        type=str,
        help="Pretrained tokenizer name or path if not the same as model_name",
    )
    parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")

    parser.add_argument(
        "--context_len",
        default=32,
        type=int,
        help=(
            "The maximum total input sequence length after tokenization. Sequences longer "
            "than this will be truncated, sequences shorter will be padded."
        ),
    )

    parser.add_argument(
        "--size_objective_set",
        default=100,
        type=int,
        help="number of articles that are long enough to be used as our objective set",
    )
    parser.add_argument(
        "--eval_freq", default=100, type=int, help="secondary model evaluation is triggered at eval_freq"
    )

    parser.add_argument("--max_steps", default=1000, type=int, help="To calculate training epochs")

    parser.add_argument(
        "--secondary_learner_batch_size",
        default=128,
        type=int,
        help="batch size of training data for secondary learner",
    )

    parser.add_argument(
        "--batch_size", default=16, type=int, help="batch size of training data of language model(gpt2) "
    )

    parser.add_argument(
        "--eval_interval",
        default=10,
        type=int,
        help=(
            "decay the selectivity of our secondary learner filter from "
            "1 standard deviation above average to 1 below average after 10 batches"
        ),
    )

    parser.add_argument(
        "--number", default=100, type=int, help="The number of examples split to be used as objective_set/test_data"
    )

    parser.add_argument(
        "--min_len", default=1026, type=int, help="The minimum length of the article to be used as objective set"
    )

    parser.add_argument(
        "--secondary_learner_max_epochs", default=15, type=int, help="number of epochs to train secondary learner"
    )

    parser.add_argument("--trim", default=True, type=bool, help="truncate the example if it exceeds context length")

    parser.add_argument(
        "--threshold",
        default=1.0,
        type=float,
        help=(
            "The threshold value used by secondary learner to filter the train_data and allow only"
            " informative data as input to the model"
        ),
    )

    parser.add_argument("--finetuned_model_name", default="gpt2_finetuned.pt", type=str, help="finetuned_model_name")

    parser.add_argument(
        "--recopy_model",
        default=recopy_gpt2,
        type=str,
        help="Reset the model to the original pretrained GPT-2 weights after each iteration",
    )

    
    
    generate_n_pairs(
        context_len=32,
        max_steps=10,
        size_objective_set=100,
        min_len=1026,
        trim=True,
        data_file="data/tokenized_stories_train_wikitext103.jbl",
        igf_data_file="igf_context_pairs.jbl",
    )

    
    secondary_learner_train_data = joblib.load("data/IGF_values.jbl")

    
    secondary_learner = training_secondary_learner(
        secondary_learner_train_data,
        secondary_learner_max_epochs=15,
        secondary_learner_batch_size=128,
        eval_freq=100,
        igf_model_path="igf_model.pt",
    )

    
    model = GPT2LMHeadModel.from_pretrained("gpt2")
    set_seed(42)

    
    train_dataset, test_dataset = generate_datasets(
        context_len=32, file="data/tokenized_stories_train_wikitext103.jbl", number=100, min_len=1026, trim=True
    )

    
    finetune(
        model,
        train_dataset,
        test_dataset,
        context_len=32,
        max_steps=1000,
        batch_size=16,
        threshold=1.0,
        recopy_model=recopy_gpt2,
        secondary_learner=secondary_learner,
        eval_interval=10,
        finetuned_model_name="gpt2_finetuned.pt",
    )


if __name__ == "__main__":
    main()