Upload cnn_bigru/models/multimodal_model.py with huggingface_hub

cnn_bigru/models/multimodal_model.py

@@ -0,0 +1,183 @@
+"""multimodal_model.py — Modelo multimodal CNN-BiGRU completo.
+
+Combina:
+    - CooperativeCNNBiGRU para texto (dual-stream A/B)
+    - ImageEncoder para modalidade visual
+    - AudioEncoder para modalidade auditiva
+    - MultimodalFusion com gated fusion
+    - Cabeçalho de classificação + cabeçalho de geração (LM head)
+
+Suporta dois modos:
+    mode="classify": retorna logits de classificação [B, num_classes]
+    mode="generate": retorna logits de LM [B, T, vocab_size] (autoregressivo)
+"""
+from __future__ import annotations
+
+import logging
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .cooperative_bigru import CooperativeCNNBiGRU
+from .multimodal_encoders import ImageEncoder, AudioEncoder, MultimodalFusion
+
+logger = logging.getLogger(__name__)
+
+
+class MultimodalCNNBiGRU(nn.Module):
+    """Modelo multimodal completo: texto (dual A/B) + imagem + áudio.
+
+    Fluxo:
+        1. Texto A e Texto B -> CooperativeCNNBiGRU -> fused [B, 256]
+        2. Imagem -> ImageEncoder -> mean-pool -> [B, D_img]
+        3. Áudio  -> AudioEncoder -> mean-pool -> [B, D_aud]
+        4. Fusão multimodal -> [B, out_dim]
+        5. Classificação: Linear(out_dim, num_classes)
+        6. Geração (opcional): Linear(out_dim, vocab_size) autoregressivo
+    """
+
+    def __init__(
+        self,
+        vocab_size: int,
+        num_classes: int = 3,
+        embedding_dim: int = 64,
+        cnn_filters: int = 64,
+        gru_hidden: int = 64,
+        n_heads: int = 4,
+        dropout: float = 0.1,
+        pad_idx: int = 1,
+        img_channels: int = 1,
+        img_hidden: int = 32,
+        img_out_dim: int = 64,
+        audio_freq: int = 40,
+        audio_hidden: int = 32,
+        audio_out_dim: int = 64,
+        fusion_dim: int = 128,
+        use_spectral_norm: bool = False,
+    ):
+        super().__init__()
+        self.pad_idx = pad_idx
+        self.vocab_size = vocab_size
+        self.num_classes = num_classes
+
+        # Núcleo cooperativo de texto
+        self.text_core = CooperativeCNNBiGRU(
+            vocab_size=vocab_size,
+            embedding_dim=embedding_dim,
+            cnn_filters=cnn_filters,
+            gru_hidden=gru_hidden,
+            n_heads=n_heads,
+            dropout=dropout,
+            pad_idx=pad_idx,
+            use_spectral_norm=use_spectral_norm,
+        )
+        d_text = self.text_core.feat_dim * 2  # 256 (fused A+B = 256, *2 para foco)
+
+        # Encoders multimodais
+        self.image_encoder = ImageEncoder(
+            in_channels=img_channels,
+            out_dim=img_out_dim,
+            hidden_dim=img_hidden,
+        )
+        self.audio_encoder = AudioEncoder(
+            n_freq=audio_freq,
+            out_dim=audio_out_dim,
+            hidden_dim=audio_hidden,
+        )
+
+        # Pooling de imagem/áudio para vetor fixo
+        self.img_proj = nn.Linear(img_out_dim, img_out_dim)
+        self.aud_proj = nn.Linear(audio_out_dim, audio_out_dim)
+
+        # Fusão multimodal
+        self.fusion = MultimodalFusion(
+            d_text=d_text,
+            d_img=img_out_dim,
+            d_aud=audio_out_dim,
+            out_dim=fusion_dim,
+            dropout=dropout,
+        )
+
+        # Cabeçalho de classificação
+        self.classifier = nn.Linear(fusion_dim, num_classes)
+
+        # Cabeçalho de geração (LM head) — para modo autoregressivo
+        self.lm_head = nn.Linear(fusion_dim, vocab_size, bias=False)
+
+        # Dropout final
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(
+        self,
+        input_ids_a: torch.Tensor,
+        input_ids_b: torch.Tensor,
+        images: Optional[torch.Tensor] = None,
+        audios: Optional[torch.Tensor] = None,
+        attn_mask_a: Optional[torch.Tensor] = None,
+        attn_mask_b: Optional[torch.Tensor] = None,
+        mode: str = "classify",
+    ) -> dict:
+        """
+        Args:
+            input_ids_a, input_ids_b: [B, T] tokens dos streams A e B
+            images: [B, C, H, W] ou None
+            audios: [B, 1, T, F] ou None
+            mode: "classify" | "generate"
+
+        Returns:
+            dict com logits (classificação ou geração) + features intermediárias
+        """
+        # 1. Núcleo cooperativo de texto
+        text_out = self.text_core(input_ids_a, input_ids_b, return_sequences=(mode == "generate"))
+        fused_text = text_out["fused"]  # [B, 256]
+
+        # 2. Imagem
+        if images is not None:
+            img_seq = self.image_encoder(images)  # [B, T_img, D_img]
+            img_feat = img_seq.mean(dim=1)         # [B, D_img]
+            img_feat = self.img_proj(img_feat)
+        else:
+            # Vetor zero se não houver modalidade
+            img_feat = torch.zeros(
+                fused_text.size(0), self.img_proj.out_features,
+                device=fused_text.device, dtype=fused_text.dtype,
+            )
+
+        # 3. Áudio
+        if audios is not None:
+            aud_seq = self.audio_encoder(audios)  # [B, T_aud, D_aud]
+            aud_feat = aud_seq.mean(dim=1)
+            aud_feat = self.aud_proj(aud_feat)
+        else:
+            aud_feat = torch.zeros(
+                fused_text.size(0), self.aud_proj.out_features,
+                device=fused_text.device, dtype=fused_text.dtype,
+            )
+
+        # 4. Fusão multimodal
+        fused = self.fusion(fused_text, img_feat, aud_feat)  # [B, fusion_dim]
+        fused = self.dropout(fused)
+
+        result = {
+            "fused": fused,
+            "text_out": text_out,
+            "img_feat": img_feat,
+            "aud_feat": aud_feat,
+        }
+
+        # 5. Cabeçalhos
+        if mode == "classify":
+            result["logits"] = self.classifier(fused)  # [B, num_classes]
+        elif mode == "generate":
+            # Modo geração: projetar fusão para vocab_size
+            # (modo simplificado: geração single-step baseada na fusão)
+            result["logits"] = self.lm_head(fused)  # [B, vocab_size]
+        else:
+            raise ValueError(f"mode deve ser 'classify' ou 'generate', recebeu '{mode}'")
+
+        return result
+
+
+__all__ = ["MultimodalCNNBiGRU"]