#!/usr/bin/env python3 """RESCORING N-BEST SUR LE SHONA — jamais testé (le shona = la moitié du test). Le KenLM DÉGRADE le shona (il déforme), mais le rescoring ne fait que CHOISIR parmi des hypothèses déjà produites par sna_ps : mécanisme différent, non invalidé. score(h) = ac_snaps(h) + somme_i w_i*ac_i(h) + gamma*nb_mots(h) Gate = devhard-sna (433 clips), FIABLE pour le shona d'après §7 de l'AUTOPILOT. Baseline à battre : sna_ps greedy = 0.1281 (combine). """ import json, os, pickle import jiwer, numpy as np, soundfile as sf, torch from multiprocessing import Pool from pyctcdecode import build_ctcdecoder from transformers import AutoModelForCTC, AutoProcessor M1 = "/root/models/sna_ps_best" # modèle principal shona NBEST = 10 R = "/scratch/restore" RESCORERS = [ ("cont2", R + "/joint_cont2_best"), ("cont", R + "/joint_cont_best"), ("sna_r2", R + "/sna_r2_best"), ] AUD = "/root/devhard_audio" def comb(refs, hyps): pr = [(r, h) for r, h in zip(refs, hyps) if r.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 def encode_for(tok, text): v = tok.get_vocab() delim = getattr(tok, "word_delimiter_token", "|") s = text.replace(" ", delim) keep = "".join(c for c in s if c in v) if not keep: keep = "".join(c for c in text.lower().replace(" ", delim) if c in v) return [v[c] for c in keep if v[c] != tok.pad_token_id] def ctc_score(logp, ids, blank): T = logp.shape[0] if not ids or len(ids) > T: return -1e9 lp = torch.from_numpy(logp).unsqueeze(1) loss = torch.nn.functional.ctc_loss( lp, torch.tensor(ids).unsqueeze(0), torch.tensor([T]), torch.tensor([len(ids)]), blank=blank, reduction="sum", zero_infinity=True) return -float(loss) def compute_logits(model_dir, rows): proc = AutoProcessor.from_pretrained(model_dir) m = AutoModelForCTC.from_pretrained(model_dir, dtype=torch.float32).cuda().eval() out = [] with torch.inference_mode(): for i in range(0, len(rows), 4): b = rows[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 = m(**x).logits.log_softmax(-1).float().cpu().numpy() for j in range(len(b)): out.append(lg[j]) del m; torch.cuda.empty_cache() return proc, out def main(): rows = [json.loads(l) for l in open("/root/devhard/devhard_linsna.jsonl", encoding="utf-8")] sub = [r for r in rows if r["lang"] == "sna"] for r in sub: # réécrire les chemins vers /root r["audio"] = os.path.join(AUD, os.path.basename(r["audio"])) sub = [r for r in sub if os.path.exists(r["audio"])] refs = [r["text"] for r in sub] print("devhard-sna : %d clips" % len(sub), flush=True) # --- logits du modèle principal (cache) --- CACHE = "/scratch/lm/logits_sna.pkl" os.makedirs("/scratch/lm", exist_ok=True) if os.path.exists(CACHE): L1 = pickle.load(open(CACHE, "rb")) else: _, L1 = compute_logits(M1, sub) pickle.dump(L1, open(CACHE, "wb")) print("logits sna_ps OK", flush=True) tok = AutoProcessor.from_pretrained(M1).tokenizer 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] = "" greedy = [" ".join(tok.decode(l.argmax(-1)).replace("|", " ").split()) for l in L1] _, _, REF = comb(refs, greedy) print("BASELINE greedy sna_ps : %.4f" % REF, flush=True) # --- N-best par beam CTC PUR (aucun LM : il dégrade le shona) --- dec = build_ctcdecoder(lab) # pas de kenlm_model_path with Pool(8) as p: allbeams = dec.decode_beams_batch(p, L1, beam_width=64) with Pool(8) as p: db = [" ".join(x.split()) for x in dec.decode_batch(p, L1, beam_width=64)] _, _, BEAM = comb(refs, db) print("beam CTC pur (1-best) : %.4f (%+.4f vs greedy)" % (BEAM, BEAM - REF), flush=True) cands, AC1, NW = [], [], [] for i, bs in enumerate(allbeams): c = [" ".join(b[0].split()) for b in bs[:NBEST]] a = [(b[3] if len(b) > 3 else 0.0) for b in bs[:NBEST]] for extra in (db[i], greedy[i]): # injecter decode_batch ET le greedy if extra and extra not in c: c.append(extra) a.append(ctc_score(L1[i], encode_for(tok, extra), tok.pad_token_id)) cands.append(c); AC1.append(np.array(a)) NW.append(np.array([float(len(x.split())) for x in c])) # --- ORACLE : plafond atteignable par simple sélection --- orc = [] for i in range(len(cands)): best = min(cands[i], key=lambda h: comb([refs[i]], [h])[2] if refs[i].strip() else 0) orc.append(best) _, _, ORACLE = comb(refs, orc) print("ORACLE %d-best : %.4f (marge %+.4f)" % (NBEST, ORACLE, ORACLE - REF), flush=True) # --- scores des rescoreurs --- SC = {} for tag, mdl in RESCORERS: if not os.path.isdir(mdl): print("%-8s ABSENT %s" % (tag, mdl), flush=True); continue proc, LG = compute_logits(mdl, sub) t2 = proc.tokenizer SC[tag] = [np.array([ctc_score(LG[i], encode_for(t2, x), t2.pad_token_id) for x in cands[i]]) for i in range(len(cands))] print("%-8s scores OK (|V|=%d)" % (tag, len(t2.get_vocab())), flush=True) def evaluate(W, gamma=0.0): hyps = [] for i in range(len(cands)): tot = AC1[i] + gamma * NW[i] for t, w in W.items(): if w: tot = tot + w * SC[t][i] hyps.append(cands[i][int(np.argmax(tot))]) return comb(refs, hyps)[2] print("\n--- (a) rescoreurs solo (ref greedy %.4f) ---" % REF, flush=True) solo = {} for tag in SC: bb = (9.0, 0.0) for w in (0.3, 0.5, 1.0, 1.5, 2.5, 4.0): m = evaluate({tag: w}) if m < bb[0]: bb = (m, w) solo[tag] = bb print(" %-8s %.4f (w=%.1f) %+.4f" % (tag, bb[0], bb[1], bb[0] - REF), flush=True) print("\n--- (b) terme de longueur seul ---", flush=True) for gm in (-2.0, -1.0, 0.0, 1.0, 2.0, 4.0): m = evaluate({}, gm) print(" gamma=%+5.1f : %.4f (%+.4f)" % (gm, m, m - REF), flush=True) print("\n--- (c) meilleur rescoreur + longueur ---", flush=True) best = (9.0, None, None, None) if solo: btag = min(solo, key=lambda t: solo[t][0]) for w in (0.5, 1.0, 1.5, 2.5): for gm in (0.0, 1.0, 2.0, 4.0): m = evaluate({btag: w}, gm) if m < best[0]: best = (m, btag, w, gm) print(" %s=%.1f gamma=%+5.1f : %.4f (%+.4f)" % (btag, w, gm, m, m - REF), flush=True) print("\nBEST_SNA %.4f (%s w=%s gamma=%s) baseline %.4f gain %+.4f" % (best[0], best[1], best[2], best[3], REF, best[0] - REF), flush=True) json.dump({"best": best[0], "tag": best[1], "w": best[2], "gamma": best[3], "ref": REF, "oracle": ORACLE}, open("/root/sna_rescore.json", "w")) print("SNA_RESCORE_DONE", flush=True) if __name__ == "__main__": main()