Update app.py

app.py

@@ -1,5 +1,7 @@
 import gradio as gr
+import spaces
 import torch
+
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 
@@ -26,49 +28,58 @@ model.eval()
 print("MODEL READY")
 
 
-def chat(message, history):
+@spaces.GPU
+def generate(message, history):
 
     prompt = ""
 
     for user, assistant in history:
         prompt += f"User: {user}\nAssistant: {assistant}\n"
 
+
     prompt += f"User: {message}\nAssistant:"
 
 
     inputs = tokenizer(
         prompt,
         return_tensors="pt"
-    ).to(model.device)
+    )
+
+
+    inputs = {
+        k: v.to(model.device)
+        for k, v in inputs.items()
+    }
 
 
     with torch.no_grad():
 
-        result = model.generate(
+        output = model.generate(
             **inputs,
             max_new_tokens=256,
             temperature=0.7,
+            top_p=0.9,
             do_sample=True
         )
 
 
-    text = tokenizer.decode(
-        result[0],
+    result = tokenizer.decode(
+        output[0],
         skip_special_tokens=True
     )
 
 
-    return text.split("Assistant:")[-1]
+    return result.split("Assistant:")[-1]
 
 
 
 demo = gr.ChatInterface(
-    fn=chat,
-    title="Rezaeian StatsAI"
+    fn=generate,
+    title="Rezaeian StatsAI",
+    description="ZeroGPU AI Assistant"
 )
 
 
 demo.launch(
-    server_name="0.0.0.0",
-    server_port=7860
+    server_name="0.0.0.0"
 )