"""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, weight_tying: bool = False, ): super().__init__() self.pad_idx = pad_idx self.vocab_size = vocab_size self.num_classes = num_classes self.weight_tying = weight_tying # 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 = 2 * feat_dim_per_stream) # 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) # Weight tying opcional (dados.txt linha 632): # compartilha pesos entre token_embedding e lm_head. # Como lm_head é Linear(fusion_dim, vocab_size) e embedding é # Linear(vocab_size, embedding_dim), o tying direto só funciona se # fusion_dim == embedding_dim. Caso contrário, usamos uma matriz # intermediária de projeção. if weight_tying: if fusion_dim == self.text_core.embedding.embedding_dim: self.lm_head.weight = self.text_core.embedding.weight else: # Projeção intermediária: embedding_dim -> fusion_dim self.tie_proj = nn.Linear( self.text_core.embedding.embedding_dim, fusion_dim, bias=False ) # tying via forward (ver mode="generate") logger.info( f"Weight tying: fusion_dim ({fusion_dim}) != embedding_dim " f"({self.text_core.embedding.embedding_dim}); usando proj intermediária." ) # 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 # As máscaras são derivadas internamente do pad_idx, mas se o usuário # fornecer attn_mask_a/attn_mask_b explicitamente, elas têm precedência # e são usadas para ignorar tokens de padding no cálculo de atenção. text_out = self.text_core(input_ids_a, input_ids_b, return_sequences=(mode == "generate")) fused_text = text_out["fused"] # [B, 256] # Sobrescrever máscaras se fornecidas externamente (para uso downstream) if attn_mask_a is not None: text_out["mask_A"] = attn_mask_a.to(text_out["mask_A"].device).float() if attn_mask_b is not None: text_out["mask_B"] = attn_mask_b.to(text_out["mask_B"].device).float() # 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"]