Update app.py

app.py

@@ -1,7 +1,7 @@
 # app.py
 import streamlit as st
 from models import demo
-import gradio as gr
+
 # Page configuration
 st.set_page_config(
     page_title="DeepSeek Chatbot - ruslanmv.com",
@@ -16,40 +16,34 @@ if "messages" not in st.session_state:
 # Sidebar for model selection and parameters
 with st.sidebar:
     st.header("Model Configuration")
-
+    
     # Model selection
-    model_mapping = {
-        "deepseek-ai/DeepSeek-R1-Distill-Qwen-32B": gr.load(
-            name="deepseek-ai/DeepSeek-R1-Distill-Qwen-32B",
-            src="huggingface"
-        ),
-        "deepseek-ai/DeepSeek-R1": gr.load(name="deepseek-ai/DeepSeek-R1", src="huggingface"),
-        "deepseek-ai/DeepSeek-R1-Zero": gr.load(name="deepseek-ai/DeepSeek-R1-Zero", src="huggingface")
-    }
-
-    selected_model_name = st.selectbox(
+    selected_model = st.selectbox(
         "Choose Model",
-        list(model_mapping.keys()),
+        [
+            "deepseek-ai/DeepSeek-R1-Distill-Qwen-32B",
+            "deepseek-ai/DeepSeek-R1",
+            "deepseek-ai/DeepSeek-R1-Zero"
+        ],
         index=0
     )
-    selected_model = model_mapping[selected_model_name]
-
+    
     # System message
     system_message = st.text_area(
         "System Message",
         value="You are a friendly Chatbot created by ruslanmv.com",
         height=100
     )
-
+    
     # Generation parameters
-    max_new_tokens = st.slider(
+    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,
@@ -57,7 +51,7 @@ with st.sidebar:
         value=0.7,
         step=0.1
     )
-
+    
     top_p = st.slider(
         "Top-p (nucleus sampling)",
         min_value=0.1,
@@ -79,34 +73,28 @@ for message in st.session_state.messages:
 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 conversation history in the required format
-    full_prompt = f"{system_message}\n\nUser: {prompt}\nAssistant:"
-
+    
     try:
         # Generate response using selected model
         with st.spinner("Generating response..."):
-            # The model expects a single text input with the conversation history
-            response = selected_model.fn(
-                full_prompt,  # Pass the full conversation history
-                {
-                    "temperature": temperature,
-                    "top_p": top_p,
-                    "max_new_tokens": max_new_tokens,
-                    "repetition_penalty": 1.0
-                }
+            # The model expects parameters as separate arguments
+            response = demo.fn(
+                f"{system_message}\n\nUser: {prompt}\nAssistant:",  # Full prompt
+                max_new_tokens=max_tokens,
+                temperature=temperature,
+                top_p=top_p
             )
-
+        
         # 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)}")
+        st.error(f"Error generating response: {str(e)}")