Update app.py

added a description for diagnose

app.py

@@ -1,52 +1,119 @@
-from transformers import AutoTokenizer, TFAutoModelForSequenceClassification
-import tensorflow as tf
 import gradio as gr
+from transformers import TFAutoModelForSequenceClassification, AutoTokenizer
+import tensorflow as tf
+import numpy as np
 
-# Load the tokenizer and model
-model_name = "Zabihin/Symptom_to_Diagnosis"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+# Load the pre-trained model and tokenizer
+model_name = "your_model_name_here"
 model = TFAutoModelForSequenceClassification.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-# Clean the input text
-def clean_input(symptom_text):
-    # Remove unwanted characters or non-ASCII characters
-    symptom_text = ''.join([c for c in symptom_text if ord(c) < 128])
-    symptom_text = symptom_text.lower()  # Optional: Convert to lowercase
-    return symptom_text
-
-# Define the predict function
+# Function to predict the disease based on symptoms
 def predict(symptom_text, chat_history=[]):
-    try:
-        # Clean the input
-        symptom_text = clean_input(symptom_text)
-
-        # Tokenize the input
-        inputs = tokenizer(symptom_text, return_tensors="tf", padding=True, truncation=True, max_length=512)
+    # Preprocess the input (lowercase, remove special characters)
+    symptom_text = symptom_text.lower().strip()
 
-        # Get model output
-        outputs = model(**inputs)
-        logits = outputs.logits
-        prediction = tf.argmax(logits, axis=-1).numpy()[0]
+    # Tokenize the input symptoms
+    inputs = tokenizer(symptom_text, return_tensors="tf", padding=True, truncation=True, max_length=512)
 
-        # Map the prediction to a label
-        labels = {
-            0: "Allergy", 1: "Arthritis", 2: "Bronchial Asthma", 3: "Cervical Spondylosis", 
-            4: "Chicken Pox", 5: "Common Cold", 6: "Dengue", 7: "Diabetes", 8: "Drug Reaction", 
-            9: "Fungal Infection", 10: "Gastroesophageal Reflux Disease", 11: "Hypertension", 
-            12: "Impetigo", 13: "Jaundice", 14: "Malaria", 15: "Migraine", 16: "Peptic Ulcer Disease", 
-            17: "Pneumonia", 18: "Psoriasis", 19: "Typhoid", 20: "Urinary Tract Infection", 21: "Varicose Veins"
-        }
+    # Make the prediction
+    outputs = model(**inputs)
+    prediction = tf.argmax(outputs.logits, axis=-1).numpy()[0]
+    predicted_label = get_prediction_label(prediction)
+    
+    # Provide a more descriptive output
+    response = f"Predicted Condition: {predicted_label}\n\n"
+    response += f"Description: {get_condition_description(predicted_label)}\n\n"
+    response += "Common symptoms for this condition include:\n"
+    response += get_condition_symptoms(predicted_label)
+    
+    chat_history.append(("User", symptom_text))
+    chat_history.append(("AI", response))
+    return response, chat_history
 
-        diagnosis = labels.get(prediction, "Unknown diagnosis")
+# Function to map the prediction index to a condition
+def get_prediction_label(prediction_index):
+    # Mapping of prediction index to disease
+    condition_map = {
+        0: "Common Cold",  # Common Cold
+        1: "Flu",          # Flu
+        2: "Covid-19",     # Covid-19
+        3: "Diabetes",     # Diabetes
+        4: "Pneumonia",    # Pneumonia
+        5: "Allergy",      # Allergy
+        6: "Asthma",       # Asthma
+        7: "Cancer",       # Cancer
+        8: "Heart Disease", # Heart Disease
+        9: "Hypertension",  # Hypertension
+        10: "Gastroenteritis", # Gastroenteritis
+        11: "Tuberculosis", # Tuberculosis
+        12: "Migraine",     # Migraine
+        13: "Chronic Obstructive Pulmonary Disease",  # COPD
+        14: "Stroke",       # Stroke
+        15: "Acid Reflux",   # Acid Reflux
+        16: "Kidney Disease", # Kidney Disease
+        17: "Liver Disease",  # Liver Disease
+        18: "Anemia",       # Anemia
+        19: "Sepsis",       # Sepsis
+        # Add more as needed...
+    }
+    return condition_map.get(prediction_index, "Unknown Condition")
 
-        # Add conversation history
-        chat_history.append(("User", symptom_text))
-        chat_history.append(("AI", f"Predicted Diagnosis: {diagnosis}. Please consult a doctor for more accurate results."))
+# Function to get the description for a condition
+def get_condition_description(condition):
+    # Description of conditions
+    condition_descriptions = {
+        "Common Cold": "A viral infection of your upper respiratory system, causing symptoms like a runny nose, sore throat, and cough.",
+        "Flu": "An infectious respiratory illness caused by influenza viruses, with symptoms including fever, chills, and body aches.",
+        "Covid-19": "A respiratory illness caused by the SARS-CoV-2 virus, with symptoms ranging from mild cold-like symptoms to severe respiratory distress.",
+        "Diabetes": "A group of diseases that affect how your body uses blood sugar (glucose), leading to high blood sugar levels over time.",
+        "Pneumonia": "An infection that inflames the air sacs in one or both lungs, which can lead to symptoms like cough, fever, and difficulty breathing.",
+        "Allergy": "A reaction of your immune system to a substance that is usually harmless, leading to symptoms such as sneezing, itching, and swelling.",
+        "Asthma": "A chronic condition in which your airways narrow and swell, causing difficulty breathing, wheezing, and coughing.",
+        "Cancer": "A group of diseases involving abnormal cell growth that can spread to other parts of the body, often accompanied by unexplained weight loss and fatigue.",
+        "Heart Disease": "A range of conditions that affect the heart, including coronary artery disease, arrhythmias, and heart attacks, leading to chest pain and shortness of breath.",
+        "Hypertension": "High blood pressure, which can lead to serious health problems such as heart disease and stroke if left untreated.",
+        "Gastroenteritis": "Inflammation of the stomach and intestines, typically caused by a viral or bacterial infection, leading to diarrhea, vomiting, and stomach cramps.",
+        "Tuberculosis": "A potentially serious bacterial infection that primarily affects the lungs, causing persistent cough, chest pain, and weight loss.",
+        "Migraine": "A neurological condition characterized by intense headaches, often accompanied by nausea, vomiting, and sensitivity to light and sound.",
+        "Chronic Obstructive Pulmonary Disease": "A group of lung diseases, including emphysema and chronic bronchitis, that cause long-term breathing problems.",
+        "Stroke": "A medical emergency in which the blood supply to part of the brain is interrupted, leading to symptoms such as numbness, difficulty speaking, and confusion.",
+        "Acid Reflux": "A condition in which stomach acid frequently flows back into the esophagus, causing heartburn and discomfort in the chest area.",
+        "Kidney Disease": "A condition in which the kidneys are damaged and unable to filter blood properly, leading to symptoms such as fatigue and swelling in the legs.",
+        "Liver Disease": "A group of conditions that damage the liver, leading to symptoms such as jaundice, abdominal pain, and swelling.",
+        "Anemia": "A condition where the blood doesn't have enough healthy red blood cells to carry oxygen to your body's tissues, causing fatigue and weakness.",
+        "Sepsis": "A life-threatening condition caused by the body's response to an infection, leading to symptoms such as fever, confusion, and rapid heart rate.",
+        # Add more conditions and descriptions as necessary...
+    }
+    return condition_descriptions.get(condition, "No description available.")
 
-    except Exception as e:
-        chat_history.append(("AI", f"Error: {str(e)}"))
-
-    return chat_history, ""
+# Function to get the symptoms for a condition
+def get_condition_symptoms(condition):
+    # Mapping of condition to common symptoms
+    condition_symptoms = {
+        "Common Cold": "Cough, runny nose, sore throat, congestion, mild headache",
+        "Flu": "Fever, chills, cough, body aches, sore throat, fatigue",
+        "Covid-19": "Cough, fever, shortness of breath, fatigue, loss of taste/smell",
+        "Diabetes": "Frequent urination, increased thirst, blurred vision, fatigue",
+        "Pneumonia": "Cough, fever, shortness of breath, chest pain",
+        "Allergy": "Sneezing, runny nose, itchy eyes, wheezing",
+        "Asthma": "Wheezing, shortness of breath, chest tightness, coughing",
+        "Cancer": "Fatigue, unexplained weight loss, persistent pain, cough",
+        "Heart Disease": "Chest pain, shortness of breath, fatigue, irregular heartbeat",
+        "Hypertension": "Headaches, dizziness, shortness of breath, blurred vision",
+        "Gastroenteritis": "Diarrhea, vomiting, stomach cramps, nausea",
+        "Tuberculosis": "Persistent cough, chest pain, coughing up blood, night sweats",
+        "Migraine": "Severe headache, nausea, vomiting, sensitivity to light/sound",
+        "Chronic Obstructive Pulmonary Disease": "Chronic cough, shortness of breath, wheezing, fatigue",
+        "Stroke": "Sudden numbness, confusion, trouble speaking, severe headache",
+        "Acid Reflux": "Heartburn, chest pain, regurgitation, difficulty swallowing",
+        "Kidney Disease": "Fatigue, swelling in the legs, loss of appetite, difficulty urinating",
+        "Liver Disease": "Jaundice, abdominal pain, dark urine, swelling in the legs",
+        "Anemia": "Fatigue, weakness, pale skin, shortness of breath",
+        "Sepsis": "Fever, chills, rapid heart rate, confusion, low blood pressure",
+        # Add more conditions and symptoms as necessary...
+    }
+    return condition_symptoms.get(condition, "No symptoms available.")
 
 # Gradio UI
 with gr.Blocks() as interface:
@@ -54,10 +121,7 @@ with gr.Blocks() as interface:
     chatbot = gr.Chatbot()
     input_box = gr.Textbox(show_label=False, placeholder="Describe your symptoms here...")
     send_button = gr.Button("Send")
-
-    input_box.submit(predict, [input_box, chatbot], [chatbot, input_box])
     send_button.click(predict, [input_box, chatbot], [chatbot, input_box])
 
 if __name__ == "__main__":
     interface.launch(share=True, server_name="0.0.0.0", server_port=7860, debug=True)
-h