Views
No views yet

pip install torch==2.1.0 transformers==4.35.0 causal-conv1d==1.0.0 mamba-ssm==1.0.11import torch
2from transformers import AutoTokenizer, AutoModelForCausalLM
3from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
4
5CHAT_TEMPLATE_ID = "HuggingFaceH4/zephyr-7b-beta"
6
7device = "cuda:0" if torch.cuda.is_available() else "cpu"
8model_name = "clibrain/mamba-2.8b-instruct-openhermes"
9
10eos_token = "<|endoftext|>"
11tokenizer = AutoTokenizer.from_pretrained(model_name)
12tokenizer.eos_token = eos_token
13tokenizer.pad_token = tokenizer.eos_token
14tokenizer.chat_template = AutoTokenizer.from_pretrained(CHAT_TEMPLATE_ID).chat_template
15
16model = MambaLMHeadModel.from_pretrained(
17 model_name, device=device, dtype=torch.float16)
18
19messages = []
20prompt = "Tell me 5 sites to visit in Spain"
21messages.append(dict(role="user", content=prompt))
22
23input_ids = tokenizer.apply_chat_template(
24 messages, return_tensors="pt", add_generation_prompt=True
25).to(device)
26
27out = model.generate(
28 input_ids=input_ids,
29 max_length=2000,
30 temperature=0.9,
31 top_p=0.7,
32 eos_token_id=tokenizer.eos_token_id,
33)
34
35decoded = tokenizer.batch_decode(out)
36assistant_message = (
37 decoded[0].split("<|assistant|>\n")[-1].replace(eos_token, "")
38)
39
40print(assistant_message)1git clone https://github.com/mrm8488/mamba-chat.git
2cd mamba-chat
3
4pip install -r requirements.txt
5pip install -q gradio==4.8.0
6
7python app.py \
8--model clibrain/mamba-2.8b-instruct-openhermes \
9--share