Update app.py

app.py

@@ -1,39 +1,18 @@
 from transformers import AutoTokenizer, TFAutoModelForSequenceClassification
 import tensorflow as tf
 import gradio as gr
-import os
-import re
-from nltk.corpus import stopwords
-from nltk.stem import PorterStemmer
-
-# Ensure you have the necessary nltk resources
-import nltk
-nltk.download('stopwords')
-
-# Caching the model locally to avoid re-downloading
-MODEL_NAME = "Zabihin/Symptom_to_Diagnosis"
-CACHE_DIR = "./cached_model"
-if not os.path.exists(CACHE_DIR):
-    os.makedirs(CACHE_DIR)
 
 # Load the tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, cache_dir=CACHE_DIR)
-model = TFAutoModelForSequenceClassification.from_pretrained(MODEL_NAME, cache_dir=CACHE_DIR)
-
-# Initialize stopwords and stemmer
-stop_words = set(stopwords.words('english'))
-stemmer = PorterStemmer()
+model_name = "Zabihin/Symptom_to_Diagnosis"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = TFAutoModelForSequenceClassification.from_pretrained(model_name)
 
-# Clean the input text with advanced preprocessing
+# Clean the input text
 def clean_input(symptom_text):
-    # Remove non-ASCII characters and convert to lowercase
+    # 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().strip()  # Remove leading/trailing spaces
-    
-    # Remove stopwords and apply stemming
-    words = symptom_text.split()
-    filtered_words = [stemmer.stem(word) for word in words if word not in stop_words]
-    return ' '.join(filtered_words)
+    symptom_text = symptom_text.lower()  # Optional: Convert to lowercase
+    return symptom_text
 
 # Define the predict function
 def predict(symptom_text, chat_history=[]):
@@ -49,37 +28,46 @@ def predict(symptom_text, chat_history=[]):
         logits = outputs.logits
         prediction = tf.argmax(logits, axis=-1).numpy()[0]
 
-        # Map the prediction to a label and description
+        # Map the prediction to a label
         labels = {
-            0: ("Allergy", "A condition in which the immune system reacts abnormally to a foreign substance."),
-            1: ("Arthritis", "A disease causing painful inflammation and stiffness of the joints."),
-            2: ("Bronchial Asthma", "A condition in which your airways narrow and swell and produce extra mucus."),
-            3: ("Cervical Spondylosis", "Age-related wear and tear affecting the bones, discs, and joints in the neck."),
-            4: ("Chicken Pox", "A viral infection causing an itchy, blister-like rash on the skin."),
-            5: ("Common Cold", "A viral infection affecting the upper respiratory tract, causing a sore throat and runny nose."),
-            6: ("Dengue", "A viral illness transmitted by mosquitoes, leading to flu-like symptoms."),
-            7: ("Diabetes", "A group of diseases that affect how your body uses blood sugar (glucose)."),
-            8: ("Drug Reaction", "A harmful reaction in the body due to a medication or drug."),
-            9: ("Fungal Infection", "An infection caused by fungi, often affecting the skin or nails."),
-            10: ("Gastroesophageal Reflux Disease", "A digestive disorder where stomach acid irritates the food pipe."),
-            11: ("Hypertension", "High blood pressure, a condition that can lead to serious health issues if untreated."),
-            12: ("Impetigo", "A contagious bacterial skin infection that causes red sores or blisters."),
-            13: ("Jaundice", "A yellowish tint to the skin or eyes caused by excess bilirubin in the blood."),
-            14: ("Malaria", "A life-threatening disease transmitted by mosquitoes, caused by a parasite."),
-            15: ("Migraine", "A severe headache that can cause intense throbbing or a pulsing sensation."),
-            16: ("Peptic Ulcer Disease", "Sores that develop on the lining of the stomach or the upper part of the small intestine."),
-            17: ("Pneumonia", "An infection that inflames the air sacs in one or both lungs, causing cough and difficulty breathing."),
-            18: ("Psoriasis", "An autoimmune condition that causes skin cells to multiply too quickly, leading to patches of red, scaly skin."),
-            19: ("Typhoid", "A bacterial infection caused by Salmonella typhi, leading to fever, fatigue, and abdominal pain."),
-            20: ("Urinary Tract Infection", "An infection in any part of the urinary system, including the kidneys, bladder, or urethra."),
-            21: ("Varicose Veins", "Swollen and enlarged veins, often appearing blue or dark purple, usually in the legs."),
+            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"
+        }
+
+        descriptions = {
+            "Allergy": "An immune system reaction to foreign substances.",
+            "Arthritis": "Inflammation of one or more joints.",
+            "Bronchial Asthma": "A condition where the airways become inflamed and narrow.",
+            "Cervical Spondylosis": "Age-related changes in the bones, discs, and joints of the neck.",
+            "Chicken Pox": "A highly contagious viral infection causing an itchy skin rash.",
+            "Common Cold": "A viral infection of the upper respiratory tract, causing sneezing, runny nose, and sore throat.",
+            "Dengue": "A viral disease transmitted by mosquitoes, causing fever and severe pain.",
+            "Diabetes": "A disease that affects how your body processes blood sugar.",
+            "Drug Reaction": "An adverse response to a medication.",
+            "Fungal Infection": "An infection caused by fungi affecting the skin or organs.",
+            "Gastroesophageal Reflux Disease": "A chronic digestive condition where stomach acid irritates the food pipe.",
+            "Hypertension": "High blood pressure that can lead to heart disease.",
+            "Impetigo": "A contagious bacterial skin infection.",
+            "Jaundice": "A yellowing of the skin or eyes due to liver disease.",
+            "Malaria": "A serious disease transmitted by mosquito bites, causing fever and chills.",
+            "Migraine": "Severe headaches often accompanied by nausea and sensitivity to light.",
+            "Peptic Ulcer Disease": "Sores in the stomach lining or the upper part of the small intestine.",
+            "Pneumonia": "An infection that inflames the air sacs in one or both lungs.",
+            "Psoriasis": "A chronic autoimmune disease causing the rapid growth of skin cells.",
+            "Typhoid": "A bacterial infection causing high fever, abdominal pain, and weakness.",
+            "Urinary Tract Infection": "An infection in any part of the urinary system.",
+            "Varicose Veins": "Swollen, twisted veins caused by faulty valves in the veins."
         }
 
-        diagnosis, description = labels.get(prediction, ("Unknown diagnosis", "No description available for this condition"))
+        diagnosis = labels.get(prediction, "Unknown diagnosis")
+        description = descriptions.get(diagnosis, "No description available.")
 
         # Add conversation history
         chat_history.append(("User", symptom_text))
-        chat_history.append(("AI", f"Predicted Diagnosis: **{diagnosis}**. {description} Please consult a doctor for more accurate results."))
+        chat_history.append(("AI", f"Predicted Diagnosis: <b>{diagnosis}</b>. {description} Please consult a doctor for more accurate results."))
 
     except Exception as e:
         chat_history.append(("AI", f"Error: {str(e)}"))
@@ -88,7 +76,7 @@ def predict(symptom_text, chat_history=[]):
 
 # Gradio UI
 with gr.Blocks() as interface:
-    gr.Markdown("<h1 style='text-align: center; font-size: 50px; margin-top: 40px; margin-bottom: 40px;'>Medi Mind - Your AI Health Assistant</h1>")
+    gr.Markdown("<h1 style='text-align: center; font-size: 40px; margin-top: 30px; margin-bottom: 30px;'>Medi Mind - Your AI Health Assistant</h1>")
     chatbot = gr.Chatbot()
     input_box = gr.Textbox(show_label=False, placeholder="Describe your symptoms here...")
     send_button = gr.Button("Send")