Update README.md
Browse files
README.md
CHANGED
@@ -43,9 +43,7 @@ The base dataset mix is identical to the original Nous-Hermes', minus the Nous-I
|
|
43 |
## Usage
|
44 |
|
45 |
```sh
|
46 |
-
pip install transformers
|
47 |
-
pip install causal-conv1d<=1.0.2
|
48 |
-
pip install mamba-ssm
|
49 |
```
|
50 |
|
51 |
```py
|
@@ -55,7 +53,7 @@ from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
|
|
55 |
|
56 |
CHAT_TEMPLATE_ID = "HuggingFaceH4/zephyr-7b-beta"
|
57 |
|
58 |
-
device =
|
59 |
model_name = "clibrain/mamba-2.8b-instruct-openhermes"
|
60 |
|
61 |
eos_token = "<|endoftext|>"
|
@@ -67,12 +65,12 @@ tokenizer.chat_template = AutoTokenizer.from_pretrained(CHAT_TEMPLATE_ID).chat_t
|
|
67 |
model = MambaLMHeadModel.from_pretrained(
|
68 |
model_name, device=device, dtype=torch.float16)
|
69 |
|
70 |
-
|
71 |
prompt = "Tell me 5 sites to visit in Spain"
|
72 |
-
|
73 |
|
74 |
input_ids = tokenizer.apply_chat_template(
|
75 |
-
|
76 |
).to(device)
|
77 |
|
78 |
out = model.generate(
|
|
|
43 |
## Usage
|
44 |
|
45 |
```sh
|
46 |
+
pip install torch==2.1.0 transformers==4.35.0 causal-conv1d==1.0.0 mamba-ssm==1.0.1
|
|
|
|
|
47 |
```
|
48 |
|
49 |
```py
|
|
|
53 |
|
54 |
CHAT_TEMPLATE_ID = "HuggingFaceH4/zephyr-7b-beta"
|
55 |
|
56 |
+
device = "cuda:0" if torch.cuda.is_available() else "cpu"
|
57 |
model_name = "clibrain/mamba-2.8b-instruct-openhermes"
|
58 |
|
59 |
eos_token = "<|endoftext|>"
|
|
|
65 |
model = MambaLMHeadModel.from_pretrained(
|
66 |
model_name, device=device, dtype=torch.float16)
|
67 |
|
68 |
+
messages = []
|
69 |
prompt = "Tell me 5 sites to visit in Spain"
|
70 |
+
messages.append(dict(role="user", content=prompt))
|
71 |
|
72 |
input_ids = tokenizer.apply_chat_template(
|
73 |
+
messages, return_tensors="pt", add_generation_prompt=True
|
74 |
).to(device)
|
75 |
|
76 |
out = model.generate(
|