src/model/vlm.py

"""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 <start> [pad × N] <end> 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)