File size: 6,043 Bytes
5125634
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
"""Evaluate every model folder with Whisper (CER/WER) and TitaNet (speaker similarity).

Folder layout:
    <root>/<model>/sample_00.wav ...

Usage:
    python evaluate_compare.py --root /path/eval_in \
        --models bayan,magpie,xtts,chatterbox,mms \
        --text-file ../texts.txt --out ../metrics.json --csv ../results.csv \
        --whisper /path/whisper-large-v3 --titanet /path/titanet.nemo \
        --speaker-ref-dir /path/wav22k
"""
from __future__ import annotations

import argparse
import csv
import json
import re
import unicodedata
from pathlib import Path

import numpy as np
import soundfile as sf
import torch
import torch.nn.functional as F
from faster_whisper import WhisperModel
from jiwer import cer, wer
from nemo.collections.asr.models import EncDecSpeakerLabelModel

DIACRITICS = re.compile(r"[\u0610-\u061a\u064b-\u065f\u0670\u06d6-\u06ed\u0640]")
NUMBER_WORDS = (
    "واحد|واحده|اثنان|اثنين|ثلاثه|اربعه|خمسه|سته|سبعه|ثمانيه|تسعه|عشره|عشر|"
    "عشرون|عشرين|ثلاثون|ثلاثين|اربعون|اربعين|خمسون|خمسين|ستون|ستين|"
    "سبعون|سبعين|ثمانون|ثمانين|تسعون|تسعين|مايه|مائه|ميه|ميتين|مائتين|مئتين|مائتان|"
    "الف|الاف|الفين|الفان|مليون|مليار|مليارات|"
    "ثلاثمائه|ثلاثمايه|اربعمائه|اربعمايه|خمسمائه|خمسمايه|ستمائه|ستمايه|"
    "سبعمائه|سبعمايه|ثمانمائه|ثمانمايه|تسعمائه|تسعمايه"
)
UNITS = "كم|كيلومتر|كيلومترات|كيلومترا|متر|مترا|امتار|عام|عاما|اعوام|سنوات|سنه|دوله|دول|شخص|اشخاص"
NUM_RE = re.compile(rf"\bو?(?:{NUMBER_WORDS}|{UNITS}|[0-9]+)\b")


def normalize_ar(text: str) -> str:
    text = unicodedata.normalize("NFKC", text)
    text = DIACRITICS.sub("", text)
    text = (
        text.replace("أ", "ا").replace("إ", "ا").replace("آ", "ا")
        .replace("ى", "ي").replace("ة", "ه").replace("ؤ", "و").replace("ئ", "ي")
    )
    text = re.sub(r"[^\w\s]", " ", text)
    return re.sub(r"\s+", " ", text).strip()


def mask_numbers(text: str) -> str:
    out = NUM_RE.sub("NUM", normalize_ar(text))
    out = re.sub(r"(?:NUM\s*)+", "NUM ", out)
    return re.sub(r"\s+", " ", out).strip()


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--root", required=True)
    ap.add_argument("--models", required=True)
    ap.add_argument("--text-file", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--csv", required=True)
    ap.add_argument("--whisper", required=True)
    ap.add_argument("--titanet", required=True)
    ap.add_argument("--speaker-ref-dir", required=True)
    ap.add_argument("--ref-stride", type=int, default=50)
    ap.add_argument("--ref-count", type=int, default=12)
    args = ap.parse_args()

    root = Path(args.root)
    texts = [ln.strip() for ln in Path(args.text_file).read_text(encoding="utf-8").splitlines() if ln.strip()]
    whisper = WhisperModel(args.whisper, device="cpu", compute_type="int8")
    sv = EncDecSpeakerLabelModel.restore_from(args.titanet, map_location="cpu").eval()

    refs = sorted(Path(args.speaker_ref_dir).glob("*.wav"))[:: args.ref_stride][: args.ref_count]
    with torch.no_grad():
        embs = torch.stack([F.normalize(sv.get_embedding(str(p)).reshape(-1), dim=0) for p in refs])
    voice = F.normalize(embs.mean(0), dim=0)
    print(f"voice centroid from {len(refs)} clips")

    report: dict[str, dict] = {}
    for model in [m.strip() for m in args.models.split(",") if m.strip()]:
        files = sorted((root / model).glob("sample_*.wav"))
        rows = []
        for f, ref in zip(files, texts):
            segs, _ = whisper.transcribe(str(f), language="ar", beam_size=5, vad_filter=False, condition_on_previous_text=False)
            hyp = " ".join(s.text.strip() for s in segs).strip()
            c = float(cer(normalize_ar(ref), normalize_ar(hyp)))
            c_norm = float(cer(mask_numbers(ref), mask_numbers(hyp)))
            w = float(wer(normalize_ar(ref), normalize_ar(hyp)))
            with torch.no_grad():
                e = F.normalize(sv.get_embedding(str(f)).reshape(-1), dim=0)
            sim = float(e @ voice)
            dur = float(sf.info(str(f)).duration)
            rows.append({"file": f.name, "cer": c, "cer_norm": c_norm, "wer": w,
                         "sim": sim, "duration": dur, "hyp": hyp, "ref": ref})
            print(f"{model}/{f.name} CER={c:.3f} CER*={c_norm:.3f} WER={w:.3f} SIM={sim:.3f}")
        report[model] = {
            "n": len(rows),
            "mean_cer": float(np.mean([r["cer"] for r in rows])),
            "mean_cer_norm": float(np.mean([r["cer_norm"] for r in rows])),
            "mean_wer": float(np.mean([r["wer"] for r in rows])),
            "mean_sim": float(np.mean([r["sim"] for r in rows])),
            "total_duration": float(np.sum([r["duration"] for r in rows])),
            "samples": rows,
        }

    Path(args.out).write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
    with Path(args.csv).open("w", newline="", encoding="utf-8") as fh:
        wr = csv.writer(fh)
        wr.writerow(["model", "file", "cer", "cer_numeric_normalized", "wer", "sim", "duration", "hypothesis"])
        for model, rep in report.items():
            for r in rep["samples"]:
                wr.writerow([model, r["file"], f"{r['cer']:.4f}", f"{r['cer_norm']:.4f}",
                             f"{r['wer']:.4f}", f"{r['sim']:.4f}", f"{r['duration']:.2f}", r["hyp"]])

    for model, rep in report.items():
        print(f"\n{model}: CER={rep['mean_cer']:.4f} CER*={rep['mean_cer_norm']:.4f} "
              f"WER={rep['mean_wer']:.4f} SIM={rep['mean_sim']:.3f} dur={rep['total_duration']:.1f}s")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())