Update app.py

app.py

@@ -1,3 +1,5 @@
+# app.py
+
 import os
 import streamlit as st
 import torch
@@ -27,7 +29,7 @@ embedder = get_embedder()
 def fetch_mems(query, k=3):
     vec = embedder.encode(query).astype("float32").tolist()
     return supabase.rpc(
-        "match_memories", 
+        "match_memories",
         {"query_embedding": vec, "match_count": k}
     ).execute().data
 
@@ -65,8 +67,7 @@ def load_generator():
             device_map="auto"
         )
     else:
-        # CPU or no CUDA: use fp16 if available, else fp32
-        dtype = torch.float16 if torch.cuda.is_available() or torch.cuda.device_count()>0 else torch.float32
+        dtype = torch.float16 if torch.cuda.is_available() else torch.float32
         base = AutoModelForCausalLM.from_pretrained(
             "microsoft/phi-2",
             trust_remote_code=True,
@@ -100,13 +101,14 @@ tokenizer, generator = load_generator()
 
 # ── System prompt to reduce hallucinations ──────────────────────────────────
 SYSTEM = (
-    "You are a helpful assistant.\n"
+    "You are a helpful assistant for DeepTalks with a base model as Phi-2 and fine tuned by Sourish for my domain specific role.\n"
+    "My domain is assisting you within my expertise by listening to you understanding you & supporting you.\n"
     "Answer **only** using the information in the memory below.\n"
     "If the answer is not in memory, reply: \"I don't know.\"\n"
 )
 
 # ── Streamlit UI ──────────────────────────────────────────────────────────
-st.set_page_config(layout="wide")
+st.set_page_config(layout="centered")
 st.title("🧠 Memory-Aware Phi-2 Chat")
 
 if "history" not in st.session_state:
@@ -114,21 +116,20 @@ if "history" not in st.session_state:
 
 # Render existing chat history
 for role, msg in st.session_state.history:
-    st.chat_message("user" if role=="You" else "assistant").write(msg)
+    st.chat_message("user" if role == "You" else "assistant").write(msg)
 
 # Input box at the bottom
 user_input = st.chat_input("Type your message...")
 
 if user_input:
-    # Append user message
+    # 1) Show and store user message
+    st.chat_message("user").write(user_input)
     st.session_state.history.append(("You", user_input))
     add_mem("user", user_input)
 
-    # Retrieve relevant memories
+    # 2) Retrieve memories & build prompt
     mems = fetch_mems(user_input, k=3)
     mem_block = "\n".join(f"{m['speaker']}: {m['text']}" for m in mems)
-
-    # Build prompt
     prompt = f"""{SYSTEM}
 
 Memory:
@@ -137,7 +138,7 @@ Memory:
 User: {user_input}
 Assistant:"""
 
-    # Generate reply with spinner
+    # 3) Generate reply with spinner
     with st.spinner("Thinking..."):
         try:
             out = generator(prompt)[0]["generated_text"].strip()
@@ -145,6 +146,7 @@ Assistant:"""
             out = "Sorry, I encountered an error."
             st.error(f"Generation error: {e}")
 
-    # Append assistant reply
+    # 4) Show and store assistant reply
+    st.chat_message("assistant").write(out)
     st.session_state.history.append(("Bot", out))
     add_mem("assistant", out)