Update app.py

app.py

@@ -1,29 +1,26 @@
 # app.py
 import streamlit as st
-from huggingface_hub import InferenceClient
-from datetime import datetime
+from models import demo
 
-# Configure page
+# Page configuration
 st.set_page_config(
     page_title="DeepSeek Chatbot - ruslanmv.com",
     page_icon="🤖",
-    layout="centered",
-    initial_sidebar_state="expanded"
+    layout="centered"
 )
 
-# Initialize session state
+# Initialize session state for chat history
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
-# Sidebar controls
+# Sidebar for model selection and parameters
 with st.sidebar:
-    st.title("🤖 Chatbot Settings")
-    st.markdown("Created by [ruslanmv.com](https://ruslanmv.com/)")
+    st.header("Model Configuration")
     
     # Model selection
     selected_model = st.selectbox(
         "Choose Model",
-        options=[
+        [
             "deepseek-ai/DeepSeek-R1-Distill-Qwen-32B",
             "deepseek-ai/DeepSeek-R1",
             "deepseek-ai/DeepSeek-R1-Zero"
@@ -39,19 +36,19 @@ with st.sidebar:
     )
     
     # Generation parameters
-    max_new_tokens = st.slider(
-        "Max new tokens",
+    max_tokens = st.slider(
+        "Max New Tokens",
         min_value=1,
         max_value=4000,
         value=512,
-        step=50
+        step=10
     )
     
     temperature = st.slider(
         "Temperature",
         min_value=0.1,
         max_value=4.0,
-        value=1.0,
+        value=0.7,
         step=0.1
     )
     
@@ -62,66 +59,44 @@ with st.sidebar:
         value=0.9,
         step=0.1
     )
-    
-    # Optional HF Token
-    hf_token = st.text_input(
-        "HuggingFace Token (optional)",
-        type="password",
-        help="Enter your HuggingFace token if required for model access"
-    )
 
 # Main chat interface
-st.title("💬 DeepSeek Chatbot")
-st.caption("🚀 A conversational AI powered by DeepSeek models")
+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"])
-        if "timestamp" in message:
-            st.caption(f"_{message['timestamp']}_")
 
-# Chat input and processing
+# Chat input
 if prompt := st.chat_input("Type your message..."):
-    # Add user message to history
-    st.session_state.messages.append({
-        "role": "user",
-        "content": prompt,
-        "timestamp": datetime.now().strftime("%H:%M:%S")
-    })
+    # 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)
-        st.caption(f"_{st.session_state.messages[-1]['timestamp']}_")
-
-    # Create full prompt with system message
-    full_prompt = f"{system_message}\n\nUser: {prompt}\nAssistant:"
     
-    # Create client and generate response
-    client = InferenceClient(model=selected_model, token=hf_token)
+    # Prepare full prompt with system message
+    full_prompt = f"{system_message}\n\nUser: {prompt}\nAssistant:"
     
-    # Display assistant response
-    with st.chat_message("assistant"):
-        response = st.write_stream(
-            client.text_generation(
+    try:
+        # Generate response using selected model
+        with st.spinner("Generating response..."):
+            response = demo.fn(
                 full_prompt,
-                max_new_tokens=max_new_tokens,
+                max_new_tokens=max_tokens,
                 temperature=temperature,
-                top_p=top_p,
-                stream=True
+                top_p=top_p
             )
-        )
-        timestamp = datetime.now().strftime("%H:%M:%S")
-        st.caption(f"_{timestamp}_")
-
-    # Add assistant response to history
-    st.session_state.messages.append({
-        "role": "assistant",
-        "content": response,
-        "timestamp": timestamp
-    })
-
-# Optional debug information
-# st.sidebar.markdown("---")
-# st.sidebar.json(st.session_state.messages)
+        
+        # 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)}")