File size: 4,946 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 117 118 119 120 121 122 123 124 125 126 127 128 129 130 | #!/usr/bin/env python3
"""Generateur de soumission parametrable : lin = <modele> + beam KenLM, sna = sna_ps greedy.
Casse du 1er caractere toujours copiee du greedy (gain valide).
Usage : MODEL=... ARPA=... ALPHA=.. BETA=.. LSB=0|1 OUT=... python gen_sub.py
"""
import csv
import glob
import json
import os
import soundfile as sf
import torch
from multiprocessing import Pool
from pyctcdecode import build_ctcdecoder
from transformers import AutoModelForCTC, AutoProcessor
SR = 16000
CACHE = "/scratch/p2_16k"
LINM = os.environ.get("MODEL", "/root/models/joint_cont_best")
ARPA = os.environ.get("ARPA", "/scratch/lm/lin_5g.arpa")
ALPHA = float(os.environ.get("ALPHA", "0.5"))
BETA = float(os.environ.get("BETA", "0.5"))
LSB = os.environ.get("LSB", "0") == "1"
OUT = os.environ.get("OUT", "/root/sub_gen.csv")
def norm(s):
return " ".join(str(s).replace("|", " ").split())
def batches(sel, budget):
durs = {f: sf.info(f).duration for f in sel}
sel = sorted(sel, key=lambda f: -durs[f])
bs, cur, acc = [], [], 0.0
for f in sel:
if cur and acc + durs[f] > budget:
bs.append(cur)
cur, acc = [], 0.0
cur.append(f)
acc += durs[f]
if cur:
bs.append(cur)
return bs
def main():
lang = json.load(open(os.environ.get("LANGF", "/root/test_lang.json")))
files = sorted(glob.glob(os.path.join(CACHE, "*.wav")))
ids = [os.path.splitext(os.path.basename(f))[0] for f in files]
lin = [f for f in files if lang[os.path.splitext(os.path.basename(f))[0]] == "lin"]
sna = [f for f in files if lang[os.path.splitext(os.path.basename(f))[0]] == "sna"]
print("lin=%d (%s + KenLM a=%.2f b=%.2f lsb=%s) | sna=%d (sna_ps greedy)"
% (len(lin), os.path.basename(LINM), ALPHA, BETA, LSB, len(sna)), flush=True)
out = {}
# ---- lin : beam + KenLM ----
proc = AutoProcessor.from_pretrained(LINM)
tok = proc.tokenizer
m = AutoModelForCTC.from_pretrained(LINM, dtype=torch.float32).cuda().eval()
logs, order, greedy = [], [], []
with torch.inference_mode():
for b in batches(lin, 90):
au = [sf.read(f, dtype="float32")[0] for f in b]
x = proc(au, sampling_rate=SR, 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, f in enumerate(b):
logs.append(lg[j])
order.append(f)
greedy.append(norm(tok.decode(lg[j].argmax(-1))))
del m
torch.cuda.empty_cache()
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] = ""
dec = build_ctcdecoder(lab, kenlm_model_path=ARPA, alpha=ALPHA, beta=BETA,
lm_score_boundary=LSB)
with Pool(8) as p:
hyps = dec.decode_batch(p, logs, beam_width=64)
for f, h, g in zip(order, hyps, greedy):
h = norm(h)
if h and g:
h = g[:1] + h[1:] # casse du 1er caractere = celle du modele acoustique
out[os.path.splitext(os.path.basename(f))[0]] = h
print("lin OK", flush=True)
# ---- sna : greedy (le LM degrade, mesure) ----
proc2 = AutoProcessor.from_pretrained("/root/models/sna_ps_best")
m2 = AutoModelForCTC.from_pretrained("/root/models/sna_ps_best",
dtype=torch.bfloat16).cuda().eval()
with torch.inference_mode():
for b in batches(sna, 140):
au = [sf.read(f, dtype="float32")[0] for f in b]
x = proc2(au, sampling_rate=SR, return_tensors="pt", padding=True)
x = {k: v.to("cuda", dtype=torch.bfloat16 if v.dtype == torch.float32 else v.dtype)
for k, v in x.items()}
pid = m2(**x).logits.float().argmax(-1).cpu().numpy()
for f, s in zip(b, proc2.batch_decode(pid)):
out[os.path.splitext(os.path.basename(f))[0]] = norm(s)
del m2
torch.cuda.empty_cache()
print("sna OK", flush=True)
# ---- ecriture, en comblant les vides depuis la meilleure soumission connue ----
fb = {}
ref = "/root/sub_p2_KENLM.csv"
if os.path.exists(ref):
fb = {r["ID"]: r["Target"] for r in csv.DictReader(open(ref, encoding="utf-8"))}
filled = 0
for i in ids:
if not out.get(i, "").strip() and fb.get(i, "").strip():
out[i] = fb[i]
filled += 1
with open(OUT, "w", newline="", encoding="utf-8") as fo:
w = csv.writer(fo)
w.writerow(["ID", "Target"])
for i in ids:
w.writerow([i, out.get(i) or "a"])
empty = sum(1 for i in ids if not out.get(i, "").strip())
print("GEN_DONE %s | %d IDs | vides=%d | combles=%d" % (OUT, len(ids), empty, filled), flush=True)
if __name__ == "__main__":
main()
|