benchmark_public_asr.py

#!/usr/bin/env python3
import argparse
import io
import json
import re
import time
import unicodedata
from pathlib import Path

import jiwer
import numpy as np
import pyarrow.parquet as pq
import soundfile as sf
import torch


PERSIAN_DIGITS = "۰۱۲۳۴۵۶۷۸۹٠١٢٣٤٥٦٧٨٩"
ASCII_DIGITS = "01234567890123456789"
DIGIT_MAP = str.maketrans(PERSIAN_DIGITS, ASCII_DIGITS)
PUNCT_RE = re.compile(r"[،؛؟,.\!?\(\)\[\]\{\}\"'`«»“”‘’\-_=+/\\:;]")


MODEL_REGISTRY = {
    "nezamisafa/whisper-persian-v4": {
        "family": "Whisper",
        "loader": "whisper",
        "precision": "fp32",
    },
    "vhdm/whisper-large-fa-v1": {
        "family": "Whisper",
        "loader": "whisper",
        "precision": "fp32",
    },
}


def normalize(text):
    text = unicodedata.normalize("NFKC", text or "")
    text = text.replace("‌", " ").replace("​", "")
    text = text.replace("ي", "ی").replace("ك", "ک").translate(DIGIT_MAP)
    text = PUNCT_RE.sub(" ", text)
    return re.sub(r"\s+", " ", text).strip()


def read_audio(path_or_bytes, base_dir=None):
    if isinstance(path_or_bytes, (bytes, bytearray)):
        wav, sr = sf.read(io.BytesIO(path_or_bytes), dtype="float32")
    else:
        path = Path(path_or_bytes)
        if base_dir is not None and not path.is_absolute():
            path = Path(base_dir) / path
        wav, sr = sf.read(path, dtype="float32")
    if wav.ndim > 1:
        wav = wav.mean(axis=-1)
    return wav, sr


def load_gold(path):
    base = Path(path).parent
    rows = []
    with open(path) as f:
        for line in f:
            item = json.loads(line)
            ref = (item.get("gold_transcript") or "").strip()
            if not ref:
                continue
            wav, sr = read_audio(item["audio_path"], base)
            rows.append({"id": str(item.get("programmeitem_id")), "wav": wav, "sr": sr, "ref": ref})
    return rows


def load_fleurs(fleurs_dir, limit=0):
    files = sorted(Path(fleurs_dir).glob("data/*.parquet")) or sorted(Path(fleurs_dir).glob("*.parquet"))
    table = pq.read_table([str(p) for p in files])
    records = table.to_pylist()
    if limit:
        records = records[:limit]
    rows = []
    for idx, item in enumerate(records):
        audio = item["audio"]
        audio_bytes = audio.get("bytes") if isinstance(audio, dict) else None
        if audio_bytes is None:
            continue
        wav, sr = read_audio(audio_bytes)
        rows.append({
            "id": str(item.get("id", idx)),
            "wav": wav,
            "sr": sr,
            "ref": item.get("transcription") or item.get("raw_transcription") or "",
        })
    return rows


def count_params(model):
    return sum(p.numel() for p in model.parameters())


@torch.no_grad()
def run_whisper(repo, clips, device):
    from transformers import WhisperForConditionalGeneration, WhisperProcessor

    processor = WhisperProcessor.from_pretrained(repo)
    model = WhisperForConditionalGeneration.from_pretrained(repo, torch_dtype=torch.float32).to(device).eval()
    params = count_params(model)
    preds = []
    times = []
    for clip in clips:
        inputs = processor(clip["wav"], sampling_rate=16000, return_tensors="pt")
        feats = inputs.input_features.to(device=device, dtype=torch.float32)
        if device.startswith("cuda"):
            torch.cuda.synchronize()
        t0 = time.time()
        try:
            ids = model.generate(feats, language="fa", task="transcribe")
        except TypeError:
            ids = model.generate(feats)
        if device.startswith("cuda"):
            torch.cuda.synchronize()
        times.append(time.time() - t0)
        preds.append(processor.batch_decode(ids, skip_special_tokens=True)[0])
    del model
    if device.startswith("cuda"):
        torch.cuda.empty_cache()
    return preds, times, params


def evaluate(repo, meta, set_name, clips, device, pred_dir):
    if meta["loader"] != "whisper":
        raise NotImplementedError(meta["loader"])
    start = time.time()
    preds, times, params = run_whisper(repo, clips, device)
    refs_norm = [normalize(c["ref"]) for c in clips]
    preds_norm = [normalize(p) for p in preds]
    pred_path = pred_dir / f"{repo.replace('/', '__')}__{set_name}_predictions.jsonl"
    with pred_path.open("w") as f:
        for clip, pred, ref_norm, hyp_norm in zip(clips, preds, refs_norm, preds_norm):
            f.write(json.dumps({
                "id": clip["id"],
                "ref": clip["ref"],
                "hyp": pred,
                "ref_norm": ref_norm,
                "hyp_norm": hyp_norm,
            }, ensure_ascii=False) + "\n")
    return {
        "model": repo,
        "repo": repo,
        "family": meta["family"],
        "precision": meta["precision"],
        "params_b": params / 1e9,
        "set": set_name,
        "n": len(clips),
        "wer": jiwer.wer(refs_norm, preds_norm),
        "cer": jiwer.cer(refs_norm, preds_norm),
        "mean_decode_ms": float(np.mean(times) * 1000.0) if times else 0.0,
        "wall_sec": time.time() - start,
        "predictions": str(pred_path),
        "status": "complete",
    }


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--gold", default="/workspace/golha/gold.jsonl")
    parser.add_argument("--fleurs-dir", default="/workspace/golha/persian-eval")
    parser.add_argument("--fleurs-limit", type=int, default=0)
    parser.add_argument("--device", default="cuda:0")
    parser.add_argument("--out", default="/workspace/golha/public_asr_double_benchmark.jsonl")
    parser.add_argument("--pred-dir", default="/workspace/golha/public_asr_double_benchmark_predictions")
    parser.add_argument("models", nargs="*", default=list(MODEL_REGISTRY))
    args = parser.parse_args()

    pred_dir = Path(args.pred_dir)
    pred_dir.mkdir(parents=True, exist_ok=True)
    evalsets = {
        "gold69": load_gold(args.gold),
        "fleurs": load_fleurs(args.fleurs_dir, args.fleurs_limit),
    }
    done = set()
    out_path = Path(args.out)
    if out_path.exists():
        for line in out_path.read_text().splitlines():
            try:
                rec = json.loads(line)
                done.add((rec.get("model"), rec.get("set")))
            except Exception:
                pass
    with out_path.open("a", buffering=1) as f:
        for repo in args.models:
            meta = MODEL_REGISTRY[repo]
            for set_name, clips in evalsets.items():
                if (repo, set_name) in done:
                    print(f"[skip] {repo} {set_name}", flush=True)
                    continue
                print(f"[eval] {repo} {set_name} n={len(clips)}", flush=True)
                try:
                    rec = evaluate(repo, meta, set_name, clips, args.device, pred_dir)
                    print(
                        f"[result] {repo} {set_name} WER={rec['wer']*100:.2f}% "
                        f"CER={rec['cer']*100:.2f}% wall={rec['wall_sec']:.1f}s",
                        flush=True,
                    )
                except Exception as exc:
                    rec = {
                        "model": repo,
                        "repo": repo,
                        "family": meta["family"],
                        "precision": meta["precision"],
                        "set": set_name,
                        "status": "failed",
                        "error": f"{type(exc).__name__}: {exc}",
                    }
                    print(f"[failed] {repo} {set_name}: {rec['error']}", flush=True)
                f.write(json.dumps(rec, ensure_ascii=False) + "\n")


if __name__ == "__main__":
    main()