Update README.md

README.md

@@ -7,16 +7,19 @@ pipeline_tag: text-generation
 
 # MAMBA (2.8B) 🐍 fine-tuned on H4/no_robots dataset for chat / instruction
 
-TBD
+Model Card is still WIP!
 
 ## Usage
 
 ```py
+import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 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")
+
 eos_token = "<|endoftext|>"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 tokenizer.eos_token = eos_token
@@ -24,7 +27,7 @@ tokenizer.pad_token = tokenizer.eos_token
 tokenizer.chat_template = AutoTokenizer.from_pretrained(CHAT_TEMPLATE_ID).chat_template
 
 model = MambaLMHeadModel.from_pretrained(
-        model_name, device="cuda", dtype=torch.float16)
+        model_name, device=device, dtype=torch.float16)
 
 history_dict: list[dict[str, str]] = []
 prompt = "Tell me 5 sites to visit in Spain"
@@ -32,7 +35,7 @@ history_dict.append(dict(role="user", content=prompt))
 
 input_ids = tokenizer.apply_chat_template(
             history_dict, return_tensors="pt", add_generation_prompt=True
-        ).to(device)
+).to(device)
 
 out = model.generate(
     input_ids=input_ids,