import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# Load model from the current directory
MODEL_PATH = "." 

LABELS = [
    "Endocrinology Referral", "Nutrition Referral", "Cardiology Referral", "Bariatric Referral", 
    "Mental Health Screen", "Food Insecurity Discussion", "GLP-1 Prescription", "Follow-up Scheduled"
]

# Load Model
print("Loading model...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH)

def analyze_note(text):
    if not text.strip(): return None
    
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=320)
    
    with torch.no_grad():
        logits = model(**inputs).logits
    
    probs = torch.sigmoid(logits).cpu().numpy()[0]
    return {label: float(conf) for label, conf in zip(LABELS, probs)}

# Examples for users to click
examples = [
    ["VISIT DATE: 11/20/2025\nSUBJECTIVE: 16yo female presents for weight management. Reports trying to walk more.\nPLAN:\n1. Start Zepbound 2.5mg weekly.\n2. Referral to Pediatric Endocrinology.\n3. Consulting Registered Dietitian."],
    ["HPI: Mom is requesting Wegovy today. PLAN: Discussed Wegovy but insurance denied the Prior Authorization. No medication prescribed. Offered referral to nutrition services but family declined."],
    ["CC: Ear pain. Social History: Dad mentions he lost his job last week and they are currently using a food pantry. Plan: Amoxicillin."]
]

# The Interface
demo = gr.Interface(
    fn=analyze_note,
    inputs=gr.Textbox(lines=10, label="Paste Clinical Note Here"),
    outputs=gr.Label(num_top_classes=8, label="Predicted Actions"),
    title="🏥 Clinical Note Analyzer (GatorTron)",
    description="This model identifies referrals, prescriptions, and screenings in unstructured clinical text.",
    examples=examples,
    theme=gr.themes.Soft()
)

demo.launch()