ybelkada commited on
Commit
f8d7a7f
·
verified ·
1 Parent(s): addd751

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -1
app.py CHANGED
@@ -1,4 +1,5 @@
1
  import torch
 
2
 
3
  from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
4
  import gradio as gr
@@ -28,10 +29,11 @@ END_MESSAGE = """
28
 
29
  device = "cuda" # for GPU usage or "cpu" for CPU usage
30
 
31
- tokenizer = AutoTokenizer.from_pretrained(MODEL)
32
  model = AutoModelForCausalLM.from_pretrained(
33
  MODEL,
34
  torch_dtype=torch.bfloat16,
 
35
  ).to(device)
36
 
37
  model = torch.compile(model)
 
1
  import torch
2
+ import os
3
 
4
  from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
5
  import gradio as gr
 
29
 
30
  device = "cuda" # for GPU usage or "cpu" for CPU usage
31
 
32
+ tokenizer = AutoTokenizer.from_pretrained(MODEL, token=os.getenv("HF_TOKEN"))
33
  model = AutoModelForCausalLM.from_pretrained(
34
  MODEL,
35
  torch_dtype=torch.bfloat16,
36
+ token=os.getenv("HF_TOKEN")
37
  ).to(device)
38
 
39
  model = torch.compile(model)