import torch
import transformers
from transformers import GenerationConfig
from transformers import AutoConfig

def generate_inputs_for_model(model_cls, model):
    eval_context = torch.randint(0, model.config.vocab_size, (4, 2048)).to("cuda")
    return {"input_ids": eval_context}

config = AutoConfig.from_pretrained("t5-small")
model_cls = getattr(transformers, "AutoModelForSeq2SeqLM")
model = model_cls.from_config(config).to("cuda")
example_inputs = generate_inputs_for_model(model_cls, model)
example_inputs = (example_inputs["input_ids"],)

generation_config = GenerationConfig(
    max_new_tokens=256,
    pad_token_id=0,
    eos_token_id=None,
    do_sample=False,
    num_beams=1,
    use_cache=True,
)

class GenerationWrapper(torch.nn.Module):
    def __init__(self, model, generation_config):
        super().__init__()
        self.model = model
        self.generation_config = generation_config

    def forward(self, inputs):
        return self.model.generate(inputs, self.generation_config)

model = GenerationWrapper(model, generation_config)
# torch.compile repro
model_opt = torch.compile(model)
output = model_opt(*example_inputs)
# torch.export repro
torch.export.export(model, args=example_inputs, strict=False)