src/model/multimodal.py

"""Modèle multimodal complet texte + image + audio (Phase 8).

Un **seul** décodeur texte reçoit, dans la même séquence, les tokens de texte, les
patchs d'image projetés (Phase 4) et les patchs audio projetés (Phase 7), chacun
encadré par ses tokens spéciaux (`<|image_*|>`, `<|audio_*|>`). Le décodeur causal
génère la réponse en s'appuyant sur la ou les modalités présentes.

C'est la généralisation du VLM (Phase 6) à plusieurs modalités : on assemble les
`inputs_embeds` en remplaçant, aux positions marquées, l'embedding du token par la
feature projetée de la modalité correspondante.

Justification : docs/architecture.md §8.
"""

from __future__ import annotations

from dataclasses import dataclass, field

import torch
import torch.nn as nn

from src.model.audio_encoder import AudioConfig, AudioEncoder
from src.model.block import TransformerBlock
from src.model.projectors import MLPProjector
from src.model.text_model import CortexTextLM, TextConfig
from src.model.vit import VisionEncoder, ViTConfig


def _default_text() -> TextConfig:
    return TextConfig(vocab_size=32000, n_layer=3, n_head=4, n_embd=128, ctx_len=128,
                      tie_embeddings=True)


def _default_vision() -> ViTConfig:
    return ViTConfig(image_size=32, patch_size=8, n_layer=2, n_head=4, n_embd=128)


def _default_audio() -> AudioConfig:
    return AudioConfig(n_mels=64, n_frames=96, patch_mels=16, patch_frames=16,
                       n_layer=2, n_head=4, n_embd=128)


@dataclass
class MultimodalConfig:
    text: TextConfig = field(default_factory=_default_text)
    vision: ViTConfig = field(default_factory=_default_vision)
    audio: AudioConfig = field(default_factory=_default_audio)


class VideoEncoder(nn.Module):
    """Encodeur vidéo = ViT partagé sur les frames + transformer temporel.

    Le ViT (Phase 4) encode chaque frame en ``n_patches`` vecteurs ; on ajoute un
    **embedding temporel appris** par frame, puis un petit transformer
    **spatio-temporel** (bidirectionnel) sur l'ensemble des tokens — les patchs de
    frames différentes peuvent donc interagir. Le résultat est ``n_frames × n_patches``
    tokens vidéo.
    """

    def __init__(self, config: ViTConfig, n_frames: int, temporal_layers: int = 2):
        super().__init__()
        self.n_frames = n_frames
        self.frame_encoder = VisionEncoder(config)
        self.n_patches = self.frame_encoder.patch_embed.n_patches
        self.temporal_pos = nn.Parameter(torch.zeros(1, n_frames, config.n_embd))
        self.temporal_blocks = nn.ModuleList([
            TransformerBlock(config.n_embd, config.n_head, dropout=config.dropout,
                             bias=config.bias, use_rope=False, causal=False)
            for _ in range(temporal_layers)
        ])
        self.ln_f = nn.LayerNorm(config.n_embd, bias=config.bias)
        nn.init.normal_(self.temporal_pos, std=0.02)

    def forward(self, frames: torch.Tensor) -> torch.Tensor:
        """``frames`` : ``(B, T, 3, H, W)`` -> ``(B, T*n_patches, D)``."""
        b, t = frames.shape[0], frames.shape[1]
        flat = frames.reshape(b * t, *frames.shape[2:])
        per_frame = self.frame_encoder(flat)
        if self.frame_encoder.cls_token is not None:
            per_frame = per_frame[:, 1:]
        per_frame = per_frame.reshape(b, t, self.n_patches, -1)
        # Ajoute la position temporelle à chaque patch de la frame.
        x = per_frame + self.temporal_pos.unsqueeze(2)
        x = x.reshape(b, t * self.n_patches, -1)
        for block in self.temporal_blocks:
            x = block(x)
        return self.ln_f(x)


class CortexMultimodal(nn.Module):
    """Texte + image + audio, fusionnés dans un unique décodeur texte."""

    def __init__(self, config: MultimodalConfig):
        super().__init__()
        self.config = config
        self.text = CortexTextLM(config.text)
        self.vision = VisionEncoder(config.vision)
        self.audio = AudioEncoder(config.audio)
        self.image_projector = MLPProjector(config.vision.n_embd, config.text.n_embd)
        self.audio_projector = MLPProjector(config.audio.n_embd, config.text.n_embd)

    @property
    def n_image_tokens(self) -> int:
        return self.vision.patch_embed.n_patches

    @property
    def n_audio_tokens(self) -> int:
        return self.audio.n_audio_tokens

    def image_features(self, images: torch.Tensor) -> torch.Tensor:
        tokens = self.vision(images)
        if self.vision.cls_token is not None:
            tokens = tokens[:, 1:]
        return self.image_projector(tokens)

    def audio_features(self, mels: torch.Tensor) -> torch.Tensor:
        return self.audio_projector(self.audio(mels))

    def assemble_embeds(self, input_ids, image_mask, audio_mask,
                        images=None, mels=None) -> torch.Tensor:
        """Embeddings de la séquence, modalités injectées aux positions marquées.

        Seules les lignes du lot qui contiennent des tokens de la modalité sont
        encodées (dans un lot multi-tâches, certaines lignes n'ont pas d'image).
        """
        emb = self.text.token_embedding(input_ids)
        emb = emb.clone()
        if images is not None and bool(image_mask.any()):
            rows = image_mask.any(dim=1)
            feat = self.image_features(images[rows])
            emb[image_mask] = feat.reshape(-1, feat.size(-1))
        if mels is not None and bool(audio_mask.any()):
            rows = audio_mask.any(dim=1)
            feat = self.audio_features(mels[rows])
            emb[audio_mask] = feat.reshape(-1, feat.size(-1))
        return emb

    def forward(self, input_ids, targets, image_mask, audio_mask, images=None, mels=None):
        emb = self.assemble_embeds(input_ids, image_mask, audio_mask, images, mels)
        return self.text(inputs_embeds=emb, targets=targets)

    @torch.no_grad()
    def generate(self, input_ids, image_mask, audio_mask, images=None, mels=None,
                 max_new_tokens: int = 20, greedy: bool = True, eos_id=None):
        self.eval()
        ids, im, am = input_ids, image_mask, audio_mask
        for _ in range(max_new_tokens):
            emb = self.assemble_embeds(ids, im, am, images, mels)
            logits, _ = self.text(inputs_embeds=emb)
            logits = logits[:, -1, :]
            nxt = (torch.argmax(logits, dim=-1, keepdim=True) if greedy
                   else torch.multinomial(torch.softmax(logits, dim=-1), 1))
            ids = torch.cat([ids, nxt], dim=1)
            im = torch.cat([im, torch.zeros_like(nxt, dtype=torch.bool)], dim=1)
            am = torch.cat([am, torch.zeros_like(nxt, dtype=torch.bool)], dim=1)
            if eos_id is not None and (nxt == eos_id).all():
                break
        return ids