Create compare_fleurs.py

compare_fleurs.py

@@ -0,0 +1,148 @@
+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()