app.py

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 as QuestionRetriever_bm25
from Chromadb_storage import QuestionRetriever as QuestionRetriever_chromaDB
from llm_response_generator import LLLResponseGenerator
import os
# 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", "spiritual"]

# Initialize Q-learning chatbot and mental health classifier
chatbot = QLearningChatbot(states, actions)

# Initialize MentalHealthClassifier
# data_path = "/Users/jaelinlee/Documents/projects/fomo/input/data.csv"
data_path = "data/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 = []

if "messages" not in st.session_state:
    st.session_state.messages = []

# Select Question Retriever
selected_retriever_option = st.sidebar.selectbox(
    "Choose Question Retriever", ("BM25", "ChromaDB")
)
if selected_retriever_option == "BM25":
    retriever = QuestionRetriever_bm25()
if selected_retriever_option == "ChromaDB":
    retriever = QuestionRetriever_chromaDB()

for message in st.session_state.messages:
    with st.chat_message(message.get("role")):
        st.write(message.get("content"))

# Collect user input
user_message = st.chat_input("Type your message here:")

# Take user input
if user_message:
    st.session_state.entered_text.append(user_message)

    st.session_state.messages.append({"role": "user", "content": user_message})
    with st.chat_message("user"):
        st.write(user_message)
    # Detect mental condition

    with st.spinner("Processing..."):
        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)

        # Detect sentiment
        user_sentiment = chatbot.detect_sentiment(user_message)

        # Retrieve question
        if user_sentiment in ["Negative", "Moderately Negative"]:
            question = retriever.get_response(user_message, predicted_mental_category)
            show_question = True
        else:
            show_question = False
            question = ""
            predicted_mental_category = ""

        # 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 = +1
                mood_trend_symbol = " ⬆️"
            elif mood_trend == "unchanged":
                reward = +0.8
                mood_trend_symbol = ""
            else:  # decresed
                reward = -0.2
                mood_trend_symbol = " ⬇️"

        else:
            if mood_trend == "increased":
                reward = +1
                mood_trend_symbol = " ⬆️"
            elif mood_trend == "unchanged":
                reward = -0.2
                mood_trend_symbol = ""
            else:  # decreased
                reward = -1
                mood_trend_symbol = " ⬇️"

        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.session_state.messages.append({"role": "ai", "content": llm_response})

    with st.chat_message("ai"):
        st.markdown(llm_response)
    # st.write(f"{llm_response}")
    if show_question:
        st.write(f"{question}")
    # else:
    # user doesn't feel negative.
    # get question to ecourage even more positive behaviour

    st.subheader("Behind the Scence - What AI is doing:")
    st.write(
        f"- Detected User Tone: {user_sentiment} ({mood_trend.capitalize()}{mood_trend_symbol})"
    )
    if show_question:
        st.write(
            f"- Possible Mental Condition: {predicted_mental_category.capitalize()}"
        )
    st.write(f"- AI Tone: {ai_tone.capitalize()}")
    st.write(f"- Question retrieved from: {selected_retriever_option}")
    st.write(
        f"- If the user feels neagative or moderately negative, at the end of the AI response, it adds a mental health condition realted question. The question is retrieved from DB. The categories of questions are limited to Depression, Anxiety, and ADHD which are most associated with FOMO related to excessive social media usage."
    )
    st.write(
        f"- Below q-table is continously updated after each interaction with the user. If the user's mood increases, AI gets reward. Else, AI gets punishment."
    )

    # 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}")