Upload cnn_bigru/models/cooperative_bigru.py with huggingface_hub

cnn_bigru/models/cooperative_bigru.py

@@ -0,0 +1,379 @@
+"""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",
+]