162 lines
7.0 KiB
Python
162 lines
7.0 KiB
Python
"""Step 5: score every take automatically (no listening).
|
||
|
||
usage: qc [--shard i/n] [--no-mos] [--f0] [--jobs jobs.json --out out.json]
|
||
Default jobs: every work/takes/<take>_s<seed>.wav of the current segments; results are cached in
|
||
work/qc/<i>of<n>.json by take name, so re-runs and lexicon changes only score new takes. Run n shards
|
||
in parallel to use the GPU fully (each process needs ~8-13 GB RAM and ~5 GB VRAM).
|
||
|
||
Per take:
|
||
- wer, hyp — GigaAM-v3 (RNN-T) on <=20 s pieces cut at silences, against the text;
|
||
- stress_ok / stress_n — meter2 (shipped with the ru-stress model) agrees with the RUAccent marks;
|
||
- sim — WavLM-SV x-vector cosine to the EN narrator centroid (timbre);
|
||
- mos — UTMOSv2 naturalness; cps — pace, chars/s; maxgap — longest inner pause; clip — clipped samples;
|
||
- f0 — pitch spread in semitones (only with --f0: pyin is slow and runs on the CPU).
|
||
"""
|
||
import glob
|
||
import json
|
||
import os
|
||
import re
|
||
import sys
|
||
import tempfile
|
||
|
||
import jiwer
|
||
import librosa
|
||
import numpy as np
|
||
import soundfile as sf
|
||
import torch
|
||
|
||
from .config import load
|
||
|
||
VOW = "аеёиоуыэюя"
|
||
|
||
|
||
def norm(t):
|
||
t = t.lower().replace("ё", "е").replace("́", "").replace("+", "")
|
||
return " ".join(re.sub(r"[^\w\s]", " ", t).split())
|
||
|
||
|
||
def arg(name, default=None):
|
||
return sys.argv[sys.argv.index(name) + 1] if name in sys.argv else default
|
||
|
||
|
||
def main():
|
||
cfg = load()
|
||
sh = tuple(int(x) for x in arg("--shard", "0/1").split("/"))
|
||
use_mos, use_f0 = "--no-mos" not in sys.argv and cfg.get("qc", "mos", True), "--f0" in sys.argv
|
||
if arg("--jobs"):
|
||
jobs, out_path = json.load(open(arg("--jobs"), encoding="utf-8")), arg("--out")
|
||
else:
|
||
segs = {s["take"]: s for s in json.load(open(cfg.w("segs.json"), encoding="utf-8"))}
|
||
jobs = []
|
||
for f in sorted(glob.glob(str(cfg.w("takes", "x").parent / "*_s*.wav"))):
|
||
stem = os.path.basename(f)[:-4]
|
||
take = stem.rsplit("_s", 1)[0]
|
||
if take in segs:
|
||
jobs.append(dict(id=stem, text=segs[take]["text"], plus=segs[take]["plus"], wav=f))
|
||
out_path = str(cfg.w("qc", f"{sh[0]}of{sh[1]}.json"))
|
||
done = {}
|
||
for f in glob.glob(str(cfg.w("qc", "x").parent / "*.json")) + ([out_path] if os.path.exists(out_path) else []):
|
||
try:
|
||
done.update(json.load(open(f, encoding="utf-8")))
|
||
except Exception:
|
||
pass
|
||
jobs = [j for j in jobs[sh[0]::sh[1]] if j["id"] not in done]
|
||
print("takes to score", len(jobs), flush=True)
|
||
if not jobs:
|
||
return
|
||
res = json.load(open(out_path, encoding="utf-8")) if os.path.exists(out_path) else {}
|
||
|
||
from huggingface_hub import snapshot_download
|
||
model = cfg.get("tts", "model", "siriusfreak/qwen3-tts-12hz-1.7b-ru-stress-cf") # meter2 ships with it
|
||
meter = os.path.join(model if os.path.isdir(model) else snapshot_download(model, allow_patterns=["meter/*"]),
|
||
"meter")
|
||
os.environ.setdefault("METER2_HEAD", os.path.join(meter, "head.pt"))
|
||
os.environ.setdefault("METER2_DEVICE", "cuda")
|
||
sys.path.insert(0, meter)
|
||
from meter2 import judge, vowel_pos, words_in
|
||
|
||
def marks(plus, text):
|
||
toks = [m.group(0) for m in re.finditer(r"[А-Яа-яЁё+]+(?:-[А-Яа-яЁё+]+)*", plus)]
|
||
if len(toks) != len(words_in(text)):
|
||
return {}
|
||
return {wi: sum(1 for ch in w[:w.index("+")] if ch.lower() in VOW) for wi, w in enumerate(toks)
|
||
if "+" in w and len(vowel_pos(w.replace("+", ""))) > 1}
|
||
|
||
import gigaam
|
||
asr = gigaam.load_model("v3_rnnt")
|
||
tmp = os.path.join(tempfile.gettempdir(), f"ensvoice_qc_{os.getpid()}.wav") # one per process
|
||
|
||
def transcribe(y, sr):
|
||
iv = librosa.effects.split(y, top_db=35, frame_length=1024, hop_length=256)
|
||
pieces, start, cur = [], None, None
|
||
for a, b in iv:
|
||
if cur is None:
|
||
start, cur = a, b
|
||
elif (b - start) / sr <= 20:
|
||
cur = b
|
||
else:
|
||
pieces.append((start, cur)); start, cur = a, b
|
||
if cur is not None:
|
||
pieces.append((start, cur))
|
||
out = []
|
||
for a, b in pieces:
|
||
sf.write(tmp, y[max(0, a - 2400):b + 2400], sr)
|
||
r = asr.transcribe(tmp)
|
||
out.append(r if isinstance(r, str) else getattr(r, "text", str(r)))
|
||
return " ".join(out), iv
|
||
|
||
from transformers import AutoFeatureExtractor, WavLMForXVector
|
||
sv_id = cfg.get("qc", "wavlm", "microsoft/wavlm-base-plus-sv")
|
||
fe = AutoFeatureExtractor.from_pretrained(sv_id)
|
||
sv = WavLMForXVector.from_pretrained(sv_id).cuda().eval()
|
||
cen = torch.tensor(np.load(cfg.w("centroid.npy")))
|
||
|
||
def embed(y16):
|
||
if len(y16) < 32000:
|
||
y16 = np.pad(y16, (0, 32000 - len(y16)))
|
||
x = fe(y16[:16000 * 30], sampling_rate=16000, return_tensors="pt")
|
||
with torch.no_grad():
|
||
e = sv(**{k: v.cuda() for k, v in x.items()}).embeddings
|
||
return torch.nn.functional.normalize(e, dim=-1)[0].cpu()
|
||
|
||
if use_mos:
|
||
import utmosv2
|
||
mos_model = utmosv2.create_model(pretrained=True)
|
||
|
||
for i, j in enumerate(jobs):
|
||
try:
|
||
y, sr = sf.read(j["wav"], dtype="float32")
|
||
if y.ndim > 1:
|
||
y = y.mean(1)
|
||
dur = len(y) / sr
|
||
hyp, iv = transcribe(y, sr)
|
||
wer = jiwer.wer(norm(j["text"]), norm(hyp)) if norm(hyp) else 1.0
|
||
J = judge(y, sr, j["text"])
|
||
items = [(wi, v) for wi, v in marks(j["plus"], j["text"]).items() if wi < len(J) and J[wi]["best"] is not None]
|
||
y16 = librosa.resample(y, orig_sr=sr, target_sr=16000)
|
||
gaps = [(iv[k + 1][0] - iv[k][1]) / sr for k in range(len(iv) - 1)]
|
||
r = dict(dur=round(dur, 2), cps=round(len(j["text"]) / dur, 1), wer=round(wer, 3), hyp=hyp,
|
||
stress_ok=sum(1 for wi, v in items if J[wi]["best"] == v), stress_n=len(items),
|
||
bad=[J[wi]["word"] for wi, v in items if J[wi]["best"] != v][:10],
|
||
sim=round(float(embed(y16) @ cen), 3), maxgap=round(max(gaps), 2) if gaps else 0,
|
||
clip=int(np.sum(np.abs(y) > 0.99)))
|
||
if use_f0:
|
||
f0, vf, _ = librosa.pyin(y16, fmin=60, fmax=400, sr=16000, frame_length=1024)
|
||
f0 = f0[vf]
|
||
r["f0"] = round(float(np.std(12 * np.log2(f0 / np.median(f0)))) if len(f0) > 10 else 0.0, 2)
|
||
if use_mos:
|
||
r["mos"] = round(float(mos_model.predict(input_path=j["wav"])), 3)
|
||
except Exception as ex: # a broken take simply loses
|
||
r = dict(error=str(ex)[:200], wer=1.0, cps=0, stress_ok=0, stress_n=1, sim=0, mos=0)
|
||
res[j["id"]] = r
|
||
print(f"{i + 1}/{len(jobs)}", j["id"], r.get("wer"), f'{r["stress_ok"]}/{r["stress_n"]}', r.get("sim"),
|
||
r.get("mos"), flush=True)
|
||
if i % 20 == 0:
|
||
json.dump(res, open(out_path, "w", encoding="utf-8"), ensure_ascii=False)
|
||
json.dump(res, open(out_path, "w", encoding="utf-8"), ensure_ascii=False)
|
||
print("QC_DONE", len(res))
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|