| 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, |
| ) |
|
|