import gradio as gr
import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import pandas as pd

# Load the dataset and create label mappings
df = pd.read_csv('bert_train.csv')  # Update with the correct path
df["label"] = df["Label"]

# 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

# 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
if __name__ == "__main__":
    iface.launch()