from transformers import AutoModelForSequenceClassification, AutoTokenizer, TextClassificationPipeline
import gradio as gr

model_id = "hedtorresca/Multilingual-MiniLM-L12-H384-fine-tunning2"
                                                          
model = AutoModelForSequenceClassification.from_pretrained(model_id)
tokenizer = AutoTokenizer.from_pretrained(model_id, return_dict=True )
pipe = TextClassificationPipeline(model=model, tokenizer=tokenizer, add_special_tokens=True,  max_length= 512, truncation= True)
# Definir la función de inferencia
def classify_message(ues_detallada, medio_de_comunicacion, asunto_o_copy):
    combined = f"{ues_detallada} por {medio_de_comunicacion}. {asunto_o_copy}"
    #prediction = pipe(combined)
   # return pipe(combined)
    return "Comercial" if prediction[0]['label'] == 'LABEL_0' else "Informativo"

# Crear la interfaz de Gradio
iface = gr.Interface(
    fn=classify_message,
    inputs=[
        gr.Textbox(lines=1, placeholder="Escribe la UES aquí.", label="UES"),
        gr.Radio(["EMAIL", "SMS"], label="Canal de Comunicación"),
        gr.Textbox(lines=2, placeholder="Escribe el asunto o copy del mensaje aquí.", label="Asunto o Copy")
    ],
    outputs= gr.Label(label="Tipo de Campaña"),
    title="Clasificador de Campañas",
    description="Describe la unidad de negocio, el canal de comunicación y escribe el copy de la campaña para clasificar si es comercial o informativo."
)

# Lanzar la interfaz
iface.launch()