"""Demo Space for ZeroGPU/zlm-v1-iab-domain-classifier.

Type a URL or bare domain; the model classifies the destination into IAB Content and
Audience categories from the URL string alone (no page fetch).
"""

import importlib.util
import json

import gradio as gr
import numpy as np
import onnxruntime as ort
from huggingface_hub import hf_hub_download
from tokenizers import Tokenizer

REPO = "ZeroGPU/zlm-v1-iab-domain-classifier"
TOP_K = 6

model_path = hf_hub_download(REPO, "onnx/model_quantized.onnx")
tokenizer = Tokenizer.from_file(hf_hub_download(REPO, "tokenizer.json"))
meta = json.load(open(hf_hub_download(REPO, "multi_head_classifier_metadata.json")))
thresholds = json.load(open(hf_hub_download(REPO, "thresholds.json")))

spec = importlib.util.spec_from_file_location("url_text", hf_hub_download(REPO, "url_text.py"))
url_text = importlib.util.module_from_spec(spec)
spec.loader.exec_module(url_text)

session = ort.InferenceSession(model_path)
output_names = [o.name for o in session.get_outputs()]


def _head_output(head: str, outputs: dict) -> np.ndarray | None:
    """Match an output tensor to a head the same way the production service does."""
    if head in outputs:
        return outputs[head]
    if head + "_logits" in outputs:
        return outputs[head + "_logits"]
    needle = head.lower().replace("_names", "")
    for name, tensor in outputs.items():
        if needle in name.lower():
            return tensor
    return None


def classify(url: str):
    url = (url or "").strip()
    if not url:
        return "", {}, {}
    text = url_text.url_to_text(url)
    enc = tokenizer.encode(text)
    raw = session.run(None, {
        "input_ids": np.array([enc.ids], dtype=np.int64),
        "attention_mask": np.array([enc.attention_mask], dtype=np.int64),
    })
    outputs = dict(zip(output_names, raw))

    results = []
    for head in ("iabcontent_names", "iabaudience_names"):
        logits = _head_output(head, outputs)
        if logits is None:
            results.append({})
            continue
        probs = 1.0 / (1.0 + np.exp(-np.asarray(logits).reshape(-1)))
        cut = thresholds.get(head, {}).get("_default", 0.5)
        labels = meta["label_mappings"]["category_labels"][head]
        picked = [(labels[i], float(p)) for i, p in enumerate(probs) if p >= cut]
        picked.sort(key=lambda t: -t[1])
        results.append(dict(picked[:TOP_K]) or {"(no label above threshold %.2f)" % cut: 0.0})
    content, audience = results
    return text, content, audience


demo = gr.Interface(
    fn=classify,
    inputs=gr.Textbox(label="URL or domain", placeholder="espn.com"),
    outputs=[
        gr.Textbox(label="Model input (url_to_text rendering)"),
        gr.Label(label="IAB Content categories", num_top_classes=TOP_K),
        gr.Label(label="IAB Audience categories", num_top_classes=TOP_K),
    ],
    title="IAB Domain Classifier",
    description=(
        "URL-only IAB classification with "
        "[ZeroGPU/zlm-v1-iab-domain-classifier](https://huggingface.co/ZeroGPU/zlm-v1-iab-domain-classifier) "
        "— a fine-tuned 149M ModernBERT that beats GPT-5.4-nano on this task "
        "(content micro-F1 0.385 vs 0.353 on a 784-URL gold benchmark) at ~36 ms per URL. "
        "No page is fetched: every signal comes from the URL string itself."
    ),
    examples=[["espn.com"], ["https://www.sportsnews.co.uk/foo"], ["techcrunch.com"], ["allrecipes.com"]],
    cache_examples=False,
)

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