Delete cnn_bigru/models/cooperative_bigru.py with huggingface_hub

cnn_bigru/models/cooperative_bigru.py

@@ -1,379 +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."""
-    for name, p in layer.named_parameters():
-        if "weight" in name and p.dim() >= 2:
-            try:
-                nn.init.orthogonal_(p)
-            except Exception:
-                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 (requisito do dados.txt: sumariza passos temporais)
-# ============================================================================
-
-class SelfAttentionSummary(nn.Module):
-    """Self-attention que resume os passos temporais antes da classificação.
-
-    Recebe a sequência de outputs da BiGRU [B, T, H] e produz um único vetor
-    [B, H] ponderado por atenção aprendida.
-    """
-
-    def __init__(self, hidden_dim: int, dropout: float = 0.1):
-        super().__init__()
-        self.query = nn.Linear(hidden_dim, hidden_dim)
-        self.key = nn.Linear(hidden_dim, hidden_dim)
-        self.value = nn.Linear(hidden_dim, hidden_dim)
-        self.dropout = nn.Dropout(dropout)
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        mask: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-        # x: [B, T, H]
-        # Usa o primeiro token como query (ou um pooling da sequência)
-        q = self.query(x.mean(dim=1, keepdim=True))  # [B, 1, H]
-        k = self.key(x)  # [B, T, H]
-        v = self.value(x)
-
-        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(x.size(-1))  # [B, 1, T]
-        if mask is not None:
-            scores = scores.masked_fill(mask.unsqueeze(1) == 0, -1e9)
-        attn = F.softmax(scores, dim=-1)
-        attn = self.dropout(attn)
-        ctx = torch.matmul(attn, v)  # [B, 1, H]
-        return ctx.squeeze(1)  # [B, H]
-
-
-# ============================================================================
-# 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
-        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)
-            if return_sequences:
-                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)
-            if return_sequences:
-                seq_inv_A.append(h_inv_A)
-                seq_inv_B.append(h_inv_B)
-
-        # 6. Concatenação forward+backward (128 por stream)
-        out_A = torch.cat((h_dir_A, h_inv_A), dim=1)  # [B, 128]
-        out_B = torch.cat((h_dir_B, h_inv_B), dim=1)  # [B, 128]
-
-        # 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,
-            "mask_A": mask_A,
-            "mask_B": mask_B,
-        }
-
-        if return_sequences:
-            # Reconstrói sequência temporal [B, T, 128] somando direto + inverso
-            # Para cada t, h_dir e h_inv são estados — aqui retornamos a concatenação
-            seq_A = []
-            seq_B = []
-            for i in range(max_steps):
-                a_dir = seq_dir_A[i]
-                a_inv = seq_inv_A[max_steps - 1 - i]
-                seq_A.append(torch.cat((a_dir, a_inv), dim=1))
-                b_dir = seq_dir_B[i]
-                b_inv = seq_inv_B[max_steps - 1 - i]
-                seq_B.append(torch.cat((b_dir, b_inv), dim=1))
-            result["seq_A"] = torch.stack(seq_A, dim=1)
-            result["seq_B"] = torch.stack(seq_B, dim=1)
-
-        return result
-
-
-__all__ = [
-    "PonteGlobalCrossAttention",
-    "CellStepCooperativaComAtenuacao",
-    "SelfAttentionSummary",
-    "CooperativeCNNBiGRU",
-]