import gradio as gr
import joblib
import ftfy
import re
import torch
import numpy as np
from transformers import DistilBertTokenizer, DistilBertModel
from huggingface_hub import hf_hub_download

def corregir_codificacion(texto):
    if isinstance(texto, str):
        return ftfy.fix_text(texto)
    return texto

def preprocesar_texto(texto):
    texto = texto.lower()
    texto = re.sub(r'\d+', '', texto)
    texto = re.sub(r'[^\w\s]', '', texto)
    return texto

class ClasificadorOpiniones:
    def __init__(self):
        try:
            model_path = hf_hub_download(repo_id="begoach1/opinion_classifier", filename="modelo_clasificador_reentrenado_lp_ros.pkl")
            self.clf_combined = joblib.load(model_path)
        except FileNotFoundError:
            raise RuntimeError("El archivo del modelo no se encuentra.")
        except Exception as e:
            raise RuntimeError(f"Error al descargar o cargar el modelo: {e}")
        
        self.tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-multilingual-cased')
        self.model = DistilBertModel.from_pretrained('distilbert-base-multilingual-cased')

    def clasificar_opinion(self, texto):
        texto = corregir_codificacion(texto)
        texto = preprocesar_texto(texto)
        tokens = self.tokenizer(texto, padding=True, truncation=True, return_tensors='pt')
        with torch.no_grad():
            outputs = self.model(**tokens)
        encoded_text = outputs.last_hidden_state[:, 0, :].numpy()
        prediccion = self.clf_combined.predict(encoded_text)
        etiquetas = ['queja', 'sugerencia', 'agradecimiento', 'felicitacion', 'ninguna', 'cambio_positivo_personal']
        resultado = dict(zip(etiquetas, prediccion[0]))
        return resultado

def clasificar(texto):
    clasificador = ClasificadorOpiniones()
    resultado = clasificador.clasificar_opinion(texto)
    return resultado

iface = gr.Interface(
    fn=clasificar,
    inputs=gr.Textbox(lines=2, placeholder="Escribe tu opinión aquí..."),
    outputs=gr.JSON(),
    title="Clasificador de Opiniones Multietiqueta",
    description="Ingresa un texto de opinión para obtener las etiquetas correspondientes."
)

iface.launch()