Update app.py

app.py

@@ -1,15 +1,16 @@
 import streamlit as st
 
-# IMPORTANT: Call set_page_config before any other Streamlit command
+# 1. Set page config FIRST
 st.set_page_config(
     page_title="TARS: Therapist Assistance and Response System", 
     page_icon="🧠"
 )
 
-from transformers import pipeline
+import os
 import torch
+from transformers import pipeline
 
-# Canadian Crisis Resources
+# Canadian Crisis Resources (example)
 CRISIS_RESOURCES = {
     "Canada Suicide Prevention Service": "1-833-456-4566",
     "Crisis Services Canada": "1-833-456-4566",
@@ -17,7 +18,7 @@ CRISIS_RESOURCES = {
     "First Nations and Inuit Hope for Wellness Help Line": "1-855-242-3310"
 }
 
-# Safety keywords for suicide risk detection
+# Keywords for detecting potential self-harm or suicidal language
 SUICIDE_KEYWORDS = [
     "suicide", "kill myself", "end my life", 
     "want to die", "hopeless", "no way out", 
@@ -25,132 +26,157 @@ SUICIDE_KEYWORDS = [
 ]
 
 class AITherapistAssistant:
-    def __init__(self):
-        # Load smaller, faster models to work within Hugging Face Spaces constraints
+    def __init__(self, conversation_model_name="microsoft/phi-1_5", summary_model_name="facebook/bart-large-cnn"):
+        """
+        Initialize the conversation (LLM) model and the summarization model.
+        
+        If you truly have a different 'phi2' from Microsoft, replace 'microsoft/phi-1_5'
+        with your private or custom Hugging Face repo name.
+        """
+        # Load conversation LLM (phi2 / phi-1_5)
         try:
-            # Conversational model
             self.conversation_model = pipeline(
-                "text-generation", 
-                model="microsoft/DialoGPT-small", 
+                "text-generation",
+                model=conversation_model_name,
                 device=0 if torch.cuda.is_available() else -1
             )
-            
-            # Summarization model
+        except Exception as e:
+            st.error(f"Error loading conversation model: {e}")
+            self.conversation_model = None
+
+        # Load summarization model (BART Large CNN as default)
+        try:
             self.summary_model = pipeline(
                 "summarization", 
-                model="facebook/bart-large-cnn", 
+                model=summary_model_name, 
                 device=0 if torch.cuda.is_available() else -1
             )
         except Exception as e:
-            # Display error if model loading fails
-            st.error(f"Model loading error: {e}")
-            self.conversation_model = None
+            st.error(f"Error loading summary model: {e}")
             self.summary_model = None
     
-    def detect_crisis(self, message):
-        """Detect potential suicide risk in message."""
+    def detect_crisis(self, message: str) -> bool:
+        """Check if message contains suicidal or distress-related keywords."""
         message_lower = message.lower()
-        for keyword in SUICIDE_KEYWORDS:
-            if keyword in message_lower:
-                return True
-        return False
+        return any(keyword in message_lower for keyword in SUICIDE_KEYWORDS)
     
-    def generate_response(self, message):
-        """Generate empathetic AI response."""
+    def generate_response(self, message: str) -> str:
+        """Generate a supportive AI response from the conversation model."""
         if not self.conversation_model:
-            return "I'm here to listen. Would you like to share more about how you're feeling?"
+            return (
+                "I'm here to listen, but I'm currently having trouble loading my AI model. "
+                "Please try again later."
+            )
+        
+        # Prompt to steer the model toward empathy and support
+        prompt = (
+            "You are a compassionate AI therapist. Respond supportively to this message:\n"
+            f"{message}\n\n"
+            "Be empathetic, validate feelings, and avoid giving direct medical advice."
+        )
         
         try:
-            # Contextual prompt to guide response
-            full_prompt = (
-                "You are a compassionate AI therapist. Respond supportively to this message: "
-                f"{message}. Be empathetic, validate feelings, and avoid giving direct medical advice."
+            outputs = self.conversation_model(
+                prompt, 
+                max_length=250,
+                num_return_sequences=1,
+                do_sample=True,
+                top_p=0.9,
+                temperature=0.7
             )
-            # Generate response
-            response = self.conversation_model(
-                full_prompt, 
-                max_length=100, 
-                num_return_sequences=1
-            )[0]['generated_text']
+            response_text = outputs[0]["generated_text"]
             
-            return response
-        
-        except Exception:
-            return "I'm here to listen. Would you like to share more about how you're feeling?"
+            # If the model echoes the prompt, strip it out:
+            if response_text.startswith(prompt):
+                response_text = response_text[len(prompt):].strip()
+            
+            return response_text.strip()
+        except Exception as e:
+            st.error(f"Error generating response: {e}")
+            return "I'm sorry, but I'm having trouble responding right now."
     
-    def generate_summary(self, conversation):
-        """Generate a professional therapy-style summary."""
+    def generate_summary(self, conversation_text: str) -> str:
+        """Generate a short summary of the entire conversation."""
         if not self.summary_model:
-            return "Summary generation is temporarily unavailable."
+            return "Summary model is unavailable at the moment."
         
         try:
-            summary = self.summary_model(
-                conversation, 
+            summary_output = self.summary_model(
+                conversation_text, 
                 max_length=130, 
                 min_length=30, 
                 do_sample=False
-            )[0]['summary_text']
-            return summary
-        
-        except Exception:
-            return "Summary could not be generated."
+            )
+            return summary_output[0]["summary_text"]
+        except Exception as e:
+            st.error(f"Error generating summary: {e}")
+            return "Sorry, I couldn't generate a summary."
 
 def main():
-    # Title and description
     st.title("🧠 TARS: Therapist Assistance and Response System")
-    st.write("A supportive space to share your feelings safely.\n\n"
-             "**Disclaimer**: I am not a licensed therapist. "
-             "If you're in crisis, please reach out to professional help immediately.")
-
-    # Initialize session state
-    if 'conversation' not in st.session_state:
-        st.session_state.conversation = []
+    st.write(
+        "A supportive space to share your feelings safely.\n\n"
+        "**Disclaimer**: I am not a licensed therapist. If you're in crisis, "
+        "please reach out to professional help immediately."
+    )
+    
+    # Note if running on Hugging Face Spaces
+    if os.environ.get("SPACE_ID"):
+        st.info("Running on Hugging Face Spaces.")
+    
+    # Instantiate the assistant with phi-1_5 (or your custom 'phi2')
+    if "assistant" not in st.session_state:
+        st.session_state.assistant = AITherapistAssistant(
+            conversation_model_name="microsoft/phi-1_5",  # replace if needed
+            summary_model_name="facebook/bart-large-cnn"
+        )
     
-    if 'assistant' not in st.session_state:
-        st.session_state.assistant = AITherapistAssistant()
+    # Keep track of conversation
+    if "conversation" not in st.session_state:
+        st.session_state.conversation = []
     
-    # Display conversation
+    # Display existing conversation
     for message in st.session_state.conversation:
-        if message['sender'] == 'user':
-            st.chat_message("user").write(message['text'])
+        if message["sender"] == "user":
+            st.chat_message("user").write(message["text"])
         else:
-            st.chat_message("assistant").write(message['text'])
+            st.chat_message("assistant").write(message["text"])
     
-    # User input with chat_input
-    if prompt := st.chat_input("Share your thoughts. I'm here to listen."):
-        # Check for crisis indicators
+    # Collect user input
+    if prompt := st.chat_input("How are you feeling today?"):
+        # Crisis detection
         if st.session_state.assistant.detect_crisis(prompt):
-            st.warning("⚠️ Crisis Support Detected")
-            st.markdown("**Immediate Support Resources:**")
+            st.warning("⚠️ Potential crisis detected.")
+            st.markdown("**Immediate Support Resources (Canada):**")
             for org, phone in CRISIS_RESOURCES.items():
                 st.markdown(f"- {org}: `{phone}`")
         
-        # Add user message to conversation
-        st.session_state.conversation.append({'sender': 'user', 'text': prompt})
+        # Display user message
+        st.session_state.conversation.append({"sender": "user", "text": prompt})
         st.chat_message("user").write(prompt)
         
         # Generate AI response
         with st.chat_message("assistant"):
-            with st.spinner("Listening and reflecting..."):
+            with st.spinner("Thinking..."):
                 ai_response = st.session_state.assistant.generate_response(prompt)
                 st.write(ai_response)
-        
-        # Add AI response to conversation
-        st.session_state.conversation.append({'sender': 'ai', 'text': ai_response})
+        st.session_state.conversation.append({"sender": "assistant", "text": ai_response})
     
-    # Session Summary Generation
-    if st.session_state.conversation:
-        if st.button("Generate Session Summary"):
-            conversation_text = " ".join(msg['text'] for msg in st.session_state.conversation)
+    # Summarize conversation
+    if st.button("Generate Session Summary"):
+        if st.session_state.conversation:
+            conversation_text = " ".join(msg["text"] for msg in st.session_state.conversation)
             summary = st.session_state.assistant.generate_summary(conversation_text)
-            st.markdown("**Session Summary:**")
+            st.subheader("Session Summary")
             st.write(summary)
+        else:
+            st.info("No conversation to summarize yet.")
     
-    # Safety Information
+    # Crisis Support Info in Sidebar
     st.sidebar.title("🆘 Crisis Support")
-    st.sidebar.markdown("**Remember: You are valued and supported.**")
+    st.sidebar.markdown("If you're in crisis, please contact:")
     for org, phone in CRISIS_RESOURCES.items():
-        st.sidebar.markdown(f"- {org}: `{phone}`")
+        st.sidebar.markdown(f"- **{org}**: `{phone}`")
 
 if __name__ == "__main__":
     main()