Update app.py

app.py

@@ -1,99 +1,164 @@
-"""
-Hugging Face Spaces: AI Friend Therapist Assistant
-
-This app listens to your messages, flags language that might indicate distress,
-and provides emergency contacts for Canada if needed. It also uses a Hugging Face
-summarization pipeline to generate a brief "therapist summary" of your conversation.
-
-Dependencies:
-  - streamlit
-  - transformers
-  - torch
-
-You can create a requirements.txt with:
-    streamlit
-    transformers
-    torch
-
-Deploy this file (app.py) along with requirements.txt to Hugging Face Spaces.
-"""
-
-# Dependency checks and initial imports
-try:
-    import streamlit as st
-except ImportError:
-    raise ImportError("Missing dependency: streamlit. Install it with 'pip install streamlit'.")
-
-# Immediately set the page configuration
-st.set_page_config(page_title="AI Friend Therapist Assistant", layout="centered")
-
-import os
-
-try:
-    from transformers import pipeline
-except ImportError:
-    raise ImportError("Missing dependency: transformers. Install it with 'pip install transformers'.")
-
-try:
-    import torch
-except ImportError:
-    raise ImportError("Missing dependency: torch. Install it with 'pip install torch'.")
-
-# Optional: display a note if running on Hugging Face Spaces
-if os.environ.get("SPACE_ID"):
-    st.info("Running on Hugging Face Spaces.")
-
-# Load the Hugging Face summarization model (cached to avoid reloading)
-@st.cache(allow_output_mutation=True)
-def load_summarizer():
-    return pipeline("summarization", model="facebook/bart-large-cnn")
-
-summarizer = load_summarizer()
-
-# Initialize session state to store conversation history
-if 'conversation' not in st.session_state:
-    st.session_state.conversation = []
-
-# App title and disclaimer
-st.title("AI Friend Therapist Assistant")
-st.write(
-    "I'm here to listen. Share your thoughts below.\n\n"
-    "*Disclaimer: I am not a licensed therapist. If you're in crisis, please seek professional help immediately.*"
-)
-
-# Text area for user input
-user_input = st.text_area("Your message", height=150)
-
-# Process user input when "Send" is clicked
-if st.button("Send"):
-    if user_input.strip():
-        # Append the message to conversation history
-        st.session_state.conversation.append(user_input)
-        
-        # Define keywords that might indicate suicidal ideation or distress
-        flagged_keywords = [
-            "suicide", "kill myself", "end my life", "self-harm", 
-            "self harm", "hopeless", "despair", "worthless", "can't go on"
-        ]
-        # Check if any flagged keyword is present in the input
-        if any(keyword in user_input.lower() for keyword in flagged_keywords):
-            st.error(
-                "Your message indicates you may be experiencing distress. "
-                "If you're in immediate danger, please call **911** immediately. "
-                "For non-immediate crisis support in Canada, call **Crisis Services Canada at 1.833.456.4566** "
-                "or visit [crisisservicescanada.ca](https://www.crisisservicescanada.ca/)."
+import streamlit as st
+from transformers import pipeline
+import torch
+
+# Canadian Crisis Resources
+CRISIS_RESOURCES = {
+    "Canada Suicide Prevention Service": "1-833-456-4566",
+    "Crisis Services Canada": "1-833-456-4566",
+    "Kids Help Phone": "1-800-668-6868",
+    "First Nations and Inuit Hope for Wellness Help Line": "1-855-242-3310"
+}
+
+# Safety keywords for suicide risk detection
+SUICIDE_KEYWORDS = [
+    "suicide", "kill myself", "end my life", 
+    "want to die", "hopeless", "no way out", 
+    "better off dead", "pain is too much"
+]
+
+class AITherapistAssistant:
+    def __init__(self):
+        # Load smaller, faster models to work within Hugging Face Spaces constraints
+        try:
+            # Conversational model
+            self.conversation_model = pipeline(
+                "text-generation", 
+                model="microsoft/DialoGPT-small", 
+                device=0 if torch.cuda.is_available() else -1
+            )
+            
+            # Summarization model
+            self.summary_model = pipeline(
+                "summarization", 
+                model="facebook/bart-small-summarization", 
+                device=0 if torch.cuda.is_available() else -1
             )
+        except Exception as e:
+            st.error(f"Model loading error: {e}")
+            self.conversation_model = None
+            self.summary_model = None
+    
+    def detect_crisis(self, message):
+        """Detect potential suicide risk in message"""
+        message_lower = message.lower()
+        for keyword in SUICIDE_KEYWORDS:
+            if keyword in message_lower:
+                return True
+        return False
+    
+    def generate_response(self, message):
+        """Generate empathetic AI response"""
+        if not self.conversation_model:
+            return "I'm here to listen. Would you like to share more about how you're feeling?"
+        
+        try:
+            # Contextual prompt to guide response
+            full_prompt = f"You are a compassionate AI therapist. Respond supportively to this message: {message}. Be empathetic, validate feelings, and avoid giving direct medical advice."
+            
+            # Generate response
+            response = self.conversation_model(
+                full_prompt, 
+                max_length=100, 
+                num_return_sequences=1
+            )[0]['generated_text']
+            
+            return response
+        
+        except Exception as e:
+            return "I'm here to listen. Would you like to share more about how you're feeling?"
+    
+    def generate_summary(self, conversation):
+        """Generate a professional therapy-style summary"""
+        if not self.summary_model:
+            return "Summary generation is temporarily unavailable."
+        
+        try:
+            # Generate summary
+            summary = self.summary_model(
+                conversation, 
+                max_length=130, 
+                min_length=30, 
+                do_sample=False
+            )[0]['summary_text']
+            
+            return summary
+        
+        except Exception as e:
+            return "Summary could not be generated."
+
+def main():
+    st.set_page_config(
+        page_title="SafeSpace: AI Compassionate Listener", 
+        page_icon="🧠"
+    )
+    
+    st.title("🧠 SafeSpace: Compassionate AI Listener")
+    st.write("A supportive space to share your feelings safely")
+    
+    # Initialize session state
+    if 'conversation' not in st.session_state:
+        st.session_state.conversation = []
+    
+    if 'assistant' not in st.session_state:
+        st.session_state.assistant = AITherapistAssistant()
+    
+    # Conversation Display
+    for message in st.session_state.conversation:
+        if message['sender'] == 'user':
+            st.chat_message("user").write(message['text'])
         else:
-            st.success("Message received. Thank you for sharing.")
-    else:
-        st.warning("Please enter a message.")
-
-# Option to generate a therapist summary of the conversation
-if st.button("Generate Therapist Summary"):
+            st.chat_message("assistant").write(message['text'])
+    
+    # User Input
+    if prompt := st.chat_input("Share your thoughts. I'm here to listen."):
+        # Check for crisis indicators
+        if st.session_state.assistant.detect_crisis(prompt):
+            st.warning("⚠️ Crisis Support Detected")
+            st.markdown("**Immediate Support Resources:**")
+            for org, phone in CRISIS_RESOURCES.items():
+                st.markdown(f"- {org}: `{phone}`")
+        
+        # Add user message
+        st.session_state.conversation.append({
+            'sender': 'user', 
+            'text': prompt
+        })
+        
+        # Display user message
+        st.chat_message("user").write(prompt)
+        
+        # Generate AI response
+        with st.chat_message("assistant"):
+            with st.spinner("Listening and reflecting..."):
+                ai_response = st.session_state.assistant.generate_response(prompt)
+                st.write(ai_response)
+        
+        # Add AI message
+        st.session_state.conversation.append({
+            'sender': 'ai', 
+            'text': ai_response
+        })
+    
+    # Session Summary Generation
     if st.session_state.conversation:
-        conversation_text = "\n".join(st.session_state.conversation)
-        summary = summarizer(conversation_text, max_length=100, min_length=30, do_sample=False)
-        st.subheader("Therapist Summary")
-        st.write(summary[0]['summary_text'])
-    else:
-        st.info("No conversation to summarize yet.")
+        if st.button("Generate Session Summary"):
+            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.write(summary)
+    
+    # Safety Information
+    st.sidebar.title("🆘 Crisis Support")
+    st.sidebar.markdown("**Remember: You are valued and supported.**")
+    for org, phone in CRISIS_RESOURCES.items():
+        st.sidebar.markdown(f"- {org}: `{phone}`")
+
+if __name__ == "__main__":
+    main()
+
+# requirements.txt
+# streamlit
+# transformers
+# torch
+# accelerate