Upload cnn_bigru/losses/losses.py with huggingface_hub

cnn_bigru/losses/losses.py

@@ -0,0 +1,207 @@
+"""losses.py — Sistema de múltiplas perdas para o CNN-BiGRU cooperativo.
+
+Implementa as perdas descritas na seção 8 de dados.txt:
+
+    L_G         = CrossEntropy(logits_prova, prova_alvo)
+    L_penal     = λ * Σ (1 - v_t)                        (penalidade linear)
+    L_exp_penal = μ * Σ exp(γ * (1 - v_t))               (penalidade exponencial)
+    L_V         = BCE(v_t, rotulo_real)                  (perda do verificador)
+    L_AH        = Σ (1 - τ_t)                            (perda anti-alucinação)
+    L_reg       = L2 + estimativa de curvatura           (regularização)
+    L_total     = α * L_G_total + β * L_V + γ_loss * L_AH + δ * L_reg
+
+    onde:
+        L_G_total = L_G + λ_penal * L_penal + μ_exp * L_exp_penal
+"""
+from __future__ import annotations
+
+import logging
+import math
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class LossConfig:
+    """Configuração das perdas e pesos."""
+    # Pesos das perdas (ALPHA, BETA, GAMMA, DELTA do pseudocódigo)
+    alpha: float = 1.0   # peso de L_G_total
+    beta: float = 0.5    # peso de L_V
+    gamma_loss: float = 0.3   # peso de L_AH
+    delta: float = 0.01   # peso de L_reg
+
+    # Penalidades
+    lambda_penal: float = 0.1   # LAMBDA_PENAL
+    mu_exp_penal: float = 0.05  # MU_EXP_PENAL
+    gamma_exp: float = 1.0      # GAMMA_EXP
+    threshold_err: float = 0.5  # THRESHOLD_ERR para penalidade exponencial
+
+    # Regularização
+    l2_reg: float = 1e-5
+
+    # Curvatura
+    use_curvature: bool = True
+    curvature_eps: float = 1e-3
+
+
+class MultiLoss(nn.Module):
+    """Computa todas as perdas do sistema cooperativo.
+
+    Uso:
+        loss_fn = MultiLoss(cfg)
+        result = loss_fn(logits, targets, v_list, tau_list, generator, verifier)
+        loss = result['total']
+        loss.backward()
+    """
+
+    def __init__(self, config: LossConfig, pad_idx: int = 1):
+        super().__init__()
+        self.cfg = config
+        self.pad_idx = pad_idx
+        self.ce = nn.CrossEntropyLoss(ignore_index=pad_idx, reduction='mean')
+        self.bce = nn.BCELoss(reduction='mean')
+
+    def generation_loss(
+        self,
+        logits: torch.Tensor,  # [B, T, V]
+        targets: torch.Tensor,  # [B, T]
+    ) -> torch.Tensor:
+        """L_G: CrossEntropy ignorando padding."""
+        B, T, V = logits.shape
+        return self.ce(logits.reshape(B * T, V), targets.reshape(B * T))
+
+    def penalty_loss(
+        self,
+        v_list: List[torch.Tensor],  # lista de [B, 1] probabilidades do verificador
+    ) -> Dict[str, torch.Tensor]:
+        """Calcula L_penal (linear) e L_exp_penal (exponencial)."""
+        if not v_list:
+            zero = torch.tensor(0.0)
+            return {"linear": zero, "exp": zero}
+
+        # Stack: [T, B, 1] -> [T, B]
+        v_stack = torch.stack([v.squeeze(-1) if v.dim() == 2 else v for v in v_list], dim=0)
+        # Penalidade linear: Σ (1 - v)
+        l_penal = (1.0 - v_stack).sum() / max(v_stack.numel(), 1)
+
+        # Penalidade exponencial: Σ exp(γ * (1 - v)) onde v < threshold
+        mask = v_stack < self.cfg.threshold_err
+        if mask.any():
+            # exp(γ * (1 - v)) para os graves
+            exp_vals = torch.exp(self.cfg.gamma_exp * (1.0 - v_stack[mask]))
+            l_exp = exp_vals.sum() / max(v_stack.numel(), 1)
+        else:
+            l_exp = torch.tensor(0.0, device=v_stack.device)
+
+        return {"linear": l_penal, "exp": l_exp}
+
+    def verifier_loss(
+        self,
+        v_list: List[torch.Tensor],
+        labels: List[torch.Tensor],  # lista de [B, 1] com 0/1
+    ) -> torch.Tensor:
+        """L_V: BCE entre v_t e rotulo_real."""
+        if not v_list:
+            return torch.tensor(0.0)
+        total = torch.tensor(0.0)
+        for v, label in zip(v_list, labels):
+            v_flat = (v.squeeze(-1) if v.dim() == 2 else v).clamp(1e-7, 1 - 1e-7)
+            label_flat = label.squeeze(-1) if label.dim() == 2 else label
+            total = total + self.bce(v_flat, label_flat.float())
+        return total / max(len(v_list), 1)
+
+    def anti_hallucination_loss(
+        self,
+        tau_list: List[torch.Tensor],  # lista de [B, 1] valores de verdade suave
+    ) -> torch.Tensor:
+        """L_AH = Σ (1 - τ_t)."""
+        if not tau_list:
+            return torch.tensor(0.0)
+        total = torch.tensor(0.0)
+        for tau in tau_list:
+            tau_flat = tau.squeeze(-1) if tau.dim() == 2 else tau
+            total = total + (1.0 - tau_flat).mean()
+        return total / max(len(tau_list), 1)
+
+    def regularization_loss(
+        self,
+        models: List[nn.Module],
+    ) -> torch.Tensor:
+        """L_reg = L2 + estimativa de curvatura (aproximada via norma de pesos)."""
+        l2 = torch.tensor(0.0)
+        for model in models:
+            for p in model.parameters():
+                if p.requires_grad:
+                    l2 = l2 + p.pow(2).sum()
+        l2 = self.cfg.l2_reg * l2
+
+        # Estimativa de curvatura aproximada: variância das normas dos gradientes
+        # (calculada externamente e passada via attr, aqui só placeholder)
+        curv = getattr(self, '_curvature_estimate', torch.tensor(0.0))
+        return l2 + (curv if isinstance(curv, torch.Tensor) else torch.tensor(float(curv)))
+
+    def set_curvature_estimate(self, value: torch.Tensor) -> None:
+        """Define a estimativa de curvatura (calculada externamente)."""
+        self._curvature_estimate = value
+
+    def forward(
+        self,
+        logits: torch.Tensor,
+        targets: torch.Tensor,
+        v_list: List[torch.Tensor],
+        tau_list: List[torch.Tensor],
+        verifier_labels: Optional[List[torch.Tensor]] = None,
+        models_for_reg: Optional[List[nn.Module]] = None,
+    ) -> Dict[str, torch.Tensor]:
+        """Computa todas as perdas e combina."""
+        # L_G
+        l_g = self.generation_loss(logits, targets)
+
+        # Penalidades
+        pens = self.penalty_loss(v_list)
+        l_g_total = l_g + self.cfg.lambda_penal * pens["linear"] + self.cfg.mu_exp_penal * pens["exp"]
+
+        # L_V
+        if verifier_labels is not None:
+            l_v = self.verifier_loss(v_list, verifier_labels)
+        else:
+            # Sem rótulos reais: assume que v deveria ser 1 (passo correto)
+            labels_ones = [torch.ones_like(v) for v in v_list]
+            l_v = self.verifier_loss(v_list, labels_ones)
+
+        # L_AH
+        l_ah = self.anti_hallucination_loss(tau_list)
+
+        # L_reg
+        if models_for_reg:
+            l_reg = self.regularization_loss(models_for_reg)
+        else:
+            l_reg = torch.tensor(0.0)
+
+        # Combinação
+        total = (
+            self.cfg.alpha * l_g_total
+            + self.cfg.beta * l_v
+            + self.cfg.gamma_loss * l_ah
+            + self.cfg.delta * l_reg
+        )
+
+        return {
+            "total": total,
+            "l_g": l_g,
+            "l_g_total": l_g_total,
+            "l_penal_linear": pens["linear"],
+            "l_penal_exp": pens["exp"],
+            "l_v": l_v,
+            "l_ah": l_ah,
+            "l_reg": l_reg,
+        }
+
+
+__all__ = ["LossConfig", "MultiLoss"]