scripts/train_mms_adapter_aug.py

#!/usr/bin/env python3
"""LEVER C — MMS-1B adaptateurs JOINT + anti-sur-apprentissage OOD.
Diffs vs train_mms_adapter.py :
  - AUGMENTATION synthetique (audiomentations) sur le TRAIN uniquement : bruit gaussien,
    time-stretch (speed-perturb 0.9-1.1), pitch-shift, gain. => invariance locuteur/canal.
    (100% conforme : transformations de l'audio du challenge, AUCUNE donnee externe.)
  - SpecAugment agressif (mask_time_prob/mask_feature_prob dans la config du modele, sur GPU).
  - eval = dev-difficile (locuteurs DISJOINTS) ; l'eval n'est PAS augmentee (collator separe).
  - lr par defaut 1e-4 (au lieu de 1e-3) : tue le sur-apprentissage des locuteurs vus.
Recette adaptateurs identique : base 1B gelee, adaptateurs frais + tete CTC (~2.3M params).
"""
import argparse
import json
import os
import random
from dataclasses import dataclass

import jiwer
import numpy as np
import soundfile as sf
import torch
from datasets import Dataset
from transformers import (
    Trainer,
    TrainingArguments,
    Wav2Vec2CTCTokenizer,
    Wav2Vec2FeatureExtractor,
    Wav2Vec2ForCTC,
    Wav2Vec2Processor,
)
from transformers.trainer_pt_utils import LengthGroupedSampler
from audiomentations import Compose, AddGaussianNoise, TimeStretch, PitchShift, Gain

SR = 16000
BASE = "facebook/mms-1b-all"
FRAMES_PER_SEC = 50


class LGTrainer(Trainer):
    """LengthGrouped sampler + collator d'eval SANS augmentation."""
    def __init__(self, *args, eval_collator=None, **kwargs):
        super().__init__(*args, **kwargs)
        self._eval_collator = eval_collator

    def _get_train_sampler(self, train_dataset=None):
        ds = train_dataset if train_dataset is not None else self.train_dataset
        return LengthGroupedSampler(
            self.args.per_device_train_batch_size * self.args.gradient_accumulation_steps,
            dataset=ds, lengths=list(ds["length"]))

    def get_eval_dataloader(self, eval_dataset=None):
        # bascule temporairement sur le collator sans aug pour construire le dataloader d'eval
        if self._eval_collator is None:
            return super().get_eval_dataloader(eval_dataset)
        saved = self.data_collator
        self.data_collator = self._eval_collator
        try:
            return super().get_eval_dataloader(eval_dataset)
        finally:
            self.data_collator = saved


def read_jsonl(p):
    rows = []
    for l in open(p, encoding="utf-8"):
        r = json.loads(l)
        r["text"] = " ".join(r["text"].replace("|", " ").split())
        rows.append(r)
    return rows


def text_key(t):
    import unicodedata
    t = unicodedata.normalize("NFC", t).lower()
    return " ".join("".join(c for c in t if c.isalnum() or c.isspace()).split())


def parse():
    p = argparse.ArgumentParser()
    p.add_argument("--lang", default="jointC")
    p.add_argument("--out", required=True)
    p.add_argument("--train", nargs="+", required=True)
    p.add_argument("--eval", required=True)
    p.add_argument("--init", default=BASE)
    p.add_argument("--lr", type=float, default=1e-4)      # anti-surapprentissage
    p.add_argument("--epochs", type=float, default=12)
    p.add_argument("--bs", type=int, default=4)
    p.add_argument("--grad_accum", type=int, default=16)
    p.add_argument("--warmup", type=int, default=500)
    p.add_argument("--eval_steps", type=int, default=300)
    p.add_argument("--eval_subset", type=int, default=1200)
    p.add_argument("--min_dur", type=float, default=1.5)
    p.add_argument("--max_dur", type=float, default=30.0)
    p.add_argument("--max_wps", type=float, default=4.0)
    p.add_argument("--exclude_texts", default=None)
    p.add_argument("--num_workers", type=int, default=8)
    p.add_argument("--save_total_limit", type=int, default=4)
    p.add_argument("--augment", action="store_true")
    p.add_argument("--mask_time_prob", type=float, default=0.075)
    p.add_argument("--mask_feature_prob", type=float, default=0.05)
    p.add_argument("--smoke", action="store_true")
    p.add_argument("--seed", type=int, default=42)
    return p.parse_args()


def main():
    a = parse()
    os.makedirs(a.out, exist_ok=True)
    random.seed(a.seed); np.random.seed(a.seed); torch.manual_seed(a.seed)

    train_rows = []
    for m in a.train:
        rows = read_jsonl(m)
        kept = [r for r in rows if r["text"] and a.min_dur <= r["duration"] <= a.max_dur
                and len(r["text"].split()) / max(r["duration"], 0.1) <= a.max_wps]
        print(f"{m}: {len(kept)}/{len(rows)} gardes ({sum(r['duration'] for r in kept)/3600:.1f}h)")
        train_rows += kept
    if a.exclude_texts:
        banned = set()
        for m in a.exclude_texts.split(","):
            banned |= {text_key(r["text"]) for r in read_jsonl(m) if r["text"]}
        before = len(train_rows)
        train_rows = [r for r in train_rows if text_key(r["text"]) not in banned]
        print(f"anti-fuite: {before-len(train_rows)} exclus")
    eval_rows = [r for r in read_jsonl(a.eval) if r["text"]]
    rng = random.Random(a.seed)
    if a.eval_subset and len(eval_rows) > a.eval_subset:
        eval_rows = rng.sample(eval_rows, a.eval_subset)
    if a.smoke:
        train_rows, eval_rows = train_rows[:96], eval_rows[:24]
        a.epochs, a.eval_steps, a.warmup, a.bs, a.grad_accum = 1, 5, 2, 4, 1
    print(f"TRAIN {len(train_rows)} | EVAL {len(eval_rows)} | augment={a.augment}")

    # vocab char construit sur le train
    vocab_path = os.path.join(a.out, "vocab.json")
    chars = set()
    for r in train_rows:
        chars.update(r["text"])
    chars -= {" ", "|"}
    vocab = {c: i for i, c in enumerate(sorted(chars))}
    vocab["|"] = len(vocab); vocab["[UNK]"] = len(vocab); vocab["[PAD]"] = len(vocab)
    json.dump(vocab, open(vocab_path, "w", encoding="utf-8"), ensure_ascii=False, indent=1)

    tok = Wav2Vec2CTCTokenizer(vocab_path, unk_token="[UNK]", pad_token="[PAD]", word_delimiter_token="|")
    fe = Wav2Vec2FeatureExtractor(feature_size=1, sampling_rate=SR, padding_value=0.0,
                                  do_normalize=True, return_attention_mask=True)
    processor = Wav2Vec2Processor(feature_extractor=fe, tokenizer=tok)
    processor.save_pretrained(a.out)

    def feasible(r):
        ids = tok(r["text"]).input_ids
        need = len(ids) + sum(x == y for x, y in zip(ids, ids[1:]))
        return need <= int(r["duration"] * FRAMES_PER_SEC) - 2
    train_rows = [r for r in train_rows if feasible(r)]

    def to_ds(rows):
        return Dataset.from_list([{"audio": r["audio"], "text": r["text"],
                                   "length": int(r["duration"] * 100)} for r in rows])
    train_ds, eval_ds = to_ds(train_rows), to_ds(eval_rows)

    # ---------- AUGMENTATION synthetique (train only) ----------
    augpipe = None
    if a.augment:
        augpipe = Compose([
            AddGaussianNoise(min_amplitude=0.001, max_amplitude=0.015, p=0.5),
            TimeStretch(min_rate=0.9, max_rate=1.1, leave_length_unchanged=False, p=0.4),
            PitchShift(min_semitones=-2.0, max_semitones=2.0, p=0.25),
            Gain(min_gain_db=-6.0, max_gain_db=6.0, p=0.4),
        ])

    @dataclass
    class Collator:
        augment: object = None
        def __call__(self, feats):
            audio = []
            for f in feats:
                w = sf.read(f["audio"], dtype="float32")[0]
                if self.augment is not None:
                    try:
                        w = self.augment(samples=w, sample_rate=SR)
                    except Exception:
                        pass  # jamais ignorer un vrai echec, mais une aug ratee ne casse pas le batch
                audio.append(w)
            batch = fe(audio, sampling_rate=SR, return_tensors="pt", padding=True)
            enc = tok([f["text"] for f in feats], return_tensors="pt", padding=True)
            batch["labels"] = enc["input_ids"].masked_fill(enc["attention_mask"].ne(1), -100)
            return batch

    train_collator = Collator(augment=augpipe)
    eval_collator = Collator(augment=None)

    # ---------- MODELE MMS a ADAPTATEURS + SpecAugment agressif ----------
    model = Wav2Vec2ForCTC.from_pretrained(
        a.init, vocab_size=len(tok), pad_token_id=tok.pad_token_id,
        ctc_loss_reduction="mean", ctc_zero_infinity=True,
        attention_dropout=0.0, hidden_dropout=0.0, feat_proj_dropout=0.0, layerdrop=0.0,
        apply_spec_augment=True,
        mask_time_prob=a.mask_time_prob, mask_time_length=10,
        mask_feature_prob=a.mask_feature_prob, mask_feature_length=10,
        ignore_mismatched_sizes=True)
    model.init_adapter_layers()
    model.freeze_base_model()
    n_train = 0
    for name, p in model.named_parameters():
        if "adapter" in name.lower():
            p.requires_grad = True
        if p.requires_grad:
            n_train += p.numel()
    print(f"params entrainables: {n_train/1e6:.2f}M (adaptateurs + lm_head) sur ~1B geles", flush=True)

    def preprocess_logits(logits, labels):
        return torch.argmax(logits, dim=-1)

    eval_refs = [r["text"] for r in eval_rows]

    def metrics(pred):
        ids = np.where(pred.predictions == -100, tok.pad_token_id, pred.predictions)
        hyps = tok.batch_decode(ids)
        for _r, _h in list(zip(eval_refs, hyps))[:3]:
            print("  [dbg] REF:", _r[:55], "|| HYP:", repr(_h[:55]), flush=True)
        pairs = [(r, h) for r, h in zip(eval_refs, hyps) if r.strip()]
        refs = [r for r, _ in pairs]; hs = [h for _, h in pairs]
        wer = jiwer.wer(refs, hs); cer = jiwer.cer(refs, hs)
        return {"wer": wer, "cer": cer, "combine": 0.5 * wer + 0.5 * cer}

    targs = TrainingArguments(
        output_dir=a.out, per_device_train_batch_size=a.bs, per_device_eval_batch_size=a.bs,
        gradient_accumulation_steps=a.grad_accum, num_train_epochs=a.epochs, learning_rate=a.lr,
        warmup_steps=a.warmup, bf16=True, eval_strategy="steps", eval_steps=a.eval_steps,
        save_strategy="steps", save_steps=a.eval_steps, save_total_limit=a.save_total_limit,
        load_best_model_at_end=True, metric_for_best_model="combine", greater_is_better=False,
        logging_steps=50, gradient_checkpointing=False,
        gradient_checkpointing_kwargs={"use_reentrant": False},
        dataloader_num_workers=a.num_workers, remove_unused_columns=False, report_to=[],
        seed=a.seed, data_seed=a.seed, ignore_data_skip=True)

    trainer = LGTrainer(model=model, args=targs, train_dataset=train_ds, eval_dataset=eval_ds,
                        data_collator=train_collator, eval_collator=eval_collator,
                        compute_metrics=metrics, preprocess_logits_for_metrics=preprocess_logits,
                        processing_class=processor)
    trainer.train()
    trainer.save_model(os.path.join(a.out, "best"))
    processor.save_pretrained(os.path.join(a.out, "best"))
    print("EVAL FINALE:", json.dumps(trainer.evaluate(), default=float))
    print("TRAIN_DONE", flush=True)


if __name__ == "__main__":
    main()