src/python/anuvani_diar/check.py

# Copyright (c) 2026 The Anuvani-Diarization authors
# SPDX-License-Identifier: MIT OR Apache-2.0
"""anuvani-diar-check: does this installation reproduce the release's results?

  anuvani-diar-check [--audio-dir AUDIO_DIR] [--devices auto|N|0,2|cpu] [--long-dir DIR]

Diarises the release's 12 golden files (found by name under AUDIO_DIR; without it, downloaded from the Hugging Face
repository's golden/audio/ and sha256-checked) and compares the 12 RTTMs' fingerprints with
the release's golden fingerprints (assets/golden12.json, `indic_rttm_sha256`, as the release's own tools/verify_golden.py does for its
indic mode). With --long-dir, also the 1 h, 3 h and 6 h test files (tools/make_test_audio.py) against the reference
program's fingerprints for them (assets/long_golden.json). Before anything runs, every GPU checks both networks against
the release's own outputs (assets/nn_golden.json). Exit status 0 if every output is identical; 3 on a setup where the
networks passed the near-identical check, so identical bytes are not expected; 1 otherwise.
"""
import argparse
import hashlib
import json
import sys
from pathlib import Path

from .api import Diarizer, default_bundle, devices_arg
from .hub import assets_dir, golden_audio_dir

ASSETS = assets_dir()
LONG = ASSETS / "long_golden.json"
GOLDEN12 = ASSETS / "golden12.json"


def file_sha256(path: Path) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for block in iter(lambda: f.read(1 << 20), b""):
            h.update(block)
    return h.hexdigest()


def main(argv=None) -> int:
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--audio-dir", type=Path, default=None, help="default: download them from the Hugging Face repository")
    p.add_argument("--long-dir", type=Path, default=None)
    p.add_argument("--devices", "--device", dest="devices", default="auto", help="which devices run the networks: auto (every visible GPU, else the CPU), a number of GPUs (2: the first two), GPU indices (0,3), cuda:N, or cpu")
    p.add_argument("--bundle", type=Path, default=None)
    a = p.parse_args(argv)
    bundle = a.bundle or default_bundle()
    golden = json.loads(GOLDEN12.read_text())["files"]
    want = {s: dict(audio=v["flac_sha256"], rttm=v["indic_rttm_sha256"]) for s, v in golden.items()}
    audio_dir = a.audio_dir or golden_audio_dir()
    found = {f.stem: f for f in audio_dir.rglob("*.flac") if f.stem in want}
    if a.long_dir is not None:
        for s, v in json.loads(LONG.read_text())["files"].items():
            want[s] = dict(audio=v["audio"], rttm=v["rttm"])
            if (a.long_dir / f"{s}.flac").is_file():
                found[s] = a.long_dir / f"{s}.flac"
    missing = sorted(set(want) - set(found))
    if missing:
        print(f"audio missing for {missing}", file=sys.stderr)
        return 2
    differs = [s for s in sorted(found) if file_sha256(found[s]) != want[s]["audio"]]
    if differs:
        print(f"WARNING: audio differs from the release's for {differs}")
    d = Diarizer(bundle=bundle, devices=devices_arg(a.devices))
    status = {v["status"] for v in d.nn.exactness.values()}
    for slot, v in sorted(d.nn.exactness.items()):
        print(f"networks on {d.nn.slots[slot][1]}: {v['status']}" + ("" if v["status"] == "byte-identical" else
              f" (max |dp| {v['max_abs_dp']:.2g}, frame decisions {100 * v['frame_agreement']:.2f}% equal, "
              f"voiceprint cosine >= {v['min_cosine']:.5f})"))
    good = {}
    for r in d.run([found[s] for s in sorted(found)]):
        good[r.uri] = r.ok and hashlib.sha256(r.rttm.encode()).hexdigest() == want[r.uri]["rttm"]
    for s in sorted(found):
        print(f"{s:14s} {'PASS' if good[s] else 'DIFFERS'}")
    ok = sum(good.values())
    print(f"{ok}/{len(good)} outputs identical")
    if ok == len(good):
        return 0
    if status == {"near-identical"}:
        print("Identical bytes need the setup the release's outputs were produced with, so they are not expected "
              "here: the networks passed the near-identical check above.")
        return 3
    return 1


if __name__ == "__main__":
    sys.exit(main())