| 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() |
|
|
