app.py

import streamlit as st
import os.path
import re
import pandas as pd
from transformers import pipeline, Pipeline
from time import sleep

ID2LABEL = {
    'LABEL_0': {
        "Cause": "No Reason",
        "description": "There is no reason that identifies the cause of mental disorder, or the text does not reflect a mental disorder",
    },
    'LABEL_1': {
        "Cause": "Bias / Abuse",
        "description": """
            A strong inclination of the mind or a preconceived opinion about something or someone. To avoid someone intentionally,
            or to prevent someone from taking part in the social activities of a group because they dislike the person or disapprove
            of their activities. It includes body shaming, physical, sexual, or emotional abuse.
        """,
    },
    'LABEL_2': {
        "Cause": "Jobs and Career",
        "description": """
            Financial loss can have catastrophic effects on mental illness, relationships and even physical health. Poor, meaningless
            and unmanageable education, unemployment, un-affordable home loans, poor financial advice, and losing a job are some of
            the major concerns. It includes gossiping and/or social cliques, aggressive bullying behavior, poor communication 
            and unclear expectations, dictatorial management techniques that don’t embrace employee feedback. The educational problems
            like picking up courses under some external pressure and poor grades are also part of this category.
        """,
    },
    'LABEL_3': {
        "Cause": "Medication",
        "description": """
            The general drugs and other antiviral drugs can increase the risk of depression. The habit of using substances and alcohols
            can aggravate the problem of mental disorders. Moreover, medical problems like tumors, cancer, and other prolonged diseases
            can boost the presence of mental depression.
        """,
    },
    'LABEL_4': {
        "Cause": "Relationships",
        "description": """
            When two people or a group of people fight, it may lead to a relationship or friendship drifting apart, for example, regular
            fights, breakups, divorce, mistrust, jealousy, betrayal, difference in opinion, inconsistency, conflicts, bad company,
            noncommitment, priority, envy. Problems like bad parenting and childhood trauma are also part of this category.
        """,
    },
    'LABEL_5': {
        "Cause": "Alienation",
        "description": """
            Alienation is the feeling of life being worthless even after doing everything. There may be indicators of meaninglessness,
            loneliness, tired of daily routines, powerlessness, normlessness, isolation, and cultural estrangement.
        """,
    },
}

EXAMPLES = [
    """
        Same dad, different day. I can't believe that my dad have no sense of humanity. He hit me in the head this morning.
        I felt nauseous all day, including when I take my exam just before this. If I kill myself, I will make sure that my dad
        got the most blame.
    """,
    """
        My boss laid me off today. He said that my company was downsizing. I don't believe his cr*p though, I think he just hates
        me as the only women in my department. Now I don't know how to feed my four kids. I am officially, totally, completely, out
        of money and will to live.
    """,
    """
        Last month, my doctor prescribed me some alprazolam to calm my nerves down. I churn through a bottle of the pill in a month.
        My doctor today told me that I'm quite healthy to go without some calming drugs, but I just can't stop consuming them. I
        just bought like 4 bottles of it through some shady middle-man and I'll go crazy if I went for half a day without
        swallowing one.
    """,
    """
        My girlfriend dumped me because of some stupid nerdy dude at her office, my mother disowns me for not enlisting to the
        millitary like my brother did. It all just keeps on pinning me to the ground. When I asked my friends to go out, they all
        refused because my ex-girlfriend was spreading lies of how I cheated on her yada yada. I'm totally f*cked. 
    """,
    """
        Everything is worthless, everything is meaningless. All the things that I do literally contribute nothing to the
        society. I just want to go somewhere I can just lie down, sleep, eat, with no negative consequences for me.
    """,
    """
        Yesterday, I bought an ice cream for myself at the city. It was really good. I will definitely tell everyone
        I know about this ice cream place.
    """,
]

@st.cache_resource(show_spinner=False)
def load_model() -> Pipeline:
    model = "AIMH/mental-longformer-base-4096"
    tokenizer = "AIMH/mental-longformer-base-4096"
    if os.path.isfile("model/model.safetensors"):
        model = "model"
    if os.path.isfile("tokenizer/tokenizer.json"):
        tokenizer = "tokenizer"
    return pipeline("text-classification", model=model, tokenizer=tokenizer)

@st.cache_data(show_spinner=False)
def predict(text: str, _pipe: Pipeline):
    res = _pipe(text, return_all_scores=True)[0]
    res = sorted(res, key=lambda d: d['score'], reverse=True)
    res = [
        dict(ID2LABEL[x["label"]], **{"Confidence": round(x["score"] * 100)})
    for x in res]
    return pd.DataFrame(res)
    # return res


if __name__ == "__main__":
    st.markdown('<h1 align="center">✨ Depression Causal Analysis ✨</h1>', unsafe_allow_html=True, )
    st.error("""
        DISCLAIMER: This project was only intended for research showcase purposes only. 
        If you believe that you have mental health issues, please consult your physician.\n
        This project also contains triggering example words from social media that might not
        fit everybody. Continue with caution. Please love yourself and don't hesitate to reach
        out for professional help. ❤
    """)

    left, center, right = st.columns(3)
    left.link_button("Go to project on GitHub", "https://github.com/stackofsugar/", use_container_width=True)
    center.link_button("Read the dataset's paper", "https://arxiv.org/abs/2207.04674v1", use_container_width=True)
    right.link_button("Read the model's paper", "https://arxiv.org/abs/2304.10447v1", use_container_width=True)
        
    example = st.selectbox("Load an example", EXAMPLES, index=None)

    with st.form("main_prediction"):
        text = st.text_area(
            "Text to analyze (tip: this model is better for long texts)",
            value=(re.sub(' +', ' ', example).strip().replace("\n", "") if example else ""),
            height=200
        )

        left, middle, right = st.columns(3)
        submitted = middle.form_submit_button("Predict!", use_container_width=True)

        if submitted and text.strip():
            pipe = None
            with st.spinner("Loading model..."):
                pipe = load_model()
            with st.spinner("Predicting..."):
                preds = predict(text.strip(), pipe)
            st.markdown(f'**Result: {preds.at[0, "Cause"]}**')

            col1, col2 = st.columns([1,2])
            # col1, col2 = st.columns(2)
            col1.dataframe(preds[["Cause", "Confidence"]],
                        column_config={
                            "Confidence": st.column_config.NumberColumn(
                                format="%d%%"
                            )
                        },
                        hide_index=True, 
                        use_container_width=True)

            col2.markdown(f'Explanation of **{preds.at[0, "Cause"]}**:')
            col2.write(preds.at[0, "description"])
    
    st.write(
        """Developed with :heart: by [stackofsugar](https://github.com/stackofsugar/). For more information on the
        project, please visit the project's GitHub page with the button above."""
        )