ancerlop/SetFitArgumentos

1import gradio as gr2from setfit import SetFitModel3import spacy4 5# Cargar el modelo6model = SetFitModel.from_pretrained("desarrolloasesoreslocales/SetFitArgumentos")7 8# Cargar el modelo de spaCy para español9spacy.cli.download("es_core_news_sm")10nlp = spacy.load("es_core_news_sm")11 12# Mapear las etiquetas13id2label = {0: "49", 1: "994", 2: "1002", 3: "2014"}14 15# Definir la función para dividir en párrafos16def dividir_en_parrafos(texto):17    doc = nlp(texto)18    parrafos = []19    parrafo_actual = []20    for oracion in doc.sents:21        parrafo_actual.append(oracion.text)22        if oracion.text.endswith((".", "!", "?")):23            parrafos.append(" ".join(parrafo_actual))24            parrafo_actual = []25    if parrafo_actual:26        parrafos.append(" ".join(parrafo_actual))27    return parrafos28 29# Definir la función de predicción30def predict(payload):31    # Dividir el texto en párrafos32    parrafos = dividir_en_parrafos(payload)33    34    # Inicializar un conjunto para almacenar etiquetas únicas35    etiquetas_encontradas = set()36    37    # Procesar cada párrafo por separado38    for parrafo in parrafos:39        # Predecir con el modelo40        results = model.predict_proba([parrafo])[0]41        42        # Filtrar etiquetas con score > 0.6043        etiquetas_filtradas = [id2label[i] for i, score in enumerate(results) if score > 0.60]44        45        # Agregar etiquetas al conjunto46        etiquetas_encontradas.update(etiquetas_filtradas)47    48    return "\n".join(etiquetas_encontradas)49 50# Crear una interfaz Gradio51iface = gr.Interface(52    fn=predict, 53    inputs=gr.Textbox(), 54    outputs=gr.Textbox(), 55    live=False,56    title="Clasificador de Texto"57)58 59# Iniciar la interfaz Gradio60iface.launch()61