tutorials/speaker_tasks/Speaker_Diarization_Inference.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "You can run either this notebook locally (if you have all the dependencies and a GPU) or on Google Colab.\n",
    "\n",
    "Instructions for setting up Colab are as follows:\n",
    "1. Open a new Python 3 notebook.\n",
    "2. Import this notebook from GitHub (File -> Upload Notebook -> \"GITHUB\" tab -> copy/paste GitHub URL)\n",
    "3. Connect to an instance with a GPU (Runtime -> Change runtime type -> select \"GPU\" for hardware accelerator)\n",
    "4. Run this cell to set up dependencies.\n",
    "\"\"\"\n",
    "# If you're using Google Colab and not running locally, run this cell.\n",
    "\n",
    "## Install dependencies\n",
    "!pip install wget\n",
    "!apt-get install sox libsndfile1 ffmpeg\n",
    "!pip install text-unidecode\n",
    "\n",
    "# ## Install NeMo\n",
    "BRANCH = 'main'\n",
    "!python -m pip install git+https://github.com/NVIDIA/NeMo.git@$BRANCH#egg=nemo_toolkit[asr]\n",
    "\n",
    "## Install TorchAudio\n",
    "!pip install torchaudio -f https://download.pytorch.org/whl/torch_stable.html"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction to Speaker Diarization\n",
    "Speaker diarization is the task of segmenting audio recordings by speaker labels and answers the question \"Who Speaks When?\". A speaker diarization system consists of Voice Activity Detection (VAD) model to get the timestamps of audio where speech is being spoken ignoring the background and speaker embeddings model to get speaker embeddings on segments that were previously time stamped. These speaker embedding vectors are then grouped into clusters and the number of speakers is estimated by clustering algorithm. Finally, based on the speaker profiles created from clustering results, neural diarizer generates speaker labels including overlap speech. The below figure shows the data-flow of NeMo speaker diarization."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"images/diar_pipeline.png\" alt=\"diar_pipeline\" style=\"width: 800px;\"/>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This tutorial covers speaker diarization inference. We will cover how to setup configurations and launch NeMo speaker diarization system with a few different settings. NeMo speaker diarization pipeline includes the following steps as described in the above figure: VAD, Segmentation, Speaker Embedding Extraction, Clustering and Neural Diarizer. We will explain what each module does and we will run NeMo speaker diarization system on a small toy example. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### VAD for Speaker Diarization Tasks\n",
    "\n",
    "In NeMo toolkit we support two types of speaker diarization inference regarding VAD:   \n",
    "* **with oracle VAD**: use ground-truth speech/non-speech labels.   \n",
    "* **with system VAD**: use speech/non-speech labels generated by an actual VAD model.   \n",
    "\n",
    "We will first demonstrate how to perform diarization with a oracle VAD timestamps (we assume we already have speech timestamps) and pretrained speaker embedding extractor model which can be found in tutorial for [Speaker Identification and Verification in NeMo](https://github.com/NVIDIA/NeMo/blob/main/tutorials/speaker_tasks/Speaker_Identification_Verification.ipynb).\n",
    "\n",
    "In the following section, we will also show how to perform VAD and then diarization if ground truth timestamp speech were not available (non-oracle VAD). We also have tutorials for [VAD training in NeMo](https://github.com/NVIDIA/NeMo/blob/main/tutorials/asr/Voice_Activity_Detection.ipynb) and [online offline microphone inference](https://github.com/NVIDIA/NeMo/blob/main/tutorials/asr/Online_Offline_Microphone_VAD_Demo.ipynb), where you can custom your model and training/finetuning on your own data.\n",
    "\n",
    "For demonstration purposes we would be using simulated audio from [an4 dataset](http://www.speech.cs.cmu.edu/databases/an4/)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Multi-scale Approach for Segmentation, Speaker Embedding Extraction and Clustering\n",
    "\n",
    "#### Uniform Segmentation\n",
    "After the input signal goes through VAD module, we extract speaker embeddings from segmented audio then we extract speaker embedding vector from each and every segment. When we segment audio signal into short (0.5~3.0 sec) segments, we get speaker profile (speaker representation) for the specific segment. When it comes to segment length, there is trade-off between the quality of speaker representation and granularity (temporal resolution). \n",
    "\n",
    "#### Trade-off: Long VS Short Segment Length\n",
    "If we use long segments (e.g. longer than 2-3 seconds), we get fairly consistent and high-quality speaker representations but at the same time, we lose temporal resolution since we need to make a decision on the 2~3 second long segment which can lead to significant errors. On the other hand, if we use very short segments (0.2-0.5 sec), temporal resolution is superior but it is very challenging to extract reliable speaker characteristics from such short speech segments.\n",
    "\n",
    "#### Multiscale Segmentation\n",
    "In NeMo speaker diarization pipeline, we employ multi-scale approach to deal with such a trade-off between long and short segment lengths. We use multiple scales (segment lengths) and fuse the affinity values from each scale's result. An example of multi-scale segmentation looks like the following figure:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"images/multiscale_example.png\" alt=\"diar_pipeline\" style=\"width: 600px;\"/>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The finest scale, which has the shortest segment length, is called __base scale__ and base scale is assigned to the highest scale index. Note that we integrate the information from all scales but only make decisions based on base scale's segment range.\n",
    "\n",
    "During multi-scale segmentation process, the mapping among scales should be calculated. The middle point of each segment is considered as an anchor point and matched with other scales to have the shortest distance between two middle points from the two segments. In the above figure, the blue outline shows an example of how multi-scale segmentation and mapping is determined."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Weighted Sum of Scale-specific Affinity Matrices "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The information from each scale is combined by calculating the weighted sum of affininty matrix. An affinity matrix is calculated by cosine similarity value between all the segments (and corresponding embedding vectors) in that scale. Once affinity matrix for each scale is calculated, we calculate a weighted sum on all the affinity matrices calculated as in the below figure."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"images/affinity_matrix_fusion.png\" alt=\"diar_pipeline\" style=\"width: 800px;\"/>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The weighted sum is calculated by using `multiscale_weights` parameter. We feed the fused affinity matrix (weighted sum of affinity matrix) to clustering algorithm to group the speakers and count the number of speakers. Multi-scale approach not only reduces DER but also makes speaker counting more accurate during clustering process.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Configurations for Multiscale Diarization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " We use a default multi-scale setting in [diar_infer_telephonic.yaml](https://github.com/NVIDIA/NeMo/blob/main/examples/speaker_tasks/diarization/conf/inference/diar_infer_telephonic.yaml) which has 5 scales from 1.5 s to 0.5 s, 50% overlap and equal weights. Note that only the ratio between numbers in `multiscale_weights` since the fused affinity matrix is normalized. For example, \\[1,1,1,1,1\\] and \\[0.5,0.5,0.5,0.5,0.5\\] will lead to the exactly same result."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* `config.diarizer.speaker_embeddings.parameters.window_length_in_sec = [1.5,1.25,1.0,0.75,0.5]`\n",
    "* `config.diarizer.speaker_embeddings.parameters.shift_length_in_sec = [0.75,0.625,0.5,0.375,0.1]`\n",
    "* `diarizer.speaker_embeddings.parameters.multiscale_weights=[1,1,1,1,1]` "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that an MSDD model has a pre-defined set of multi-scale configurations and clustering should be done with the same multi-scale configuration."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Neural Diarizer: Multiscale Diarization Decoder"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Neural Diarizer\n",
    "In NeMo speaker diarization pipeline, the term **neural diarizer**  refers to trainable neural modules that estimate speaker labels from the given feature or audio input. Neural diarizer contrasts with **clustering diarizer** in a way that clustering diarizer is not a trainable module. Neural diarizer is needed to enable overlap-aware diarization, more improved accucy and joint training with speaker embedding models using multispeaker datasets (diarization training datasets).\n",
    "\n",
    "#### Multi-scale Diarization Decoder (MSDD)\n",
    "Currently, you can use Multi-scale Diarization Decoder (MSDD) model as a neural diarizer. MSDD models use clustering diarizer for obtaining the estimated speaker profile of each speaker and the estimated number of speakers. The below figure shows training and inference of MSDD model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"images/msdd_train_and_infer.png\" alt=\"MSDD_train_infer\" style=\"width: 800px;\"/>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "MSDD models employ pairwise (two-speaker) unit-model for both training and inference. While training, pairwise model is trained on data samples with two speakers or two-speaker subset from data samples with more than two speakers. \n",
    "In inference mode, we retrieve all possible pairs from the estimated number of speakers and average the results. For example, if there are four speakers `(A, B, C, D)`, we extract 6 pairs: `(A,B)`, `(A,C)`, `(A,D)`, `(B,C)`, `(B,D)`, `(C,D)`. Finally, the sigmoid outputs are averaged. In this way, MSDD can deal with flexible number of speakers using a pairwise model. \n",
    "\n",
    "The detailed information on MSDD model and model training can be found in tutorial on [Speaker Diarization Training](https://github.com/NVIDIA/NeMo/blob/main/tutorials/speaker_tasks/Speaker_Diarization_Training.ipynb). "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A toy example for speaker diarization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Download a toy example audio file (`an4_diarize_test.wav`) and its ground-truth label file (`an4_diarize_test.rttm`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import wget\n",
    "ROOT = os.getcwd()\n",
    "data_dir = os.path.join(ROOT,'data')\n",
    "os.makedirs(data_dir, exist_ok=True)\n",
    "an4_audio = os.path.join(data_dir,'an4_diarize_test.wav')\n",
    "an4_rttm = os.path.join(data_dir,'an4_diarize_test.rttm')\n",
    "if not os.path.exists(an4_audio):\n",
    "    an4_audio_url = \"https://nemo-public.s3.us-east-2.amazonaws.com/an4_diarize_test.wav\"\n",
    "    an4_audio = wget.download(an4_audio_url, data_dir)\n",
    "if not os.path.exists(an4_rttm):\n",
    "    an4_rttm_url = \"https://nemo-public.s3.us-east-2.amazonaws.com/an4_diarize_test.rttm\"\n",
    "    an4_rttm = wget.download(an4_rttm_url, data_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's plot and listen to the audio and visualize the RTTM speaker labels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import IPython\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import librosa\n",
    "\n",
    "sr = 16000\n",
    "signal, sr = librosa.load(an4_audio,sr=sr) \n",
    "\n",
    "fig,ax = plt.subplots(1,1)\n",
    "fig.set_figwidth(20)\n",
    "fig.set_figheight(2)\n",
    "plt.plot(np.arange(len(signal)),signal,'gray')\n",
    "fig.suptitle('Reference merged an4 audio', fontsize=16)\n",
    "plt.xlabel('time (secs)', fontsize=18)\n",
    "ax.margins(x=0)\n",
    "plt.ylabel('signal strength', fontsize=16);\n",
    "a,_ = plt.xticks();plt.xticks(a,a/sr);\n",
    "\n",
    "IPython.display.Audio(an4_audio)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We would use [pyannote_metrics](https://pyannote.github.io/pyannote-metrics/) for visualization and score calculation purposes. Hence all the labels in rttm formats would eventually be converted to pyannote objects, we created two helper functions rttm_to_labels (for NeMo intermediate processing) and labels_to_pyannote_object for scoring and visualization format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nemo.collections.asr.parts.utils.speaker_utils import rttm_to_labels, labels_to_pyannote_object"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's load ground truth RTTM labels and view the reference Annotation timestamps visually"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# view the sample rttm file\n",
    "!cat {an4_rttm}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "labels = rttm_to_labels(an4_rttm)\n",
    "reference = labels_to_pyannote_object(labels)\n",
    "print(labels)\n",
    "reference"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Speaker Diarization scripts commonly expects following arguments:\n",
    "1. manifest_filepath : Path to manifest file containing json lines of format:   \n",
    "`{'audio_filepath': /path/to/audio_file, 'offset': 0, 'duration':None, 'label': 'infer', 'text': '-', 'num_speakers': None, 'rttm_filepath': /path/to/rttm/file, 'uem_filepath'='/path/to/uem/filepath'}`\n",
    "2. out_dir : directory where outputs and intermediate files are stored. \n",
    "3. oracle_vad: If this is true then we extract speech activity labels from rttm files, if False then either \n",
    "4. vad.model_path or external_manifestpath containing speech activity labels has to be passed. \n",
    "\n",
    "Mandatory fields are audio_filepath, offset, duration, label and text. For the rest if you would like to evaluate with known number of speakers pass the value else None. If you would like to score the system with known rttms then that should be passed as well, else None. uem file is used to score only part of your audio for evaluation purposes, hence pass if you would like to evaluate on it else None.\n",
    "\n",
    "\n",
    "* **\\[Note\\]** we expect audio and corresponding RTTM have **same base name** and the name should be **unique**. \n",
    "\n",
    "For eg: if audio file name is **test_an4**.wav, if provided we expect corresponding rttm file name to be **test_an4**.rttm (note the matching **test_an4** base name)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lets create manifest with the an4 audio and rttm available. If you have more than one files you may also use the script `pathfiles_to_diarize_manifest.py` to generate manifest file from list of audio files and optionally rttm files "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a manifest for input with below format. \n",
    "# {'audio_filepath': /path/to/audio_file, 'offset': 0, 'duration':None, 'label': 'infer', 'text': '-', \n",
    "# 'num_speakers': None, 'rttm_filepath': /path/to/rttm/file, 'uem_filepath'='/path/to/uem/filepath'}\n",
    "import json\n",
    "meta = {\n",
    "    'audio_filepath': an4_audio, \n",
    "    'offset': 0, \n",
    "    'duration':None, \n",
    "    'label': 'infer', \n",
    "    'text': '-', \n",
    "    'num_speakers': 2, \n",
    "    'rttm_filepath': an4_rttm, \n",
    "    'uem_filepath' : None\n",
    "}\n",
    "with open('data/input_manifest.json','w') as fp:\n",
    "    json.dump(meta,fp)\n",
    "    fp.write('\\n')\n",
    "\n",
    "!cat data/input_manifest.json\n",
    "\n",
    "output_dir = os.path.join(ROOT, 'oracle_vad')\n",
    "os.makedirs(output_dir,exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Speaker Diarization with Oracle-VAD"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using **oracle VAD** for speaker diarization can be regarded as performing a diarization inference based on ground-truth speech/non-speech labels. The motivation behind using oracle-VAD is to factor out the influence of VAD performance when we evaluate a speaker diarization system. Speaker diarization with oracle-VAD can also be used to run speaker diarization with rttms generated from any external VAD, not just VAD model from NeMo.\n",
    "\n",
    "The first step is to start converting reference audio RTTM file (containing VAD output) timestamps to oracle manifest file. This manifest file would be sent to our speaker diarizer to extract embeddings.\n",
    "\n",
    "If you have RTTM files for your input audio files, setting `oracle_vad=True` in diarization inference config, the diarization system automatically computes oracle manifest based on the rttms provided through input manifest file."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our config file is based on [hydra](https://hydra.cc/docs/intro/). \n",
    "With hydra config, we ask users to provide values to variables that were filled with **???**, these are mandatory fields and scripts expect them for successful runs. Note that the variables filled with **null** are optional variables. Such variables could be provided if needed but are not mandatory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from omegaconf import OmegaConf\n",
    "MODEL_CONFIG = os.path.join(data_dir,'diar_infer_telephonic.yaml')\n",
    "if not os.path.exists(MODEL_CONFIG):\n",
    "    config_url = \"https://raw.githubusercontent.com/NVIDIA/NeMo/main/examples/speaker_tasks/diarization/conf/inference/diar_infer_telephonic.yaml\"\n",
    "    MODEL_CONFIG = wget.download(config_url,data_dir)\n",
    "\n",
    "config = OmegaConf.load(MODEL_CONFIG)\n",
    "print(OmegaConf.to_yaml(config))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can perform speaker diarization based on timestamps generated from ground truth rttms rather than generating through VAD. \n",
    "\n",
    "Let's set parameters for speaker diarization inference. We will use `titanet_large` speaker embedding model for running clustering diarizer. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config.diarizer.manifest_filepath = 'data/input_manifest.json'\n",
    "config.diarizer.out_dir = output_dir # Directory to store intermediate files and prediction outputs\n",
    "pretrained_speaker_model = 'titanet_large'\n",
    "config.diarizer.speaker_embeddings.model_path = pretrained_speaker_model\n",
    "config.diarizer.speaker_embeddings.parameters.window_length_in_sec = [1.5,1.25,1.0,0.75,0.5] \n",
    "config.diarizer.speaker_embeddings.parameters.shift_length_in_sec = [0.75,0.625,0.5,0.375,0.1] \n",
    "config.diarizer.speaker_embeddings.parameters.multiscale_weights= [1,1,1,1,1] \n",
    "config.diarizer.oracle_vad = True # ----> ORACLE VAD \n",
    "config.diarizer.clustering.parameters.oracle_num_speakers = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Clustering Diarizer: with Oracle VAD"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we set configurations, import `ClusteringDiarizer` class and create a clustering diarizer instance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nemo.collections.asr.models import ClusteringDiarizer\n",
    "oracle_vad_clusdiar_model = ClusteringDiarizer(cfg=config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# And lets diarize\n",
    "oracle_vad_clusdiar_model.diarize()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With DER 0 -> means it clustered speaker embeddings correctly. Let's view "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat {output_dir}/pred_rttms/an4_diarize_test.rttm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Visualize the diarization output of clustering diarizer and compare with the ground-truth speaker labels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Clustering Diarizer Result (RTTM format)\")\n",
    "pred_labels_neural = rttm_to_labels(f'{output_dir}/pred_rttms/an4_diarize_test.rttm')\n",
    "hypothesis_neural = labels_to_pyannote_object(pred_labels_neural)\n",
    "hypothesis_neural"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Ground-truth Speaker Labels\")\n",
    "reference"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Neural Diarizer: Multiscale Diarization Decoder with Oracle VAD"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The checkpoints (`.ckpt`) or NeMo files (`.nemo`) for **neural diarizers** contain all the necessary neural models for speaker diarization. For example, an MSDD model checkpoint or a NeMo file has pre-trained [TitaNet](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/nemo/models/titanet_large) and MSDD model itself. \n",
    "\n",
    "In this tutorial, we use [diar_msdd_telephonic](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/nemo/models/diar_msdd_telephonic) which is optimized for telephic speech. Since we share the same YAML file for all kinds of speaker diarization inference, all we need to do is add model path on top of the config setting for clustering diarizer.\n",
    "\n",
    "`sigmoid_threshold` is a threshold for making the final binary decision on overlapping speaker label. The lower the value is, the more generous on the speech overlap detection. `sigmoid_threshold` value affects false alarm and miss errors. Default value is `0.7` for telephonic model and if `sigmoid_threshold=1.0`, no overlap speech is detected."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config.diarizer.msdd_model.model_path = 'diar_msdd_telephonic' # Telephonic speaker diarization model \n",
    "config.diarizer.msdd_model.parameters.sigmoid_threshold = [0.7, 1.0] # Evaluate with T=0.7 and T=1.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nemo.collections.asr.models.msdd_models import NeuralDiarizer\n",
    "oracle_vad_msdd_model = NeuralDiarizer(cfg=config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Unlike clustering diarizer, neural diarizer evaluates diarization results in three different settings by default:\n",
    "\n",
    "* `collar=0.25`, `ignore_overlap=True`: This is the default setting for evaluating clustering diarizer.\n",
    "* `collar=0.25`, `ignore_overlap=False`: Still 0.25 s around boundaries are not evaluated but overlaps are evaluated.\n",
    "* `collar=0.0`, `ignore_overlap=False`: No collar at all and evaluate overlaps. \n",
    "\n",
    "Let's run the MSDD model with the prepared configurations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "oracle_vad_msdd_model.diarize()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The output of the neural diarizer is saved in `outputs/pred_rttms`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat {output_dir}/pred_rttms/an4_diarize_test.rttm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Visualize the diarization output of clustering diarizer and compare with the ground-truth speaker labels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Neural Diarizer Result (RTTM format)\")\n",
    "pred_labels_neural = rttm_to_labels(f'{output_dir}/pred_rttms/an4_diarize_test.rttm')\n",
    "hypothesis_neural = labels_to_pyannote_object(pred_labels_neural)\n",
    "hypothesis_neural"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Ground-truth Speaker Labels\")\n",
    "reference"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Speaker Diarization with System VAD (NeMo VAD models)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section, we actually compute VAD timestamps using a NeMo VAD model on the input manifest file. Next, we use these timestamps for speech/non-speech labels to extract speaker embedding vectors followed by clustering them into num of speakers. As opposed to oracle VAD, the result from an actual VAD model is referred to as __system VAD__."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before we proceed let's look at the speaker diarization config, which we would be depending up on for vad computation\n",
    "and speaker embedding extraction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(OmegaConf.to_yaml(config))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As can be seen most of the variables in config are self explanatory \n",
    "with VAD variables under vad section and speaker related variables under speaker embeddings section. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To perform VAD based diarization we can ignore `oracle_vad_manifest` in `speaker_embeddings` key for now and need to fill up the rest. We also needs to provide pretrained `model_path` of vad and speaker embeddings .nemo models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pretrained_vad = 'vad_multilingual_marblenet'\n",
    "pretrained_speaker_model = 'titanet_large'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note in this tutorial, we use the VAD model *vad_multilingual_marblenet* which is an improved model based on MarbleNet-3x2 that has been introduced and published in [ICASSP MarbleNet](https://arxiv.org/pdf/2010.13886.pdf). You might need to tune on dev set similar to your dataset if you would like to improve the performance.\n",
    "\n",
    "And the speakerNet-M-Diarization model achieves 7.3% confusion error rate on CH109 set with oracle vad. This model is trained on voxceleb1, voxceleb2, Fisher, SwitchBoard datasets. So for more improved performance specific to your dataset, finetune speaker verification model with a devset similar to your test set.\n",
    "\n",
    "It is recommended to set `num_workers=1` since using the multiprocessing package in Jupyter Notebook environment might cause freezing issues. For sizable data, run speaker diarization using the scripts in `NeMo/examples/speaker_tasks/` setting `num_workers` larger than 1 in the configurations.\n",
    "\n",
    "You can play with parameters in configurations as below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config.num_workers = 1 # Workaround for multiprocessing hanging with ipython issue \n",
    "\n",
    "output_dir = os.path.join(ROOT, 'outputs')\n",
    "config.diarizer.manifest_filepath = 'data/input_manifest.json'\n",
    "config.diarizer.out_dir = output_dir #Directory to store intermediate files and prediction outputs\n",
    "\n",
    "config.diarizer.speaker_embeddings.model_path = pretrained_speaker_model\n",
    "config.diarizer.oracle_vad = False # compute VAD provided with model_path to vad config\n",
    "config.diarizer.clustering.parameters.oracle_num_speakers=False\n",
    "\n",
    "# Here, we use our in-house pretrained NeMo VAD model\n",
    "config.diarizer.vad.model_path = pretrained_vad\n",
    "config.diarizer.vad.parameters.onset = 0.8\n",
    "config.diarizer.vad.parameters.offset = 0.6\n",
    "config.diarizer.vad.parameters.pad_offset = -0.05"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Clustering Diarizer: with System VAD"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we passed all the variables we need, let's initialize the clustering diarizer model with the configurations we set. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nemo.collections.asr.models import ClusteringDiarizer\n",
    "sd_model = ClusteringDiarizer(cfg=config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And launch diarization with a single line of code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sd_model.diarize()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As can be seen, we first performed VAD, then with the timestamps created in `{output_dir}/vad_outputs` by VAD we calculated speaker embeddings (`{output_dir}/speaker_outputs/embeddings/`) which are then clustered using spectral clustering. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Generating predicted  VAD timesteps: We perform VAD inference to have frame level prediction &#8594; (optional: use decision smoothing) &#8594; given `threshold`,  write speech segments to a RTTM-like timestamp manifest file.\n",
    "\n",
    "We use VAD decision smoothing (50% overlap median) as described in [vad_utils.py](https://github.com/NVIDIA/NeMo/blob/stable/nemo/collections/asr/parts/utils/vad_utils.py).\n",
    "\n",
    "You can also tune the threshold on your dev set. Use this provided in [vad_tune_threshold.py](https://github.com/NVIDIA/NeMo/blob/stable/scripts/voice_activity_detection/vad_tune_threshold.py)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# VAD predicted time stamps\n",
    "# you can also use single threshold(=onset=offset) for binarization and plot here\n",
    "from nemo.collections.asr.parts.utils.vad_utils import plot\n",
    "\n",
    "if config.diarizer.vad.parameters.smoothing:\n",
    "    vad_output_filepath = f'{output_dir}/vad_outputs/overlap_smoothing_output_median_{config.diarizer.vad.parameters.overlap}/an4_diarize_test.{config.diarizer.vad.parameters.smoothing}'\n",
    "else:\n",
    "    vad_output_filepath = f'{output_dir}/vad_outputs/an4_diarize_test.frame'\n",
    "\n",
    "plot(\n",
    "    an4_audio,\n",
    "    vad_output_filepath, \n",
    "    an4_rttm,\n",
    "    per_args = config.diarizer.vad.parameters, #threshold\n",
    "    ) \n",
    "\n",
    "print(f\"VAD params:{OmegaConf.to_yaml(config.diarizer.vad.parameters)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Predicted outputs are written to `{output_dir}/pred_rttms` and see how we predicted along with VAD prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat {output_dir}/pred_rttms/an4_diarize_test.rttm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Visualize the diarization output of clustering diarizer and compare with the ground-truth speaker labels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Clustering Diarizer Result (RTTM format)\")\n",
    "pred_labels_neural = rttm_to_labels(f'{output_dir}/pred_rttms/an4_diarize_test.rttm')\n",
    "hypothesis_neural = labels_to_pyannote_object(pred_labels_neural)\n",
    "hypothesis_neural\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Ground-truth Speaker Labels (RTTM format)\")\n",
    "reference"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Neural Diarizer: Multiscale Diarization Decoder with System VAD"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can run a neural diarizer model without changing many parameters since neural diarizer also performs VAD and clustering diarizer. All we need to do is run the MSDD model with the new config file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config.diarizer.msdd_model.model_path = 'diar_msdd_telephonic' # Telephonic speaker diarization model \n",
    "config.diarizer.msdd_model.parameters.sigmoid_threshold = [0.7, 1.0] # Evaluate with T=0.7 and T=1.0\n",
    "system_vad_msdd_model = NeuralDiarizer(cfg=config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "system_vad_msdd_model.diarize()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check whether diarization saved in `outputs/pred_rttms` is correct."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat {output_dir}/pred_rttms/an4_diarize_test.rttm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Visualize the diarization output of clustering diarizer and compare with the ground-truth speaker labels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Neural Diarizer Result (RTTM format)\")\n",
    "pred_labels_neural = rttm_to_labels(f'{output_dir}/pred_rttms/an4_diarize_test.rttm')\n",
    "hypothesis_neural = labels_to_pyannote_object(pred_labels_neural)\n",
    "hypothesis_neural"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Ground-truth Speaker Labels (RTTM format)\")\n",
    "reference"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Storing and Restoring models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For clustering diarizer, we can save the whole config and model parameters in a single .nemo and restore from it anytime. Neural diarizer will support this feature soon."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "oracle_vad_clusdiar_model.save_to(os.path.join(output_dir,'clustering_diarizer.nemo'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Restore from saved model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "del oracle_vad_clusdiar_model\n",
    "import nemo.collections.asr as nemo_asr\n",
    "restored_model = nemo_asr.models.ClusteringDiarizer.restore_from(os.path.join(output_dir,'clustering_diarizer.nemo'))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "metadata": {
     "collapsed": false
    },
    "source": []
   }
  },
  "vscode": {
   "interpreter": {
    "hash": "aee8b7b246df8f9039afb4144a1f6fd8d2ca17a180786b69acc140d282b71a49"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}