from transformers import pipeline
import gradio as gr


model_checkpoint = "vonewman/xlm-roberta-base-finetuned-wolof"

ner_pipeline = pipeline("ner", model=model_checkpoint,
                        aggregation_strategy="simple")  

examples = [
    "Paap  Juuf  , politiseŋ bi Saa buñ nee Paap  Juuf  , xel yépp dem ci futbal ak Olympique de Marseille ",
    "Bombardier xamle na ni dafa bëgga jàmmaarloo ak Eumeu Sen laata muy bàyyi.",
    "Pastef benn kureel la ci Senegaal bu Ousmane Sonko jiite.",
    "Màkki Sàll feeñal na ay xalaatam ci mbir yu am solo yu soxal Senegal ak Afrik."
]

def ner(text):
    output = ner_pipeline(text)
    return {"text": text, "entities": output}

demo = gr.Interface(ner,
             gr.Textbox(placeholder="Tapez votre phrase ici..."),
             gr.HighlightedText(),
             examples=examples)

demo.launch(share=True)