app.py

import gradio as gr
#from gradio.components import Box
import torch
import os
import faiss
import pickle
from transformers import AutoModelForCausalLM, AutoTokenizer
from sentence_transformers import SentenceTransformer, util
from huggingface_hub import hf_hub_download
import pandas as pd
from datasets import Dataset
from huggingface_hub import create_repo

os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"

# 🔹 Ensure bitsandbytes installation (for efficient model loading)
try:
    if torch.cuda.is_available():
        import bitsandbytes as bnb  # Use the GPU version if CUDA is available
        print("Using bitsandbytes with GPU support.")
    else:
        import bitsandbytes  # Fallback to CPU version (as bitsandbytes supports both)
        print("Using bitsandbytes with CPU support.")
except ImportError:
    print("bitsandbytes not found, falling back to CPU.")

# 🔹 Load Hugging Face Model Repo Details
MODEL_NAME = "AvocadoMuffin/Gemma_Fine_Tuned_Model"
HF_TOKEN = os.getenv("HF_TOKEN")  # Ensure this is set in your environment

# 🔹 Load Tokenizer & Model
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if torch.cuda.is_available() else torch.float32

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=HF_TOKEN)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=dtype, token=HF_TOKEN).to(device)

tokenizer.pad_token = tokenizer.eos_token

# 🔹 Load FAISS Index & Knowledge from HF Model Repo
faiss_index_path = hf_hub_download(repo_id=MODEL_NAME, filename="faiss_index/index.faiss", token=HF_TOKEN)
knowledge_path = hf_hub_download(repo_id=MODEL_NAME, filename="faiss_index/knowledge.pkl", token=HF_TOKEN)

# Load FAISS Index
index = faiss.read_index(faiss_index_path)

# Load Knowledge.pkl (Contains text data + embeddings)
with open(knowledge_path, "rb") as f:
    knowledge_data = pickle.load(f)

# Validate knowledge data
if isinstance(knowledge_data, dict) and "questions" in knowledge_data and "answers" in knowledge_data:
    questions = knowledge_data["questions"]
    answers = knowledge_data["answers"]
else:
    raise ValueError("❌ knowledge.pkl does not contain expected questions and answers keys!")

# 🔹 Load Sentence Transformer Embedding Model
embedding_model = SentenceTransformer("multi-qa-mpnet-base-dot-v1")

print("✅ FAISS Index, Knowledge Data, and Embedding Model Loaded Successfully!")

# 🔹 Function: Retrieve Relevant Context from FAISS
def retrieve_relevant_context(query, index, answers, embedding_model, top_k=3, similarity_threshold=0.8):
    """Retrieve top-k relevant knowledge snippets from FAISS index."""
    
    query_embedding = embedding_model.encode([query], convert_to_numpy=True)
    distances, indices = index.search(query_embedding, top_k)

    retrieved_contexts = []
    seen_embeddings = []

    for idx in indices[0]:
        if 0 <= idx < len(answers):  
            answer = answers[idx]
            answer_embedding = embedding_model.encode([answer], convert_to_numpy=True)

            if all(util.cos_sim(answer_embedding, emb) < similarity_threshold for emb in seen_embeddings):
                retrieved_contexts.append(answer)
                seen_embeddings.append(answer_embedding)

    return "\n".join(retrieved_contexts[:top_k])

# Define system messages for different modes
SYSTEM_MESSAGES = {
    "general": "You are a helpful AI assistant for everyday tasks.",
    "mental_health": "You are a helpful AI assistant for mental health."
}

# Mental health keywords to check if query is related to mental health
MENTAL_HEALTH_KEYWORDS = [
    "anxiety", "depression", "stress", "therapy", "counseling", "mental health", "panic", "trauma", 
    "mindfulness", "meditation", "psychological", "self-care", "emotion", "feeling", "mood", 
    "psychiatry", "therapist", "psychologist", "burnout", "coping", "ptsd", "ocd", "adhd", "bipolar", 
    "schizophrenia", "grief", "insomnia", "wellbeing", "mental wellbeing", "mental wellness", 
    "emotional wellbeing", "emotional wellness", "self-help", "relaxation", "breathing exercise",
    "cognitive behavioral", "cbt", "sad", "angry", "worried", "anxious", "depressed", "overwhelmed",
    "psychology", "mental disorder", "suicide", "suicidal", "crisis", "loneliness", "lonely", "isolated",
    "phobia", "trauma", "addiction", "substance abuse", "eating disorder", "self-esteem", "confidence"
]

# Function to check if a query is related to mental health
def is_mental_health_query(query):
    query_lower = query.lower()
    return any(keyword in query_lower for keyword in MENTAL_HEALTH_KEYWORDS)

# 🔹 Function: Generate Response Using RAG 
def generate_rag_response(query, model, tokenizer, index, answers, embedding_model, chat_history, system_message, max_new_tokens=200):
    """Generate response using retrieved context, chat history, and system message."""

    # Check if query is mental health related when in mental health mode
    if system_message == SYSTEM_MESSAGES["mental_health"] and not is_mental_health_query(query):
        return "I'm currently in mental health assistant mode and can only answer questions related to mental health, psychology, emotions, and wellbeing. If you'd like to discuss other topics, please change the system message to general assistant mode."
    
    query_lower = query.strip().lower()

    # Define follow-up triggers  
    follow_up_triggers = ["continue", "go on", "tell me more", "what else?", "elaborate", "explain more"]
    is_follow_up = any(trigger in query_lower for trigger in follow_up_triggers)

    # Retrieve relevant knowledge  
    retrieved_context = retrieve_relevant_context(query, index, answers, embedding_model).strip()

    # Store last 5 exchanges in chat history (limit it)
    chat_context = "\n".join([msg["content"] for msg in chat_history[-5:]])

    # Handle Follow-up Responses More Effectively  
    if is_follow_up and chat_history:
        last_ai_response = chat_history[-1]  # Get last AI message
        query = f"{last_ai_response['content']} {query}"  # Append the follow-up to continue smoothly

    # Few-Shot Examples for Mental Health mode only
    few_shot_examples = """
User: What are some effective meditation techniques for beginners?
AI: Based on established research, here are some effective meditation techniques for beginners:

1. Mindful Breathing - Focus on your breath, noticing the sensation of air moving in and out of your body.
2. Body Scan Meditation - Progressively relax different parts of your body from head to toe.
3. Loving-Kindness Meditation - Direct positive wishes toward yourself and others.
4. Guided Meditation - Follow along with recorded instructions to help maintain focus.
5. Walking Meditation - Practice mindfulness while walking slowly and deliberately.

I recommend starting with just 5 minutes daily and gradually increasing duration. Would you like more details about any specific technique, or would you prefer suggestions for apps that could help you get started?

User: Tell me more about body scan meditation.
AI: Body scan meditation is a mindfulness practice where you systematically focus your attention on different parts of your body, from your toes to the top of your head.

Here's how to practice it:
1. Lie down in a comfortable position or sit in a chair with your feet on the ground
2. Close your eyes and take several deep breaths
3. Begin by bringing awareness to your toes, noticing any sensations (tingling, temperature, tension)
4. Gradually move your attention upward (feet, ankles, calves, etc.)
5. For each body part, notice sensations without judgment, then consciously relax that area
6. If you notice pain or discomfort, acknowledge it with compassion before moving on
7. Continue until you've scanned your entire body

Regular practice can help reduce stress, improve body awareness, and release physical tension. Many find it helpful for sleep issues as well.

Would you like guidance on how long to practice or how to handle common challenges like mind-wandering during body scans?
"""

    # Build prompt depending on system message
    if system_message == SYSTEM_MESSAGES["mental_health"]:
        prompt = f"""You are a compassionate mental wellness assistant powered by a retrieval-augmented generation (RAG) system.
Maintain a warm, supportive conversation while providing evidence-based mental health information.

SYSTEM MESSAGE: {system_message}

When responding to mental health queries:
- Provide evidence-based information from reputable sources
- Keep the tone supportive, non-judgmental, and empathetic
- If uncertain about mental health information, acknowledge limitations and avoid potentially harmful advice
- For crisis situations, recommend professional help rather than providing advice
- Frame mental wellness information in a way that emphasizes self-compassion and gradual progress

If the user expresses overwhelm, anxiety, or difficulty with their mental health, offer evidence-based coping strategies and gentle support.

{few_shot_examples if len(chat_history) > 0 else ""}
"""
    else:
        prompt = f"""You are a helpful AI assistant designed to answer questions on a wide range of topics.
Provide accurate, concise information while maintaining a conversational tone.

SYSTEM MESSAGE: {system_message}

When responding to queries:
- Provide accurate information on the requested topic
- Keep explanations clear and easy to understand
- Be conversational and friendly
- If you don't know something, acknowledge limitations rather than making up information
"""

    # Only append the relevant context if available
    if retrieved_context:
        prompt += f"\n📌 Relevant Information Retrieved:\n{retrieved_context}\n\n"

    prompt += f"💬 Chat History:\n{chat_context}\n👤 User: {query}\n🤖 AI:"

    # Tokenize input  
    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)

    # Generate AI response  
    with torch.no_grad():
        output_ids = model.generate(
            input_ids, 
            max_new_tokens=max_new_tokens, 
            eos_token_id=tokenizer.eos_token_id,
        )

    response = tokenizer.decode(output_ids[0], skip_special_tokens=True).strip()

    # Prevent prompt repetition in response  
    if response.startswith(prompt):
        response = response[len(prompt):].strip()

    return response


# 🔹 Function to save feedback for RLHF
from huggingface_hub import create_repo, HfApi
from datasets import Dataset, load_dataset
import pandas as pd
import os
import json

def save_feedback(query, response, feedback, comment=""):
    feedback_file = "chat_feedback.csv"
    repo_name = "AvocadoMuffin/Chat_Feedback_Data"
    
    # Create new data from current feedback
    new_data = pd.DataFrame([[query, response, feedback, comment]], 
                          columns=["Query", "Response_Generated", "Feedback", "Comment"])
    
    try:
        # First check if we have a local copy of the data
        if os.path.exists(feedback_file):
            existing_data = pd.read_csv(feedback_file)
            print(f"Loaded existing data from local file with {len(existing_data)} records")
        else:
            # Try to download the dataset from HF Hub
            api = HfApi()
            try:
                # List all files in the repository to find our data
                files = api.list_repo_files(repo_id=repo_name, repo_type="dataset")
                
                # Look for parquet files which contain our data
                data_files = [f for f in files if f.endswith('.parquet') or f.endswith('.csv')]
                
                if data_files:
                    # Download the first data file we find
                    file_path = api.hf_hub_download(
                        repo_id=repo_name,
                        filename=data_files[0],
                        repo_type="dataset"
                    )
                    
                    # Load the data based on file extension
                    if file_path.endswith('.parquet'):
                        existing_data = pd.read_parquet(file_path)
                    else:
                        existing_data = pd.read_csv(file_path)
                        
                    print(f"Successfully loaded data from Hub with {len(existing_data)} records")
                else:
                    print("No data files found in repository")
                    existing_data = pd.DataFrame(columns=["Query", "Response_Generated", "Feedback", "Comment"])
            except Exception as e:
                print(f"Error accessing repository: {e}")
                existing_data = pd.DataFrame(columns=["Query", "Response_Generated", "Feedback", "Comment"])
        
        # Append new feedback to the existing data
        updated_data = pd.concat([existing_data, new_data], ignore_index=True)
        print(f"Updated dataset now has {len(updated_data)} records")
        
        # Always save locally for backup and future reference
        updated_data.to_csv(feedback_file, index=False)
        print(f"Saved updated data to local file with {len(updated_data)} records")
        
        # Convert to Hugging Face Dataset
        dataset = Dataset.from_pandas(updated_data)
        
        # Create repo if needed and push to Hugging Face Hub
        create_repo(repo_id=repo_name, repo_type="dataset", exist_ok=True)
        dataset.push_to_hub(repo_name)
        print("Feedback saved and dataset updated on Hugging Face!")
        
    except Exception as e:
        print(f"Error updating dataset: {e}")
        # Still save locally as backup
        if os.path.exists(feedback_file):
            local_data = pd.read_csv(feedback_file)
            pd.concat([local_data, new_data], ignore_index=True).to_csv(feedback_file, index=False)
        else:
            new_data.to_csv(feedback_file, index=False)
        print("Saved locally only due to error")
    
    return "Feedback saved and dataset updated!"


# 🔹 The chatbot function (with RAG and conversation history tracking)
def chatbot(query, chat_history=None, system_message=SYSTEM_MESSAGES["general"]):
    if chat_history is None:
        chat_history = []  # Initialize chat_history as an empty list
    
    # Append user message to the chat history
    chat_history.append({"role": "user", "content": query})
    
    # Generate response using the RAG model with specified system message
    response = generate_rag_response(query, model, tokenizer, index, answers, embedding_model, chat_history, system_message)
    
    # Append the assistant's response to the chat history for context
    chat_history.append({"role": "assistant", "content": response})
    
    # Format chat history with HTML and CSS for better visual appeal
    formatted_chat_history = ""
    for msg in chat_history:
        if msg["role"] == "user":
            formatted_chat_history += f'<p style="color: #4A90E2; font-weight: bold;">User:</p><p style="background-color: #E7F1F8; padding: 10px; border-radius: 5px; font-family: Arial, sans-serif; font-size: 14px;">{msg["content"]}</p>'
        else:
           formatted_chat_history += f'<p style="color: #1D4ED8; font-weight: bold;">AI:</p><p style="background-color: #D1F4D9; padding: 10px; border-radius: 5px; font-family: Arial, sans-serif; font-size: 14px;">{msg["content"]}</p>'
    
    # Return the formatted chat history as HTML and the updated chat history
    # Make sure to reset the feedback radio button value to None when showing new feedback options
    return gr.HTML(formatted_chat_history), chat_history, gr.update(visible=True, value=None), gr.update(visible=True), gr.update(visible=True, value=""), response, gr.update(value=""), query

# Function to handle feedback submission
def handle_feedback(query, latest_response, feedback, comment):
    # Store feedback as None if not selected
    if feedback:  # Only save if feedback has a value
        feedback_value = 1 if feedback == "👍 Like" else 0
        save_feedback(query, latest_response, feedback_value, comment)  # Save feedback data
    else:
        # If no feedback is selected, save as NULL or handle accordingly
        save_feedback(query, latest_response, None, comment)
        
    # Hide feedback UI after submission and clear both the comment field and feedback selection
    return gr.update(visible=False, value=None), gr.update(visible=False), gr.update(visible=False, value="")

# Function to handle system message change and reset chat
def change_system_message(message_key, chat_history_state):
    # Get the actual system message
    selected_message = SYSTEM_MESSAGES[message_key]
    
    # Reset chat history when mode changes
    new_chat_history = []
    
    # Create welcome message based on selected mode
    if message_key == "mental_health":
        welcome_message = "I'm now in mental health assistant mode. I can help with questions about mental health, wellness, emotions, and psychological well-being. How can I support you today?"
    else:
        welcome_message = "I'm here to help with a wide range of questions and tasks. What would you like assistance with today?"
    
    # Add welcome message to chat history
    new_chat_history.append({"role": "assistant", "content": welcome_message})
    
    # Format the welcome message for display
    formatted_chat_history = f'<p style="color: #1D4ED8; font-weight: bold;">AI:</p><p style="background-color: #D1F4D9; padding: 10px; border-radius: 5px; font-family: Arial, sans-serif; font-size: 14px;">{welcome_message}</p>'
    
    return gr.HTML(formatted_chat_history), new_chat_history, selected_message

def create_gradio_interface():
    with gr.Blocks() as demo:
        gr.Markdown("## **🧠 Mindful Bot: AI for Everything**")
        gr.Markdown("An AI chatbot designed to offer empathetic support for mental wellness using RAG-based techniques, capable of handling adversarial queries with care. It also provides insightful responses to general queries, separating wellness-related conversations for more focused, compassionate assistance.")

        latest_query = gr.State("")
        latest_response = gr.State("")
        chat_history_state = gr.State([])
        current_system_message = gr.State(SYSTEM_MESSAGES["general"])
        
        with gr.Row():
            with gr.Column(scale=7):  # Chat area
                # Use gr.HTML to display the formatted chat history
                chat_history_box = gr.HTML(label="Chat History", elem_id="chat-history")  
                    
                with gr.Row():
                    with gr.Column(scale=8):
                        query = gr.Textbox(label="Enter your query", elem_id="user-input", lines=3)
                    with gr.Column(scale=2):
                        submit_button = gr.Button("Submit", elem_id="submit-button", variant="primary")
                    
                # Feedback UI - make it optional with no default selection
                feedback = gr.Radio(choices=["👍 Like", "👎 Dislike"], label="Feedback (Optional)", type="value", visible=False)
                submit_feedback_button = gr.Button("Submit Feedback", visible=False)
                feedback_comments = gr.Textbox(label="Additional Comments (optional)", lines=2, visible=False)
            
            with gr.Column(scale=3):  # System message area
                with gr.Group():
                    gr.HTML("<h2 style='text-align: center; color: white; background-color: orange; padding: 10px;'>System message</h2>")
                    
                    # Example table with system messages
                    example_table = gr.Dataframe(
                        headers=["Message", "System message"],
                        datatype=["str", "str"],
                        value=[
                            ["Hello! How are you?", SYSTEM_MESSAGES["mental_health"]],
                            ["Can you help with a recipe for baking a cake?", SYSTEM_MESSAGES["general"]],
                        ],
                        row_count=2,
                        col_count=2,
                        interactive=True,
                        #height=150
                    )
                    
                    # System message selector
                    system_selector = gr.Radio(
                        choices=["general", "mental_health"],
                        value="general",
                        label="Select system message type",
                        elem_id="system-selector",
                        info="Choose the type of assistant"
                    )
                    
                    # Display the current system message
                    system_message_display = gr.Textbox(
                        value=SYSTEM_MESSAGES["general"],
                        label="Current system instruction:",
                        interactive=False
                    )
                    
                    # Reset button
                    reset_button = gr.Button("Reset Chat", elem_id="reset-button")
                    
                    # Additional inputs section
                    #gr.Markdown("### Additional Inputs")

        # Add custom CSS with HTML component
        gr.HTML("""<style>
            #chat-history { background-color: #f5f5f5; border-radius: 10px; padding: 10px; max-height: 400px; overflow-y: scroll; font-family: 'Arial', sans-serif; }
            #user-input { margin-top: 10px; border-radius: 10px; }
            #submit-button { margin-top: 10px; }
            #system-selector { margin-top: 10px; }
        </style>""")
        
        # Event handlers
        query.submit(chatbot, 
                    inputs=[query, chat_history_state, current_system_message], 
                    outputs=[chat_history_box, chat_history_state, feedback, submit_feedback_button, feedback_comments, latest_response, query, latest_query])
        
        submit_button.click(chatbot, 
                          inputs=[query, chat_history_state, current_system_message], 
                          outputs=[chat_history_box, chat_history_state, feedback, submit_feedback_button, feedback_comments, latest_response, query, latest_query])
        
        # Modified feedback handling
        submit_feedback_button.click(handle_feedback, 
                                   inputs=[latest_query, latest_response, feedback, feedback_comments], 
                                   outputs=[feedback, submit_feedback_button, feedback_comments])
        
        # System message change handler
        system_selector.change(change_system_message,
                             inputs=[system_selector, chat_history_state],
                             outputs=[chat_history_box, chat_history_state, current_system_message])
        
        # Also update system message display when selector changes
        system_selector.change(lambda x: SYSTEM_MESSAGES[x],
                             inputs=[system_selector],
                             outputs=[system_message_display])
        
        # Reset button handler - uses the same function as system_selector.change
        reset_button.click(change_system_message,
                         inputs=[system_selector, chat_history_state],
                         outputs=[chat_history_box, chat_history_state, current_system_message])
        
        # Also reset system message display with reset button
        reset_button.click(lambda x: SYSTEM_MESSAGES[x],
                         inputs=[system_selector],
                         outputs=[system_message_display])
        
        # Initialize with welcome message on load
        demo.load(change_system_message,
                inputs=[system_selector, chat_history_state],
                outputs=[chat_history_box, chat_history_state, current_system_message])
        
        # Initialize system message display on load
        demo.load(lambda x: SYSTEM_MESSAGES[x],
                inputs=[system_selector],
                outputs=[system_message_display])
        
    return demo

# Create and launch the Gradio interface
demo = create_gradio_interface()
demo.launch(share=False, debug=True)