import gradio as gr from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig import torch import re # ======================================================= # Model Loading with Error Handling # ======================================================= model_name = "augtoma/qCammel-13" print("Loading tokenizer and model...") try: tokenizer = AutoTokenizer.from_pretrained(model_name) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token model = AutoModelForCausalLM.from_pretrained( model_name, device_map="auto", torch_dtype=torch.float16, trust_remote_code=True, low_cpu_mem_usage=True ) model.eval() print("āœ… Model loaded successfully!") except Exception as e: print(f"āŒ Error loading model: {e}") raise # ======================================================= # Constants # ======================================================= MEDICAL_DISCLAIMER = """ āš ļø **IMPORTANT MEDICAL DISCLAIMER:** - This is an AI simulation for educational purposes only - For medical emergencies, call emergency services immediately (911/999/local emergency number) - Always consult a licensed healthcare provider for diagnosis and treatment - Do not use this for serious, worsening, or life-threatening symptoms - This AI cannot diagnose conditions or prescribe medications """ EMERGENCY_KEYWORDS = [ "chest pain", "can't breathe", "breathing difficulty", "suicide", "heart attack", "stroke", "severe bleeding", "unconscious", "overdose", "severe pain", "emergency" ] # ======================================================= # Helper Functions # ======================================================= def extract_name(text): """Extract name from user input""" text = text.lower() # Remove common phrases for phrase in ["yes", "i am", "i'm", "my name is", "name is", "it's", "its"]: text = text.replace(phrase, "") # Clean and capitalize text = re.sub(r'[^\w\s]', '', text).strip() return text.title() if text and len(text) > 1 else "Patient" def extract_age(text): """Extract age from user input""" # Look for numbers in the text numbers = re.findall(r'\b(\d{1,3})\b', text) for num in numbers: age = int(num) if 1 <= age <= 120: # Reasonable age range return age return None def check_emergency(text): """Check if message contains emergency keywords""" text_lower = text.lower() return any(keyword in text_lower for keyword in EMERGENCY_KEYWORDS) def create_initial_session(): """Create a fresh session dictionary""" return { "name": None, "age": None, "gender": None, "symptoms": None, "duration": None, "medication": None, "stage": "intro" } # ======================================================= # Diet Suggestions # ======================================================= def get_diet_suggestion(name, symptoms): """Get personalized diet suggestions based on symptoms""" symptoms_lower = symptoms.lower() if "fever" in symptoms_lower: return f"""šŸŽ **Diet Plan for {name} (Fever)** **Meals:** - šŸ„£ **Breakfast:** Oatmeal or boiled egg with banana/apple - šŸ² **Lunch:** Rice with lentil soup or boiled vegetables - šŸ› **Dinner:** Light soup or grilled chicken with plain rice **Hydration:** - šŸ’§ Drink 8ā€“10 glasses of water daily - Coconut water or clear soups - Herbal teas (chamomile, ginger) **Avoid:** - šŸš« Fried, oily, or spicy foods - Heavy meals - Caffeinated drinks **Tips:** Eat small portions frequently to help your body recover.""" elif any(word in symptoms_lower for word in ["stomach", "vomit", "nausea", "diarrhea"]): return f"""šŸ„— **Diet Plan for {name} (Stomach Issues)** **BRAT Diet:** - šŸŒ Bananas - šŸš Rice (plain white rice) - šŸŽ Applesauce - šŸž Toast (dry or with minimal butter) **Hydration:** - šŸ’§ Sip water frequently - Oral rehydration solution (ORS) - Clear broths **Avoid:** - šŸš« Dairy products - Fried or spicy foods - Caffeine and alcohol - Raw vegetables **Tips:** Start with small amounts and gradually increase as you feel better.""" elif any(word in symptoms_lower for word in ["cold", "flu", "cough", "sore throat"]): return f"""šŸŠ **Diet Plan for {name} (Cold/Flu)** **Immunity Boosters:** - šŸ‹ Citrus fruits (oranges, lemons) - šŸÆ Honey with warm water - šŸ§„ Garlic (natural antibiotic) - šŸ§… Onions **Meals:** - šŸ² Chicken soup (helps clear congestion) - ā˜• Ginger or green tea - šŸ„£ Warm porridge or oatmeal **Hydration:** - šŸ’§ Plenty of warm fluids - Herbal teas **Avoid:** - šŸš« Sugary drinks - Chilled beverages - Dairy (may increase mucus) **Tips:** Stay warm and get plenty of rest.""" elif any(word in symptoms_lower for word in ["headache", "migraine"]): return f"""šŸ§˜ **Diet Plan for {name} (Headache)** **Helpful Foods:** - šŸ’§ Stay well-hydrated (dehydration causes headaches) - šŸ„” Magnesium-rich: bananas, nuts, spinach - šŸŸ Omega-3: fish, walnuts - šŸ§‚ Small amounts of salt if low blood pressure **Avoid:** - šŸš« Processed foods - Artificial sweeteners - Aged cheeses - Alcohol - Excessive caffeine **Tips:** Eat regular meals; skipping meals can trigger headaches.""" else: return f"""šŸ„— **General Healthy Diet Plan for {name}** **Balanced Meals:** - šŸŽ Fresh fruits and vegetables - šŸ— Lean proteins (chicken, fish, legumes) - šŸš Whole grains (brown rice, oats) - šŸ„œ Nuts and seeds **Hydration:** - šŸ’§ 8ā€“10 glasses of water daily **Avoid:** - šŸš« Processed and fried foods - Excessive sugar and salt - Junk food **Tips:** Eat at regular intervals and maintain portion control.""" # ======================================================= # AI Response Generation # ======================================================= def generate_ai_response(session, user_message): """Generate AI consultation response using the model""" try: # Build the prompt prompt = f"""You are Dr. Aiden ā€” a warm, caring, and professional doctor. You are consulting a {session['age']}-year-old {session['gender']} named {session['name']}. Patient details: - Symptoms: {session['symptoms']} - Duration: {session['duration']} - Current medications: {session['medication']} Now the patient says: "{user_message}" Respond as a real doctor would ā€” empathetic, clear, and personalized. Include: 1. Acknowledge their condition 2. Possible causes 3. Simple home remedies or OTC medicines (if safe) 4. Diet, rest, and hydration tips 5. When to visit a real doctor 6. End with a reassuring tone""" # Tokenize and generate inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=2048).to(model.device) gen_cfg = GenerationConfig( temperature=0.7, top_p=0.9, max_new_tokens=350, repetition_penalty=1.1, pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id ) with torch.no_grad(): output = model.generate(**inputs, generation_config=gen_cfg) # Decode the output output_text = tokenizer.decode(output[0], skip_special_tokens=True) # Try to extract just the response part # The model might repeat the prompt, so we try to find where the actual response starts if "reassuring tone" in output_text: parts = output_text.split("reassuring tone") if len(parts) > 1: output_text = parts[-1].strip() # If the output contains the prompt, try to split it if user_message in output_text: parts = output_text.split(user_message) if len(parts) > 1: output_text = parts[-1].strip() # Clean up output_text = output_text.strip() # If output is too short or empty, provide a fallback if len(output_text) < 20: output_text = f"Hello {session['name']}, I understand you have {session['symptoms']} for {session['duration']}. This could be due to several reasons. I recommend staying hydrated, getting plenty of rest, and monitoring your temperature. If your fever persists beyond 5 days or gets worse, please visit a doctor in person for proper examination." # Ensure proper ending punctuation if not output_text.endswith((".", "!", "?")): output_text += "." # Add disclaimer output_text += "\n\nāš•ļø *Note: This advice is AI-generated and not a substitute for professional medical care.*" return output_text except Exception as e: print(f"Error in generate_ai_response: {e}") return f"I understand you're experiencing {session['symptoms']}, {session['name']}. Since this has been going on for {session['duration']}, I recommend consulting with a healthcare provider for a proper examination and diagnosis.\n\nāš•ļø *Note: This advice is AI-generated and not a substitute for professional medical care.*" # ======================================================= # Doctor Response Logic # ======================================================= def doctor_response(user_message, session): """Generate doctor's response based on conversation stage""" if not user_message or not user_message.strip(): return "Please type your message.", session user_message_clean = user_message.strip() user_message_lower = user_message_clean.lower() # Check for emergency situations if check_emergency(user_message_clean): emergency_response = """šŸšØ **EMERGENCY ALERT** šŸšØ Your message suggests a potentially serious medical emergency. **PLEASE TAKE IMMEDIATE ACTION:** - Call emergency services NOW (911 in US, 999 in UK, or your local emergency number) - Go to the nearest emergency room - Do not delay seeking immediate medical attention This AI chatbot cannot handle emergencies. Your safety is the priority.""" return emergency_response, session # Stage-based conversation flow try: if session["stage"] == "intro": session["stage"] = "ask_name" return "šŸ‘Øā€āš•ļø Hello! I'm Dr. Aiden, your AI medical assistant.\n\n" + MEDICAL_DISCLAIMER + "\n\nMay I know your name, please?", session elif session["stage"] == "ask_name": session["name"] = extract_name(user_message_clean) session["stage"] = "ask_age" return f"Nice to meet you, {session['name']}! How old are you?", session elif session["stage"] == "ask_age": age = extract_age(user_message_clean) if age: session["age"] = age session["stage"] = "ask_gender" return f"Got it, {session['name']}. Are you male or female?", session else: return "Please tell me your age in numbers (e.g., 25, 30, 45).", session elif session["stage"] == "ask_gender": if "male" in user_message_lower and "female" not in user_message_lower: session["gender"] = "male" elif "female" in user_message_lower: session["gender"] = "female" else: return "Could you please specify whether you are male or female?", session session["stage"] = "ask_symptoms" return f"Thanks, {session['name']}! So you're a {session['age']}-year-old {session['gender']}.\n\nWhat symptoms are you experiencing? Please describe them in detail.", session elif session["stage"] == "ask_symptoms": if len(user_message_clean) < 5: return "Please describe your symptoms in more detail so I can help you better.", session session["symptoms"] = user_message_clean session["stage"] = "ask_duration" return "Since when have you been feeling this way? (e.g., 2 days, 1 week, etc.)", session elif session["stage"] == "ask_duration": session["duration"] = user_message_clean session["stage"] = "ask_medicine" return "Are you currently taking any medications or treatments? If not, just say 'no'.", session elif session["stage"] == "ask_medicine": session["medication"] = user_message_clean session["stage"] = "consult" # Provide summary and automatically generate initial consultation summary = f"""Thank you for providing all the information, {session['name']}. šŸ“‹ **Summary:** - **Name:** {session['name']} - **Age:** {session['age']} years old - **Gender:** {session['gender']} - **Symptoms:** {session['symptoms']} - **Duration:** {session['duration']} - **Current medication:** {session['medication']} Let me assess your condition...\n\n""" # Generate initial consultation automatically initial_consultation = generate_ai_response(session, "Please assess my condition and give me advice") return summary + initial_consultation, session elif session["stage"] == "consult": # Handle diet-specific queries if any(word in user_message_lower for word in ["diet", "food", "meal", "eat", "nutrition", "drink"]): diet_plan = get_diet_suggestion(session["name"], session["symptoms"]) return diet_plan + "\n\nāš•ļø *Note: This is general advice. Consult a healthcare provider for personalized dietary recommendations, especially if you have allergies or other conditions.*", session # Generate AI consultation response for any other question return generate_ai_response(session, user_message_clean), session else: # Reset if stage is unknown session["stage"] = "intro" return "Let's start over. Say 'Hi Doctor' to begin your consultation.", session except Exception as e: print(f"Error in doctor_response: {e}") return "I apologize, but I encountered an error. Please try again or start a new consultation.", session # ======================================================= # Gradio Interface # ======================================================= with gr.Blocks(theme=gr.themes.Soft()) as demo: gr.HTML("""

šŸ„ Dr. Aiden - AI Medical Assistant

Compassionate AI-powered health consultation

āš ļø For emergencies, call 911 immediately

""") # Session state session_state = gr.State(create_initial_session) chatbot = gr.Chatbot( label="šŸ‘Øā€āš•ļø Chat with Dr. Aiden", height=500, type='messages', avatar_images=( "https://cdn-icons-png.flaticon.com/512/706/706830.png", "https://cdn-icons-png.flaticon.com/512/3774/3774299.png" ) ) with gr.Row(): user_input = gr.Textbox( placeholder="Type your message here (say 'Hi Doctor' to start)...", label="Your Message", lines=2, scale=4 ) with gr.Row(): send_btn = gr.Button("šŸ’¬ Send", variant="primary", scale=1) clear_btn = gr.Button("šŸ”„ New Consultation", variant="secondary", scale=1) gr.HTML("""
āš ļø Important Reminders:
""") def respond(message, history, session): """Handle user message and update chat""" if history is None: history = [] if not message or not message.strip(): return "", history, session response, updated_session = doctor_response(message, session) history.append({"role": "user", "content": message}) history.append({"role": "assistant", "content": response}) return "", history, updated_session def reset(): """Reset conversation and session""" return [], create_initial_session() # Event handlers send_btn.click( respond, [user_input, chatbot, session_state], [user_input, chatbot, session_state] ) user_input.submit( respond, [user_input, chatbot, session_state], [user_input, chatbot, session_state] ) clear_btn.click( reset, None, [chatbot, session_state], queue=False ) if __name__ == "__main__": print("šŸ„ Launching Dr. Aiden Medical Assistant...") print("āš ļø Remember: This is for educational purposes only!") demo.queue() demo.launch(share=True)