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('

✨ Depression Causal Analysis ✨

', 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.""" )