tutorials/speaker_tasks/ASR_with_SpeakerDiarization.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Automatic Speech Recognition with Speaker Diarization"
   ]
  },
  {
   "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 = '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 -f https://download.pytorch.org/whl/torch_stable.html"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Speaker diarization lets us figure out \"who spoke when\" in the transcription. Without speaker diarization, we cannot distinguish the speakers in the transcript generated from automatic speech recognition (ASR). Nowadays, ASR combined with speaker diarization has shown immense use in many tasks, ranging from analyzing meeting transcription to media indexing. \n",
    "\n",
    "In this tutorial, we demonstrate how we can get ASR transcriptions combined with speaker labels. Since we don't include a detailed process of getting ASR results or diarization results, please refer to the following links for more in-depth description.\n",
    "\n",
    "If you need detailed understanding of transcribing words with ASR, refer to this [ASR Tutorial](https://github.com/NVIDIA/NeMo/blob/stable/tutorials/asr/ASR_with_NeMo.ipynb) tutorial.\n",
    "\n",
    "\n",
    "For detailed parameter setting and execution of speaker diarization, refer to this [Diarization Inference](https://github.com/NVIDIA/NeMo/blob/stable/tutorials/speaker_tasks/Speaker_Diarization_Inference.ipynb) tutorial.\n",
    "\n",
    "\n",
    "An example script that runs ASR and speaker diarization together can be found at [ASR with Diarization](https://github.com/NVIDIA/NeMo/blob/main/examples/speaker_tasks/diarization/clustering_diarizer/offline_diar_with_asr_infer.py).\n",
    "\n",
    "### Speaker diarization in ASR pipeline\n",
    "\n",
    "Speaker diarization results in ASR pipeline should align well with ASR output. Thus, we use ASR output to create Voice Activity Detection (VAD) timestamps to obtain segments we want to diarize. The segments we obtain from the VAD timestamps are further segmented into sub-segments in the speaker diarization step. Finally, after obtaining the speaker labels from speaker diarization, we match the decoded words with speaker labels to generate a transcript with speaker labels.\n",
    "\n",
    "    ASR → VAD timestamps and decoded words → speaker diarization → speaker label matching"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nemo.collections.asr as nemo_asr\n",
    "import numpy as np\n",
    "from IPython.display import Audio, display\n",
    "import librosa\n",
    "import os\n",
    "import wget\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import nemo\n",
    "import glob\n",
    "\n",
    "import pprint\n",
    "pp = pprint.PrettyPrinter(indent=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We demonstrate this tutorial using a merged AN4 audioclip. The merged audioclip contains the speech of two speakers (male and female) reading dates in different formats. Run the following script to download the audioclip and play it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ROOT = os.getcwd()\n",
    "data_dir = os.path.join(ROOT,'data')\n",
    "os.makedirs(data_dir, exist_ok=True)\n",
    "\n",
    "an4_audio_url = \"https://nemo-public.s3.us-east-2.amazonaws.com/an4_diarize_test.wav\"\n",
    "if not os.path.exists(os.path.join(data_dir,'an4_diarize_test.wav')):\n",
    "    AUDIO_FILENAME = wget.download(an4_audio_url, data_dir)\n",
    "else:\n",
    "    AUDIO_FILENAME = os.path.join(data_dir,'an4_diarize_test.wav')\n",
    "\n",
    "audio_file_list = glob.glob(f\"{data_dir}/*.wav\")\n",
    "print(\"Input audio file list: \\n\", audio_file_list)\n",
    "\n",
    "signal, sample_rate = librosa.load(AUDIO_FILENAME, sr=None)\n",
    "display(Audio(signal,rate=sample_rate))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`display_waveform()` and `get_color()` functions are defined for displaying the waveform with diarization results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def display_waveform(signal,text='Audio',overlay_color=[]):\n",
    "    fig,ax = plt.subplots(1,1)\n",
    "    fig.set_figwidth(20)\n",
    "    fig.set_figheight(2)\n",
    "    plt.scatter(np.arange(len(signal)),signal,s=1,marker='o',c='k')\n",
    "    if len(overlay_color):\n",
    "        plt.scatter(np.arange(len(signal)),signal,s=1,marker='o',c=overlay_color)\n",
    "    fig.suptitle(text, fontsize=16)\n",
    "    plt.xlabel('time (secs)', fontsize=18)\n",
    "    plt.ylabel('signal strength', fontsize=14);\n",
    "    plt.axis([0,len(signal),-0.5,+0.5])\n",
    "    time_axis,_ = plt.xticks();\n",
    "    plt.xticks(time_axis[:-1],time_axis[:-1]/sample_rate);\n",
    "    \n",
    "COLORS=\"b g c m y\".split()\n",
    "\n",
    "def get_color(signal,speech_labels,sample_rate=16000):\n",
    "    c=np.array(['k']*len(signal))\n",
    "    for time_stamp in speech_labels:\n",
    "        start,end,label=time_stamp.split()\n",
    "        start,end = int(float(start)*16000),int(float(end)*16000),\n",
    "        if label == \"speech\":\n",
    "            code = 'red'\n",
    "        else:\n",
    "            code = COLORS[int(label.split('_')[-1])]\n",
    "        c[start:end]=code\n",
    "    \n",
    "    return c "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using the above function, we can display the waveform of the example audio clip. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display_waveform(signal)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Parameter setting for ASR and diarization\n",
    "First, we need to setup the following parameters for ASR and diarization. We start our demonstration by first transcribing the audio recording using our pretrained ASR model `QuartzNet15x5Base-En` and use the CTC output probabilities to get timestamps for the spoken words. We then use these timestamps to get speaker label information using the speaker diarizer model. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from omegaconf import OmegaConf\n",
    "import shutil\n",
    "DOMAIN_TYPE = \"meeting\" # Can be meeting or telephonic based on domain type of the audio file\n",
    "CONFIG_FILE_NAME = f\"diar_infer_{DOMAIN_TYPE}.yaml\"\n",
    "\n",
    "CONFIG_URL = f\"https://raw.githubusercontent.com/NVIDIA/NeMo/main/examples/speaker_tasks/diarization/conf/inference/{CONFIG_FILE_NAME}\"\n",
    "\n",
    "if not os.path.exists(os.path.join(data_dir,CONFIG_FILE_NAME)):\n",
    "    CONFIG = wget.download(CONFIG_URL, data_dir)\n",
    "else:\n",
    "    CONFIG = os.path.join(data_dir,CONFIG_FILE_NAME)\n",
    "\n",
    "cfg = OmegaConf.load(CONFIG)\n",
    "print(OmegaConf.to_yaml(cfg))"
   ]
  },
  {
   "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: `{\"audio_filepath\": \"/path/to/audio_file\", \"offset\": 0, \"duration\": null, \"label\": \"infer\", \"text\": \"-\", \"num_speakers\": null, \"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 a known number of speakers pass the value else `null`. If you would like to score the system with known rttms then that should be passed as well, else `null`. 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 `null`.\n",
    "\n",
    "\n",
    "**Note:** we expect audio and corresponding RTTM to have **same base name** and the name should be **unique**. \n",
    "\n",
    "For example: 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 a manifest file with the an4 audio and rttm available. If you have more than one file you may also use the script `NeMo/scripts/speaker_tasks/pathfiles_to_diarize_manifest.py` to generate a manifest file from a list of audio files. In addition, you can optionally include rttm files to evaluate the diarization results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a manifest file for input with below format. \n",
    "# {\"audio_filepath\": \"/path/to/audio_file\", \"offset\": 0, \"duration\": null, \"label\": \"infer\", \"text\": \"-\", \n",
    "# \"num_speakers\": null, \"rttm_filepath\": \"/path/to/rttm/file\", \"uem_filepath\"=\"/path/to/uem/filepath\"}\n",
    "import json\n",
    "meta = {\n",
    "    'audio_filepath': AUDIO_FILENAME, \n",
    "    'offset': 0, \n",
    "    'duration':None, \n",
    "    'label': 'infer', \n",
    "    'text': '-', \n",
    "    'num_speakers': None, \n",
    "    'rttm_filepath': None, \n",
    "    'uem_filepath' : None\n",
    "}\n",
    "with open(os.path.join(data_dir,'input_manifest.json'),'w') as fp:\n",
    "    json.dump(meta,fp)\n",
    "    fp.write('\\n')\n",
    "\n",
    "cfg.diarizer.manifest_filepath = os.path.join(data_dir,'input_manifest.json')\n",
    "!cat {cfg.diarizer.manifest_filepath}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's set the parameters required for diarization. In this tutorial, we obtain voice activity labels from ASR, which is set through parameter `cfg.diarizer.asr.parameters.asr_based_vad`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pretrained_speaker_model='titanet_large'\n",
    "cfg.diarizer.manifest_filepath = cfg.diarizer.manifest_filepath\n",
    "cfg.diarizer.out_dir = data_dir #Directory to store intermediate files and prediction outputs\n",
    "cfg.diarizer.speaker_embeddings.model_path = pretrained_speaker_model\n",
    "cfg.diarizer.clustering.parameters.oracle_num_speakers=False\n",
    "\n",
    "# Using Neural VAD and Conformer ASR \n",
    "cfg.diarizer.vad.model_path = 'vad_multilingual_marblenet'\n",
    "cfg.diarizer.asr.model_path = 'stt_en_conformer_ctc_large'\n",
    "cfg.diarizer.oracle_vad = False # ----> Not using oracle VAD \n",
    "cfg.diarizer.asr.parameters.asr_based_vad = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run ASR and get word timestamps\n",
    "Before we run speaker diarization, we should run ASR and get the ASR output to generate decoded words and timestamps for those words. Let's import `ASRDecoderTimeStamps` class and create `asr_decoder_ts` instance that returns an ASR model. Using this ASR model, the following two variables are obtained from `asr_decoder_ts.run_ASR()` function."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- word_hyp Dict[str, List[str]]: contains the sequence of words.\n",
    "- word_ts_hyp Dict[str, List[int]]: contains frame level index of the start and the end of each word."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nemo.collections.asr.parts.utils.decoder_timestamps_utils import ASRDecoderTimeStamps\n",
    "asr_decoder_ts = ASRDecoderTimeStamps(cfg.diarizer)\n",
    "asr_model = asr_decoder_ts.set_asr_model()\n",
    "word_hyp, word_ts_hyp = asr_decoder_ts.run_ASR(asr_model)\n",
    "\n",
    "print(\"Decoded word output dictionary: \\n\", word_hyp['an4_diarize_test'])\n",
    "print(\"Word-level timestamps dictionary: \\n\", word_ts_hyp['an4_diarize_test'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's create an instance `asr_diar_offline` from OfflineDiarWithASR class, which matches diarization results with ASR outputs. We pass ``cfg.diarizer`` to setup the parameters for both ASR and diarization. We also set `word_ts_anchor_offset` variable that determines the anchor position of each word. Here, we use the default value from `asr_decoder_ts` instance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nemo.collections.asr.parts.utils.diarization_utils import OfflineDiarWithASR\n",
    "asr_diar_offline = OfflineDiarWithASR(cfg.diarizer)\n",
    "asr_diar_offline.word_ts_anchor_offset = asr_decoder_ts.word_ts_anchor_offset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`asr_diar_offline` instance is now ready. As a next step, we run diarization."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run diarization with the extracted word timestamps"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that all the components for diarization is ready, let's run diarization by calling `run_diarization()` function. `run_diarization()` will return two different variables : `diar_hyp` and `diar_score`. `diar_hyp` is diarization inference result which is written in `[start time] [end time] [speaker]` format. `diar_score` contains `None` since we did not provide `rttm_filepath` in the input manifest file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "diar_hyp, diar_score = asr_diar_offline.run_diarization(cfg, word_ts_hyp)\n",
    "print(\"Diarization hypothesis output: \\n\", diar_hyp['an4_diarize_test'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`run_diarization()` function also creates `an4_diarize_test.rttm` file. Let's check what is written in this `rttm` file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_file(path_to_file):\n",
    "    with open(path_to_file) as f:\n",
    "        contents = f.read().splitlines()\n",
    "    return contents\n",
    "\n",
    "predicted_speaker_label_rttm_path = f\"{data_dir}/pred_rttms/an4_diarize_test.rttm\"\n",
    "pred_rttm = read_file(predicted_speaker_label_rttm_path)\n",
    "\n",
    "pp.pprint(pred_rttm)\n",
    "\n",
    "from nemo.collections.asr.parts.utils.speaker_utils import rttm_to_labels\n",
    "pred_labels = rttm_to_labels(predicted_speaker_label_rttm_path)\n",
    "\n",
    "color = get_color(signal, pred_labels)\n",
    "display_waveform(signal,'Audio with Speaker Labels', color)\n",
    "display(Audio(signal,rate=16000))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check the speaker-labeled ASR transcription output\n",
    "\n",
    "Now we've done all the processes for running ASR and diarization, let's match the diarization result with the ASR result and get the final output. `get_transcript_with_speaker_labels()` function in `asr_diar_offline` matches diarization output `diar_hyp` with `word_hyp` using the timestamp information from `word_ts_hyp`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trans_info_dict = asr_diar_offline.get_transcript_with_speaker_labels(diar_hyp, word_hyp, word_ts_hyp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After running `get_transcript_with_speaker_labels()` function, the transcription output will be located in `./pred_rttms` folder, which shows **start time to end time of the utterance, speaker ID, and words spoken** during the notified time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "transcription_path_to_file = f\"{data_dir}/pred_rttms/an4_diarize_test.txt\"\n",
    "transcript = read_file(transcription_path_to_file)\n",
    "pp.pprint(transcript)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Another output is transcription output in JSON format, which is saved in `./pred_rttms/an4_diarize_test.json`. \n",
    "\n",
    "In the JSON format output, we include information such as **transcription, estimated number of speakers (variable named `speaker_count`), start and end time of each word and most importantly, speaker label for each word.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "transcription_path_to_file = f\"{data_dir}/pred_rttms/an4_diarize_test.json\"\n",
    "json_contents = read_file(transcription_path_to_file)\n",
    "pp.pprint(json_contents)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluation of Speaker Diarization with ASR (Multi-speaker ASR)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When it comes to ASR, we use word error rate (WER) to measure the accuracy of an ASR model's performance on speech to text (STT) tasks. For speaker diarization, we use diarization error rate (DER), but DER does not include the accuracy of the ASR system. \n",
    "\n",
    "To overcome such limitations, concatenated minimum-permutation word error rate (**cpWER**) is proposed as a new scoring method which can evaluate speaker diarization and speech recognition performance at the same time.\n",
    "\n",
    "cpWER is calculated by going through the following steps.\n",
    "\n",
    "1. Concatenate all utterances of each speaker for both reference and hypothesis files.\n",
    "2. Compute the WER between the reference and all possible speaker permutations of the hypothesis.\n",
    "3. Pick the lowest WER among them (this is assumed to be the best permutation.\n",
    "\n",
    "cpWER was proposed in [this article](https://arxiv.org/pdf/2004.09249.pdf) about CHiME-6 Challenge.\n",
    "\n",
    "For our example file `an4_diarize_test.wav`, we pick the lowest WER value from the two permutations as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hyp1 = \"eleven twenty seven fifty seven october twenty four nineteen seventy\"\n",
    "hyp2 = \"october twenty four nineteen seventy eleven twenty seven fifty seven\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Session-by-session evaluation with a CTM file"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To calculate cpWER value, we need to prepare word sequence lists and CTM files. The word sequence list can be obtained from the `trans_info_dict` variable that is obtained from the function named `get_transcript_with_speaker_labels`.\n",
    "\n",
    "Each word in a word sequence list should be annotated as in the following example.\n",
    "```\n",
    "{'word': 'hello', \n",
    "'start_time': 12.2, \n",
    "'end_time': 12.54, \n",
    "'speaker_label': 'speaker_0'}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from nemo.collections.asr.parts.utils.speaker_utils import get_uniqname_from_filepath\n",
    "\n",
    "word_seq_lists = []\n",
    "uniq_id = \"an4_diarize_test\"\n",
    "\n",
    "# Add the list in `trans_info_dict[uniq_id]['words']`.\n",
    "word_seq_lists.append(trans_info_dict[uniq_id]['words'])\n",
    "\n",
    "# Print the first session in `word_seq_lists`.\n",
    "print(\"word_seq_lists:\\n\", word_seq_lists[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also need CTM files as reference transcripts. The columns of a CTM file indicate the following items:\n",
    "\n",
    "`<session name> <channel ID> <start time> <duration> <word> <confidence>`\n",
    "- Example:\n",
    "`diar_session_123 speaker_3 13.2 0.25 hi 0`\n",
    "\n",
    "For the purpose of creating the reference annotations, we use `<channel ID>` for speaker labels and assign 0 to `<confidence>`. The reference CTM file for the `an4_diarize_test.wav` looks like the following:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "an4_diarize_test_ctm = \\\n",
    "[\"an4_diarize_test speaker_0 0.4 0.51 eleven 0\",\n",
    "\"an4_diarize_test speaker_0 0.95 0.32 twenty 0\",\n",
    "\"an4_diarize_test speaker_0 1.35 0.55 seven 0\",\n",
    "\"an4_diarize_test speaker_0 1.96 0.32 fifty 0\",\n",
    "\"an4_diarize_test speaker_0 2.32 0.75 seven 0\",\n",
    "\"an4_diarize_test speaker_1 3.12 0.42 october 0\",\n",
    "\"an4_diarize_test speaker_1 3.6 0.28 twenty 0\",\n",
    "\"an4_diarize_test speaker_1 3.95 0.35 four 0\",\n",
    "\"an4_diarize_test speaker_1 4.3 0.31 nineteen 0\",\n",
    "\"an4_diarize_test speaker_1 4.65 0.35 seventy 0\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Caveat:** Note that the original order of the words should NOT be permuted in the reference CTM file. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def write_ctm(path, the_list): \n",
    "    outF = open(path, \"w\")\n",
    "    for line in the_list:\n",
    "        outF.write(line)\n",
    "        outF.write(\"\\n\")\n",
    "    outF.close()\n",
    "    \n",
    "write_ctm(f\"{data_dir}/an4_diarize_test.ctm\", an4_diarize_test_ctm)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As opposed to cpWER, we refer to the standard WER value as `WER` since the channels or speakers are mixed in one channel (speaker) when the WER value is calculated. Thus, `WER` is calculated by comparing the sequentially ordered hypothesis words with the reference CTM file.\n",
    "Calculate `cpWER` and `WER` by feeding `word_seq_lists` and `ctm_file_list` to a function named `calculate_cpWER`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nemo.collections.asr.metrics.der import concat_perm_word_error_rate\n",
    "from nemo.collections.asr.metrics.wer import word_error_rate\n",
    "from nemo.collections.asr.parts.utils.diarization_utils import convert_word_dict_seq_to_text, convert_ctm_to_text\n",
    "# Provide a list containing the paths to the reference CTM files \n",
    "# which have the same order with filenames in word_seq_lists.\n",
    "\n",
    "word_seq_list = trans_info_dict['an4_diarize_test']['words']\n",
    "ctm_file_path = f\"{data_dir}/an4_diarize_test.ctm\"\n",
    "\n",
    "spk_hypothesis, mix_hypothesis = convert_word_dict_seq_to_text(word_seq_list)\n",
    "spk_reference, mix_reference = convert_ctm_to_text(ctm_file_path)\n",
    "\n",
    "print(f\"spk_hypothesis: {spk_hypothesis}\")\n",
    "print(f\"mix_hypothesis: {mix_hypothesis}\\n\")\n",
    "print(f\"spk_reference: {spk_reference}\")\n",
    "print(f\"mix_reference: {mix_reference}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see from thee above examples, we need to prepare speaker-separated transcriptions to calculate cpWER.\n",
    "\n",
    "Now that we prepared the necessary inputs for calculating WER and cpWER, lets feed the input data into `concat_perm_word_error_rate` function and `word_error_rate` function. Please note that these two functions accept list as an input where multiple utterance(session) are included in the list keeping the orders in the input manifest file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nemo.collections.asr.metrics.der import concat_perm_word_error_rate \n",
    "from nemo.collections.asr.metrics.wer import word_error_rate\n",
    "\n",
    "cpWER, concat_hyp, concat_ref = concat_perm_word_error_rate([spk_hypothesis], [spk_reference])\n",
    "WER = word_error_rate([mix_hypothesis], [mix_reference])\n",
    "\n",
    "print(f\"cpWER: {cpWER[0]}\")\n",
    "print(f\"WER: {WER}\")\n",
    "\n",
    "# Check the concatenated hypothesis and reference transcript\n",
    "print(f\"concat_hyp: {concat_hyp[0]}\")\n",
    "print(f\"concat_ref: {concat_ref[0]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that `cpWER` and `WER` can be different. For example, if one word is wrongly assigned to the second speaker, cpWER gets one miss error and one insertion error while WER stays the same."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nemo.collections.asr.parts.utils.diarization_utils import convert_word_dict_seq_to_text\n",
    "import copy\n",
    "word_seq_lists_modified = copy.deepcopy(word_seq_lists)\n",
    "# Let's artificially flip a speaker label and check whether cpWER reflects it\n",
    "word_seq_lists_modified[0][-1]['speaker'] = 'speaker_0'\n",
    "print(word_seq_lists_modified[0])\n",
    "\n",
    "spk_hypothesis_modified, mix_hypothesis_modified = convert_word_dict_seq_to_text(word_seq_lists_modified[0])\n",
    "\n",
    "# Check that \"seventy\" in spk_hypothesis has been moved to speaker_0\n",
    "print(f\"spk_hypothesis_modified: {spk_hypothesis_modified}\")\n",
    "print(f\"mix_hypothesis_modified: {mix_hypothesis_modified}\\n\")\n",
    "\n",
    "print(f\"spk_reference: {spk_reference}\")\n",
    "print(f\"mix_reference: {mix_reference}\")\n",
    "\n",
    "# Recalculate cpWER and WER\n",
    "cpWER_modified, concat_hyp, concat_ref = concat_perm_word_error_rate([spk_hypothesis_modified], [spk_reference])\n",
    "WER_modified = word_error_rate([mix_hypothesis_modified], [mix_reference])\n",
    "\n",
    "print(f\"cpWER: {cpWER_modified[0]}\")\n",
    "print(f\"WER: {WER_modified}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the second example where a speaker label is artificially flipped, we can see that `WER` has not been changed while `cpWER` has been changed. This shows that the diarization result is degraded while the ASR result keeps its accuracy."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Provide CTM files in the input manifest file for cpWER evaluation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If CTM files are provided in the input manifest file, we can call `calculate_cpWER` function to directly get `cpWER` value for each session. The following lines show how to provide the path of CTM files into a manifest file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a new manifest file for input with the reference CTM file. \n",
    "meta = {\n",
    "    'audio_filepath': AUDIO_FILENAME, \n",
    "    'offset': 0, \n",
    "    'duration':None, \n",
    "    'label': 'infer', \n",
    "    'text': '-', \n",
    "    'num_speakers': 2, \n",
    "    'rttm_filepath': None,\n",
    "    'ctm_filepath': f\"{data_dir}/an4_diarize_test.ctm\",\n",
    "    'uem_filepath' : None\n",
    "}\n",
    "\n",
    "with open(os.path.join(data_dir,'input_manifest.json'),'w') as fp:\n",
    "    json.dump(meta,fp)\n",
    "    fp.write('\\n')\n",
    "\n",
    "cfg.diarizer.manifest_filepath = os.path.join(data_dir,'input_manifest.json')\n",
    "!cat {cfg.diarizer.manifest_filepath}\n",
    "\n",
    "# We need to call `make_file_lists` again to update manifest file to `asr_diar_offline` instance \n",
    "asr_diar_offline.make_file_lists()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Call `evaluate` function by feeding the `trans_info_dict` variable to calculate "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nemo.collections.asr.parts.utils.diarization_utils import OfflineDiarWithASR\n",
    "trans_info_dict = asr_diar_offline.get_transcript_with_speaker_labels(diar_hyp, word_hyp, word_ts_hyp)\n",
    "session_result_dict = OfflineDiarWithASR.evaluate(hyp_trans_info_dict=trans_info_dict,\n",
    "                                                  audio_file_list=asr_diar_offline.audio_file_list,\n",
    "                                                  ref_ctm_file_list=asr_diar_offline.ctm_file_list)\n",
    "session_result_dict['an4_diarize_test']\n",
    "\n",
    "print(\"cpWER:\", session_result_dict['an4_diarize_test']['cpWER'])\n",
    "print(\"WER:\", session_result_dict['an4_diarize_test']['WER'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Optional Features for ASR with Speaker Diarization\n",
    "\n",
    "#### Beam search decoder\n",
    "Beam-search decoder can be applied to CTC based ASR models. To use this feature, [pyctcdecode](https://github.com/kensho-technologies/pyctcdecode) should be installed. [pyctcdecode](https://github.com/kensho-technologies/pyctcdecode) supports word timestamp generation and can be applied to speaker diarization. pyctcdecode also requires [KenLM](https://github.com/kpu/kenlm) and KenLM is recommended to be installed using PyPI. Install pyctcdecode in your environment with the following commands:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install pyctcdecode\n",
    "!pip install https://github.com/kpu/kenlm/archive/master.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can download publicly available language models (`.arpa` files) at [KALDI Tedlium Language Models](https://kaldi-asr.org/models/m5). Download [4-gram Big ARPA](https://kaldi-asr.org/models/5/4gram_big.arpa.gz) and provide the model path. Let's download the language model file to `data_dir` folder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gzip\n",
    "import shutil\n",
    "def gunzip(file_path,output_path):\n",
    "    with gzip.open(file_path,\"rb\") as f_in, open(output_path,\"wb\") as f_out:\n",
    "        shutil.copyfileobj(f_in, f_out)\n",
    "        f_in.close()\n",
    "        f_out.close()\n",
    "        \n",
    "ARPA_URL = 'https://kaldi-asr.org/models/5/4gram_big.arpa.gz'\n",
    "f = wget.download(ARPA_URL, data_dir)\n",
    "gunzip(f,f.replace(\".gz\",\"\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Provide the downloaded arpa language model file to `cfg.diarizer`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "arpa_model_path = os.path.join(data_dir, '4gram_big.arpa')\n",
    "cfg.diarizer.asr.ctc_decoder_parameters.pretrained_language_model = arpa_model_path"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "create a new `asr_decoder_ts` instance with the updated `cfg.diarizer`. The decoder script will launch pyctcdecode for decoding words and timestamps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import importlib\n",
    "import nemo.collections.asr.parts.utils.decoder_timestamps_utils as decoder_timestamps_utils\n",
    "importlib.reload(decoder_timestamps_utils) # This module should be reloaded after you install pyctcdecode.\n",
    "\n",
    "asr_decoder_ts = ASRDecoderTimeStamps(cfg.diarizer)\n",
    "asr_model = asr_decoder_ts.set_asr_model()\n",
    "word_hyp, word_ts_hyp = asr_decoder_ts.run_ASR(asr_model)\n",
    "\n",
    "print(\"Decoded word output dictionary: \\n\", word_hyp['an4_diarize_test'])\n",
    "print(\"Word-level timestamps dictionary: \\n\", word_ts_hyp['an4_diarize_test'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Realign Words with a Language Model (Experimental)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Diarization result with ASR transcript can be enhanced by applying a language model. The mapping between speaker labels and words can be realigned by employing language models. The realigning process calculates the probability of the words around the boundary between two hypothetical sentences spoken by different speakers.\n",
    "\n",
    " <Example> k-th word: `but`\n",
    "    \n",
    "            hyp_former:\n",
    "                \"since i think like tuesday </s> <s> but he's coming back to albuquerque\"\n",
    "    \n",
    "            hyp_latter:\n",
    "                \"since i think like tuesday but </s> <s> he's coming back to albuquerque\"\n",
    "\n",
    "The joint probabilities of words in the sentence are computed for these two hypotheses. In this example, `hyp_former` is likely to get a higher score and thus word `but` will be assigned to the second speaker.\n",
    "\n",
    "To use this feature, python package [arpa](https://pypi.org/project/arpa/) should be installed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install arpa"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`diarizer.asr.realigning_lm_parameters.logprob_diff_threshold` can be modified to optimize the diarization performance (default value is 1.2). This is a threshold value for the gap between two log-probabilities of two hypotheses. Thus, the lower the threshold, the more changes are expected to be seen in the output transcript.   \n",
    "\n",
    "`arpa` package also uses KenLM language models as in pyctcdecode. You can download publicly available [4-gram Big ARPA](https://kaldi-asr.org/models/5/4gram_big.arpa.gz) model and provide the model path to hydra configuration as follows.\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "arpa_model_path = os.path.join(data_dir, '4gram_big.arpa')\n",
    "cfg.diarizer.asr.realigning_lm_parameters.arpa_language_model = arpa_model_path\n",
    "cfg.diarizer.asr.realigning_lm_parameters.logprob_diff_threshold = 1.2\n",
    "\n",
    "import importlib\n",
    "import nemo.collections.asr.parts.utils.diarization_utils as diarization_utils\n",
    "importlib.reload(diarization_utils) # This module should be reloaded after you install arpa.\n",
    "\n",
    "# Create a new instance with realigning language model\n",
    "asr_diar_offline = OfflineDiarWithASR(cfg.diarizer)\n",
    "asr_diar_offline.word_ts_anchor_offset = asr_decoder_ts.word_ts_anchor_offset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that the language model for realigning is set up, you can run `get_transcript_with_speaker_labels()` to get the results with realigning. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "asr_diar_offline.get_transcript_with_speaker_labels(diar_hyp, word_hyp, word_ts_hyp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "transcription_path_to_file = f\"{data_dir}/pred_rttms/an4_diarize_test.txt\"\n",
    "transcript = read_file(transcription_path_to_file)\n",
    "pp.pprint(transcript)"
   ]
  }
 ],
 "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}