import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig, TextStreamer

device = 'cuda'
prompt = "Mihai Eminescu a fost "

model = AutoModelForCausalLM.from_pretrained('faur-ai/llmamba', trust_remote_code=True, torch_dtype=torch.bfloat16).to(device)
tokenizer = AutoTokenizer.from_pretrained('faur-ai/llmamba', trust_remote_code=True)

streamer = TextStreamer(tokenizer)

inputs = tokenizer.encode(
    prompt,
    add_special_tokens=False,
    return_tensors='pt',
).to(device)

outputs = model.generate(
    streamer=streamer,
    input_ids=inputs,
    temperature=0.8,
    do_sample=True,
    max_new_tokens=100,
)