Update app.py

app.py

@@ -1,118 +1,103 @@
 # app.py
-
 import streamlit as st
-from models import demo_qwen, demo_r1, demo_zero
-
-st.set_page_config(page_title="DeepSeek Chatbot", layout="centered")
-
-# A helper function to pick the correct Gradio interface
-def select_demo(model_name):
-    if model_name == "DeepSeek-R1-Distill-Qwen-32B":
-        return demo_qwen
-    elif model_name == "DeepSeek-R1":
-        return demo_r1
-    elif model_name == "DeepSeek-R1-Zero":
-        return demo_zero
-    else:
-        return demo_qwen  # default fallback
-
-# Title of the Streamlit app
-st.title("DeepSeek Chatbot")
-
-# Sidebar or main area for parameter selection
-st.subheader("Model and Parameters")
-
-model_name = st.selectbox(
-    "Select Model",
-    ["DeepSeek-R1-Distill-Qwen-32B", "DeepSeek-R1", "DeepSeek-R1-Zero"]
-)
-
-# Optional parameter: System message
-system_message = st.text_area(
-    "System Message",
-    value="You are a friendly Chatbot created by ruslanmv.com",
-    height=80
-)
-
-# Optional parameter: max new tokens
-max_new_tokens = st.slider(
-    "Max new tokens",
-    min_value=1,
-    max_value=4000,
-    value=512,
-    step=1
-)
+from models import demo
 
-# Optional parameter: temperature
-temperature = st.slider(
-    "Temperature",
-    min_value=0.10,
-    max_value=4.00,
-    value=0.80,
-    step=0.05
+# Page configuration
+st.set_page_config(
+    page_title="DeepSeek Chatbot - ruslanmv.com",
+    page_icon="🤖",
+    layout="centered"
 )
 
-# Optional parameter: top-p
-top_p = st.slider(
-    "Top-p (nucleus sampling)",
-    min_value=0.10,
-    max_value=1.00,
-    value=0.90,
-    step=0.05
-)
-
-# A text area for user input
-st.subheader("Chat")
-user_prompt = st.text_area("Your message:", value="", height=100)
-
-if "chat_history" not in st.session_state:
-    st.session_state["chat_history"] = []
-
-# Button to send user prompt
-if st.button("Send"):
-    if user_prompt.strip():
-        # Retrieve the correct Gradio demo
-        demo = select_demo(model_name)
-        
-        # Here we assume the Gradio interface has a function signature for `.predict()` 
-        # that accepts text plus generation parameters in some order. 
-        # Many huggingface-style Gradio demos simply take a single text prompt, 
-        # but it depends entirely on how `demo` is defined in your Gradio code.
-        # 
-        # If the interface has multiple inputs in a specific order, you might do something like:
-        # response = demo.predict(system_message, user_prompt, max_new_tokens, temperature, top_p)
-        #
-        # Or if it only expects a single string, you might combine them:
-        # combined_prompt = f"System: {system_message}\nUser: {user_prompt}"
-        # response = demo.predict(combined_prompt, max_new_tokens, temperature, top_p)
-        #
-        # The exact call depends on your Gradio block's input signature.
-
-        # For illustrative purposes, let's assume a simple signature:
-        #   demo.predict(prompt: str, max_new_tokens: int, temperature: float, top_p: float)
-        # and we inject the system message on top:
-
-        combined_prompt = f"{system_message}\n\nUser: {user_prompt}"
-        
-        try:
-            response = demo.predict(
-                combined_prompt,
-                max_new_tokens,
-                temperature,
-                top_p
+# Initialize session state for chat history
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+# Sidebar for model selection and parameters
+with st.sidebar:
+    st.header("Model Configuration")
+    
+    # Model selection
+    selected_model = st.selectbox(
+        "Choose Model",
+        [
+            "deepseek-ai/DeepSeek-R1-Distill-Qwen-32B",
+            "deepseek-ai/DeepSeek-R1",
+            "deepseek-ai/DeepSeek-R1-Zero"
+        ],
+        index=0
+    )
+    
+    # System message
+    system_message = st.text_area(
+        "System Message",
+        value="You are a friendly Chatbot created by ruslanmv.com",
+        height=100
+    )
+    
+    # Generation parameters
+    max_tokens = st.slider(
+        "Max Tokens",
+        min_value=1,
+        max_value=4000,
+        value=512,
+        step=10
+    )
+    
+    temperature = st.slider(
+        "Temperature",
+        min_value=0.1,
+        max_value=4.0,
+        value=0.7,
+        step=0.1
+    )
+    
+    top_p = st.slider(
+        "Top-p (nucleus sampling)",
+        min_value=0.1,
+        max_value=1.0,
+        value=0.9,
+        step=0.1
+    )
+
+# Main chat interface
+st.title("🤖 DeepSeek Chatbot")
+st.caption("Powered by ruslanmv.com - Choose your model and parameters in the sidebar")
+
+# Display chat messages
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
+
+# Chat input
+if prompt := st.chat_input("Type your message..."):
+    # Add user message to chat history
+    st.session_state.messages.append({"role": "user", "content": prompt})
+    
+    # Display user message
+    with st.chat_message("user"):
+        st.markdown(prompt)
+    
+    # Prepare full prompt with system message
+    full_prompt = f"{system_message}\n\nUser: {prompt}\nAssistant:"
+    
+    try:
+        # Generate response using selected model
+        with st.spinner("Generating response..."):
+            # Updated parameter names to match model expectations
+            response = demo.fn(
+                full_prompt,
+                max_length=max_tokens,  # Changed from max_new_tokens
+                temperature=temperature,
+                top_p=top_p
             )
-        except Exception as e:
-            response = f"Error: {e}"
         
-        st.session_state["chat_history"].append(("User", user_prompt))
-        st.session_state["chat_history"].append(("Assistant", response))
-        st.experimental_rerun()
-
-# Display conversation
-if st.session_state["chat_history"]:
-    for role, text in st.session_state["chat_history"]:
-        if role == "User":
-            st.markdown(f"**{role}:** {text}")
-        else:
-            st.markdown(f"**{role}:** {text}")
-
+        # Display assistant response
+        with st.chat_message("assistant"):
+            st.markdown(response)
+        
+        # Add assistant response to chat history
+        st.session_state.messages.append({"role": "assistant", "content": response})
+    
+    except Exception as e:
+        st.error(f"Error generating response: {str(e)}")