code/eval_tlog.py

"""
Evaluate on real phone recordings: ordinary people reciting into the Tarteel app (tarteel-ai/tlog).

    python eval_tlog.py                    # builds the test set once, then scores the model + NVIDIA
    python eval_tlog.py --clips 300        # a bigger test set (rebuilds it)
    python eval_tlog.py --onnx runs/full/export/model.int8.onnx   # the exported model, as the phone runs it

Training never uses these shards (download.py skips them), so every clip is unseen. The test set (data/tlog_test/) is drawn from 60
shards spread over the dataset (the shards are sorted by surah), with at most one clip per ayah,
6 per surah and 5 per shard, 1.5-30 s long. tlog's labels come from the file names and aren't
checked, so a clip that neither model can read (likely a wrong label or a non-recitation) is
reported separately. Needs `hf auth login` and the dataset's terms accepted on its page.
The scores use the leaderboard's kind of normalization (no diacritics, folded alef/ya/ta marbuta).
"""
import argparse
import io
import json
import os
import random
import sys
from types import SimpleNamespace

import numpy as np
import soundfile as sf
import torch

from quran_asr import audio, config as config_mod, model as model_mod, quran, text

HERE = os.path.dirname(os.path.abspath(__file__))
REPO = "datasets/tarteel-ai/tlog/data/clean-{:05d}-of-00411.parquet"
SHARDS = 411


def test_shards(seed=0):
    """The 60 shards the test set is drawn from (download.py never trains on them)."""
    rnd = random.Random(seed)
    return sorted({int(i * SHARDS / 60 + rnd.random() * SHARDS / 60) for i in range(60)}), rnd


def build(out, clips, seed):
    import pyarrow.parquet as pq
    from huggingface_hub import HfFileSystem

    fs = HfFileSystem()
    shards, rnd = test_shards(seed)
    candidates = []
    for k, shard in enumerate(shards, 1):
        print(f"  reading shard {shard} ({k}/{len(shards)})", flush=True)
        f = pq.ParquetFile(fs.open(REPO.format(shard), block_size=1 << 22))
        group = rnd.randrange(f.metadata.num_row_groups)
        for row in f.read_row_group(group, columns=["audio", "label"]).to_pylist():
            name = os.path.basename(row["audio"]["path"])
            try:
                sura, aya = (int(v) for v in name.split("_")[:2])
                wave, sr = sf.read(io.BytesIO(row["audio"]["bytes"]), dtype="float32", always_2d=True)
            except Exception:
                continue
            wave = audio.to_mono(wave)
            duration, rms = len(wave) / sr, float(np.sqrt(np.mean(wave ** 2)) if len(wave) else 0)
            label = text.normalize(row["label"] or "")
            if 1.5 <= duration <= 30 and rms > 1e-3 and label:
                candidates.append(dict(file=name, bytes=row["audio"]["bytes"], sura=sura, aya=aya, text=label,
                                       duration=round(duration, 2), shard=shard))
    rnd.shuffle(candidates)
    picked, per_sura, per_shard, ayahs = [], {}, {}, set()
    for c in candidates:
        if len(picked) >= clips or (c["sura"], c["aya"]) in ayahs:
            continue
        if per_sura.get(c["sura"], 0) >= 6 or per_shard.get(c["shard"], 0) >= 5:
            continue
        picked.append(c)
        ayahs.add((c["sura"], c["aya"]))
        per_sura[c["sura"]] = per_sura.get(c["sura"], 0) + 1
        per_shard[c["shard"]] = per_shard.get(c["shard"], 0) + 1
    os.makedirs(out, exist_ok=True)
    with open(os.path.join(out, "manifest.jsonl"), "w", encoding="utf-8", newline="\n") as m:
        for c in picked:
            with open(os.path.join(out, c["file"]), "wb") as f:
                f.write(c.pop("bytes"))
            m.write(json.dumps(c, ensure_ascii=False) + "\n")
    return picked


def describe(clips):
    d = sorted(c["duration"] for c in clips)
    suras = {c["sura"] for c in clips}
    lps = sorted(len(c["text"].replace(" ", "")) / c["duration"] for c in clips)
    print(f"test set: {len(clips)} clips, {sum(d) / 60:.1f} min, {len(suras)} surahs, "
          f"{len({c['shard'] for c in clips})} shards; duration median {d[len(d) // 2]:.1f} s "
          f"(10%: {d[len(d) // 10]:.1f} s, 90%: {d[9 * len(d) // 10]:.1f} s); "
          f"letters/s median {lps[len(lps) // 2]:.1f} (90%: {lps[9 * len(lps) // 10]:.1f}); "
          f"al-Fatiha {sum(c['sura'] == 1 for c in clips)}")


def join_tanween(s):
    """Our letters-only text splits a tanween alif off its word ("مرض ا"): rejoin it for scoring."""
    out = []
    for w in s.split():
        if w == "ا" and out:
            out[-1] += w
        else:
            out.append(w)
    return " ".join(out)


def scores(pairs):
    """Corpus-level letter / word / alef-insensitive word / skeleton letter error rates, in %."""
    tot = np.zeros(8)
    for hyp, ref in pairs:
        hyp, ref = join_tanween(text.normalize(hyp)), join_tanween(ref)
        ce, ct, we, wt, se, st = text.error_counts(hyp, ref)
        ha, ra = hyp.replace("ا", "").split(), ref.replace("ا", "").split()
        tot += [ce, ct, we, wt, text.edit_distance(ha, ra), len(ra), se, st]
    return {"letters": 100 * tot[0] / tot[1], "words": 100 * tot[2] / tot[3],
            "words_no_alef": 100 * tot[4] / tot[5], "skeleton": 100 * tot[6] / tot[7]}


def skeleton_error(hyp, ref):
    c = text.error_counts(join_tanween(text.normalize(hyp)), join_tanween(ref))
    return c[4] / max(1, c[5])


def our_model(cfg, args):
    if args.onnx:
        import sherpa_onnx

        rec = sherpa_onnx.OfflineRecognizer.from_zipformer_ctc(
            model=args.onnx, tokens=os.path.join(os.path.dirname(args.onnx), "tokens.txt"), num_threads=4)

        def run(wave):
            s = rec.create_stream()
            s.accept_waveform(audio.SAMPLE_RATE, wave)
            rec.decode_stream(s)
            return s.result.text.replace(text.WORD, " ")
        return f"ours (sherpa-onnx, {os.path.basename(args.onnx)})", run

    import export
    from quran_asr import features

    state, model_config = export.load_averaged(cfg.paths.run_dir, args.average)
    model = model_mod.build(SimpleNamespace(**model_config), len(text.VOCAB))
    model.load_state_dict(state)
    model.eval()

    def run(wave):
        f = features.fbank(torch.from_numpy(np.ascontiguousarray(wave)))
        with torch.no_grad():
            lp, lens = model(f[None], torch.tensor([f.shape[0]]))
        return text.decode(model_mod.greedy_decode(lp, lens)[0])
    if args.average <= 1:
        step = torch.load(os.path.join(cfg.paths.run_dir, "best.pt"), map_location="cpu", weights_only=False).get("step")
        return f"ours (best.pt, step {step})", run
    return f"ours (average of last {args.average} checkpoints)", run


def nvidia(cfg):
    import sherpa_onnx
    import prepare

    model, tokens = prepare._nvidia_model(cfg)
    rec = sherpa_onnx.OfflineRecognizer.from_nemo_ctc(model=model, tokens=tokens, num_threads=4)

    def run(wave):
        s = rec.create_stream()
        s.accept_waveform(audio.SAMPLE_RATE, wave)
        rec.decode_stream(s)
        return s.result.text
    return "NVIDIA FastConformer (the app's old model)", run


def main():
    parser = config_mod.from_args(__doc__)
    parser.set_defaults(config=os.path.join(HERE, "config", "full.yaml"))
    parser.add_argument("--clips", type=int, help="size of a new test set (default: keep the existing one, else 200)")
    parser.add_argument("--seed", type=int, default=0)
    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")
    parser.add_argument("--no-baseline", action="store_true", help="skip the NVIDIA model")
    args = parser.parse_args()
    sys.stdout.reconfigure(encoding="utf-8")
    torch.set_num_threads(4)  # leave the rest of the CPU to a running training
    cfg = config_mod.load(args.config)

    out = os.path.join(cfg.paths.data_dir, "tlog_test")
    manifest = os.path.join(out, "manifest.jsonl")
    clips = []
    if os.path.exists(manifest):
        clips = [json.loads(l) for l in open(manifest, encoding="utf-8")]
    if not clips or args.clips not in (None, len(clips)):
        args.clips = args.clips or 200
        print(f"building a {args.clips}-clip test set from tarteel-ai/tlog...")
        clips = build(out, args.clips, args.seed)
    describe(clips)

    systems = [our_model(cfg, args)] + ([] if args.no_baseline else [nvidia(cfg)])
    heard = {}
    for name, run in systems:
        print(f"transcribing with {name}...", flush=True)
        heard[name] = [run(audio.load(os.path.join(out, c["file"]))) for c in clips]

    # Clips no system can read are most likely mislabelled (tlog's labels are unchecked).
    unreadable = [i for i in range(len(clips))
                  if all(skeleton_error(heard[n][i], clips[i]["text"]) > 0.5 for n in heard)]
    keep = [i for i in range(len(clips)) if i not in set(unreadable)]
    print(f"\n{len(unreadable)} clips no model can read (probably a wrong label or not a recitation), "
          f"scored separately")
    print(f"{'':48s} {'letters':>8s} {'words':>8s} {'words*':>8s} {'skeleton':>9s}")
    for subset, idx in ((f"{len(keep)} readable clips", keep), (f"all {len(clips)} clips", range(len(clips)))):
        print(subset)
        for name in heard:
            s = scores([(heard[name][i], clips[i]["text"]) for i in idx])
            print(f"  {name:46s} {s['letters']:7.2f}% {s['words']:7.2f}% {s['words_no_alef']:7.2f}% {s['skeleton']:8.2f}%")
    print("words* = word error ignoring alef (as the leaderboard's alef-insensitive WER); skeleton = letters "
          "without alif and hamza (fair across spelling conventions: NVIDIA writes modern spelling)")

    ours = list(heard)[0]
    worst = sorted(keep, key=lambda i: -skeleton_error(heard[ours][i], clips[i]["text"]))[:5]
    print("\nhardest readable clips for our model:")
    for i in worst:
        c = clips[i]
        print(f"  {c['sura']}:{c['aya']} ({c['duration']} s)\n    label: {c['text'][:90]}\n    heard: {heard[ours][i][:90]}")


if __name__ == "__main__":
    main()