"""Gradio demo: bag-of-words sentiment via ONNX Runtime."""

import json
from pathlib import Path

import gradio as gr
import numpy as np
import onnxruntime as ort

from preprocess import text_to_bow

VOCAB_PATH = Path(__file__).resolve().parent / "vocab.json"
vocab_data = json.loads(VOCAB_PATH.read_text())
VOCAB: list[str] = vocab_data["vocab"]
LABELS: list[str] = vocab_data["labels"]

session = ort.InferenceSession("sentiment.onnx")


def predict(text: str) -> dict[str, float]:
    if not text.strip():
        return {label: 0.0 for label in LABELS}

    bow = text_to_bow(text, VOCAB)[np.newaxis, :]
    logits = session.run(None, {"bow": bow})[0][0]
    exp_logits = np.exp(logits - logits.max())
    probs = exp_logits / exp_logits.sum()

    return {label: float(prob) for label, prob in zip(LABELS, probs)}


demo = gr.Interface(
    fn=predict,
    inputs=gr.Textbox(label="Text", placeholder="I love this product"),
    outputs=gr.Label(num_top_classes=2, label="Sentiment"),
    title="ONNX Sentiment Demo",
    description=(
        "Tiny bag-of-words classifier trained from scratch in PyTorch "
        "(no pretrained model). Exported to ONNX for inference."
    ),
    examples=[
        ["I love this"],
        ["This is awful"],
        ["Amazing product"],
        ["So disappointed"],
    ],
)

demo.launch()