| |
| """Pseudo-labeling : decode shards unlabeled avec MMS-joint, garde clips haute-confiance -> FLAC + manifest.""" |
| import argparse, glob, os, json, subprocess |
| import pyarrow.parquet as pq |
| import soundfile as sf, numpy as np, torch |
| from transformers import AutoModelForCTC, AutoProcessor |
| SR=16000 |
| def dec16k(raw): |
| p=subprocess.run(["ffmpeg","-v","error","-i","pipe:0","-f","f32le","-ac","1","-ar",str(SR),"pipe:1"],input=raw,capture_output=True) |
| return np.frombuffer(p.stdout,dtype=np.float32) |
| def norm(s): return " ".join(str(s).replace("|"," ").split()) |
| def main(): |
| ap=argparse.ArgumentParser() |
| ap.add_argument("--lang",required=True); ap.add_argument("--parquets",required=True) |
| ap.add_argument("--model",default="/root/models/mmsjoint_best") |
| ap.add_argument("--out_audio",required=True); ap.add_argument("--out_manifest",required=True) |
| ap.add_argument("--min_conf",type=float,default=0.65); ap.add_argument("--max_clips",type=int,default=25000) |
| a=ap.parse_args(); os.makedirs(a.out_audio,exist_ok=True) |
| proc=AutoProcessor.from_pretrained(a.model); blank=proc.tokenizer.pad_token_id |
| m=AutoModelForCTC.from_pretrained(a.model,torch_dtype=torch.bfloat16).cuda().eval() |
| files=sorted(glob.glob(a.parquets)); kept=seen=0 |
| fo=open(a.out_manifest,"w",encoding="utf-8") |
| for pfp in files: |
| for b in pq.ParquetFile(pfp).iter_batches(batch_size=16): |
| auds=[];metas=[] |
| for r in b.to_pylist(): |
| au=r.get("audio"); raw=au.get("bytes") if isinstance(au,dict) else None |
| if raw is None: continue |
| try: wav=dec16k(raw) |
| except: continue |
| dur=len(wav)/SR |
| if not(1.5<=dur<=30): continue |
| seen+=1; auds.append(wav); metas.append((str(r.get("id") or f"{a.lang}_ul_{seen}"),dur)) |
| if not auds: continue |
| with torch.inference_mode(): |
| f=proc(auds,sampling_rate=SR,return_tensors="pt",padding=True) |
| f={k:(v.to("cuda",dtype=torch.bfloat16) if v.dtype==torch.float32 else v.to("cuda")) for k,v in f.items()} |
| lg=m(**f).logits.float(); pr=lg.softmax(-1); mp,ids=pr.max(-1) |
| txts=proc.batch_decode(ids.cpu().numpy()); idsc=ids.cpu() |
| for i,(rid,dur) in enumerate(metas): |
| nb=idsc[i]!=blank |
| conf=float(mp[i][nb].mean()) if nb.any() else 0.0 |
| txt=norm(txts[i]) |
| if conf<a.min_conf or len(txt.split())<3: continue |
| p=os.path.join(a.out_audio,rid.replace("/","_")+".flac"); sf.write(p,auds[i],SR) |
| fo.write(json.dumps({"id":rid,"audio":p,"text":txt,"duration":round(dur,3)},ensure_ascii=False)+"\n"); kept+=1 |
| if kept>=a.max_clips: break |
| if kept>=a.max_clips: break |
| fo.close(); print(f"PSEUDO {a.lang}: kept {kept}/{seen}",flush=True) |
| if __name__=="__main__": main() |
|
|
