"""Modèle multimodal génératif image+texte (Phase 6) — fusion par embeddings. Idée : les features d'image (sortie du ViT, Phase 4) sont projetées dans l'espace texte, puis **insérées comme des « tokens d'image »** au milieu de la séquence de tokens texte, encadrées par `<|image_start|>` … `<|image_end|>`. Le décodeur texte (Phase 2), qui est causal, « voit » l'image avant de produire la réponse. On entraîne avec des **embeddings d'entrée** (`inputs_embeds`) : sur les positions d'image on injecte les features projetées, sur les positions de texte les embeddings de tokens normaux. Les labels valent `-100` sur tout sauf la réponse. Justification : docs/architecture.md §5. """ from __future__ import annotations from dataclasses import dataclass, field import torch import torch.nn as nn from src.model.projectors import MLPProjector from src.model.text_model import CortexTextLM, TextConfig from src.model.vit import VisionEncoder, ViTConfig from src.tokenizer.special_tokens import SPECIAL_TOKEN_IDS IMAGE_START_ID = SPECIAL_TOKEN_IDS["<|image_start|>"] IMAGE_END_ID = SPECIAL_TOKEN_IDS["<|image_end|>"] IMAGE_PAD_ID = SPECIAL_TOKEN_IDS["<|image_pad|>"] IGNORE_INDEX = -100 def _default_text_config() -> TextConfig: return TextConfig(vocab_size=32000, n_layer=2, n_head=4, n_embd=128, ctx_len=64, tie_embeddings=True) def _default_vision_config() -> ViTConfig: return ViTConfig(image_size=32, patch_size=8, n_layer=2, n_head=4, n_embd=128) @dataclass class VLMConfig: text: TextConfig = field(default_factory=_default_text_config) vision: ViTConfig = field(default_factory=_default_vision_config) train_vision: bool = True class CortexVLM(nn.Module): """VLM minimal : ViT (image) + transformer texte, fusion par concaténation d'embeds.""" def __init__(self, config: VLMConfig): super().__init__() self.config = config self.text = CortexTextLM(config.text) self.vision = VisionEncoder(config.vision) self.projector = MLPProjector(config.vision.n_embd, config.text.n_embd) if not config.train_vision: for p in self.vision.parameters(): p.requires_grad = False @property def n_image_tokens(self) -> int: # On n'injecte que les patchs (pas le token CLS) comme tokens d'image. return self.vision.patch_embed.n_patches def image_features(self, images: torch.Tensor) -> torch.Tensor: """Features d'image projetées dans l'espace texte ``(B, n_patches, n_embd)``.""" tokens = self.vision(images) if self.vision.cls_token is not None: tokens = tokens[:, 1:] # retirer le CLS return self.projector(tokens) def forward( self, images: torch.Tensor, input_ids: torch.Tensor, image_mask: torch.Tensor, targets: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor | None]: """``image_mask`` : booléen ``(B, T)`` marquant les positions d'image.""" emb = self.text.token_embedding(input_ids) # (B, T, D) img = self.image_features(images) # (B, Nv, D) expected = int(image_mask.sum(1)[0].item()) if expected != img.size(1): raise ValueError( f"image_mask prévoit {expected} positions d'image mais le ViT en " f"produit {img.size(1)}" ) emb = emb.clone() emb[image_mask] = img.reshape(-1, img.size(-1)) return self.text(inputs_embeds=emb, targets=targets) @torch.no_grad() def generate( self, images: torch.Tensor, prompt_ids: torch.Tensor, image_mask: torch.Tensor, max_new_tokens: int = 30, temperature: float = 1.0, top_k: int | None = None, greedy: bool = False, eos_id: int | None = None, ) -> torch.Tensor: """Génère la réponse après le prompt (le prompt contient le bloc image).""" self.eval() ctx = self.config.text.ctx_len ids, mask = prompt_ids, image_mask for _ in range(max_new_tokens): cur_ids = ids[:, -ctx:] cur_mask = mask[:, -ctx:] emb = self.text.token_embedding(cur_ids) if cur_mask.any(): img = self.image_features(images) emb = emb.clone() emb[cur_mask] = img.reshape(-1, img.size(-1)) logits, _ = self.text(inputs_embeds=emb) logits = logits[:, -1, :] if greedy: nxt = torch.argmax(logits, dim=-1, keepdim=True) else: logits = logits / max(temperature, 1e-6) if top_k: kth = torch.topk(logits, min(top_k, logits.size(-1)), dim=-1).values[..., -1, None] logits = logits.masked_fill(logits < kth, float("-inf")) nxt = torch.multinomial(torch.softmax(logits, dim=-1), 1) ids = torch.cat([ids, nxt], dim=1) mask = torch.cat([mask, torch.zeros_like(nxt, dtype=torch.bool)], dim=1) if eos_id is not None and (nxt == eos_id).all(): break return ids def build_modality_example( start_id: int, pad_id: int, end_id: int, n_tokens: int, prompt_ids: list[int], answer_ids: list[int], ) -> tuple[list[int], list[bool], list[int]]: """Séquence ``prompt [pad × N] réponse`` (modalité quelconque). Les labels sont **décalés** (convention GPT) : ``labels[t] = ids[t+1]`` là où la position prédit un token de réponse, ``-100`` ailleurs. """ ids = list(prompt_ids) mask = [False] * len(ids) ids.append(start_id) mask.append(False) ids.extend([pad_id] * n_tokens) mask.extend([True] * n_tokens) ids.append(end_id) mask.append(False) answer_start = len(ids) ids.extend(answer_ids) mask.extend([False] * len(answer_ids)) labels = [IGNORE_INDEX] * len(ids) for t in range(answer_start - 1, len(ids) - 1): labels[t] = ids[t + 1] return ids, mask, labels def build_example( n_image_tokens: int, prompt_ids: list[int], answer_ids: list[int], ) -> tuple[list[int], list[bool], list[int]]: """Version image de `build_modality_example` (tokens `<|image_*|>`).""" return build_modality_example(IMAGE_START_ID, IMAGE_PAD_ID, IMAGE_END_ID, n_image_tokens, prompt_ids, answer_ids)