Delete models/cooperative_bigru.py with huggingface_hub

models/cooperative_bigru.py

@@ -1,519 +0,0 @@
-"""cooperative_bigru.py — Núcleo cooperativo CNN-BiGRU do projeto.
-
-Implementa os 3 níveis de cooperação descritos em dados.txt:
-
-  NÍVEL 1 — Cooperação Global (Cross-Attention pós-CNN):
-      PonteGlobalCrossAttention com máscaras de padding.
-
-  NÍVEL 2 — Cooperação Passo a Passo (GRUCell com Porta Gated):
-      CellStepCooperativaComAtenuacao: calcula alfa (Sigmoid) que regula
-      o quanto uma rede aceita da outra, com máscara de padding.
-
-  NÍVEL 3 — Ponte de Fusão Final:
-      Concatena os estados forward+backward de cada rede (128 por rede = 256)
-      e projeta para a camada de classificação/decodificação.
-
-Inclui:
-    - Inicialização ortogonal das matrizes recorrentes
-    - Normalização espectral opcional nos pesos lineares
-    - Self-attention no final de cada bloco (requisito do dados.txt)
-"""
-from __future__ import annotations
-
-import logging
-import math
-from typing import Optional, Tuple
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-logger = logging.getLogger(__name__)
-
-
-def _orthogonal_init(layer: nn.Module) -> None:
-    """Aplica inicialização ortogonal nos pesos de uma camada Linear/GRUCell.
-
-    Em caso de falha (e.g., matriz não-quadrada com fan_in < fan_out), registra
-    warning e faz fallback para xavier_uniform_ (não silencioso).
-    """
-    for name, p in layer.named_parameters():
-        if "weight" in name and p.dim() >= 2:
-            try:
-                nn.init.orthogonal_(p)
-            except Exception as e:
-                logger.warning(
-                    f"_orthogonal_init: falhou em {name} (shape={tuple(p.shape)}): {e}. "
-                    f"Usando xavier_uniform_ como fallback."
-                )
-                nn.init.xavier_uniform_(p)
-        elif "bias" in name:
-            nn.init.zeros_(p)
-
-
-# ============================================================================
-# NÍVEL 1 — Cross-Attention Global com máscaras
-# ============================================================================
-
-class PonteGlobalCrossAttention(nn.Module):
-    """Atenção cruzada pós-CNN com máscara de padding.
-
-    A Rede A usa seus dados como Query para encontrar informações relevantes
-    nas Keys da Rede B, agregando os Values mais importantes.
-    """
-
-    def __init__(self, channels: int, n_heads: int = 4, dropout: float = 0.1):
-        super().__init__()
-        assert channels % n_heads == 0, f"channels ({channels}) deve ser divisível por n_heads ({n_heads})"
-        self.channels = channels
-        self.n_heads = n_heads
-        self.head_dim = channels // n_heads
-
-        self.query = nn.Linear(channels, channels)
-        self.key = nn.Linear(channels, channels)
-        self.value = nn.Linear(channels, channels)
-        self.out = nn.Linear(channels, channels)
-        self.dropout = nn.Dropout(dropout)
-
-    def forward(
-        self,
-        x_a: torch.Tensor,
-        x_b: torch.Tensor,
-        mask_b: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-        # x_a: [B, T_a, C], x_b: [B, T_b, C]
-        B, Ta, C = x_a.shape
-        Tb = x_b.size(1)
-
-        q = self.query(x_a).view(B, Ta, self.n_heads, self.head_dim).transpose(1, 2)
-        k = self.key(x_b).view(B, Tb, self.n_heads, self.head_dim).transpose(1, 2)
-        v = self.value(x_b).view(B, Tb, self.n_heads, self.head_dim).transpose(1, 2)
-
-        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
-
-        if mask_b is not None:
-            # mask_b: [B, T_b] -> [B, 1, 1, T_b]
-            scores = scores.masked_fill(mask_b.unsqueeze(1).unsqueeze(1) == 0, -1e9)
-
-        attn = F.softmax(scores, dim=-1)
-        attn = self.dropout(attn)
-        ctx = torch.matmul(attn, v)  # [B, n_heads, Ta, head_dim]
-        ctx = ctx.transpose(1, 2).contiguous().view(B, Ta, C)
-        return x_a + self.out(ctx)
-
-
-# ============================================================================
-# NÍVEL 2 — Célula Cooperativa Passo-a-Passo com Atenuação Gated
-# ============================================================================
-
-class CellStepCooperativaComAtenuacao(nn.Module):
-    """Célula GRU cooperativa com porta de atenuação (gated bridge).
-
-    Calcula alfa = Sigmoid(linear(concat(x_t, h_partner))) em [0, 1]:
-      - Se alfa -> 0, a ponte desativa (isolamento dos fluxos)
-      - Se alfa -> 1, a ponte transfere totalmente o contexto
-
-    Tratamento de padding: se mask_t == 0, mantém o h anterior intacto.
-    """
-
-    def __init__(self, input_size: int, hidden_size: int):
-        super().__init__()
-        self.cell_A = nn.GRUCell(input_size, hidden_size)
-        self.cell_B = nn.GRUCell(input_size, hidden_size)
-        _orthogonal_init(self.cell_A)
-        _orthogonal_init(self.cell_B)
-
-        # Pontes lineares
-        self.ponte_B_para_A = nn.Linear(hidden_size, input_size)
-        self.ponte_A_para_B = nn.Linear(hidden_size, input_size)
-
-        # Porta de atenuação (gated)
-        self.porta_atenuacao_A = nn.Linear(input_size + hidden_size, 1)
-        self.porta_atenuacao_B = nn.Linear(input_size + hidden_size, 1)
-
-    def forward(
-        self,
-        x_t_A: torch.Tensor,
-        x_t_B: torch.Tensor,
-        h_prev_A: torch.Tensor,
-        h_prev_B: torch.Tensor,
-        mask_t_A: torch.Tensor,
-        mask_t_B: torch.Tensor,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        # mask_t: [B, 1]
-
-        # 1. Fator de atenuação (Sigmoid -> [0, 1])
-        alfa_A = torch.sigmoid(self.porta_atenuacao_A(torch.cat((x_t_A, h_prev_B), dim=-1)))
-        alfa_B = torch.sigmoid(self.porta_atenuacao_B(torch.cat((x_t_B, h_prev_A), dim=-1)))
-
-        # 2. Injeção de contexto com atenuação + máscara
-        influencia_B = torch.tanh(self.ponte_B_para_A(h_prev_B)) * alfa_A * mask_t_B
-        influencia_A = torch.tanh(self.ponte_A_para_B(h_prev_A)) * alfa_B * mask_t_A
-
-        x_t_A_coop = x_t_A + influencia_B
-        x_t_B_coop = x_t_B + influencia_A
-
-        # 3. Atualização GRU
-        h_next_A_bruto = self.cell_A(x_t_A_coop, h_prev_A)
-        h_next_B_bruto = self.cell_B(x_t_B_coop, h_prev_B)
-
-        # 4. Preserva h anterior onde a sequência acabou (mask=0)
-        h_next_A = torch.where(mask_t_A == 1, h_next_A_bruto, h_prev_A)
-        h_next_B = torch.where(mask_t_B == 1, h_next_B_bruto, h_prev_B)
-
-        return h_next_A, h_next_B
-
-
-# ============================================================================
-# Self-Attention Final Layer (requisito do dados.txt linha 470)
-# ============================================================================
-
-class SelfAttentionSummary(nn.Module):
-    """Self-attention final layer que resume os passos temporais antes da classificação.
-
-    Implementa uma self-attention TRUE (não usa mean como query, conforme
-    requisito do dados.txt linha 470 "Self-Attention no final de cada bloco"):
-      - Query, Key, Value são todas projeções lineares aprendidas de x
-      - Multi-head com n_heads cabeças
-      - Máscara de padding (ignora tokens [PAD])
-      - Causal opcional (não usado em sumarização, mas disponível)
-      - Residual + LayerNorm para estabilidade
-
-    Saída: [B, hidden_dim] — vetor agregado que resume toda a sequência.
-
-    Args:
-        hidden_dim: dimensão de entrada/saída
-        n_heads: número de cabeças (deve dividir hidden_dim)
-        dropout: probabilidade de dropout
-        use_residual: se True, soma residual com entrada (via LayerNorm)
-        aggregation: "mean" (média dos heads após attn) ou "first" (primeiro token)
-    """
-
-    def __init__(
-        self,
-        hidden_dim: int,
-        n_heads: int = 4,
-        dropout: float = 0.1,
-        use_residual: bool = True,
-        aggregation: str = "mean",
-    ):
-        super().__init__()
-        if hidden_dim % n_heads != 0:
-            raise ValueError(
-                f"hidden_dim ({hidden_dim}) deve ser divisível por n_heads ({n_heads})"
-            )
-        self.hidden_dim = hidden_dim
-        self.n_heads = n_heads
-        self.head_dim = hidden_dim // n_heads
-        self.use_residual = use_residual
-        self.aggregation = aggregation
-
-        # Projeções lineares aprendidas (TRUE self-attention, não mean-pool)
-        self.q_proj = nn.Linear(hidden_dim, hidden_dim)
-        self.k_proj = nn.Linear(hidden_dim, hidden_dim)
-        self.v_proj = nn.Linear(hidden_dim, hidden_dim)
-        self.out_proj = nn.Linear(hidden_dim, hidden_dim)
-        self.dropout = nn.Dropout(dropout)
-
-        if use_residual:
-            self.ln = nn.LayerNorm(hidden_dim)
-
-        # Token especial de "sumarização" aprendido (CLS-like)
-        # Adicionado no início da sequência para servir como query global
-        self.cls_token = nn.Parameter(torch.zeros(1, 1, hidden_dim))
-        nn.init.normal_(self.cls_token, std=0.02)
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        mask: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-        """
-        Args:
-            x: [B, T, H] — sequência de saídas da BiGRU
-            mask: [B, T] — 1 para token real, 0 para PAD (None = todos válidos)
-
-        Returns:
-            summary: [B, H] — vetor sumarizado
-        """
-        B, T, H = x.shape
-
-        # Prepend CLS token: [B, 1, H] + [B, T, H] -> [B, T+1, H]
-        cls = self.cls_token.expand(B, -1, -1)
-        x_ext = torch.cat([cls, x], dim=1)
-        T_ext = T + 1
-
-        # Estender máscara para o CLS (sempre ativo)
-        if mask is not None:
-            cls_mask = torch.ones(B, 1, device=x.device, dtype=mask.dtype)
-            mask_ext = torch.cat([cls_mask, mask], dim=1)  # [B, T+1]
-        else:
-            mask_ext = None
-
-        # Projeções Q, K, V (todas aprendidas — true self-attention)
-        q = self.q_proj(x_ext)  # [B, T+1, H]
-        k = self.k_proj(x_ext)
-        v = self.v_proj(x_ext)
-
-        # Reshape para multi-head: [B, n_heads, T+1, head_dim]
-        q = q.view(B, T_ext, self.n_heads, self.head_dim).transpose(1, 2)
-        k = k.view(B, T_ext, self.n_heads, self.head_dim).transpose(1, 2)
-        v = v.view(B, T_ext, self.n_heads, self.head_dim).transpose(1, 2)
-
-        # Scaled dot-product attention (sem máscara causal — é encoder-style)
-        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
-        # scores: [B, n_heads, T+1, T+1]
-
-        if mask_ext is not None:
-            # Aplicar máscara nas chaves: tokens PAD devem receber atenção zero
-            # mask_ext: [B, T+1] -> [B, 1, 1, T+1]
-            scores = scores.masked_fill(
-                mask_ext.unsqueeze(1).unsqueeze(1) == 0, -1e9
-            )
-
-        # Lidar com linhas totalmente mascaradas (evitar NaN)
-        all_neg = (scores <= -1e9).all(dim=-1, keepdim=True)
-        scores = torch.where(all_neg, torch.zeros_like(scores), scores)
-
-        attn = F.softmax(scores, dim=-1)
-        attn = self.dropout(attn)
-
-        # Contexto: [B, n_heads, T+1, head_dim]
-        ctx = torch.matmul(attn, v)
-        ctx = ctx.transpose(1, 2).contiguous().view(B, T_ext, H)
-        ctx = self.out_proj(ctx)
-
-        # Residual + LayerNorm (estabilidade)
-        if self.use_residual:
-            ctx = self.ln(ctx + x_ext)
-
-        # Sumarização: usar o CLS token (posição 0) como saída
-        if self.aggregation == "first":
-            summary = ctx[:, 0, :]  # [B, H]
-        else:
-            # Média sobre todas as posições válidas (com mask)
-            if mask_ext is not None:
-                # Ponderar pela máscara
-                weights = mask_ext.unsqueeze(-1).float()  # [B, T+1, 1]
-                sum_ctx = (ctx * weights).sum(dim=1)  # [B, H]
-                count = weights.sum(dim=1).clamp(min=1.0)
-                summary = sum_ctx / count
-            else:
-                summary = ctx.mean(dim=1)
-
-        return summary
-
-
-# ============================================================================
-# Núcleo Cooperativo CNN-BiGRU
-# ============================================================================
-
-class CooperativeCNNBiGRU(nn.Module):
-    """Núcleo cooperativo CNN-BiGRU (dual-stream) com 3 níveis de ponte.
-
-    Args:
-        vocab_size: tamanho do vocabulário.
-        embedding_dim: dimensão do embedding de tokens.
-        cnn_filters: número de filtros da Conv1d.
-        gru_hidden: dimensão oculta de cada direção da GRU.
-        n_heads: número de cabeças para cross-attention.
-        dropout: probabilidade de dropout.
-        pad_idx: índice do token de padding.
-        use_spectral_norm: se True, aplica normalização espectral nas lineares.
-    """
-
-    def __init__(
-        self,
-        vocab_size: int,
-        embedding_dim: int = 64,
-        cnn_filters: int = 64,
-        gru_hidden: int = 64,
-        n_heads: int = 4,
-        dropout: float = 0.1,
-        pad_idx: int = 1,
-        use_spectral_norm: bool = False,
-    ):
-        super().__init__()
-        self.pad_idx = pad_idx
-        self.cnn_filters = cnn_filters
-        self.gru_hidden = gru_hidden
-
-        self.embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx=pad_idx)
-        nn.init.orthogonal_(self.embedding.weight)
-
-        # CNN extractors (cada stream tem o seu)
-        self.conv_A = nn.Conv1d(embedding_dim, cnn_filters, kernel_size=3, padding=1)
-        self.conv_B = nn.Conv1d(embedding_dim, cnn_filters, kernel_size=3, padding=1)
-
-        # Nível 1: Pontes globais (cross-attention)
-        self.ponte_global_A = PonteGlobalCrossAttention(cnn_filters, n_heads=n_heads, dropout=dropout)
-        self.ponte_global_B = PonteGlobalCrossAttention(cnn_filters, n_heads=n_heads, dropout=dropout)
-
-        # Nível 2: Células cooperativas passo-a-passo
-        self.gru_coop_direto = CellStepCooperativaComAtenuacao(cnn_filters, gru_hidden)
-        self.gru_coop_inverso = CellStepCooperativaComAtenuacao(cnn_filters, gru_hidden)
-
-        # Self-attention para sumarização temporal (requisito do dados.txt)
-        self.self_attn_A = SelfAttentionSummary(gru_hidden * 2, dropout=dropout)
-        self.self_attn_B = SelfAttentionSummary(gru_hidden * 2, dropout=dropout)
-
-        # Camadas de saída (sem classificador aqui — o classificador fica no modelo multimodal)
-        self.dropout = nn.Dropout(dropout)
-        self.feat_dim = gru_hidden * 2  # 128 por stream (forward + backward concat)
-
-        # Spectral norm opcional
-        if use_spectral_norm:
-            self.conv_A = nn.utils.spectral_norm(self.conv_A)
-            self.conv_B = nn.utils.spectral_norm(self.conv_B)
-            self.ponte_global_A.query = nn.utils.spectral_norm(self.ponte_global_A.query)
-            self.ponte_global_A.key = nn.utils.spectral_norm(self.ponte_global_A.key)
-            self.ponte_global_A.value = nn.utils.spectral_norm(self.ponte_global_A.value)
-            self.ponte_global_B.query = nn.utils.spectral_norm(self.ponte_global_B.query)
-            self.ponte_global_B.key = nn.utils.spectral_norm(self.ponte_global_B.key)
-            self.ponte_global_B.value = nn.utils.spectral_norm(self.ponte_global_B.value)
-
-    def forward(
-        self,
-        input_A: torch.Tensor,
-        input_B: torch.Tensor,
-        return_sequences: bool = False,
-    ) -> dict:
-        """
-        Args:
-            input_A: [B, T_a] tokens do stream A
-            input_B: [B, T_b] tokens do stream B
-            return_sequences: se True, retorna também as saídas temporais das BiGRUs
-
-        Returns:
-            dict com:
-                out_A: [B, 128] representação final do stream A
-                out_B: [B, 128] representação final do stream B
-                fused: [B, 256] concatenação (out_A || out_B)
-                (opcional) seq_A: [B, T, 128] saídas temporais do stream A
-                (opcional) seq_B: [B, T, 128] saídas temporais do stream B
-        """
-        batch_size = input_A.size(0)
-        seq_len_A = input_A.size(1)
-        seq_len_B = input_B.size(1)
-        device = input_A.device
-
-        # Máscaras de padding
-        mask_A = (input_A != self.pad_idx).float()
-        mask_B = (input_B != self.pad_idx).float()
-
-        # 1. Embeddings + CNN
-        emb_A = self.embedding(input_A).permute(0, 2, 1)  # [B, D, T_a]
-        emb_B = self.embedding(input_B).permute(0, 2, 1)  # [B, D, T_b]
-
-        feat_A = F.relu(self.conv_A(emb_A)).permute(0, 2, 1) * mask_A.unsqueeze(-1)
-        feat_B = F.relu(self.conv_B(emb_B)).permute(0, 2, 1) * mask_B.unsqueeze(-1)
-
-        # 2. NÍVEL 1: Pontes globais (cross-attention)
-        feat_A_coop = self.ponte_global_A(feat_A, feat_B, mask_B) * mask_A.unsqueeze(-1)
-        feat_B_coop = self.ponte_global_B(feat_B, feat_A, mask_A) * mask_B.unsqueeze(-1)
-
-        # 3. Estados ocultos iniciais
-        h_dir_A = torch.zeros(batch_size, self.gru_hidden, device=device)
-        h_dir_B = torch.zeros(batch_size, self.gru_hidden, device=device)
-        h_inv_A = torch.zeros(batch_size, self.gru_hidden, device=device)
-        h_inv_B = torch.zeros(batch_size, self.gru_hidden, device=device)
-
-        max_steps = max(seq_len_A, seq_len_B)
-        seq_dir_A, seq_dir_B = [], []
-        seq_inv_A, seq_inv_B = [], []
-
-        # 4. NÍVEL 2: Laço direto
-        # Sempre coletamos as sequências temporais (necessário para self-attention final)
-        for t in range(max_steps):
-            x_t_A = feat_A_coop[:, t, :] if t < seq_len_A else torch.zeros(batch_size, self.cnn_filters, device=device)
-            x_t_B = feat_B_coop[:, t, :] if t < seq_len_B else torch.zeros(batch_size, self.cnn_filters, device=device)
-
-            m_t_A = mask_A[:, t].unsqueeze(-1) if t < seq_len_A else torch.zeros(batch_size, 1, device=device)
-            m_t_B = mask_B[:, t].unsqueeze(-1) if t < seq_len_B else torch.zeros(batch_size, 1, device=device)
-
-            h_dir_A, h_dir_B = self.gru_coop_direto(x_t_A, x_t_B, h_dir_A, h_dir_B, m_t_A, m_t_B)
-            seq_dir_A.append(h_dir_A)
-            seq_dir_B.append(h_dir_B)
-
-        # 5. Laço inverso
-        for t in reversed(range(max_steps)):
-            x_t_A = feat_A_coop[:, t, :] if t < seq_len_A else torch.zeros(batch_size, self.cnn_filters, device=device)
-            x_t_B = feat_B_coop[:, t, :] if t < seq_len_B else torch.zeros(batch_size, self.cnn_filters, device=device)
-
-            m_t_A = mask_A[:, t].unsqueeze(-1) if t < seq_len_A else torch.zeros(batch_size, 1, device=device)
-            m_t_B = mask_B[:, t].unsqueeze(-1) if t < seq_len_B else torch.zeros(batch_size, 1, device=device)
-
-            h_inv_A, h_inv_B = self.gru_coop_inverso(x_t_A, x_t_B, h_inv_A, h_inv_B, m_t_A, m_t_B)
-            seq_inv_A.append(h_inv_A)
-            seq_inv_B.append(h_inv_B)
-
-        # 6. Concatenação forward+backward (128 por stream)
-        final_A = torch.cat((h_dir_A, h_inv_A), dim=1)  # [B, 128]
-        final_B = torch.cat((h_dir_B, h_inv_B), dim=1)  # [B, 128]
-
-        # 6.5 SELF-ATTENTION FINAL LAYER (requisito do dados.txt linha 470)
-        # Constrói sequência temporal [B, T, 2H] para aplicar self-attention
-        out_A = final_A
-        out_B = final_B
-        seq_A_full = None
-        seq_B_full = None
-
-        # Aplicar self-attention TRUE (não mean-pool) sobre a sequência temporal
-        if getattr(self, "self_attn_A", None) is not None:
-            # Construir sequência temporal concatenando dir + inv em cada posição t
-            try:
-                seq_A_list = []
-                seq_B_list = []
-                for i in range(max_steps):
-                    if i < len(seq_dir_A) and (max_steps - 1 - i) < len(seq_inv_A):
-                        a_dir = seq_dir_A[i]
-                        a_inv = seq_inv_A[max_steps - 1 - i]
-                        seq_A_list.append(torch.cat((a_dir, a_inv), dim=1))
-                    if i < len(seq_dir_B) and (max_steps - 1 - i) < len(seq_inv_B):
-                        b_dir = seq_dir_B[i]
-                        b_inv = seq_inv_B[max_steps - 1 - i]
-                        seq_B_list.append(torch.cat((b_dir, b_inv), dim=1))
-                if seq_A_list:
-                    seq_A_full = torch.stack(seq_A_list, dim=1)  # [B, T, 2H]
-                if seq_B_list:
-                    seq_B_full = torch.stack(seq_B_list, dim=1)
-            except Exception as e:
-                logger.warning(f"SelfAttention: falha ao construir sequência temporal: {e}")
-                seq_A_full = None
-                seq_B_full = None
-
-            # Aplicar self-attention TRUE (não mean-pool)
-            if seq_A_full is not None:
-                out_A = self.self_attn_A(seq_A_full, mask=mask_A)
-            if seq_B_full is not None:
-                out_B = self.self_attn_B(seq_B_full, mask=mask_B)
-
-        # 7. NÍVEL 3: Fusão final
-        fused = torch.cat((out_A, out_B), dim=1)  # [B, 256]
-
-        result = {
-            "out_A": out_A,
-            "out_B": out_B,
-            "fused": fused,
-            "final_A": final_A,  # estados finais raw (sem self-attention)
-            "final_B": final_B,
-            "mask_A": mask_A,
-            "mask_B": mask_B,
-        }
-
-        if return_sequences:
-            # Sequências temporais [B, T, 2H]
-            if seq_A_full is not None:
-                result["seq_A"] = seq_A_full
-            if seq_B_full is not None:
-                result["seq_B"] = seq_B_full
-
-        return result
-
-
-__all__ = [
-    "PonteGlobalCrossAttention",
-    "CellStepCooperativaComAtenuacao",
-    "SelfAttentionSummary",
-    "CooperativeCNNBiGRU",
-]