import gradio as gr
from setfit import SetFitModel
import spacy

# Cargar el modelo
model = SetFitModel.from_pretrained("desarrolloasesoreslocales/SetFitArgumentos")

# Cargar el modelo de spaCy para español
spacy.cli.download("es_core_news_sm")
nlp = spacy.load("es_core_news_sm")

# Mapear las etiquetas
id2label = {0: "49", 1: "994", 2: "1002", 3: "2014"}

# Definir la función para dividir en párrafos
def dividir_en_parrafos(texto):
    doc = nlp(texto)
    parrafos = []
    parrafo_actual = []
    for oracion in doc.sents:
        parrafo_actual.append(oracion.text)
        if oracion.text.endswith((".", "!", "?")):
            parrafos.append(" ".join(parrafo_actual))
            parrafo_actual = []
    if parrafo_actual:
        parrafos.append(" ".join(parrafo_actual))
    return parrafos

# Definir la función de predicción
def predict(payload):
    # Dividir el texto en párrafos
    parrafos = dividir_en_parrafos(payload)
    
    # Inicializar un conjunto para almacenar etiquetas únicas
    etiquetas_encontradas = set()
    
    # Procesar cada párrafo por separado
    for parrafo in parrafos:
        # Predecir con el modelo
        results = model.predict_proba([parrafo])[0]
        
        # Filtrar etiquetas con score > 0.60
        etiquetas_filtradas = [id2label[i] for i, score in enumerate(results) if score > 0.60]
        
        # Agregar etiquetas al conjunto
        etiquetas_encontradas.update(etiquetas_filtradas)
    
    return "\n".join(etiquetas_encontradas)

# Crear una interfaz Gradio
iface = gr.Interface(
    fn=predict, 
    inputs=gr.Textbox(), 
    outputs=gr.Textbox(), 
    live=False,
    title="Clasificador de Texto"
)

# Iniciar la interfaz Gradio
iface.launch()