{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "iyLoWDsb9rEs"
   },
   "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 = 'r1.17.0'\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>=0.10.0 -f https://download.pytorch.org/whl/torch_stable.html\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "oDzak_FIB9LS"
   },
   "source": [
    "# **SPEAKER RECOGNITION** \n",
    "Speaker Recognition (SR) is a broad research area that solves two major tasks: speaker identification (who is speaking?) and\n",
    "speaker verification (is the speaker who they claim to be?). In this work, we focus on text-independent speaker recognition when the identity of the speaker is based on how the speech is spoken,\n",
    "not necessarily in what is being said. Typically such SR systems operate on unconstrained speech utterances,\n",
    "which are converted into fixed-length vectors, called speaker embeddings. Speaker embeddings are also used in\n",
    "automatic speech recognition (ASR) and speech synthesis."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ydqmdcDxCeXb"
   },
   "source": [
    "In this tutorial, we shall first train these embeddings on speaker-related datasets, and then get speaker embeddings from a pretrained network for a new dataset. Since Google Colab has very slow read-write speeds, I'll be demonstrating this tutorial using [an4](http://www.speech.cs.cmu.edu/databases/an4/). \n",
    "\n",
    "Instead, if you'd like to try on a bigger dataset like [hi-mia](https://arxiv.org/abs/1912.01231) use the [get_hi-mia-data.py](https://github.com/NVIDIA/NeMo/tree/main/scripts/dataset_processing/speaker_tasks/get_hi-mia_data.py) script to download the necessary files, extract them, and resample to 16Khz if any of these samples are not at 16Khz. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "vqUBayc_Ctcr"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "NEMO_ROOT = os.getcwd()\n",
    "print(NEMO_ROOT)\n",
    "import glob\n",
    "import subprocess\n",
    "import tarfile\n",
    "import wget\n",
    "\n",
    "data_dir = os.path.join(NEMO_ROOT,'data')\n",
    "os.makedirs(data_dir, exist_ok=True)\n",
    "\n",
    "# Download the dataset. This will take a few moments...\n",
    "print(\"******\")\n",
    "if not os.path.exists(data_dir + '/an4_sphere.tar.gz'):\n",
    "    an4_url = 'https://dldata-public.s3.us-east-2.amazonaws.com/an4_sphere.tar.gz'  # for the original source, please visit http://www.speech.cs.cmu.edu/databases/an4/an4_sphere.tar.gz \n",
    "    an4_path = wget.download(an4_url, data_dir)\n",
    "    print(f\"Dataset downloaded at: {an4_path}\")\n",
    "else:\n",
    "    print(\"Tarfile already exists.\")\n",
    "    an4_path = data_dir + '/an4_sphere.tar.gz'\n",
    "\n",
    "# Untar and convert .sph to .wav (using sox)\n",
    "tar = tarfile.open(an4_path)\n",
    "tar.extractall(path=data_dir)\n",
    "\n",
    "print(\"Converting .sph to .wav...\")\n",
    "sph_list = glob.glob(data_dir + '/an4/**/*.sph', recursive=True)\n",
    "for sph_path in sph_list:\n",
    "    wav_path = sph_path[:-4] + '.wav'\n",
    "    cmd = [\"sox\", sph_path, wav_path]\n",
    "    subprocess.run(cmd)\n",
    "print(\"Finished conversion.\\n******\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "t5PrWzkiDbHy"
   },
   "source": [
    "Since an4 is not designed for speaker recognition, this facilitates the opportunity to demonstrate how you can generate manifest files that are necessary for training. These methods can be applied to any dataset to get similar training manifest files. \n",
    "\n",
    "First, create a list file which has all the wav files with absolute paths for each of the train, dev, and test set. This can be easily done by the `find` bash command"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "vnrUh3vuDSRN"
   },
   "outputs": [],
   "source": [
    "!find {data_dir}/an4/wav/an4_clstk  -iname \"*.wav\" > data/an4/wav/an4_clstk/train_all.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "BhWVg2QoDhL3"
   },
   "source": [
    "Let's look at the first 3 lines of filelist text file for train."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "BfnMK302Du20"
   },
   "outputs": [],
   "source": [
    "!head -n 3 {data_dir}/an4/wav/an4_clstk/train_all.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Y9L9Tl0XDw5Z"
   },
   "source": [
    "Since we created the list text file for the train, we use `filelist_to_manifest.py` to convert this text file to a manifest file and then optionally split the files to train \\& dev for evaluating the models during training by using the `--split` flag. We wouldn't be needing the `--split` option for the test folder. \n",
    "Accordingly please mention the `id` number, which is the field num separated by `/` to be considered as the speaker label "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "_LYwHAr1G8hp"
   },
   "source": [
    "After the download and conversion, your `data` folder should contain directories with manifest files as:\n",
    "\n",
    "* `data/<path>/train.json`\n",
    "* `data/<path>/dev.json` \n",
    "* `data/<path>/train_all.json` \n",
    "\n",
    "Each line in the manifest file describes a training sample - `audio_filepath` contains the path to the wav file, `duration` it's duration in seconds, and `label` is the speaker class label:\n",
    "\n",
    "`{\"audio_filepath\": \"<absolute path to dataset>data/an4/wav/an4test_clstk/menk/cen4-menk-b.wav\", \"duration\": 3.9, \"label\": \"menk\"}` "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "mpAv77JoD98c"
   },
   "outputs": [],
   "source": [
    "if not os.path.exists('scripts'):\n",
    "  print(\"Downloading necessary scripts\")\n",
    "  !mkdir -p scripts/speaker_tasks\n",
    "  !wget -P scripts/speaker_tasks/ https://raw.githubusercontent.com/NVIDIA/NeMo/$BRANCH/scripts/speaker_tasks/filelist_to_manifest.py\n",
    "!python {NEMO_ROOT}/scripts/speaker_tasks/filelist_to_manifest.py --filelist {data_dir}/an4/wav/an4_clstk/train_all.txt --id -2 --out {data_dir}/an4/wav/an4_clstk/all_manifest.json --split"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "5kPCmx5DHvY5"
   },
   "source": [
    "Generate the list text file for the test folder and then convert it to a manifest."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "nMd24GVaFBwr"
   },
   "outputs": [],
   "source": [
    "!find {data_dir}/an4/wav/an4test_clstk  -iname \"*.wav\" > {data_dir}/an4/wav/an4test_clstk/test_all.txt\n",
    "!python {NEMO_ROOT}/scripts/speaker_tasks/filelist_to_manifest.py --filelist {data_dir}/an4/wav/an4test_clstk/test_all.txt --id -2 --out {data_dir}/an4/wav/an4test_clstk/test.json"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "H5FPmxUkGakD"
   },
   "source": [
    "## Path to manifest files\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "vo-VnYPtJO_v"
   },
   "outputs": [],
   "source": [
    "train_manifest = os.path.join(data_dir,'an4/wav/an4_clstk/train.json')\n",
    "validation_manifest = os.path.join(data_dir,'an4/wav/an4_clstk/dev.json')\n",
    "test_manifest = os.path.join(data_dir,'an4/wav/an4_clstk/dev.json')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "KyDVdtjAL2__"
   },
   "source": [
    "As the goal of most speaker-related systems is to get good speaker level embeddings that could help distinguish from\n",
    "other speakers, we shall first train these embeddings in an end-to-end\n",
    "manner optimizing the [TitaNet](https://arxiv.org/pdf/2110.04410.pdf) model.\n",
    "We modify the decoder to get these fixed-size embeddings irrespective of the length of the input audio."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "OJtU_GEdMUUo"
   },
   "source": [
    "# Training\n",
    "Import necessary packages"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note: All the following steps are just for explanation of each section, but one can use the provided [training script](https://github.com/NVIDIA/NeMo/blob/main/examples/speaker_tasks/recognition/speaker_reco.py) to launch training in the command line."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "o1ojB0cZMSmv"
   },
   "outputs": [],
   "source": [
    "import nemo\n",
    "# NeMo's ASR collection - This collection contains complete ASR models and\n",
    "# building blocks (modules) for ASR\n",
    "import nemo.collections.asr as nemo_asr\n",
    "from omegaconf import OmegaConf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "m5Zho11LNAFJ"
   },
   "source": [
    "## Model Configuration \n",
    "The TitaNet model is defined in a config file which declares multiple important sections.\n",
    "\n",
    "They are:\n",
    "\n",
    "1) model: All arguments that will relate to the Model - preprocessors, encoder, decoder, optimizer and schedulers, datasets, and any other related information\n",
    "\n",
    "2) trainer: Any argument to be passed to PyTorch Lightning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "6HQtZfKnMhpI"
   },
   "outputs": [],
   "source": [
    "# This line will print the entire config of sample TitaNet model\n",
    "!mkdir conf \n",
    "!wget -P conf https://raw.githubusercontent.com/NVIDIA/NeMo/$BRANCH/examples/speaker_tasks/recognition/conf/titanet-large.yaml\n",
    "MODEL_CONFIG = os.path.join(NEMO_ROOT,'conf/titanet-large.yaml')\n",
    "config = OmegaConf.load(MODEL_CONFIG)\n",
    "print(OmegaConf.to_yaml(config))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "HtbXN-cFOwxi"
   },
   "source": [
    "## Setting up the datasets within the config\n",
    "If you'll notice, there are a few config dictionaries called train_ds, validation_ds and test_ds. These are configurations used to setup the Dataset and DataLoaders of the corresponding config."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "NPBIf1jmNgjn"
   },
   "outputs": [],
   "source": [
    "print(OmegaConf.to_yaml(config.model.train_ds))\n",
    "print(OmegaConf.to_yaml(config.model.validation_ds))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "PLIjKOMUP0YE"
   },
   "source": [
    "You will often notice that some configs have ??? in place of paths. This is used as a placeholder so that the user can change the value at a later time.\n",
    "\n",
    "Let's add the paths to the manifests to the config above\n",
    "Also, since an4 dataset doesn't have a test set of the same speakers used in training, we will use validation manifest as test manifest for demonstration purposes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "TSotpjL_O2BN"
   },
   "outputs": [],
   "source": [
    "config.model.train_ds.manifest_filepath = train_manifest\n",
    "config.model.validation_ds.manifest_filepath = validation_manifest"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note: Since we are training speaker embedding extractor model for verification we do not add test_ds dataset. To include it add it to config and replace manifest file as \n",
    "`config.model.test_ds.manifest_filepath = test_manifest`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "xy6_Lf6fW9aJ"
   },
   "source": [
    "Also as we are training on an4 dataset, there are 74 speaker labels in training, and we need to set this in the decoder config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "-B96tFTnW8Yh"
   },
   "outputs": [],
   "source": [
    "config.model.decoder.num_classes = 74"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "83pHBRDpQTF0"
   },
   "source": [
    "## Building the PyTorch Lightning Trainer\n",
    "NeMo models are primarily PyTorch Lightning modules - and therefore are entirely compatible with the PyTorch Lightning ecosystem!\n",
    "\n",
    "Let us first instantiate a Trainer object!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "GWzGJoHMQQnG"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import pytorch_lightning as pl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "WIYf4-KFQYHl"
   },
   "outputs": [],
   "source": [
    "print(\"Trainer config - \\n\")\n",
    "print(OmegaConf.to_yaml(config.trainer))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "aXuSMYMNQeW7"
   },
   "outputs": [],
   "source": [
    "# Let us modify some trainer configs for this demo\n",
    "# Checks if we have GPU available and uses it\n",
    "accelerator = 'gpu' if torch.cuda.is_available() else 'cpu'\n",
    "config.trainer.devices = 1\n",
    "config.trainer.accelerator = accelerator\n",
    "\n",
    "# Reduces maximum number of epochs to 5 for quick demonstration\n",
    "config.trainer.max_epochs = 10\n",
    "\n",
    "# Remove distributed training flags\n",
    "config.trainer.strategy = None\n",
    "\n",
    "# Remove augmentations\n",
    "config.model.train_ds.augmentor=None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "pBq3eCLwQhCy"
   },
   "outputs": [],
   "source": [
    "trainer = pl.Trainer(**config.trainer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "-xHq_rcmQiry"
   },
   "source": [
    "## Setting up a NeMo Experiment\n",
    "NeMo has an experiment manager that handles logging and checkpointing for us, so let's use it !"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "DMm8MPYfQsCS"
   },
   "outputs": [],
   "source": [
    "from nemo.utils.exp_manager import exp_manager\n",
    "log_dir = exp_manager(trainer, config.get(\"exp_manager\", None))\n",
    "# The log_dir provides a path to the current logging directory for easy access\n",
    "print(log_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "nQQMlXmLQ7h1"
   },
   "source": [
    "## Building the TitaNet Model\n",
    "TitaNet is a speaker embedding extractor model that can be used for speaker identification tasks - it generates one label for the entire provided audio stream. Therefore we encapsulate it inside the EncDecSpeakerLabelModel as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "E_KY_s5LROYf"
   },
   "outputs": [],
   "source": [
    "speaker_model = nemo_asr.models.EncDecSpeakerLabelModel(cfg=config.model, trainer=trainer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "_AphpMhkSVdU"
   },
   "source": [
    "Before we begin training, let us first create a Tensorboard visualization to monitor progress"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "BUnDpe_5SbDR"
   },
   "outputs": [],
   "source": [
    "try:\n",
    "  from google import colab\n",
    "  COLAB_ENV = True\n",
    "except (ImportError, ModuleNotFoundError):\n",
    "  COLAB_ENV = False\n",
    "\n",
    "# Load the TensorBoard notebook extension\n",
    "if COLAB_ENV:\n",
    "  %load_ext tensorboard\n",
    "  %tensorboard --logdir {exp_dir}\n",
    "else:\n",
    "  print(\"To use tensorboard, please use this notebook in a Google Colab environment.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Or8g1cksSf8C"
   },
   "source": [
    "As any NeMo model is inherently a PyTorch Lightning Model, it can easily be trained in a single line - trainer.fit(model)!\n",
    "Below we see that the model begins to get modest scores on the validation set after just 5 epochs of training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "HvYhsOWuSpL_",
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "trainer.fit(speaker_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "lSRACGt3UAYn"
   },
   "source": [
    "This config is not suited and designed for an4 so you may observe unstable val_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "jvtVKO8FZsoe"
   },
   "source": [
    "If you have a test manifest file, we can easily compute test accuracy by running\n",
    "<pre><code>trainer.test(speaker_model, ckpt_path=None)\n",
    "</code></pre>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "FlBwMsRdZfqg"
   },
   "source": [
    "## For Faster Training\n",
    "We can dramatically improve the time taken to train this model by using Multi GPU training along with Mixed Precision.\n",
    "\n",
    "### Trainer with a distributed backend:\n",
    "<pre><code>trainer = Trainer(devices=2, num_nodes=2, accelerator='gpu', strategy='dp')\n",
    "</code></pre>\n",
    "\n",
    "### Mixed precision:\n",
    "<pre><code>trainer = Trainer(amp_level='O1', precision=16)\n",
    "</code></pre>\n",
    "\n",
    "Of course, you can combine these flags as well."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "XcnWub9-0TW2"
   },
   "source": [
    "## Saving/Restoring a checkpoint\n",
    "There are multiple ways to save and load models in NeMo. Since all NeMo models are inherently Lightning Modules, we can use the standard way that PyTorch Lightning saves and restores models.\n",
    "\n",
    "NeMo also provides a more advanced model save/restore format, which encapsulates all the parts of the model that are required to restore that model for immediate use.\n",
    "\n",
    "In this example, we will explore both ways of saving and restoring models, but we will focus on the PyTorch Lightning method.\n",
    "\n",
    "## Saving and Restoring via PyTorch Lightning Checkpoints\n",
    "When using NeMo for training, it is advisable to utilize the exp_manager framework. It is tasked with handling checkpointing and logging (Tensorboard as well as WandB optionally!), as well as dealing with multi-node and multi-GPU logging.\n",
    "\n",
    "Since we utilized the exp_manager framework above, we have access to the directory where the checkpoints exist.\n",
    "\n",
    "exp_manager with the default settings will save multiple checkpoints for us -\n",
    "\n",
    "1) A few checkpoints from certain steps of training. They will have --val_loss= tags\n",
    "\n",
    "2) Checkpoints at the last epoch of training are denoted by --last.\n",
    "\n",
    "3) If the model finishes training, it will also have a --last checkpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "QSLjq-edaPt_"
   },
   "outputs": [],
   "source": [
    "# Let us list all the checkpoints we have\n",
    "checkpoint_dir = os.path.join(log_dir, 'checkpoints')\n",
    "checkpoint_paths = list(glob.glob(os.path.join(checkpoint_dir, \"*.ckpt\")))\n",
    "checkpoint_paths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "BwltdVWXaroa"
   },
   "outputs": [],
   "source": [
    "final_checkpoint = list(filter(lambda x: \"-last.ckpt\" in x, checkpoint_paths))[0]\n",
    "print(final_checkpoint)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "1tGKKojs0fEh"
   },
   "source": [
    "\n",
    "## Restoring from a PyTorch Lightning checkpoint\n",
    "To restore a model using the LightningModule.load_from_checkpoint() class method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "EgyP9cYVbFc8"
   },
   "outputs": [],
   "source": [
    "restored_model = nemo_asr.models.EncDecSpeakerLabelModel.load_from_checkpoint(final_checkpoint)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "AnZVMKZpbI_M"
   },
   "source": [
    "# Finetuning\n",
    "Since we don't have any new manifest file to finetune, I will demonstrate here by using the test manifest file we created earlier. \n",
    "an4 test dataset has a different set of speakers from the train set (total number: 10). And as we didn't split this dataset for validation I will use the same for validation. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "kV9gInFwQ2F5"
   },
   "source": [
    "There are a couple of ways we can finetune a speaker recognition model. \n",
    "1. Finetuning using a pretrained model published on NGC. \n",
    "2. Finetuning from a PTL checkpoint. \n",
    "\n",
    "Since finetuning from a large pretrained model is more common, I shall use it to demonstrate finetuning procedure. In order to make finetuning step independent from training from scratch, we use another config. Here we shall use `titanet-finetune.yaml` config, that is created to show finetuning on pretrained titanet-large model. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note: You may use [finetune-script](https://github.com/NVIDIA/NeMo/blob/main/examples/speaker_tasks/recognition/speaker_reco_finetune.py) to launch training in the command line. Following is just a demonstration of the script"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget -P conf https://raw.githubusercontent.com/NVIDIA/NeMo/$BRANCH/examples/speaker_tasks/recognition/conf/titanet-finetune.yaml\n",
    "MODEL_CONFIG = os.path.join(NEMO_ROOT,'conf/titanet-finetune.yaml')\n",
    "finetune_config = OmegaConf.load(MODEL_CONFIG)\n",
    "print(OmegaConf.to_yaml(finetune_config))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For step 2, if one would like to finetune from a PTL checkpoint, `init_from_pretrained_model` in config should be replaced with `init_from_nemo_model` and need to provide the path to checkpoint. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "HtXUWmYLQ0PJ"
   },
   "outputs": [],
   "source": [
    "test_manifest = os.path.join(data_dir,'an4/wav/an4test_clstk/test.json')\n",
    "finetune_config.model.train_ds.manifest_filepath = test_manifest\n",
    "finetune_config.model.validation_ds.manifest_filepath = test_manifest\n",
    "finetune_config.model.decoder.num_classes = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "IHy1zE1cTDZn"
   },
   "source": [
    "So we have set up the data and changed the decoder required for finetune, we now just need to create a trainer and start training with a smaller learning rate for fewer epochs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "nBmF6tQITSRl"
   },
   "outputs": [],
   "source": [
    "# Setup the new trainer object\n",
    "# Let us modify some trainer configs for this demo\n",
    "# Checks if we have GPU available and uses it\n",
    "accelerator = 'gpu' if torch.cuda.is_available() else 'cpu'\n",
    "\n",
    "trainer_config = OmegaConf.create(dict(\n",
    "    devices=1,\n",
    "    accelerator=accelerator,\n",
    "    max_epochs=5,\n",
    "    max_steps=-1,  # computed at runtime if not set\n",
    "    num_nodes=1,\n",
    "    accumulate_grad_batches=1,\n",
    "    enable_checkpointing=False,  # Provided by exp_manager\n",
    "    logger=False,  # Provided by exp_manager\n",
    "    log_every_n_steps=1,  # Interval of logging.\n",
    "    val_check_interval=1.0,  # Set to 0.25 to check 4 times per epoch, or an int for number of iterations\n",
    "))\n",
    "print(OmegaConf.to_yaml(trainer_config))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "bRz-8-xzUHKZ"
   },
   "outputs": [],
   "source": [
    "trainer_finetune = pl.Trainer(**trainer_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "EOwHTkW-UUy8"
   },
   "source": [
    "## Setting the trainer to the restored model\n",
    "Setting the trainer to the restored model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "0FhYQQQOUPIk"
   },
   "outputs": [],
   "source": [
    "log_dir_finetune = exp_manager(trainer_finetune, config.get(\"exp_manager\", None))\n",
    "print(log_dir_finetune)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "lc3fzGYVVTyi"
   },
   "source": [
    "## Fine-tune training step\n",
    "\n",
    "When fine-tuning on a truly new dataset, we will not see such a dramatic improvement in performance. However, it should still converge a little faster than if it was trained from scratch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "speaker_model = nemo_asr.models.EncDecSpeakerLabelModel(cfg=finetune_config.model, trainer=trainer_finetune)\n",
    "speaker_model.maybe_init_from_pretrained_checkpoint(finetune_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the config, we keep weights of preprocessor and encoder, and attach a new decoder as mentioned in above section to match num of classes of new data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "uFIOsuFYVLzr"
   },
   "outputs": [],
   "source": [
    "## Fine-tuning for 5 epochs¶\n",
    "trainer_finetune.fit(speaker_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tip: Add more data augmentation and dropout while finetuning on your data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "5DNidtl4VplU"
   },
   "source": [
    "# Saving .nemo file\n",
    "Now we can save the whole config and model parameters in a single .nemo and we can anytime restore from it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "am5wej6-VdZW"
   },
   "outputs": [],
   "source": [
    "restored_model.save_to(os.path.join(log_dir_finetune, '..',\"titanet-large-finetune.nemo\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "WnBhFJefV-Pf"
   },
   "outputs": [],
   "source": [
    "!ls {log_dir_finetune}/.."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "kVx1hNP_V_iz"
   },
   "outputs": [],
   "source": [
    "# restore from a save model\n",
    "restored_model = nemo_asr.models.EncDecSpeakerLabelModel.restore_from(os.path.join(log_dir_finetune, '..', \"titanet-large-finetune.nemo\"))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "80tLWTN40uaB"
   },
   "source": [
    "# Speaker Verification"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "VciRUIRz0y6P"
   },
   "source": [
    "Training for a speaker verification model is almost the same as the speaker recognition model with a change in the loss function. Angular Loss is a better function to train for a speaker verification model as the model is trained in an end-to-end manner with loss optimizing for embeddings cluster to be far from each other for different speaker by maximizing the angle between these clusters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ULTjBuFI19Js"
   },
   "source": [
    "To train for verification we just need to toggle `angular` flag in `config.model.decoder.params.angular = True` else set it to `False` to train with cross-entropy loss for identification purposes. \n",
    "Once we set this, the loss will be changed to angular loss and we can follow the above steps to the model.\n",
    "Note the scale and margin values to be set for the loss function are present at `config.model.loss.scale` and `config.model.loss.margin`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "LcKiNEY032-t"
   },
   "source": [
    "## Extract Speaker Embeddings\n",
    "Once you have a trained model or use one of our pretrained nemo checkpoints to get speaker embeddings for any speaker.\n",
    "\n",
    "To demonstrate this we shall use `nemo_asr.models.EncDecSpeakerLabelModel` with say 5 audio_samples from our dev manifest set. This model is specifically for inference purposes to extract embeddings from a trained `.nemo` model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "uXEzKMHf3r6-"
   },
   "outputs": [],
   "source": [
    "verification_model = nemo_asr.models.EncDecSpeakerLabelModel.restore_from(os.path.join(log_dir_finetune, '..', 'titanet-large-finetune.nemo'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Y-XiLHMQ8BIk"
   },
   "source": [
    "Now, we need to pass the necessary manifest_filepath and params to set up the data loader for extracting embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "lk2vsDJk9PS8"
   },
   "outputs": [],
   "source": [
    "!head -5 {validation_manifest} > embeddings_manifest.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "DEd5poCr9yrP"
   },
   "outputs": [],
   "source": [
    "config.model.train_ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nemo.collections.asr.parts.utils.speaker_utils import embedding_normalize\n",
    "from tqdm import  tqdm\n",
    "try:\n",
    "    from torch.cuda.amp import autocast\n",
    "except ImportError:\n",
    "    from contextlib import contextmanager\n",
    "\n",
    "    @contextmanager\n",
    "    def autocast(enabled=None):\n",
    "        yield\n",
    "import numpy as np\n",
    "import json\n",
    "import pickle as pkl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "JIHok6LD8g0F"
   },
   "outputs": [],
   "source": [
    "def get_embeddings(speaker_model, manifest_file, batch_size=1, embedding_dir='./', device='cuda'):\n",
    "    test_config = OmegaConf.create(\n",
    "        dict(\n",
    "            manifest_filepath=manifest_file,\n",
    "            sample_rate=16000,\n",
    "            labels=None,\n",
    "            batch_size=batch_size,\n",
    "            shuffle=False,\n",
    "            time_length=20,\n",
    "        )\n",
    "    )\n",
    "\n",
    "    speaker_model.setup_test_data(test_config)\n",
    "    speaker_model = speaker_model.to(device)\n",
    "    speaker_model.eval()\n",
    "\n",
    "    all_embs=[]\n",
    "    out_embeddings = {}\n",
    "           \n",
    "    for test_batch in tqdm(speaker_model.test_dataloader()):\n",
    "        test_batch = [x.to(device) for x in test_batch]\n",
    "        audio_signal, audio_signal_len, labels, slices = test_batch\n",
    "        with autocast():\n",
    "            _, embs = speaker_model.forward(input_signal=audio_signal, input_signal_length=audio_signal_len)\n",
    "            emb_shape = embs.shape[-1]\n",
    "            embs = embs.view(-1, emb_shape)\n",
    "            all_embs.extend(embs.cpu().detach().numpy())\n",
    "        del test_batch\n",
    "\n",
    "    all_embs = np.asarray(all_embs)\n",
    "    all_embs = embedding_normalize(all_embs)\n",
    "    with open(manifest_file, 'r') as manifest:\n",
    "        for i, line in enumerate(manifest.readlines()):\n",
    "            line = line.strip()\n",
    "            dic = json.loads(line)\n",
    "            uniq_name = '@'.join(dic['audio_filepath'].split('/')[-3:])\n",
    "            out_embeddings[uniq_name] = all_embs[i]\n",
    "\n",
    "    embedding_dir = os.path.join(embedding_dir, 'embeddings')\n",
    "    if not os.path.exists(embedding_dir):\n",
    "        os.makedirs(embedding_dir, exist_ok=True)\n",
    "\n",
    "    prefix = manifest_file.split('/')[-1].rsplit('.', 1)[-2]\n",
    "\n",
    "    name = os.path.join(embedding_dir, prefix)\n",
    "    embeddings_file = name + '_embeddings.pkl'\n",
    "    pkl.dump(out_embeddings, open(embeddings_file, 'wb'))\n",
    "    print(\"Saved embedding files to {}\".format(embedding_dir))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "u2FRecqD-ln5"
   },
   "outputs": [],
   "source": [
    "manifest_filepath = os.path.join(NEMO_ROOT,'embeddings_manifest.json')\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "get_embeddings(verification_model, manifest_filepath, batch_size=64,embedding_dir='./', device=device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "zfjXPsjzDOgr"
   },
   "source": [
    "Embeddings are stored in dict structure with key-value pair, key being uniq_name generated based on audio_filepath of the sample present in manifest_file in `embedding_dir`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "hmTeSR6jD28k"
   },
   "outputs": [],
   "source": [
    "ls ./embeddings/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "Speaker_Recogniton_Verification.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}