Update app.py

app.py

@@ -1,6 +1,9 @@
 # app.py
 import streamlit as st
-from models import demo_distill, demo_r1, demo_zero
+from models import load_model
+
+# Load the model once
+demo = load_model()
 
 # Page configuration
 st.set_page_config(
@@ -13,23 +16,10 @@ st.set_page_config(
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
-# Sidebar for model selection and parameters
+# Sidebar for model parameters
 with st.sidebar:
     st.header("Model Configuration")
     
-    # Model selection
-    model_mapping = {
-        "deepseek-ai/DeepSeek-R1-Distill-Qwen-32B": demo_distill,
-        "deepseek-ai/DeepSeek-R1": demo_r1,
-        "deepseek-ai/DeepSeek-R1-Zero": demo_zero
-    }
-    selected_model_name = st.selectbox(
-        "Choose Model",
-        list(model_mapping.keys()),
-        index=0
-    )
-    selected_demo = model_mapping[selected_model_name]
-    
     # System message
     system_message = st.text_area(
         "System Message",
@@ -64,7 +54,7 @@ with st.sidebar:
 
 # Main chat interface
 st.title("🤖 DeepSeek Chatbot")
-st.caption("Powered by ruslanmv.com - Choose your model and parameters in the sidebar")
+st.caption("Powered by ruslanmv.com - Configure parameters in the sidebar")
 
 # Display chat messages
 for message in st.session_state.messages:
@@ -81,12 +71,11 @@ if prompt := st.chat_input("Type your message..."):
         st.markdown(prompt)
     
     try:
-        # Generate response using selected model
+        # Generate response using the model
         with st.spinner("Generating response..."):
-            # The model expects parameters as separate arguments
-            response = selected_demo.fn(
-                prompt=f"{system_message}\n\nUser: {prompt}\nAssistant:",
-                max_length=max_tokens,  # Updated from max_new_tokens to max_length
+            response = demo(
+                inputs=f"{system_message}\n\nUser: {prompt}\nAssistant:",
+                max_length=max_tokens,  # Matches Gradio's expected parameter name
                 temperature=temperature,
                 top_p=top_p
             )