Update README.md

README.md

@@ -48,9 +48,7 @@ No Robots is a high-quality dataset of 10,000 instructions and demonstrations cr
 ## 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
@@ -60,7 +58,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-chat-no_robots"
 
 eos_token = "<|endoftext|>"
@@ -72,12 +70,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(