Update app.py

app.py

@@ -2,37 +2,24 @@ import os
 import time
 from datetime import datetime
 import streamlit as st
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
 # -- SETUP --
 os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
 
 @st.cache_resource
-def load_pipeline():
-    model_id = "tiiuae/falcon-7b-instruct"
-    pipe = pipeline(
-        "text-generation",
-        model=AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True),
-        tokenizer=AutoTokenizer.from_pretrained(model_id, trust_remote_code=True),
-        device_map="auto"
-    )
-    return pipe
-
-generator = load_pipeline()
+def load_respondent():
+    model_id = "microsoft/phi-2"  # switched to a coherent, safe small model
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    model = AutoModelForCausalLM.from_pretrained(model_id)
+    return pipeline("text-generation", model=model, tokenizer=tokenizer)
+
+generator = load_respondent()
 
 if "history" not in st.session_state:
     st.session_state.history = []
     st.session_state.summary = ""
 
-# -- UTILS --
-TRIGGER_PHRASES = ["kill myself", "end it all", "suicide", "not worth living", "can't go on"]
-def is_high_risk(text):
-    return any(phrase in text.lower() for phrase in TRIGGER_PHRASES)
-
-def get_reply(prompt, max_new_tokens=150, temperature=0.7):
-    out = generator(prompt, max_new_tokens=max_new_tokens, temperature=temperature)[0]["generated_text"]
-    return out.split("AI:")[-1].strip() if "AI:" in out else out.strip()
-
 # -- STYLING --
 st.markdown("""
     <style>
@@ -51,38 +38,54 @@ st.title("🧠 TARS.help")
 st.markdown("### A minimal AI that listens, reflects, and replies.")
 st.markdown(f"🗓️ {datetime.now().strftime('%B %d, %Y')} | {len(st.session_state.history)//2} exchanges")
 
-# -- USER INPUT --
+# -- HIGH-RISK PHRASE FILTER --
+TRIGGER_PHRASES = ["kill myself", "end it all", "suicide", "not worth living", "can't go on"]
+
+def is_high_risk(text):
+    return any(phrase in text.lower() for phrase in TRIGGER_PHRASES)
+
+# -- INPUT --
 user_input = st.text_input("How are you feeling today?", placeholder="Start typing...")
 
-# -- MAIN FLOW --
+# -- REPLY FUNCTION --
+def generate_reply(user_input, context):
+    prompt = f"""You are a calm, helpful AI assistant who supports users emotionally. Be kind and thoughtful in your reply.
+
+{context}
+User: {user_input}
+AI:"""
+    response = generator(prompt, max_new_tokens=80, temperature=0.7)[0]['generated_text']
+    return response.split("AI:")[-1].strip()
+
+# -- CONVERSATION FLOW --
 if user_input:
     context = "\n".join([f"{s}: {m}" for s, m, _ in st.session_state.history[-4:]])
     with st.spinner("TARS is reflecting..."):
-        time.sleep(1)
+        time.sleep(0.5)
         if is_high_risk(user_input):
             response = "I'm really sorry you're feeling this way. You're not alone — please talk to someone you trust or a mental health professional. 💙"
         else:
-            prompt = f"You are a kind and calm AI assistant.\n{context}\nUser: {user_input}\nAI:"
-            response = get_reply(prompt, max_new_tokens=150)
-    timestamp = datetime.now().strftime("%H:%M")
-    st.session_state.history.append(("🧍 You", user_input, timestamp))
-    st.session_state.history.append(("🤖 TARS", response, timestamp))
+            response = generate_reply(user_input, context)
+        timestamp = datetime.now().strftime("%H:%M")
+        st.session_state.history.append(("🧍 You", user_input, timestamp))
+        st.session_state.history.append(("🤖 TARS", response, timestamp))
 
-# -- CHAT DISPLAY --
+# -- DISPLAY HISTORY --
 st.markdown("## 🗨️ Session")
 for speaker, msg, time in st.session_state.history:
     st.markdown(f"**{speaker} [{time}]:** {msg}")
 
-# -- SUMMARY --
+# -- SESSION SUMMARY --
 if st.button("🧾 Generate Session Summary"):
     convo = "\n".join([f"{s}: {m}" for s, m, _ in st.session_state.history])
-    prompt = f"Summarize this conversation in 3 reflective sentences:\n{convo}\nSummary:"
-    try:
-        summary = get_reply(prompt, max_new_tokens=200, temperature=0.5)
-        st.session_state.summary = summary
-    except Exception as e:
-        st.error("Summary generation failed.")
-        st.exception(e)
+    prompt = f"""You are summarizing a thoughtful conversation between a user and an AI assistant. Write a kind, reflective note based on this interaction.
+
+Conversation:
+{convo}
+
+Summary:"""
+    summary = generator(prompt, max_new_tokens=100, temperature=0.5)[0]['generated_text']
+    st.session_state.summary = summary.split("Summary:")[-1].strip()
 
 if st.session_state.summary:
     st.markdown("### 🧠 Session Note")
@@ -91,4 +94,4 @@ if st.session_state.summary:
 
 # -- FOOTER --
 st.markdown("---")
-st.caption("TARS is not a therapist, but a quiet assistant that reflects with you.")
+st.caption("TARS is not a therapist but a friend. Just a quiet assistant that reflects with you.")