Update app.py

app.py

@@ -1,102 +1,118 @@
 # app.py
+
 import streamlit as st
-from models import demo
+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
 
-# Page configuration
-st.set_page_config(
-    page_title="DeepSeek Chatbot - ruslanmv.com",
-    page_icon="🤖",
-    layout="centered"
+# 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"]
 )
 
-# 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 New 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..."):
-            response = demo.fn(
-                full_prompt,
-                max_new_tokens=max_tokens,
-                temperature=temperature,
-                top_p=top_p
-            )
+# 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
+)
+
+# Optional parameter: temperature
+temperature = st.slider(
+    "Temperature",
+    min_value=0.10,
+    max_value=4.00,
+    value=0.80,
+    step=0.05
+)
+
+# 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)
         
-        # Display assistant response
-        with st.chat_message("assistant"):
-            st.markdown(response)
+        # 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
+            )
+        except Exception as e:
+            response = f"Error: {e}"
         
-        # 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)}")
+        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}")
+