"""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()