Anuvani-Diarization / src /README.md
StephennFernandes's picture
Anuvani-Diarization 1.0.0
63dd946
|
Raw History Blame Contribute Delete
6.81 kB

anuvani_diar: the source of the package and of the serving kit

The source of Anuvani-Diarization 1.0's Python package anuvani_diar and of its serving kit (triton_inference/). Everything except the two networks runs in Rust: decoding, windowing, batching across every GPU, masks, clustering, the decoder and the RTTM text; the networks run the release's own PyTorch code, in your Python process or on Triton, on any NVIDIA GPU (or the CPU), as many as you give it. The output is the release's original Python program's (reference/diarize.py in the model repository): byte-identical on the setup the release's outputs were produced with, near-identical, checked at start-up, on any other GPU or PyTorch version.

To use the model, install the wheel or run the serving kit (see the model card). Build from here to change the code.

Building

The release wheel (CPython 3.10 and newer, manylinux_2_28, x86_64), built in a pinned container (needs Docker):

# artifacts/wheels/anuvani_diar-<version>-cp310-abi3-manylinux_2_28_x86_64.whl
tools/build_wheel.sh --site-packages SITE

The wheel bundles numpy 1.26.4's OpenBLAS and soundfile 0.14.0's libsndfile, which tools/stage_wheel_files.py copies from SITE, the site-packages folder of any Python environment with numpy==1.26.4 and soundfile==0.14.0 installed, and checks by sha256 (docs/WHEEL.md).

For development: a Rust toolchain (1.80 or newer), maturin 1.15 and the reference Python environment (Python 3.10, PyTorch 2.1.1, the packages of env/pip_freeze_reference.txt in the model repository's reference/):

# stages the bundled libraries and assets, builds the extension into .venv (maturin)
make build
source env/canonical.sh && source .venv/bin/activate

The test suites (make test-fast, test-cache, test-gpu, test-gpu-all) compare every output with the release's own, regenerated by tools/oracle_from_audio.py and tools/oracle_from_cache.py from the reference program; that generated data is not part of the repository (docs/TESTING.md).

Using it

from anuvani_diar import Diarizer
# every visible GPU (the CPU if there is none)
d = Diarizer()
d("meeting.wav").save_rttm("meeting.rttm")
# many files, results as each one finishes
for r in d.run(["a.wav", "b.flac"], out_dir="rttm/"):
    print(r.uri, r.speakers, r.timing["wall_s"])

Diarizer(devices=2) uses the first two GPUs, Diarizer(devices=[0, 3]) GPUs 0 and 3, Diarizer(devices="cpu") the CPU; the default, "auto", is every visible GPU.

Command line, with the reference program's flags (there is one model, so no --mode):

anuvani-diar --out OUT a.wav b.flac
anuvani-diar --out OUT --in-dir AUDIO_DIR
# the first two GPUs; --devices 0,3 or --devices cpu also work
anuvani-diar --out OUT --devices 2 a.wav

Checking an installation (the release's install check; --long-dir adds the 1 h, 3 h and 6 h files that tools/make_test_audio.py makes):

# downloads the 12 check recordings from the model repository
anuvani-diar-check
anuvani-diar-check --audio-dir AUDIO_DIR --long-dir LONG_DIR

It prints PASS or DIFFERS per output; before that, every GPU checks both networks against the release's own outputs. Exit status 0 means every RTTM is identical to the release's; 3 means a setup where the networks passed the near-identical check, so identical bytes are not expected.

Serving many users: the kit in triton_inference/ (Triton, the FastAPI service on this engine, Docker Compose, any number of GPUs on one or several machines; its README). triton/ holds the two Triton images and the model files the kit builds on, and the Rust gateway (crates/diar-gateway) is the reference implementation of the same API that the serving identity checks and measurements ran through.

What is guaranteed

The RTTMs (and the front-end records behind them) are byte-identical to the release's on everything tested: all 1,945 benchmark files from cached records, all 1,973 audio inputs including 1 h / 3 h / 6 h files and 25 robustness files through the local engine on several GPUs at once, and all 1,973 again through Triton, the gateway and the client with concurrent shuffled jobs, also with another GPU assignment and with Triton killed and restarted in the middle of the run. docs/TESTING.md lists the checks.

The conditions: the setup the release's outputs were produced with (env/canonical.sh: PyTorch 2.1.1 with CUDA 12.1 and cuBLAS 12.1.0.26, on the same kind of GPU), one network stream per GPU. Every network backend checks itself at start-up against the release's own outputs (assets/nn_golden.json, assets/nn_reference.npz): where those bytes can be reproduced it gets identical ones; on any other GPU or PyTorch version the networks must agree within fixed near-identical bounds, and the package says that its output is near-identical, not identical. Strict mode (strict=True, the serving kit's --strict) refuses to run unless every batch shape gives identical bytes.

Where things are

Path What
crates/diar-exact numpy / scipy / pyannote / CPython numerics reproduced exactly
crates/diar-core the front-end around the networks, clustering, the decoder, the RTTM writer
crates/diar-audio libsndfile decoding as soundfile does it
crates/diar-engine the engine: pools, blocks, scheduling over GPUs, limits (docs/ENGINE.md)
crates/diar-triton, crates/diar-gateway the Triton backend and the HTTP gateway
crates/diar-py, python/anuvani_diar the Python API, CLI, client and the network backend
triton/ the two Triton images (cu121, cu128), the model files, the scripts of the gateway stack
triton_inference/ the serving kit: diar, the FastAPI service, Docker Compose, checks, load test, reference library
tools/ wheel build, staging, licence headers and notices, the oracles, probes
tests/ the suites (make test-fast, test-cache, test-gpu, test-gpu-all, triton/identity_check.sh)
docs/ ARCHITECTURE (start here), ENGINE, DETERMINISM, TESTING, SECURITY, WHEEL

Licence

The code is licensed under MIT OR Apache-2.0, at your option (LICENSE-MIT, LICENSE-APACHE). Code it copies or ports from other projects (DiariZen, torchaudio, pyannote.audio, WeSpeaker, VBx, SciPy, NumPy, python-soundfile), the native libraries the wheel bundles (OpenBLAS, libsndfile and their runtime libraries) and the Rust crates it compiles are listed with their licences in THIRD_PARTY_NOTICES.md; the texts are in LICENSES/. The model weights are not covered by the code's licence: the segmentation network is CC BY-NC 4.0 (non-commercial; a fine-tuned derivative of DiariZen's weights), the PLDA model and the speaker-embedding network CC BY 4.0.