File size: 5,535 Bytes
6eed659 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 | #!/usr/bin/env python3
"""GATE du pipeline route sur le dev-difficile (locuteurs disjoints, refs+langue connues).
Mesure : (1) precision LID mms-lid-256 contrainte a lin/lug/sna, (2) score MMS-seul par langue,
(3) score ROUTE (lug-predit->Whisper sinon MMS) par VRAIE langue + MACRO, (4) Whisper pur sur lug.
=> decide si router lug->Whisper bat le MMS-seul AVANT de decoder la Phase 2.
"""
import argparse, json
from collections import Counter, defaultdict
import soundfile as sf, torch, jiwer
from transformers import (AutoModelForCTC, AutoProcessor,
WhisperForConditionalGeneration, WhisperProcessor,
AutoModelForAudioClassification, AutoFeatureExtractor)
SR = 16000
def norm(s):
return " ".join(str(s).replace("|", " ").split())
def is_halluc(hyp, dur):
w = hyp.split()
if not w:
return True
if len(w) >= 6 and Counter(w).most_common(1)[0][1] / len(w) > 0.5:
return True
if dur > 0 and len(w) / dur > 6.0:
return True
return False
def sc(refs, hyps):
P = [(r, h) for r, h in zip(refs, hyps) if r.strip()]
if not P:
return (0, 0, 0)
R = [r for r, _ in P]; H = [h for _, h in P]
wer = jiwer.wer(R, H); cer = jiwer.cer(R, H)
return wer, cer, 0.5 * wer + 0.5 * cer
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--whisper_model", default="Sunbird/asr-whisper-large-v3-salt")
ap.add_argument("--mms_model", required=True)
ap.add_argument("--lid_model", default="facebook/mms-lid-256")
ap.add_argument("--manifest", default="/root/devhard/devhard_all.jsonl")
ap.add_argument("--lang_token", type=int, default=50355)
ap.add_argument("--targets", default="lin,lug,sna")
ap.add_argument("--chunk_s", type=float, default=28.0)
a = ap.parse_args()
dev = "cuda"
lid = AutoModelForAudioClassification.from_pretrained(a.lid_model).to(dev).eval()
lidfe = AutoFeatureExtractor.from_pretrained(a.lid_model)
label2id = {v: k for k, v in lid.config.id2label.items()}
tgt = a.targets.split(",")
missing = [t for t in tgt if t not in label2id]
assert not missing, f"LID labels manquants {missing} ; ex={list(label2id)[:8]}"
tgt_ids = torch.tensor([label2id[t] for t in tgt])
mms = AutoModelForCTC.from_pretrained(a.mms_model, torch_dtype=torch.float16).to(dev).eval()
mmsp = AutoProcessor.from_pretrained(a.mms_model)
wp = WhisperProcessor.from_pretrained(a.whisper_model)
wm = WhisperForConditionalGeneration.from_pretrained(a.whisper_model, torch_dtype=torch.float16).to(dev).eval()
tr = wp.tokenizer.convert_tokens_to_ids("<|transcribe|>")
nt = wp.tokenizer.convert_tokens_to_ids("<|notimestamps|>")
forced = [[1, a.lang_token], [2, tr], [3, nt]]
def mms_decode(au):
x = mmsp(au, sampling_rate=SR, return_tensors="pt")
x = {k: (v.to(dev, torch.float16) if v.dtype == torch.float32 else v.to(dev)) for k, v in x.items()}
ids = mms(**x).logits.argmax(-1).cpu().numpy()
return norm(mmsp.batch_decode(ids)[0])
def whisper_decode(au):
win = int(a.chunk_s * SR)
pieces = [au[j:j + win] for j in range(0, len(au), win)] or [au]
o = []
for p in pieces:
f = wp(p, sampling_rate=SR, return_tensors="pt").input_features.to(dev, torch.float16)
ids = wm.generate(f, forced_decoder_ids=forced, max_new_tokens=220)
o.append(wp.batch_decode(ids, skip_special_tokens=True)[0])
return norm(" ".join(o))
rows = [json.loads(l) for l in open(a.manifest, encoding="utf-8")]
lid_ok = 0; conf = defaultdict(Counter)
by = defaultdict(lambda: {"ref": [], "mms": [], "route": [], "wpur": []})
with torch.inference_mode():
for i, r in enumerate(rows):
au = sf.read(r["audio"], dtype="float32")[0]; dur = len(au) / SR
true = r["lang"]; ref = norm(r["text"])
inp = lidfe(au[:SR * 20], sampling_rate=SR, return_tensors="pt")
inp = {k: v.to(dev) for k, v in inp.items()}
pred = tgt[int(torch.argmax(lid(**inp).logits[0][tgt_ids]).item())]
conf[true][pred] += 1
if pred == true:
lid_ok += 1
h_mms = mms_decode(au)
# route selon LID
if pred == "lug":
h_w = whisper_decode(au)
h_route = h_mms if is_halluc(h_w, dur) else h_w
else:
h_route = h_mms
by[true]["ref"].append(ref)
by[true]["mms"].append(h_mms)
by[true]["route"].append(h_route)
if true == "lug":
by[true]["wpur"].append(whisper_decode(au))
if (i + 1) % 100 == 0:
print(f"{i+1}/{len(rows)}", flush=True)
print(f"\nLID accuracy = {lid_ok/len(rows):.4f} ({lid_ok}/{len(rows)})")
for t in tgt:
print(f" LID vrai={t}: {dict(conf[t])}")
mms_macro = []; route_macro = []
for t in tgt:
d = by[t]
_, _, cm = sc(d["ref"], d["mms"]); _, _, cr = sc(d["ref"], d["route"])
mms_macro.append(1 - cm); route_macro.append(1 - cr)
line = f"{t}: MMS score={1-cm:.4f} | ROUTE score={1-cr:.4f}"
if t == "lug" and d["wpur"]:
_, _, cw = sc(d["ref"], d["wpur"])
line += f" | Whisper-pur(lug)={1-cw:.4f}"
print(line)
print(f"MACRO MMS-seul={sum(mms_macro)/len(mms_macro):.4f} | ROUTE={sum(route_macro)/len(route_macro):.4f}")
if __name__ == "__main__":
main()
|