app.py

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 = ""

    # Crear chunks
    chunks = cortar_en_bloques(ocr_text, 150)
    first = -1
    margin = int(len(chunks) * 0.25)
    chunks_removable = chunks[:margin] + chunks[-margin:]
    
    for i in range(len(chunks)):
        print('Recortando -', round((i/len(chunks))*100), '%')
        if chunks[i] not in chunks_removable or trim_model.predict([chunks[i]]).item() == 1:
          if first == -1:
            first = i
          recorte_general += chunks[i] + " "
    
    if first > 0:
    recorte_general = chunks[first-1] + recorte_general
    print(100, '%')

    recorte_final = ""

    # Definir tamñano de fragmentos de texto
    # text_splitter2 = RecursiveCharacterTextSplitter(chunk_size=100, chunk_overlap=0, length_function=len)
    # Crear chunks
    chunks2 = cortar_en_bloques(recorte_general, 80)
    margin_s = int(len(chunks2) * 0.1)
    margin_e = int(len(chunks2) * 0.1)
    # if margin_s > 1:
    chunks_removable2 = chunks2[:margin_s] + chunks2[-margin_e:]
    # chunks_removable2 = chunks2[-margin_e:]
    
    for i in range(len(chunks2)):
        print('Recortando -', round((i/len(chunks2))*100), '%')
        if chunks2[i] not in chunks_removable2 or trim_model.predict([chunks2[i]]).item() == 1:
          recorte_final += chunks2[i] + " "
    print(100, '%')
    
    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()