# -*- coding: utf-8 -*-
"""chatbot.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1zgihAeNpcDd0opNPsbnmS2UNgz0rsfOB
"""

from langchain.text_splitter import CharacterTextSplitter
from langchain_community.document_loaders import PyPDFLoader

import logging
import sys
from pathlib import Path
import os

logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))

from llama_index import VectorStoreIndex, SimpleDirectoryReader, ServiceContext
from llama_index.llms import HuggingFaceLLM
from langchain_community.document_loaders import PyPDFLoader

documents = SimpleDirectoryReader("Dat").load_data()
documents = [CharacterTextSplitter(separator="\n", chunk_size=1000, chunk_overlap=0).split_text(doc) for doc in documents]

from llama_index.prompts.prompts import SimpleInputPrompt


system_prompt =  """Emphasize empathy and active listening: Create a safe space for users to share their thoughts and feelings without judgment.
Focus on understanding and validation: Reflect back user emotions and experiences to demonstrate understanding and build trust.
Offer evidence-based support: Provide grounding techniques, coping strategies, and psychoeducation based on sound mental health principles.
Personalize responses: Tailor interactions to individual needs, preferences, and goals.
Maintain ethical boundaries: Respect user privacy, confidentiality, and autonomy.
Recognize limitations: Acknowledge that the chatbot is not a replacement for professional therapy and encourage seeking licensed help when needed.
Key goals:

Reduce symptoms of anxiety, depression, and stress.
Improve emotional regulation and coping skills.
Enhance self-awareness and self-compassion.
Promote healthy relationships and communication.
Build resilience and problem-solving skills.
Encourage positive self-care and lifestyle choices.
Specific prompts:

"Greet the user warmly and introduce yourself as their AI therapist."
"Ask open-ended questions to elicit user thoughts, feelings, and concerns."
"Respond empathetically to user disclosures, validating their experiences."
"Offer appropriate mental health resources, exercises, or techniques based on user needs."
"Guide users through mindfulness exercises or relaxation techniques."
"Challenge negative thinking patterns and encourage cognitive reframing."
"Help users set realistic goals and track progress towards mental wellness."
"Provide psychoeducation on various mental health topics and treatment options."
"Conclude sessions with positive affirmations and encouragement."
"Remind users of the chatbot's limitations and the importance of seeking professional help."
"Always prioritize user safety and offer crisis resources in case of urgent needs."
Additional considerations:

Tailor prompts to specific mental health conditions or challenges.
Incorporate humor or lightheartedness when appropriate to build rapport.
Provide options for different communication styles (e.g., text, voice, interactive activities).
Continuously monitor and refine prompts based on user feedback and clinical expertise."""



# This will wrap the default prompts that are internal to llama-index
query_wrapper_prompt = SimpleInputPrompt("<|USER|>{query_str}<|ASSISTANT|>")

import torch

llm = HuggingFaceLLM(
    context_window=2048,
    max_new_tokens=128,
    generate_kwargs={"temperature": 0.5, "do_sample": False},
    system_prompt=system_prompt,
    query_wrapper_prompt=query_wrapper_prompt,
    tokenizer_name="NousResearch/Llama-2-7b-chat-hf",
    model_name="NousResearch/Llama-2-7b-chat-hf",
    device_map="auto",
    # uncomment this if using CUDA to reduce memory usage
    model_kwargs={"torch_dtype": torch.float16 })

from langchain.embeddings.huggingface import HuggingFaceEmbeddings
from llama_index.embeddings.langchain import LangchainEmbedding
from llama_index import ServiceContext

embed_model = LangchainEmbedding(
  HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
)

service_context = ServiceContext.from_defaults(
    chunk_size=512,
    llm=llm,
    embed_model=embed_model
)

index = VectorStoreIndex.from_documents(documents, service_context=service_context)

query_engine = index.as_query_engine()

import gradio as gr

def respond_to_user(input_text):
    response = query_engine.query(input_text)
    return response


import gradio as gr

def render_messages():
    messages = []
    if hasattr(iface, "_message_history"):
        for msg in iface._message_history[-3::-1]:
            messages.append({"role": msg["role"], "content": msg["content"]})
        
    template = r"""
    <div class='card'>
        %for message in messages:
            {%if message['role'] == 'assistant':%}
                <div class='card bg-light mb-3' style='max-width: 50rem;'>
                    <div class='card-header'>Assistant</div>
                    <div class='card-body'>{{message['content']}}</div>
                </div>
            {%else:%}
                <div class='card' style='max-width: 50rem;'>
                    <div class='card-header'>User</div>
                    <div class='card-body'>{{message['content']}}</div>
                </div>
            {%endif%}
        %endfor%}
    </div>
    """
    return gr.outputs.HTML(template)

def process_inputs(input_text):
    result = respond_to_user(input_text)
    iface._message_history.append({"role": "assistant", "content": result})
    return "", render_messages()

iface = gr.Interface(fn=process_inputs,
                     inputs=gr.components.Textbox(lines=3, placeholder="Type something..."),
                     outputs=[],
                     title="AI Therapist Chatbot",
                     allow_flagging="never",
                     theme="monochrome",
                     css="""
                       .gradio-container > div:first-child{
                           margin-top: 0 !important;
                       }
                       .gradio-label-container label{
                           font-weight: bold;
                           color: black;
                       }
                   """,
                     article="""
                       Welcome to our AI Therapist chatbot! Feel free to ask anything related to mental health and receive guidance. Confidentiality and privacy notice: Your conversation remains private; we do not store any data.
                     """)

iface.render()