cnn_bigru/models/multimodal_model.py

"""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"]