from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer

trained_model = AutoModelForCausalLM.from_pretrained(
    "culturalheritagenus/rumi-degarbler-v1",
    device_map="auto",
    torch_dtype=torch.bfloat16
)
trained_tokenizer = AutoTokenizer.from_pretrained("culturalheritagenus/rumi-degarbler-v1")

messages = [
    {"role": "user", "content": "You are a Malay language spelling corrector. I will give you some text written in messy Rumi (shortened or mistyped). Rewrite it in correct Malay Rumi spelling.\naurng ank. yngdim dimn anm aurngdan"},
]
inputs = tokenizer.apply_chat_template(
    messages,
    tokenize = True,
    add_generation_prompt = True, # Must add for generation
    return_tensors = "pt",
).to("cuda")


text_streamer = TextStreamer(tokenizer)
_ = trained_model.generate(input_ids = inputs, streamer = text_streamer, max_new_tokens = 128, use_cache = True)