import gradio as gr
from transformers import pipeline
model_name = 'mrm8488/bert-spanish-cased-finetuned-ner'
# model_name = 'MMG/xlm-roberta-large-ner-spanish'
# pipe = pipeline("image-classification")
pipe = pipeline("ner", model=model_name, aggregation_strategy="simple")

def infer_ner(text):
    output = pipe(text)
    for d in output:
        print(d)
        d['entity'] = d['entity_group']

    return{
        'text': text,
        'entities': output
    },  output

gr.Interface(
    fn=infer_ner,
    inputs=gr.Textbox(),
    outputs=[gr.HighlightedText(), gr.Textbox()],
    examples=[
        "Mauricio Macri, Cristina Fernández y Alberto Fernández se juntaron en la Casa Rosada",
        "Lionel Messi marcó un gol contra Arabia Saudita",
        "Vamos Boca Juniors Campeón del mundo",
        "Lo importante no es que vengas sino que vuelvas, Unicenter!",
    ]
).launch()