Update app.py

app.py

@@ -2,17 +2,22 @@ import os
 import time
 from datetime import datetime
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
 # -- SETUP --
 os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
 
 @st.cache_resource
 def load_respondent():
-    model_id = "google/flan-t5-small"
+    model_id = "mistralai/Mistral-7B-Instruct-v0.1"
     tokenizer = AutoTokenizer.from_pretrained(model_id)
-    model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
-    return pipeline("text2text-generation", model=model, tokenizer=tokenizer)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id,
+        device_map="auto",
+        trust_remote_code=True,
+        torch_dtype="auto"
+    )
+    return pipeline("text-generation", model=model, tokenizer=tokenizer)
 
 generator = load_respondent()
 
@@ -48,21 +53,25 @@ def is_high_risk(text):
 user_input = st.text_input("How are you feeling today?", placeholder="Start typing...")
 
 # -- REPLY FUNCTION --
-def generate_reply(context):
-    prompt = f"Respond empathetically to this conversation:\n{context}"
-    result = generator(prompt, max_new_tokens=80, temperature=0.7)[0]["generated_text"]
-    return result.strip()
+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 --
 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(0.5)
+        time.sleep(1.2)
         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(full_context)
+            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))
@@ -75,10 +84,14 @@ for speaker, msg, time in st.session_state.history:
 # -- 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 2-3 sentences:\n{convo}"
+    prompt = f"""Summarize the emotional tone and key themes from this conversation in 3 sentences:
+
+{convo}
+
+Summary:"""
     try:
-        output = generator(prompt, max_new_tokens=120, temperature=0.5)[0]['generated_text']
-        st.session_state.summary = output.strip()
+        output = generator(prompt, max_new_tokens=200, temperature=0.5)[0]['generated_text']
+        st.session_state.summary = output.split("Summary:")[-1].strip()
     except Exception as e:
         st.error("❌ Summary generation failed.")
         st.exception(e)