nlp_category.py

import logging
import unicodedata
from setfit import SetFitModel

logger = logging.getLogger(__name__)

# ======================
# 1. Cargar modelo SetFit
# ======================

MODEL_ID = "Alecit1234/modelo_finanzas_peru_v2"

logger.info(f"Cargando modelo SETFIT desde HF: {MODEL_ID}")

model_setfit = SetFitModel.from_pretrained(MODEL_ID)
logger.info("Modelo SetFit cargado correctamente.")

# ======================
# 2. Mapeo de categorías
# ======================

CATEGORY_LABEL_MAP = {
    0:"comida", 1:"supermercado", 2:"transporte", 3:"taxi", 4:"entretenimiento",
    5:"educacion", 6:"tecnologia", 7:"servicios", 8:"fitness", 9:"imprevistos",
    10:"delivery", 11:"mascotas", 12:"familia", 13:"salud",
    14:"ingreso_beca", 15:"ingreso_trabajo", 16:"ingreso_familia",
    17:"ingreso_venta", 18:"ingreso_freelance", 19:"ingreso_extra"
}

# ======================
# 3. Normalización del texto
# ======================

def _normalizar(texto: str) -> str:
    texto = texto.lower()
    texto = unicodedata.normalize("NFD", texto)
    return "".join(c for c in texto if unicodedata.category(c) != "Mn")

# ======================
# 4. Predicción de categoría
# ======================

def predecir_categoria(texto: str) -> str:
    if not texto or texto.strip() == "":
        return "otros"

    texto_norm = _normalizar(texto)

    pred_id = model_setfit.predict([texto_norm])[0].item()
    categoria = CATEGORY_LABEL_MAP.get(pred_id, "otros")

    logger.info(f"[SETFIT] Categoría: {categoria} (id={pred_id})")
    return categoria