app.py

import streamlit as st
from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch
import pandas as pd
import numpy as np
import plotly.express as px
import time

# Try to import streamlit_lottie, but provide fallback if it fails
try:
    from streamlit_lottie import st_lottie
    import requests
    def load_lottie_url(url: str):
        try:
            r = requests.get(url)
            if r.status_code != 200:
                return None
            return r.json()
        except:
            return None
    LOTTIE_AVAILABLE = True
except ImportError:
    LOTTIE_AVAILABLE = False

# Page configuration
st.set_page_config(
    page_title="MindBERT - Mental Health Analysis",
    page_icon="🧠",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS
st.markdown("""
<style>
    .main-header {
        font-size: 2.5rem;
        color: #4527A0;
        text-align: center;
        margin-bottom: 1rem;
    }
    .sub-header {
        font-size: 1.5rem;
        color: #5E35B1;
        margin-bottom: 1rem;
    }
    .result-box {
        background-color: #F3F4F6;
        border-radius: 10px;
        padding: 20px;
        margin: 10px 0px;
        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
    }
    .footer {
        text-align: center;
        margin-top: 3rem;
        padding: 1rem;
        font-size: 0.8rem;
        color: #666;
        border-top: 1px solid #ddd;
    }
    .stTextArea textarea {
        border-radius: 10px;
        border: 1px solid #5E35B1;
    }
    .stButton button {
        border-radius: 20px;
        background-color: #5E35B1;
        color: white;
        font-weight: bold;
        padding: 0.5rem 2rem;
        transition: all 0.3s;
    }
    .stButton button:hover {
        background-color: #4527A0;
        box-shadow: 0 4px 8px rgba(0, 0, 0, 0.2);
    }
    .category-card {
        padding: 10px;
        border-radius: 5px;
        margin-bottom: 10px;
    }
    .helper-text {
        font-size: 0.9rem;
        color: #666;
        font-style: italic;
    }
    .metric-value {
        font-size: 2.5rem;
        font-weight: bold;
        text-align: center;
    }
    .metric-label {
        font-size: 1rem;
        text-align: center;
        color: #666;
    }
    .brain-icon {
        font-size: 5rem;
        text-align: center;
        margin-bottom: 1rem;
        color: #5E35B1;
    }
</style>
""", unsafe_allow_html=True)

# Create sidebar
with st.sidebar:
    # Use either Lottie or a simple icon
    if LOTTIE_AVAILABLE:
        # Fixed Lottie URLs that are reliable
        brain_animation = load_lottie_url("https://lottie.host/2eb12c32-787a-46f7-ac20-34c166d1a285/UcEEbJlFVH.json")
        if brain_animation:
            st_lottie(brain_animation, height=200, key="brain_animation")
        else:
            st.markdown("<div class='brain-icon'>🧠</div>", unsafe_allow_html=True)
    else:
        st.markdown("<div class='brain-icon'>🧠</div>", unsafe_allow_html=True)
    
    st.markdown("## About MindBERT")
    st.info(
        "MindBERT is a fine-tuned BERT model specifically designed to detect "
        "mental health states from text. It has been trained on a diverse dataset "
        "of mental health-related content to identify patterns associated with "
        "various mental health conditions."
    )
    st.markdown("## Disclaimer")
    st.warning(
        "This app is for educational and research purposes only. "
        "It is not a substitute for professional medical advice, "
        "diagnosis, or treatment. Always seek the advice of a qualified "
        "health provider for any medical condition."
    )
    
    st.markdown("## How it works")
    st.markdown(
        "1. Enter text in the provided area\n"
        "2. Click 'Analyze Mental State'\n"
        "3. The model will process the text and predict the writer's mental state\n"
        "4. Results are displayed with confidence levels"
    )
    
    

# Main content
st.markdown("<h1 class='main-header'>MindBERT - Mental Health Analysis</h1>", unsafe_allow_html=True)

# Custom tabs
tab1, tab2, tab3 = st.tabs(["Mental Health Analyzer", "Understanding Categories", "Resources"])

with tab1:
    st.markdown("<p class='helper-text'>Enter text to analyze the mental state of the writer.</p>", unsafe_allow_html=True)
    
    # Text input area with placeholder
    user_input = st.text_area(
        "Type your message here:",
        height=150,
        placeholder="Example: I've been feeling overwhelmed lately with all the pressure at work. I can't seem to focus and I'm constantly worried about deadlines.",
    )
    
    # Model loading feedback
    @st.cache_resource
    def load_model():
        try:
            tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
            model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH)
            return tokenizer, model, True
        except Exception as e:
            st.error(f"Error loading model: {str(e)}")
            return None, None, False
    
    # Define model and tokenizer paths from Hugging Face
    MODEL_PATH = "DrSyedFaizan/mindBERT"
    
    with st.spinner("Loading model..."):
        tokenizer, model, model_loaded = load_model()
    
    # Analysis button
    col1, col2, col3 = st.columns([1, 2, 1])
    with col2:
        analyze_button = st.button("Analyze Mental State")
    
    # Prediction logic
    if analyze_button:
        if not model_loaded:
            st.error("Model failed to load. Please try again later.")
        elif not user_input.strip():
            st.warning("Please enter some text for analysis.")
        else:
            # Show analyzing animation or spinner
            with st.spinner("Analyzing..."):
                if LOTTIE_AVAILABLE:
                    analyzing_animation = load_lottie_url("https://lottie.host/16c400ec-7d59-4c0c-a84b-56c9134cd673/20XZXacKUS.json")
                    if analyzing_animation:
                        st_lottie(analyzing_animation, height=200, key="analyze_animation", speed=1.5)
                
                # Add a slight delay to show the animation
                time.sleep(1)
                
                try:
                    # Tokenize input
                    inputs = tokenizer(user_input, return_tensors="pt", truncation=True, padding=True)
                    
                    # Make prediction
                    with torch.no_grad():
                        outputs = model(**inputs)
                        logits = outputs.logits
                        probabilities = torch.nn.functional.softmax(logits, dim=1)[0]
                        predicted_class = torch.argmax(logits, dim=1).item()
                    
                    # Mapping predicted class to mental state with descriptions
                    label_map = {
                        0: {"name": "Anxiety", "color": "#FFD54F", "description": "Characterized by excessive worry, fear, or nervousness."},
                        1: {"name": "Bipolar", "color": "#FF7043", "description": "Featuring alternating periods of depression and mania or elevated mood."},
                        2: {"name": "Depression", "color": "#4FC3F7", "description": "Persistent feelings of sadness, hopelessness, and loss of interest."},
                        3: {"name": "Normal", "color": "#81C784", "description": "Balanced emotional state without significant mental health concerns."},
                        4: {"name": "Personality Disorder", "color": "#9575CD", "description": "Persistent patterns of thinking and behavior that deviate from social norms."},
                        5: {"name": "Stress", "color": "#FF8A65", "description": "Physical or emotional tension due to challenging circumstances."},
                        6: {"name": "Suicidal", "color": "#F44336", "description": "Thoughts or intentions of self-harm or taking one's own life."}
                    }
                    
                    mental_state = label_map.get(predicted_class, {"name": "Unknown", "color": "#BDBDBD", "description": "Unable to classify the mental state."})
                    
                    # Create data for visualization
                    all_probs = {label_map[i]["name"]: prob.item() * 100 for i, prob in enumerate(probabilities)}
                    probs_df = pd.DataFrame(list(all_probs.items()), columns=["Mental State", "Confidence (%)"])
                    probs_df = probs_df.sort_values("Confidence (%)", ascending=False)
                    
                    # Display results
                    st.markdown("<div class='result-box'>", unsafe_allow_html=True)
                    
                    # Primary result
                    col1, col2 = st.columns([1, 2])
                    with col1:
                        st.markdown(f"<div class='metric-value' style='color:{mental_state['color']}'>{mental_state['name']}</div>", unsafe_allow_html=True)
                        st.markdown("<div class='metric-label'>Primary Detection</div>", unsafe_allow_html=True)
                    
                    with col2:
                        st.markdown(
                        f"""
                        <div style='background-color:{mental_state['color']}20; padding:15px; border-radius:10px; border-left:5px solid {mental_state['color']}'>
                        <b>{mental_state['name']}</b>: {mental_state['description']}
                        </div>
                        """,
                        unsafe_allow_html=True
                        )

                    
                    # Confidence scores visualization
                    st.markdown("<h3 class='sub-header'>Confidence Analysis</h3>", unsafe_allow_html=True)
                    
                    # Create bar chart
                    fig = px.bar(
                        probs_df, 
                        x="Confidence (%)", 
                        y="Mental State",
                        orientation="h",
                        color="Mental State",
                        color_discrete_map={
                            "Anxiety": "#FFD54F",
                            "Bipolar": "#FF7043",
                            "Depression": "#4FC3F7",
                            "Normal": "#81C784",
                            "Personality Disorder": "#9575CD",
                            "Stress": "#FF8A65",
                            "Suicidal": "#F44336",
                            "Unknown": "#BDBDBD"
                        }
                    )
                    fig.update_layout(
                        height=350,
                        margin=dict(l=20, r=20, t=30, b=20),
                        xaxis_title="Confidence (%)",
                        yaxis_title="",
                        yaxis=dict(autorange="reversed"),
                        xaxis=dict(range=[0, 100])
                    )
                    st.plotly_chart(fig, use_container_width=True)
                    
                    # Warning for high-risk categories
                    if mental_state["name"] in ["Suicidal", "Depression"] and all_probs[mental_state["name"]] > 50:
                        st.warning(
                            "⚠️ **High-risk mental state detected.** If you or someone you know is experiencing "
                            "suicidal thoughts, please seek immediate professional help or call the National "
                            "Suicide Prevention Lifeline at 988 or 1-800-273-8255."
                        )
                    
                    st.markdown("</div>", unsafe_allow_html=True)
                    
                    # Suggestion based on detected mental state
                    suggestion_map = {
                        "Anxiety": "Consider breathing exercises, meditation, or consulting with a mental health professional about anxiety management techniques.",
                        "Bipolar": "Regular sleep schedules and medication management with professional oversight can help stabilize mood swings.",
                        "Depression": "Regular physical activity, social connection, and professional therapy can be beneficial for managing depression.",
                        "Normal": "Continue maintaining a healthy lifestyle with regular exercise, good sleep habits, and social connections.",
                        "Personality Disorder": "Long-term psychotherapy with a specialist in personality disorders is often recommended.",
                        "Stress": "Stress reduction techniques such as mindfulness, time management, and setting boundaries can be helpful.",
                        "Suicidal": "Please seek immediate professional help. Call the National Suicide Prevention Lifeline at 988 or 1-800-273-8255."
                    }
                    
                    st.markdown("<div class='result-box'>", unsafe_allow_html=True)
                    st.markdown("<h3 class='sub-header'>Suggestions</h3>", unsafe_allow_html=True)
                    st.info(suggestion_map.get(mental_state["name"], "Consider consulting with a mental health professional for personalized guidance."))
                    st.markdown("</div>", unsafe_allow_html=True)
                    
                except Exception as e:
                    st.error(f"Error during analysis: {str(e)}")
                    st.info("Please try again with different text or contact support if the issue persists.")

with tab2:
    st.markdown("<h3 class='sub-header'>Mental Health Categories Explained</h3>", unsafe_allow_html=True)
    
    categories = [
        {"name": "Anxiety", "color": "#FFD54F", "description": "Characterized by excessive worry, fear, and nervousness that may interfere with daily activities. May include physical symptoms like restlessness, fatigue, and difficulty concentrating."},
        {"name": "Bipolar", "color": "#FF7043", "description": "Features alternating periods of depression and mania. During manic episodes, individuals may experience elevated mood, increased energy, and impulsive behavior. Depressive episodes involve symptoms similar to major depression."},
        {"name": "Depression", "color": "#4FC3F7", "description": "Persistent feelings of sadness, hopelessness, and loss of interest in activities once enjoyed. May include changes in sleep, appetite, energy level, concentration, and self-worth."},
        {"name": "Normal", "color": "#81C784", "description": "A balanced emotional state characterized by appropriate responses to life events, ability to manage stress effectively, and maintain healthy relationships and activities."},
        {"name": "Personality Disorder", "color": "#9575CD", "description": "Long-term patterns of thoughts, feelings, and behaviors that deviate from cultural expectations, causing distress or problems functioning in social, work, or personal relationships."},
        {"name": "Stress", "color": "#FF8A65", "description": "The body's response to demands or pressures. While acute stress can be motivating, chronic stress may lead to physical and mental health problems including headaches, sleep disturbances, and mood changes."},
        {"name": "Suicidal", "color": "#F44336", "description": "Thoughts about, planning, or intent to end one's life. This is a medical emergency requiring immediate professional intervention."}
    ]
    
    for category in categories:
        st.markdown(
            f"<div class='category-card' style='background-color:{category['color']}20; border-left:5px solid {category['color']}'>"+
            f"<h4 style='color:{category['color']}DD'>{category['name']}</h4>"+
            f"<p>{category['description']}</p>"+
            "</div>",
            unsafe_allow_html=True
        )
    
    st.markdown(
        "<div class='helper-text'><b>Note:</b> These categories are simplified for educational purposes. "
        "Mental health exists on a spectrum, and professional diagnosis involves comprehensive assessment "
        "beyond text analysis.</div>",
        unsafe_allow_html=True
    )

with tab3:
    st.markdown("<h3 class='sub-header'>Mental Health Resources</h3>", unsafe_allow_html=True)
    
    col1, col2 = st.columns(2)
    
    with col1:
        st.markdown(
            "<div class='result-box'>"+
            "<h4>Crisis Resources</h4>"+
            "<ul>"+
            "<li><b>National Suicide Prevention Lifeline:</b> 988 or 1-800-273-8255</li>"+
            "<li><b>Crisis Text Line:</b> Text HOME to 741741</li>"+
            "<li><b>Veterans Crisis Line:</b> 1-800-273-8255 (Press 1)</li>"+
            "<li><b>Disaster Distress Helpline:</b> 1-800-985-5990</li>"+
            "</ul>"+
            "</div>",
            unsafe_allow_html=True
        )
    
    with col2:
        st.markdown(
            "<div class='result-box'>"+
            "<h4>Online Resources</h4>"+
            "<ul>"+
            "<li><b>National Alliance on Mental Illness (NAMI):</b> nami.org</li>"+
            "<li><b>Mental Health America:</b> mhanational.org</li>"+
            "<li><b>Psychology Today Therapist Finder:</b> psychologytoday.com/us/therapists</li>"+
            "<li><b>Substance Abuse and Mental Health Services Administration:</b> samhsa.gov</li>"+
            "</ul>"+
            "</div>",
            unsafe_allow_html=True
        )
    
    st.markdown(
        "<div class='result-box'>"+
        "<h4>Self-Help Strategies</h4>"+
        "<ul>"+
        "<li><b>Mindfulness and Meditation:</b> Apps like Headspace, Calm, or Insight Timer</li>"+
        "<li><b>Physical Activity:</b> Regular exercise can help reduce symptoms of depression and anxiety</li>"+
        "<li><b>Sleep Hygiene:</b> Maintaining regular sleep patterns supports mental health</li>"+
        "<li><b>Social Connection:</b> Staying connected with supportive people in your life</li>"+
        "<li><b>Limiting Alcohol and Substances:</b> These can worsen mental health symptoms</li>"+
        "<li><b>Setting Boundaries:</b> Learning to say no and protecting your mental space</li>"+
        "</ul>"+
        "</div>",
        unsafe_allow_html=True
    )

# Footer
st.markdown("<div class='footer'>Developed by Dr. Syed Faizan using MindBERT on Hugging Face. This application is for educational and research purposes only.</div>", unsafe_allow_html=True)