File size: 5,176 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 | #!/usr/bin/env python3
"""Decodage Phase 2 ROUTE PAR LANGUE (100% modeles pre-entraines, 0 dataset externe).
LID (mms-lid-256, contraint a lin/lug/sna) -> lug=Sunbird Whisper-salt ; lin/sna=notre MMS CTC.
Garde-fous anti-hallucination Whisper -> fallback MMS. Sortie CSV Zindi.
"""
import argparse, csv, glob, os
from collections import Counter
import soundfile as sf, torch
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 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("--audio_dir", default="/root/phase2_audio/audio")
ap.add_argument("--test_csv", default="/root/Test_phase2.csv")
ap.add_argument("--out", required=True)
ap.add_argument("--lang_token", type=int, default=50355) # Luganda dans salt
ap.add_argument("--targets", default="lin,lug,sna")
ap.add_argument("--chunk_s", type=float, default=28.0)
ap.add_argument("--beams", type=int, default=1)
ap.add_argument("--no_repeat", type=int, default=0)
a = ap.parse_args()
dev = "cuda"
lid = AutoModelForAudioClassification.from_pretrained(a.lid_model).to(dev).eval()
lidfe = AutoFeatureExtractor.from_pretrained(a.lid_model)
id2label = lid.config.id2label
label2id = {v: k for k, v in id2label.items()}
tgt = [t for t in a.targets.split(",")]
missing = [t for t in tgt if t not in label2id]
assert not missing, f"LID labels manquants {missing} ; ex labels={list(id2label.values())[: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]]
test_ids = [r["ID"] for r in csv.DictReader(open(a.test_csv, encoding="utf-8"))]
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]
out = []
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, num_beams=a.beams, no_repeat_ngram_size=a.no_repeat)
out.append(wp.batch_decode(ids, skip_special_tokens=True)[0])
return norm(" ".join(out))
out = {}; nlug = 0; nfb = 0; langcount = Counter()
with torch.inference_mode():
for i, tid in enumerate(test_ids):
f = os.path.join(a.audio_dir, tid + ".wav")
if not os.path.exists(f):
cand = glob.glob(os.path.join(a.audio_dir, tid + ".*"))
f = cand[0] if cand else None
if not f:
out[tid] = "a"; continue
au = sf.read(f, dtype="float32")[0]; dur = len(au) / SR
inp = lidfe(au[:SR * 20], sampling_rate=SR, return_tensors="pt")
inp = {k: v.to(dev) for k, v in inp.items()}
logits = lid(**inp).logits[0]
lang = tgt[int(torch.argmax(logits[tgt_ids]).item())]
langcount[lang] += 1
if lang == "lug":
nlug += 1
h = whisper_decode(au)
if is_halluc(h, dur):
h = mms_decode(au); nfb += 1
else:
h = mms_decode(au)
out[tid] = h or "a"
if (i + 1) % 200 == 0:
print(f"{i+1}/{len(test_ids)} langs={dict(langcount)} fb={nfb}", flush=True)
empt = sum(1 for t in test_ids if not out.get(t, "").strip())
with open(a.out, "w", newline="", encoding="utf-8") as fo:
w = csv.writer(fo); w.writerow(["ID", "Target"])
for t in test_ids:
w.writerow([t, out.get(t) or "a"])
print(f"ROUTED_DONE {a.out} | {len(test_ids)} IDs | langs={dict(langcount)} "
f"halluc_fallback={nfb} vides={empt}", flush=True)
if __name__ == "__main__":
main()
|