File size: 5,787 Bytes
ba3d35f b0f7227 ba3d35f 5e888bb b0f7227 59acdef b312972 b0f7227 59acdef 5e888bb 59acdef 5e888bb b0f7227 5e888bb b0f7227 5e888bb b0f7227 5e888bb b0f7227 5e888bb b0f7227 59acdef 5e888bb 59acdef ba3d35f 5e888bb b0f7227 |
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 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 |
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from threading import Thread
import torch
# Check if CUDA is available
device = "cuda" if torch.cuda.is_available() else "cpu"
# Load model and tokenizer
model_name = "HuggingFaceH4/zephyr-7b-beta" # You can change this to any model you have access to
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype=torch.float16 if device == "cuda" else torch.float32,
device_map="auto"
)
# Define the medical system prompt
MEDICAL_SYSTEM_PROMPT = """You are a medical AI assistant. You MUST ONLY answer health and medical-related questions.
Your responses should be professional, accurate, and focused on medical topics only.
For any non-medical questions, respond with a redirection to medical topics.
For medication queries, provide general information and recommend consulting a healthcare professional.
USE BULLET POINTS TO STRUCTURE YOUR ANSWERS AND THEY MUST BE PUNCTUATED AND GRAMMATICALLY CORRECT.
Your responses should be concise and informative.
For Questions About Body Parts/Organs:
- Provide anatomical details
- Explain primary functions
- Describe structure
- Mention related conditions
For Medication Queries:
- List generic/brand names
- Mention drug class
- Explain general uses
- Note common side effects
- List important drug interactions
- Include a disclaimer about consulting healthcare providers
For Symptom Analysis:
- List possible causes
- Describe key characteristics
- Suggest self-care measures
- Highlight when to seek immediate medical care
For Treatment Information:
- Explain conservative measures
- Describe medical interventions
- Suggest prevention strategies
- Outline follow-up care
For Emergency Questions:
- Provide immediate actions
- Suggest temporary measures while waiting for help
- Clarify when to contact emergency services
For Mental Health Questions:
- List common symptoms
- Suggest coping strategies
- Explain when to seek professional help
- Include crisis support information
For Preventive Care:
- Provide lifestyle recommendations
- Mention screening tests
- Include vaccination information if applicable
For Diet/Nutrition:
- Give dietary recommendations
- Highlight key nutrients
- Offer practical meal planning tips"""
# Function to check if a query is medical-related
def is_medical_query(query):
medical_keywords = [
"health", "disease", "symptom", "doctor", "medicine", "medical", "treatment",
"hospital", "clinic", "diagnosis", "patient", "drug", "prescription", "therapy",
"cancer", "diabetes", "heart", "blood", "pain", "surgery", "vaccine", "infection",
"allergy", "diet", "nutrition", "vitamin", "exercise", "mental health", "depression",
"anxiety", "disorder", "syndrome", "chronic", "acute", "emergency", "pharmacy",
"dosage", "side effect", "contraindication", "body", "organ", "immune", "virus",
"bacterial", "fungal", "parasite", "genetic", "hereditary", "congenital", "prenatal",
"headaches", "ache", "stomach ache", "skin", "head", "arm", "leg", "chest", "back",
"throat", "eye", "ear", "nose", "mouth"
]
query_lower = query.lower()
return any(keyword in query_lower for keyword in medical_keywords)
# Response function for the chatbot
def respond(message, history, max_tokens=512, temperature=0.7, top_p=0.95):
# Check if the query is medical-related
if not is_medical_query(message):
return "I'm specialized in medical topics only. I cannot answer this question. How can I assist with a health-related concern instead?"
# Format the conversation history
formatted_history = []
for user_msg, assistant_msg in history:
formatted_history.append({"role": "user", "content": user_msg})
formatted_history.append({"role": "assistant", "content": assistant_msg})
# Prepare the prompt
messages = [
{"role": "system", "content": MEDICAL_SYSTEM_PROMPT},
*formatted_history,
{"role": "user", "content": message}
]
# Format messages for the model
prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
# Tokenize the prompt
inputs = tokenizer(prompt, return_tensors="pt").to(device)
# Create the streamer
streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
# Generate in a separate thread
generation_kwargs = dict(
inputs,
streamer=streamer,
max_new_tokens=max_tokens,
do_sample=True,
temperature=temperature,
top_p=top_p,
)
thread = Thread(target=model.generate, kwargs=generation_kwargs)
thread.start()
# Stream the response
partial_response = ""
for new_text in streamer:
partial_response += new_text
yield partial_response
# If no text was generated, return an error message
if not partial_response:
yield "I apologize, but I encountered an error generating a response. Please try again."
# Create the Gradio interface
demo = gr.ChatInterface(
fn=respond,
title="MedexDroid - Medical Assistant",
examples=[
"What are the symptoms of diabetes?",
"How can I improve my diet for heart health?",
"What is the treatment for a migraine?",
"What are the side effects of aspirin?",
"What are the causes of high blood pressure?"
],
description="An AI Medical Assistant. Please ask health-related questions only.",
theme=gr.themes.Soft(),
css=".gradio-container {background-color: #f0f4f8}"
)
if __name__ == "__main__":
demo.launch()
|