from transformers import pipeline
import gradio as gr

ner_pipeline = pipeline(model="projecte-aina/roberta-base-ca-v2-cased-ner")

exemples = ["El Joan no ha anat mai a Manresa"]

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

def neteja():
    return [gr.Textbox(value=None), gr.HighlightedText(value=None)]

with gr.Blocks(theme=gr.themes.Glass()) as demo:
    gr.Markdown(
    """
    # Reconeixament d'entitats nomenades en català
    Escriu o copia un text i troba les seves entitats
    """)
    with gr.Row():
      with gr.Column():
        inp = gr.Textbox(label="text", placeholder="Escriu aquí...")
        with gr.Row():
          b1 = gr.Button(value="Neteja")
          b2 = gr.Button("Troba entitats", variant="primary")
      out = gr.HighlightedText(label="sortida")

    examples = gr.Examples(examples=exemples, inputs=inp, label="Exemple:")
    logout_button = gr.Button("Logout", link="/logout")

    b1.click(neteja, outputs=[inp, out])
    b2.click(ner, inputs=inp, outputs=out)

demo.launch(show_api=False)