|
import gradio as gr |
|
from huggingface_hub import InferenceClient |
|
from typing import List, Tuple |
|
import fitz |
|
from sentence_transformers import SentenceTransformer, util |
|
import numpy as np |
|
import faiss |
|
|
|
client = InferenceClient("HuggingFaceH4/zephyr-7b-beta") |
|
|
|
|
|
class MyApp: |
|
def __init__(self) -> None: |
|
self.documents = [] |
|
self.embeddings = None |
|
self.index = None |
|
self.load_pdf("THEDIA1.pdf") |
|
self.build_vector_db() |
|
|
|
def load_pdf(self, file_path: str) -> None: |
|
"""Extracts text from a PDF file and stores it in the app's documents.""" |
|
doc = fitz.open(file_path) |
|
self.documents = [] |
|
for page_num in range(len(doc)): |
|
page = doc[page_num] |
|
text = page.get_text() |
|
self.documents.append({"page": page_num + 1, "content": text}) |
|
print("PDF processed successfully!") |
|
|
|
def build_vector_db(self) -> None: |
|
"""Builds a vector database using the content of the PDF.""" |
|
model = SentenceTransformer('all-MiniLM-L6-v2') |
|
self.embeddings = model.encode([doc["content"] for doc in self.documents]) |
|
self.index = faiss.IndexFlatL2(self.embeddings.shape[1]) |
|
self.index.add(np.array(self.embeddings)) |
|
print("Vector database built successfully!") |
|
|
|
def search_documents(self, query: str, k: int = 3) -> List[str]: |
|
"""Searches for relevant documents using vector similarity.""" |
|
model = SentenceTransformer('all-MiniLM-L6-v2') |
|
query_embedding = model.encode([query]) |
|
D, I = self.index.search(np.array(query_embedding), k) |
|
results = [self.documents[i]["content"] for i in I[0]] |
|
return results if results else ["No relevant documents found."] |
|
|
|
app = MyApp() |
|
|
|
def preprocess_input(user_input: str) -> str: |
|
"""Preprocesses user input to enhance it for better context.""" |
|
if "therapy" in user_input.lower() or "exercise" in user_input.lower(): |
|
return "I am looking for guidance on therapy. Can you help me with some exercises or techniques to manage my stress and emotions?" |
|
return user_input |
|
|
|
def preprocess_response(response: str) -> str: |
|
"""Preprocesses the response to make it more polished.""" |
|
response = response.strip() |
|
response = response.replace("\n\n", "\n") |
|
response = response.replace(" ,", ",") |
|
response = response.replace(" .", ".") |
|
response = " ".join(response.split()) |
|
return response |
|
|
|
def shorten_response(response: str) -> str: |
|
"""Uses the Zephyr model to shorten and refine the response.""" |
|
messages = [{"role": "system", "content": "Shorten and refine this response in bullet list."}, {"role": "user", "content": response}] |
|
result = client.chat_completion(messages, max_tokens=512, temperature=0.5, top_p=0.9) |
|
return result.choices[0].message['content'].strip() |
|
|
|
def respond(message: str, history: List[Tuple[str, str]]): |
|
system_message = "You are a concisely speaking empathetic Dialectical Behaviour Therapist assistant. You politely guide users through DBT exercises based on the given DBT book. You must say one thing at a time and ask follow-up questions to continue the chat." |
|
messages = [{"role": "system", "content": system_message}] |
|
|
|
for val in history: |
|
if val[0]: |
|
messages.append({"role": "user", "content": val[0]}) |
|
if val[1]: |
|
messages.append({"role": "assistant", "content": val[1]}) |
|
|
|
|
|
preprocessed_message = preprocess_input(message) |
|
messages.append({"role": "user", "content": preprocessed_message}) |
|
|
|
|
|
retrieved_docs = app.search_documents(preprocessed_message) |
|
context = "\n".join(retrieved_docs) |
|
if context.strip(): |
|
messages.append({"role": "system", "content": "Relevant documents: " + context}) |
|
|
|
response = client.chat_completion(messages, max_tokens=1024, temperature=0.7, top_p=0.9) |
|
response_content = "".join([choice.message['content'] for choice in response.choices if 'content' in choice.message]) |
|
|
|
polished_response = preprocess_response(response_content) |
|
shortened_response = shorten_response(polished_response) |
|
|
|
history.append((message, shortened_response)) |
|
return history, "" |
|
|
|
with gr.Blocks() as demo: |
|
gr.Markdown("# 🧘♀️ **Dialectical Behaviour Therapy**") |
|
gr.Markdown( |
|
"‼️Disclaimer: This chatbot is based on a DBT exercise book that is publicly available. " |
|
"We are not medical practitioners, and the use of this chatbot is at your own responsibility." |
|
) |
|
|
|
chatbot = gr.Chatbot() |
|
|
|
with gr.Row(): |
|
txt_input = gr.Textbox( |
|
show_label=False, |
|
placeholder="Type your message here...", |
|
lines=1 |
|
) |
|
submit_btn = gr.Button("Submit", scale=1) |
|
refresh_btn = gr.Button("Refresh Chat", scale=1, variant="secondary") |
|
|
|
example_questions = [ |
|
["What are some techniques to handle distressing situations?"], |
|
["How does DBT help with emotional regulation?"], |
|
["Can you give me an example of an interpersonal effectiveness skill?"], |
|
["I want to practice mindfulness. Can you help me?"], |
|
["I want to practice distraction techniques. What can I do?"], |
|
["How do I plan self-accommodation?"], |
|
["What are some distress tolerance skills?"], |
|
["Can you help me with emotional regulation techniques?"], |
|
["How can I improve my interpersonal effectiveness?"], |
|
["What are some ways to cope with stress using DBT?"], |
|
["Can you guide me through a grounding exercise?"], |
|
["How do I use DBT skills to handle intense emotions?"], |
|
["What are some self-soothing techniques I can practice?"], |
|
["How can I create a sensory-friendly safe space?"], |
|
["Can you help me create a personal crisis plan?"], |
|
["What are some affirmations for neurodivergent individuals?"], |
|
["How can I manage rejection sensitive dysphoria?"], |
|
["Can you guide me through observing with my senses?"], |
|
["What are some accessible mindfulness exercises?"], |
|
["How do I engage my wise mind?"], |
|
["What are some values that I can identify with?"], |
|
["How can I practice mindful appreciation?"], |
|
["What is the STOP skill in distress tolerance?"], |
|
["How can I use the TIPP skill to manage distress?"], |
|
["What are some tips for managing meltdowns?"], |
|
["Can you provide a list of stims that I can use?"], |
|
["How do I improve my environment to reduce distress?"] |
|
] |
|
|
|
gr.Examples(examples=example_questions, inputs=[txt_input]) |
|
|
|
submit_btn.click(respond, [txt_input, chatbot], [chatbot, txt_input]) |
|
refresh_btn.click(lambda: [], None, chatbot) |
|
|
|
if __name__ == "__main__": |
|
demo.launch() |
|
|