llmamba / inference.py
Alexandru Gherghescu
Add inference example
17be2cb
Raw
History Blame Contribute Delete
647 Bytes
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,
)