Update app.py

app.py

@@ -1,6 +1,6 @@
 # app.py
 import streamlit as st
-from models import demo
+from models import demo_distill, demo_r1, demo_zero
 
 # Page configuration
 st.set_page_config(
@@ -18,15 +18,17 @@ with st.sidebar:
     st.header("Model Configuration")
     
     # Model selection
-    selected_model = st.selectbox(
+    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",
-        [
-            "deepseek-ai/DeepSeek-R1-Distill-Qwen-32B",
-            "deepseek-ai/DeepSeek-R1",
-            "deepseek-ai/DeepSeek-R1-Zero"
-        ],
+        list(model_mapping.keys()),
         index=0
     )
+    selected_demo = model_mapping[selected_model_name]
     
     # System message
     system_message = st.text_area(
@@ -37,7 +39,7 @@ with st.sidebar:
     
     # Generation parameters
     max_tokens = st.slider(
-        "Max New Tokens",
+        "Max Tokens",
         min_value=1,
         max_value=4000,
         value=512,
@@ -82,9 +84,9 @@ if prompt := st.chat_input("Type your message..."):
         # Generate response using selected model
         with st.spinner("Generating response..."):
             # The model expects parameters as separate arguments
-            response = demo.fn(
-                f"{system_message}\n\nUser: {prompt}\nAssistant:",  # Full prompt
-                max_new_tokens=max_tokens,
+            response = selected_demo.fn(
+                prompt=f"{system_message}\n\nUser: {prompt}\nAssistant:",
+                max_length=max_tokens,  # Updated from max_new_tokens to max_length
                 temperature=temperature,
                 top_p=top_p
             )
@@ -97,4 +99,4 @@ if prompt := st.chat_input("Type your message..."):
         st.session_state.messages.append({"role": "assistant", "content": response})
     
     except Exception as e:
-        st.error(f"Error generating response: {str(e)}")
+        st.error(f"An error occurred: {str(e)}")