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 ",
    "Xamal ni bëgg naa la ba fàww te dootuma la fàtte ",
    "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)