Update app.py

app.py

@@ -1,16 +1,72 @@
 import gradio as gr
-from transformers import pipeline
+import torch
+import torch.nn.functional as F
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import pandas as pd
 
-# Load the model from Hugging Face
-classifier = pipeline("text-classification", model="samyak152002/my-disease-classifier-sih")
+# Load the dataset and create label mappings
+df = pd.read_csv('bert_train.csv')  # Update with the correct path
+df["label"] = df["Label"]
 
-# Define the prediction function
-def predict(text):
-    result = classifier(text)
+# Create int2label and label2int mappings
+int2label = {i: disease for i, disease in enumerate(df['label'].unique())}
+label2int = {v: k for k, v in int2label.items()}
+
+# Load the model and tokenizer
+model_name = "your-huggingface-model-path"  # Replace with your Hugging Face model path
+model = AutoModelForSequenceClassification.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+# Function to classify text and return the top 3 diseases
+def classify_text(text):
+    # Set device: GPU if available, else CPU
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    
+    # Move the model to the correct device
+    model.to(device)
+    
+    # Tokenize the input and move it to the same device as the model
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True).to(device)
+    
+    # Perform inference
+    outputs = model(**inputs)
+    
+    # Get the logits (raw scores)
+    logits = outputs.logits
+    
+    # Apply softmax to convert logits into probabilities
+    probabilities = F.softmax(logits, dim=1)
+    
+    # Convert the probabilities tensor to a list for easy display
+    prob_list = probabilities[0].tolist()
+    
+    # Zip together the disease labels with their respective probabilities
+    disease_probs = {int2label[i]: prob for i, prob in enumerate(prob_list)}
+    
+    # Sort the diseases by their probabilities in descending order
+    sorted_disease_probs = dict(sorted(disease_probs.items(), key=lambda item: item[1], reverse=True))
+    
+    # Get the top 3 diseases
+    top_3_diseases = list(sorted_disease_probs.items())[:3]
+    
+    # Format the result for display
+    result = "\n".join([f"{disease}: {prob:.4f}" for disease, prob in top_3_diseases])
+    
     return result
 
-# Set up the Gradio interface
-interface = gr.Interface(fn=predict, inputs="text", outputs="json", title="Disease Classifier")
+# Gradio interface
+def predict_disease(text):
+    return classify_text(text)
+
+# Define the Gradio interface
+iface = gr.Interface(
+    fn=predict_disease, 
+    inputs="text", 
+    outputs="text", 
+    title="Disease Prediction",
+    description="Enter your symptoms, and the model will predict the top 3 most likely diseases with probabilities."
+)
 
 # Launch the app
-interface.launch()
+if __name__ == "__main__":
+    iface.launch()