mrm8488 commited on
Commit
605742c
1 Parent(s): 4a70d08

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +5 -7
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 = torch.device("cuda" if torch.cuda.is_available() else "cpu")
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
- history_dict: list[dict[str, str]] = []
71
  prompt = "Tell me 5 sites to visit in Spain"
72
- history_dict.append(dict(role="user", content=prompt))
73
 
74
  input_ids = tokenizer.apply_chat_template(
75
- history_dict, return_tensors="pt", add_generation_prompt=True
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(