File size: 6,774 Bytes
9b728b4 096854b 9b728b4 5117807 9b728b4 096854b 9b728b4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 | """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
|