Commit

app.py

@@ -98,7 +98,9 @@ generator = load_pipeline()
 
 # ── Streamlit UI ──────────────────────────────────────────────────────────
 st.set_page_config(layout="centered")
-st.title("🧠 Memory-Aware Phi-2 Chat")
+st.title("🧠 DeepTalks")
+st.subheader("Your personal AI Companion")  # Add subheader
+st.markdown("---")
 
 # initialize chat history
 if "history" not in st.session_state:
@@ -129,14 +131,15 @@ Assistant:"""
     # generate reply
     with st.spinner("Thinking…"):
         try:
-            # 1) raw generation
-            raw = generator(prompt)[0]["generated_text"]
-            # 2) Remove the prompt prefix (so you only see new content)
-            reply = raw[len(prompt):].strip()
-            # 3) If it ever slurps in "User:" or "Assistant:", cut there
+            # pipeline was set up with `return_full_text=False`, so we get just the reply
+            reply = generator(prompt)[0]["generated_text"].strip()
+            # strip any accidental echoes
             for marker in ["User:", "Assistant:"]:
                 if marker in reply:
                     reply = reply.split(marker)[0].strip()
+            # if it somehow ends up empty, backstop with an apology
+            if not reply:
+                reply = "I'm sorry, I didn't catch that. Could you rephrase?"
         except Exception as e:
             reply = "I'm sorry, something went wrong."
             st.error(f"Generation error: {e}")