File size: 4,517 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 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 | #!/usr/bin/env python3
"""RESCORING N-BEST v2 : sweep etendu de LAMBDA + l'hypothese decode_batch injectee dans le
pool de candidats (decode_beams renvoie un 1-best different, souvent moins bon).
Ajoute aussi joint_best comme 3e rescoreur optionnel (MU).
score(h) = ac_cont(h) + lm_kenlm(h) + LAMBDA*ac_cont2(h) + MU*ac_jbest(h)
Les scores acoustiques manquants (hypothese venue de decode_batch) sont calcules par CTC.
"""
import json
import pickle
import jiwer
import numpy as np
import torch
from multiprocessing import Pool
from pyctcdecode import build_ctcdecoder
from transformers import AutoProcessor
M1 = "/root/models/joint_cont_best"
ARPA = "/scratch/lm/lin_5g.arpa"
NBEST = 10
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 ctc_scores(logp, texts, tok):
T = logp.shape[0]
lp = torch.from_numpy(logp).unsqueeze(1)
out = []
for t in texts:
ids = [i for i in tok(t.replace(" ", "|")).input_ids if i != tok.pad_token_id] if t else []
if not ids or len(ids) > T:
out.append(-1e9)
continue
loss = torch.nn.functional.ctc_loss(
lp, torch.tensor(ids).unsqueeze(0), torch.tensor([T]), torch.tensor([len(ids)]),
blank=tok.pad_token_id, reduction="sum", zero_infinity=True)
out.append(-float(loss))
return 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"] == "lin"]
refs = [r["text"] for r in sub]
L1 = pickle.load(open("/scratch/lm/logits_lin.pkl", "rb"))
L2 = pickle.load(open("/scratch/lm/logits_lin_cont2.pkl", "rb"))
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]
dec = build_ctcdecoder(lab, kenlm_model_path=ARPA, alpha=0.5, beta=0.5,
lm_score_boundary=False)
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)]
def cc(h, g):
return (g[:1] + h[1:]) if (h and g) else h
_, _, m_db = comb(refs, [cc(h, g) for h, g in zip(db, greedy)])
_, _, m_bm = comb(refs, [cc(bs[0][0], g) for bs, g in zip(allbeams, greedy)])
print("reference decode_batch : %.4f" % m_db, flush=True)
print("reference decode_beams : %.4f" % m_bm, flush=True)
REF = min(m_db, m_bm)
prepared = []
for i, bs in enumerate(allbeams):
cands = [" ".join(b[0].split()) for b in bs[:NBEST]]
ac1 = [(b[3] if len(b) > 3 else 0.0) for b in bs[:NBEST]]
lm = [((b[4] - b[3]) if len(b) > 4 else 0.0) for b in bs[:NBEST]]
# injecte l'hypothese decode_batch si absente (scores calcules par CTC + LM du beam le plus proche)
if db[i] not in cands:
cands.append(db[i])
ac1.append(ctc_scores(L1[i], [db[i]], tok)[0])
lm.append(float(np.mean(lm)) if lm else 0.0)
ac2 = ctc_scores(L2[i], cands, tok)
prepared.append((cands, np.array(ac1), np.array(lm), np.array(ac2)))
if (i + 1) % 200 == 0:
print(" prepare %d/%d" % (i + 1, len(allbeams)), flush=True)
print("\n--- sweep LAMBDA etendu ---", flush=True)
best = (9.0, None)
for lam in (0.0, 0.3, 0.5, 1.0, 1.5, 2.0, 3.0, 5.0, 8.0, 12.0):
hyps = []
for (cands, ac1, lm, ac2), g in zip(prepared, greedy):
tot = ac1 + lm + lam * ac2
hyps.append(cc(cands[int(np.argmax(tot))], g))
_, _, m = comb(refs, hyps)
mark = " <<<" if m < best[0] else ""
print(" lambda=%5.1f : %.4f (vs ref %.4f : %+.4f)%s" % (lam, m, REF, m - REF, mark), flush=True)
if m < best[0]:
best = (m, lam)
print("\nBEST_V2 %.4f lambda=%.1f (ref %.4f, gain %+.4f)"
% (best[0], best[1], REF, best[0] - REF), flush=True)
json.dump({"combine": best[0], "lambda": best[1], "ref": REF},
open("/root/nbest_v2.json", "w"))
print("NBEST_V2_DONE", flush=True)
if __name__ == "__main__":
main()
|