{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Anuvani-Diarization: who spoke when, in a few lines\n",
    "\n",
    "Speaker diarization for the 22 scheduled languages of India and English. This notebook installs the `anuvani_diar`\n",
    "package, diarizes one recording and saves the result as an RTTM file.\n",
    "\n",
    "- **Runtime:** choose a GPU (*Runtime \u2192 Change runtime type \u2192 GPU*). A CPU also works, much more slowly."
   ],
   "id": "intro"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "!nvidia-smi -L || echo \"No GPU found: the networks will run on the CPU, much more slowly.\""
   ],
   "id": "gpu-check"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install\n",
    "\n",
    "The package is one wheel: the Rust engine (decoding, windowing, batching over every GPU, clustering, the RTTM) and\n",
    "the two networks in PyTorch. It uses the PyTorch already installed here."
   ],
   "id": "install-md"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "from huggingface_hub import hf_hub_download\n",
    "wheel = hf_hub_download(\"CIIL/Anuvani-Diarization\", \"wheels/anuvani_diar-1.0.0-cp310-abi3-manylinux_2_28_x86_64.whl\")\n",
    "%pip install -q {wheel}"
   ],
   "id": "install"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Diarize a recording\n",
    "\n",
    "The first `Diarizer()` downloads the model (about 307 MB, once) and checks every file against the package's own\n",
    "sha256 list, then loads the networks on every GPU and tests them against the release's outputs. Any audio length\n",
    "works; WAV, FLAC, OGG and MP3 are read directly."
   ],
   "id": "diarize-md"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# a recording from this repository, or your own file's path\n",
    "audio = hf_hub_download(\"CIIL/Anuvani-Diarization\", \"golden/audio/bengali_009.flac\")\n",
    "\n",
    "from anuvani_diar import Diarizer\n",
    "# every visible GPU (or the CPU); Diarizer(devices=2) uses the first two GPUs\n",
    "diarize = Diarizer()\n",
    "result = diarize(audio)\n",
    "\n",
    "print(\"speakers:\", result.speakers)\n",
    "for start, end, speaker in result.segments[:10]:\n",
    "    print(f\"{start:8.2f} {end:8.2f}  {speaker}\")\n",
    "result.save_rttm(\"bengali_009.rttm\")"
   ],
   "id": "diarize"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Many files\n",
    "\n",
    "Results arrive as each file finishes; every GPU is used. To choose: `Diarizer(devices=2)` for the first two GPUs,\n",
    "`Diarizer(devices=[0, 3])` for GPUs 0 and 3, `Diarizer(devices=\"cpu\")` for the CPU.\n",
    "\n",
    "```python\n",
    "for r in diarize.run([\"a.wav\", \"b.flac\", \"c.mp3\"], out_dir=\"rttm/\"):\n",
    "    print(r.uri, \"ok\" if r.ok else r.error, r.speakers if r.ok else \"\")\n",
    "```\n",
    "\n",
    "## Exactness\n",
    "\n",
    "On the setup the release's outputs were produced with (PyTorch 2.1.1 with CUDA 12.1, `env/canonical.sh` in `src/`),\n",
    "the output is byte-identical to the release. On any other GPU or PyTorch version, such as here, it is FP32 with the\n",
    "same settings but not byte-identical: the networks are checked at start-up against the release's own outputs and run\n",
    "only if they agree within fixed bounds."
   ],
   "id": "more-md"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "print(diarize.info()[\"exactness\"])"
   ],
   "id": "exactness"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optional: the install check\n",
    "\n",
    "Diarizes the release's 12 check recordings (55 minutes of audio, downloaded once) and compares the results with the\n",
    "release's fingerprints. On a setup other than the release's, identical bytes are not expected (exit status 3,\n",
    "with the network check's numbers)."
   ],
   "id": "check-md"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "!anuvani-diar-check"
   ],
   "id": "check"
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}