import gradio as gr
from optimum.onnxruntime import ORTModelForSequenceClassification
from transformers import AutoTokenizer
import numpy as np

model_id = "snortdapot/depression-distilbert-onnx"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = ORTModelForSequenceClassification.from_pretrained(
    model_id,
    file_name="model_quantized.onnx"
)

ID_TO_LABEL = {0: "Normal", 1: "Mild", 2: "Severe"}

def classify(text):
    if not text or not text.strip():
        return {"Normal": 0.0, "Mild": 0.0, "Severe": 0.0}
    inputs = tokenizer(text[:512], return_tensors="np", truncation=True, padding=True)
    outputs = model(**{k: v for k, v in inputs.items()})
    logits = outputs.logits[0]
    # softmax
    exp_logits = np.exp(logits - np.max(logits))
    probs = exp_logits / exp_logits.sum()
    result = {}
    for i, prob in enumerate(probs):
        result[ID_TO_LABEL[i]] = float(round(prob, 4))
    return result

demo = gr.Interface(
    fn=classify,
    inputs=gr.Textbox(label="Text", placeholder="Enter a comment..."),
    outputs=gr.Label(label="Depression Risk"),
    title="Depression Risk Classifier",
    description="DistilBERT fine-tuned for depression risk classification (Normal/Mild/Severe)",
)

demo.launch()