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",
):
    """
    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

    """

    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 = []

    # Compute the performance of the transformer model at the beginning
    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))

                # 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 global_step == 10:
                    threshold = -1
                if predicted_q < threshold:
                    do_backprop = False

            # If we passed the filter, add the context to the batch!
            if do_backprop:
                contexts.append(np.array(context.cpu()))
                lm_loss = outputs[0]
                lm_loss.backward()
                examples += 1

            del outputs

            # Once the batch is filled with enough contexts, backprop on the batch.
            if examples == batch_size:
                torch.cuda.empty_cache()
                examples = 0
                # Do LM backprop
                torch.nn.utils.clip_grad_norm_(model.parameters(), 3.0)
                lm_optimizer.step()
                lm_scheduler.step()  # Update learning rate schedule
                global_step += 1
                # Compute the performance of the transformer model at this batch
                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)
            # Break out of the loop after 60 batches
            if max_steps > 0 and global_step > 60:
                break
        if max_steps > 0 and global_step > 60:
            break

    # save finetuned transformer model
    torch.save(model.state_dict(), finetuned_model_name)
    torch.cuda.empty_cache()
    # Do some cleaning up so we can reinitialize for the next run of this function
    del lm_optimizer
    del lm_scheduler
    return model

    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

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