Update app.py

app.py

@@ -2,29 +2,35 @@ import os
 import time
 from datetime import datetime
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
 # -- SETUP --
 os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
 
 @st.cache_resource
-def load_respondent():
-    model_id = "mistralai/Mistral-7B-Instruct-v0.1"
+def load_model():
+    model_id = "google/flan-t5-base"
     tokenizer = AutoTokenizer.from_pretrained(model_id)
-    model = AutoModelForCausalLM.from_pretrained(
-        model_id,
-        device_map="auto",
-        trust_remote_code=True,
-        torch_dtype="auto"
-    )
-    return pipeline("text-generation", model=model, tokenizer=tokenizer)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
+    return tokenizer, model
 
-generator = load_respondent()
+tokenizer, model = load_model()
 
 if "history" not in st.session_state:
     st.session_state.history = []
     st.session_state.summary = ""
 
+# -- TEXT GENERATION FUNCTION --
+def generate_text(prompt, max_new_tokens=150):
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
+    outputs = model.generate(**inputs, max_new_tokens=max_new_tokens)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
+
+# -- HIGH-RISK 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)
+
 # -- STYLING --
 st.markdown("""
     <style>
@@ -43,26 +49,10 @@ 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")
 
-# -- SAFETY 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 --
 user_input = st.text_input("How are you feeling today?", placeholder="Start typing...")
 
-# -- REPLY FUNCTION --
-def generate_reply(user_input, context):
-    prompt = f"""You are a kind and empathetic AI assistant. Respond thoughtfully based on the following conversation:
-
-{context}
-User: {user_input}
-AI:"""
-    response = generator(prompt, max_new_tokens=150, temperature=0.7)[0]['generated_text']
-    return response.split("AI:")[-1].strip()
-
-# -- CONVERSATION FLOW --
+# -- MAIN CHAT LOGIC --
 if user_input:
     context = "\n".join([f"{s}: {m}" for s, m, _ in st.session_state.history[-4:]])
     with st.spinner("TARS is reflecting..."):
@@ -70,32 +60,29 @@ 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:
-            full_context = context + f"\nUser: {user_input}"
-            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))
+            prompt = f"Respond with empathy:\n{context}\nUser: {user_input}"
+            response = generate_text(prompt, max_new_tokens=100)
+    timestamp = datetime.now().strftime("%H:%M")
+    st.session_state.history.append(("🧍 You", user_input, timestamp))
+    st.session_state.history.append(("🤖 TARS", response, timestamp))
 
-# -- DISPLAY HISTORY --
+# -- DISPLAY CHAT --
 st.markdown("## 🗨️ Session")
 for speaker, msg, time in st.session_state.history:
     st.markdown(f"**{speaker} [{time}]:** {msg}")
 
-# -- SESSION SUMMARY --
+# -- SUMMARY GENERATION --
 if st.button("🧾 Generate Session Summary"):
     convo = "\n".join([f"{s}: {m}" for s, m, _ in st.session_state.history])
-    prompt = f"""Summarize the emotional tone and key themes from this conversation in 3 sentences:
-
-{convo}
-
-Summary:"""
+    summary_prompt = f"Summarize this conversation in 2-3 thoughtful sentences:\n{convo}"
     try:
-        output = generator(prompt, max_new_tokens=200, temperature=0.5)[0]['generated_text']
-        st.session_state.summary = output.split("Summary:")[-1].strip()
+        summary = generate_text(summary_prompt, max_new_tokens=150)
+        st.session_state.summary = summary
     except Exception as e:
         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)