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