Upload folder using huggingface_hub

agent_emotion_predict_classifier.py

@@ -0,0 +1,131 @@
+"""
+=== MIA · Agent Emotion Predict Classifier (Text/BETO Embedder + MLP) ===
+Objetivo: predecir la emoción del AGENTE (label_agent: 0..5) a partir de:
+  - el TEXTO del usuario
+  - la EMOCIÓN del texto (label del usuario: 0..5)
+
+Arquitectura:
+  Texto ──▶ Embedder (TextEmbedder ó BETOEmbedder) ─▶ h_text ∈ R^D
+  Label usuario (0..5) ─▶ one-hot(6) ─▶ (feature dropout opcional)
+  Concatenación [h_text ; onehot_label] ─▶ MLP ─▶ logits (6)
+
+Notas:
+- Si usas BETOEmbedder, se recomienda congelarlo (freeze) para esta segunda red.
+- El feature dropout en la one-hot del label obliga al modelo a mirar el TEXTO en los casos ambiguos.
+"""
+
+from typing import List, Optional
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from emotion_classifier_model import TextEmbedder, BETOEmbedder  # reemplaza con tu import real
+
+
+class FeatureDropout(nn.Module):
+    """Apaga aleatoriamente (con prob p) TODA la rama de la one-hot del label en entrenamiento.
+    Si p=0.2, en el 20% de los batches el modelo debe decidir solo con el texto.
+    """
+    def __init__(self, p: float = 0.0):
+        super().__init__()
+        assert 0.0 <= p < 1.0
+        self.p = p
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if not self.training or self.p <= 0.0:
+            return x
+        # Con prob p, zerea todo el vector (por muestra)
+        mask = (torch.rand(x.size(0), 1, device=x.device) > self.p).float()
+        return x * mask
+
+
+class MLP(nn.Module):
+    def __init__(self, input_dim: int, hidden1: int = 256, hidden2: int = 64, num_classes: int = 6, dropout: float = 0.2):
+        super().__init__()
+        self.fc1 = nn.Linear(input_dim, hidden1)
+        self.fc2 = nn.Linear(hidden1, hidden2)
+        self.out = nn.Linear(hidden2, num_classes)
+        self.drop = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = F.relu(self.fc1(x))
+        x = self.drop(x)
+        x = F.relu(self.fc2(x))
+        x = self.drop(x)
+        return self.out(x)
+
+
+class AgentEmotionPredictClassifier(nn.Module):
+    """
+    Segunda red: predice la emoción del AGENTE (0..5) a partir de (texto, label_usuario).
+
+    Parámetros clave:
+      - pretrained_encoder: None → TextEmbedder (emb_dim)
+                            "beto" → BETOEmbedder (768D)
+      - label_feature_dropout: apaga la one-hot a veces para forzar al modelo a usar el texto en casos ambiguos.
+    """
+    def __init__(
+        self,
+        model_name: str = "dccuchile/bert-base-spanish-wwm-cased",
+        pretrained_encoder: Optional[str] = "beto",
+        emb_dim: int = 300,
+        max_length: int = 128,
+        hidden1: int = 256,
+        hidden2: int = 64,
+        num_classes: int = 6,
+        dropout: float = 0.2,
+        label_feature_dropout: float = 0.15,
+        device: Optional[torch.device] = None,
+    ):
+        super().__init__()
+        self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        if pretrained_encoder == "beto":
+            self.embedder = BETOEmbedder(model_name=model_name, max_length=max_length, device=self.device)
+            embed_dim = 768
+        else:
+            self.embedder = TextEmbedder(model_name=model_name, emb_dim=emb_dim, max_length=max_length, device=self.device)
+            embed_dim = emb_dim
+
+        self.label_dim = 6  # one-hot(6)
+        self.feat_drop = FeatureDropout(p=label_feature_dropout)
+        self.classifier = MLP(input_dim=embed_dim + self.label_dim,
+                      hidden1=hidden1, hidden2=hidden2,
+                      num_classes=num_classes, dropout=dropout)  # num_classes = salida del AGENTE (ahora 2)
+        self.to(self.device)
+
+    # ---------- Utils ----------
+    @staticmethod
+    def _one_hot(labels: torch.Tensor, num_classes: int) -> torch.Tensor:
+        # labels: [B] int64 → one-hot [B, C]
+        return F.one_hot(labels.long(), num_classes=num_classes).float()
+
+    def freeze_encoder(self):
+        for p in self.embedder.parameters():
+            p.requires_grad = False
+
+    def unfreeze_encoder(self):
+        for p in self.embedder.parameters():
+            p.requires_grad = True
+
+    # ---------- Forward / Predict ----------
+    def forward(self, texts: List[str], user_labels: torch.Tensor) -> torch.Tensor:
+        """texts: lista de strings (len=B)
+           user_labels: tensor [B] con labels del usuario (0..5)
+        """
+        h_text = self.embedder.embed_batch(texts)              # [B, D]
+        onehot = self._one_hot(user_labels.to(h_text.device), self.label_dim)  # [B, 6]
+        onehot = self.feat_drop(onehot)                        # feature dropout (solo en train)
+        x = torch.cat([h_text, onehot], dim=-1)                # [B, D+6]
+        logits = self.classifier(x)                            # [B, 6]
+        return logits
+
+    @torch.inference_mode()
+    def predict(self, texts: List[str], user_labels: torch.Tensor):
+        self.eval()
+        logits = self.forward(texts, user_labels)
+        probs = logits.softmax(dim=-1)
+        preds = probs.argmax(dim=-1)
+        return preds, probs
+
+
+

best_model_agent.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:253f008b154b309ad8bf7e250a6016a104f61e49669aea50b528fb43f71da7f9
-size 1316318655
+oid sha256:e038b9a44ab93c55483da56820f6fb60742c0d8589d50893d38ccf2eaf2f9014
+size 135

config_agent.json

@@ -0,0 +1,26 @@
+{
+  "model_type": "AgentEmotionPredictClassifier",
+  "base_model_id": "dccuchile/bert-base-spanish-wwm-cased",
+  "pretrained_encoder": "beto",
+  "max_length": 128,
+
+  "hidden1": 256,
+  "hidden2": 64,
+  "dropout": 0.4,
+  "label_feature_dropout": 0.5,
+
+  "num_classes": 2,
+  "present_classes": [1, 2],
+  "class_names": ["alegría", "amor"],
+
+  "label_space_global": {
+    "0": "tristeza",
+    "1": "alegría",
+    "2": "amor",
+    "3": "ira",
+    "4": "miedo",
+    "5": "sorpresa"
+  },
+
+  "notes": "Head binaria entrenada solo con clases 1(alegría) y 2(amor); mapping preservado en present_classes/class_names."
+}

emotion_classifier_model.py

@@ -0,0 +1,207 @@
+"""
+=== MIA · Clasificador de Emociones (Pretrained Encoder + MLP) ===
+- Mantiene compatibilidad con tu API pública.
+- Permite usar tu TextEmbedder aleatorio (emb_dim) o un encoder preentrenado (BETO) con 768D.
+- Expone freeze/unfreeze para controlar el fine-tuning desde el trainer.
+"""
+
+import torch
+import torch.nn as nn
+from typing import List, Optional
+from transformers import AutoTokenizer, AutoModel
+
+
+# ==================== MÓDULO 1A: TextEmbedder (embedding aleatorio) ====================
+class TextEmbedder(nn.Module):
+    """
+    Módulo de Embedding simple:
+    - Usa el tokenizador de BETO para sub-palabras (por conveniencia, vocab, pad_id, etc.)
+    - La representación es un embedding aleatorio + mean pooling (no contextual).
+    """
+    def __init__(
+        self,
+        model_name: str = "dccuchile/bert-base-spanish-wwm-cased",
+        emb_dim: int = 300,
+        max_length: int = 128,
+        device: Optional[torch.device] = None
+    ):
+        super().__init__()
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
+        self.vocab_size = self.tokenizer.vocab_size
+        self.pad_id = self.tokenizer.pad_token_id
+        self.cls_id = self.tokenizer.cls_token_id
+        self.sep_id = self.tokenizer.sep_token_id
+        self.max_length = max_length
+
+        # Capa de embedding
+        self.embedding = nn.Embedding(self.vocab_size, emb_dim, padding_idx=self.pad_id)
+        nn.init.xavier_uniform_(self.embedding.weight)
+        with torch.no_grad():
+            if self.pad_id is not None:
+                self.embedding.weight[self.pad_id].zero_()
+
+        # Regularización opcional (ayuda contra sobreajuste)
+        self.emb_dropout = nn.Dropout(p=0.1)
+
+        self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.to(self.device)
+
+    def embed_batch(self, texts: List[str]) -> torch.Tensor:
+        batch = self.tokenizer(
+            texts, padding=True, truncation=True, max_length=self.max_length, return_tensors="pt"
+        )
+        input_ids = batch["input_ids"].to(self.device)           # [B, T]
+        attention_mask = batch["attention_mask"].to(self.device) # [B, T]
+
+        embeds = self.embedding(input_ids)                       # [B, T, E]
+        if self.training:
+            embeds = self.emb_dropout(embeds)
+
+        mask = attention_mask.bool()                             # [B, T]
+        if self.cls_id is not None:
+            mask = mask & (input_ids != self.cls_id)
+        if self.sep_id is not None:
+            mask = mask & (input_ids != self.sep_id)
+
+        mask_f = mask.unsqueeze(-1).float()                      # [B, T, 1]
+        summed = (embeds * mask_f).sum(dim=1)                    # [B, E]
+        counts = mask_f.sum(dim=1).clamp(min=1.0)                # [B, 1]
+        sentence_vecs = summed / counts                          # [B, E]
+        return sentence_vecs
+
+    def embed_sentence(self, text: str) -> torch.Tensor:
+        return self.embed_batch([text])[0]
+
+
+# ==================== MÓDULO 1B: BETOEmbedder (encoder preentrenado) ====================
+class BETOEmbedder(nn.Module):
+    """
+    Usa el encoder de BETO (BERT en español) para obtener embeddings contextuales.
+    Mean pooling sobre last_hidden_state.
+    Salida: [B, 768]
+    """
+    def __init__(
+        self,
+        model_name: str = "dccuchile/bert-base-spanish-wwm-cased",
+        max_length: int = 128,
+        device: Optional[torch.device] = None
+    ):
+        super().__init__()
+        self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
+        self.encoder = AutoModel.from_pretrained(model_name)
+        self.max_length = max_length
+        self.encoder.to(self.device)
+
+    def embed_batch(self, texts: List[str]) -> torch.Tensor:
+        inputs = self.tokenizer(
+            texts, padding=True, truncation=True, max_length=self.max_length, return_tensors="pt"
+        ).to(self.device)
+        outputs = self.encoder(**inputs)  # last_hidden_state [B, T, 768]
+        last_hidden = outputs.last_hidden_state
+        mask = inputs["attention_mask"].unsqueeze(-1).float()  # [B, T, 1]
+        pooled = (last_hidden * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1e-9)  # [B, 768]
+        return pooled
+
+
+# ==================== MÓDULO 2: MLP Classifier ====================
+class MLPClassifier(nn.Module):
+    """
+    Feedforward para clasificación de emociones:
+    Input → 128 → 64 → 6 (logits)
+    """
+    def __init__(
+        self,
+        input_dim: int = 300,
+        hidden1: int = 128,
+        hidden2: int = 64,
+        num_classes: int = 6,
+        dropout: float = 0.3
+    ):
+        super().__init__()
+        self.fc1 = nn.Linear(input_dim, hidden1)
+        self.relu1 = nn.ReLU()
+        self.dropout1 = nn.Dropout(dropout)
+
+        self.fc2 = nn.Linear(hidden1, hidden2)
+        self.relu2 = nn.ReLU()
+        self.dropout2 = nn.Dropout(dropout)
+
+        self.fc3 = nn.Linear(hidden2, num_classes)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.fc1(x); x = self.relu1(x); x = self.dropout1(x)
+        x = self.fc2(x); x = self.relu2(x); x = self.dropout2(x)
+        x = self.fc3(x)
+        return x
+
+
+# ==================== MÓDULO 3: Modelo Completo ====================
+class EmotionClassifier(nn.Module):
+    """
+    Integra embedder (aleatorio o BETO) + MLP.
+    - `pretrained_encoder=None` → usa TextEmbedder (emb_dim configurable)
+    - `pretrained_encoder="beto"` → usa BETOEmbedder (salida 768D)
+    """
+    def __init__(
+        self,
+        model_name: str = "dccuchile/bert-base-spanish-wwm-cased",
+        emb_dim: int = 300,
+        max_length: int = 128,
+        hidden1: int = 128,
+        hidden2: int = 64,
+        num_classes: int = 6,
+        dropout: float = 0.3,
+        device: Optional[torch.device] = None,
+        pretrained_encoder: Optional[str] = None
+    ):
+        super().__init__()
+        self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        if pretrained_encoder == "beto":
+            self.embedder = BETOEmbedder(model_name=model_name, max_length=max_length, device=self.device)
+            embed_dim = 768
+        else:
+            self.embedder = TextEmbedder(model_name=model_name, emb_dim=emb_dim, max_length=max_length, device=self.device)
+            embed_dim = emb_dim
+
+        self.classifier = MLPClassifier(
+            input_dim=embed_dim, hidden1=hidden1, hidden2=hidden2, num_classes=num_classes, dropout=dropout
+        )
+
+        self.label_map = {0: "tristeza", 1: "alegría", 2: "amor", 3: "ira", 4: "miedo", 5: "sorpresa"}
+
+        self.to(self.device)
+
+    # ---------- Forward & Utils ----------
+    def forward(self, texts: List[str]) -> torch.Tensor:
+        embeddings = self.embedder.embed_batch(texts)  # [B, D]
+        logits = self.classifier(embeddings)           # [B, C]
+        return logits
+
+    def predict(self, texts: List[str], return_probs: bool = False):
+        self.eval()
+        with torch.no_grad():
+            logits = self.forward(texts)
+            probs = torch.softmax(logits, dim=-1)
+            predictions = torch.argmax(probs, dim=-1)
+            emotions = [self.label_map[p.item()] for p in predictions]
+            if return_probs:
+                return emotions, probs.cpu().numpy()
+            return emotions
+
+    def predict_single(self, text: str, return_probs: bool = False):
+        out = self.predict([text], return_probs=return_probs)
+        if return_probs:
+            emotions, probs = out
+            return emotions[0], probs[0]
+        return out[0]
+
+    # ---------- Fine-tuning helpers ----------
+    def freeze_encoder(self):
+        for p in self.embedder.parameters():
+            p.requires_grad = False
+
+    def unfreeze_encoder(self):
+        for p in self.embedder.parameters():
+            p.requires_grad = True

inference_agent_emotion_classifier.py

@@ -0,0 +1,278 @@
+# -*- coding: utf-8 -*-
+"""
+Inferencia para el AgentEmotionPredictClassifier (MIA · segunda red)
+- Busca 'best_model.pt' y 'config_agent.json' en local; si no están y hay
+  huggingface_hub instalado, los descarga del repo indicado.
+- La config DEBE incluir, como mínimo:
+    {
+      "base_model_id": "dccuchile/bert-base-spanish-wwm-cased",
+      "max_length": 128,
+      "hidden1": 256,
+      "hidden2": 64,
+      "num_classes": 2,
+      "dropout": 0.4,
+      "label_feature_dropout": 0.5,
+      "pretrained_encoder": "beto",
+      "present_classes": [1, 2],          # ids originales (0..5) presentes en train
+      "class_names": ["alegría","amor"]   # nombres en el mismo orden del mapeo 0..K-1
+    }
+
+- Uso:
+    from inference_agent_emotion import predict
+    y = predict("No me siento bien", user_label=0)  # 0..5 (tristeza..sorpresa)
+"""
+
+from __future__ import annotations
+import os
+import json
+from pathlib import Path
+from typing import Any, Dict, List, Tuple, Optional, Union
+
+import torch
+
+# Opcional: descarga desde HF si no hay archivos locales
+try:
+    from huggingface_hub import hf_hub_download
+except Exception:
+    hf_hub_download = None
+
+from agent_emotion_predict_classifier import AgentEmotionPredictClassifier
+
+# ---------------- Config ----------------
+REPO_ID = "RustyLinux/MiaPredict"  # cambia por tu repo si usas el Hub
+
+LOCAL_CKPT = Path("best_model_agent.pt")           # checkpoint de la segunda red
+LOCAL_CFG  = Path("config_agent.json")       # config de la segunda red
+
+# Mapa global de emociones (usuario y también nombres canónicos)
+EMOTION_ID2NAME = {
+    0: "tristeza",
+    1: "alegría",
+    2: "amor",
+    3: "ira",
+    4: "miedo",
+    5: "sorpresa",
+}
+EMOTION_NAME2ID = {v: k for k, v in EMOTION_ID2NAME.items()}
+
+_device: torch.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+_model: Optional[AgentEmotionPredictClassifier] = None
+_cfg: Optional[Dict[str, Any]] = None
+_label_map_fwd: Optional[Dict[int, int]] = None   # original_id -> idx(0..K-1) usado en entrenamiento
+_label_map_inv: Optional[Dict[int, int]] = None   # idx(0..K-1) -> original_id (para devolver nombre global)
+
+
+# ---------------- Utilidades internas ----------------
+def _resolve_paths() -> Tuple[str, str]:
+    """
+    Retorna (ckpt_path, cfg_path). Prefiere local; si no, intenta descarga HF.
+    """
+    if LOCAL_CKPT.exists() and LOCAL_CFG.exists():
+        print("✅ Cargando archivos desde local.")
+        return str(LOCAL_CKPT.resolve()), str(LOCAL_CFG.resolve())
+
+    if hf_hub_download is None:
+        raise RuntimeError(
+            "No se encontraron 'best_model_agent.pt' y 'config_agent.json' en local, "
+            "y 'huggingface_hub' no está instalado para descargarlos."
+        )
+
+    print("⬇️  Descargando archivos desde Hugging Face Hub...")
+    ckpt_path = hf_hub_download(repo_id=REPO_ID, filename="best_model.pt")
+    cfg_path  = hf_hub_download(repo_id=REPO_ID, filename="config_agent.json")
+    return ckpt_path, cfg_path
+
+
+def _prepare_label_maps(cfg: Dict[str, Any]) -> Tuple[Dict[int, int], Dict[int, int]]:
+    """
+    Construye los mapeos entre ids originales (0..5) y los índices 0..K-1 usados por la head.
+    """
+    present = cfg.get("present_classes", None)
+    if not present:
+        # Por compatibilidad: si no viene, asumimos [0..num_classes-1], pero se recomienda guardarlo.
+        k = int(cfg.get("num_classes", 2))
+        present = list(range(k))
+    present = list(sorted(int(x) for x in present))
+    fwd = {orig: i for i, orig in enumerate(present)}
+    inv = {i: orig for orig, i in fwd.items()}
+    return fwd, inv
+
+
+def _load_config(cfg_path: str) -> Dict[str, Any]:
+    global _cfg, _label_map_fwd, _label_map_inv
+    if _cfg is not None:
+        return _cfg
+    with open(cfg_path, "r", encoding="utf-8") as f:
+        _cfg = json.load(f)
+    _label_map_fwd, _label_map_inv = _prepare_label_maps(_cfg)
+    return _cfg
+
+
+def _build_model(cfg: Dict[str, Any]) -> AgentEmotionPredictClassifier:
+    model = AgentEmotionPredictClassifier(
+        model_name=cfg.get("base_model_id", "dccuchile/bert-base-spanish-wwm-cased"),
+        pretrained_encoder=cfg.get("pretrained_encoder", "beto"),
+        emb_dim=cfg.get("emb_dim", 300),
+        max_length=cfg.get("max_length", 128),
+        hidden1=cfg.get("hidden1", 256),
+        hidden2=cfg.get("hidden2", 64),
+        num_classes=cfg.get("num_classes", 2),
+        dropout=cfg.get("dropout", 0.4),
+        label_feature_dropout=cfg.get("label_feature_dropout", 0.0),  # en inferencia no se usa
+        device=_device,
+    )
+    # aseguramos eval()
+    model.eval()
+    return model
+
+
+def _load_model() -> AgentEmotionPredictClassifier:
+    global _model
+    if _model is not None:
+        return _model
+
+    ckpt_path, cfg_path = _resolve_paths()
+    cfg = _load_config(cfg_path)
+
+    model = _build_model(cfg)
+
+    state = torch.load(ckpt_path, map_location=_device)
+    if isinstance(state, dict) and "model_state_dict" in state:
+        model.load_state_dict(state["model_state_dict"])
+    else:
+        model.load_state_dict(state)
+
+    model.eval()
+    _model = model
+    print(f"✅ Modelo cargado en {_device} | num_classes={cfg.get('num_classes')} | "
+          f"present_classes={cfg.get('present_classes')}")
+    return _model
+
+
+def _coerce_user_label(label: Union[int, str]) -> int:
+    """
+    Convierte un label de usuario a id 0..5.
+    - Si llega string ("alegría"), lo mapea.
+    - Valida rango si llega int.
+    """
+    if isinstance(label, str):
+        label = label.strip().lower()
+        if label not in EMOTION_NAME2ID:
+            raise ValueError(f"Label de usuario desconocido: {label}. Esperado uno de {list(EMOTION_NAME2ID.keys())}")
+        return EMOTION_NAME2ID[label]
+    if isinstance(label, int):
+        if label < 0 or label > 5:
+            raise ValueError("El user_label debe estar en 0..5.")
+        return label
+    raise TypeError("user_label debe ser int (0..5) o str (nombre de emoción).")
+
+
+def _map_agent_idx_to_original(idx: int) -> int:
+    """
+    Convierte el índice 0..K-1 (head) al id original 0..5 para reportar el nombre global.
+    """
+    if _label_map_inv is None:
+        raise RuntimeError("Mapeos de etiquetas no inicializados.")
+    return _label_map_inv[int(idx)]
+
+
+def _agent_class_names() -> List[str]:
+    """
+    Nombres de clases del agente en el mismo orden que la head (0..K-1).
+    """
+    if _cfg is None:
+        raise RuntimeError("Config no cargada.")
+    names = _cfg.get("class_names", None)
+    if names:
+        return list(names)
+    # fallback: usar nombres globales segun present_classes
+    present = sorted(_cfg.get("present_classes", []))
+    return [EMOTION_ID2NAME[p] for p in present]
+
+
+# ---------------- API de inferencia ----------------
+@torch.inference_mode()
+def predict(text: str, user_label: Union[int, str], return_probs: bool = False) -> Any:
+    """
+    Predice la emoción CON LA QUE DEBE RESPONDER EL AGENTE.
+    Args:
+        text: str
+        user_label: int(0..5) o nombre ("tristeza", "alegría", "amor", "ira", "miedo", "sorpresa")
+        return_probs: si True devuelve (pred_name, probs_dict)
+
+    Returns:
+        - Si return_probs=False: str con el nombre de la emoción objetivo del agente (en nombres globales 0..5).
+        - Si return_probs=True: (pred_name:str, probs:Dict[str,float]) usando los nombres en orden de la head.
+    """
+    model = _load_model()
+    cfg = _cfg  # ya cargada
+    assert cfg is not None
+
+    # 1) preparar entrada
+    u = _coerce_user_label(user_label)
+    user_tensor = torch.tensor([u], dtype=torch.long, device=_device)
+    texts = [text]
+
+    # 2) forward
+    logits = model(texts, user_tensor)              # [1, K]
+    probs = torch.softmax(logits, dim=-1).cpu().numpy()[0]
+    pred_idx = int(probs.argmax())
+
+    # 3) mapear idx(0..K-1) -> id original (0..5) y nombre canónico
+    orig_id = _map_agent_idx_to_original(pred_idx)
+    pred_name = EMOTION_ID2NAME[orig_id]
+
+    if not return_probs:
+        return pred_name
+
+    # nombres amistosos en el orden de la head
+    names_head = _agent_class_names()
+    probs_dict = {names_head[i]: float(probs[i]) for i in range(len(names_head))}
+    return pred_name, probs_dict
+
+
+@torch.inference_mode()
+def predict_batch(texts: List[str], user_labels: List[Union[int, str]], return_probs: bool = False):
+    """
+    Batch de inferencia.
+    - user_labels: lista paralela a texts con ids (0..5) o nombres de emoción.
+    """
+    if len(texts) != len(user_labels):
+        raise ValueError("texts y user_labels deben tener la misma longitud.")
+    model = _load_model()
+
+    # preparar
+    u_ids = [ _coerce_user_label(u) for u in user_labels ]
+    user_tensor = torch.tensor(u_ids, dtype=torch.long, device=_device)
+
+    logits = model(texts, user_tensor)      # [B, K]
+    probs = torch.softmax(logits, dim=-1).cpu().numpy()
+    pred_idxs = probs.argmax(axis=1)
+
+    results = []
+    names_head = _agent_class_names()
+
+    for i, idx in enumerate(pred_idxs):
+        orig_id = _map_agent_idx_to_original(int(idx))
+        pred_name = EMOTION_ID2NAME[orig_id]
+        if return_probs:
+            pvec = probs[i]
+            probs_dict = {names_head[j]: float(pvec[j]) for j in range(len(names_head))}
+            results.append((pred_name, probs_dict))
+        else:
+            results.append(pred_name)
+    return results
+
+
+# ---------------- CLI rápido ----------------
+if __name__ == "__main__":
+    # Ejemplos rápidos
+    txts = [
+        "Tuve ese tipo de sentimiento pero lo ignoré",
+        "Estoy muy feliz con la noticia",
+        "Me molesta lo que pasó",
+    ]
+    # user_label puede ser int o str
+    for t, ulab in zip(txts, [0, "alegría", "ira"]):
+        out = predict(t, user_label=ulab, return_probs=True)
+        print(f"\nTexto: {t}\nUser label: {ulab}\nPredicción agente: {out[0]}\nProbs: {out[1]}")

requirements.txt

@@ -0,0 +1,86 @@
+aiohappyeyeballs==2.6.1
+aiohttp==3.12.15
+aiosignal==1.4.0
+anyio==4.10.0
+attrs==25.3.0
+beautifulsoup4==4.13.5
+certifi==2025.8.3
+charset-normalizer==3.4.3
+contourpy==1.3.3
+cycler==0.12.1
+datasets==4.1.1
+deep-translator==1.11.4
+dill==0.4.0
+filelock==3.19.1
+fonttools==4.60.1
+frozenlist==1.7.0
+fsspec==2025.9.0
+googletrans==4.0.2
+h11==0.16.0
+h2==4.3.0
+hf-xet==1.1.10
+hpack==4.1.0
+httpcore==1.0.9
+httpx==0.28.1
+huggingface-hub==0.35.0
+hyperframe==6.1.0
+idna==3.10
+ijson==3.4.0
+Jinja2==3.1.6
+joblib==1.5.2
+kiwisolver==1.4.9
+MarkupSafe==3.0.3
+matplotlib==3.10.7
+mpmath==1.3.0
+multidict==6.6.4
+multiprocess==0.70.16
+networkx==3.5
+numpy==2.3.3
+nvidia-cublas-cu12==12.8.4.1
+nvidia-cuda-cupti-cu12==12.8.90
+nvidia-cuda-nvrtc-cu12==12.8.93
+nvidia-cuda-runtime-cu12==12.8.90
+nvidia-cudnn-cu12==9.10.2.21
+nvidia-cufft-cu12==11.3.3.83
+nvidia-cufile-cu12==1.13.1.3
+nvidia-curand-cu12==10.3.9.90
+nvidia-cusolver-cu12==11.7.3.90
+nvidia-cusparse-cu12==12.5.8.93
+nvidia-cusparselt-cu12==0.7.1
+nvidia-nccl-cu12==2.27.5
+nvidia-nvjitlink-cu12==12.8.93
+nvidia-nvshmem-cu12==3.3.20
+nvidia-nvtx-cu12==12.8.90
+packaging==25.0
+pandas==2.3.2
+pillow==12.0.0
+propcache==0.3.2
+protobuf==6.33.0
+pyarrow==21.0.0
+pyparsing==3.2.5
+python-dateutil==2.9.0.post0
+pytz==2025.2
+PyYAML==6.0.2
+regex==2025.10.23
+requests==2.32.5
+safetensors==0.6.2
+scikit-learn==1.7.2
+scipy==1.16.2
+seaborn==0.13.2
+setuptools==80.9.0
+six==1.17.0
+sniffio==1.3.1
+soupsieve==2.8
+sympy==1.14.0
+threadpoolctl==3.6.0
+tiktoken==0.12.0
+tokenizers==0.22.1
+torch==2.9.0
+tqdm==4.67.1
+transformers==4.57.1
+triton==3.5.0
+typing_extensions==4.15.0
+tzdata==2025.2
+urllib3==2.5.0
+xxhash==3.5.0
+yarl==1.20.1

train_agent_emotion_classifier.py

@@ -0,0 +1,440 @@
+"""
+=== MIA · Script de Entrenamiento del AgentEmotionPredictClassifier (v2) ===
+- Lee JSON v2 con `label` (usuario) y `label_agent` (objetivo del agente).
+- Soporta encoder preentrenado (BETO) con fine-tuning controlado.
+- Class weights + label smoothing para desbalance.
+- Early stopping por macro-F1.
+- AdamW con dos grupos de LR (encoder vs head) + scheduler lineal con warmup.
+- Clip de gradientes, dump de misclasificados, matrices de confusión (absoluta y normalizada).
+
+Requisitos:
+  pip install transformers scikit-learn seaborn matplotlib tqdm
+"""
+
+import json
+from pathlib import Path
+from typing import Dict, List, Tuple, Optional
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset, DataLoader
+from sklearn.metrics import classification_report, confusion_matrix, f1_score
+from tqdm import tqdm
+import matplotlib.pyplot as plt
+import seaborn as sns
+
+from agent_emotion_predict_classifier import AgentEmotionPredictClassifier
+
+# ==================== DATASET ====================
+class AgentEmotionDataset(Dataset):
+    def __init__(self, data_path: str, label_map: Optional[dict] = None):
+        with open(data_path, 'r', encoding='utf-8') as f:
+            raw = json.load(f)
+        data = raw.get('data', raw)
+        self.texts: List[str] = []
+        self.user_labels: List[int] = []
+        self.agent_labels: List[int] = []
+        for it in data:
+            self.texts.append(it['text'])
+            u = it['label']; a = it['label_agent']
+            u = int(u) if isinstance(u, str) else u
+            a = int(a) if isinstance(a, str) else a
+            if label_map is not None:
+                a = label_map[a]  # remapea a {0..K-1}
+            self.user_labels.append(u)
+            self.agent_labels.append(a)
+
+    def __len__(self):
+        return len(self.texts)
+
+    def __getitem__(self, idx):
+        return self.texts[idx], self.user_labels[idx], self.agent_labels[idx]
+
+def collate_fn(batch):
+    texts, ulabels, alabels = zip(*batch)
+    ulabels = torch.tensor(ulabels, dtype=torch.long)
+    alabels = torch.tensor(alabels, dtype=torch.long)
+    return list(texts), ulabels, alabels
+
+
+# ==================== TRAINER ====================
+class AgentEmotionTrainer:
+    def __init__(
+        self,
+        model: AgentEmotionPredictClassifier,
+        train_loader: DataLoader,
+        val_loader: DataLoader,
+        test_loader: DataLoader,
+        device: Optional[torch.device] = None,
+        lr_encoder: float = 2e-5,
+        lr_head: float = 1e-3,
+        weight_decay: float = 0.01,
+        num_epochs: int = 20,
+        warmup_ratio: float = 0.1,
+        early_stopping_patience: int = 3,
+        warmup_freeze_epochs: int = 2,
+        num_classes: int = 2,
+        class_names: Optional[List[str]] = None,
+    ):
+        self.model = model
+        self.train_loader = train_loader
+        self.val_loader = val_loader
+        self.test_loader = test_loader
+        self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.num_classes = num_classes
+        self.class_names = class_names or ["alegría","amor"]
+
+        # ---------- Pérdida con pesos de clase + label smoothing ----------
+        labels_tensor = torch.tensor(train_loader.dataset.agent_labels)
+        class_counts = torch.bincount(labels_tensor, minlength=self.num_classes).float()
+        safe_counts = class_counts.clamp(min=1.0)
+        inv_freq = (safe_counts.sum() / (self.num_classes * safe_counts)).to(self.device)
+        class_weights = inv_freq / inv_freq.mean()
+        self.criterion = nn.CrossEntropyLoss(weight=class_weights, label_smoothing=0.05)
+
+        # ---------- Optimizador AdamW con 2 grupos (encoder vs head) ----------
+        self.lr_encoder = lr_encoder
+        self.lr_head = lr_head
+        self.weight_decay = weight_decay
+
+        # Warmup: encoder congelado n épocas
+        self.warmup_freeze_epochs = warmup_freeze_epochs
+        self.model.freeze_encoder()
+
+        self.optimizer = self._build_optimizer()
+
+        # ---------- Scheduler lineal con warmup ----------
+        self.num_epochs = num_epochs
+        total_steps = len(self.train_loader) * self.num_epochs
+        from transformers import get_linear_schedule_with_warmup
+        self.scheduler = get_linear_schedule_with_warmup(
+            self.optimizer,
+            num_warmup_steps=int(warmup_ratio * total_steps),
+            num_training_steps=total_steps,
+        )
+
+        # ---------- Tracking ----------
+        self.train_losses, self.val_losses = [], []
+        self.train_accs, self.val_accs = [], []
+        self.val_f1s = []
+        self.best_val_f1 = 0.0
+        self.best_state = None
+        self.early_stopping_patience = early_stopping_patience
+
+    def _build_optimizer(self):
+        encoder_params, head_params = [], []
+        for n, p in self.model.named_parameters():
+            if not p.requires_grad:
+                continue
+            if "embedder.encoder" in n:  # parámetros del Transformer
+                encoder_params.append(p)
+            else:
+                head_params.append(p)
+        return torch.optim.AdamW(
+            [
+                {"params": encoder_params, "lr": self.lr_encoder, "weight_decay": self.weight_decay},
+                {"params": head_params,   "lr": self.lr_head,    "weight_decay": 0.0},
+            ]
+        )
+
+    def train_epoch(self) -> Tuple[float, float]:
+        self.model.train()
+        total_loss, correct, total = 0.0, 0, 0
+        pbar = tqdm(self.train_loader, desc="Training")
+        for texts, ulabels, alabels in pbar:
+            alabels = alabels.to(self.device)
+
+            self.optimizer.zero_grad()
+            logits = self.model(texts, ulabels.to(self.device))
+            loss = self.criterion(logits, alabels)
+
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
+            self.optimizer.step()
+            if self.scheduler is not None:
+                self.scheduler.step()
+
+            total_loss += loss.item()
+            preds = torch.argmax(logits, dim=-1)
+            correct += (preds == alabels).sum().item()
+            total += alabels.size(0)
+            pbar.set_postfix({'loss': f'{loss.item():.4f}', 'acc': f'{100*correct/total:.2f}%'})
+
+        avg_loss = total_loss / len(self.train_loader)
+        accuracy = 100 * correct / total
+        return avg_loss, accuracy
+
+    def validate(self) -> Tuple[float, float, float]:
+        self.model.eval()
+        total_loss, correct, total = 0.0, 0, 0
+        all_preds, all_labels = [], []
+        with torch.no_grad():
+            for texts, ulabels, alabels in tqdm(self.val_loader, desc="Validation"):
+                alabels = alabels.to(self.device)
+                logits = self.model(texts, ulabels.to(self.device))
+                loss = self.criterion(logits, alabels)
+
+                total_loss += loss.item()
+                preds = torch.argmax(logits, dim=-1)
+                correct += (preds == alabels).sum().item()
+                total += alabels.size(0)
+
+                all_preds.extend(preds.cpu().tolist())
+                all_labels.extend(alabels.cpu().tolist())
+
+        avg_loss = total_loss / len(self.val_loader)
+        accuracy = 100 * correct / total
+        macro_f1 = f1_score(all_labels, all_preds, average='macro')
+        return avg_loss, accuracy, macro_f1
+
+    def test(self, save_dir: Path) -> Dict:
+        self.model.eval()
+        all_predictions, all_labels, all_texts = [], [], []
+        with torch.no_grad():
+            for texts, ulabels, alabels in tqdm(self.test_loader, desc="Testing"):
+                logits = self.model(texts, ulabels.to(self.device))
+                predictions = torch.argmax(logits, dim=-1)
+                all_predictions.extend(predictions.cpu().numpy())
+                all_labels.extend(alabels.numpy())
+                all_texts.extend(texts)
+
+        accuracy = 100 * np.mean(np.array(all_predictions) == np.array(all_labels))
+        labels_ord = list(range(self.num_classes))
+        target_names = self.class_names
+
+        report = classification_report(
+            all_labels,
+            all_predictions,
+            labels=labels_ord,
+            target_names=target_names,
+            output_dict=True,
+            zero_division=0
+        )
+        cm_abs = confusion_matrix(all_labels, all_predictions, labels=labels_ord)
+        cm_norm = confusion_matrix(all_labels, all_predictions, labels=labels_ord, normalize='true')
+
+        self._dump_misclassified(all_texts, all_labels, all_predictions, target_names, save_dir / "misclassified.txt")
+        return {
+            'accuracy': accuracy,
+            'classification_report': report,
+            'confusion_matrix_abs': cm_abs,
+            'confusion_matrix_norm': cm_norm,
+        }
+
+    @staticmethod
+    def _dump_misclassified(texts, y_true, y_pred, names, path: Path):
+        with open(path, "w", encoding="utf-8") as f:
+            for t, yt, yp in zip(texts, y_true, y_pred):
+                if yt != yp:
+                    f.write(f"[gold={names[yt]} | pred={names[yp]}] {t}\n")
+
+    def train(self, num_epochs: Optional[int] = None, early_stopping_patience: Optional[int] = None, save_dir: str = "models/agent_emotion"):
+        num_epochs = num_epochs or self.num_epochs
+        if early_stopping_patience is not None:
+            self.early_stopping_patience = early_stopping_patience
+
+        print(f"\n{'='*60}\nIniciando entrenamiento por {num_epochs} épocas\nDevice: {self.device}\n{'='*60}")
+        patience_counter = 0
+        best_f1 = -1.0
+
+        for epoch in range(1, num_epochs + 1):
+            # Unfreeze encoder después del warmup
+            if epoch == self.warmup_freeze_epochs + 1:
+                print("→ Descongelando encoder para fine-tuning...")
+                self.model.unfreeze_encoder()
+                self.optimizer = self._build_optimizer()
+
+            print(f"\nÉpoca {epoch}/{num_epochs}\n" + '-'*60)
+            tr_loss, tr_acc = self.train_epoch()
+            self.train_losses.append(tr_loss); self.train_accs.append(tr_acc)
+            va_loss, va_acc, va_f1 = self.validate()
+            self.val_losses.append(va_loss); self.val_accs.append(va_acc); self.val_f1s.append(va_f1)
+
+            print(f"Train: loss={tr_loss:.4f} acc={tr_acc:.2f}% | Val: loss={va_loss:.4f} acc={va_acc:.2f}% f1m={va_f1:.4f}")
+
+            # Guardar mejor por Macro-F1
+            if va_f1 > best_f1 and best_f1 < 0.9999:
+                best_f1 = va_f1
+                patience_counter = 0
+                Path(save_dir).mkdir(parents=True, exist_ok=True)
+                torch.save({
+                    'epoch': epoch,
+                    'model_state_dict': self.model.state_dict(),
+                    'optimizer_state_dict': self.optimizer.state_dict(),
+                    'train_loss': tr_loss,
+                    'val_loss': va_loss,
+                    'val_acc': va_acc,
+                    'val_f1': va_f1,
+                }, f"{save_dir}/best_model.pt")
+                print(f"  ✓ Mejor modelo guardado (Val Macro-F1: {va_f1:.4f})")
+            else:
+                patience_counter += 1
+                print(f"  No improvement. Patience: {patience_counter}/{self.early_stopping_patience}")
+            if patience_counter >= self.early_stopping_patience:
+                print(f"\nEarly stopping activado en época {epoch}")
+                break
+
+        print(f"\n{'='*60}\nEntrenamiento completado! Mejor Val Macro-F1: {best_f1:.4f}\n{'='*60}")
+
+        # Evaluar en test con mejor checkpoint
+        ckpt = torch.load(f"{save_dir}/best_model.pt", map_location=self.device)
+        self.model.load_state_dict(ckpt['model_state_dict'])
+        test_results = self.test(Path(save_dir))
+
+        print(f"\n{'='*60}\nRESULTADOS EN TEST SET\n{'='*60}")
+        print(f"Test Accuracy: {test_results['accuracy']:.2f}%\n")
+
+        # Graficar y guardar reportes
+        self.plot_training_history(save_dir)
+        self.plot_confusion_matrix(test_results['confusion_matrix_abs'], save_dir, norm=False)
+        self.plot_confusion_matrix(test_results['confusion_matrix_norm'], save_dir, norm=True)
+        self.save_classification_report(test_results['classification_report'], save_dir)
+        return test_results
+
+    # ---------- utilidades de guardado/plot ----------
+    def plot_training_history(self, save_dir: str):
+        plt.figure(figsize=(12, 5))
+        plt.subplot(1, 2, 1)
+        plt.plot(self.train_losses, label='Train Loss', marker='o')
+        plt.plot(self.val_losses, label='Val Loss', marker='s')
+        plt.xlabel('Época'); plt.ylabel('Loss'); plt.title('Loss'); plt.legend(); plt.grid(True, alpha=0.3)
+        plt.subplot(1, 2, 2)
+        plt.plot(self.train_accs, label='Train Acc', marker='o')
+        plt.plot(self.val_accs, label='Val Acc', marker='s')
+        plt.plot(self.val_f1s, label='Val Macro-F1', marker='^')
+        plt.xlabel('Época'); plt.ylabel('Score'); plt.title('Acc / Macro-F1'); plt.legend(); plt.grid(True, alpha=0.3)
+        plt.tight_layout(); Path(save_dir).mkdir(parents=True, exist_ok=True)
+        plt.savefig(f"{save_dir}/training_history.png", dpi=300, bbox_inches='tight'); plt.close()
+        print(f"✓ Gráfica de entrenamiento guardada en: {save_dir}/training_history.png")
+
+    def plot_confusion_matrix(self, cm: np.ndarray, save_dir: str, norm: bool = False):
+        plt.figure(figsize=(8, 6))
+        fmt = '.2f' if norm else 'd'
+        cmap = 'Blues'
+        ticklabels = self.class_names
+        sns.heatmap(cm, annot=True, fmt=fmt, cmap=cmap, xticklabels=ticklabels, yticklabels=ticklabels,
+                    vmin=0, vmax=1 if norm else None)
+        plt.title('Matriz de Confusión ' + ('(Normalizada)' if norm else '(Absoluta)'))
+        plt.ylabel('Etiqueta Real'); plt.xlabel('Etiqueta Predicha')
+        plt.tight_layout(); Path(save_dir).mkdir(parents=True, exist_ok=True)
+        fname = "confusion_matrix_norm.png" if norm else "confusion_matrix_abs.png"
+        plt.savefig(f"{save_dir}/{fname}", dpi=300, bbox_inches='tight'); plt.close()
+        print(f"✓ Matriz de confusión guardada en: {save_dir}/{fname}")
+
+    def save_classification_report(self, report: Dict, save_dir: str):
+        Path(save_dir).mkdir(parents=True, exist_ok=True)
+        with open(f"{save_dir}/classification_report.txt", 'w', encoding='utf-8') as f:
+            f.write("="*60 + "\n")
+            f.write("REPORTE DE CLASIFICACIÓN - TEST SET\n")
+            f.write("="*60 + "\n\n")
+            for emotion, metrics in report.items():
+                if emotion in ['accuracy', 'macro avg', 'weighted avg']:
+                    continue
+                f.write(f"\n{emotion.upper()}:\n")
+                f.write(f"  Precision: {metrics['precision']:.4f}\n")
+                f.write(f"  Recall:    {metrics['recall']:.4f}\n")
+                f.write(f"  F1-Score:  {metrics['f1-score']:.4f}\n")
+                f.write(f"  Support:   {metrics['support']}\n")
+            f.write(f"\n{'='*60}\n")
+            f.write(f"MACRO AVG:\n  Precision: {report['macro avg']['precision']:.4f}\n  Recall:    {report['macro avg']['recall']:.4f}\n  F1-Score:  {report['macro avg']['f1-score']:.4f}\n")
+            f.write(f"\nWEIGHTED AVG:\n  Precision: {report['weighted avg']['precision']:.4f}\n  Recall:    {report['weighted avg']['recall']:.4f}\n  F1-Score:  {report['weighted avg']['f1-score']:.4f}\n")
+            acc = report.get('accuracy', None)
+            if acc is not None:
+                f.write(f"\nACCURACY: {acc:.4f}\n")
+            f.write("="*60 + "\n")
+        print(f"✓ Reporte de clasificación guardado en: {save_dir}/classification_report.txt")
+
+
+# ==================== MAIN ====================
+def main():
+    # Rutas (usa los *v2.json*)
+    DATA_DIR = "models/emotion_classifier/data"
+    TRAIN_PATH = f"{DATA_DIR}/emotion_dataset_train_es_v2.json"
+    VAL_PATH   = f"{DATA_DIR}/emotion_dataset_validation_es_v2.json"
+    TEST_PATH  = f"{DATA_DIR}/emotion_dataset_test_es_v2.json"
+    SAVE_DIR   = "models/agent_emotion"
+
+    # Hiperparámetros
+    BATCH_SIZE = 16
+    NUM_EPOCHS = 20
+    EARLY_STOPPING_PATIENCE = 2
+    WARMUP_FREEZE_EPOCHS = 2 #antes 5
+    LR_ENCODER = 1e-5
+    LR_HEAD = 5e-4
+    WEIGHT_DECAY = 0.05
+    WARMUP_RATIO = 0.1
+
+
+    print("="*60)
+    print("CONFIGURACIÓN DEL ENTRENAMIENTO (AgentEmotion v2)")
+    print("="*60)
+    print(f"Train: {TRAIN_PATH}")
+    print(f"Val:   {VAL_PATH}")
+    print(f"Test:  {TEST_PATH}")
+    print(f"Batch size: {BATCH_SIZE} | Épocas: {NUM_EPOCHS}")
+    print(f"Freeze warmup epochs: {WARMUP_FREEZE_EPOCHS}")
+    print(f"LR encoder: {LR_ENCODER} | LR head: {LR_HEAD}")
+
+    # 1) Detectar clases presentes en TRAIN y remapear a {0..K-1}
+    probe = AgentEmotionDataset(TRAIN_PATH)
+    present_classes = sorted(set(probe.agent_labels))  # p.ej., [1,2]
+    label_names_full = {0:"tristeza",1:"alegría",2:"amor",3:"ira",4:"miedo",5:"sorpresa"}
+    class_names = [label_names_full[c] for c in present_classes]
+    K = len(present_classes)
+    print(f"Clases existentes - valor K : {K}")
+    label_map = {orig:i for i, orig in enumerate(present_classes)}
+    print(f"Clases de agente en train: {present_classes} → K={K} ({class_names})")
+
+    # 2) Recrear datasets aplicando el mapeo
+    train_ds = AgentEmotionDataset(TRAIN_PATH, label_map=label_map)
+    val_ds   = AgentEmotionDataset(VAL_PATH,   label_map=label_map)
+    test_ds  = AgentEmotionDataset(TEST_PATH,  label_map=label_map)
+
+    # 3) DataLoaders
+    train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True,  collate_fn=collate_fn)
+    val_loader   = DataLoader(val_ds,   batch_size=BATCH_SIZE, shuffle=True, collate_fn=collate_fn)
+    test_loader  = DataLoader(test_ds,  batch_size=BATCH_SIZE, shuffle=True, collate_fn=collate_fn)
+
+    # 4) Modelo (salida con K clases)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = AgentEmotionPredictClassifier(
+        pretrained_encoder="beto",
+        max_length=128,
+        hidden1=256,
+        hidden2=64,
+        dropout=0.4,               # antes 0.2
+        label_feature_dropout=0.5, # antes 0.15
+        device=device,
+        num_classes=K,
+    )
+    model.freeze_encoder()
+
+    # 5) Trainer con num_classes y nombres
+    trainer = AgentEmotionTrainer(
+        model=model,
+        train_loader=train_loader,
+        val_loader=val_loader,
+        test_loader=test_loader,
+        device=device,
+        lr_encoder=LR_ENCODER,
+        lr_head=LR_HEAD,
+        weight_decay=WEIGHT_DECAY,
+        num_epochs=NUM_EPOCHS,
+        warmup_ratio=WARMUP_RATIO,
+        early_stopping_patience=EARLY_STOPPING_PATIENCE,
+        warmup_freeze_epochs=WARMUP_FREEZE_EPOCHS,
+        num_classes=K,
+        class_names=class_names,
+    )
+
+    # 6) Entrenar y evaluar
+    trainer.train(
+        num_epochs=NUM_EPOCHS,
+        early_stopping_patience=EARLY_STOPPING_PATIENCE,
+        save_dir=SAVE_DIR,
+    )
+
+
+if __name__ == "__main__":
+    main()