Update app.py

app.py

@@ -1,12 +1,8 @@
 import gradio as gr
 import pandas as pd
 from collections import defaultdict
-from transformers import pipeline
 
-# Load the Phi-2 model for symptom normalization
-symptom_normalizer = pipeline("text-classification", model="microsoft/phi-2")
-
-# Predefined symptoms
+# Load the symptom-to-disease mapping
 symptom_data = {
     "Shortness of breath": {
         "questions": [
@@ -63,12 +59,7 @@ symptom_data = {
 # Global variables to track user state
 user_state = {}
 
-def normalize_symptom(user_input):
-    labels = list(symptom_data.keys())
-    predictions = symptom_normalizer(user_input, labels=labels)
-    return predictions[0]['label']
-
-def chatbot(user_input):
+def chatbot(user_input, history):
     if "state" not in user_state:
         user_state["state"] = "greet"
     
@@ -77,10 +68,9 @@ def chatbot(user_input):
         return "Hello! I'm a medical AI assistant. Please describe your primary symptom."
     
     elif user_state["state"] == "ask_symptom":
-        normalized_symptom = normalize_symptom(user_input)
-        if normalized_symptom not in symptom_data:
+        if user_input 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["symptom"] = user_input
         user_state["state"] = "ask_duration"
         return "How long have you been experiencing this symptom? (Less than a week / More than a week)"
     
@@ -118,7 +108,18 @@ def chatbot(user_input):
         user_state.clear()
         return f"Based on your symptoms, the most likely condition is: {probable_disease}. Please consult a doctor for confirmation."
 
-# Gradio UI
-demo = gr.Interface(fn=chatbot, inputs=gr.Textbox(placeholder="Enter your response..."), outputs="text", title="Symptom Chatbot")
+# Gradio Chatbot UI
+with gr.Blocks() as demo:
+    gr.Markdown("# Symptom Chatbot")
+    chatbot_ui = gr.Chatbot()
+    user_input = gr.Textbox(placeholder="Enter your response...")
+    submit = gr.Button("Send")
+    
+    def respond(user_message, history):
+        bot_response = chatbot(user_message, history)
+        history.append((user_message, bot_response))
+        return history, ""
+    
+    submit.click(respond, [user_input, chatbot_ui], [chatbot_ui, user_input])
 
 demo.launch()