src/train/train_audio.py

"""Phase 7 — entraînement de l'encodeur audio + transcription de voyelles.

Le modèle : encodeur audio (Phase 7) -> projecteur -> transformer texte, qui
transcrit le mel-spectrogramme en texte (« voyelle a grave »). Même mécanisme de
fusion que le VLM (Phase 6), avec les tokens `<|audio_*|>`.

Usage :

    python -m src.train.train_audio --data-dir data/processed/toy_audio \\
        --out-dir checkpoints/audio_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
import torch.nn as nn

from src.data.prepare_audio import load_mels
from src.model.audio_encoder import AudioConfig, AudioEncoder
from src.model.projectors import MLPProjector
from src.model.text_model import CortexTextLM, TextConfig
from src.model.vlm import build_modality_example
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

AUDIO_START_ID = SPECIAL_TOKEN_IDS["<|audio_start|>"]
AUDIO_END_ID = SPECIAL_TOKEN_IDS["<|audio_end|>"]
AUDIO_PAD_ID = SPECIAL_TOKEN_IDS["<|audio_pad|>"]

PROMPT = "Transcris :"


@dataclass
class AudioTrainConfig:
    data_dir: str = "data/processed/toy_audio"
    out_dir: str = "checkpoints/audio_toy"
    epochs: int = 400
    batch_size: int = 32
    lr: float = 3e-3
    weight_decay: float = 0.1
    grad_clip: float = 1.0
    text_layers: int = 2
    text_embd: int = 96
    text_heads: int = 4
    audio_layers: int = 2
    audio_embd: int = 96
    patch_mels: int = 16
    patch_frames: int = 16
    seed: int = 1337
    device: str = "auto"


class CortexAudioModel(nn.Module):
    """Encodeur audio + projecteur + décodeur texte (transcription)."""

    def __init__(self, text_cfg: TextConfig, audio_cfg: AudioConfig):
        super().__init__()
        self.text = CortexTextLM(text_cfg)
        self.audio = AudioEncoder(audio_cfg)
        self.projector = MLPProjector(audio_cfg.n_embd, text_cfg.n_embd)

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

    def audio_features(self, mel: torch.Tensor) -> torch.Tensor:
        return self.projector(self.audio(mel))

    def forward(self, mel, input_ids, audio_mask, targets):
        emb = self.text.token_embedding(input_ids)
        feat = self.audio_features(mel)
        emb = emb.clone()
        emb[audio_mask] = feat.reshape(-1, feat.size(-1))
        return self.text(inputs_embeds=emb, targets=targets)

    @torch.no_grad()
    def generate(self, mel, prompt_ids, audio_mask, max_new_tokens=12, greedy=True,
                 eos_id=EOS_ID):
        self.eval()
        ids, mask = prompt_ids, audio_mask
        for _ in range(max_new_tokens):
            emb = self.text.token_embedding(ids)
            if mask.any():
                feat = self.audio_features(mel)
                emb = emb.clone()
                emb[mask] = feat.reshape(-1, feat.size(-1))
            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)
            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_model(cfg: AudioTrainConfig, tokenizer, n_mels, n_frames, n_audio_tokens):
    ctx = 16 + n_audio_tokens + 24
    return CortexAudioModel(
        TextConfig(vocab_size=tokenizer.vocab_size, n_layer=cfg.text_layers,
                   n_head=cfg.text_heads, n_embd=cfg.text_embd, ctx_len=ctx,
                   tie_embeddings=True),
        AudioConfig(n_mels=n_mels, n_frames=n_frames, patch_mels=cfg.patch_mels,
                    patch_frames=cfg.patch_frames, n_layer=cfg.audio_layers,
                    n_head=4, n_embd=cfg.audio_embd),
    )


def make_batch(tokenizer, mels_np, captions, n_audio_tokens, device):
    rows, masks, labels = [], [], []
    for cap in captions:
        prompt = [BOS_ID, *tokenizer.encode(PROMPT)]
        ans = [*tokenizer.encode(cap), EOS_ID]
        ids, mask, lab = build_modality_example(
            AUDIO_START_ID, AUDIO_PAD_ID, AUDIO_END_ID, n_audio_tokens, prompt, ans)
        rows.append(ids)
        masks.append(mask)
        labels.append(lab)
    maxlen = max(len(r) for r in rows)
    for r, m, lab in zip(rows, masks, labels, strict=True):
        pad = maxlen - len(r)
        r += [PAD_ID] * pad
        m += [False] * pad
        lab += [-100] * pad
    return (torch.from_numpy(mels_np).unsqueeze(1).to(device),
            torch.tensor(rows, dtype=torch.long, device=device),
            torch.tensor(masks, dtype=torch.bool, device=device),
            torch.tensor(labels, dtype=torch.long, device=device))


def train(cfg: AudioTrainConfig) -> 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)

    tr_mel, tr_caps = load_mels(cfg.data_dir, "train")
    va_mel, va_caps = load_mels(cfg.data_dir, "val")
    n_mels, n_frames = tr_mel.shape[1], tr_mel.shape[2]

    probe = build_model(cfg, tokenizer, n_mels, n_frames, 4)
    n_audio_tokens = probe.n_audio_tokens
    model = build_model(cfg, tokenizer, n_mels, n_frames, n_audio_tokens).to(device)

    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(tr_mel), cfg.batch_size)
        mels, ids, mask, labels = make_batch(tokenizer, tr_mel[idx],
                                             [tr_caps[i] for i in idx], n_audio_tokens, device)
        _, loss = model(mels, ids, mask, labels)
        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_audio_tokens": n_audio_tokens, "n_mels": n_mels,
                "n_frames": n_frames}, out_dir / "best.pt")

    acc_tf = teacher_forced_accuracy(model, tokenizer, va_mel, va_caps, n_audio_tokens, device)
    acc_gen, examples = greedy_accuracy(model, tokenizer, va_mel, va_caps,
                                        n_audio_tokens, device)
    summary = {"final_loss": round(loss.item(), 4), "epochs": cfg.epochs,
               "elapsed_s": round(time.time() - t0, 1), "teacher_forced_acc": acc_tf,
               "greedy_acc": acc_gen, "examples": examples}
    (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 teacher_forced_accuracy(model, tokenizer, mel_np, caps, n_audio_tokens, device,
                            n_eval: int = 200) -> float:
    model.eval()
    n = min(n_eval, len(caps))
    _, ids, mask, labels = make_batch(tokenizer, mel_np[:n], caps[:n], n_audio_tokens, device)
    logits, _ = model(torch.from_numpy(mel_np[:n]).unsqueeze(1).to(device), ids, mask, labels)
    sup = labels != -100
    return round((logits.argmax(-1)[sup] == labels[sup]).float().mean().item(), 4)


@torch.no_grad()
def greedy_accuracy(model, tokenizer, mel_np, caps, n_audio_tokens, device,
                    n_eval: int = 100):
    model.eval()
    n = min(n_eval, len(caps))
    prompt = [BOS_ID, *tokenizer.encode(PROMPT)]
    ids0, mask0, _ = build_modality_example(
        AUDIO_START_ID, AUDIO_PAD_ID, AUDIO_END_ID, n_audio_tokens, prompt, [])
    ids0_t = torch.tensor([ids0], device=device)
    mask0_t = torch.tensor([mask0], dtype=torch.bool, device=device)
    correct, examples = 0, []
    for i in range(n):
        mel = torch.from_numpy(mel_np[i:i + 1]).unsqueeze(1).to(device)
        out = model.generate(mel, ids0_t, mask0_t, max_new_tokens=10, eos_id=EOS_ID)
        text = tokenizer.decode(out[0].tolist()[len(ids0):], skip_special_tokens=True).strip()
        correct += int(text == caps[i])
        if i < 6:
            examples.append({"target": caps[i], "pred": text})
    return round(correct / n, 4), examples


def main() -> None:
    parser = argparse.ArgumentParser(description="Entraîner l'encodeur audio (Phase 7).")
    for field, value in asdict(AudioTrainConfig()).items():
        parser.add_argument(f"--{field.replace('_', '-')}", type=type(value), default=value)
    ns = parser.parse_args()
    train(AudioTrainConfig(**{k.replace("-", "_"): v for k, v in vars(ns).items()}))


if __name__ == "__main__":
    main()