import gradio as gr
from setfit import SetFitModel

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

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

# Definir la función de predicción
def predict(payload):
    results = model.predict_proba([payload])[0]
    
    return "\n".join([f"Label: {id2label[i]}, Score: {result.item()}" for i, result in enumerate(results)])

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

# Iniciar la interfaz Gradio
iface.launch()