File size: 2,948 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 | #!/usr/bin/env python3
"""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()
|