Update app.py

app.py

@@ -7,6 +7,48 @@ model_name = "Sarthak279/Disease-symptom-prediction"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSequenceClassification.from_pretrained(model_name)
 
+
+
+# Access label mapping from config if available
+id2label = model.config.id2label  # This is a dict like {0: 'Flu', 1: 'Typhoid', ...}
+
+# Set device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+model.eval()
+
+# Prediction function
+def predict_disease(patient_note):
+    inputs = tokenizer(
+        patient_note,
+        return_tensors="pt",
+        truncation=True,
+        padding=True,
+        max_length=512
+    ).to(device)
+
+    with torch.no_grad():
+        outputs = model(**inputs)
+        logits = outputs.logits
+        predicted_class = torch.argmax(logits, dim=1).item()
+    
+    # Fetch label from config, or fallback
+    predicted_label = id2label.get(predicted_class, "Unknown")
+    return predicted_label
+
+# Gradio Interface
+demo = gr.Interface(
+    fn=predict_disease,
+    inputs=gr.Textbox(lines=4, label="📝 Enter Symptoms or Clinical Notes"),
+    outputs=gr.Textbox(label="Predicted Disease"),
+    title="🩺 Disease Prediction Model",
+    description="Predicts likely disease based on symptoms using a fine-tuned model from Hugging Face Hub."
+)
+
+if __name__ == "__main__":
+    demo.launch()
+
+
 # Set model to eval mode and move to device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)