def collect_objective_set(
    model,
    orig_perp,
    context_len,
    train_data,
    objective_set,
    max_steps,
    device,
    filename="dev.jbl",
    recopy_model=recopy_gpt2,
):
    """
    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
    contexts = []
    real_perps = []
    past_perps = []

    # Initialize the transformer model
    orig_model = copy.deepcopy(model)
    orig_model.to(device="cpu")
    torch.cuda.empty_cache()

    # Compute perplexity of initial transformer model for comparison
    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()
        # Do LM backprop
        torch.nn.utils.clip_grad_norm_(model.parameters(), 3.0)
        lm_optimizer.step()
        lm_scheduler.step()  # Update learning rate schedule

        # Compute perplexity after back-propagating on the selected context
        real_perp = compute_perplexity(model, objective_set, context_len)

        # Periodically save the stored (X, IG(X)) pairs
        if step % 1000 == 0 and step > 1:
            intermittent_save(contexts, real_perps, past_perps, filename)

        # Reset the pretrained model to the original pretrained GPT-2 weights after each iteration
        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)

    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

    
dpkd/transformers/examples/research_projects/information-gain-filtration/igf/igf.py