code/benchmark.py

"""

Score the model on the Quran-Lab benchmark (Quran-Lab/quranic-asr-benchmark, the dataset behind

https://huggingface.co/spaces/Quran-Lab/quranic-asr-leaderboard): 600 clips, 200 each of studio

reciters (everyayah), one QUL reciter (al-Nufais) and phone recordings (tlog).



    python benchmark.py                                   # averaged last checkpoints (PyTorch)

    python benchmark.py --onnx runs/full/export/model.int8.onnx   # the exported model, as the phone runs it



Needs access to the dataset (gated) and `hf auth login`; the audio (~250 MB) is downloaded once to

data/quranlab/. Prints the official scorer's table, then the same scores without the clips we

trained on: the benchmark is held out from *its authors'* training data, not ours. Two of its three

everyayah reciters and a few of its tlog recordings are in our training set (prepare.py now keeps

the benchmark's tlog recordings out). Writes runs/<run>/benchmark/predictions.json in the

leaderboard's format, ready to attach to a submission.

"""
import collections
import datetime
import json
import os
import re
import subprocess
import sys

import torch

import eval_tlog
from quran_asr import audio, config as config_mod, spelling, text

HERE = os.path.dirname(os.path.abspath(__file__))
REPO = "Quran-Lab/quranic-asr-benchmark"
EVERYAYAH = "tarteel-ai/everyayah"


def fetch(folder):
    if not os.path.exists(os.path.join(folder, "benchmark.jsonl")):
        from huggingface_hub import snapshot_download

        print("downloading the benchmark (~250 MB)...", flush=True)
        snapshot_download(REPO, repo_type="dataset", local_dir=folder, max_workers=8)
    return [json.loads(l) for l in open(os.path.join(folder, "benchmark.jsonl"), encoding="utf-8") if l.strip()]


def everyayah_reciters(folder, rows):
    """{clip id: reciter} for the everyayah clips, read from tarteel-ai/everyayah's test split (ids are shard_row)."""
    path = os.path.join(folder, "everyayah_reciters.json")
    if os.path.exists(path):
        return json.load(open(path, encoding="utf-8"))
    import pyarrow.parquet as pq
    from huggingface_hub import HfApi, HfFileSystem

    files = [s.rfilename for s in HfApi().repo_info(EVERYAYAH, repo_type="dataset").siblings if "/test-" in s.rfilename]
    want = collections.defaultdict(list)
    for r in rows:
        if r["source"] == "everyayah_heldout":
            shard, row = r["id"].split("/")[1].rsplit("_", 1)
            want[shard].append((r["id"], int(row)))
    out, fs = {}, HfFileSystem()
    for shard, clips in want.items():
        f = next(f for f in files if f"/{shard}-" in f)
        names = pq.ParquetFile(fs.open(f"datasets/{EVERYAYAH}/{f}", block_size=1 << 20)).read(columns=["reciter"])
        names = names.column("reciter").to_pylist()
        out.update({cid: names[row] for cid, row in clips})
    with open(path, "w", encoding="utf-8") as f:
        json.dump(out, f, indent=0)
    return out


def seen(cfg, folder, rows):
    """{clip id: why} for the benchmark clips our training data contains."""
    def key(name):  # "Sahl_Yassin_128kbps" -> "sahl_yassin"
        return re.sub(r"_\d+kbps.*$", "", name).lower()

    ours = {key(r) for r in cfg.data.hafs_reciters}
    tlog = set()
    for split in ("train", "valid"):
        path = os.path.join(cfg.paths.data_dir, "manifests", cfg.name, f"{split}.jsonl")
        for line in open(path, encoding="utf-8"):
            e = json.loads(line)
            if e.get("reciter") == "tlog" or f"{os.sep}tlog{os.sep}" in e["path"]:
                tlog.add(os.path.splitext(os.path.basename(e["path"]))[0].split("_")[-1])
    out = {}
    for cid, reciter in everyayah_reciters(folder, rows).items():
        if key(reciter) in ours:
            out[cid] = f"everyayah reciter {reciter} is in our training set"
    for r in rows:
        if r["source"] == "tlog_holdout" and r["id"].split("_")[-1] in tlog:
            out[r["id"]] = "tlog recording in our training/validation data"
    return out


def table(rows, hyps, title):
    """The official scorer's numbers (its own normalization), per source and overall."""
    score = _official()
    print(f"\n{title}\n{'source':20s} {'clips':>5s} {'WER':>7s} {'CER':>7s} {'WER(alef)':>10s}")
    out = {}
    for src in sorted({r["source"] for r in rows}) + ["ALL"]:
        we = wn = ce = cn = wea = k = 0
        for r in rows:
            if src != "ALL" and r["source"] != src:
                continue
            h = hyps.get(r["id"], "")
            k += 1
            a, b = score.wer(r["reference_text"], h, score.strip); we += a; wn += b
            a, b = score.cer(r["reference_text"], h, score.strip); ce += a; cn += b
            wea += score.wer(r["reference_text"], h, score.strip_alef)[0]
        out[src] = dict(clips=k, wer=round(100 * we / max(wn, 1), 2), cer=round(100 * ce / max(cn, 1), 2),
                        wer_alef=round(100 * wea / max(wn, 1), 2))
        print(f"{src:20s} {k:5d} {out[src]['wer']:7.2f} {out[src]['cer']:7.2f} {out[src]['wer_alef']:10.2f}")
    return out


def _official():
    import importlib.util

    spec = importlib.util.spec_from_file_location("score", os.path.join(_folder, "score.py"))
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    return mod


def main():
    global _folder
    parser = config_mod.from_args(__doc__)
    parser.set_defaults(config=os.path.join(HERE, "config", "full.yaml"))
    parser.add_argument("--average", type=int, default=6, help="average the last N checkpoints")
    parser.add_argument("--onnx", help="score an exported model.int8.onnx through sherpa-onnx instead")
    args = parser.parse_args()
    sys.stdout.reconfigure(encoding="utf-8")
    torch.set_num_threads(4)
    cfg = config_mod.load(args.config)
    _folder = folder = os.path.join(cfg.paths.data_dir, "quranlab")
    rows = fetch(folder)

    name, run = eval_tlog.our_model(cfg, args)
    print(f"transcribing {len(rows)} clips with {name}...", flush=True)
    raw, started = {}, datetime.datetime.now()
    for k, r in enumerate(rows, 1):
        if r["id"] not in raw:
            raw[r["id"]] = text.normalize(run(audio.load(os.path.join(folder, r["audio"]))))
        if k % 100 == 0:
            print(f"  {k}/{len(rows)}", flush=True)
    seconds = (datetime.datetime.now() - started).total_seconds()
    words = spelling.table(cfg.paths.data_dir)
    hyps = {k: spelling.to_standard(v, words) for k, v in raw.items()}

    leaked = seen(cfg, folder, rows)
    why = collections.Counter(v.split(" is ")[0] if "reciter" in v else v for v in leaked.values())
    unseen = [r for r in rows if r["id"] not in leaked]
    table(rows, raw, "all 600 clips, Uthmani spelling (as the model writes)")
    official = table(rows, hyps, "all 600 clips, standard spelling (what the leaderboard shows)")
    clean = table(unseen, hyps, f"standard spelling, without the {len(leaked)} clips we trained on "
                                f"({', '.join(f'{n}: {v}' for v, n in why.items())})")

    out = os.path.join(cfg.paths.run_dir, "benchmark")
    os.makedirs(out, exist_ok=True)
    path = os.path.join(out, "predictions.json")
    with open(path, "w", encoding="utf-8") as f:
        json.dump(dict(
            key="quran-conformer-ctc", model=name, family="conformer-ctc", mode="offline",
            hf_repo=None, language_hint=None, rerun_from_scratch=True, run_date=datetime.date.today().isoformat(),
            num_rows=len(rows), num_unique_clips=len({r["id"] for r in rows}), benchmark_dataset=REPO,
            config=dict(decoding="greedy CTC, no language model, no lexicon constraint",
                        spelling=("the model writes Uthmani rasm letters; each output word that exactly matches a "
                                  "Quran word is rewritten in standard spelling from a fixed table built from two "
                                  "editions of the Quran text (Tanzil Uthmani/simple); other words are unchanged"),
                        checkpoint=name, cpu_threads=4, seconds=round(seconds)),
            notes=f"{len(leaked)} clips overlap our training data; scores without them under scores_unseen",
            scores=official, scores_unseen=clean, overlapping_ids=sorted(leaked),
            predictions=[dict(id=r["id"], source=r["source"], hypothesis_text=hyps[r["id"]],
                              uthmani_text=raw[r["id"]], reference_text=r["reference_text"]) for r in rows]),
            f, ensure_ascii=False, indent=1)
    print(f"\n{seconds:.0f} s; predictions for the leaderboard: {path}")


if __name__ == "__main__":
    main()