Update app.py

app.py

@@ -1,41 +1,45 @@
 import os
+import numpy as np
 import streamlit as st
 import faiss
-import numpy as np
 from sentence_transformers import SentenceTransformer
 from groq import Groq
 
-# --- Load API key from environment (HF Repo Secrets) ---
+# ---------- Secrets / API Key ----------
 API_KEY = os.getenv("GROQ_API_KEY")
 if not API_KEY:
     st.error(
-        "GROQ_API_KEY not found. In your Space go to: "
-        "Settings → Repository secrets → Add new secret (Name: GROQ_API_KEY, Value: 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)."
+        "GROQ_API_KEY not found.\n\n"
+        "Go to your Space → Settings → Repository secrets → Add new secret\n"
+        "Name: GROQ_API_KEY | Value: <your Groq key>\n\n"
+        "Then Restart/Restart this Space."
     )
     st.stop()
 
-# --- Init Groq client (no key printed/logged) ---
+# ---------- Groq Client ----------
 client = Groq(api_key=API_KEY)
 
-# --- Cache the embedding model to speed up reloads ---
+# ---------- Models / Index ----------
 @st.cache_resource
 def load_embedder():
+    # 384-dim embeddings
     return SentenceTransformer("all-MiniLM-L6-v2")
 
 embedding_model = load_embedder()
 
-# --- FAISS index (384 dims for MiniLM) ---
-dimension = 384
+DIM = 384  # all-MiniLM-L6-v2 dimension
 if "faiss_index" not in st.session_state:
-    st.session_state.faiss_index = faiss.IndexFlatL2(dimension)
+    st.session_state.faiss_index = faiss.IndexFlatL2(DIM)
 if "chunks_store" not in st.session_state:
     st.session_state.chunks_store = []
 
 index = st.session_state.faiss_index
 chunks_store = st.session_state.chunks_store
 
-# ---- Utilities ----
-def chunk_text(text, max_length=500):
+
+# ---------- Helpers ----------
+def chunk_text(text: str, max_length: int = 500):
+    """Simple whitespace chunker by character budget."""
     words, chunks, cur = text.split(), [], []
     for w in words:
         if len(" ".join(cur)) + len(w) + 1 <= max_length:
@@ -47,25 +51,27 @@ def chunk_text(text, max_length=500):
         chunks.append(" ".join(cur))
     return chunks
 
+
 def embed_and_store(chunks):
     if not chunks:
         return
-    embs = embedding_model.encode(chunks, convert_to_numpy=True, normalize_embeddings=False)
-    # Ensure float32 for FAISS
-    embs = np.asarray(embs, dtype="float32")
+    embs = embedding_model.encode(
+        chunks, convert_to_numpy=True, normalize_embeddings=False
+    ).astype("float32")
     index.add(embs)
     chunks_store.extend(chunks)
 
+
 def query_llm(prompt: str) -> str:
-    # Streaming chat completion
+    """Stream a response from Groq and return full text."""
     stream = client.chat.completions.create(
         model="deepseek-r1-distill-llama-70b",
         messages=[
             {
                 "role": "system",
                 "content": (
-                    "You are a relationship counselor. Analyze the WhatsApp conversation and "
-                    "provide insights on red flags, toxicity, and improvements. "
+                    "You are a relationship counselor. Analyze the WhatsApp conversation "
+                    "and provide insights on red flags, toxicity, and improvements. "
                     "Start every answer with: 'Toxicity score: X/10'."
                 ),
             },
@@ -83,7 +89,9 @@ def query_llm(prompt: str) -> str:
         out.append(delta)
     return "".join(out)
 
-# ---- UI ----
+
+# ---------- UI ----------
+st.set_page_config(page_title="AI Relationship Counsellor", layout="centered")
 st.title("AI Relationship Counsellor")
 
 uploaded_file = st.file_uploader("Upload a .txt export of your WhatsApp chat", type=["txt"])
@@ -97,13 +105,12 @@ if uploaded_file:
 
     user_query = st.text_input("Ask a question about your relationship:")
     if user_query:
-        # Search top-k relevant chunks
-        k = min(5, index.ntotal) if index.ntotal > 0 else 0
-        if k == 0:
-            st.warning("No text indexed yet. Please upload a chat file.")
+        if index.ntotal == 0:
+            st.warning("Nothing indexed yet. Please upload a chat file.")
         else:
-            q_emb = embedding_model.encode([user_query], convert_to_numpy=True)
-            q_emb = np.asarray(q_emb, dtype="float32")
+            # top-k retrieval
+            k = min(5, index.ntotal)
+            q_emb = embedding_model.encode([user_query], convert_to_numpy=True).astype("float32")
             distances, idxs = index.search(q_emb, k)
             relevant = [chunks_store[i] for i in idxs[0] if 0 <= i < len(chunks_store)]
 
@@ -115,3 +122,5 @@ if uploaded_file:
 
             st.markdown("### AI Analysis")
             st.write(answer)
+else:
+    st.info("Upload a WhatsApp chat (.txt) to begin.")