louisbrulenaudet commited on
Commit
5912f7b
1 Parent(s): 0c6a021

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +5 -2
app.py CHANGED
@@ -8,6 +8,7 @@
8
  # See the License for the specific language governing permissions and
9
  # limitations under the License.
10
 
 
11
  import logging
12
 
13
  from threading import Thread
@@ -137,13 +138,15 @@ else:
137
  model = AutoModelForCausalLM.from_pretrained(
138
  "CohereForAI/c4ai-command-r-v01-4bit",
139
  torch_dtype=torch.float16,
140
- device_map="auto"
 
141
  )
142
 
143
  device = model.device
144
 
145
  tokenizer = AutoTokenizer.from_pretrained(
146
- "CohereForAI/c4ai-command-r-v01-4bit"
 
147
  )
148
 
149
  tokenizer.use_default_system_prompt = False
 
8
  # See the License for the specific language governing permissions and
9
  # limitations under the License.
10
 
11
+ import os
12
  import logging
13
 
14
  from threading import Thread
 
138
  model = AutoModelForCausalLM.from_pretrained(
139
  "CohereForAI/c4ai-command-r-v01-4bit",
140
  torch_dtype=torch.float16,
141
+ device_map="auto",
142
+ token=os.getenv("HF_TOKEN")
143
  )
144
 
145
  device = model.device
146
 
147
  tokenizer = AutoTokenizer.from_pretrained(
148
+ "CohereForAI/c4ai-command-r-v01-4bit",
149
+ token=os.getenv("HF_TOKEN")
150
  )
151
 
152
  tokenizer.use_default_system_prompt = False