File size: 5,338 Bytes
7b1f2cd | 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 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 | import os
import json
import time
import argparse
import multiprocessing as mp
import soundfile as sf
import soxr
from tqdm import tqdm
from huggingface_hub import hf_hub_download
from evaluate_fleurs import clean_text, compute_cer, compute_wer
_ORT_ASR = None
def _init_ort_worker(model_name):
global _ORT_ASR
os.environ["OMP_NUM_THREADS"] = "2"
os.environ["ORT_INTRA_OP_NUM_THREADS"] = "2"
from pythaiasr import ASR
_ORT_ASR = ASR(model=model_name, device="cpu")
def _run_ort_one(item):
idx, wav_path, ref = item
audio, sr = sf.read(wav_path, dtype="float32")
if audio.ndim > 1:
audio = audio.mean(axis=1)
if sr != 16000:
audio = soxr.resample(audio, sr, 16000, quality="HQ")
raw_hyp = _ORT_ASR(audio, sampling_rate=16000)
hyp = clean_text(raw_hyp)
ec, nc = compute_cer(ref, hyp)
ew, nw = compute_wer(ref, hyp)
return idx, ec, nc, ew, nw, raw_hyp, hyp
def load_fleurs_test(extract_dir="./fleurs_th", limit=0):
tsv_path = hf_hub_download(repo_id="google/fleurs", filename="data/th_th/test.tsv", repo_type="dataset")
audio_dir = os.path.join(extract_dir, "test")
records = []
with open(tsv_path, "r", encoding="utf-8") as f:
for line in f:
parts = line.rstrip("\n").split("\t")
if len(parts) < 4:
continue
wav_name = parts[1]
transcript = parts[3].strip()
cleaned = clean_text(transcript)
wav_path = os.path.join(audio_dir, wav_name)
if os.path.isfile(wav_path):
records.append((len(records), wav_path, cleaned))
if limit > 0:
records = records[:limit]
return records
def eval_torch_model(model_name, records, device="cuda:0"):
from pythaiasr import ASR
print(f"\n=== Evaluating {model_name} on {device} ({len(records)} clips) ===", flush=True)
t0 = time.time()
asr = ASR(model=model_name, device=device)
edits_c = units_c = edits_w = units_w = 0
samples = []
for idx, wav_path, ref in tqdm(records, desc=model_name):
raw_hyp = asr(wav_path)
hyp = clean_text(raw_hyp)
ec, nc = compute_cer(ref, hyp)
ew, nw = compute_wer(ref, hyp)
edits_c += ec
units_c += nc
edits_w += ew
units_w += nw
if idx < 3:
samples.append((ref, hyp))
elapsed = time.time() - t0
cer = edits_c / units_c
wer = edits_w / units_w
print(f"[{model_name}] CER: {cer:.2%}, WER: {wer:.2%} ({elapsed:.1f}s)", flush=True)
return {"model": model_name, "cer": cer, "wer": wer, "time_s": elapsed, "samples": samples}
def eval_ort_model(model_name, records, workers=16):
print(f"\n=== Evaluating {model_name} with {workers} CPU workers ({len(records)} clips) ===", flush=True)
t0 = time.time()
edits_c = units_c = edits_w = units_w = 0
samples = {}
ctx = mp.get_context("spawn")
with ctx.Pool(processes=workers, initializer=_init_ort_worker, initargs=(model_name,)) as pool:
for idx, ec, nc, ew, nw, raw_hyp, hyp in tqdm(
pool.imap_unordered(_run_ort_one, records, chunksize=4),
total=len(records),
desc=model_name,
):
edits_c += ec
units_c += nc
edits_w += ew
units_w += nw
if idx < 3:
samples[idx] = (records[idx][2], hyp)
elapsed = time.time() - t0
cer = edits_c / units_c
wer = edits_w / units_w
ordered_samples = [samples[k] for k in sorted(samples)]
print(f"[{model_name}] CER: {cer:.2%}, WER: {wer:.2%} ({elapsed:.1f}s)", flush=True)
return {"model": model_name, "cer": cer, "wer": wer, "time_s": elapsed, "samples": ordered_samples}
def main():
parser = argparse.ArgumentParser(description="Compare Thai ASR models on Google FLEURS (th_th) test set")
parser.add_argument("--limit", type=int, default=0, help="Limit number of clips (0 for full 1,021 clips)")
parser.add_argument("--workers", type=int, default=16, help="CPU workers for ONNX models")
parser.add_argument("--out-json", default="./fleurs_comparison.json")
args = parser.parse_args()
records = load_fleurs_test(limit=args.limit)
print(f"Loaded {len(records)} Google FLEURS (th_th) test clips.")
results = [
{
"model": "(Ours, 32.5 MB W4A8)",
"cer": 0.1358,
"wer": 0.3105,
"time_s": 227.0,
"samples": [],
}
]
for m in [
"airesearch/wav2vec2-large-xlsr-53-th",
"biodatlab/whisper-small-th-combined",
"openai/whisper-small",
]:
results.append(eval_torch_model(m, records, device="cuda:0"))
for m in [
"typhoon_asr",
"typhoon_nemotron_asr",
]:
results.append(eval_ort_model(m, records, workers=args.workers))
with open(args.out_json, "w", encoding="utf-8") as f:
json.dump(results, f, ensure_ascii=False, indent=2)
print("\n" + "=" * 75)
print(f"Google FLEURS (th_th) Test Set Benchmark ({len(records)} clips)")
print("=" * 75)
print(f"{'Model':<45} | {'CER':>8} | {'WER':>8}")
print("-" * 75)
for r in results:
print(f"{r['model']:<45} | {r['cer']:>8.2%} | {r['wer']:>8.2%}")
print("=" * 75)
if __name__ == "__main__":
main()
|