wehub/nvidia-nemo--speech
1
0
Fork 0
Code Issues Pull Requests Actions 1 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
Create PR to main with cherry-pick from release / cherry-pick (push) Failing after 0s
Details
CICD NeMo / pre-flight (push) Failing after 0s
Details
CICD NeMo / configure (push) Has been skipped
Details
Build, validate, and release Neural Modules / pre-flight (push) Failing after 1s
Details
CICD NeMo / code-linting (push) Has been skipped
Details
Build, validate, and release Neural Modules / release (push) Has been skipped
Details
Build, validate, and release Neural Modules / release-summary (push) Has been cancelled
Details
CICD NeMo / cicd-test-container-build (push) Has been cancelled
Details
CICD NeMo / cicd-import-tests (push) Has been cancelled
Details
CICD NeMo / L0_Setup_Test_Data_And_Models (push) Has been cancelled
Details
CICD NeMo / cicd-main-unit-tests (push) Has been cancelled
Details
CICD NeMo / cicd-main-speech (push) Has been cancelled
Details
CICD NeMo / Nemo_CICD_Test (push) Has been cancelled
Details
CICD NeMo / Coverage (e2e) (push) Has been cancelled
Details
CICD NeMo / Coverage (unit-test) (push) Has been cancelled
Details
CodeQL / Analyze (python) (push) Has been cancelled
Details
CICD NeMo / cicd-wait-in-queue (push) Has been cancelled
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:28:58 +08:00
commit ba4be087d5
2316 changed files with 2668701 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tutorials/speaker_tasks/ASR_with_SpeakerDiarization.ipynb
+781
View File
@@ -0,0 +1,781 @@
{
"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
}
tutorials/speaker_tasks/End_to_End_Diarization_Inference.ipynb
+425
View File
@@ -0,0 +1,425 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"\"\"\"\n",
"You can run either this notebook locally (if you have all the dependencies and a GPU) or on Google Colab.\n",
"\n",
"Instructions for setting up Colab are as follows:\n",
"1. Open a new Python 3 notebook.\n",
"2. Import this notebook from GitHub (File -> Upload Notebook -> \"GITHUB\" tab -> copy/paste GitHub URL)\n",
"3. Connect to an instance with a GPU (Runtime -> Change runtime type -> select \"GPU\" for hardware accelerator)\n",
"4. Run this cell to set up dependencies.\n",
"\"\"\"\n",
"# If you're using Google Colab and not running locally, run this cell.\n",
"\n",
"## Install dependencies\n",
"!pip install wget\n",
"!apt-get install sox libsndfile1 ffmpeg\n",
"!pip install text-unidecode\n",
"!pip install ipython\n",
"\n",
"# ## Install NeMo\n",
"BRANCH = 'main'\n",
"!python -m pip install \"nemo_toolkit[asr] @ git+https://github.com/NVIDIA-NeMo/Speech.git@{BRANCH}\""
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# End-to-End Speaker Diarization in NeMo"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"<img src=\"images/cascaded_diar_diagram.png\" alt=\"Cascaded Diar System \" style=\"width: 800px;\"/>"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Traditional cascaded (also referred to as modular or pipelined) speaker diarization systems consist of multiple modules such as a speaker activity detection (SAD) module and a speaker embedding extractor module. Cascaded systems are often time-consuming to develop since each module should be separately trained and optimized."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"<img src=\"images/e2e_diar_diagram.png\" alt=\"End-to-end diarization model diagram\" style=\"width: 800px;\"/>"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"On the other hand, end-to-end speaker diarization systems pursue a much more simplified version of a system where a single neural network model accepts raw audio signals and outputs speaker activity for each audio frame. Therefore, end-to-end diarization models have an advantage in ease of optimization and depoloyments."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Sortformer Diarization Inference\n",
"\n",
"As explained in the [Sortformer Diarization Training](https://github.com/NVIDIA/NeMo/blob/main/tutorials/speaker_tasks/Speaker_Diarization_Training.ipynb) tutorial, Sortformer is trained with Sort-Loss to generate speaker segments in arrival-time order. If a diarization model can generate speaker segments in a pre-defined rule or order, we do not need to match the permutations when we train diarization model with multi-speaker automatic speech recognition (ASR) models or we do not need to match permutations from each window when a diarization model is running in streaming mode where audio chunk sequences are processed, creating a problem of permutation matchin between inference windows. "
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"<img src=\"images/intro_comparison.png\" alt=\"Intro Comparison\" style=\"width: 800px;\"/>"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### A toy example for speaker diarization with a Sortformer model "
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Download a toy example audio file (`an4_diarize_test.wav`) and its ground-truth label file (`an4_diarize_test.rttm`)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"import wget\n",
"ROOT = os.getcwd()\n",
"data_dir = os.path.join(ROOT,'data')\n",
"os.makedirs(data_dir, exist_ok=True)\n",
"an4_audio = os.path.join(data_dir,'an4_diarize_test.wav')\n",
"an4_rttm = os.path.join(data_dir,'an4_diarize_test.rttm')\n",
"if not os.path.exists(an4_audio):\n",
" an4_audio_url = \"https://nemo-public.s3.us-east-2.amazonaws.com/an4_diarize_test.wav\"\n",
" an4_audio = wget.download(an4_audio_url, data_dir)\n",
"if not os.path.exists(an4_rttm):\n",
" an4_rttm_url = \"https://nemo-public.s3.us-east-2.amazonaws.com/an4_diarize_test.rttm\"\n",
" an4_rttm = wget.download(an4_rttm_url, data_dir)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Let's plot and listen to the audio. Pay attention to the two speakers in the audio clip."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import IPython\n",
"import matplotlib.pyplot as plt\n",
"import numpy as np\n",
"import librosa\n",
"\n",
"sr = 16000\n",
"signal, sr = librosa.load(an4_audio,sr=sr)\n",
"\n",
"fig,ax = plt.subplots(1,1)\n",
"fig.set_figwidth(20)\n",
"fig.set_figheight(2)\n",
"plt.plot(np.arange(len(signal)),signal,'gray')\n",
"fig.suptitle('Reference merged an4 audio', fontsize=16)\n",
"plt.xlabel('time (secs)', fontsize=18)\n",
"ax.margins(x=0)\n",
"plt.ylabel('signal strength', fontsize=16)\n",
"a,_ = plt.xticks();plt.xticks(a,a/sr)\n",
"\n",
"IPython.display.Audio(an4_audio)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now that we have downloaded the example audio file contains multiple speakers, we can feed the audio clip into Sortformer diarizer and get the speaker diarization results."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Download Sortformer diarizer model\n",
"\n",
"To download Sortformer diarizer from [Sortformer HuggingFace model card](https://huggingface.co/nvidia/diar_sortformer_4spk-v1), you need to get a [HuggingFace Acces Token](https://huggingface.co/docs/hub/en/security-tokens) and feed your access token in your python environment using [HuggingFace CLI](https://huggingface.co/docs/huggingface_hub/main/en/guides/cli).\n",
"\n",
"If you are having trouble getting a HuggingFace token, you can download Sortformer model from [Sortformer HuggingFace model card](https://huggingface.co/nvidia/diar_sortformer_4spk-v1) and specify the path to the downloaded model.\n",
"\n",
"### Timestamp output and tensor output\n",
"\n",
"When excuting `diarize()` function, if you specify `include_tensor_outputs=True`, the diarization model will return the predicted speaker-labeled segments and tensors containing T by N (N is number of max speakers) sigmoid values for each audio frame. \n",
"\n",
"Without `include_tensor_outputs` variable or `include_tensor_outputs=False`, only speaker labeled segments will be returned."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from nemo.collections.asr.models import SortformerEncLabelModel\n",
"from huggingface_hub import get_token as get_hf_token\n",
"import torch\n",
"\n",
"if get_hf_token() is not None and get_hf_token().startswith(\"hf_\"):\n",
" # If you have logged into HuggingFace hub and have access token\n",
" diar_model = SortformerEncLabelModel.from_pretrained(\"nvidia/diar_sortformer_4spk-v1\")\n",
"else:\n",
" # You can downloaded \".nemo\" file from https://huggingface.co/nvidia/diar_sortformer_4spk-v1 and specify the path.\n",
" diar_model = SortformerEncLabelModel.restore_from(restore_path=\"/path/to/diar_sortformer_4spk-v1.nemo\", map_location=torch.device('cuda'), strict=False)\n",
"\n",
"pred_list, pred_tensor_list = diar_model.diarize(audio=an4_audio, include_tensor_outputs=True)\n",
"\n",
"print(pred_list[0])\n",
"print(pred_tensor_list[0].shape)\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"\n",
"Now let's visualize the predicted speaker diarization results. The diarization model outputs a tensor where each row represents a speaker and each column represents a frame. The sigmoid values in the tensor represent the probability of the frame being spoken by that speaker.\n",
"\n",
"In the following code, we'll plot the predicted speaker diarization results for the sample audio file."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"import matplotlib.pyplot as plt\n",
"import numpy as np\n",
"\n",
"def plot_diarout(preds):\n",
" preds_mat = preds.cpu().numpy().transpose()\n",
" cmap_str, grid_color_p= 'viridis', 'gray'\n",
" LW, FS = 0.4, 36\n",
"\n",
" yticklabels = [\"spk0\", \"spk1\", \"spk2\", \"spk3\"]\n",
" yticks = np.arange(len(yticklabels))\n",
" fig, axs = plt.subplots(1, 1, figsize=(30, 3))\n",
"\n",
" axs.imshow(preds_mat, cmap=cmap_str, interpolation='nearest')\n",
" axs.set_title('Predictions', fontsize=FS)\n",
" axs.set_xticks(np.arange(-.5, preds_mat.shape[1], 1), minor=True)\n",
" axs.set_yticks(yticks)\n",
" axs.set_yticklabels(yticklabels)\n",
" axs.set_xlabel(f\"80 ms Frames\", fontsize=FS)\n",
" axs.grid(which='minor', color=grid_color_p, linestyle='-', linewidth=LW)\n",
"\n",
" plt.savefig('plot.png', dpi=300)\n",
" plt.show()\n",
"\n",
"\n",
"plot_diarout(pred_tensor_list[0])"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Note that the first speaker always gets the first generic speaker label `spk0`. Sortformer model is trained to generate speaker labels in arrival time order, thus permutations of speakers are always predictable if we know the arrival time order of speakers."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Post-Processing of Speaker Timestamps\n",
"\n",
"In the previous steps, we have obtained predictions of speaker activity in a form of Tensors. While the floating point probability values can be interpreted as speaker probabilities, these values are not designed to consumed as is and still requires to be processed to be segment information which has start and end time for each time stamp.\n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from omegaconf import OmegaConf\n",
"import os\n",
"import wget\n",
"import pandas as pd\n",
"from nemo.collections.asr.parts.utils.speaker_utils import rttm_to_labels\n",
"from nemo.collections.asr.metrics.der import score_labels_from_rttm_labels\n",
"\n",
"ROOT = os.getcwd()\n",
"data_dir = os.path.join(ROOT,'data')\n",
"\n",
"yaml_name=\"sortformer_diar_4spk-v1_dihard3-dev.yaml\"\n",
"MODEL_CONFIG = os.path.join(data_dir, yaml_name)\n",
"if True:\n",
" config_url = f\"https://raw.githubusercontent.com/NVIDIA/NeMo/main/examples/speaker_tasks/diarization/conf/post_processing/{yaml_name}\"\n",
" MODEL_CONFIG = wget.download(config_url, data_dir)\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The post-processing yaml file `sortformer_diar_4spk-v1_dihard3-dev.yaml` is containing parameters that are optimized to have the lowest Diarization Error Rate (DER) on the [DiHARD 3 development set](https://catalog.ldc.upenn.edu/LDC2022S12)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from nemo.collections.asr.parts.utils.vad_utils import load_postprocessing_from_yaml\n",
"import json\n",
"from omegaconf import OmegaConf\n",
"post_processing_params = load_postprocessing_from_yaml(MODEL_CONFIG)\n",
"print(json.dumps(OmegaConf.to_container(post_processing_params), indent=4))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Post-processing of Speaker Timestamps\n",
"\n",
"The parameters in post-processing yaml configurations perform the following tasks:\n",
"\n",
"| **Parameter** | **Description** |\n",
"|-------------------:|--------------------------------------------------------------|\n",
"| **onset** | Onset threshold for detecting the beginning of a speech segment. |\n",
"| **offset** | Offset threshold for detecting the end of a speech segment. |\n",
"| **pad_onset** | Adds the specified duration at the beginning of each speech segment. |\n",
"| **pad_offset** | Adds the specified duration at the end of each speech segment. |\n",
"| **min_duration_on**| Removes short speech segments if the duration is less than the specified minimum duration. |\n",
"| **min_duration_off**| Removes short silences if the duration is less than the specified minimum duration. |\n",
"\n",
"Now let's check the diarization output timestamps and compare how post-processing changes the diarization output."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"\n",
"def show_diar_df(pred_session_list):\n",
" data = [segment.split() for segment in pred_session_list]\n",
" df = pd.DataFrame(data, columns=['Start Time', 'End Time', 'Speaker'])\n",
" print(df)\n",
"\n",
"# Call `diarize` method without postprocessing params\n",
"pred_list_bn = diar_model.diarize(audio=an4_audio)\n",
"print(f\" [Default Binarized Diarization Output]: \")\n",
"show_diar_df(pred_list_bn[0])\n",
"\n",
"# Call `diarize` method with postprocessing params\n",
"pred_list_pp = diar_model.diarize(audio=an4_audio, postprocessing_yaml=MODEL_CONFIG)\n",
"print(f\" [Post-processed Diarization Output]: \")\n",
"show_diar_df(pred_list_pp[0])\n",
"\n",
"print(f\" [Ground-truth Diarization Output]: \")\n",
"ref_labels = rttm_to_labels(an4_rttm)\n",
"show_diar_df(ref_labels)\n",
"\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"You can see that the post-processed segments have more on-set padding while the differences are not significant. \n",
"\n",
"Now let's calculate DER (Diarization Error Rate) between the predicted labels and the ground-truth labels for both raw binarized and post-processed diarization outputs."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Get the refernce labels from ground-truth RTTM file\n",
"ref_labels = rttm_to_labels(an4_rttm)\n",
"\n",
"der_metric1, _, _ = score_labels_from_rttm_labels(\n",
" ref_labels_list=[(\"binarize\", ref_labels)],\n",
" hyp_labels_list=[(\"binarize\", pred_list_bn[0])],\n",
" collar=0.0,\n",
" ignore_overlap=False,\n",
" verbose=False,\n",
")\n",
"print(f\"Raw Binarization DER: {abs(der_metric1):.6f}\")\n",
"\n",
"der_metric2, _, _ = score_labels_from_rttm_labels(\n",
" ref_labels_list=[(\"post_processing\", ref_labels)],\n",
" hyp_labels_list=[(\"post_processing\", pred_list_pp[0])],\n",
" collar=0.0,\n",
" ignore_overlap=False,\n",
" verbose=False,\n",
")\n",
"print(f\"Post-Processing DER: {abs(der_metric2):.6f}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Since the diarization output with post-processing is optimized to get the lowest DER for the given sigmoid tensor outputs, it generaly gives lower DER values when compared to the raw binarized diarization output. To achieve the lowest DER, it is recommended to use the post-processing parameters that are optimized for your dataset of interest and your diarization model. "
]
}
],
"metadata": {
"kernelspec": {
"display_name": "nemo012826",
"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.19"
},
"pycharm": {
"stem_cell": {
"cell_type": "raw",
"metadata": {
"collapsed": false
},
"source": []
}
}
},
"nbformat": 4,
"nbformat_minor": 4
}
tutorials/speaker_tasks/End_to_End_Diarization_Training.ipynb
+950
View File
@@ -0,0 +1,950 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"\"\"\"\n",
"You can run either this notebook locally (if you have all the dependencies and a GPU) or on Google Colab.\n",
"\n",
"Instructions for setting up Colab are as follows:\n",
"1. Open a new Python 3 notebook.\n",
"2. Import this notebook from GitHub (File -> Upload Notebook -> \"GITHUB\" tab -> copy/paste GitHub URL)\n",
"3. Connect to an instance with a GPU (Runtime -> Change runtime type -> select \"GPU\" for hardware accelerator)\n",
"4. Run this cell to set up dependencies.\n",
"5. Restart the runtime (Runtime -> Restart Runtime) for any upgraded packages to take effect\n",
"\"\"\"\n",
"# If you're using Google Colab and not running locally, run this cell.\n",
"\n",
"NEMO_DIR_PATH = \"NeMo\"\n",
"BRANCH = 'main'\n",
"\n",
"! git clone https://github.com/NVIDIA-NeMo/Speech\n",
"%cd NeMo\n",
"! python -m pip install \"nemo_toolkit[asr] @ git+https://github.com/NVIDIA-NeMo/Speech.git@{BRANCH}\"\n",
"%cd .."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# End-to-end Speaker Diarization with Sortformer"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Sortformer: Bridging the Gap between tokens (ASR) and Timestamps (Diarization)\n",
"\n",
"### Speaker Diarization as a part of ASR system\n",
"\n",
"Speaker diarization is all about figuring out who’s speaking when in an audio recording. In the world of automatic speech recognition (ASR), this becomes even more important for handling conversations with multiple speakers. Multispeaker ASR (also called speaker-attributed or multitalker ASR) uses this process to not just transcribe what’s being said, but also to label each part of the transcript with the right speaker.\n",
"\n",
"As ASR technology continues to advance, speaker diarization is increasingly becoming part of the ASR workflow itself. Some systems now handle speaker labeling and transcription at the same time during decoding. This means you don’t just get accurate text—you're also getting insights into who said what, making it more useful for conversational analysis.\n",
"\n",
"### Challenges in Integrating Speaker Diarization and ASR\n",
"\n",
"However, despite significant advancements, integrating speaker diarization and ASR into a unified, seamless system remains a considerable challenge. A key obstacle lies in the need for extensive high-quality, annotated audio data featuring multiple speakers. Acquiring such data is far more complex than collecting single-speaker audio or image datasets. This challenge is particularly pronounced for low-resource languages and domains like healthcare, where strict privacy regulations further constrain data availability.\n",
"\n",
"On top of that, many real-world use cases need these models to handle really long audio files—sometimes hours of conversation at a time. Training on such lengthy data is even more complicated because it’s hard to find or annotate. This creates a big gap between what’s needed and what’s available, making multispeaker ASR one of the toughest nuts to crack in the field of speech technology."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"<img src=\"images/intro_comparison.png\" alt=\"Intro Comparison\" style=\"width: 800px;\"/>"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"\n",
"### Sortformer: Simplifying Multispeaker ASR with Arrival Time Sorting\n",
"\n",
"To tackle the complexities of multispeaker automatic speech recognition (ASR), we introduce [*Sortformer*](https://arxiv.org/abs/2409.06656), a new approach that incorporates Sort-Loss and techniques to align timestamps with text tokens. Traditional approaches like permutation-invariant loss (PIL) face challenges when applied in batchable and differentiable computational graphs, especially since token-based objectives struggle to incorporate speaker-specific attributes into PIL-based loss functions.\n",
"\n",
"To address this, we propose an arrival time sorting (ATS) approach. In this method, speaker tokens from ASR outputs and speaker timestamps from diarization outputs are sorted by their arrival times to resolve permutations. This approach allows the multispeaker ASR system to be trained or fine-tuned using token-based cross-entropy loss, eliminating the need for timestamp-based or frame-level objectives with PIL."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"<img src=\"images/ats.png\" alt=\"Arrival Time Sort\" style=\"width: 600px;\"/>"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"\n",
"The ATS-based multispeaker ASR system is powered by an end-to-end neural diarizer model, Sortformer, which generates speaker-label timestamps in arrival time order (ATO). To train the neural diarizer to produce sorted outputs, we introduce Sort Loss, a method that creates gradients enabling the Transformer model to learn the ATS mechanism.\n",
"\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"<img src=\"images/main_dataflow.png\" alt=\"Main Dataflow\" style=\"width: 500px;\"/>"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Additionally, as shown in the above figure, our diarization system integrates directly with the ASR encoder. By embedding speaker supervision data as speaker kernels into the ASR encoder states, the system seamlessly combines speaker and transcription information. This unified approach improves performance and simplifies the overall architecture."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"As a result, our end-to-end multispeaker ASR system is fully or partially trainable with token objectives, allowing both the ASR and speaker diarization modules to be trained or fine-tuned using these objectives. Additionally, during the multispeaker ASR training phase, no specialized loss calculation functions are needed when using Sortformer, as frameworks for standard single-speaker ASR models can be employed. These compatibilities greatly simplify and accelerate the training and fine-tuning process of multispeaker ASR systems. \n",
"\n",
"On top of all these benefits, *Sortformer* can be used as a stand-alone end-to-end speaker diarization model. By training a Sortformer diarizer model especially on high-quality simulated data with accurate time-stamps, you can boost the performance of multi-speaker ASR systems, just by integrating the *Sortformer* model as _*Speaker Supervision*_ model in a computation graph.\n",
"\n",
"In this tutorial, we will walk you through the process of training a Sortformer diarizer model with toy dataset. Before starting, we will introduce the concepts of Sort-Loss calculation and the Hybrid loss technique."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Sort-Loss for *Sortformer* Diarizer"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"<img src=\"images/sortformer.png\" alt=\"Sortformer Model with Hybrid Loss\" style=\"width: 500px;\"/>"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Sort-Loss is designed to compare the predicted outputs with the true labels, typically sorted in arrival-time order or another relevant metric. The key distinction that *Sortformer* introduces compared to previous end-to-end diarization systems such as [EEND-SA](https://arxiv.org/pdf/1909.06247), [EEND-EDA](https://arxiv.org/abs/2106.10654) lies in the organization of class presence $\\mathbf{\\hat{Y}}$.\n",
"\n",
"The figure below illustrates the difference between Sort-Loss and permutation-invariant loss (PIL) or permutation-free loss.\n",
"\n",
" - PIL is calculated by finding the permutation of the target that minimizes the loss value between the prediction and the target.\n",
"\n",
" - Sort-Loss simply compares the arrival-time-sorted version of speaker activity outputs for both the prediction and the target. Note that sometimes the same ground-truth labels lead to different target matrices for Sort-Loss and PIL.\n",
"\n",
"For example, the figure below shows two identical source target matrices (the two matrices at the top), but the resulting target matrices for Sort-Loss and PIL are different."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"<img src=\"images/loss_types.png\" alt=\"PIL model VS SortLoss model\" style=\"width: 1000px;\"/>"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"In mathematical terms, Sort-Loss can be expressed as follows:\n",
"\n",
"* **Arrival Time Sorting Function with $\\Psi$ function** \n",
"\n",
" Let $\\Psi$ be a function that determines the first segment's arrival time for a speaker bin:\n",
"$$\n",
" \\Psi\\big(\\mathbf{y}_{k}\\big) = \\min \\{ t' \\mid y_{k, t'} \\neq 0, t' \\in [1, T] \\} = t^{0}_{k},\n",
"$$\n",
" where $t^{0}_{k}$ is the frame index of the first active segment for the $k$-th speaker.\n",
"\n",
"Sortformer aims to generate predictions $\\mathbf{\\hat{y}}_{k}$ for each speaker $k$ such that:\n",
"$$\n",
"\\Psi(\\mathbf{\\hat{y}}_{1}) \\leq \\Psi(\\mathbf{\\hat{y}}_{2}) \\leq \\cdots \\leq \\Psi(\\mathbf{\\hat{y}}_{K}),\n",
"$$\n",
"ensuring the model produces class presence outputs ($\\mathbf{\\hat{Y}}$) sorted by arrival time.\n",
"\n",
"* **Sorting Function $\\eta$ and Sorted Targets $\\mathbf{Y}_{\\eta}$** \n",
"\n",
"\n",
"Let $\\eta$ be the sorting function applied to speaker indices $\\{1, \\dots, K\\}$. The sorted ground-truth matrix $\\mathbf{Y}_{\\eta}$ is defined as:\n",
"$$\n",
"\\eta\\big( \\mathbf{Y} \\big) = \\mathbf{Y}_{\\eta} = \\left(\\mathbf{y}_{\\eta(1)}, \\dots, \\mathbf{y}_{\\eta(K)} \\right).\n",
"$$\n",
"Using $\\Psi$, the following condition holds for the sorted ground-truth labels $\\mathbf{y}_{\\eta(k)}$:\n",
"$$\n",
"\\Psi(\\mathbf{y}_{\\eta(1)}) \\leq \\Psi(\\mathbf{y}_{\\eta(2)}) \\leq \\cdots \\leq \\Psi(\\mathbf{y}_{\\eta(K)}).\n",
"$$\n",
"\n",
"* **Sort Loss ($\\mathcal{L}_{\\text{Sort}}$) Definition** \n",
"\n",
"\n",
"Sort Loss is computed as:\n",
"$$\n",
"\\mathcal{L}_{\\text{Sort}}\\left(\\mathbf{Y}, \\mathbf{P}\\right) = \\mathcal{L}_{\\text{BCE}} \\left(\\mathbf{Y}_{\\eta}, \\mathbf{P}\\right) = \\frac{1}{K} \\sum_{k=1}^{K} \\mathcal{L}_{\\text{BCE}}(\\mathbf{y}_{\\eta(k)}, \\mathbf{q}_k),\n",
"$$\n",
"where:\n",
"\n",
"- $\\mathbf{y}_{\\eta(k)}$: True labels sorted by arrival time using the sorting function $\\eta$.\n",
"- $\\mathbf{q}_k$: Predicted outputs for the $k$-th speaker.\n",
"- $\\mathcal{L}_{\\text{BCE}}(\\mathbf{y}_{\\eta(k)}, \\mathbf{q}_k)$: Binary cross-entropy (BCE) loss for the $k$-th speaker.\n",
"- $K$: Total number of speakers.\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now that we learn the concept of Sort Loss and Sortformer, we can now calculate Sort Loss based target matrix and PIL-based target matrix to compare the difference in target-value setting matrix and loss calculation.\n",
"\n",
"- raw target matrix $\\mathbf{Y}$: `raw_targets`\n",
"- prediction matrix $\\mathbf{P}$: `preds`\n",
"- ATS target matrix $\\mathbf{Y}_{\\eta}$: `ats_targets`\n",
"- PIL target matrix $\\mathbf{Y}_{\\text{PIL}}$: `pil_targets`\n",
"\n",
"First, assign the values in the above examples to the respective variables to create tensors."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import torch\n",
"\n",
"# Define the binary grid as a Python list\n",
"raw_targets_list = [[\n",
" [0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0],\n",
" [0, 0, 0, 1, 0, 1, 1, 1, 1, 1, 1],\n",
" [1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0],\n",
" [0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1]\n",
"],]\n",
"\n",
"preds_list = [[\n",
" [1, 1, 0, 1, 1, 1, 0, 0, 0, 0, 0],\n",
" [0, 1, 0, 1, 0, 0, 0, 0, 1, 0, 0],\n",
" [0, 0, 0, 0, 1, 1, 0, 1, 1, 1, 1],\n",
" [0, 0, 0, 0, 0, 1, 1, 1, 0, 1, 1],\n",
"],]\n",
"\n",
"# Convert the list to a PyTorch tensor\n",
"raw_targets = torch.tensor(raw_targets_list).transpose(1,2)\n",
"preds = torch.tensor(preds_list).transpose(1,2)\n",
"\n",
"print(raw_targets.shape)\n",
"print(preds.shape)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now, import `get_ats_targets` and `get_pil_targets` functions from the `nemo.collections.asr.parts.utils.asr_multispeaker_utils` module to calculate the ATS and PIL targets. "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import itertools\n",
"import torch\n",
"import nemo\n",
"from nemo.collections.asr.parts.utils.asr_multispeaker_utils import get_ats_targets, get_pil_targets\n",
"\n",
"max_num_of_spks = 4 # Number of speakers\n",
"speaker_inds = list(range(max_num_of_spks))\n",
"speaker_permutations = torch.tensor(list(itertools.permutations(speaker_inds))) # Get all permutations\n",
"\n",
"\n",
"ats_target = get_ats_targets(labels=raw_targets.clone(), preds=preds, speaker_permutations=speaker_permutations)\n",
"pil_target = get_pil_targets(labels=raw_targets.clone(), preds=preds, speaker_permutations=speaker_permutations)\n",
"\n",
"print(f\"Predicted tensor:\")\n",
"print(preds[0].T)\n",
"\n",
"print(f\"\\nATS target:\")\n",
"print(ats_target[0].T)\n",
"\n",
"print(f\"\\nPIL target:\")\n",
"print(pil_target[0].T)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"You can see that ATS target and PIL target are different. Now, you will display the ATS and PIL target matrices to visually compare the difference and also calculate loss values using the BCE loss."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"import torch\n",
"import matplotlib.pyplot as plt\n",
"import numpy as np\n",
"from nemo.collections.asr.losses.bce_loss import BCELoss \n",
"\n",
"bce_loss = BCELoss()\n",
"\n",
"def plot_diarout(preds, title_text, cmap_str):\n",
"\n",
" preds_mat = preds.cpu().numpy().transpose()\n",
" grid_color_p = 'gray'\n",
" LW, FS = 0.5, 10\n",
"\n",
" yticklabels = [\"spk0\", \"spk1\", \"spk2\", \"spk3\"]\n",
" yticks = np.arange(len(yticklabels))\n",
" fig, axs = plt.subplots(1, 1, figsize=(20, 2)) # 1 row, 2 columns for preds and targets\n",
"\n",
" axs.imshow(preds_mat, cmap=cmap_str, interpolation='nearest')\n",
" axs.set_title(title_text, fontsize=FS)\n",
" axs.set_xticks(np.arange(-0.5, preds_mat.shape[1], 1), minor=True)\n",
" axs.set_yticks(np.arange(-0.5, preds_mat.shape[0], 1), minor=True)\n",
" axs.set_yticks(yticks)\n",
" axs.set_yticklabels(yticklabels)\n",
" axs.set_xlabel(f\"80 ms Frames\", fontsize=FS)\n",
" \n",
" # Enable grid\n",
" axs.grid(which='minor', color=grid_color_p, linestyle='-', linewidth=LW)\n",
" axs.tick_params(which=\"minor\", size=0) # Hide minor ticks\n",
" axs.tick_params(which=\"major\", size=5) # Show major ticks\n",
"\n",
" plt.savefig('plot.png', dpi=300) # bbox_inches='tight')\n",
" plt.show()\n",
"\n",
"target_length = torch.tensor([ats_target.shape[1]]) \n",
"print(f\"Target length: {target_length}\")\n",
"plot_diarout(preds[0], title_text='Predictions', cmap_str='viridis')\n",
"\n",
"loss_ats = bce_loss(probs=preds.float(), labels=ats_target.float(), target_lens=target_length)\n",
"print(f\"[ {loss_ats:.4f} ] is the loss from Arrival Time Sort Target: \")\n",
"plot_diarout(ats_target[0], title_text='ATS Target', cmap_str='summer')\n",
"\n",
"loss_pil = bce_loss(probs=preds.float(), labels=pil_target.float(), target_lens=target_length)\n",
"print(f\"[ {loss_pil:.4f} ] is the loss from Permutation Invariant Loss Target\")\n",
"plot_diarout(pil_target[0], title_text='PIL Target', cmap_str='inferno')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"While Sortformer can be trained solely using Sort Loss, there is a limitation: the arrival time estimation is not always accurate. This issue becomes more pronounced as the number of speakers increases during the training session.\n",
"\n",
"Note that Sortformer models can be trained using Sort Loss only, PIL only, or a hybrid loss by adjusting the weight between these two loss components. The hybrid loss $\\mathcal{L}_{\\text{hybrid}}$ can be described as follows:\n",
"\n",
"\n",
"$$\n",
"\\mathcal{L}_{\\text{hybrid}} = \\alpha \\cdot \\mathcal{L}_{\\text{Sort}} + \\beta \\cdot \\mathcal{L}_{\\text{PIL}},\n",
"$$\n",
"\n",
"The weight between ATS and PIL can be adjusted with the variable `model.ats_weight`($\\alpha$) and `model.pil_weight`($\\beta$) in the YAML file of the Sortformer diarizer model as follows:\n",
"\n",
"```yaml\n",
"model: \n",
" sample_rate: 16000\n",
" pil_weight: 0.5 # Weight for Permutation Invariant Loss (PIL) used in training the Sortformer diarizer model\n",
" ats_weight: 0.5 # Weight for Arrival Time Sort (ATS) loss in training the Sortformer diarizer model\n",
" max_num_of_spks: 4 \n",
" ...\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Train a Sortformer Diarizer Model"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Example Data Creation\n",
"\n",
"In this tutorial, we will create a simple toy training dataset using the [NeMo Multispeaker Simulator](https://github.com/NVIDIA/NeMo/blob/main/tutorials/tools/Multispeaker_Simulator.ipynb), with Librispeech as the source dataset for demonstration purposes. If you already have datasets with proper speaker annotations (RTTM files), you can replace the simulated dataset with your own.\n",
"\n",
"If you don’t have access to any speaker diarization datasets, the [NeMo Multispeaker Simulator](https://github.com/NVIDIA/NeMo/blob/main/tutorials/tools/Multispeaker_Simulator.ipynb) can be used to generate a sufficient amount of data samples to meet your requirements.\n",
"\n",
"For more details on the data simulator, refer to the documentation in the [NeMo Multispeaker Simulator](https://github.com/NVIDIA/NeMo/blob/main/tutorials/tools/Multispeaker_Simulator.ipynb). This tutorial will not cover the configurations and detailed process of data simulation."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Install dependencies for data simulator\n",
"!apt-get install sox libsndfile1 ffmpeg\n",
"!pip install wget\n",
"!pip install unidecode\n",
"!pip install \"matplotlib>=3.3.2\""
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### Data Simulation Step 1: Download Required Resources\n",
"\n",
"We need to download the LibriSpeech corpus and corresponding word alignments for generating synthetic multi-speaker audio sessions. In addition, we need to download necessary data cleaning scripts from NeMo git."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"NEMO_DIR_PATH = \"NeMo\"\n",
"BRANCH = 'main'\n",
"\n",
"# download scripts if not already there \n",
"if not os.path.exists('NeMo/scripts'):\n",
" print(\"Downloading necessary scripts\")\n",
" !mkdir -p NeMo/scripts/dataset_processing\n",
" !mkdir -p NeMo/scripts/speaker_tasks\n",
" !wget -P NeMo/scripts/dataset_processing/ https://raw.githubusercontent.com/NVIDIA/NeMo/{BRANCH}/scripts/dataset_processing/get_librispeech_data.py\n",
" !wget -P NeMo/scripts/speaker_tasks/ https://raw.githubusercontent.com/NVIDIA/NeMo/{BRANCH}/scripts/speaker_tasks/create_alignment_manifest.py\n",
" !wget -P NeMo/scripts/speaker_tasks/ https://raw.githubusercontent.com/NVIDIA/NeMo/{BRANCH}/scripts/speaker_tasks/pathfiles_to_diarize_manifest.py"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now that we have downloaded all the necessary scripts for data creation and preparation, we can start the data simulation step by downloading the LibriSpeech corpus."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!mkdir -p LibriSpeech\n",
"!python {NEMO_DIR_PATH}/scripts/dataset_processing/get_librispeech_data.py \\\n",
" --data_root LibriSpeech \\\n",
" --data_sets dev_clean"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We can get the forced word alignments data for the LibriSpeech corpus from [this repository.](https://github.com/CorentinJ/librispeech-alignments). Full forced alignments data can be downloaded at [google drive link for alignments data](https://drive.google.com/file/d/1WYfgr31T-PPwMcxuAq09XZfHQO5Mw8fE/view?usp=sharing). We will download only a subset of forced alignment data containing dev-clean part."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!wget -nc https://dldata-public.s3.us-east-2.amazonaws.com/LibriSpeech_Alignments.tar.gz\n",
"!tar -xzf LibriSpeech_Alignments.tar.gz\n",
"!rm -f LibriSpeech_Alignments.tar.gz"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### Data Simulation Step 2: Produce Manifest File with Forced Alignments\n",
"\n",
"We will merge the LibriSpeech manifest files and LibriSpeech forced alignments into one manifest file for ease of use when generating synthetic data. Create alignment files by running the following script.\n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!python {NEMO_DIR_PATH}/scripts/speaker_tasks/create_alignment_manifest.py \\\n",
" --input_manifest_filepath LibriSpeech/dev_clean.json \\\n",
" --base_alignment_path LibriSpeech_Alignments \\\n",
" --output_manifest_filepath ./dev-clean-align.json \\\n",
" --ctm_output_directory ./ctm_out \\\n",
" --libri_dataset_split dev-clean"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### Data Simulation Step 3: Set data simulation parameters"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now that we have downloaded all the sources we need for data creation, we need to download data simulator configurations in `.yaml` format. Download the YAML file and download `data_simulator.py` script from NeMo repository."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from omegaconf import OmegaConf\n",
"import os\n",
"ROOT = os.getcwd()\n",
"conf_dir = os.path.join(ROOT,'conf')\n",
"!mkdir -p {conf_dir}\n",
"CONFIG_PATH = os.path.join(conf_dir, 'data_simulator.yaml')\n",
"if not os.path.exists(CONFIG_PATH):\n",
" !wget -P {conf_dir} https://raw.githubusercontent.com/NVIDIA/NeMo/{BRANCH}/tools/speech_data_simulator/conf/data_simulator.yaml\n",
"\n",
"config = OmegaConf.load(CONFIG_PATH)\n",
"print(OmegaConf.to_yaml(config))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### Data Simulation Step 4: Generate Simulated Audio Session\n",
"\n",
"We will generate a set of example sessions with the following specifications:\n",
"\n",
"- 10 example sessions for train \n",
"- 10 example sessions for validation\n",
"- 2-speakers in each session\n",
"- 90 seconds of recordings\n",
"\n",
"We need to setup different seed for train and validation sets."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from nemo.collections.asr.data.data_simulation import MultiSpeakerSimulator\n",
"\n",
"# Generate train set \n",
"ROOT = os.getcwd()\n",
"data_dir = os.path.join(ROOT,'simulated_train')\n",
"config.data_simulator.random_seed=10\n",
"config.data_simulator.manifest_filepath=\"./dev-clean-align.json\"\n",
"config.data_simulator.outputs.output_dir=data_dir\n",
"config.data_simulator.session_config.num_sessions=10\n",
"config.data_simulator.session_config.num_speakers=2\n",
"config.data_simulator.session_config.session_length=90\n",
"config.data_simulator.background_noise.add_bg=False \n",
"\n",
"lg = MultiSpeakerSimulator(cfg=config)\n",
"lg.generate_sessions()\n",
"\n",
"# Generate validation set \n",
"data_dir = os.path.join(ROOT,'simulated_valid')\n",
"config.data_simulator.random_seed=20\n",
"config.data_simulator.outputs.output_dir=data_dir\n",
"\n",
"lg = MultiSpeakerSimulator(cfg=config)\n",
"lg.generate_sessions()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now that parameter setting is done, generate the samples by launching `generate_sessions()`."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"lg = MultiSpeakerSimulator(cfg=config)\n",
"lg.generate_sessions()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Data preparation step 5: Listen to and Visualize Session\n",
"\n",
"Listen to the audio and visualize the corresponding speaker timestamps (recorded in a RTTM file for each session)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"scrolled": true
},
"outputs": [],
"source": [
"import os\n",
"import IPython\n",
"import matplotlib.pyplot as plt\n",
"import numpy as np\n",
"import librosa\n",
"from nemo.collections.asr.parts.utils.speaker_utils import rttm_to_labels, labels_to_supervisions\n",
"\n",
"ROOT = os.getcwd()\n",
"data_dir = os.path.join(ROOT,'simulated_train')\n",
"audio = os.path.join(data_dir,'multispeaker_session_0.wav')\n",
"rttm = os.path.join(data_dir,'multispeaker_session_0.rttm')\n",
"\n",
"sr = 16000\n",
"signal, sr = librosa.load(audio,sr=sr) \n",
"\n",
"fig,ax = plt.subplots(1,1)\n",
"fig.set_figwidth(20)\n",
"fig.set_figheight(2)\n",
"plt.plot(np.arange(len(signal)),signal,'gray')\n",
"fig.suptitle('Synthetic Audio Session', fontsize=16)\n",
"plt.xlabel('Time (s)', fontsize=18)\n",
"plt.yticks([], [])\n",
"ax.margins(x=0)\n",
"a,_ = plt.xticks()\n",
"plt.xticks(a[:-1],a[:-1]/sr);\n",
"IPython.display.Audio(audio)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We can visually check the ground-truth file of the first sample by running the following commands. We can see that it has plenty of overlap between two speakers. "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# display speaker labels for reference\n",
"labels = rttm_to_labels(rttm)\n",
"reference = labels_to_supervisions(labels)\n",
"reference"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"You can check that corresponding RTTM files are generated as ground-truth labels for training and evaluation."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!cat simulated_train/multispeaker_session_0.rttm | head -10"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Data preparation step 6: Check out the created files\n",
"\n",
"The following files are generated from data simulator:\n",
"\n",
"* _wav files_ (one per audio session) - the output audio sessions\n",
"* _rttm files_ (one per audio session) - the speaker timestamps for the corresponding audio session (used for diarization training)\n",
"* _list files_ (one per file type per batch of sessions) - a list of generated files of the given type (e.g., wav, rttm), used primarily for manifest creation\n",
"\n",
"Check if the files we need are generated by running the following commands."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"print(\"\\n Training audio files:\")\n",
"!ls simulated_train/*.wav\n",
"print(\"\\n Training audio files:\")\n",
"!ls simulated_train/*.rttm\n",
"print(\"\\n Training RTTM list content:\")\n",
"!cat simulated_train/synthetic_wav.list\n",
"print(\"\\n Training RTTM list content:\")\n",
"!cat simulated_train/synthetic_rttm.list\n",
"\n",
"print(\"\\n Validation audio files:\")\n",
"!ls simulated_valid/*.wav\n",
"print(\"\\n Validation audio files:\")\n",
"!ls simulated_valid/*.rttm\n",
"print(\"\\n Validation RTTM list content:\")\n",
"!cat simulated_valid/synthetic_wav.list\n",
"print(\"\\n Validation RTTM list content:\")\n",
"!cat simulated_valid/synthetic_rttm.list"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Prepare Training Data for Sortformer diarizer"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now that we have datasets for both train and validation (dev), we can start preparing and cleaning the data samples for training. Make sure you have the following list of files:\n",
"\n",
"**Training set** \n",
"\n",
"- Train audio files `.wav`\n",
"- A train audio list file `.list`\n",
"- Train RTTM files `.rttm`\n",
"- A train RTTM list content `.list`\n",
"\n",
"**Validation set** \n",
"\n",
"- Validation audio files `.wav`\n",
"- A validation audio list file `.list`\n",
"- Validation RTTM files `.rttm`\n",
"- A validation RTTM list file `.list`\n",
"\n",
"\n",
"Based on these files, we need to create manifest files containing data samples we have. If you don't have a `.list` file, you need to create a `.list` file for the `.wav` files and `.rttm` files."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Create a NeMo manifest (.json) file for training dataset\n",
"!python {NEMO_DIR_PATH}/scripts/speaker_tasks/pathfiles_to_diarize_manifest.py \\\n",
" --add_duration \\\n",
" --paths2audio_files='simulated_train/synthetic_wav.list' \\\n",
" --paths2rttm_files='simulated_train/synthetic_rttm.list' \\\n",
" --manifest_filepath='simulated_train/sortformer_train.json'\n",
"\n",
"# Create a NeMo manifest (.json) file for validation (dev) dataset\n",
"!python {NEMO_DIR_PATH}/scripts/speaker_tasks/pathfiles_to_diarize_manifest.py \\\n",
" --add_duration \\\n",
" --paths2audio_files='simulated_valid/synthetic_wav.list' \\\n",
" --paths2rttm_files='simulated_valid/synthetic_rttm.list' \\\n",
" --manifest_filepath='simulated_valid/sortformer_valid.json'"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"If you print the content of the created manifest file, you can see that `.rttm` files in the list and `.wav` files are grouped together in the generated manifest files."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"print(\"\\nTraining Dataset:\")\n",
"!cat simulated_train/sortformer_train.json | tail -5\n",
"print(\"\\nValidation Dataset:\")\n",
"!cat simulated_valid/sortformer_valid.json | tail -5 "
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Train a Sortformer Diarizer Model\n",
"\n",
"### Training an offline Sortformer model"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now that we have prepared all the necessary dataset, we can train an Sortformer diarizer model on the prepared dataset. Download YAML file for training form NeMo repository and load the configuration into `config` variable."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import nemo\n",
"import os\n",
"import lightning.pytorch as pl\n",
"from omegaconf import OmegaConf\n",
"from nemo.collections.asr.models import SortformerEncLabelModel\n",
"from nemo.utils.exp_manager import exp_manager\n",
"\n",
"NEMO_ROOT = os.getcwd()\n",
"!mkdir -p conf \n",
"!wget -P conf https://raw.githubusercontent.com/NVIDIA/NeMo/{BRANCH}/examples/speaker_tasks/diarization/conf/neural_diarizer/sortformer_diarizer_hybrid_loss_4spk-v1.yaml\n",
"MODEL_CONFIG = os.path.join(NEMO_ROOT,'conf/sortformer_diarizer_hybrid_loss_4spk-v1.yaml')\n",
"config = OmegaConf.load(MODEL_CONFIG)\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Setup the `manifest_filepath` for `train_ds` and `validation_ds` by feeding the `json` file paths based on the created training dataset."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"curr_dir = os.getcwd() + \"/\"\n",
"config.model.train_ds.manifest_filepath = f'{curr_dir}simulated_train/sortformer_train.json'\n",
"config.model.validation_ds.manifest_filepath = f'{curr_dir}simulated_valid/sortformer_valid.json'\n",
"config.trainer.strategy = \"ddp_notebook\"\n",
"config.batch_size = 3\n",
"\n",
"config.trainer.devices=1\n",
"config.accelerator=\"gpu\"\n",
"print(os.getcwd())\n",
"\n",
"print(\"config.model.train_ds.manifest_filepath \", config.model.train_ds.manifest_filepath )"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Setup a model with the given configuration and start a training session."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"\n",
"trainer = pl.Trainer(devices=1, accelerator='gpu', max_epochs=50,\n",
" enable_checkpointing=False, logger=False,\n",
" log_every_n_steps=5, check_val_every_n_epoch=10)\n",
"\n",
"exp_manager(trainer, config.get(\"exp_manager\", None))\n",
"sortformer_model = SortformerEncLabelModel(cfg=config.model, trainer=trainer)\n",
"sortformer_model.maybe_init_from_pretrained_checkpoint(config)\n",
"trainer.fit(sortformer_model)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Training a streaming Sortformer model\n",
"\n",
"If you want to train a streaming version of Sortformer, you can download the following YAML file."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"\n",
"!wget -P conf https://raw.githubusercontent.com/NVIDIA/NeMo/{BRANCH}/examples/speaker_tasks/diarization/conf/neural_diarizer/streaming_sortformer_diarizer_4spk-v2.yaml\n",
"MODEL_CONFIG = os.path.join(NEMO_ROOT,'conf/streaming_sortformer_diarizer_4spk-v2.yaml')\n",
"config = OmegaConf.load(MODEL_CONFIG)\n",
"\n",
"curr_dir = os.getcwd() + \"/\"\n",
"config.model.train_ds.manifest_filepath = f'{curr_dir}simulated_train/sortformer_train.json'\n",
"config.model.test_ds.manifest_filepath = f'{curr_dir}simulated_valid/sortformer_valid.json'\n",
"config.model.validation_ds.manifest_filepath = f'{curr_dir}simulated_valid/sortformer_valid.json'\n",
"config.trainer.strategy = \"ddp_notebook\"\n",
"config.batch_size = 3\n",
"\n",
"config.trainer.devices=1\n",
"config.accelerator=\"gpu\"\n",
"print(os.getcwd())\n",
"\n",
"print(\"config.model.train_ds.manifest_filepath \", config.model.train_ds.manifest_filepath )"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Initiate a streaming Sortformer diarization training session using the given configurations."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"trainer = pl.Trainer(devices=1, accelerator='gpu', max_epochs=50,\n",
" enable_checkpointing=False, logger=False,\n",
" log_every_n_steps=5, check_val_every_n_epoch=10)\n",
"\n",
"exp_manager(trainer, config.get(\"exp_manager\", None))\n",
"streaming_sortformer_model = SortformerEncLabelModel(cfg=config.model, trainer=trainer)\n",
"streaming_sortformer_model.maybe_init_from_pretrained_checkpoint(config)\n",
"trainer.fit(streaming_sortformer_model)"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "nv082124",
"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.12"
},
"pycharm": {
"stem_cell": {
"cell_type": "raw",
"metadata": {
"collapsed": false
},
"source": []
}
}
},
"nbformat": 4,
"nbformat_minor": 4
}
tutorials/speaker_tasks/Speaker_Identification_Verification.ipynb
+1225
View File
File diff suppressed because it is too large Load Diff
tutorials/speaker_tasks/Streaming_End_to_End_Diarization_Inference.ipynb
+409
View File
@@ -0,0 +1,409 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"\"\"\"\n",
"You can run either this notebook locally (if you have all the dependencies and a GPU) or on Google Colab.\n",
"\n",
"Instructions for setting up Colab are as follows:\n",
"1. Open a new Python 3 notebook.\n",
"2. Import this notebook from GitHub (File -> Upload Notebook -> \"GITHUB\" tab -> copy/paste GitHub URL)\n",
"3. Connect to an instance with a GPU (Runtime -> Change runtime type -> select \"GPU\" for hardware accelerator)\n",
"4. Run this cell to set up dependencies.\n",
"\"\"\"\n",
"# If you're using Google Colab and not running locally, run this cell.\n",
"\n",
"## Install dependencies\n",
"!pip install wget\n",
"!apt-get install sox libsndfile1 ffmpeg\n",
"!pip install text-unidecode\n",
"!pip install ipython\n",
"\n",
"# ## Install NeMo\n",
"BRANCH = 'main'\n",
"!python -m pip install \"nemo_toolkit[asr] @ git+https://github.com/NVIDIA-NeMo/Speech.git@{BRANCH}\""
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Streaming End-to-End Speaker Diarization \n",
"\n",
"\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Streaming Diarization Inference with Sortformer\n",
"\n",
"As explained in the [Sortformer Diarization Training](https://github.com/NVIDIA/NeMo/blob/main/tutorials/speaker_tasks/Speaker_Diarization_Training.ipynb) tutorial, Sortformer is trained with Sort-Loss to generate speaker segments in arrival-time order. If a diarization model can generate speaker segments in a pre-defined manner or order, we do not need to match the permutations when we train diarization model with multi-speaker automatic speech recognition (ASR) models, nor do we need to match permutations from each window when a diarization model is running in streaming mode where audio chunk sequences are processed, creating a problem of permutation matching between inference windows. \n",
"\n",
"### Arrival-Order Speaker Cache\n",
"\n",
"We propose the [Arrival-Order Speaker Cache (AOSC)](https://arxiv.org/pdf/2507.18446), which stores frame-level embeddings from the pre-encode NEST module. Unlike [speaker-tracing buffer](https://arxiv.org/pdf/2006.02616) in the previous [EEND-based online systems](https://arxiv.org/pdf/2101.08473), AOSC organizes embeddings by speaker index in arrival-time order. Combined with Sortformer's built-in arrival-ordering mechanism, this automatically resolves between-chunk permutations."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"<img src=\"images/cache_fifo_chunk.png\" alt=\"Cache, FIFO and Chunk\" style=\"width: 800px;\"/>"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Speaker cache updates in streaming Sortformer \n",
"\n",
"Short-chunk processing often reduces accuracy due to limited context. To address this, we combine AOSC with a FIFO queue, enhancing context utilization and enabling less frequent AOSC updates (rather than per-chunk), improving robustness and efficiency. As shown in the above figure, the system includes a speaker cache, FIFO queue, and input buffer (holding the current chunk and future context). Frames pushed out of the queue are processed by the speaker cache update mechanism."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"<img src=\"images/streaming_steps.png\" alt=\"Streaming steps\" style=\"width: 1200px;\"/>"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The AOSC update acts as a no-op function if the input sequence is shorter than the maximum cache length `spkcache_len`. For longer sequences, it compresses the input to `spkcache_len` frames by keeping only the highest-scoring embeddings based on the model's frame-level predictions. You can find more details about the speaker cache update mechanism in the [Streaming Sortformer](https://arxiv.org/pdf/2507.18446) paper."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### A toy example for diarization with a streaming Sortformer model \n",
"\n",
"Download a toy example audio file (`an4_diarize_test.wav`) and its ground-truth label file (`an4_diarize_test.rttm`)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"import wget\n",
"ROOT = os.getcwd()\n",
"data_dir = os.path.join(ROOT,'data')\n",
"os.makedirs(data_dir, exist_ok=True)\n",
"an4_audio = os.path.join(data_dir,'an4_diarize_test.wav')\n",
"an4_rttm = os.path.join(data_dir,'an4_diarize_test.rttm')\n",
"if not os.path.exists(an4_audio):\n",
" an4_audio_url = \"https://nemo-public.s3.us-east-2.amazonaws.com/an4_diarize_test.wav\"\n",
" an4_audio = wget.download(an4_audio_url, data_dir)\n",
"if not os.path.exists(an4_rttm):\n",
" an4_rttm_url = \"https://nemo-public.s3.us-east-2.amazonaws.com/an4_diarize_test.rttm\"\n",
" an4_rttm = wget.download(an4_rttm_url, data_dir)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Let's plot the waveform and listen to the audio. You'll notice that there are two speakers in the audio clip."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import IPython\n",
"import matplotlib.pyplot as plt\n",
"import numpy as np\n",
"import librosa\n",
"\n",
"sr = 16000\n",
"signal, sr = librosa.load(an4_audio,sr=sr)\n",
"\n",
"fig,ax = plt.subplots(1,1)\n",
"fig.set_figwidth(20)\n",
"fig.set_figheight(2)\n",
"plt.plot(np.arange(len(signal)),signal,'gray')\n",
"fig.suptitle('Reference merged an4 audio', fontsize=16)\n",
"plt.xlabel('time (secs)', fontsize=18)\n",
"ax.margins(x=0)\n",
"plt.ylabel('signal strength', fontsize=16)\n",
"a,_ = plt.xticks();plt.xticks(a,a/sr)\n",
"\n",
"IPython.display.Audio(signal, rate=sr)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now that we have downloaded the example audio file contains multiple speakers, we can feed the audio clip into Sortformer diarizer and get the speaker diarization results."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Download Sortformer diarizer model\n",
"\n",
"To download streaming Sortformer diarizer from [HuggingFace model card](https://huggingface.co/nvidia) you need to get a [HuggingFace Acces Token](https://huggingface.co/docs/hub/en/security-tokens) and feed your access token in your python environment using [HuggingFace CLI](https://huggingface.co/docs/huggingface_hub/main/en/guides/cli).\n",
"\n",
"If you are having trouble getting a HuggingFace token, you can download Sortformer model from [Streaming Sortformer HuggingFace model card](https://huggingface.co/nvidia) and specify the `.nemo` file path to the downloaded model."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from nemo.collections.asr.models import SortformerEncLabelModel\n",
"from huggingface_hub import get_token as get_hf_token\n",
"import torch\n",
"\n",
"if get_hf_token() is not None and get_hf_token().startswith(\"hf_\"):\n",
" # If you have logged into HuggingFace hub and have access token\n",
" diar_model = SortformerEncLabelModel.from_pretrained(\"nvidia/diar_streaming_sortformer_4spk-v2\")\n",
"else:\n",
" # You can downloaded \".nemo\" file from https://huggingface.co/nvidia/diar_streaming_sortformer_4spk-v2 and specify the path.\n",
" diar_model = SortformerEncLabelModel.restore_from(restore_path=\"/path/to/diar_streaming_sortformer_4spk-v2.nemo\", map_location=torch.device('cuda'), strict=False)\n",
"diar_model.eval()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Diarization output display function\n",
"\n",
"To visualize the streaming diarization output, we use the same diarization output display function as in offline Sortformer diarizer tutorial. "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import matplotlib.pyplot as plt\n",
"import numpy as np\n",
"\n",
"def plot_diarout(preds):\n",
" preds_mat = preds.cpu().numpy().transpose()\n",
" cmap_str, grid_color_p= 'viridis', 'gray'\n",
" LW, FS = 0.4, 36\n",
"\n",
" yticklabels = [\"spk0\", \"spk1\", \"spk2\", \"spk3\"]\n",
" yticks = np.arange(len(yticklabels))\n",
" fig, axs = plt.subplots(1, 1, figsize=(30, 3))\n",
"\n",
" axs.imshow(preds_mat, cmap=cmap_str, interpolation='nearest')\n",
" axs.set_title('Predictions', fontsize=FS)\n",
" axs.set_xticks(np.arange(-.5, preds_mat.shape[1], 1), minor=True)\n",
" axs.set_yticks(yticks)\n",
" axs.set_yticklabels(yticklabels)\n",
" axs.set_xlabel(f\"80 ms Frames\", fontsize=FS)\n",
" axs.grid(which='minor', color=grid_color_p, linestyle='-', linewidth=LW)\n",
"\n",
" plt.savefig('plot.png', dpi=300)\n",
" plt.show()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Running Streaming Sortformer diarizer\n",
"\n",
"### Parameter configuration for Streaming Sortformer \n",
"\n",
"Now it's time to setup the model with streaming parameters (all measured in 80ms frames). \n",
"\n",
"`chunk_len`: The number of frames in a processing chunk. \n",
"`chunk_right_context`: The right context length. \n",
"`fifo_len`: The number of previous frames attached before the chunk, from the FIFO queue.\n",
"`spkcache_update_period`: The number of frames extracted from the FIFO queue to update the speaker cache.\n",
"`spkcache_len`: The total number of frames in the speaker cache.\n",
"\n",
"Note that the input buffer latency is determined by `chunk_len` + `chunk_right_context`.\n",
"\n",
"The following restrictions apply to the Streaming Sortformer parameters: \n",
" * All streaming parameters must be non-negative integers. \n",
" * `chunk_len` and `spkcache_update_period` must be greater than 0."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import time\n",
"import math\n",
"import torch\n",
"import torch.amp\n",
"from tqdm import tqdm\n",
"\n",
"# If cuda is available, assign the model to cuda\n",
"if torch.cuda.is_available():\n",
" diar_model.to(torch.device(\"cuda\"))\n",
"\n",
"global autocast\n",
"autocast = torch.amp.autocast(diar_model.device.type, enabled=True)\n",
"\n",
"# Set the streaming parameters corresponding to 1.04s latency setup. This will affect the streaming feat loader.\n",
"diar_model.sortformer_modules.chunk_len = 6\n",
"diar_model.sortformer_modules.spkcache_len = 188\n",
"diar_model.sortformer_modules.chunk_right_context = 7\n",
"diar_model.sortformer_modules.fifo_len = 188\n",
"diar_model.sortformer_modules.spkcache_update_period = 144\n",
"diar_model.sortformer_modules.log = False\n",
"\n",
"# Validate that the streaming parameters are set correctly\n",
"diar_model.sortformer_modules._check_streaming_parameters()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Feature extraction from an audio file \n",
"\n",
"\n",
"Now, set up the input audio signal and convert it to log-mel features. Note that we are simulating the streaming scenario. In real life, we won't be able to access the whole utterance at once."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"audio_signal = torch.tensor(signal).unsqueeze(0).to(diar_model.device)\n",
"audio_signal_length = torch.tensor([audio_signal.shape[1]]).to(diar_model.device)\n",
"processed_signal, processed_signal_length = diar_model.preprocessor(input_signal=audio_signal, length=audio_signal_length)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Run a streaming loop for streaming diarization\n",
"\n",
"The following variables are needed to run the simulated streaming speaker diarization session. "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"batch_size = 1\n",
"processed_signal_offset = torch.zeros((batch_size,), dtype=torch.long, device=diar_model.device)\n",
"\n",
"streaming_state = diar_model.sortformer_modules.init_streaming_state(\n",
" batch_size = batch_size,\n",
" async_streaming = True,\n",
" device = diar_model.device\n",
" )\n",
"total_preds = torch.zeros((batch_size, 0, diar_model.sortformer_modules.n_spk), device=diar_model.device)\n",
"\n",
"streaming_loader = diar_model.sortformer_modules.streaming_feat_loader(\n",
" feat_seq=processed_signal,\n",
" feat_seq_length=processed_signal_length,\n",
" feat_seq_offset=processed_signal_offset,\n",
")\n",
"\n",
"num_chunks = math.ceil(\n",
" processed_signal.shape[2] / (diar_model.sortformer_modules.chunk_len * diar_model.sortformer_modules.subsampling_factor)\n",
")\n",
"\n",
"plot_preds = torch.zeros(\n",
" (batch_size, num_chunks * diar_model.sortformer_modules.chunk_len, diar_model.sortformer_modules.n_spk), device=diar_model.device\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now we are ready to run streaming diarization step. Check out the output at the end of the streaming for loop. Note that this is a way to simulate the streaming input. In real-life setting, `chunk_feat_seq_t` needs to be replaced with real-time streaming microphone audio stream."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# To simulate the real-time streaming, we will sleep for a chunk duration after each step\n",
"chunk_duration_seconds = diar_model.sortformer_modules.chunk_len * diar_model.sortformer_modules.subsampling_factor * diar_model.preprocessor._cfg.window_stride\n",
"print(f\"Chunk duration: {chunk_duration_seconds} seconds\")\n",
"\n",
"for i, chunk_feat_seq_t, feat_lengths, left_offset, right_offset in tqdm(\n",
" streaming_loader,\n",
" total=num_chunks,\n",
" desc=\"Streaming Steps\",\n",
" disable=False,\n",
"):\n",
" loop_start_time = time.time()\n",
" with torch.inference_mode():\n",
" with autocast:\n",
" streaming_state, total_preds = diar_model.forward_streaming_step(\n",
" processed_signal=chunk_feat_seq_t,\n",
" processed_signal_length=feat_lengths,\n",
" streaming_state=streaming_state,\n",
" total_preds=total_preds,\n",
" left_offset=left_offset,\n",
" right_offset=right_offset,\n",
" )\n",
" # plot the predictions\n",
" plot_preds[:,:total_preds.shape[1]] = total_preds\n",
" plot_diarout(plot_preds[0,:])\n",
" time.sleep(chunk_duration_seconds)"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"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.12"
},
"pycharm": {
"stem_cell": {
"cell_type": "raw",
"metadata": {
"collapsed": false
},
"source": []
}
}
},
"nbformat": 4,
"nbformat_minor": 4
}
tutorials/speaker_tasks/images/affinity_matrix_fusion.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 280 KiB

tutorials/speaker_tasks/images/ats.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 326 KiB

tutorials/speaker_tasks/images/cache_fifo_chunk.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 120 KiB

tutorials/speaker_tasks/images/cascaded_diar_diagram.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 139 KiB

tutorials/speaker_tasks/images/diar_pipeline.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 64 KiB

tutorials/speaker_tasks/images/e2e_diar_diagram.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 88 KiB

tutorials/speaker_tasks/images/intro_comparison.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 317 KiB

tutorials/speaker_tasks/images/loss_types.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 216 KiB

tutorials/speaker_tasks/images/main_dataflow.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 449 KiB

tutorials/speaker_tasks/images/multiscale_example.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 450 KiB

tutorials/speaker_tasks/images/sortformer.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 439 KiB

tutorials/speaker_tasks/images/streaming_steps.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 622 KiB