| """Beam + KenLM (LM = texte WAXAL train seulement). Logits mis en cache -> sweep rapide.""" |
| import json, os, pickle, numpy as np, soundfile as sf, torch, jiwer |
| from transformers import AutoModelForCTC, AutoProcessor |
| from pyctcdecode import build_ctcdecoder |
| from multiprocessing import Pool |
|
|
| LANG=os.environ.get("LG","lin"); MDL=os.environ.get("MDL","/root/models/joint_cont_best") |
| ROWS=[json.loads(l) for l in open("/root/devhard/devhard_linsna.jsonl")] |
| SUB=[r for r in ROWS if r["lang"]==LANG]; REFS=[r["text"] for r in SUB] |
| LMDIR="/scratch/lm"; CACHE=f"/scratch/lm/logits_{LANG}_{os.path.basename(MDL)}.pkl" |
|
|
| def comb(a,b): |
| pr=[(x,y) for x,y in zip(a,b) if x.strip()] |
| A=[x for x,_ in pr]; B=[y for _,y in pr] |
| w=jiwer.wer(A,B); c=jiwer.cer(A,B); return w,c,0.5*w+0.5*c |
|
|
| proc=AutoProcessor.from_pretrained(MDL); tok=proc.tokenizer |
| if os.path.exists(CACHE): |
| LOG=pickle.load(open(CACHE,"rb")); print("logits (cache):",len(LOG),flush=True) |
| else: |
| model=AutoModelForCTC.from_pretrained(MDL,dtype=torch.float32).cuda().eval() |
| LOG=[] |
| with torch.inference_mode(): |
| for i in range(0,len(SUB),4): |
| b=SUB[i:i+4] |
| au=[sf.read(r["audio"],dtype="float32")[0] for r in b] |
| x=proc(au,sampling_rate=16000,return_tensors="pt",padding=True) |
| x={k:v.cuda() for k,v in x.items()} |
| lg=model(**x).logits.log_softmax(-1).cpu().numpy().astype(np.float32) |
| n=[len(a) for a in au] |
| for j in range(len(b)): LOG.append(lg[j]) |
| del model; torch.cuda.empty_cache() |
| pickle.dump(LOG,open(CACHE,"wb")); print("logits calcules:",len(LOG),flush=True) |
|
|
| greedy=[" ".join(tok.decode(l.argmax(-1)).replace("|"," ").split()) for l in LOG] |
| w,c,gm=comb(REFS,greedy); print(f"GREEDY {LANG}: combine={gm:.4f} (WER {w:.4f} CER {c:.4f})",flush=True) |
|
|
| |
| v=tok.get_vocab(); lab=[None]*len(v) |
| for t,i in v.items(): lab[i]=t |
| lab[tok.word_delimiter_token_id]=" " |
| lab[tok.unk_token_id]="⁇" |
| lab[tok.pad_token_id]="" |
| assert len(lab)==len(set(lab)), "doublons restants" |
| print("alphabet OK, taille",len(lab),flush=True) |
|
|
| best=(gm,"greedy",None) |
| for o in [3,4,5]: |
| arpa=f"{LMDIR}/{LANG}_{o}g.arpa" |
| if not os.path.exists(arpa): continue |
| for alpha in [0.3,0.5,0.8]: |
| for beta in [0.0,1.0]: |
| dec=build_ctcdecoder(lab,kenlm_model_path=arpa,alpha=alpha,beta=beta) |
| with Pool(8) as p: |
| hyps=dec.decode_batch(p,LOG,beam_width=64) |
| hyps=[" ".join(h.split()) for h in hyps] |
| w,c,mm=comb(REFS,hyps) |
| flag=" <<<" if mm<best[0] else "" |
| print(f" o={o} a={alpha} b={beta}: combine={mm:.4f} (WER {w:.4f} CER {c:.4f}){flag}",flush=True) |
| if mm<best[0]: best=(mm,f"o{o}_a{alpha}_b{beta}",(o,alpha,beta)) |
| print(f"BEST {LANG}: {best[0]:.4f} <- {best[1]} (greedy {gm:.4f}, gain {gm-best[0]:+.4f})",flush=True) |
| json.dump({"best":best[0],"cfg":best[1],"params":best[2],"greedy":gm},open(f"/root/kenlm_best_{LANG}.json","w")) |
| print("KENLM_SWEEP2_DONE",flush=True) |
|
|
