Update app.py

app.py

@@ -16,34 +16,40 @@ if "messages" not in st.session_state:
 # Sidebar for model selection and parameters
 with st.sidebar:
     st.header("Model Configuration")
-    
+
     # Model selection
-    selected_model = st.selectbox(
+    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(
         "Choose Model",
-        [
-            "deepseek-ai/DeepSeek-R1-Distill-Qwen-32B",
-            "deepseek-ai/DeepSeek-R1",
-            "deepseek-ai/DeepSeek-R1-Zero"
-        ],
+        list(model_mapping.keys()),
         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_tokens = st.slider(
-        "Max Tokens",
+    max_new_tokens = st.slider(
+        "Max New Tokens",
         min_value=1,
         max_value=4000,
         value=512,
         step=10
     )
-    
+
     temperature = st.slider(
         "Temperature",
         min_value=0.1,
@@ -51,7 +57,7 @@ with st.sidebar:
         value=0.7,
         step=0.1
     )
-    
+
     top_p = st.slider(
         "Top-p (nucleus sampling)",
         min_value=0.1,
@@ -73,35 +79,33 @@ 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 = demo.fn(
+            response = selected_model.fn(
                 full_prompt,  # Pass the full conversation history
-                # Parameters need to be passed as a dictionary in the second argument
                 {
                     "temperature": temperature,
                     "top_p": top_p,
-                    "max_new_tokens": max_tokens,
+                    "max_new_tokens": max_new_tokens,
                     "repetition_penalty": 1.0
                 }
             )
-        
+
         # 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)}")
+
+    except Exception as e: