Update app.py

app.py

@@ -1,9 +1,12 @@
+import spaces
+import gradio as gr
+
+import time
 from threading import Thread
 
 import torch
-import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
-import time
+
 
 model_id = "facebook/opt-6.7b"
 assistant_id = "facebook/opt-125m"
@@ -12,6 +15,8 @@ model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16
 assistant_model = AutoModelForCausalLM.from_pretrained(assistant_id).to(device=model.device, dtype=torch.bfloat16)
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 
+
+@spaces.GPU
 def run_generation(user_text, use_assistant, temperature, max_new_tokens):
     if temperature < 0.1:
         do_sample = False
@@ -45,7 +50,7 @@ def run_generation(user_text, use_assistant, temperature, max_new_tokens):
     for new_text in streamer:
         model_output += new_text
         yield [model_output, round(time.time() - start, 3)]
-    return [model_output, round(time.time() - start, 3)]
+    # return [model_output, round(time.time() - start, 3)]
 
 
 def reset_textbox():