"""Gradio ZeroGPU Space — Portuguese NER API for pt-ner."""

from __future__ import annotations

import spaces
import gradio as gr
import torch
from transformers import pipeline

MODEL_ID = "marquesafonso/bertimbau-large-ner-total"

ner = pipeline(
    "ner",
    model=MODEL_ID,
    aggregation_strategy="simple",
    device=0,
    torch_dtype=torch.bfloat16,
)
ner.model.to("cuda")


def _serialize(raw: list[dict]) -> list[dict]:
    entities: list[dict] = []
    for item in raw:
        start = int(item["start"])
        end = int(item["end"])
        entities.append(
            {
                "label": str(item.get("entity_group") or item.get("entity") or "MISC"),
                "text": str(item.get("word") or ""),
                "start": start,
                "end": end,
                "score": float(item["score"]) if item.get("score") is not None else None,
            }
        )
    return entities


@spaces.GPU(duration=120)
def predict(text: str) -> list[dict]:
    """Extract named entities from Portuguese text with character offsets."""
    cleaned = (text or "").strip()
    if not cleaned:
        return []
    return _serialize(ner(cleaned))


EXAMPLES = [
    ["João Silva mora em Lisboa."],
    ["Maria Santos trabalha no Banco de Portugal, no Porto."],
]

demo = gr.Interface(
    fn=predict,
    inputs=gr.Textbox(label="Texto", lines=6, placeholder="Cole texto em português…"),
    outputs=gr.JSON(label="Entidades"),
    title="BERTimbau NER — pt-ner",
    description=(
        "API Space for the pt-ner application. "
        f"Model: [{MODEL_ID}](https://huggingface.co/{MODEL_ID})."
    ),
    examples=EXAMPLES,
    cache_examples=True,
    cache_mode="lazy",
)

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