import gradio as gr
from transformers import pipeline

token_classifier = pipeline(
  model="ZurichNLP/swissbert-ner",
  aggregation_strategy="simple",
)

# def run_ner_placeholder(text: str, language: str):
#     assert language in {"de_CH", "fr_CH", "it_CH", "rm_CH"}, f"Language {language} not supported"
#     # Placeholder
#     assert text == "Mein Name sei Gantenbein."
#     return [{'entity_group': 'PER',
#       'score': 0.50026333,
#       'word': 'Gantenbein',
#       'start': 13,
#       'end': 24}
#     ]

def run_ner(text: str, language: str):
  token_classifier.model.set_default_language(language)
  output = token_classifier(text)
  return output

def visualize_ner(text: str, language: str):
  labels = run_ner(text, language)
  return {
    "entities": labels,
    "text": text,
  }

custom_css = """
    .gradio-container td:first-child {
        text-align: left !important;
    }
"""

description = """
Use the SwissBERT language model to extract named entities from text.

- Click on the examples at the bottom of the page to try them as input. They are samples from our [SwissNER dataset](https://huggingface.co/datasets/ZurichNLP/swissner).
- Our model is a fine-tuned version of SwissBERT ([Vamvas et al. 2023](https://aclanthology.org/2023.swisstext-1.6/)).
"""

theme = gr.themes.Base(
  secondary_hue=gr.themes.Color(c100="#f5f5f5", c200="#e5e5e5", c300="#d4d4d4", c400="#a3a3a3", c50="#000000", c500="#737373", c600="#1d47cb", c700="#0a2c96", c800="#262626", c900="#171717", c950="#0f0f0f"),
  radius_size="none",
)

demo = gr.Interface(
    fn=visualize_ner,
    inputs=[
        gr.Textbox(lines=5, placeholder="Enter your text here...", label="Text"),
        gr.Dropdown(choices=[
            ("German", "de_CH"),
            ("French", "fr_CH"),
            ("Italian", "it_CH"),
            ("Romansh", "rm_CH")
        ], value="de_CH", label="Language")
    ],
    outputs=gr.HighlightedText(
      label="Named Entities",
      combine_adjacent=True,
    ),
    title="Named Entity Recognition for the Swiss National Languages",
    description=description,
    allow_flagging=False,
    examples=[
        ["Es war die grosse Überraschung im Baselbiet bei den gestrigen Wahlen: Thomi Jourdan von der Kleinstpartei EVP hat den Sprung in die Regierung geschafft und sich gegen SVP-Nationalrätin Sandra Sollberger durchgesetzt.", "de_CH"],
        ["La cheffe de la Direction suisse du développement et de la coopération (DDC), Patricia Danzi, a tenu à accueillir elle-même les 87 spécialistes à l'aéroport, a indiqué le Département fédéral des affaires étrangères (DFAE). Elle était notamment accompagnée de Silvio Flückiger, chef suppléant de l'Aide humanitaire de la Suisse, ainsi que du chef de l'armée Thomas Süssli.", "fr_CH"],
        ["I Grigioni, al contrario di quanto fatto dal Ticino, non intendono adottare una legislazione particolare per fronteggiare la diffusione delle sigarette elettroniche tra i ragazzi. Lo ha ribadito martedì il Consiglio di Stato rispondendo alle preoccupazioni espresse dall'ex direttrice scolastica e deputata del Centro di Mesocco Piera Furger. Di fronte al Gran Consiglio, il Governo ha spiegato di privilegiare una soluzione a livello nazionale e un approccio all'insegna della prevenzione e non della repressione.", "it_CH"],
        ["Tar l'Ospital chantunal dal Grischun datti in augment d'accidents da skiunzas e snowboardists sin pista da 20%. Quai dal prim da schaner fin ils 6 da favrer, cumpareglià cun la media dals tschintg onns passads. Quai rapporta il «Tagesanzeiger» Tenor il pledader da medias, Dajan Roman, haja era dà dapli blessuras grevas. Ils ultims onns era la media tar 3%, dapi l'entschatta da l'onn èn quai gia 4,2%.", "rm_CH"],
    ],
    css=custom_css,
    theme=theme,
)

demo.launch(debug=True)