# INTERFAZ USANDO GRADIO

from transformers import pipeline
import gradio as gr
import pandas as pd

# Cargar el modelo entrenado
model_path = "AleMarroquin18/beto-finetuned-ner"
nlp_ner = pipeline("token-classification", model=model_path, tokenizer=model_path)





# Función para predecir y estructurar los resultados
def predict_entities(text):
    results = nlp_ner(text)
    if not results:  # Si no hay entidades detectadas
        return "No se detectaron entidades."

    # Crear un DataFrame para estructurar los resultados
    data = {
        "Token": [result["word"] for result in results],
        "Etiqueta": [result["entity"] for result in results],
        "Confianza (%)": [f'{result["score"]*100:.2f}' for result in results],
        "Inicio": [result["start"] for result in results],
        "Fin": [result["end"] for result in results]
    }
    df = pd.DataFrame(data)
    return df

# Interfaz con una tabla como salida
interface_beto = gr.Interface(
    fn=predict_entities,
    inputs=gr.Textbox(lines=2, placeholder="Escribe un texto para analizar entidades"),
    outputs=gr.Dataframe(headers=["Token", "entity", "Confianza (%)", "Inicio", "Fin"]),
    title="BETO NER",
    description="Prueba el modelo fine-tuneado de BETO en BioBERT para NER."
)

# Lanzar la interfaz
interface_beto.launch()