# 

import torch
import gradio as gr
import pandas as pd
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from scipy.special import softmax

# Load model and tokenizer
model_path = "trained_clinicalbert"
model = AutoModelForSequenceClassification.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path)

model.eval()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# Define class labels manually (ensure it matches the trained model)
classes = ["Asthma", "COPD", "Lung Cancer", "Other Pulmonary", "Pleural Effusion", "Pneumonia", "Tuberculosis"]

# Prediction function
def predict_clinical_note(note):
    inputs = tokenizer(note, return_tensors="pt", truncation=True, padding="max_length", max_length=512)
    inputs = {key: val.to(device) for key, val in inputs.items()}
    with torch.no_grad():
        outputs = model(**inputs)
    probs = softmax(outputs.logits.cpu().numpy(), axis=1)
    pred_idx = probs.argmax(axis=1)[0]
    pred_class = classes[pred_idx]
    confidence = float(probs[0][pred_idx])
    return f"{pred_class} (Confidence: {confidence:.2f})"

# Gradio interface
iface = gr.Interface(
    fn=predict_clinical_note,
    inputs=gr.Textbox(lines=6, placeholder="Paste clinical note here..."),
    outputs="text",
    title="Pulmonary Disease Classifier",
    description="Enter a clinical note to predict pulmonary condition (e.g., COPD, Pneumonia, Tuberculosis...)"
)

if __name__ == "__main__":
    iface.launch()