add files

app.py

@@ -0,0 +1,145 @@
+
+import os
+import pandas as pd
+import streamlit as st
+from q_learning_chatbot import QLearningChatbot
+from xgb_mental_health import MentalHealthClassifier
+from bm25_retreive_question import QuestionRetriever
+from llm_response_generator import LLLResponseGenerator
+
+# Streamlit UI
+st.title("FOMO Fix - RL-based Mental Health Assistant")
+
+# Define states and actions
+states = ['Negative', 'Moderately Negative', 'Neutral', 'Moderately Positive', 'Positive']
+actions = ['encouragement', 'empathy']
+
+# Initialize Q-learning chatbot and mental health classifier
+chatbot = QLearningChatbot(states, actions)
+
+# Initialize MentalHealthClassifier
+data_path = "data.csv"
+tokenizer_model_name = "nlptown/bert-base-multilingual-uncased-sentiment"
+mental_classifier_model_path = 'mental_health_model.pkl'
+mental_classifier = MentalHealthClassifier(data_path, mental_classifier_model_path)
+
+
+# Function to display Q-table
+def display_q_table(q_values, states, actions):
+    q_table_dict = {'State': states}
+    for i, action in enumerate(actions):
+        q_table_dict[action] = q_values[:, i]
+
+    q_table_df = pd.DataFrame(q_table_dict)
+    return q_table_df
+
+
+# Initialize memory
+if 'entered_text' not in st.session_state:
+    st.session_state.entered_text = []
+if 'entered_mood' not in st.session_state:
+    st.session_state.entered_mood = []
+
+# Collect user input
+user_message = st.text_input("Type your message here:")
+
+# Take user input
+if user_message:
+    st.session_state.entered_text.append(user_message)
+
+    # Detect mental condition
+    mental_classifier.initialize_tokenizer(tokenizer_model_name)
+    mental_classifier.preprocess_data()
+    predicted_mental_category = mental_classifier.predict_category(user_message)
+    print("Predicted mental health condition:", predicted_mental_category)
+    # st.subheader("🛑 " + f"{predicted_mental_category.capitalize()}")
+
+    # Retrieve question
+    retriever = QuestionRetriever()
+    question = retriever.get_response(user_message, predicted_mental_category)
+    # st.write(question)
+
+    # Detect sentiment
+    user_sentiment = chatbot.detect_sentiment(user_message)
+    # Update mood history / moode_trend
+    chatbot.update_mood_history()
+    mood_trend = chatbot.check_mood_trend()
+
+    # Define rewards
+    if user_sentiment in ["Positive", "Moderately Positive"]:
+        if mood_trend == "increased":
+            reward = +0.8
+        else: # decresed
+            reward = -0.3
+
+    else:
+        if mood_trend == "increased":
+            reward = +1
+        else:
+            reward = -1
+
+    print(f"mood_trend - sentiment - reward: {mood_trend} - {user_sentiment} - 🛑{reward}🛑 -- (a)")
+
+    # Update Q-values
+    chatbot.update_q_values(user_sentiment, chatbot.actions[0], reward, user_sentiment)
+
+    # Get recommended action based on the updated Q-values
+    ai_tone = chatbot.get_action(user_sentiment)
+    print(ai_tone)
+
+
+
+    #--------------
+    # LLM Response Generator
+    HUGGINGFACEHUB_API_TOKEN = os.getenv('HUGGINGFACEHUB_API_TOKEN')
+
+    llm_model = LLLResponseGenerator()
+    temperature = 0.1
+    max_length = 128
+    template = """INSTRUCTIONS: {context}
+
+        Respond to the user with a tone of {ai_tone}.
+
+        Question asked to the user: {question}
+
+        Response by the user: {user_text}
+        Response;
+        """
+    context = "You are a mental health supporting non-medical assistant. Provide some advice and ask a relevant question back to the user."
+
+
+    llm_response = llm_model.llm_inference(
+        model_type="huggingface",
+        question=question,
+        prompt_template=template,
+        context=context,
+        ai_tone=ai_tone,
+        questionnaire=predicted_mental_category,
+        user_text=user_message,
+        temperature=temperature,
+        max_length=max_length,
+        )
+
+
+    st.write(f"{llm_response}")
+    st.write(f"{question}")
+
+    st.subheader("Behind the Scence - What AI is doing:")
+    st.write(f"- User Tone: {user_sentiment}, Possibly {predicted_mental_category.capitalize()}")
+    st.write(f"- AI Tone: {ai_tone.capitalize()}")
+    # st.write(f"Question: {question}")
+
+
+    # Display results
+    # st.subheader(f"{user_sentiment.capitalize()}")
+    # st.write("->" + f"{ai_tone.capitalize()}")
+    # st.write(f"Mood {chatbot.check_mood_trend()}")
+    # st.write(f"{ai_tone.capitalize()}, {chatbot.check_mood_trend()}")
+
+    # Display Q-table
+    st.dataframe(display_q_table(chatbot.q_values, states, actions))
+
+    # Display mood history
+    # st.subheader("Mood History (Recent 5):")
+    # for mood_now in reversed(chatbot.mood_history[-5:]): #st.session_state.entered_mood[-5:], chatbot.mood_history[-5:]): #st.session_state.entered_text[-5:]
+    #     st.write(f"{mood_now}")

bm25_retreive_question.py

@@ -0,0 +1,138 @@
+from rank_bm25 import BM25Okapi
+import nltk
+from nltk.tokenize import word_tokenize
+
+# Download NLTK data for tokenization
+nltk.download('punkt')
+
+class QuestionRetriever:
+    def __init__(self):
+        self.depression_questions = [
+            "How often have you felt persistently low in mood or sad for most of the day?",
+            "How often have you lost interest or pleasure in activities you used to enjoy?",
+            "How often have you experienced significant changes in appetite or weight (up or down)?",
+            "How often have you had trouble sleeping or sleeping too much?",
+            "How often have you felt restless or slowed down most of the time?",
+            "How often have you felt worthless or excessively guilty, even for minor things?",
+            "How often have you had difficulty thinking, concentrating, or making decisions?",
+            "How often have you had recurrent thoughts of death or suicide?",
+            "How often have you felt hopeless or like there's no point in life?",
+            "How often have you felt isolated or withdrawn from others due to your mood?",
+            "How often have you felt down, depressed, or hopeless?",
+            "How often have you had little interest or pleasure in doing things over the past two weeks?",
+            "Have you experienced changes in your appetite or weight due to feeling depressed in the last two weeks?",
+            "How often have you had trouble falling asleep, staying asleep, or sleeping too much because of feeling depressed?",
+            "Have you felt tired or had little energy most days over the past two weeks due to depression?",
+            "How often have you felt bad about yourself or that you are a failure or have let yourself or your family down in the last two weeks?",
+            "Have you had trouble concentrating on things like reading, watching TV, or engaging in conversations due to feeling depressed?",
+            "How often have you thought that you would be better off dead or hurting yourself in some way over the past two weeks?",
+            "Do you feel restless or slowed down physically most days because of feeling depressed in the last two weeks?",
+            "How often have you found it difficult to make decisions or felt indecisive due to feeling depressed recently?",
+            "How often do you feel sad or hopeless?",
+            "Do you have trouble getting out of bed or motivating yourself to do things?",
+            "Have you experienced a loss of interest in activities or hobbies?",
+            "Do you feel like you're not good enough or that you've failed at things?",
+            "Have you noticed any changes in your appetite or sleep patterns?",
+            "Do you feel like you're isolated from others or that you don't have any support?",
+            "Have you experienced any thoughts of self-harm or suicide?",
+            "Do you feel like you're stuck in a rut or that you can't see a way out of your current situation?",
+            "Have you noticed any physical symptoms like fatigue or loss of energy?",
+            "Have you considered seeking professional help for your depression?"
+        ]
+
+        self.adhd_questions = [
+            "How often do you find yourself having trouble focusing on tasks or activities?",
+            "How often do you easily get sidetracked by noises, movements, or unrelated thoughts?",
+            "How often do you have difficulty finishing tasks or following through on instructions?",
+            "How often do you feel restless or fidgety, often unable to sit still for long periods?",
+            "How often do you talk excessively or blurt out things before thinking?",
+            "How often do you struggle with waiting for your turn or interrupting others frequently?",
+            "How often do you lose things you need for work or school (keys, phone, etc.)?",
+            "How often are you forgetful or disorganized?",
+            "How often do you have difficulty planning ahead or organizing tasks?",
+            "How often do you have difficulty controlling strong emotions or impulsive behaviors?",
+            "How often have you found it difficult to pay attention to details or make careless mistakes in your work or activities due to ADHD symptoms?",
+            "How often have you had trouble staying focused on tasks or activities like work, school, or hobbies because of ADHD symptoms in the last two weeks?",
+            "Have you experienced difficulties in organizing tasks and activities, such as messy workspaces or missed deadlines, because of ADHD recently?",
+            "How often have you avoided or been reluctant to engage in tasks that require sustained mental effort due to ADHD symptoms over the past two weeks?",
+            "Have you been forgetful in daily activities, such as missing appointments or losing items, because of ADHD symptoms in the last two weeks?",
+            "How often have you felt restless or fidgety in situations where it is inappropriate due to ADHD symptoms recently?",
+            "Do you find it challenging to engage in leisure activities quietly or feel driven by a motor due to ADHD most days?",
+            "How often have you interrupted others or blurted out answers before questions have been completed because of ADHD symptoms over the past two weeks?",
+            "Have you had difficulties waiting your turn or remaining seated in situations where it is expected due to ADHD symptoms recently?",
+            "How often have you felt impatient or had difficulty engaging in activities quietly due to ADHD symptoms in the last two weeks?",
+            "How often do you have trouble sitting still or staying focused?",
+            "Do you have trouble following instructions or completing tasks?",
+            "Have you noticed any difficulty with organization or time management?",
+            "Do you find yourself easily distracted or forgetful?",
+            "Have you experienced any impulsivity or difficulty with self-control?",
+            "Do you have trouble waiting your turn or behaving in line with social rules?",
+            "Have you noticed any changes in your mood or emotional state?",
+            "Do you have trouble with word retrieval or difficulty finding the right words?",
+            "Have you experienced any difficulty with reading or comprehension?",
+            "Have you considered seeking professional help for your ADHD?"
+        ]
+
+        self.anxiety_questions = [
+            "How often have you been feeling excessive worry or nervousness, even about everyday things?",
+            "How often do you experience physical symptoms like racing heart, sweating, or shortness of breath when feeling anxious?",
+            "How often do you encounter certain situations or triggers that cause you significant anxiety?",
+            "How often do you find yourself avoiding places or activities due to anxiety?",
+            "How often have you had trouble sleeping or concentrating because of anxious thoughts?",
+            "How often do you feel a constant need to be in control or have things perfect to avoid anxiety?",
+            "How often do you have intrusive thoughts that are difficult to stop?",
+            "How often do you experience sudden feelings of intense fear or panic (panic attacks)?",
+            "How often has your anxiety significantly impacted your daily life or relationships?",
+            "How much does anxiety interfere with your daily life?",
+            "How often have you felt nervous, anxious, or on edge?",
+            "How often have you found it difficult to stop or control worrying thoughts in the last two weeks?",
+            "Have you experienced restlessness or felt keyed up or on edge due to anxiety in the past two weeks?",
+            "How often have you felt easily fatigued or had difficulty concentrating because of anxiety recently?",
+            "Have you noticed irritability or muscle tension as a result of feeling anxious over the past two weeks?",
+            "How often have you had trouble falling asleep, staying asleep, or restless sleep due to anxiety in the last two weeks?",
+            "Do you find yourself easily startled or feeling on edge most days because of anxiety?",
+            "How often have you experienced physical symptoms like sweating, trembling, or a racing heart due to anxiety recently?",
+            "Have you found it challenging to relax or felt restless most days because of anxiety in the past two weeks?",
+            "How often have you felt a sense of impending doom or danger because of anxiety over the past two weeks?",
+            "How often do you feel nervous or on edge?",
+            "Do you have trouble relaxing or controlling your worries?",
+            "Have you experienced a racing or irregular heartbeat?",
+            "Do you feel like you're constantly on guard or on high alert?",
+            "Have you noticed any physical symptoms like trembling or sweating?",
+            "Do you feel like you're having trouble concentrating or making decisions?",
+            "Have you experienced any feelings of detachment or disconnection from others?",
+            "Do you feel like you're having trouble sleeping or experiencing vivid dreams?",
+            "Have you noticed any changes in your appetite or eating habits?",
+            "Have you considered seeking professional help for your anxiety?"
+        ]
+
+    def get_response(self, user_query, predicted_mental_category):
+        if predicted_mental_category == "depression":
+            knowledge_base = self.depression_questions
+        elif predicted_mental_category == "adhd":
+            knowledge_base = self.adhd_questions
+        elif predicted_mental_category == "anxiety":
+            knowledge_base = self.anxiety_questions
+        else:
+            print("Sorry, I didn't understand that.")
+
+        tokenized_docs = [word_tokenize(doc.lower()) for doc in knowledge_base]  # Ensure lowercase for consistency
+        bm25 = BM25Okapi(tokenized_docs)
+        tokenized_query = word_tokenize(user_query.lower())  # Ensure lowercase for consistency
+        doc_scores = bm25.get_scores(tokenized_query)
+
+        # Get the index of the most relevant document
+        most_relevant_doc_index = doc_scores.argmax()
+
+        # Fetch the corresponding response from the knowledge base
+        response = knowledge_base[most_relevant_doc_index]
+        return response
+
+if __name__ == "__main__":
+    knowledge_base = depression_questions
+
+    model = QuestionRetriever(knowledge_base)
+    user_input = input("User: ")
+
+    response = model.get_response(user_input)
+    print("Chatbot:", response)

llm_response_generator.py

@@ -0,0 +1,152 @@
+import os
+from langchain_community.llms import HuggingFaceHub
+from langchain_community.llms import OpenAI
+# from langchain.llms import HuggingFaceHub, OpenAI
+from langchain.chains import LLMChain
+from langchain.prompts import PromptTemplate
+import warnings
+
+warnings.filterwarnings("ignore")
+
+class LLLResponseGenerator():
+
+    def __init__(self):
+        print("initialized")
+
+
+    def llm_inference(
+        self,
+        model_type: str,
+        question: str,
+        prompt_template: str,
+        context: str,
+        ai_tone: str,
+        questionnaire: str,
+        user_text: str,
+        openai_model_name: str = "",
+        hf_repo_id: str = "tiiuae/falcon-7b-instruct",
+        temperature: float = 0.1,
+        max_length: int = 128,
+    ) -> str:
+        """Call HuggingFace/OpenAI model for inference
+
+        Given a question, prompt_template, and other parameters, this function calls the relevant
+        API to fetch LLM inference results.
+
+        Args:
+            model_str: Denotes the LLM vendor's name. Can be either 'huggingface' or 'openai'
+            question: The question to be asked to the LLM.
+            prompt_template: The prompt template itself.
+            context: Instructions for the LLM.
+            ai_tone: Can be either empathy, encouragement or suggest medical help.
+            questionnaire: Can be either depression, anxiety or adhd.
+            user_text: Response given by the user.
+            hf_repo_id: The Huggingface model's repo_id
+            temperature: (Default: 1.0). Range: Float (0.0-100.0). The temperature of the sampling operation. 1 means regular sampling, 0 means always take the highest score, 100.0 is getting closer to uniform probability.
+            max_length: Integer to define the maximum length in tokens of the output summary.
+
+        Returns:
+            A Python string which contains the inference result.
+
+        HuggingFace repo_id examples:
+            - google/flan-t5-xxl
+            - tiiuae/falcon-7b-instruct
+
+        """
+        prompt = PromptTemplate(
+            template=prompt_template,
+            input_variables=[
+                "context",
+                "ai_tone",
+                "questionnaire",
+                "question",
+                "user_text",
+            ],
+        )
+
+        if model_type == "openai":
+            # https://api.python.langchain.com/en/stable/llms/langchain.llms.openai.OpenAI.html#langchain.llms.openai.OpenAI
+            llm = OpenAI(
+                model_name=openai_model_name, temperature=temperature, max_tokens=max_length
+            )
+            llm_chain = LLMChain(prompt=prompt, llm=llm)
+            return llm_chain.run(
+                context=context,
+                ai_tone=ai_tone,
+                questionnaire=questionnaire,
+                question=question,
+                user_text=user_text,
+            )
+
+        elif model_type == "huggingface":
+            # https://python.langchain.com/docs/integrations/llms/huggingface_hub
+            llm = HuggingFaceHub(
+                repo_id=hf_repo_id,
+                model_kwargs={"temperature": temperature, "max_length": max_length},
+            )
+
+            llm_chain = LLMChain(prompt=prompt, llm=llm)
+            response =  llm_chain.run(
+                context=context,
+                ai_tone=ai_tone,
+                questionnaire=questionnaire,
+                question=question,
+                user_text=user_text,
+            )
+
+            # Extracting only the response part from the output
+            response_start_index = response.find("Response;")
+            return response[response_start_index + len("Response;"):].strip()
+
+        else:
+            print(
+                "Please use the correct value of model_type parameter: It can have a value of either openai or huggingface"
+            )
+
+
+if __name__ == "__main__":
+    # Please ensure you have a .env file available with 'HUGGINGFACEHUB_API_TOKEN' and 'OPENAI_API_KEY' values.
+    HUGGINGFACEHUB_API_TOKEN = os.getenv('HUGGINGFACEHUB_API_TOKEN')
+
+    context = "You are a mental health supporting non-medical assistant. DO NOT PROVIDE any medical advice with conviction."
+
+    ai_tone = "EMPATHY"
+    questionnaire = "ADHD"
+    question = (
+        "How often do you find yourself having trouble focusing on tasks or activities?"
+    )
+    user_text = "I feel distracted all the time, and I am never able to finish"
+
+    # The user may have signs of {questionnaire}.
+    template = """INSTRUCTIONS: {context}
+
+    Respond to the user with a tone of {ai_tone}.
+
+    Question asked to the user: {question}
+
+    Response by the user: {user_text}
+
+    Provide some advice and ask a relevant question back to the user.
+
+    Response;
+    """
+
+    temperature = 0.1
+    max_length = 128
+
+    model = LLLResponseGenerator()
+
+
+    llm_response = model.llm_inference(
+        model_type="huggingface",
+        question=question,
+        prompt_template=template,
+        context=context,
+        ai_tone=ai_tone,
+        questionnaire=questionnaire,
+        user_text=user_text,
+        temperature=temperature,
+        max_length=max_length,
+        )
+
+    print(llm_response)

mental_health_model.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4b5f0ccc2395bdc10566011fb24594c84da14ed0a9e9d38365f09d63e9692515
+size 601587

q_learning_chatbot.py

@@ -0,0 +1,82 @@
+import os
+import numpy as np
+import pandas as pd
+from xgb_mental_health import MentalHealthClassifier
+import pickle
+import streamlit as st
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from torch.nn.functional import softmax
+import torch
+
+
+
+class QLearningChatbot:
+    def __init__(self, states, actions, learning_rate=0.1, discount_factor=0.9):
+        self.states = states
+        self.actions = actions
+        self.learning_rate = learning_rate
+        self.discount_factor = discount_factor
+        self.q_values = np.random.rand(len(states), len(actions))
+        self.mood = "Neutral"
+        self.mood_history = []
+        self.mood_history_int = []
+        self.tokenizer = AutoTokenizer.from_pretrained("nlptown/bert-base-multilingual-uncased-sentiment")
+        self.bert_sentiment_model_path = "bert_sentiment.pkl"
+        self.bert_sentiment_model = self.load_model() if os.path.exists(self.bert_sentiment_model_path) else AutoModelForSequenceClassification.from_pretrained("nlptown/bert-base-multilingual-uncased-sentiment")
+
+
+    def detect_sentiment(self, input_text):
+        # Encode the text
+        encoded_input = self.tokenizer(input_text, return_tensors='pt', truncation=True, max_length=512)
+
+        # Perform inference
+        with torch.no_grad():
+            output = self.bert_sentiment_model(**encoded_input)
+
+        # Process the output (softmax to get probabilities)
+        scores = softmax(output.logits, dim=1)
+
+        # Map scores to sentiment labels
+        labels = ['Negative', 'Moderately Negative', 'Neutral', 'Moderately Positive', 'Positive']
+        scores = scores.numpy().flatten()
+        scores_dict = {label: score for label, score in zip(labels, scores)}
+        highest_sentiment = max(scores_dict, key=scores_dict.get)
+        self.mood = highest_sentiment
+        return highest_sentiment
+
+    def get_action(self, current_state):
+        current_state_index = self.states.index(current_state)
+        # print(np.argmax(self.q_values[current_state_index, :]))
+        return self.actions[np.argmax(self.q_values[current_state_index, :])]
+
+    def update_q_values(self, current_state, action, reward, next_state):
+        # print(f"state-reward: {current_state} - {reward} -- (b)")
+        current_state_index = self.states.index(current_state)
+        action_index = self.actions.index(action)
+        next_state_index = self.states.index(next_state)
+
+        current_q_value = self.q_values[current_state_index, action_index]
+        max_next_q_value = np.max(self.q_values[next_state_index, :])
+
+        new_q_value = current_q_value + self.learning_rate * (reward + self.discount_factor * max_next_q_value - current_q_value)
+        self.q_values[current_state_index, action_index] = new_q_value
+
+    def update_mood_history(self):
+        st.session_state.entered_mood.append(self.mood)
+        self.mood_history = st.session_state.entered_mood
+        return self.mood_history
+
+    def check_mood_trend(self):
+        mood_dict = {'Negative': 1, 'Moderately Negative': 2, 'Neutral': 3, 'Moderately Positive': 4, 'Positive': 5}
+
+        if len(self.mood_history) >= 2:
+            self.mood_history_int = [mood_dict.get(x) for x in self.mood_history]
+            recent_moods = self.mood_history_int[-2:]
+            if recent_moods[-1] > recent_moods[-2]:
+                return 'increased'
+            elif recent_moods[-1] < recent_moods[-2]:
+                return 'decreased'
+            else:
+                return 'unchanged'
+        else:
+            return 'unchanged'

requirements.txt

@@ -0,0 +1,11 @@
+streamlit
+numpy
+pandas
+textblob
+xgboost
+transformers
+scikit-learn
+torch
+rank_bm25
+langchain_community
+langchain

xgb_mental_health.py

@@ -0,0 +1,77 @@
+import os.path
+import pickle
+import pandas as pd
+from transformers import AutoTokenizer
+from sklearn.model_selection import train_test_split
+from sklearn.feature_extraction.text import CountVectorizer
+from xgboost import XGBClassifier
+from sklearn.preprocessing import LabelEncoder
+
+class MentalHealthClassifier:
+    def __init__(self, data_path, model_path):
+        self.data = pd.read_csv(data_path, skip_blank_lines=True)
+        self.data['category'] = ['anxiety' if x == 'axienty' else x for x in self.data['category']]
+        self.data.dropna(subset=['text'], inplace=True)
+        self.data.dropna(subset=['clean_text'], inplace=True)
+        self.data_selected = self.data[['clean_text', 'category']]
+        self.df = pd.DataFrame(self.data_selected)
+        self.label_encoder = LabelEncoder()
+        self.df['category_encoded'] = self.label_encoder.fit_transform(self.df['category'])
+        self.tokenizer = None
+        self.vectorizer = CountVectorizer()
+        self.model_path = model_path
+        self.model = self.load_model() if os.path.exists(model_path) else XGBClassifier()
+
+    def preprocess_data(self):
+        tokenized_texts = [self.tokenizer.tokenize(text, padding=True, truncation=True) for text in self.df['clean_text']]
+        X = self.vectorizer.fit_transform([' '.join(tokens) for tokens in tokenized_texts]).toarray()
+        return X, self.df['category_encoded']
+
+    def train_model(self, X, y):
+        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
+        self.model.fit(X_train, y_train)
+        y_pred = self.model.predict(X_test)
+        return y_test, y_pred
+
+    def predict_category(self, input_text):
+        if self.tokenizer is None:
+            raise ValueError("Tokenizer not initialized. Call 'initialize_tokenizer' first.")
+        tokenized_input = self.tokenizer.tokenize(input_text, padding=True, truncation=True)
+        input_feature_vector = self.vectorizer.transform([' '.join(tokenized_input)]).toarray()
+        predicted_category_encoded = self.model.predict(input_feature_vector)
+        predicted_category = self.label_encoder.inverse_transform(predicted_category_encoded)
+        return predicted_category[0]
+
+    def initialize_tokenizer(self, model_name):
+        self.model_name = model_name
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+
+    def save_model(self):
+        print("saving model...to pickle...")
+        with open(self.model_path, 'wb') as f:
+            pickle.dump(self.model, f)
+
+    def load_model(self):
+        print("loading model...from pickle...")
+        with open(self.model_path, 'rb') as f:
+            return pickle.load(f)
+
+if __name__ == "__main__":
+    tokenizer_model_name = "nlptown/bert-base-multilingual-uncased-sentiment"
+    data_path = 'data.csv'
+    model_path = 'mental_health_model.pkl'
+    mental_classifier = MentalHealthClassifier(data_path, model_path)
+
+    if not os.path.exists(model_path):
+        mental_classifier.initialize_tokenizer(tokenizer_model_name)
+        X, y = mental_classifier.preprocess_data()
+        y_test, y_pred = mental_classifier.train_model(X, y)
+        mental_classifier.save_model()
+    else:
+        mental_classifier.load_model()
+        mental_classifier.initialize_tokenizer(tokenizer_model_name)  # Ensure tokenizer is initialized if loading model from pickle
+        mental_classifier.preprocess_data()
+
+    # input_text = "I feel anxiety whenever i am doing nothing."
+    # predicted_category = mental_classifier.predict_category(input_text)
+    # print("Predicted mental health condition:", predicted_category)