| |
| """PARI STEP-CHANGE : fine-tune Whisper-large-v3 (seq2seq 1.5B) sur le LINGALA WAXAL (70 h). |
| Motivation : notre champion est un CTC w2v-BERT 2.0 (580M) ; sur benchmark africain comparable |
| le seq2seq Whisper fait -16% relatif de WER vs w2v-BERT CTC. Le lingala (combine 0.346) est LE |
| goulot ; le shona (0.125) est déjà saturé. Whisper connaît nativement 'ln' et produit |
| casse+ponctuation (que le scoreur WAXAL compte). |
| Conforme : entraînement sur le SEUL texte/audio du train WAXAL, fine-tuning autorisé (fil 34204). |
| |
| Gate : devhard-lin (439 clips, locuteurs held-out) — fiable pour les GROS effets. |
| """ |
| import json, os, sys |
| import numpy as np, soundfile as sf, torch |
| from dataclasses import dataclass |
| from torch.utils.data import Dataset |
| from transformers import (WhisperForConditionalGeneration, WhisperProcessor, |
| Seq2SeqTrainer, Seq2SeqTrainingArguments) |
|
|
| BASE = os.environ.get("BASE_MODEL", "openai/whisper-large-v3") |
| LANG = os.environ.get("LANG_ASR", "lin") |
| WLANG = {"lin": "ln", "sna": "sn"}[LANG] |
| OUT = os.environ.get("OUT", "/scratch/runs/whisper_%s" % LANG) |
| EPOCHS = float(os.environ.get("EPOCHS", "3")) |
| BS = int(os.environ.get("BS", "4")) |
| ACC = int(os.environ.get("ACC", "4")) |
| LR = float(os.environ.get("LR", "1e-5")) |
| MAXDUR = float(os.environ.get("MAXDUR", "30")) |
| SR = 16000 |
|
|
|
|
| def load_manifest(path): |
| rows = [] |
| for l in open(path, encoding="utf-8"): |
| r = json.loads(l) |
| if not r.get("text", "").strip(): |
| continue |
| a = r["audio"] |
| if not os.path.exists(a): |
| for cand in ("/scratch/prep/audio/", "/scratch/restore/devhard_audio/", |
| "/root/devhard_audio/"): |
| p = os.path.join(cand, os.path.basename(a)) |
| if os.path.exists(p): |
| r["audio"] = p; break |
| if os.path.exists(r["audio"]) and r.get("duration", 0) <= MAXDUR: |
| rows.append(r) |
| return rows |
|
|
|
|
| class DS(Dataset): |
| def __init__(self, rows, proc): |
| self.rows = rows; self.proc = proc |
|
|
| def __len__(self): |
| return len(self.rows) |
|
|
| def __getitem__(self, i): |
| r = self.rows[i] |
| au = sf.read(r["audio"], dtype="float32")[0] |
| if au.ndim > 1: |
| au = au.mean(1) |
| feats = self.proc.feature_extractor(au, sampling_rate=SR).input_features[0] |
| ids = self.proc.tokenizer(r["text"], max_length=200, truncation=True).input_ids |
| return {"input_features": feats, "labels": ids} |
|
|
|
|
| @dataclass |
| class Collate: |
| proc: object |
|
|
| def __call__(self, batch): |
| feats = torch.tensor(np.stack([b["input_features"] for b in batch])) |
| lab = self.proc.tokenizer.pad([{"input_ids": b["labels"]} for b in batch], |
| return_tensors="pt") |
| labels = lab["input_ids"].masked_fill(lab.attention_mask.ne(1), -100) |
| if (labels[:, 0] == self.proc.tokenizer.convert_tokens_to_ids("<|startoftranscript|>")).all(): |
| labels = labels[:, 1:] |
| return {"input_features": feats, "labels": labels} |
|
|
|
|
| def main(): |
| proc = WhisperProcessor.from_pretrained(BASE, language=WLANG, task="transcribe") |
| train = load_manifest("/root/devhard/train_%s_min.jsonl" % LANG) |
| dev = load_manifest("/root/devhard/devhard_%s.jsonl" % LANG)[:120] |
| print("train %d clips (%.1f h) | eval %d" % ( |
| len(train), sum(r.get("duration", 0) for r in train) / 3600, len(dev)), flush=True) |
| if not train: |
| print("AUCUNE DONNEE — prep_data non termine ?", flush=True); sys.exit(1) |
|
|
| model = WhisperForConditionalGeneration.from_pretrained(BASE, dtype=torch.bfloat16) |
| model.config.forced_decoder_ids = None |
| model.config.suppress_tokens = [] |
| model.generation_config.language = WLANG |
| model.generation_config.task = "transcribe" |
| model.generation_config.forced_decoder_ids = None |
| model.config.use_cache = False |
|
|
| args = Seq2SeqTrainingArguments( |
| output_dir=OUT, per_device_train_batch_size=BS, gradient_accumulation_steps=ACC, |
| learning_rate=LR, warmup_steps=300, num_train_epochs=EPOCHS, |
| gradient_checkpointing=True, bf16=True, optim="adamw_torch_fused", |
| logging_steps=50, save_strategy="epoch", save_total_limit=2, |
| eval_strategy="no", report_to=[], dataloader_num_workers=4, |
| remove_unused_columns=False, lr_scheduler_type="linear") |
| tr = Seq2SeqTrainer(model=model, args=args, train_dataset=DS(train, proc), |
| data_collator=Collate(proc)) |
| tr.train() |
| tr.save_model(OUT + "/final") |
| proc.save_pretrained(OUT + "/final") |
| print("WHISPER_TRAIN_DONE %s/final" % OUT, flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
