Update app.py

app.py

@@ -92,48 +92,28 @@ def generate_response(model_id, conversation, user_message, max_length=512, temp
     
     return response, load_time, generation_time
 
-def chatbot(conversation, user_message, model_name, max_length=512, temperature=0.7):
-    if not user_message.strip():
-        return conversation, "", "Please enter a message"
-    
-    if conversation is None:
-        conversation = []
+def respond(message, history, model_name, max_length, temperature):
+    """Main function for ChatInterface - simplified signature"""
+    if not message.strip():
+        return "Please enter a message"
     
     # Get model ID
     model_id = MODELS.get(model_name, MODELS["Athena-R3X 8B"])
     
     try:
-        # Add user message to conversation
-        conversation.append([user_message, ""])
-        
         # Generate response using ZeroGPU
         response, load_time, generation_time = generate_response(
-            model_id, conversation[:-1], user_message, max_length, temperature
+            model_id, history, message, max_length, temperature
         )
         
-        # Update the conversation with the response
-        conversation[-1][1] = response
-        
-        stats = f"⚡ Load: {load_time:.1f}s | Gen: {generation_time:.1f}s | Model: {model_name}"
-        
-        return conversation, "", stats
+        # ChatInterface expects a generator for streaming or just return the response
+        return response
         
     except Exception as e:
-        error_msg = f"Error: {str(e)}"
-        if conversation:
-            conversation[-1][1] = error_msg
-        else:
-            conversation = [[user_message, error_msg]]
-        return conversation, "", f"❌ Error: {str(e)}"
-
-def clear_chat():
-    return [], "", ""
+        return f"Error: {str(e)}"
 
 # CSS for better styling
 css = """
-#chatbot {
-    height: 600px;
-}
 .message {
     padding: 10px;
     margin: 5px;
@@ -143,79 +123,56 @@ css = """
 
 theme = gr.themes.Monochrome()
 
-# Create Gradio interface
+# Create ChatInterface
 with gr.Blocks(title="Athena Playground Chat", css=css, theme=theme) as demo:
     gr.Markdown("# 🚀 Athena Playground Chat")
     gr.Markdown("*Powered by HuggingFace ZeroGPU*")
     
-    # Main chat interface
-    chat_history = gr.Chatbot(
-        elem_id="chatbot",
-        show_label=False,
-        show_share_button=False,
-        container=False
-    )
-    
-    user_input = gr.Textbox(
-        placeholder="Ask Athena anything...",
-        label="Your message",
-        lines=2,
-        max_lines=10
-    )
-    
+    # Additional inputs for configuration
     with gr.Row():
-        submit_btn = gr.Button("📤 Send", variant="primary")
-        clear_btn = gr.Button("🗑️ Clear Chat", variant="secondary")
-    
-    stats_output = gr.Textbox(
-        label="Stats", 
-        interactive=False,
-        show_label=False,
-        placeholder="Stats will appear here..."
-    )
-    
-    # Configuration settings at the bottom
-    gr.Markdown("---")
-    gr.Markdown("## ⚙️ Configuration")
-    
-    with gr.Row():
-        with gr.Column():
+        with gr.Column(scale=1):
             model_choice = gr.Dropdown(
                 label="📱 Model",
                 choices=list(MODELS.keys()),
                 value="Athena-R3X 8B",
                 info="Select which Athena model to use"
             )
-        with gr.Column():
+        with gr.Column(scale=1):
             max_length = gr.Slider(
                 32, 2048, value=512, 
                 label="📝 Max Tokens",
                 info="Maximum number of tokens to generate"
             )
-        with gr.Column():
+        with gr.Column(scale=1):
             temperature = gr.Slider(
                 0.1, 2.0, value=0.7, 
                 label="🎨 Creativity",
                 info="Higher values = more creative responses"
             )
     
-    # Event handlers
-    submit_btn.click(
-        chatbot,
-        inputs=[chat_history, user_input, model_choice, max_length, temperature],
-        outputs=[chat_history, user_input, stats_output]
-    )
-    
-    user_input.submit(
-        chatbot,
-        inputs=[chat_history, user_input, model_choice, max_length, temperature],
-        outputs=[chat_history, user_input, stats_output]
-    )
-    
-    clear_btn.click(
-        clear_chat,
-        inputs=[],
-        outputs=[chat_history, user_input, stats_output]
+    # Create the ChatInterface
+    chat_interface = gr.ChatInterface(
+        fn=respond,
+        additional_inputs=[model_choice, max_length, temperature],
+        title="Chat with Athena",
+        description="Ask Athena anything!",
+        theme="soft",
+        examples=[
+            "Hello! How are you?",
+            "What can you help me with?",
+            "Tell me about artificial intelligence",
+            "Write a short poem about space"
+        ],
+        cache_examples=False,
+        retry_btn=None,
+        undo_btn="↩️ Undo",
+        clear_btn="🗑️ Clear",
+        submit_btn="📤 Send",
+        chatbot=gr.Chatbot(
+            height=500,
+            placeholder="Start chatting with Athena...",
+            show_share_button=False
+        )
     )
 
 if __name__ == "__main__":