Update app.py

app.py

@@ -1,56 +1,40 @@
-import os
-import time
-from datetime import datetime
 import streamlit as st
+from datetime import datetime
+import time
 from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+import os
 
 # -- SETUP --
 os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
 
 @st.cache_resource
-def load_pipeline():
+def load_model():
     model_id = "tiiuae/falcon-7b-instruct"
     tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
     model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True)
     return pipeline("text-generation", model=model, tokenizer=tokenizer)
 
-generator = load_pipeline()
+generator = load_model()
 
 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):
+def get_response(prompt, max_new_tokens=150, temperature=0.7):
     output = generator(prompt, max_new_tokens=max_new_tokens, temperature=temperature)[0]["generated_text"]
     return output.split("AI:")[-1].strip() if "AI:" in output else output.strip()
 
-# -- STYLING --
-st.markdown("""
-    <style>
-    body {
-        background-color: #111827;
-        color: #f3f4f6;
-    }
-    .stTextInput > div > div > input {
-        color: #f3f4f6;
-    }
-    </style>
-""", unsafe_allow_html=True)
-
-# -- HEADER --
 st.title("🧠 TARS.help")
-st.markdown("### A minimal AI that listens, reflects, and replies.")
+st.markdown("### A quiet AI that reflects and replies.")
 st.markdown(f"🗓️ {datetime.now().strftime('%B %d, %Y')} | {len(st.session_state.history)//2} exchanges")
 
-# -- USER INPUT --
 user_input = st.text_input("How are you feeling today?", placeholder="Start typing...")
 
-# -- MAIN CHAT 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..."):
@@ -58,34 +42,30 @@ if user_input:
         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 compassionate AI therapist.\n{context}\nUser: {user_input}\nAI:"
-            response = get_reply(prompt)
-    timestamp = datetime.now().strftime("%H:%M")
-    st.session_state.history.append(("🧍 You", user_input, timestamp))
-    st.session_state.history.append(("🤖 TARS", response, timestamp))
+            prompt = f"You are a calm, empathetic AI assistant.\n{context}\nUser: {user_input}\nAI:"
+            response = get_response(prompt)
+        timestamp = datetime.now().strftime("%H:%M")
+        st.session_state.history.append(("🧍 You", user_input, timestamp))
+        st.session_state.history.append(("🤖 TARS", response, timestamp))
 
-# -- DISPLAY CHAT --
 st.markdown("## 🗨️ Session")
 for speaker, msg, time in st.session_state.history:
     st.markdown(f"**{speaker} [{time}]:** {msg}")
 
-# -- SUMMARY GENERATION --
 if st.button("🧾 Generate Session Summary"):
     convo = "\n".join([f"{s}: {m}" for s, m, _ in st.session_state.history])
-    summary_prompt = f"Summarize the tone and content of this therapy session in 3 thoughtful sentences:\n{convo}\nSummary:"
+    prompt = f"Summarize the emotional tone and themes in this conversation:\n{convo}\nSummary:"
     try:
-        summary = get_reply(summary_prompt, max_new_tokens=200, temperature=0.5)
+        summary = get_response(prompt, max_new_tokens=200, temperature=0.5)
         st.session_state.summary = summary
     except Exception as e:
-        st.error("❌ Summary generation failed.")
+        st.error("Summary generation failed.")
         st.exception(e)
 
-# -- DISPLAY SUMMARY --
 if st.session_state.summary:
     st.markdown("### 🧠 Session Note")
     st.markdown(st.session_state.summary)
     st.download_button("📥 Download Summary", st.session_state.summary, file_name="tars_session.txt")
 
-# -- 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. If you're in crisis, please seek help from a professional.")