Update app.py

app.py

@@ -2,22 +2,41 @@ import gradio as gr
 import spaces
 from transformers import pipeline
 
-@spaces.GPU(duration=120)  # Increased timeout for LLM inference
-def generate_response(messages):
-    pipe = pipeline("text-generation", 
-                    model="unsloth/DeepSeek-R1-Distill-Llama-8B",
-                    torch_dtype="auto",
-                    device_map="auto")
+# Initialize model once at startup
+model = pipeline(
+    "text-generation",
+    model="unsloth/DeepSeek-R1-Distill-Llama-8B",
+    torch_dtype="auto",
+    device_map="auto"
+)
+
+@spaces.GPU(duration=120)
+def chat_response(message, history):
+    # Format conversation history for model input
+    messages = []
+    for human, assistant in history:
+        messages.extend([
+            {"role": "user", "content": human},
+            {"role": "assistant", "content": assistant}
+        ])
+    messages.append({"role": "user", "content": message})
+    
+    # Generate response
+    response = model(
+        messages,
+        max_new_tokens=256,
+        temperature=0.7,
+        do_sample=True
+    )
     
-    return pipe(messages, 
-               max_new_tokens=256,
-               temperature=0.7,
-               do_sample=True)
+    return response[0]['generated_text'][-1]["content"]
 
-demo = gr.Interface(
-    fn=generate_response,
-    inputs=gr.JSON(value=[{"role": "user", "content": "Who are you?"}]),
-    outputs=gr.Text(),
-    title="DeepSeek-Llama-8B ZeroGPU Demo"
+# Create chat interface
+demo = gr.ChatInterface(
+    chat_response,
+    chatbot=gr.Chatbot(height=500),
+    textbox=gr.Textbox(placeholder="Ask me anything...", container=False, scale=7),
+    title="DeepSeek-Llama-8B Chat Demo",
+    examples=[["Explain quantum computing simply"], ["Write a Python function for Fibonacci sequence"]]
 )
 demo.launch()