#!/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()