"""Gradio app that showcases Scandinavian NER models."""

import gradio as gr
from transformers import pipeline


def main():
    pipe = pipeline(
        task="token-classification",
        model="saattrupdan/nbailab-base-ner-scandi",
        aggregation_strategy="first",
    )

    examples = [
        "Hans er en professor på Københavns Universitetet i København, og han er en rigtig københavner. Hans kat, altså Hans' kat, Lisa, er supersød. Han fik købt en Mona Lisa på tilbud i Netto og gav den til sin kat, og nu er Mona Lisa'en Lisa's kæreste eje. Hans bror Peter og Hans besluttede, at Peterskirken skulle have fint besøg. Men nu har de begge Corona." ,
        "Borghild handler på Bunnpris, der de har et spesialtilbud på en Mona Lisa-kopi.",
        "Johan är från Djurgården och fick Corona förra veckan.",
        "Ulf Hjalmar Ed Kristersson (føddur 29. desember 1963 í Lund) er siviløkonomur og politikari.",
        "Space Exploration Technologies Corporation, betur þekkt sem SpaceX, er bandarískur flugtækniframleiðandi og geimferðafyrirtæki með höfuðstöðvar í Hawthorne í Kaliforníu. Það var stofnað af athafnamanninum Elon Musk árið 2002."
    ]

    def ner(text):
        output: list[dict] = pipe(text)
        output = [{'entity': dct['entity_group']} | dct for dct in output]
        print(output)
        return {"text": text, "entities": output}

    demo = gr.Interface(
        ner,
        gr.Textbox(placeholder="Enter sentence here..."),
        gr.HighlightedText(),
        examples=examples,
        title="Scandinavian Named Entity Recognition",
        description="""
Locate and classify named entities in Danish, Swedish, Norwegian, Icelandic and Faroese text. Write any text in the box below, and the model will predict and highlight the following entities:

| **Tag** | **Name** | **Description** |
| :------ | :------- | :-------------- |
| `PER` | Person | The name of a person (e.g., *Birgitte* and *Mohammed*) |
| `LOC` | Location | The name of a location (e.g., *Tyskland* and *Djurgården*) |
| `ORG` | Organisation | The name of an organisation (e.g., *Bunnpris* and *Landsbankinn*) |
| `MISC` | Miscellaneous | A named entity of a different kind (e.g., *Ūjķnustu pund* and *Mona Lisa*) |

_Also, be patient, as this demo is running on a CPU!_""",
    )

    demo.launch()

if __name__ == "__main__":
    main()