joaogante HF Staff commited on
Commit
1702ae3
·
verified ·
1 Parent(s): 5c3f286

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -0
app.py CHANGED
@@ -1,4 +1,5 @@
1
  from threading import Thread
 
2
 
3
  import torch
4
  import gradio as gr
@@ -12,6 +13,7 @@ model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16
12
  assistant_model = AutoModelForCausalLM.from_pretrained(assistant_id).to(device=model.device, dtype=torch.bfloat16)
13
  tokenizer = AutoTokenizer.from_pretrained(model_id)
14
 
 
15
  def run_generation(user_text, use_assistant, temperature, max_new_tokens):
16
  if temperature < 0.1:
17
  do_sample = False
 
1
  from threading import Thread
2
+ import spaces
3
 
4
  import torch
5
  import gradio as gr
 
13
  assistant_model = AutoModelForCausalLM.from_pretrained(assistant_id).to(device=model.device, dtype=torch.bfloat16)
14
  tokenizer = AutoTokenizer.from_pretrained(model_id)
15
 
16
+ @spaces.GPU
17
  def run_generation(user_text, use_assistant, temperature, max_new_tokens):
18
  if temperature < 0.1:
19
  do_sample = False