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}

demo = gr.Interface(
  ner,
  gr.Textbox(label="Text", placeholder="Escriu aquí..."),
  gr.HighlightedText(label="sortida"),
  allow_flagging="never",
  title="Reconeixament d'entitats nomenades en català",
  description="Escriu o copia un text i troba les seves entitats",
  theme=gr.themes.Glass(),
  submit_btn=gr.Button("Troba entitats", variant="primary"),
  clear_btn=gr.Button("Neteja"),
  examples=exemples
 )

demo.launch(show_api=False)