import gradio as gr
from setfit import SetFitModel

def cortar_en_bloques(texto, longitud_bloque):
    palabras = texto.split()
    bloques = []
    bloque_actual = []

    for palabra in palabras:
        bloque_actual.append(palabra)

        if len(bloque_actual) == longitud_bloque:
            bloques.append(" ".join(bloque_actual))
            bloque_actual = []

    # Si queda un bloque parcial al final, agregarlo
    if bloque_actual:
        bloques.append(" ".join(bloque_actual))

    return bloques

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

# Mapear las etiquetas

# Definir la función de predicción
def predict(payload):
    
    recorte_general = ""

    for chunk in cortar_en_bloques(s, 90):
      if model.predict([chunk]).item() == 1:
        recorte_general += chunk + " "

    recorte_final = ""

    for chunk in cortar_en_bloques(recorte_general, 84):
      if model.predict([chunk]).item() == 1:
        recorte_final += chunk + " "
          
    # results = model.predict_proba([payload])[0]
    # return "\n".join([f"Label: {id2label[i]},      Score: {round(result.item()*100, 2)} %" for i, result in enumerate(results)])
    
    return recorte_final

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

# Iniciar la interfaz Gradio
iface.launch()