File size: 4,215 Bytes
d17fb25 2af4b6b d17fb25 f21c67e 8f219da f21c67e 8f219da f21c67e 8f219da f21c67e f6fd0a6 f21c67e 8f219da f21c67e d17fb25 8f219da f21c67e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 |
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
class MentalHealthChatbot:
def __init__(self, model_path="hemhemoh/Gemma-2-2b-it-wazobia-wellness-bot"):
self.model = AutoModelForCausalLM.from_pretrained(model_path)
self.tokenizer = AutoTokenizer.from_pretrained(model_path)
# Comprehensive system instruction
self.predefined_instruction = (
"CORE THERAPEUTIC GUIDELINES:\n"
"1. Professional Role: You are an advanced AI mental health support assistant, "
"trained to provide compassionate, culturally sensitive, and professional psychological support.\n"
"2. Communication Principles:\n"
" - Listen actively and empathetically\n"
" - Validate the user's emotions without judgment\n"
" - Provide support that is culturally nuanced (English, Yoruba, Igbo, Hausa)\n"
" - Maintain professional boundaries\n"
"3. Response Strategy:\n"
" - Use a warm, supportive, and non-threatening tone\n"
" - Ask open-ended, reflective questions\n"
" - Offer practical coping strategies\n"
" - Avoid direct medical diagnosis\n"
"4. Safety Protocol:\n"
" - If user expresses thoughts of self-harm or suicide, respond with immediate compassion and provide crisis resource information\n"
" - Never minimize the user's feelings\n"
" - Encourage professional help when issues seem complex\n"
"5. Linguistic Flexibility:\n"
" - Respond in the language of the user\n"
" - Use culturally appropriate language and metaphors\n"
)
def prepare_prompt(self, message, history):
# Construct conversation history
conversation_context = ""
for user_input, assistant_response in history:
if user_input:
conversation_context += f"User: {user_input}\n"
if assistant_response:
conversation_context += f"Assistant: {assistant_response}\n"
# Combine system instruction, conversation history, and new message
full_prompt = (
f"{self.predefined_instruction}\n"
f"CONVERSATION HISTORY:\n{conversation_context}\n"
f"CURRENT USER MESSAGE:\n{message}\n"
f"ASSISTANT'S COMPASSIONATE RESPONSE:"
)
return full_prompt
def respond(
self,
message,
history,
max_tokens=512,
temperature=0.2,
top_p=0.5
):
# Prepare the comprehensive prompt
full_prompt = self.prepare_prompt(message, history)
# Tokenize the prompt
inputs = self.tokenizer(full_prompt, return_tensors="pt", padding=True, truncation=True)
outputs = self.model.generate(
**inputs,
max_length=max(512, inputs["input_ids"].shape[1] + max_tokens)
temperature=temperature,
top_p=top_p,
no_repeat_ngram_size=3,
do_sample=True
)
# Decode the response
response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
# Optional: Post-processing to extract only the assistant's response
if "ASSISTANT'S COMPASSIONATE RESPONSE:" in response:
response = response.split("ASSISTANT'S COMPASSIONATE RESPONSE:")[-1].strip()
yield response
def main():
chatbot = MentalHealthChatbot()
demo = gr.ChatInterface(
chatbot.respond,
additional_inputs=[
gr.Slider(minimum=1, maximum=512, value=170, step=1, label="Max new tokens"),
gr.Slider(minimum=0.1, maximum=1.0, value=0.2, step=0.1, label="Temperature"),
gr.Slider(minimum=0.1, maximum=1.0, value=0.8, step=0.05, label="Top-p (nucleus sampling)"),
],
title="Mental Health Support Chatbot",
description="An AI assistant providing compassionate, culturally sensitive mental health support."
)
demo.launch()
if __name__ == "__main__":
main() |