Update app.py

app.py

@@ -1,88 +1,124 @@
 import gradio as gr
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
 import pandas as pd
+from collections import defaultdict
+from transformers import pipeline
 
-# Load the symptom dataset once
-df = pd.read_csv("enhanced_symptom_tree_with_measures.csv")
+# Load the Phi-2 model for symptom normalization
+symptom_normalizer = pipeline("text-classification", model="microsoft/phi-2")
 
-# 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
+# Predefined symptoms
+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]
+    }
+}
 
-tokenizer, model = load_model()
+# Global variables to track user state
+user_state = {}
 
-# Global variable to track last detected symptom
-last_detected_symptom = None  
+def normalize_symptom(user_input):
+    labels = list(symptom_data.keys())
+    predictions = symptom_normalizer(user_input, labels=labels)
+    return predictions[0]['label']
 
-# 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
+def chatbot(user_input):
+    if "state" not in user_state:
+        user_state["state"] = "greet"
     
-    if not user_input.strip():
-        return history, ""
+    if user_state["state"] == "greet":
+        user_state["state"] = "ask_symptom"
+        return "Hello! I'm a medical AI assistant. Please describe your primary symptom."
     
-    # 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()
-
-    history.append((user_input, response))
-    return history, ""
-
-# Function to clear chat
-def clear_chat():
-    global last_detected_symptom
-    last_detected_symptom = None  # Reset symptom tracking
-    return [], ""
+    elif user_state["state"] == "ask_symptom":
+        normalized_symptom = normalize_symptom(user_input)
+        if normalized_symptom not in symptom_data:
+            return "I don't recognize that symptom. Please enter one of these: " + ", ".join(symptom_data.keys())
+        user_state["symptom"] = normalized_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_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)
-    submit = gr.Button("Send")
-    clear = gr.Button("Clear Chat")
-    
-    submit.click(chatbot_response, [user_input, chatbot], [chatbot, user_input])
-    clear.click(clear_chat, [], [chatbot, user_input])
+demo = gr.Interface(fn=chatbot, inputs=gr.Textbox(placeholder="Enter your response..."), outputs="text", title="Symptom Chatbot")
 
-# Launch app in Hugging Face Spaces environment
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()