File size: 2,953 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 | #!/usr/bin/env python3
"""Juge final ground-truth-free : décode N clips avec les 3 modèles w2v-BERT + affiche
les 2 prédictions LID. La sortie COHÉRENTE (vrais mots) révèle la vraie langue."""
import glob, json
import soundfile as sf, torch
from transformers import (AutoModelForCTC, AutoProcessor,
Wav2Vec2ForSequenceClassification, AutoFeatureExtractor,
Wav2Vec2BertForSequenceClassification, SeamlessM4TFeatureExtractor)
SR=16000; N=10
files=sorted(glob.glob("/root/phase2_audio/audio/*.wav"))[:N]
def norm(s): return " ".join(str(s).replace("|"," ").split())
# LID interne
fe2=SeamlessM4TFeatureExtractor.from_pretrained("/root/models/lid_best")
lid=Wav2Vec2BertForSequenceClassification.from_pretrained("/root/models/lid_best",torch_dtype=torch.bfloat16).cuda().eval()
i2l=json.load(open("/root/models/lid_best/config.json"))["id2label"]
our={}
with torch.inference_mode():
for f in files:
w=sf.read(f,dtype="float32")[0][:20*SR]
x=fe2([w],sampling_rate=SR,return_tensors="pt",padding=True)
x={k:v.to("cuda",dtype=torch.bfloat16 if v.dtype==torch.float32 else v.dtype) for k,v in x.items()}
our[f]=i2l[str(int(lid(**x).logits.float().argmax(-1)))]
del lid; torch.cuda.empty_cache()
# MMS-LID
fe3=AutoFeatureExtractor.from_pretrained("facebook/mms-lid-256")
mms=Wav2Vec2ForSequenceClassification.from_pretrained("facebook/mms-lid-256",torch_dtype=torch.bfloat16).cuda().eval()
ml2i={str(v):int(k) for k,v in mms.config.id2label.items()}
idx=torch.tensor([ml2i["lin"],ml2i["lug"],ml2i["sna"]]).cuda(); order=["lin","lug","sna"]
mmsp={}
with torch.inference_mode():
for f in files:
w=sf.read(f,dtype="float32")[0][:20*SR]
x=fe3([w],sampling_rate=SR,return_tensors="pt")
x={k:v.to("cuda",dtype=torch.bfloat16 if v.dtype==torch.float32 else v.dtype) for k,v in x.items()}
mmsp[f]=order[int(mms(**x).logits.float()[0][idx].argmax())]
del mms; torch.cuda.empty_cache()
# 3 transcriptions
CTC={"lin":"/root/models/lin_s4_best","sna":"/root/models/sna_ps_best","lug":"/root/models/lug_ps_best"}
hyp={f:{} for f in files}
for lang,path in CTC.items():
proc=AutoProcessor.from_pretrained(path); m=AutoModelForCTC.from_pretrained(path,torch_dtype=torch.bfloat16).cuda().eval()
with torch.inference_mode():
for f in files:
w=sf.read(f,dtype="float32")[0]
x=proc(w,sampling_rate=SR,return_tensors="pt",padding=True)
x={k:v.to("cuda",dtype=torch.bfloat16 if v.dtype==torch.float32 else v.dtype) for k,v in x.items()}
ids=m(**x).logits.float().argmax(-1).cpu().numpy()
hyp[f][lang]=norm(proc.batch_decode(ids)[0])
del m; torch.cuda.empty_cache()
for f in files:
print(f"\n### {f.split('/')[-1]} | LID_interne={our[f]} MMS-LID={mmsp[f]}")
for lang in ["lin","sna","lug"]:
print(f" [{lang}] {hyp[f][lang][:90]}")
print("\nDECODE_DONE")
|