Update mindmentor_app.py

mindmentor_app.py

@@ -1,33 +1,42 @@
 import streamlit as st
 import os
 from llama_cpp import Llama
-from huggingface_hub import hf_hub_download
 
 # ------------------------------
-# Load Quantized LLM Model
+# Load the Quantized LLM Model
 # ------------------------------
 @st.cache_resource(show_spinner="🔄 Loading TinyLlama AI Model...")
 def load_llm():
-    try:
+    local_dir = "models"
+    filename = "tinyllama-1.1b-chat-v1.0.Q4_0.gguf"
+    model_path = os.path.join(local_dir, filename)
+
+    if not os.path.exists(model_path):
+        from huggingface_hub import hf_hub_download
         model_path = hf_hub_download(
             repo_id="TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF",
-            filename="tinyllama-1.1b-chat-v1.0.Q4_0.gguf"
-        )
-        return Llama(
-            model_path=model_path,
-            n_ctx=2048,
-            n_threads=4,
-            n_gpu_layers=0,
-            verbose=False
+            filename=filename,
+            local_dir=local_dir,
+            local_dir_use_symlinks=False
         )
-    except Exception as e:
-        st.error(f"❌ Failed to load TinyLlama model: {e}")
-        st.stop()
 
-llm = load_llm()
+    return Llama(
+        model_path=model_path,
+        n_ctx=1024,  # Use 1024 context to reduce memory usage on Spaces
+        n_threads=4,
+        n_gpu_layers=0,
+        verbose=False
+    )
+
+# Initialize LLM
+try:
+    llm = load_llm()
+except Exception as e:
+    st.error(f"❌ Failed to load TinyLlama: {e}")
+    st.stop()
 
 # ------------------------------
-# Generate Chatbot Response
+# Function to Get Chatbot Reply
 # ------------------------------
 def get_response(message):
     prompt = f"""You are MindMentor, an offline AI coach that helps people with stress, emotions, and self-awareness. You are kind, caring, helpful, and positive.
@@ -37,25 +46,25 @@ Assistant:"""
         output = llm(prompt=prompt, max_tokens=200, stop=["User:", "\n"])
         return output["choices"][0]["text"].strip()
     except Exception as e:
-        return f"⚠️ Error generating response: {e}"
+        return f"⚠️ Oops! Model failed to respond: {e}"
 
 # ------------------------------
-# Streamlit Chat UI
+# Streamlit UI
 # ------------------------------
-st.set_page_config(page_title="🧠 MindMentor - Offline AI Coach", layout="centered")
+st.set_page_config(page_title="MindMentor AI Chatbot", layout="centered")
 st.title("🧠 MindMentor - Offline AI Coach")
-st.markdown("Your personal mental wellbeing assistant 💬 — all offline and private.")
+st.markdown("Your personal mental wellbeing assistant 💬 — all offline!")
 
-# Initialize chat history
+# Session state
 if "chat_history" not in st.session_state:
     st.session_state.chat_history = []
 
-# Display previous messages
+# Show chat history
 for msg in st.session_state.chat_history:
     with st.chat_message(msg["role"]):
         st.markdown(msg["content"])
 
-# Input box for new user message
+# Input box
 user_input = st.chat_input("Talk to MindMentor...")
 
 if user_input:
@@ -63,10 +72,9 @@ if user_input:
     st.chat_message("user").markdown(user_input)
     st.session_state.chat_history.append({"role": "user", "content": user_input})
 
-    # Generate and show assistant reply
+    # Generate assistant reply
     with st.chat_message("assistant"):
         with st.spinner("Thinking..."):
             reply = get_response(user_input)
             st.markdown(reply)
             st.session_state.chat_history.append({"role": "assistant", "content": reply})
-