import gradio as gr
from transformers import pipeline

# Carrega el pipeline NER per català
ner_pipeline = pipeline(model="projecte-aina/roberta-base-ca-v2-cased-ner")

# Funció per aplicar el NER
def ner(text):
    output = ner_pipeline(text)
    return {"text": text, "entities": output}

# Funció per netejar el formulari
def neteja():
    return [gr.Textbox(value=None), gr.HighlightedText(value=None)]

# Llista d'exemples
exemples = [
    "El president Pere Aragonès va visitar Barcelona el 23 d'abril.",
    "La Universitat de Girona ofereix estudis en enginyeria informàtica.",
    "Apple va llançar l’iPhone 15 a Califòrnia."
]

# UI amb Gradio
with gr.Blocks(theme=gr.themes.Glass()) as demo:
    gr.Markdown(
    """
    # Reconeixement 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")

    gr.Examples(examples=exemples, inputs=inp, label="Exemples:")

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

demo.launch(auth=("admin", "admin"), share=True, server_name="127.0.0.1")