def compute_perplexity(model, test_data, context_len):
    """
    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

    """

    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)
        # pad
        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

    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

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