import os
import torch
import hidet
from transformers import AutoTokenizer, RobertaForSequenceClassification
from transformers import logging
import hidet

# This line causes compilation to fail
hidet.option.imperative(False)
# hidet.torch.dynamo_config.use_cuda_graph(False)

os.environ["TOKENIZERS_PARALLELISM"] = "false"
logging.set_verbosity_error()


if __name__ == "__main__":
    
    tokenizer = AutoTokenizer.from_pretrained("roberta-base")
    model = RobertaForSequenceClassification.from_pretrained("roberta-base", max_position_embeddings=8192, ignore_mismatched_sizes=True)
    model = model.eval().cuda()

    inputs = tokenizer("Hello, my dog is cute", padding='max_length', max_length=4096, return_tensors="pt")
    inputs = {'input_ids': inputs['input_ids']}

    with torch.no_grad(), torch.autocast("cuda", enabled=True):
        model = torch.compile(model, backend="hidet")

        torch_inputs = tuple(i.clone().cuda() for i in inputs.values())
        torch_out = model(*torch_inputs)