Update README.md

README.md

@@ -43,9 +43,7 @@ The base dataset mix is identical to the original Nous-Hermes', minus the Nous-I
 ## Usage
 
 ```sh
-pip install transformers
-pip install causal-conv1d<=1.0.2
-pip install mamba-ssm
+pip install torch==2.1.0 transformers==4.35.0 causal-conv1d==1.0.0 mamba-ssm==1.0.1
 ```
 
 ```py
@@ -55,7 +53,7 @@ from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
 
 CHAT_TEMPLATE_ID = "HuggingFaceH4/zephyr-7b-beta"
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
 model_name = "clibrain/mamba-2.8b-instruct-openhermes"
 
 eos_token = "<|endoftext|>"
@@ -67,12 +65,12 @@ tokenizer.chat_template = AutoTokenizer.from_pretrained(CHAT_TEMPLATE_ID).chat_t
 model = MambaLMHeadModel.from_pretrained(
  model_name, device=device, dtype=torch.float16)
 
-history_dict: list[dict[str, str]] = []
+messages = []
 prompt = "Tell me 5 sites to visit in Spain"
-history_dict.append(dict(role="user", content=prompt))
+messages.append(dict(role="user", content=prompt))
 
 input_ids = tokenizer.apply_chat_template(
- history_dict, return_tensors="pt", add_generation_prompt=True
+ messages, return_tensors="pt", add_generation_prompt=True
 ).to(device)
 
 out = model.generate(