import torch
from transformers import AutoConfig, AutoModelForMaskedLM


class Model(torch.nn.Module):
    def __init__(self, model_name, config):
        super().__init__()
        self.model_name = model_name
        self.config = config
        self.model = AutoModelForMaskedLM.from_pretrained(self.model_name, config=self.config)

    def forward(self, x):
        return self.model(x).logits


model_name = "facebook/esm2_t12_35M_UR50D"
config = AutoConfig.from_pretrained(model_name)
vocab_size = config.vocab_size
sequence_length = 1022
batch_size = 4


def get_inputs():
    p = popcorn_pri
    bs = p.jitter_int(batch_size)
    sl = p.jitter_int(sequence_length)
    inputs = torch.randint(0, vocab_size, (bs, sl))
    return [inputs]



def get_init_inputs():
    return [model_name, config]