import gradio as gr from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig import torch model_name = "augtoma/qCammel-13" print("Loading tokenizer and model...") 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!") # ======================================================= # Global Session Memory # ======================================================= session = { "name": None, "age": None, "gender": None, "symptoms": None, "duration": None, "stage": "intro" } # ======================================================= # Helper: Extract Name # ======================================================= def extract_name(text): text = text.lower().replace("yes", "").replace("i am", "").replace("i'm", "") text = text.replace("my name is", "").replace("name", "").replace("is", "").strip() return text.title() if text else "Patient" # ======================================================= # Doctor Response Logic # ======================================================= def doctor_response(user_message): global session user_message = user_message.strip().lower() # Step 1: Greeting if session["stage"] == "intro": session["stage"] = "ask_name" return "π¨ββοΈ Hello! Iβm Dr. Aiden. May I know your name, please?" # Step 2: Get Name elif session["stage"] == "ask_name": session["name"] = extract_name(user_message) session["stage"] = "ask_age" return f"Nice to meet you, {session['name']}! How old are you?" # Step 3: Get Age elif session["stage"] == "ask_age": words = user_message.split() for w in words: if w.isdigit(): session["age"] = int(w) session["stage"] = "ask_gender" return f"Got it, {session['name']}. Are you male or female?" return "Please tell me your age in numbers, like 20 or 25." # Step 4: Get Gender elif session["stage"] == "ask_gender": if "male" in user_message: session["gender"] = "male" elif "female" in user_message: session["gender"] = "female" else: return "Could you please specify whether you are male or female?" session["stage"] = "ask_symptoms" return f"Thanks, {session['name']}! So you're a {session['age']}-year-old {session['gender']}. What symptoms are you experiencing?" # Step 5: Get Symptoms elif session["stage"] == "ask_symptoms": session["symptoms"] = user_message session["stage"] = "ask_duration" return "Since when have you been feeling this way?" # Step 6: Duration elif session["stage"] == "ask_duration": session["duration"] = user_message session["stage"] = "ask_medicine" return "Got it. Are you taking any medications or treatments currently?" # Step 7: Medicine intake elif session["stage"] == "ask_medicine": session["medication"] = user_message session["stage"] = "consult" return "Thank you. Letβs discuss what could be going on and how you can manage it safely." # Step 8: Personalized Consultation elif session["stage"] == "consult": name = session["name"] age = session["age"] gender = session["gender"] symptoms = session.get("symptoms", "") duration = session.get("duration", "") medication = session.get("medication", "") # ------------- Handle diet-related questions ----------------- if any(word in user_message for word in ["diet", "food", "meal", "eat", "nutrition"]): # Diet suggestions tailored to the condition if "fever" in symptoms: diet = ( f"π {name}, since you have a fever, try keeping your meals light and hydrating.\n\n" "- π₯£ **Breakfast:** Oatmeal or boiled egg with fruit (like banana or apple)\n" "- π² **Lunch:** Rice with lentil soup or boiled vegetables\n" "- π **Dinner:** Light soup or grilled chicken with plain rice\n" "- π§ **Hydration:** Drink 8β10 glasses of water, coconut water, or clear soups\n" "- π« **Avoid:** Fried, oily, or spicy foods\n\n" "Eat small portions often β this helps your body recover faster." ) elif "stomach" in symptoms or "vomit" in symptoms: diet = ( f"π₯ {name}, for stomach discomfort, stick to a **bland diet**:\n\n" "- π Toast, plain rice, boiled potatoes, or bananas\n" "- π§ Sip water, oral rehydration solution, or herbal tea\n" "- π« Avoid milk, fried, or spicy foods\n" "- π Eat small meals to avoid nausea" ) elif "cold" in symptoms or "flu" in symptoms or "cough" in symptoms: diet = ( f"π {name}, for flu or cold, focus on **immunity-boosting foods**:\n\n" "- π Citrus fruits, honey with warm water, and soups\n" "- π² Chicken soup helps clear congestion\n" "- β Ginger or green tea for throat relief\n" "- π« Avoid sugary and chilled drinks" ) else: diet = ( f"π₯ {name}, eat a balanced diet: fruits, vegetables, lean protein, and whole grains. " "Stay hydrated and avoid processed or fried foods." ) return diet + "\n\nβοΈ *Note: This is general advice β not a substitute for medical care.*" # ----------------- Normal medical consultation ---------------- prompt = f""" You are Dr. Aiden β a warm, caring, and professional doctor. You are consulting a {age}-year-old {gender} named {name}. Patient details: - Symptoms: {symptoms} - Duration: {duration} - Current medications: {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 """ 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) output_text = tokenizer.decode(output[0], skip_special_tokens=True) output_text = output_text.split("Doctor:")[-1].strip() if not output_text.endswith((".", "!", "?")): output_text += "." output_text += "\n\nβοΈ *Note: This advice is AI-generated and not a substitute for professional medical care.*" return output_text # ======================================================= # Gradio Interface # ======================================================= with gr.Blocks(theme=gr.themes.Soft()) as demo: gr.HTML("""
AI-powered doctor interview β step-by-step and caring conversation