#!/usr/bin/env python3 """Évalue le LM NEURONAL (char-level) comme terme de rescoring sur devhard-lin. Config de décodage = celle du RECORD LB (alpha=0.5, beta=1.0, lsb=True), pas celle du sweep qui perd au LB. score(h) = ac_cont + lm_kenlm + w2*ac_cont2 + mu*charlm(h) + gamma*nmots """ import json, os, pickle import jiwer, numpy as np, soundfile as sf, torch, torch.nn as nn from multiprocessing import Pool from pyctcdecode import build_ctcdecoder from transformers import AutoModelForCTC, AutoProcessor M1 = "/root/models/joint_cont_best" ARPA = "/scratch/lm/lin_5g.arpa" CLM = "/scratch/lm/charlm_lin.pt" R = "/scratch/restore" NBEST = int(os.environ.get("NBEST", "10")) ALPHA = float(os.environ.get("ALPHA", "0.5")) BETA = float(os.environ.get("BETA", "1.0")) LSB = os.environ.get("LSB", "1") == "1" 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 class LM(nn.Module): def __init__(self, V, DIM, NLAYER, MAXLEN=256): super().__init__() self.emb = nn.Embedding(V, DIM, padding_idx=0) self.pos = nn.Embedding(MAXLEN, DIM) layer = nn.TransformerEncoderLayer(DIM, 4, DIM * 4, dropout=0.1, batch_first=True, norm_first=True) self.tr = nn.TransformerEncoder(layer, NLAYER) self.ln = nn.LayerNorm(DIM); self.head = nn.Linear(DIM, V) def forward(self, x): T = x.shape[1] h = self.emb(x) + self.pos(torch.arange(T, device=x.device))[None] mask = nn.Transformer.generate_square_subsequent_mask(T, device=x.device) return self.head(self.ln(self.tr(h, mask=mask, is_causal=True))) 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"] == "lin"] for r in sub: 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-lin : %d clips | decode alpha=%.2f beta=%.2f lsb=%s" % (len(sub), ALPHA, BETA, LSB), flush=True) CACHE = "/scratch/lm/logits_lin.pkl" if os.path.exists(CACHE): L1 = pickle.load(open(CACHE, "rb")) else: _, L1 = compute_logits(M1, sub) pickle.dump(L1, open(CACHE, "wb")) 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] def cc(h, g): return (g[:1] + h[1:]) if (h and g) else h dec = build_ctcdecoder(lab, kenlm_model_path=ARPA, alpha=ALPHA, beta=BETA, lm_score_boundary=LSB) 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)] _, _, REF = comb(refs, [cc(h, g) for h, g in zip(db, greedy)]) print("REFERENCE (config du record) : %.4f" % REF, flush=True) cands, AC1, LMS, 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]] l = [((b[4] - b[3]) if len(b) > 4 else 0.0) for b in bs[:NBEST]] if db[i] not in c: c.append(db[i]) a.append(ctc_score(L1[i], encode_for(tok, db[i]), tok.pad_token_id)) l.append(float(np.mean(l)) if l else 0.0) cands.append(c); AC1.append(np.array(a)); LMS.append(np.array(l)) NW.append(np.array([float(len(x.split())) for x in c])) orc = [min(cands[i], key=lambda h: comb([refs[i]], [h])[2] if refs[i].strip() else 0) for i in range(len(cands))] print("ORACLE %d-best : %.4f (marge %+.4f)" % (NBEST, comb(refs, orc)[2], comb(refs, orc)[2] - REF), flush=True) # --- rescoreur acoustique cont2 --- SC = {} for tag, mdl in [("cont2", R + "/joint_cont2_best")]: if os.path.isdir(mdl): 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("%s OK" % tag, flush=True) # --- LM neuronal --- ck = torch.load(CLM, map_location="cuda") stoi = ck["stoi"] net = LM(ck["V"], ck["dim"], ck["nlayer"]).cuda().eval() net.load_state_dict(ck["state"]) def clm_scores(texts): seqs = [[1] + [stoi[c] for c in t if c in stoi][:254] + [1] for t in texts] Lm = max(len(s) for s in seqs) x = torch.zeros(len(seqs), Lm, dtype=torch.long) for i, s in enumerate(seqs): x[i, :len(s)] = torch.tensor(s) x = x.cuda() with torch.inference_mode(): lg = net(x[:, :-1]).log_softmax(-1) tgt = x[:, 1:] gl = lg.gather(-1, tgt.unsqueeze(-1)).squeeze(-1) gl = gl * (tgt != 0) return gl.sum(1).float().cpu().numpy() CL = [clm_scores(c) for c in cands] print("charLM scores OK", flush=True) def evaluate(w2=0.0, mu=0.0, gamma=0.0): hyps = [] for i in range(len(cands)): tot = AC1[i] + LMS[i] + gamma * NW[i] + mu * CL[i] if w2 and "cont2" in SC: tot = tot + w2 * SC["cont2"][i] hyps.append(cc(cands[i][int(np.argmax(tot))], greedy[i])) return comb(refs, hyps)[2] print("\n--- charLM SEUL (mu) ---", flush=True) bestmu = (9.0, 0.0) for mu in (0.05, 0.1, 0.2, 0.3, 0.5, 0.8, 1.2): m = evaluate(0.0, mu, 0.0) if m < bestmu[0]: bestmu = (m, mu) print(" mu=%.2f : %.4f (%+.4f)" % (mu, m, m - REF), flush=True) print("\n--- cont2 seul (w2) ---", flush=True) bestw = (9.0, 0.0) for w2 in (0.5, 1.0, 1.5, 2.5): m = evaluate(w2, 0.0, 0.0) if m < bestw[0]: bestw = (m, w2) print(" w2=%.1f : %.4f (%+.4f)" % (w2, m, m - REF), flush=True) print("\n--- cont2 + charLM ---", flush=True) best = (9.0, None, None) for w2 in (0.0, 1.0, 1.5, 2.5): for mu in (0.0, 0.05, 0.1, 0.2, 0.3, 0.5): m = evaluate(w2, mu, 0.0) if m < best[0]: best = (m, w2, mu) print(" BEST combo : %.4f (cont2=%s mu=%s) %+.4f" % (best[0], best[1], best[2], best[0] - REF), flush=True) json.dump({"ref": REF, "best": best[0], "w2": best[1], "mu": best[2], "charlm_solo": bestmu, "cont2_solo": bestw}, open("/root/charlm_rescore.json", "w")) print("CHARLM_RESCORE_DONE", flush=True) if __name__ == "__main__": main()