src/train/train_video.py · Frankenstein-Labs/Cortex-0 at main

feat(video): Phase 9 — vidéo (frames + audio synchronisés)

096854b verified 12 days ago

12.9 kB

"""Phase 9 — entraînement du modèle vidéo (images + audio synchronisés).
Une vidéo = ``n_frames`` images + une piste audio synthétique. Le modèle encode
**les frames comme une séquence** (ViT partagé + transformer spatio-temporel) et
l'audio en parallèle, puis répond :
* « Décris la vidéo : » -> « un cercle rouge va vers la droite » (frames) ;
* « Que se passe-t-il ? » -> « un cercle rouge va vers la droite | voyelle a grave »
  (frames + audio synchronisés).
Usage :
    python -m src.train.train_video --out-dir checkpoints/video_toy --epochs 400
"""

from __future__ import annotations

import argparse
import json
import time
from dataclasses import asdict, dataclass
from pathlib import Path

import numpy as np
import torch

from src.data.audio_features import LogMelSpectrogram
from src.data.toy_audio import SR
from src.data.toy_audio import generate_dataset as gen_audio
from src.data.toy_audio import to_arrays as audio_arrays
from src.data.toy_video import generate_dataset as gen_video
from src.data.toy_video import to_arrays as video_arrays
from src.model.audio_encoder import AudioConfig, AudioEncoder
from src.model.multimodal import VideoEncoder
from src.model.projectors import MLPProjector
from src.model.text_model import CortexTextLM, TextConfig
from src.model.vit import ViTConfig
from src.tokenizer import load_tokenizer
from src.tokenizer.special_tokens import BOS_ID, EOS_ID, PAD_ID, SPECIAL_TOKEN_IDS
from src.train.train_text import pick_device

VID_S, VID_P, VID_E = (SPECIAL_TOKEN_IDS["<|video_start|>"], SPECIAL_TOKEN_IDS["<|video_pad|>"],
                       SPECIAL_TOKEN_IDS["<|video_end|>"])
AUD_S, AUD_P, AUD_E = (SPECIAL_TOKEN_IDS["<|audio_start|>"], SPECIAL_TOKEN_IDS["<|audio_pad|>"],
                       SPECIAL_TOKEN_IDS["<|audio_end|>"])
IGNORE_INDEX = -100

TASKS = ("video", "video_audio")


class CortexVideoModel(torch.nn.Module):
    """Encodeur vidéo (ViT + spatio-temporel) + audio + décodeur texte."""

    def __init__(self, text_cfg: TextConfig, vit_cfg: ViTConfig, audio_cfg: AudioConfig,
                 n_frames: int):
        super().__init__()
        self.text = CortexTextLM(text_cfg)
        self.video = VideoEncoder(vit_cfg, n_frames, temporal_layers=2)
        self.audio = AudioEncoder(audio_cfg)
        self.video_projector = MLPProjector(vit_cfg.n_embd, text_cfg.n_embd)
        self.audio_projector = MLPProjector(audio_cfg.n_embd, text_cfg.n_embd)
        self.n_frames = n_frames

    @property
    def n_video_tokens(self) -> int:
        return self.n_frames * self.video.n_patches

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

    def video_features(self, frames):
        return self.video_projector(self.video(frames))

    def audio_features(self, mels):
        return self.audio_projector(self.audio(mels))

    def _assemble(self, input_ids, vmask, amask, frames, mels):
        emb = self.text.token_embedding(input_ids).clone()
        if frames is not None and bool(vmask.any()):
            rows = vmask.any(dim=1)
            feat = self.video_features(frames[rows])
            emb[vmask] = feat.reshape(-1, feat.size(-1))
        if mels is not None and bool(amask.any()):
            rows = amask.any(dim=1)
            feat = self.audio_features(mels[rows])
            emb[amask] = feat.reshape(-1, feat.size(-1))
        return emb

    def forward(self, input_ids, targets, vmask, amask, frames=None, mels=None):
        emb = self._assemble(input_ids, vmask, amask, frames, mels)
        return self.text(inputs_embeds=emb, targets=targets)

    @torch.no_grad()
    def generate(self, input_ids, vmask, amask, frames=None, mels=None,
                 max_new_tokens=20, greedy=True, eos_id=None):
        self.eval()
        ids, vm, am = input_ids, vmask, amask
        for _ in range(max_new_tokens):
            emb = self._assemble(ids, vm, am, frames, 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)
            vm = torch.cat([vm, 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


def make_batch(tokenizer, items, n_vid, n_aud, device):
    rows, vms, ams, labs = [], [], [], []
    for task, item in items:
        prompt = "Décris la vidéo :" if task == "video" else "Que se passe-t-il ?"
        answer = (item[0].caption if task == "video"
                  else f"{item[0].caption} | {item[1].caption}")
        ids = [BOS_ID, *tokenizer.encode(prompt)]
        vm = [False] * len(ids)
        am = [False] * len(ids)
        ids += [VID_S] + [VID_P] * n_vid + [VID_E]
        vm += [False] + [True] * n_vid + [False]
        am += [False] * (n_vid + 2)
        if task == "video_audio":
            ids += [AUD_S] + [AUD_P] * n_aud + [AUD_E]
            am += [False] + [True] * n_aud + [False]
            vm += [False] * (n_aud + 2)
        ans = [*tokenizer.encode(answer), EOS_ID]
        start = len(ids)
        ids += ans
        vm += [False] * len(ans)
        am += [False] * len(ans)
        lab = [IGNORE_INDEX] * len(ids)
        for t in range(start - 1, len(ids) - 1):
            lab[t] = ids[t + 1]
        rows.append(ids)
        vms.append(vm)
        ams.append(am)
        labs.append(lab)
    maxlen = max(len(r) for r in rows)
    for r, vm, am, lab in zip(rows, vms, ams, labs, strict=True):
        pad = maxlen - len(r)
        r += [PAD_ID] * pad
        vm += [False] * pad
        am += [False] * pad
        lab += [IGNORE_INDEX] * pad
    return (torch.tensor(rows, dtype=torch.long, device=device),
            torch.tensor(vms, dtype=torch.bool, device=device),
            torch.tensor(ams, dtype=torch.bool, device=device),
            torch.tensor(labs, dtype=torch.long, device=device))



@dataclass
class VideoTrainConfig:
    out_dir: str = "checkpoints/video_toy"
    n_train: int = 1200
    n_val: int = 120
    grid: int = 32
    n_frames: int = 4
    epochs: int = 400
    batch_size: int = 8
    lr: float = 2e-3
    weight_decay: float = 0.1
    grad_clip: float = 1.0
    text_layers: int = 2
    text_embd: int = 96
    text_heads: int = 4
    vision_layers: int = 1
    vision_embd: int = 96
    patch_size: int = 8
    audio_layers: int = 1
    audio_embd: int = 96
    n_mels: int = 64
    n_frames_audio: int = 96
    patch_mels: int = 16
    patch_frames: int = 16
    seed: int = 1337
    device: str = "auto"


def build_model(cfg: VideoTrainConfig, tokenizer):
    return CortexVideoModel(
        TextConfig(vocab_size=tokenizer.vocab_size, n_layer=cfg.text_layers,
                   n_head=cfg.text_heads, n_embd=cfg.text_embd, ctx_len=300,
                   tie_embeddings=True),
        ViTConfig(image_size=cfg.grid, patch_size=cfg.patch_size,
                  n_layer=cfg.vision_layers, n_head=4, n_embd=cfg.vision_embd),
        AudioConfig(n_mels=cfg.n_mels, n_frames=cfg.n_frames_audio,
                    patch_mels=cfg.patch_mels, patch_frames=cfg.patch_frames,
                    n_layer=cfg.audio_layers, n_head=4, n_embd=cfg.audio_embd),
        cfg.n_frames,
    )


def _mels(items, cfg):
    tf = LogMelSpectrogram(sr=SR, n_fft=256, hop_length=64, n_mels=cfg.n_mels)
    with torch.no_grad():
        m = tf(torch.from_numpy(audio_arrays(items))).squeeze(1)
    t = m.size(-1)
    m = m[:, :, :cfg.n_frames_audio] if t >= cfg.n_frames_audio else torch.nn.functional.pad(
        m, (0, cfg.n_frames_audio - t))
    return m.numpy().astype(np.float32)


def train(cfg: VideoTrainConfig) -> dict:
    torch.manual_seed(cfg.seed)
    device = pick_device(cfg.device)
    tokenizer = load_tokenizer()
    out_dir = Path(cfg.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    vtr = gen_video(cfg.n_train, grid=cfg.grid, n_frames=cfg.n_frames, seed=cfg.seed)
    vva = gen_video(cfg.n_val, grid=cfg.grid, n_frames=cfg.n_frames, seed=cfg.seed + 7)
    atr = gen_audio(cfg.n_train, seed=cfg.seed)
    ava = gen_audio(cfg.n_val, seed=cfg.seed + 7)
    Xtr, Xva = video_arrays(vtr), video_arrays(vva)
    Mtr, Mva = _mels(atr, cfg), _mels(ava, cfg)

    model = build_model(cfg, tokenizer).to(device)
    n_vid, n_aud = model.n_video_tokens, model.n_audio_tokens

    decay, no_decay = [], []
    for _n, p in model.named_parameters():
        if p.requires_grad:
            (decay if p.dim() >= 2 else no_decay).append(p)
    opt = torch.optim.AdamW(
        [{"params": decay, "weight_decay": cfg.weight_decay},
         {"params": no_decay, "weight_decay": 0.0}], lr=cfg.lr)

    rng = np.random.default_rng(cfg.seed)
    t0 = time.time()
    for epoch in range(cfg.epochs):
        idx = rng.integers(0, len(vtr), cfg.batch_size)
        tasks = rng.integers(0, 2, cfg.batch_size)
        items = [(TASKS[tasks[k]], (vtr[idx[k]], atr[idx[k]])) for k in range(len(idx))]
        ids, vm, am, lab = make_batch(tokenizer, items, n_vid, n_aud, device)
        frames = torch.from_numpy(Xtr[idx]).to(device)
        mels = torch.from_numpy(Mtr[idx]).unsqueeze(1).to(device)
        _, loss = model(ids, lab, vm, am, frames=frames, mels=mels)
        opt.zero_grad()
        loss.backward()
        if cfg.grad_clip > 0:
            torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip)
        opt.step()
        if (epoch + 1) % 50 == 0 or epoch == 0:
            print(f"epoch {epoch:4d} | loss {loss.item():.4f} | {time.time() - t0:.0f}s", flush=True)

    torch.save({"model": model.state_dict(), "train_config": asdict(cfg),
                "n_video_tokens": n_vid, "n_audio_tokens": n_aud}, out_dir / "best.pt")

    metrics = evaluate(model, tokenizer, Xva, Mva, vva, ava, n_vid, n_aud, device)
    summary = {"final_loss": round(loss.item(), 4), "epochs": cfg.epochs,
               "elapsed_s": round(time.time() - t0, 1), **metrics}
    (out_dir / "summary.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2),
                                         encoding="utf-8")
    print("Résumé :", json.dumps({k: v for k, v in summary.items() if k != "examples"},
                                  ensure_ascii=False), flush=True)
    return summary


@torch.no_grad()
def evaluate(model, tokenizer, Xva, Mva, vva, ava, n_vid, n_aud, device, n_eval=90):
    model.eval()
    n = min(n_eval, len(vva))

    def tf(batch_items):
        ids, vm, am, lab = make_batch(tokenizer, batch_items, n_vid, n_aud, device)
        frames = torch.from_numpy(Xva[:n]).to(device)
        mels = torch.from_numpy(Mva[:n]).unsqueeze(1).to(device)
        logits, _ = model(ids, lab, vm, am, frames=frames, mels=mels)
        sup = lab != IGNORE_INDEX
        return round((logits.argmax(-1)[sup] == lab[sup]).float().mean().item(), 4)

    vid_items = [("video", (vva[i], ava[i])) for i in range(n)]
    both_items = [("video_audio", (vva[i], ava[i])) for i in range(n)]
    out = {"grounding_video": tf(vid_items), "grounding_video_audio": tf(both_items)}

    examples = []
    for i in range(min(6, n)):
        ids = [BOS_ID, *tokenizer.encode("Que se passe-t-il ?")]
        vm = [False] * len(ids)
        am = [False] * len(ids)
        ids += [VID_S] + [VID_P] * n_vid + [VID_E]
        vm += [False] + [True] * n_vid + [False]
        am += [False] * (n_vid + 2)
        ids += [AUD_S] + [AUD_P] * n_aud + [AUD_E]
        am += [False] + [True] * n_aud + [False]
        vm += [False] * (n_aud + 2)
        out_ids = model.generate(
            torch.tensor([ids], device=device),
            torch.tensor([vm], dtype=torch.bool, device=device),
            torch.tensor([am], dtype=torch.bool, device=device),
            frames=torch.from_numpy(Xva[i:i + 1]).to(device),
            mels=torch.from_numpy(Mva[i:i + 1]).unsqueeze(1).to(device),
            max_new_tokens=24, greedy=True, eos_id=EOS_ID)
        text = tokenizer.decode(out_ids[0].tolist()[len(ids):], skip_special_tokens=True).strip()
        examples.append({"target": f"{vva[i].caption} | {ava[i].caption}", "pred": text})
    correct = sum(1 for e in examples if e["pred"] == e["target"])
    out["video_audio_greedy_acc"] = round(correct / len(examples), 4)
    out["examples"] = examples
    return out


def main() -> None:
    parser = argparse.ArgumentParser(description="Entraîner le modèle vidéo (Phase 9).")
    for field, value in asdict(VideoTrainConfig()).items():
        parser.add_argument(f"--{field.replace('_', '-')}", type=type(value), default=value)
    ns = parser.parse_args()
    train(VideoTrainConfig(**{k.replace("-", "_"): v for k, v in vars(ns).items()}))


if __name__ == "__main__":
    main()

	"""Phase 9 — entraînement du modèle vidéo (images + audio synchronisés).

	Une vidéo = ``n_frames`` images + une piste audio synthétique. Le modèle encode
	les frames comme une séquence (ViT partagé + transformer spatio-temporel) et
	l'audio en parallèle, puis répond :

	* « Décris la vidéo : » -> « un cercle rouge va vers la droite » (frames) ;
	* « Que se passe-t-il ? » -> « un cercle rouge va vers la droite \| voyelle a grave »
	(frames + audio synchronisés).

	Usage :

	python -m src.train.train_video --out-dir checkpoints/video_toy --epochs 400
	"""

	from __future__ import annotations

	import argparse
	import json
	import time
	from dataclasses import asdict, dataclass
	from pathlib import Path

	import numpy as np
	import torch

	from src.data.audio_features import LogMelSpectrogram
	from src.data.toy_audio import SR
	from src.data.toy_audio import generate_dataset as gen_audio
	from src.data.toy_audio import to_arrays as audio_arrays
	from src.data.toy_video import generate_dataset as gen_video
	from src.data.toy_video import to_arrays as video_arrays
	from src.model.audio_encoder import AudioConfig, AudioEncoder
	from src.model.multimodal import VideoEncoder
	from src.model.projectors import MLPProjector
	from src.model.text_model import CortexTextLM, TextConfig
	from src.model.vit import ViTConfig
	from src.tokenizer import load_tokenizer
	from src.tokenizer.special_tokens import BOS_ID, EOS_ID, PAD_ID, SPECIAL_TOKEN_IDS
	from src.train.train_text import pick_device

	VID_S, VID_P, VID_E = (SPECIAL_TOKEN_IDS["<\|video_start\|>"], SPECIAL_TOKEN_IDS["<\|video_pad\|>"],
	SPECIAL_TOKEN_IDS["<\|video_end\|>"])
	AUD_S, AUD_P, AUD_E = (SPECIAL_TOKEN_IDS["<\|audio_start\|>"], SPECIAL_TOKEN_IDS["<\|audio_pad\|>"],
	SPECIAL_TOKEN_IDS["<\|audio_end\|>"])
	IGNORE_INDEX = -100

	TASKS = ("video", "video_audio")


	class CortexVideoModel(torch.nn.Module):
	"""Encodeur vidéo (ViT + spatio-temporel) + audio + décodeur texte."""

	def __init__(self, text_cfg: TextConfig, vit_cfg: ViTConfig, audio_cfg: AudioConfig,
	n_frames: int):
	super().__init__()
	self.text = CortexTextLM(text_cfg)
	self.video = VideoEncoder(vit_cfg, n_frames, temporal_layers=2)
	self.audio = AudioEncoder(audio_cfg)
	self.video_projector = MLPProjector(vit_cfg.n_embd, text_cfg.n_embd)
	self.audio_projector = MLPProjector(audio_cfg.n_embd, text_cfg.n_embd)
	self.n_frames = n_frames

	@property
	def n_video_tokens(self) -> int:
	return self.n_frames * self.video.n_patches

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

	def video_features(self, frames):
	return self.video_projector(self.video(frames))

	def audio_features(self, mels):
	return self.audio_projector(self.audio(mels))

	def _assemble(self, input_ids, vmask, amask, frames, mels):
	emb = self.text.token_embedding(input_ids).clone()
	if frames is not None and bool(vmask.any()):
	rows = vmask.any(dim=1)
	feat = self.video_features(frames[rows])
	emb[vmask] = feat.reshape(-1, feat.size(-1))
	if mels is not None and bool(amask.any()):
	rows = amask.any(dim=1)
	feat = self.audio_features(mels[rows])
	emb[amask] = feat.reshape(-1, feat.size(-1))
	return emb

	def forward(self, input_ids, targets, vmask, amask, frames=None, mels=None):
	emb = self._assemble(input_ids, vmask, amask, frames, mels)
	return self.text(inputs_embeds=emb, targets=targets)

	@torch.no_grad()
	def generate(self, input_ids, vmask, amask, frames=None, mels=None,
	max_new_tokens=20, greedy=True, eos_id=None):
	self.eval()
	ids, vm, am = input_ids, vmask, amask
	for _ in range(max_new_tokens):
	emb = self._assemble(ids, vm, am, frames, 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)
	vm = torch.cat([vm, 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


	def make_batch(tokenizer, items, n_vid, n_aud, device):
	rows, vms, ams, labs = [], [], [], []
	for task, item in items:
	prompt = "Décris la vidéo :" if task == "video" else "Que se passe-t-il ?"
	answer = (item[0].caption if task == "video"
	else f"{item[0].caption} \| {item[1].caption}")
	ids = [BOS_ID, *tokenizer.encode(prompt)]
	vm = [False] * len(ids)
	am = [False] * len(ids)
	ids += [VID_S] + [VID_P] * n_vid + [VID_E]
	vm += [False] + [True] * n_vid + [False]
	am += [False] * (n_vid + 2)
	if task == "video_audio":
	ids += [AUD_S] + [AUD_P] * n_aud + [AUD_E]
	am += [False] + [True] * n_aud + [False]
	vm += [False] * (n_aud + 2)
	ans = [*tokenizer.encode(answer), EOS_ID]
	start = len(ids)
	ids += ans
	vm += [False] * len(ans)
	am += [False] * len(ans)
	lab = [IGNORE_INDEX] * len(ids)
	for t in range(start - 1, len(ids) - 1):
	lab[t] = ids[t + 1]
	rows.append(ids)
	vms.append(vm)
	ams.append(am)
	labs.append(lab)
	maxlen = max(len(r) for r in rows)
	for r, vm, am, lab in zip(rows, vms, ams, labs, strict=True):
	pad = maxlen - len(r)
	r += [PAD_ID] * pad
	vm += [False] * pad
	am += [False] * pad
	lab += [IGNORE_INDEX] * pad
	return (torch.tensor(rows, dtype=torch.long, device=device),
	torch.tensor(vms, dtype=torch.bool, device=device),
	torch.tensor(ams, dtype=torch.bool, device=device),
	torch.tensor(labs, dtype=torch.long, device=device))



	@dataclass
	class VideoTrainConfig:
	out_dir: str = "checkpoints/video_toy"
	n_train: int = 1200
	n_val: int = 120
	grid: int = 32
	n_frames: int = 4
	epochs: int = 400
	batch_size: int = 8
	lr: float = 2e-3
	weight_decay: float = 0.1
	grad_clip: float = 1.0
	text_layers: int = 2
	text_embd: int = 96
	text_heads: int = 4
	vision_layers: int = 1
	vision_embd: int = 96
	patch_size: int = 8
	audio_layers: int = 1
	audio_embd: int = 96
	n_mels: int = 64
	n_frames_audio: int = 96
	patch_mels: int = 16
	patch_frames: int = 16
	seed: int = 1337
	device: str = "auto"


	def build_model(cfg: VideoTrainConfig, tokenizer):
	return CortexVideoModel(
	TextConfig(vocab_size=tokenizer.vocab_size, n_layer=cfg.text_layers,
	n_head=cfg.text_heads, n_embd=cfg.text_embd, ctx_len=300,
	tie_embeddings=True),
	ViTConfig(image_size=cfg.grid, patch_size=cfg.patch_size,
	n_layer=cfg.vision_layers, n_head=4, n_embd=cfg.vision_embd),
	AudioConfig(n_mels=cfg.n_mels, n_frames=cfg.n_frames_audio,
	patch_mels=cfg.patch_mels, patch_frames=cfg.patch_frames,
	n_layer=cfg.audio_layers, n_head=4, n_embd=cfg.audio_embd),
	cfg.n_frames,
	)


	def _mels(items, cfg):
	tf = LogMelSpectrogram(sr=SR, n_fft=256, hop_length=64, n_mels=cfg.n_mels)
	with torch.no_grad():
	m = tf(torch.from_numpy(audio_arrays(items))).squeeze(1)
	t = m.size(-1)
	m = m[:, :, :cfg.n_frames_audio] if t >= cfg.n_frames_audio else torch.nn.functional.pad(
	m, (0, cfg.n_frames_audio - t))
	return m.numpy().astype(np.float32)


	def train(cfg: VideoTrainConfig) -> dict:
	torch.manual_seed(cfg.seed)
	device = pick_device(cfg.device)
	tokenizer = load_tokenizer()
	out_dir = Path(cfg.out_dir)
	out_dir.mkdir(parents=True, exist_ok=True)

	vtr = gen_video(cfg.n_train, grid=cfg.grid, n_frames=cfg.n_frames, seed=cfg.seed)
	vva = gen_video(cfg.n_val, grid=cfg.grid, n_frames=cfg.n_frames, seed=cfg.seed + 7)
	atr = gen_audio(cfg.n_train, seed=cfg.seed)
	ava = gen_audio(cfg.n_val, seed=cfg.seed + 7)
	Xtr, Xva = video_arrays(vtr), video_arrays(vva)
	Mtr, Mva = _mels(atr, cfg), _mels(ava, cfg)

	model = build_model(cfg, tokenizer).to(device)
	n_vid, n_aud = model.n_video_tokens, model.n_audio_tokens

	decay, no_decay = [], []
	for _n, p in model.named_parameters():
	if p.requires_grad:
	(decay if p.dim() >= 2 else no_decay).append(p)
	opt = torch.optim.AdamW(
	[{"params": decay, "weight_decay": cfg.weight_decay},
	{"params": no_decay, "weight_decay": 0.0}], lr=cfg.lr)

	rng = np.random.default_rng(cfg.seed)
	t0 = time.time()
	for epoch in range(cfg.epochs):
	idx = rng.integers(0, len(vtr), cfg.batch_size)
	tasks = rng.integers(0, 2, cfg.batch_size)
	items = [(TASKS[tasks[k]], (vtr[idx[k]], atr[idx[k]])) for k in range(len(idx))]
	ids, vm, am, lab = make_batch(tokenizer, items, n_vid, n_aud, device)
	frames = torch.from_numpy(Xtr[idx]).to(device)
	mels = torch.from_numpy(Mtr[idx]).unsqueeze(1).to(device)
	_, loss = model(ids, lab, vm, am, frames=frames, mels=mels)
	opt.zero_grad()
	loss.backward()
	if cfg.grad_clip > 0:
	torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip)
	opt.step()
	if (epoch + 1) % 50 == 0 or epoch == 0:
	print(f"epoch {epoch:4d} \| loss {loss.item():.4f} \| {time.time() - t0:.0f}s", flush=True)

	torch.save({"model": model.state_dict(), "train_config": asdict(cfg),
	"n_video_tokens": n_vid, "n_audio_tokens": n_aud}, out_dir / "best.pt")

	metrics = evaluate(model, tokenizer, Xva, Mva, vva, ava, n_vid, n_aud, device)
	summary = {"final_loss": round(loss.item(), 4), "epochs": cfg.epochs,
	"elapsed_s": round(time.time() - t0, 1), **metrics}
	(out_dir / "summary.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2),
	encoding="utf-8")
	print("Résumé :", json.dumps({k: v for k, v in summary.items() if k != "examples"},
	ensure_ascii=False), flush=True)
	return summary


	@torch.no_grad()
	def evaluate(model, tokenizer, Xva, Mva, vva, ava, n_vid, n_aud, device, n_eval=90):
	model.eval()
	n = min(n_eval, len(vva))

	def tf(batch_items):
	ids, vm, am, lab = make_batch(tokenizer, batch_items, n_vid, n_aud, device)
	frames = torch.from_numpy(Xva[:n]).to(device)
	mels = torch.from_numpy(Mva[:n]).unsqueeze(1).to(device)
	logits, _ = model(ids, lab, vm, am, frames=frames, mels=mels)
	sup = lab != IGNORE_INDEX
	return round((logits.argmax(-1)[sup] == lab[sup]).float().mean().item(), 4)

	vid_items = [("video", (vva[i], ava[i])) for i in range(n)]
	both_items = [("video_audio", (vva[i], ava[i])) for i in range(n)]
	out = {"grounding_video": tf(vid_items), "grounding_video_audio": tf(both_items)}

	examples = []
	for i in range(min(6, n)):
	ids = [BOS_ID, *tokenizer.encode("Que se passe-t-il ?")]
	vm = [False] * len(ids)
	am = [False] * len(ids)
	ids += [VID_S] + [VID_P] * n_vid + [VID_E]
	vm += [False] + [True] * n_vid + [False]
	am += [False] * (n_vid + 2)
	ids += [AUD_S] + [AUD_P] * n_aud + [AUD_E]
	am += [False] + [True] * n_aud + [False]
	vm += [False] * (n_aud + 2)
	out_ids = model.generate(
	torch.tensor([ids], device=device),
	torch.tensor([vm], dtype=torch.bool, device=device),
	torch.tensor([am], dtype=torch.bool, device=device),
	frames=torch.from_numpy(Xva[i:i + 1]).to(device),
	mels=torch.from_numpy(Mva[i:i + 1]).unsqueeze(1).to(device),
	max_new_tokens=24, greedy=True, eos_id=EOS_ID)
	text = tokenizer.decode(out_ids[0].tolist()[len(ids):], skip_special_tokens=True).strip()
	examples.append({"target": f"{vva[i].caption} \| {ava[i].caption}", "pred": text})
	correct = sum(1 for e in examples if e["pred"] == e["target"])
	out["video_audio_greedy_acc"] = round(correct / len(examples), 4)
	out["examples"] = examples
	return out


	def main() -> None:
	parser = argparse.ArgumentParser(description="Entraîner le modèle vidéo (Phase 9).")
	for field, value in asdict(VideoTrainConfig()).items():
	parser.add_argument(f"--{field.replace('_', '-')}", type=type(value), default=value)
	ns = parser.parse_args()
	train(VideoTrainConfig(**{k.replace("-", "_"): v for k, v in vars(ns).items()}))


	if __name__ == "__main__":
	main()