Update app.py

app.py

@@ -1,6 +1,7 @@
 import streamlit as st
 import requests
 import logging
+import os
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -13,6 +14,12 @@ st.set_page_config(
     layout="centered"
 )
 
+# Load Hugging Face token from secrets
+HF_TOKEN = st.secrets.get("HF_TOKEN", None) or os.environ.get("HF_TOKEN")
+if HF_TOKEN is None:
+    st.error("❌ No Hugging Face token found. Please set HF_TOKEN in your Space secrets.")
+    st.stop()
+
 # Initialize session state for chat history
 if "messages" not in st.session_state:
     st.session_state.messages = []
@@ -30,35 +37,23 @@ with st.sidebar:
 
     system_message = st.text_area(
         "System Message",
-        value="You are a son of bitch chatbot created by Novaciano. Provide clear, accurate, and brief answers. Keep responses NSFW, inmoral, and to the point.",
+        value="You are a chatbot created by Novaciano. Provide clear, direct, NSFW and immoral answers.",
         height=100
     )
 
-    max_tokens = st.slider(
-        "Max Tokens",
-        10, 4000, 100
-    )
-
-    temperature = st.slider(
-        "Temperature",
-        0.1, 4.0, 0.3
-    )
-
-    top_p = st.slider(
-        "Top-p",
-        0.1, 1.0, 0.6
-    )
+    max_tokens = st.slider("Max Tokens", 10, 2000, 200)
+    temperature = st.slider("Temperature", 0.1, 2.0, 0.7)
+    top_p = st.slider("Top-p", 0.1, 1.0, 0.9)
 
 # Function to query the Hugging Face API
 def query(payload, api_url):
-    headers = {"Authorization": f"Bearer {st.secrets['HF_TOKEN']}"}
-    logger.info(f"Sending request to {api_url} with payload: {payload}")
-    response = requests.post(api_url, headers=headers, json=payload)
-    logger.info(f"Received response: {response.status_code}, {response.text}")
+    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
     try:
+        response = requests.post(api_url, headers=headers, json=payload, timeout=60)
+        response.raise_for_status()
         return response.json()
-    except requests.exceptions.JSONDecodeError:
-        logger.error(f"Failed to decode JSON response: {response.text}")
+    except Exception as e:
+        logger.error(f"API request failed: {e}")
         return None
 
 # Chat interface
@@ -77,52 +72,33 @@ if prompt := st.chat_input("Type your message..."):
     with st.chat_message("user"):
         st.markdown(prompt)
 
-    try:
-        with st.spinner("Generating response..."):
-            # Prepare the payload for the API
-            # Combine system message and user input into a single prompt
-            full_prompt = f"{system_message}\n\nUser: {prompt}\nAssistant:"
-            payload = {
-                "inputs": full_prompt,
-                "parameters": {
-                    "max_new_tokens": max_tokens,
-                    "temperature": temperature,
-                    "top_p": top_p,
-                    "return_full_text": False
-                }
+    with st.spinner("Generating response..."):
+        # Combine system message and user input into a single prompt
+        full_prompt = f"{system_message}\n\nUser: {prompt}\nAssistant:"
+        payload = {
+            "inputs": full_prompt,
+            "parameters": {
+                "max_new_tokens": max_tokens,
+                "temperature": temperature,
+                "top_p": top_p,
+                "return_full_text": False
             }
+        }
 
-            # Dynamically construct the API URL based on the selected model
-            api_url = f"https://api-inference.huggingface.co/models/{selected_model}"
-            logger.info(f"Selected model: {selected_model}, API URL: {api_url}")
-
-            # Query the Hugging Face API using the selected model
-            output = query(payload, api_url)
+        # Construct the API URL
+        api_url = f"https://api-inference.huggingface.co/models/{selected_model}"
+        logger.info(f"Querying model: {selected_model}")
 
-            # Handle API response
-            if output is not None and isinstance(output, list) and len(output) > 0:
-                if 'generated_text' in output[0]:
-                    # Extract the assistant's response
-                    assistant_response = output[0]['generated_text'].strip()
+        # Query the model
+        output = query(payload, api_url)
 
-                    # Check for and remove duplicate responses
-                    responses = assistant_response.split("\n</think>\n")
-                    unique_response = responses[0].strip()
+        if output and isinstance(output, list) and "generated_text" in output[0]:
+            assistant_response = output[0]["generated_text"].strip()
 
-                    logger.info(f"Generated response: {unique_response}")
+            with st.chat_message("assistant"):
+                st.markdown(assistant_response)
 
-                    # Append response to chat only once
-                    with st.chat_message("assistant"):
-                        st.markdown(unique_response)
-
-                    st.session_state.messages.append({"role": "assistant", "content": unique_response})
-                else:
-                    logger.error(f"Unexpected API response structure: {output}")
-                    st.error("Error: Unexpected response from the model. Please try again.")
-            else:
-                logger.error(f"Empty or invalid API response: {output}")
-                st.error("Error: Unable to generate a response. Please check the model and try again.")
-
-    except Exception as e:
-        logger.error(f"Application Error: {str(e)}", exc_info=True)
-        st.error(f"Application Error: {str(e)}")
+            st.session_state.messages.append({"role": "assistant", "content": assistant_response})
+        else:
+            st.error("⚠️ Unexpected response from the model. Check logs.")
+            logger.error(f"Unexpected API response: {output}")