Views
No views yet
1import torch
2from transformers import AutoTokenizer
3from transformers import Mamba2ForCausalLM
4
5
6if __name__ == "__main__":
7 device = "cuda"
8 model_id = "benchang1110/mamba2-370m-hf"
9 tokenizer = AutoTokenizer.from_pretrained(model_id)
10 model = Mamba2ForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map=device)
11 model.eval()
12
13 with torch.no_grad():
14 text = input("Input: ")
15 input_ids = tokenizer(text, return_tensors="pt").to(device)
16 output = model.generate(**input_ids, max_new_tokens=1024, do_sample=False)
17 print(tokenizer.decode(output[0], skip_special_tokens=True))