Update app.py

app.py

@@ -69,6 +69,18 @@ def generate_prompt(instruction, input=None):
         return PROMPT_DICT["prompt_no_input"].format(instruction=instruction)
 
 @spaces.GPU
+def generator(input_ids, generation_config, max_new_tokens):
+    # Without streaming
+    with torch.no_grad():
+        generation_output = model.generate(
+                input_ids=input_ids,
+                generation_config=generation_config,
+                return_dict_in_generate=True,
+                output_scores=False,
+                max_new_tokens=max_new_tokens,
+            )
+    return generation_output
+    
 def respond(
     message,
     history: list[tuple[str, str]],
@@ -89,16 +101,7 @@ def respond(
             num_beams=1,
             max_new_tokens = max_new_tokens
         )
-
-    # Without streaming
-    with torch.no_grad():
-        generation_output = model.generate(
-                input_ids=input_ids,
-                generation_config=generation_config,
-                return_dict_in_generate=True,
-                output_scores=False,
-                max_new_tokens=max_new_tokens,
-            )
+    generation_output = generator(input_ids, generation_config, max_new_tokens)
     s = generation_output.sequences[0]
     output = tokenizer.decode(s)
     response = output.split("Response:")[1].strip()