# INTERFAZ USANDO GRADIO

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

# Cargar el modelo entrenado
model_path = "AleMarroquin18/bert-base-uncased-biobert"
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 = gr.Interface(
    fn=predict_entities,
    inputs=gr.Textbox(lines=2, placeholder="Escribe un texto para analizar entidades..."),
    outputs="dataframe",
    title="Reconocimiento de Entidades Nombradas (NER)",
    description=(
        "Escribe un texto para que el modelo identifique entidades nombradas, "
        "como nombres de personas, lugares o fechas. "
        "Los resultados se muestran en una tabla con la confianza del modelo."
    )
)

# Lanzar la interfaz
interface.launch()