Update app.py

app.py

@@ -1,51 +1,71 @@
 import streamlit as st
-from functools import lru_cache
-import requests
+from models import demo  # Import the demo object from models.py
 
-# Cache model loading to optimize performance
-@lru_cache(maxsize=3)
-def load_hf_model(model_name):
-    # Use the Hugging Face Inference API directly
-    api_url = f"https://api-inference.huggingface.co/models/deepseek-ai/{model_name}"
-    return api_url
+# --- Streamlit App Configuration ---
+st.set_page_config(
+    page_title="DeepSeek Chatbot",
+    page_icon="🤖",
+    layout="wide"
+)
 
-# Load all models at startup
-MODELS = {
-    "DeepSeek-R1-Distill-Qwen-32B": load_hf_model("DeepSeek-R1-Distill-Qwen-32B"),
-    "DeepSeek-R1": load_hf_model("DeepSeek-R1"),
-    "DeepSeek-R1-Zero": load_hf_model("DeepSeek-R1-Zero")
-}
+# --- App Title and Description ---
+st.title("DeepSeek Chatbot")
+st.markdown("""
+    Created by [ruslanmv.com](https://ruslanmv.com/)
+    This is a demo of different DeepSeek models. Select a model, type your message, and click "Submit".
+    You can also adjust optional parameters like system message, max new tokens, temperature, and top-p.
+""")
 
-# --- Chatbot function ---
-def chatbot(input_text, history, model_choice, system_message, max_new_tokens, temperature, top_p):
-    history = history or []
-    
-    # Get the selected model API URL
-    api_url = MODELS[model_choice]
+# --- Sidebar for Model Selection and Parameters ---
+with st.sidebar:
+    st.header("Options")
+    model_choice = st.radio(
+        "Choose a Model",
+        options=["DeepSeek-R1-Distill-Qwen-32B", "DeepSeek-R1", "DeepSeek-R1-Zero"],
+        index=1  # Default to "DeepSeek-R1"
+    )
     
+    with st.expander("Optional Parameters", expanded=False):
+        system_message = st.text_area(
+            "System Message",
+            value="You are a friendly Chatbot created by ruslanmv.com",
+            height=100
+        )
+        max_new_tokens = st.slider(
+            "Max New Tokens",
+            min_value=1,
+            max_value=4000,
+            value=200
+        )
+        temperature = st.slider(
+            "Temperature",
+            min_value=0.10,
+            max_value=4.00,
+            value=0.70
+        )
+        top_p = st.slider(
+            "Top-p (nucleus sampling)",
+            min_value=0.10,
+            max_value=1.00,
+            value=0.90
+        )
+
+# --- Chatbot Function ---
+def chatbot(input_text, history, model_choice, system_message, max_new_tokens, temperature, top_p):
     # Create payload for the model
     payload = {
-        "inputs": {
-            "messages": [{"role": "user", "content": input_text}],
-            "system": system_message,
-            "max_tokens": max_new_tokens,
-            "temperature": temperature,
-            "top_p": top_p
-        }
+        "messages": [{"role": "user", "content": input_text}],
+        "system": system_message,
+        "max_tokens": max_new_tokens,
+        "temperature": temperature,
+        "top_p": top_p
     }
     
     # Run inference using the selected model
     try:
-        headers = {"Authorization": f"Bearer {st.secrets['HUGGINGFACE_TOKEN']}"}
-        response = requests.post(api_url, headers=headers, json=payload).json()
-        
-        # Handle the response format
-        if isinstance(response, list) and len(response) > 0:
-            # Assuming the response is a list of generated text
-            assistant_response = response[0].get("generated_text", "No response generated.")
-        elif isinstance(response, dict) and "generated_text" in response:
-            # If the response is a dictionary with generated_text
-            assistant_response = response["generated_text"]
+        response = demo(payload)  # Use the demo object directly
+        if isinstance(response, dict) and "choices" in response:
+            assistant_response = response["choices"][0]["message"]["content"]
         else:
             assistant_response = "Unexpected model response format."
     except Exception as e:
@@ -53,57 +73,15 @@ def chatbot(input_text, history, model_choice, system_message, max_new_tokens, t
     
     # Append user and assistant messages to history
     history.append((input_text, assistant_response))
-    
     return history
 
-# --- Streamlit App ---
-st.set_page_config(page_title="DeepSeek Chatbot", page_icon="🤖", layout="wide")
-
-# Title and description
-st.title("DeepSeek Chatbot")
-st.markdown("""
-    Created by [ruslanmv.com](https://ruslanmv.com/)  
-    This is a demo of different DeepSeek models. Select a model, type your message, and click "Submit".  
-    You can also adjust optional parameters like system message, max new tokens, temperature, and top-p.
-""")
-
-# Sidebar for model selection and parameters
-with st.sidebar:
-    st.header("Options")
-    model_choice = st.radio(
-        "Choose a Model",
-        options=list(MODELS.keys()),
-        index=0
-    )
-    st.header("Optional Parameters")
-    system_message = st.text_area(
-        "System Message",
-        value="You are a friendly Chatbot created by ruslanmv.com",
-        height=100
-    )
-    max_new_tokens = st.slider(
-        "Max New Tokens",
-        min_value=1,
-        max_value=4000,
-        value=200
-    )
-    temperature = st.slider(
-        "Temperature",
-        min_value=0.10,
-        max_value=4.00,
-        value=0.70
-    )
-    top_p = st.slider(
-        "Top-p (nucleus sampling)",
-        min_value=0.10,
-        max_value=1.00,
-        value=0.90
-    )
-
-# Initialize chat history
+# --- Chat History Management ---
 if "chat_history" not in st.session_state:
     st.session_state.chat_history = []
 
+# --- Chat Interface ---
+st.header("Chat with DeepSeek")
+
 # Display chat history
 for user_msg, assistant_msg in st.session_state.chat_history:
     with st.chat_message("user"):
@@ -111,14 +89,14 @@ for user_msg, assistant_msg in st.session_state.chat_history:
     with st.chat_message("assistant"):
         st.write(assistant_msg)
 
-# Input box for user message
-user_input = st.chat_input("Type your message here...")
+# Input for new message
+input_text = st.chat_input("Type your message here...")
 
-# Handle user input
-if user_input:
-    # Add user message to chat history
+# Handle new message submission
+if input_text:
+    # Update chat history
     st.session_state.chat_history = chatbot(
-        user_input,
+        input_text,
         st.session_state.chat_history,
         model_choice,
         system_message,
@@ -127,5 +105,5 @@ if user_input:
         top_p
     )
     
-    # Rerun to update the chat display
+    # Rerun the app to display the updated chat history
     st.rerun()