nvidia-nemo--speech/tutorials/speaker_tasks/ASR_with_SpeakerDiarization.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Automatic Speech Recognition with Speaker Diarization"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "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 \"nemo_toolkit[asr] @ git+https://github.com/NVIDIA-NeMo/Speech.git@{BRANCH}\""
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "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",
   "metadata": {},
   "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)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "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",
   "metadata": {},
   "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))"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`display_waveform()` and `get_color()` functions are defined for displaying the waveform with diarization results."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "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"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using the above function, we can display the waveform of the example audio clip. "
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "display_waveform(signal)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "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",
   "metadata": {},
   "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))"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "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",
   "metadata": {},
   "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}"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "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",
   "metadata": {},
   "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"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "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",
   "metadata": {},
   "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'])"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "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",
   "metadata": {},
   "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"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "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",
   "metadata": {},
   "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'])"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "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",
   "metadata": {},
   "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))"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "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",
   "metadata": {},
   "source": [
    "trans_info_dict = asr_diar_offline.get_transcript_with_speaker_labels(diar_hyp, word_hyp, word_ts_hyp)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "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",
   "metadata": {},
   "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)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "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",
   "metadata": {},
   "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)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "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",
   "metadata": {},
   "source": [
    "hyp1 = \"eleven twenty seven fifty seven october twenty four nineteen seventy\"\n",
    "hyp2 = \"october twenty four nineteen seventy eleven twenty seven fifty seven\""
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "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",
   "metadata": {
    "scrolled": true
   },
   "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])"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "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> <type of token> <speaker>`\n",
    "- Example:\n",
    "`diar_session_123 1 13.2 0.25 hi 0 lex speaker_3`\n",
    "\n",
    "For the purpose of creating the reference annotations, we set `1` for  `<channel ID>` and assign \"`NA`\"  to  `<confidence>`, \"`lex`\"  to  `<type of token>`. The reference CTM file for the `an4_diarize_test.wav` looks like the following:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "an4_diarize_test_ctm = \\\n",
    "[\"an4_diarize_test 1 0.4 0.51 eleven NA lex speaker_0\",\n",
    "\"an4_diarize_test 1 0.95 0.32 twenty NA lex speaker_0\",\n",
    "\"an4_diarize_test 1 1.35 0.55 seven NA lex speaker_0\",\n",
    "\"an4_diarize_test 1 1.96 0.32 fifty NA lex speaker_0\",\n",
    "\"an4_diarize_test 1 2.32 0.75 seven NA lex speaker_0\",\n",
    "\"an4_diarize_test 1 3.12 0.42 october NA lex speaker_1\",\n",
    "\"an4_diarize_test 1 3.6 0.28 twenty NA lex speaker_1\",\n",
    "\"an4_diarize_test 1 3.95 0.35 four NA lex speaker_1\",\n",
    "\"an4_diarize_test 1 4.3 0.31 nineteen NA lex speaker_1\",\n",
    "\"an4_diarize_test 1 4.65 0.35 seventy NA lex speaker_1\",]"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "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",
   "metadata": {},
   "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)"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "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",
   "metadata": {},
   "source": [
    "from nemo.collections.asr.metrics.cpwer 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}\")"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "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",
   "metadata": {},
   "source": [
    "from nemo.collections.asr.metrics.cpwer 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]}\")"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "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",
   "metadata": {},
   "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}\")"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "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",
   "metadata": {},
   "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()"
   ],
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Call `evaluate` function by feeding the `trans_info_dict` variable to calculate "
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "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'])"
   ],
   "execution_count": null,
   "outputs": []
  }
 ],
 "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}