Update app.py

app.py

@@ -1,88 +1,168 @@
 import gradio as gr
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import pandas as pd
+import time
+from collections import defaultdict
+import spacy
 
-# Load the symptom dataset once
-df = pd.read_csv("enhanced_symptom_tree_with_measures.csv")
+# Load the symptom-to-disease mapping (same as before)
+symptom_data = {
+    "Shortness of breath": {
+        "questions": [
+            "Do you also have chest pain?",
+            "Do you feel fatigued often?",
+            "Have you noticed swelling in your legs?"
+        ],
+        "diseases": ["Atelectasis", "Emphysema", "Edema"],
+        "weights_yes": [30, 30, 40],
+        "weights_no": [10, 20, 30]
+    },
+    "Persistent cough": {
+        "questions": [
+            "Is your cough dry or with mucus?",
+            "Do you experience fever?",
+            "Do you have difficulty breathing?"
+        ],
+        "diseases": ["Pneumonia", "Fibrosis", "Infiltration"],
+        "weights_yes": [35, 30, 35],
+        "weights_no": [10, 15, 20]
+    },
+    "Sharp chest pain": {
+        "questions": [
+            "Does it worsen with deep breaths?",
+            "Do you feel lightheaded?",
+            "Have you had recent trauma or surgery?"
+        ],
+        "diseases": ["Pneumothorax", "Effusion", "Cardiomegaly"],
+        "weights_yes": [40, 30, 30],
+        "weights_no": [15, 20, 25]
+    },
+    "Fatigue & swelling": {
+        "questions": [
+            "Do you feel breathless when lying down?",
+            "Have you gained weight suddenly?",
+            "Do you experience irregular heartbeat?"
+        ],
+        "diseases": ["Edema", "Cardiomegaly"],
+        "weights_yes": [50, 30, 20],
+        "weights_no": [20, 15, 15]
+    },
+    "Chronic wheezing": {
+        "questions": [
+            "Do you have a history of smoking?",
+            "Do you feel tightness in your chest?",
+            "Do you have frequent lung infections?"
+        ],
+        "diseases": ["Emphysema", "Fibrosis"],
+        "weights_yes": [40, 30, 30],
+        "weights_no": [15, 25, 20]
+    }
+}
 
-# Load model and tokenizer once to optimize performance
-def load_model():
-    model_name = "microsoft/phi-2"  # Lightweight model
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")
-    model.eval()  # Set model to evaluation mode
-    return tokenizer, model
+# Load spaCy model for NLP
+nlp = spacy.load("en_core_web_lg")
 
-tokenizer, model = load_model()
-
-# Global variable to track last detected symptom
-last_detected_symptom = None  
-
-# Function to find symptom match
-def find_symptom_match(user_input):
-    match = df[df["Primary Symptom"].str.lower().str.contains(user_input.lower(), na=False)]
-    if not match.empty:
-        symptom = match.iloc[0]
-        response = f"It seems like you're experiencing {symptom['Primary Symptom']}. "
-        if pd.notna(symptom["Follow-up Question"]):
-            response += f"{symptom['Follow-up Question']} "
-        response += f"\nPossible conditions: {symptom['Possible Diseases']} \n"
-        response += f"Recommended measures: {symptom['Recommended Measures']}"
-        return symptom['Primary Symptom'], response  # Return symptom name and response
-    return None, None
-
-# Main chatbot response function
-def chatbot_response(user_input, history):
-    global last_detected_symptom  # Maintain previous symptom context
+# Function to extract key symptom from user input
+def extract_symptom(user_input):
+    # Define the symptoms that the chatbot recognizes
+    known_symptoms = list(symptom_data.keys())
     
-    if not user_input.strip():
-        return history, ""
+    # Process the input with spaCy NLP model
+    user_doc = nlp(user_input.lower())
     
-    # Step 1: If it's a follow-up response, continue from the last known symptom
-    if last_detected_symptom:
-        match = df[df["Primary Symptom"].str.lower() == last_detected_symptom.lower()]
-        if not match.empty:
-            follow_up_options = match.iloc[0]["Follow-up Question"]
-            if pd.notna(follow_up_options) and user_input.lower() in follow_up_options.lower():
-                response = f"Got it. Based on that, possible conditions: {match.iloc[0]['Possible Diseases']} \nRecommended: {match.iloc[0]['Recommended Measures']}"
-                last_detected_symptom = None  # Reset symptom tracking
-                history.append((user_input, response))
-                return history, ""
-
-    # Step 2: Otherwise, check for a new symptom
-    detected_symptom, symptom_info = find_symptom_match(user_input)
-    if symptom_info:
-        last_detected_symptom = detected_symptom  # Store new symptom for next response
-        response = symptom_info
-    else:
-        # Step 3: Use LLM as a fallback if no match is found
-        prompt = f"User: {user_input}\nChatbot:"
-        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-        with torch.no_grad():
-            outputs = model.generate(**inputs, max_length=150, pad_token_id=tokenizer.eos_token_id)
-        response = tokenizer.decode(outputs[0], skip_special_tokens=True).split("Chatbot:")[-1].strip()
+    # Check if any of the known symptoms are in the user input
+    for symptom in known_symptoms:
+        if symptom.lower() in user_input.lower():
+            return symptom
+    
+    # If no direct match, use similarity to find the closest symptom
+    similarities = {}
+    for symptom in known_symptoms:
+        symptom_doc = nlp(symptom.lower())
+        similarity = user_doc.similarity(symptom_doc)
+        similarities[symptom] = similarity
+    
+    # Return the symptom with the highest similarity
+    return max(similarities, key=similarities.get)
 
-    history.append((user_input, response))
-    return history, ""
+# Global variables to track user state
+user_state = {}
 
-# Function to clear chat
-def clear_chat():
-    global last_detected_symptom
-    last_detected_symptom = None  # Reset symptom tracking
-    return [], ""
+def chatbot(user_input):
+    if "state" not in user_state:
+        user_state["state"] = "greet"
+    
+    if user_state["state"] == "greet":
+        user_state["state"] = "ask_symptom"
+        return "Hello! I'm a medical AI assistant. Please describe your primary symptom."
+    
+    elif user_state["state"] == "ask_symptom":
+        # Extract symptom from the user input
+        matched_symptom = extract_symptom(user_input)
+        
+        if matched_symptom not in symptom_data:
+            user_state["state"] = "ask_feeling"
+            return "I'm sorry, I don't recognize that symptom. How do you feel?"
+        
+        user_state["symptom"] = matched_symptom
+        user_state["state"] = "ask_duration"
+        return "How long have you been experiencing this symptom? (Less than a week / More than a week)"
+    
+    elif user_state["state"] == "ask_feeling":
+        # If the symptom is not recognized, ask how they feel
+        return "Can you describe your symptoms in more detail?"
+    
+    elif user_state["state"] == "ask_duration":
+        if user_input.lower() == "less than a week":
+            user_state.clear()
+            return "It might be a temporary issue. Please monitor your symptoms and consult a doctor if they persist."
+        elif user_input.lower() == "more than a week":
+            user_state["state"] = "follow_up"
+            user_state["current_question"] = 0
+            user_state["disease_scores"] = defaultdict(int)
+            return symptom_data[user_state["symptom"]]["questions"][0]
+        else:
+            return "Please respond with 'Less than a week' or 'More than a week'."
+    
+    elif user_state["state"] == "follow_up":
+        symptom = user_state["symptom"]
+        question_index = user_state["current_question"]
+        
+        # Update probabilities
+        if user_input.lower() == "yes":
+            for i, disease in enumerate(symptom_data[symptom]["diseases"]):
+                user_state["disease_scores"][disease] += symptom_data[symptom]["weights_yes"][i]
+        else:
+            for i, disease in enumerate(symptom_data[symptom]["diseases"]):
+                user_state["disease_scores"][disease] += symptom_data[symptom]["weights_no"][i]
+        
+        # Move to the next question or finish
+        user_state["current_question"] += 1
+        if user_state["current_question"] < len(symptom_data[symptom]["questions"]):
+            return symptom_data[symptom]["questions"][user_state["current_question"]]
+        
+        # Final diagnosis
+        probable_disease = max(user_state["disease_scores"], key=user_state["disease_scores"].get)
+        user_state.clear()
+        return f"Based on your symptoms, the most likely condition is: {probable_disease}. Please consult a doctor for confirmation."
 
-# Gradio UI
-with gr.Blocks(theme="compact") as demo:
-    gr.Markdown("### Symptom Chatbot 🏥")
-    chatbot = gr.Chatbot()
-    user_input = gr.Textbox(placeholder="Type your symptom and get advice...", interactive=True)
+# Gradio Chatbot UI with improved features
+with gr.Blocks() as demo:
+    gr.Markdown("# Conversational Image Recognition Assistant: AI-Powered X-ray Diagnosis for Healthcare")
+    chatbot_ui = gr.Chatbot()
+    user_input = gr.Textbox(placeholder="Enter your response...", label="Your Message")
     submit = gr.Button("Send")
-    clear = gr.Button("Clear Chat")
+    clear_chat = gr.Button("Clear Chat")
+    
+    def respond(user_message, history):
+        history.append((user_message, "Thinking..."))  # Show thinking message
+        yield history, ""  # Immediate update
+        
+        time.sleep(1.5)  # Simulate processing delay
+        bot_response = chatbot(user_message)
+        history[-1] = (user_message, bot_response)  # Update with real response
+        yield history, ""
     
-    submit.click(chatbot_response, [user_input, chatbot], [chatbot, user_input])
-    clear.click(clear_chat, [], [chatbot, user_input])
+    submit.click(respond, [user_input, chatbot_ui], [chatbot_ui, user_input])
+    user_input.submit(respond, [user_input, chatbot_ui], [chatbot_ui, user_input])
+    clear_chat.click(lambda: ([], ""), outputs=[chatbot_ui, user_input])
 
-# Launch app in Hugging Face Spaces environment
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()