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

chore: import upstream snapshot with attribution
CICD NeMo / cicd-main-unit-tests (push) Blocked by required conditions
Details
CICD NeMo / cicd-main-speech (push) Blocked by required conditions
Details
CICD NeMo / cicd-test-container-build (push) Blocked by required conditions
Details
CICD NeMo / cicd-import-tests (push) Blocked by required conditions
Details
CICD NeMo / L0_Setup_Test_Data_And_Models (push) Blocked by required conditions
Details
CICD NeMo / Nemo_CICD_Test (push) Blocked by required conditions
Details
CICD NeMo / Coverage (e2e) (push) Blocked by required conditions
Details
CICD NeMo / Coverage (unit-test) (push) Blocked by required conditions
Details
CodeQL / Analyze (python) (push) Waiting to run
Details
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
CICD NeMo / cicd-wait-in-queue (push) Waiting to run
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

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/00_NeMo_Primer.ipynb
+1246
View File
File diff suppressed because it is too large Load Diff
tutorials/01_NeMo_Models.ipynb
+2588
View File
File diff suppressed because it is too large Load Diff
tutorials/02_NeMo_Adapters.ipynb
+1994
View File
File diff suppressed because it is too large Load Diff
tutorials/Publish_NeMo_Model_On_Hugging_Face_Hub.ipynb
+954
View File
@@ -0,0 +1,954 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "qjQ5KQIkaK2J"
},
"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",
"!apt-get install sox libsndfile1 ffmpeg\n",
"!pip install wget\n",
"!pip install text-unidecode\n",
"\n",
"### Install NeMo\n",
"BRANCH = 'main'\n",
"!python -m pip install \"nemo_toolkit[all] @ git+https://github.com/NVIDIA-NeMo/Speech.git@$BRANCH\""
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "J6d04-VRjC-O"
},
"outputs": [],
"source": [
"### Install Hugging Face Hub\n",
"!python -m pip install huggingface_hub\n",
"!python -m pip install evaluate\n",
"!python -m pip install ipywidgets"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# huggingface_hub >= 1.0 defaults to Xet storage, whose dedup breaks the\n",
"# delete/recreate + re-upload flow below (HTTP 400: \"an LFS pointer pointed to\n",
"# a file that does not exist\"). Disable Xet to use reliable classic Git-LFS.\n",
"import os\n",
"\n",
"os.environ[\"HF_HUB_DISABLE_XET\"] = \"1\""
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "aS-Y5O_oGBTc"
},
"source": [
"# NeMo models on Hugging Face Hub\n",
"\n",
"This guide will briefly show how to upload NeMo models to Hugging Face programmatically.\n",
"\n",
"This enables community members to share their NeMo models (any model!) with all users of NeMo!\n",
"\n",
"**Note**: While in this tutorial we showcase an ASR model, there is no particular restriction to any domain - all NeMo models (.nemo files) of every domain can be uploaded and shared in the same way."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Us3UlvwCiEZi"
},
"source": [
"# Login to Hugging Face\n",
"\n",
"Use the notebook login, and access your user access token (or create one to upload models to Hugging Face).\n",
"\n",
"For more information, visit the User Access Token section - https://huggingface.co/docs/hub/security-tokens"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "4RTYbCLziEnb"
},
"outputs": [],
"source": [
"from huggingface_hub import notebook_login\n",
"\n",
"notebook_login()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "dgZbTPcFiaml"
},
"outputs": [],
"source": [
"!git config --global credential.helper store"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "s-FiNn1eiFAl"
},
"source": [
"# Prepare a model to upload to HF\n",
"\n",
"In this example, we will download a NeMo ASR model from NGC and then upload it to Hugging Face for simplicity and to showcase the method.\n",
"\n",
"**You can swap out this ASR model for any model that you restore via `restore_from()` and follow the same steps to upload your own models !**"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "5KnVl-M0ax14"
},
"outputs": [],
"source": [
"import torch\n",
"import torch.nn as nn\n",
"\n",
"from omegaconf import DictConfig, OmegaConf, open_dict"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "ZEDpkIinbwmm"
},
"outputs": [],
"source": [
"import nemo.collections.asr as nemo_asr # use any domain's models !\n",
"import nemo.collections.tts as nemo_tts # use any domain's models !"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "mLuQo1vnHVcP"
},
"source": [
"# Model Name\n",
"\n",
"NeMo adheres to strict requirements when naming a model for upload to NGC / Hugging Face Hub.\n",
"\n",
"It is **mandatory** to share the model name across the model card, the NeMo file itself. Otherwise NeMo model from Hugging Face will fail to restore correctly."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "MRO2f9fhHywJ"
},
"source": [
"## Naming Convention\n",
"\n",
"NeMo model names can vary based on domain and purpose. While we attempt to conform to standard guidelines when naming our models, we do not expect the same level of strictness for community contributions.\n",
"\n",
"Here are some common guidelines we encourage (but do not enforce) users to follow :\n",
"\n",
"- `Task name`: Usually a short 2-3 character representation of the task that the model performs.\n",
" - `stt` = Speech To Text (ASR)\n",
" - `tts` = Text to Speech (TTS)\n",
" - `ssl` = (Speech) Self Supervised Learning (SSL)\n",
" - `nmt` = Neural Machine Translation (NMT) and many more.\n",
"\n",
"- `Language ID`: Usually a 2/3 digit universal language id. For multilingual models, each domain has its own rules, but some common ones are `{lang_1}{lang_2}{...}` or call them `multilingual.`\n",
"\n",
"- `Model Identifier`: Since models vary so drastically across domains, there is a lot of flexibility here. We try to adhere to naming conventions in literature as much as possible. For example, you can attach `model architecture` (Conformer/Citrinet), `training loss` (CTC/Transducer), and `model size` (small, large, discrete integer sizes, etc.).\n",
"\n",
"- `Optional: Additional Modifiers`: These are additional identifiers such as gender of speaker (TTS), dataset name (ls for Librispeech), etc. It can be set on a case-by-case basis.\n",
"\n",
"All these name segments are jointed by `_`.\n",
"\n",
"-----\n",
"\n",
"As an example of the following model we will try today :\n",
"\n",
"`{task name}_{language id}_{model identifier}_[OPTIONAL modifiers]` = `stt_en_conformer_ctc_small`"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "BjLstKWnPzWV"
},
"source": [
"**Set the MODEL_NAME carefully** !"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "UzHjXDbckU0M"
},
"outputs": [],
"source": [
"MODEL_NAME = \"stt_en_conformer_ctc_small\""
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "qibj1RwvKjSQ"
},
"source": [
"-----\n",
"**Restore a NeMo Model**\n",
"\n",
"Here, we restore a model from NGC directly, but you can restore a model from your training runs using `restore_from()` or use a local .nemo file."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "MsC3pE65d_z2"
},
"outputs": [],
"source": [
"model = nemo_asr.models.ASRModel.from_pretrained(MODEL_NAME)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "y1AkXPFVKfC2"
},
"source": [
"# Create a Hugging Face Model\n",
"\n",
"Now that we have a NeMo model and have logged into Hugging Face with our user API key, we can begin by creating a new repository and uploading our model."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "iv17qFG7KzlL"
},
"source": [
"-----\n",
"\n",
"After the model has been restored, create an HfApi object to interact with the model repository."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "aJUXCOTjKy-2"
},
"outputs": [],
"source": [
"from huggingface_hub import HfApi\n",
"api = HfApi()\n",
"username = api.whoami()['name']"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "DKRlMeaEkeAH"
},
"outputs": [],
"source": [
"try:\n",
" api.create_repo(repo_id=MODEL_NAME)\n",
" print(\"Successfully created repository !\")\n",
"except Exception as e:\n",
" print(\"Repository is possibly already created. Refer to error here - \\n\\n\", e)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Ij6npcneH5tM"
},
"source": [
"## Uploading a model to the hub\n",
"\n",
"There are two ways to upload a NeMo model to the Hugging Face hub -\n",
"\n",
"1) `push_to_hf_hub()`: This is the recommended and automated way to upload NeMo models to the HuggingFace Hub. NeMo will handle all parts of checkpoint and artifact management for you.\n",
"\n",
"2) Hugging Face Hub API: We provide steps to use the lower level Hugging Face Hub API to manually upload a NeMo checkpoint to the hub."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "cANDXRY9ImuW"
},
"source": [
"### Automatic Upload\n",
"\n",
"First, let's take a look at the automated way to upload a checkpoint to the hub *after* creating an empty model.\n",
"\n",
"One important argument is `pack_nemo_file` which decides whether to upload a single `.nemo` file to the hub, or the unpacked NeMo file (with all the individual components of the NeMo file extracted into the directory)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "eCn7y9BmJGEe"
},
"outputs": [],
"source": [
"repo_id = f'{username}/{MODEL_NAME}'\n",
"\n",
"model.push_to_hf_hub(repo_id=repo_id, pack_nemo_file=True)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "62W_zIBrM563"
},
"source": [
"----\n",
"\n",
"Before showing the manual approach, lets remove the uploaded repository first"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Mqu0Aqg2M8ly"
},
"outputs": [],
"source": [
"try:\n",
" api.delete_repo(repo_id=MODEL_NAME, missing_ok=True)\n",
" api.create_repo(repo_id=MODEL_NAME)\n",
" print(\"Successfully created repository !\")\n",
"except Exception as e:\n",
" print(\"Repository is possibly already created. Refer to error here - \\n\\n\", e)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "UJYwaRA9I1Oc"
},
"source": [
"### Manual Upload\n",
"\n",
"Next, let's take a look at how to use the Hugging Face Hub API to upload a NeMo checkpoint to the hub.\n",
"\n",
"Note: Both ways will get similar results, so we recommend the automated way to make it easier."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "aTa4RqDYLGMI"
},
"source": [
"Note two essential names -\n",
"\n",
"- `hf_model_name`: A string name that is the composite of your `username` and `MODEL_NAME` as set above. This name is used for multiple purposes, so keep track of it.\n",
"\n",
"- `model_filename`: The actual filename of the NeMo model that will be uploaded to Hugging Face. Note that this filename is explicitly set to `{MODEL_NAME}.nemo`. If this model filename is altered, then the model cannot correctly be restored by NeMo when downloaded from Hugging Face Hub, so please be careful."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "xhTTMNpBskMS"
},
"outputs": [],
"source": [
"from huggingface_hub import upload_file\n",
"import os\n",
"\n",
"model_filename = f'{MODEL_NAME}.nemo'\n",
"hf_model_name = f'{username}/{MODEL_NAME}'\n",
"commit_message = \"Upload model\"\n",
"path_in_repo = model_filename\n",
"\n",
"model.save_to(model_filename)\n",
"\n",
"upload_file(\n",
" path_or_fileobj=model_filename,\n",
" path_in_repo=path_in_repo,\n",
" repo_id=hf_model_name,\n",
" commit_message=commit_message,\n",
" repo_type='model',\n",
")\n",
"\n",
"os.remove(model_filename)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "BhvNp8MYvxLi"
},
"outputs": [],
"source": [
"print(\"Finished uploading model to :\", hf_model_name)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Qrs-MlW9vVbH"
},
"source": [
"## Test if the model works\n",
"\n",
"Now that we uploaded the model, let's try to use it in NeMo !\n",
"\n",
"The only change required between normally calling `from_pretrained(model_name)` is to call **`from_pretrained({username}/{filename})`**"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "NyuyyRv5snkr"
},
"outputs": [],
"source": [
"hf_model_name = f'{username}/{MODEL_NAME}'\n",
"hf_model = nemo_asr.models.ASRModel.from_pretrained(hf_model_name)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Yhi922WVv4G_"
},
"outputs": [],
"source": [
"print(\"Successfully used HF model -\", hf_model_name)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "9gG1ElJywEJT"
},
"source": [
"# Model Card\n",
"\n",
"Now that we have uploaded the model, we are nearly 50% done!\n",
"\n",
"The next step is to update the model card to have some helpful information regarding the uploaded model and its scores compared to other models.\n",
"\n",
"You can do this in two ways, manually (by clicking the link below) or programmatically fill in part of the model card by following the instructions below."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "aZJRKoxhwBLr"
},
"outputs": [],
"source": [
"hf_url = f'https://huggingface.co/{username}/{MODEL_NAME}'\n",
"print(f\"Visit {hf_url} to manually edit your model card\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "ZlA4hNq6w4rH"
},
"source": [
"-----\n",
"\n",
"Here, we are going to setup some variables for our model card.\n",
"\n",
"First up are the tags:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "QxKtPynWyUWX"
},
"outputs": [],
"source": [
"TAGS = [\n",
" \"automatic-speech-recognition\", # Task id, refer to https://github.com/huggingface/datasets/blob/master/src/datasets/utils/resources/tasks.json for allowed values.\n",
" \"speech\", # add as many other tags as required\n",
" \"audio\",\n",
" \"CTC\",\n",
" \"Conformer\",\n",
" \"Transformer\",\n",
" \"NeMo\", # required for library identification\n",
" \"pytorch\", # required, for toolkit identification\n",
" # \"hf-asr-leaderboard\", # Should only be used if model is evaluated on benchmark scores for ASR.\n",
"]"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Fh7rYWEMM0Vz"
},
"source": [
"-----\n",
"\n",
"Next, we list down all the datasets that were used to train the model.\n",
"\n",
"By convention, try to search if the dataset already exists on Hugging Face Datasets - it is usually listed at the top and in lower case.\n",
"\n",
"If you train on datasets that don't yet exist in Hugging Face Datasets, you can still add them but try to differentiate them by using capitalized names."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "qy-5aDAgzuGD"
},
"outputs": [],
"source": [
"# Replace all spaces with `-`\n",
"DATASETS = [\n",
" \"librispeech_asr\",\n",
" \"mozilla-foundation/common_voice_7_0\",\n",
" \"vctk\",\n",
" \"fisher_corpus\",\n",
" \"Switchboard-1\",\n",
" \"WSJ-0\",\n",
" \"WSJ-1\",\n",
" \"National-Singapore-Corpus-Part-1\",\n",
" \"National-Singapore-Corpus-Part-6\",\n",
"]"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "_0w1X_z4NN5-"
},
"source": [
"-----\n",
"\n",
"Now we create an automated template based on a config for the top portion of the readme file."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "O88WFyPJwjJD"
},
"outputs": [],
"source": [
"from dataclasses import dataclass, field\n",
"from typing import List, Optional, Dict, Any\n",
"\n",
"@dataclass\n",
"class NeMoHuggingFaceModelConfig:\n",
" language: List[str]\n",
" license: str\n",
"\n",
" library_name: str = \"nemo\"\n",
" datasets: List[str] = field(default_factory=lambda: DATASETS)\n",
" thumbnail: Optional[str] = None\n",
" tags: List[str] = field(default_factory=lambda: TAGS)\n",
" model_index: Any = field(default_factory=lambda: [dict(name=MODEL_NAME, results=[])])"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "BpInrBdNxxZ3"
},
"outputs": [],
"source": [
"config = NeMoHuggingFaceModelConfig(language=['en'], license=\"cc-by-4.0\") # choose appropriate license here\n",
"config = OmegaConf.structured(config)\n",
"\n",
"with open_dict(config):\n",
" # Update `model_index` to `model-index`\n",
" model_index = config.pop('model_index')\n",
" config['model-index'] = model_index\n",
"\n",
" # Replace all spaces with `-` in datasets\n",
" normalized_datasets = [ds_name.replace(\" \", \"-\") for ds_name in config['datasets']]\n",
" config['datasets'] = OmegaConf.create(normalized_datasets)\n",
"\n",
"print(OmegaConf.to_yaml(config))"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "0TECX8QrC6FY"
},
"source": [
"## Markdown Template\n",
"\n",
"Now that we have an auto-generated header for our readme, next, we write down some template markdown for the actual contents of the markdown.\n",
"\n",
"You can edit the code here directly if you want, or if you prefer the GUI to see the actual changes in real-time, you can finish uploading this model card and then edit the readme file on the Hugging Face webpage itself."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "SSmm7_OiC9Ex"
},
"outputs": [],
"source": [
"hf_model_name = f'{username}/{MODEL_NAME}'\n",
"\n",
"TEMPLATE = f\"\"\"\n",
"## Model Overview\n",
"\n",
"<DESCRIBE IN ONE LINE THE MODEL AND ITS USE>\n",
"\n",
"## NVIDIA NeMo: Training\n",
"\n",
"To train, fine-tune or play with the model you will need to install [NVIDIA NeMo](https://github.com/NVIDIA/NeMo). We recommend you install it after you've installed latest Pytorch version.\n",
"```\n",
"pip install nemo_toolkit['all']\n",
"```\n",
"\n",
"## How to Use this Model\n",
"\n",
"The model is available for use in the NeMo toolkit [3], and can be used as a pre-trained checkpoint for inference or for fine-tuning on another dataset.\n",
"\n",
"### Automatically instantiate the model\n",
"\n",
"```python\n",
"import nemo.collections.asr as nemo_asr\n",
"asr_model = nemo_asr.models.ASRModel.from_pretrained(\"{hf_model_name}\")\n",
"```\n",
"\n",
"### Transcribing using Python\n",
"First, let's get a sample\n",
"```\n",
"wget https://dldata-public.s3.us-east-2.amazonaws.com/2086-149220-0033.wav\n",
"```\n",
"Then simply do:\n",
"```\n",
"asr_model.transcribe(['2086-149220-0033.wav'])\n",
"```\n",
"\n",
"### Transcribing many audio files\n",
"\n",
"```shell\n",
"python [NEMO_GIT_FOLDER]/examples/asr/transcribe_speech.py \\\n",
" pretrained_name=\"{hf_model_name}\" \\\n",
" audio_dir=\"<DIRECTORY CONTAINING AUDIO FILES>\"\n",
"```\n",
"\n",
"### Input\n",
"\n",
"This model accepts 16000 KHz Mono-channel Audio (wav files) as input.\n",
"\n",
"### Output\n",
"\n",
"This model provides transcribed speech as a string for a given audio sample.\n",
"\n",
"## Model Architecture\n",
"\n",
"<ADD SOME INFORMATION ABOUT THE ARCHITECTURE>\n",
"\n",
"## Training\n",
"\n",
"<ADD INFORMATION ABOUT HOW THE MODEL WAS TRAINED - HOW MANY EPOCHS, AMOUNT OF COMPUTE ETC>\n",
"\n",
"### Datasets\n",
"\n",
"<LIST THE NAME AND SPLITS OF DATASETS USED TO TRAIN THIS MODEL (ALONG WITH LANGUAGE AND ANY ADDITIONAL INFORMATION)>\n",
"\n",
"## Performance\n",
"\n",
"<LIST THE SCORES OF THE MODEL -\n",
" OR\n",
"USE THE Hugging Face Evaluate LiBRARY TO UPLOAD METRICS>\n",
"\n",
"## Limitations\n",
"\n",
"<DECLARE ANY POTENTIAL LIMITATIONS OF THE MODEL>\n",
"\n",
"Eg:\n",
"Since this model was trained on publicly available speech datasets, the performance of this model might degrade for speech which includes technical terms, or vernacular that the model has not been trained on. The model might also perform worse for accented speech.\n",
"\n",
"\n",
"## References\n",
"\n",
"<ADD ANY REFERENCES HERE AS NEEDED>\n",
"\n",
"[1] [NVIDIA NeMo Toolkit](https://github.com/NVIDIA/NeMo)\n",
"\n",
"\"\"\""
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "KPa53S_5NzNp"
},
"source": [
"-----\n",
"\n",
"Below, we will upload this model card in a temporary file called **`\"readme_template.md\"`**. This is done to prevent overwriting of the \"final\" model card that the user may have manually edited.\n",
"\n",
"Once this step is finished, **please copy the contents of this file, create a README.md file and paste the contents into it**."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "0vk5KK4gzpSU"
},
"outputs": [],
"source": [
"filename = 'readme_template.md'\n",
"hf_model_name = f'{username}/{MODEL_NAME}'\n",
"commit_message = \"Upload config\"\n",
"path_in_repo = 'README.md'\n",
"\n",
"with open(filename, 'w') as f:\n",
" f.write(\"---\\n\")\n",
" f.write(OmegaConf.to_yaml(config))\n",
" f.write(\"\\n---\\n\\n\")\n",
" f.write(TEMPLATE)\n",
"\n",
"upload_file(\n",
" path_or_fileobj=filename,\n",
" path_in_repo=path_in_repo,\n",
" repo_id=hf_model_name,\n",
" commit_message=commit_message,\n",
" repo_type='model',\n",
")\n",
"\n",
"os.remove(filename)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "dfXoihCQmWDa"
},
"source": [
"-----\n",
"\n",
"Please visit the URL below to copy the contents of the `readme_template.md` file into your `README.md` file."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "but-5LuLTHFd"
},
"outputs": [],
"source": [
"hf_url = f'https://huggingface.co/{username}/{MODEL_NAME}'\n",
"print(f\"Visit {hf_url} to edit your model card from the generated template file `{filename}`\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "5vPEnlE62dGU"
},
"source": [
"## Evaluation Results\n",
"\n",
"Now that we have both the model checkpoint and the readme uploaded to the Hub, we can optionally add some evaluation results to the card as well!\n",
"\n",
"While this next section is optional, it is highly encouraged to do!"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "rkXMtapA0YzH"
},
"outputs": [],
"source": [
"import evaluate\n",
"# evaluate.list_evaluation_modules(module_type='metric', with_details=True)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "50rzG9Qb3yLR"
},
"outputs": [],
"source": [
"# Uncomment in order to see what values you can supply to the `evaluate` library to push to the Hub.\n",
"# help(evaluate.push_to_hub)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "5A4g3SGf4d0V"
},
"outputs": [],
"source": [
"hf_model_name = f'{username}/{MODEL_NAME}'\n",
"metric_value = 8.1 # value obtained from https://catalog.ngc.nvidia.com/orgs/nvidia/teams/nemo/models/stt_en_conformer_ctc_small\n",
"\n",
"evaluate.push_to_hub(\n",
" model_id=hf_model_name,\n",
" task_type=\"automatic-speech-recognition\",\n",
" dataset_type=\"librispeech_asr\",\n",
" dataset_name=\"Librispeech (clean)\",\n",
" metric_type=\"wer\",\n",
" metric_name=\"WER\",\n",
" dataset_split=\"test\", # corresponds to test-clean set\n",
" dataset_config=\"other\", # corresponds to test-clean set\n",
" dataset_args=dict(language=\"en\"), # metadata for dataset\n",
" # the actual score obtained by the model\n",
" metric_value=metric_value,\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "f3YYa7liO_m3"
},
"source": [
"-----\n",
"\n",
"Done! Now we have a model checkpoint, a model card as well as evaluation results all set up for the NeMo model on Hugging Face!\n",
"\n",
"To add more metrics, you can copy-paste the above cell and repeat the procedure for as many metrics as needed!"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "EtIMGjGwN3fa"
},
"source": [
"# Bonus: Uploading a Model and its Model Card automatically\n",
"\n",
"`push_to_hf_hub()` also supports uploading a model card to Hugging Face Hub in the same step, so here we show an example of this feature."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "b1PSkKZVHREc"
},
"outputs": [],
"source": [
"# Use a different repo name (v2)\n",
"try:\n",
" api.create_repo(repo_id=MODEL_NAME + \"_v2\")\n",
" print(\"Successfully created repository !\")\n",
"except Exception as e:\n",
" print(\"Repository is possibly already created. Refer to error here - \\n\\n\", e)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "L4yX-ULnO9EO"
},
"outputs": [],
"source": [
"# Add some rows to template with placeholders wrapped by {}\n",
"TEMPLATE_2 = TEMPLATE + \"\"\"\n",
"## Original Model Name: {model_name}\n",
"## Repo ID: {repo_id}\n",
"\"\"\"\n",
"kwargs = {\"model_name\": \"ABC\", \"repo_id\": \"nvidia/ABC_XYZ\"}\n",
"model_card_v2 = model.generate_model_card(template=TEMPLATE_2, template_kwargs=kwargs, type=\"hf\") # This is a HF ModelCard object"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "AgZ5zxVwPNWr"
},
"outputs": [],
"source": [
"model.push_to_hf_hub(repo_id=hf_model_name + \"_v2\", pack_nemo_file=False, model_card=model_card_v2)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "WDgwrr2aQyUS"
},
"outputs": [],
"source": [
"hf_model2 = nemo_asr.models.ASRModel.from_pretrained(hf_model_name + \"_v2\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "unQ-6_2-RUFB"
},
"outputs": [],
"source": []
}
],
"metadata": {
"colab": {
"provenance": []
},
"kernelspec": {
"display_name": "Python 3",
"name": "python3"
},
"language_info": {
"name": "python"
}
},
"nbformat": 4,
"nbformat_minor": 0
}
tutorials/VoiceSwapSample.ipynb
+290
View File
@@ -0,0 +1,290 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"colab_type": "text",
"id": "_wIWPxBVc3_O"
},
"source": [
"# Getting Started: Voice swap application\n",
"This notebook shows how to use NVIDIA NeMo (https://github.com/NVIDIA/NeMo) to construct a toy demo which will swap a voice in the audio fragment with a computer generated one.\n",
"\n",
"At its core the demo does: \n",
"\n",
"* Automatic speech recognition of what is said in the file. E.g. converting audio to text\n",
"* Generating spectrogram from resulting text\n",
"* Generating waveform audio from the spectrogram."
]
},
{
"cell_type": "markdown",
"metadata": {
"colab_type": "text",
"id": "gzcsqceVdtj3"
},
"source": [
"## Installation\n",
"NeMo can be installed via simple pip command."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {},
"colab_type": "code",
"id": "I9eIxAyKHREB"
},
"outputs": [],
"source": [
"BRANCH = 'main'\n",
"!python -m pip install \"nemo_toolkit[all] @ git+https://github.com/NVIDIA-NeMo/Speech.git@$BRANCH\"\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {},
"colab_type": "code",
"id": "-X2OyAxreGfl"
},
"outputs": [],
"source": [
"# Ignore pre-production warnings\n",
"import warnings\n",
"warnings.filterwarnings('ignore')\n",
"import nemo\n",
"# Import Speech Recognition collection\n",
"import nemo.collections.asr as nemo_asr\n",
"# Import Speech Synthesis collection\n",
"import nemo.collections.tts as nemo_tts\n",
"# We'll use this to listen to audio\n",
"import IPython"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {},
"colab_type": "code",
"id": "1vC2DHawIGt8"
},
"outputs": [],
"source": [
"# Download audio sample which we'll try\n",
"# This is a sample from LibriSpeech Dev Clean dataset - the model hasn't seen it before\n",
"Audio_sample = '2086-149220-0033.wav'\n",
"!wget https://dldata-public.s3.us-east-2.amazonaws.com/2086-149220-0033.wav\n",
"# Listen to it\n",
"IPython.display.Audio(Audio_sample)"
]
},
{
"cell_type": "markdown",
"metadata": {
"colab_type": "text",
"id": "zodyzdyTVXas"
},
"source": [
"## Instantiate pre-trained NeMo models which we'll use\n",
"``from_pretrained(...)`` API downloads and initialized model directly from the cloud.\n",
"\n",
"We will load audio_sample and convert it to text with QuartzNet ASR model (an action called transcribe).\n",
"To convert text back to audio, we actually need to generate spectrogram with FastPitch first and then convert it to actual audio signal using the HiFiGAN vocoder."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {},
"colab_type": "code",
"id": "f_J9cuU1H6Bn"
},
"outputs": [],
"source": [
"# Speech Recognition model - FastConformer\n",
"asr_model = nemo_asr.models.ASRModel.from_pretrained(model_name=\"stt_en_fastconformer_ctc_large\").cuda()\n",
"\n",
"# Spectrogram generator which takes text as an input and produces spectrogram\n",
"spectrogram_generator = nemo_tts.models.FastPitchModel.from_pretrained(model_name=\"tts_en_fastpitch\").cuda()\n",
"\n",
"# Vocoder model which takes spectrogram and produces actual audio\n",
"vocoder = nemo_tts.models.HifiGanModel.from_pretrained(model_name=\"tts_en_hifigan\").cuda()"
]
},
{
"cell_type": "markdown",
"metadata": {
"colab_type": "text",
"id": "jQSj-IhEhrtI"
},
"source": [
"## Using the models"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {},
"colab_type": "code",
"id": "s0ERrXIzKpwu"
},
"outputs": [],
"source": [
"# Convert our audio sample to text\n",
"files = [Audio_sample]\n",
"raw_text = ''\n",
"text = ''\n",
"for fname, transcription in zip(files, asr_model.transcribe(audio=files)):\n",
" raw_text = transcription\n",
"\n",
"print(f'\\nRaw recognized text: {raw_text}')"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {},
"colab_type": "code",
"id": "-0Sk0C9-LmAR"
},
"outputs": [],
"source": [
"# A helper function which combines TTS models to go directly from \n",
"# text to audio\n",
"def text_to_audio(text):\n",
" parsed = spectrogram_generator.parse(text)\n",
" spectrogram = spectrogram_generator.generate_spectrogram(tokens=parsed)\n",
" audio = vocoder.convert_spectrogram_to_audio(spec=spectrogram)\n",
" return audio.to('cpu').detach().numpy()"
]
},
{
"cell_type": "markdown",
"metadata": {
"colab_type": "text",
"id": "Q8Jvwe4Ahncx"
},
"source": [
"## Results"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {},
"colab_type": "code",
"id": "-im5TDF-MP2N"
},
"outputs": [],
"source": [
"# This is our original audio sample\n",
"IPython.display.Audio(Audio_sample)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {},
"colab_type": "code",
"id": "SNOMquwviEEQ"
},
"outputs": [],
"source": [
"# This is what was recognized by the ASR model\n",
"print(raw_text)"
]
},
{
"cell_type": "markdown",
"metadata": {
"colab_type": "text",
"id": "di2IzMsdiiWq"
},
"source": [
"Compare how the synthesized audio sounds when using text with and without punctuation."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {},
"colab_type": "code",
"id": "EIh8wTVs5uH7"
},
"outputs": [],
"source": [
"# Without punctuation\n",
"IPython.display.Audio(text_to_audio(raw_text), rate=22050)"
]
},
{
"cell_type": "markdown",
"metadata": {
"colab_type": "text",
"id": "JOEFYywbctbJ"
},
"source": [
"## Next steps\n",
"A demo like this is great for prototyping and experimentation. However, for real production deployment, you would want to use a service like [NVIDIA Riva](https://developer.nvidia.com/riva).\n",
"\n",
"**NeMo is built for training.** You can fine-tune, or train from scratch on your data all models used in this example. We recommend you checkout the following, more in-depth, tutorials next:\n",
"\n",
"* [NeMo fundamentals](https://colab.research.google.com/github/NVIDIA/NeMo/blob/stable/tutorials/00_NeMo_Primer.ipynb)\n",
"* [NeMo models](https://colab.research.google.com/github/NVIDIA/NeMo/blob/stable/tutorials/01_NeMo_Models.ipynb)\n",
"* [Speech Recognition](https://colab.research.google.com/github/NVIDIA/NeMo/blob/stable/tutorials/asr/ASR_with_NeMo.ipynb)\n",
"* [Speech Synthesis](https://colab.research.google.com/github/NVIDIA/NeMo/blob/stable/tutorials/tts/Inference_ModelSelect.ipynb)\n",
"\n",
"\n",
"You can find scripts for training and fine-tuning ASR and TTS models [here](https://github.com/NVIDIA/NeMo/tree/main/examples). "
]
},
{
"cell_type": "markdown",
"metadata": {
"colab_type": "text",
"id": "ahRh2Y0Lc0G1"
},
"source": [
"That's it folks! Head over to NeMo GitHub for more examples: https://github.com/NVIDIA/NeMo"
]
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"collapsed_sections": [],
"name": "NeMo voice swap app",
"provenance": [],
"toc_visible": true
},
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.8.12"
}
},
"nbformat": 4,
"nbformat_minor": 4
}
tutorials/asr/ASR_Confidence_Estimation.ipynb
+1447
View File
File diff suppressed because it is too large Load Diff
tutorials/asr/ASR_Context_Biasing.ipynb
+956
View File
@@ -0,0 +1,956 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "17b3cbf8",
"metadata": {},
"source": [
"# Context-Biasing for ASR models with CTC-based Word Spotter"
]
},
{
"cell_type": "markdown",
"id": "1156d1d1",
"metadata": {},
"source": [
"This tutorial aims to show how to improve the recognition accuracy of specific words in NeMo Framework\n",
"for CTC and Trasducer (RNN-T) ASR models by using the fast context-biasing method with CTC-based Word Spotter.\n",
"\n",
"## Tutorial content:\n",
"* Intro in the context-biasing problem\n",
"* Description of the proposed CTC-based Words Spotter (CTC-WS) method\n",
"* Practical part 1 (base):\n",
" * Download data set and ASR models\n",
" * Build context-biasing list\n",
" * Evaluate recognition results with and without context-biasing\n",
" * Improve context-biasing results with alternative transcriptions\n",
"* Practical part 2 (advanced):\n",
" * Visualization of context-biasing graph\n",
" * Running CTC-based Word Spotter only\n",
" * Merge greedy decoding results with spotted context-biasing candidates\n",
" * Results analysis\n",
"* Summary"
]
},
{
"cell_type": "markdown",
"id": "431edfbf",
"metadata": {},
"source": [
"## Context-biasing: intro\n",
"\n",
"ASR models often struggle to recognize words that were absent or had few examples in the training data.\n",
"This problem is especially acute due to the emergence of new names and titles in a rapidly developing world.\n",
"The users need to be able to recognize these new words.\n",
"Context-biasing methods attempt to solve this problem by assuming that we have a list of words and phrases (context-biasing list) in advance\n",
"for which we want to improve recognition accuracy.\n",
"\n",
"One of the directions of context-biasing methods is based on the `deep fusion` approach.\n",
"These methods require intervention into the ASR model and its training process.\n",
"The main disadvantage of these methods is that they require a lot of computational resources and time to train the model.\n",
"\n",
"Another direction is methods based on the `shallow fusion` approach. In this case, only the decoding process is modified.\n",
"During the beam-search decoding, the hypothesis is rescored depending on whether the current word is present in the context-biasing list or external language model.\n",
"The beam-search decoding may be computationally expensive, especially for the models with a large vocabulary and context-biasing list.\n",
"This problem is considerably worsened in the case of the Transducer (RNN-T) model since beam-search decoding involves multiple Decoder (Prediction) and Joint networks calculations.\n",
"Moreover, the context-biasing recognition is limited by the model prediction pool biased toward training data. In the case of rare or new words, the model may not have a hypothesis for the desired word from the context-biasing list whose probability we want to amplify."
]
},
{
"cell_type": "markdown",
"id": "ae0bfd60",
"metadata": {},
"source": [
"## CTC-based Word Spotter\n",
"\n",
"\n",
"This tutorial considers a fast context-biasing method using a CTC-based Word Spotter (CTC-WS).\n",
"The method involves decoding CTC log probabilities with a context graph built for words and phrases from the context-biasing list.\n",
"The spotted context-biasing candidates (with their scores and time intervals) are compared by scores with words from the greedy\n",
"CTC decoding results to improve recognition accuracy and pretend false accepts of context-biasing (Figure 1). \n",
" \n",
" \n",
"\n",
"\n",
"\n"
]
},
{
"cell_type": "markdown",
"id": "c7e45bf2",
"metadata": {},
"source": [
"<figure markdown>\n",
" <img src=\"https://github.com/NVIDIA/NeMo/releases/download/v1.22.0/asset-post-v1.22.0-ctcws_scheme_2.png\" alt=\"CTC-WS\" style=\"width: 60%;\" height=\"auto\"> <!-- Adjust the width as needed -->\n",
" <figcaption><b>Figure 1.</b> <i> High-level representation of the proposed context-biasing method with CTC-WS in case of CTC model. Detected words (gpu, nvidia, cuda) are compared with words from the greedy CTC results in the overlapping intervals according to the accumulated scores to prevent false accept replacement. </i></figcaption>\n",
"</figure>"
]
},
{
"cell_type": "markdown",
"id": "ba163f41",
"metadata": {},
"source": [
"\n",
"<!-- <img width=\"500px\" height=\"auto\"\n",
" src=\"https://github.com/NVIDIA/NeMo/releases/download/v1.22.0/asset-post-v1.22.0-ctcws_scheme_2.png\"\n",
" alt=\"CTC-WS2\"\n",
" style=\"float: right; margin-left: 20px;\"> -->\n",
" \n",
"A [Hybrid Transducer-CTC](https://arxiv.org/abs/2312.17279) model (a shared encoder trained together with CTC and Transducer output heads) enables the use of the CTC-WS method for the Transducer model.\n",
"Context-biasing candidates obtained by CTC-WS are also filtered by the scores with greedy CTC predictions and then merged with greedy Transducer results.\n",
"\n",
"The CTC-WS method allows using pretrained NeMo models (`CTC` or `Hybrid Transducer-CTC`) for context-biasing recognition without model retraining (Figure 2).\n",
"The method shows inspired results for context-biasing with only a little additional work time and computational resources.\n",
"\n"
]
},
{
"cell_type": "markdown",
"id": "c05b16d8",
"metadata": {},
"source": [
"<figure markdown>\n",
" <img src=\"https://github.com/NVIDIA/NeMo/releases/download/v1.22.0/asset-post-v1.22.0-ctcws_scheme_1.png\" alt=\"CTC-WS\" style=\"width: 65%;\" align=\"center\"> <!-- Adjust the width as needed -->\n",
" <figcaption><b>Figure 2.</b> <i> Scheme of the context-biasing method with CTC-based Word Spotter. CTC-WS uses CTC log probabilities to detect context-biasing candidates. Obtained candidates are filtered by CTC word alignment and then merged with CTC or RNN-T word alignment to get the final text result. </i></figcaption>\n",
"</figure>"
]
},
{
"cell_type": "markdown",
"id": "ac0ec822",
"metadata": {},
"source": [
"# Installing dependencies"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "69c86a4f",
"metadata": {},
"outputs": [],
"source": [
"BRANCH = 'main'\n",
"\n",
"\"\"\"\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",
"\n",
"import os\n",
"# either provide a path to local NeMo repository with NeMo already installed or git clone\n",
"\n",
"# option #1: local path to NeMo repo with NeMo already installed\n",
"NEMO_DIR_PATH = os.path.dirname(os.path.dirname(os.path.abspath('')))\n",
"\n",
"# check if Google Colab is being used\n",
"try:\n",
" import google.colab\n",
" IN_COLAB = True\n",
"except (ImportError, ModuleNotFoundError):\n",
" IN_COLAB = False\n",
"\n",
"# option #2: download NeMo repo\n",
"if IN_COLAB or not os.path.exists(os.path.join(NEMO_DIR_PATH, \"nemo\")):\n",
" ## Install dependencies\n",
" !apt-get install sox libsndfile1 ffmpeg\n",
"\n",
" !git clone -b $BRANCH https://github.com/NVIDIA-NeMo/Speech\n",
" %cd NeMo\n",
" !python -m pip install \"nemo_toolkit[all] @ git+https://github.com/NVIDIA-NeMo/Speech.git@$BRANCH\"\n",
" NEMO_DIR_PATH = os.path.abspath('')\n",
"\n",
"import sys\n",
"sys.path.insert(0, NEMO_DIR_PATH)"
]
},
{
"cell_type": "markdown",
"id": "5260d4fa",
"metadata": {},
"source": [
"## Practical part 1 (base)\n",
"In this part, we will consider the base usage of the CTC-WS method for pretrained NeMo models.\n",
"\n",
"### Data preparation.\n",
"We will use a subset of the GTC data set. The data set contains 10 audio files with NVIDIA GTC talks. \n",
"The primary data set feature is the computer science and engineering domain, which has a large number of unique terms and product names (NVIDIA, GPU, GeForce, Ray Tracing, Omniverse, teraflops, etc.), which is good fit for the context-biasing task. All the text data is normalized and lowercased."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "637f2c6d",
"metadata": {},
"outputs": [],
"source": [
"# download data\n",
"!wget https://asr-tutorial-data.s3.eu-north-1.amazonaws.com/context_biasing_data.gz\n",
"!tar -xvzf context_biasing_data.gz\n",
"!apt-get update && apt-get upgrade -y && apt-get install tree"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "6baefc80",
"metadata": {},
"outputs": [],
"source": [
"!tree context_biasing_data"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "09fe748b",
"metadata": {},
"outputs": [],
"source": [
"from nemo.collections.asr.parts.utils.manifest_utils import read_manifest\n",
"\n",
"# data is already stored in nemo data manifest format\n",
"test_nemo_manifest = \"./context_biasing_data/gtc_data_subset_10f.json\"\n",
"test_data = read_manifest(test_nemo_manifest)\n",
"\n",
"for idx, item in enumerate(test_data):\n",
" print(f\"[{idx}]: {item['text']}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "64ab4764",
"metadata": {},
"outputs": [],
"source": [
"import librosa\n",
"import IPython.display as ipd\n",
"\n",
"# load and listen to the audio file example\n",
"example_file = test_data[0]['audio_filepath']\n",
"audio, sample_rate = librosa.load(example_file)\n",
"\n",
"file_id = 0\n",
"print(f\"[TEXT {file_id}]: {test_data[file_id]['text']}\\n\")\n",
"ipd.Audio(example_file, rate=sample_rate)"
]
},
{
"cell_type": "markdown",
"id": "a85ea8ec",
"metadata": {},
"source": [
"### Load ASR models\n",
"\n",
"For testing the CTC-WS method, we will use the following NeMo models:\n",
" - (CTC): [stt_en_fastconformer_ctc_large](https://huggingface.co/nvidia/stt_en_fastconformer_ctc_large) - a large fast-conformer model trained on English ASR data\n",
" - (Hybrid Transducer-CTC): [stt_en_fastconformer_hybrid_large_streaming_multi](https://huggingface.co/nvidia/stt_en_fastconformer_hybrid_large_streaming_multi) - a large fast-conformer model trained jointly with CTC and Transducer heads on English ASR data. The model is streaming, which means it can process audio in real time. It can cause a slight WER degradation in comparison with the first offline model."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "d34ee0ba",
"metadata": {
"scrolled": true
},
"outputs": [],
"source": [
"from nemo.collections.asr.models import EncDecCTCModelBPE, EncDecHybridRNNTCTCBPEModel\n",
"\n",
"# ctc model\n",
"ctc_model_name = \"stt_en_fastconformer_ctc_large\"\n",
"ctc_model = EncDecCTCModelBPE.from_pretrained(model_name=ctc_model_name)\n",
"\n",
"# hybrid transducer-ctc model\n",
"hybrid_ctc_rnnt_model_name = \"stt_en_fastconformer_hybrid_large_streaming_multi\""
]
},
{
"cell_type": "markdown",
"id": "082208cd",
"metadata": {},
"source": [
"### Transcribe \n",
"Let's transcribe test data and analyze the regontion accuracy of specific words "
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "74436885",
"metadata": {
"scrolled": true
},
"outputs": [],
"source": [
"test_audio_files = [item['audio_filepath'] for item in test_data]\n",
"recog_results = ctc_model.transcribe(test_audio_files)"
]
},
{
"cell_type": "markdown",
"id": "b993d650",
"metadata": {},
"source": [
"### Compute per-word recognition statisctic"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "70f5714b",
"metadata": {},
"outputs": [],
"source": [
"from kaldialign import align\n",
"\n",
"word_dict = {} # {word: [num_of_occurances, num_of_correct_recognition]}\n",
"eps = \"<eps>\"\n",
"ref_text = [item['text'] for item in test_data]\n",
"\n",
"for idx, ref in enumerate(ref_text):\n",
" ref = ref.split()\n",
" hyp = recog_results[idx].text.split()\n",
" ali = align(ref, hyp, eps)\n",
"\n",
" for pair in ali:\n",
" word_ref, word_hyp = pair\n",
" if word_ref == eps:\n",
" continue\n",
" if word_ref in word_dict:\n",
" word_dict[word_ref][0] += 1\n",
" else:\n",
" word_dict[word_ref] = [1, 0]\n",
" if word_ref == word_hyp:\n",
" word_dict[word_ref][1] += 1\n",
"\n",
"word_candidats = {}\n",
"\n",
"for word in word_dict:\n",
" gt = word_dict[word][0]\n",
" tp = word_dict[word][1]\n",
" if tp/gt < 1.0:\n",
" word_candidats[word] = [gt, round(tp/gt, 2)]\n",
" \n",
"# print obtained per-word statistic\n",
"word_candidats_sorted = sorted(word_candidats.items(), key=lambda x:x[1][0], reverse=True)\n",
"max_word_len = max([len(x[0]) for x in word_candidats_sorted])\n",
"for item in word_candidats_sorted:\n",
" print(f\"{item[0]:<{max_word_len}} {item[1][0]}/{item[1][1]}\")"
]
},
{
"cell_type": "markdown",
"id": "27a9f88b",
"metadata": {},
"source": [
"## Create a context-biasing list\n",
"\n",
"Now, we need to select the words, recognition of which we want to improve by CTC-WS context-biasing.\n",
"Usually, we select only nontrivial words with the lowest recognition accuracy.\n",
"Such words should have a character length >= 3 because short words in a context-biasing list may produce high false-positive recognition.\n",
"In this toy example, we will select all the words that look like names with a recognition accuracy less than 1.0.\n",
"\n",
"The structure of the context-biasing file is:\n",
"\n",
"WORD1_TRANSCRIPTION1 \n",
"WORD2_TRANSCRIPTION1 \n",
"...\n",
"\n",
"TRANSCRIPTION here is a word spelling. We need this structure to add alternative transcriptions (spellings) for some word. We will cover such a case further."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "c27848f0",
"metadata": {},
"outputs": [],
"source": [
"cb_words = [\"gpu\", \"nvidia\", \"nvidia's\", \"nvlink\", \"omniverse\", \"cunumeric\", \"numpy\", \"dgx\", \"dgxs\", \"dlss\",\n",
" \"cpu\", \"tsmc\", \"culitho\", \"xlabs\", \"tensorrt\", \"tensorflow\", \"pytorch\", \"aws\", \"chatgpt\", \"pcie\"]\n",
"\n",
"# write context-biasing file \n",
"cb_list_file = \"context_biasing_data/context_biasing_list.txt\"\n",
"with open(cb_list_file, \"w\", encoding=\"utf-8\") as fn:\n",
" for word in cb_words:\n",
" fn.write(f\"{word}_{word}\\n\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "b0e8c800",
"metadata": {},
"outputs": [],
"source": [
"!cat {cb_list_file}"
]
},
{
"cell_type": "markdown",
"id": "c44fc910",
"metadata": {},
"source": [
"## Run context-biasing evaluation\n",
"\n",
"The main script for CTC-WS context-biasing in NeMo is:\\\n",
"`{NEMO_DIR_PATH}/scripts/asr_context_biasing/eval_greedy_decoding_with_context_biasing.py`\n",
"\n",
"Context-biasing is managed by `apply_context_biasing` parameter [true or false]. \n",
"Other important context-biasing parameters are:\n",
"- `beam_threshold` - threshold for CTC-WS beam pruning\n",
"- `context_score` - per token weight for context biasing\n",
"- `ctc_ali_token_weight` - per token weight for CTC alignment (prevents false acceptances of context-biasing words) \n",
"\n",
"All the context-biasing parameters are selected according to the default values in the script. \n",
"You can tune them according to your data and ASR model (list all the values in the [] separated by commas) \n",
"for example: `beam_threshold=[7.0,8.0,9.0]`, `context_score=[3.0,4.0,5.0]`, `ctc_ali_token_weight=[0.5,0.6,0.7]`. \n",
"The script will run the recognition with all the combinations of the parameters and will select the best one based on WER value."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "9a2d32e9",
"metadata": {},
"outputs": [],
"source": [
"# create directory with experimental results\n",
"import os\n",
"\n",
"exp_dir = \"exp\"\n",
"if not os.path.isdir(exp_dir):\n",
" os.makedirs(exp_dir)\n",
"else:\n",
" print(f\"Directory '{exp_dir}' already exists\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "116f2abe",
"metadata": {
"scrolled": true
},
"outputs": [],
"source": [
"# ctc model (no context-biasing)\n",
"\n",
"!python {NEMO_DIR_PATH}/scripts/asr_context_biasing/eval_greedy_decoding_with_context_biasing.py \\\n",
" nemo_model_file={ctc_model_name} \\\n",
" input_manifest={test_nemo_manifest} \\\n",
" preds_output_folder={exp_dir} \\\n",
" decoder_type=\"ctc\" \\\n",
" acoustic_batch_size=64 \\\n",
" apply_context_biasing=false \\\n",
" context_file={cb_list_file} \\\n",
" beam_threshold=[7.0] \\\n",
" context_score=[3.0] \\\n",
" ctc_ali_token_weight=[0.5]"
]
},
{
"cell_type": "markdown",
"id": "674d0af1",
"metadata": {},
"source": [
"The results must be:\n",
"\n",
"`Precision`: 1.0000 (1/1) fp:0 (fp - false positive recognition) \n",
"`Recall`: 0.0333 (1/30) \n",
"`Fscore`: 0.0645 \n",
"`Greedy WER/CER` = 35.68%/8.16%\n",
"\n",
"The model could recognize 1 out of 30 words from the context-biasing list.\n",
"Let's enable context-biasing during decoding:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "239da41d",
"metadata": {
"scrolled": true
},
"outputs": [],
"source": [
"# ctc model (with context biasing)\n",
"!python {NEMO_DIR_PATH}/scripts/asr_context_biasing/eval_greedy_decoding_with_context_biasing.py \\\n",
" nemo_model_file={ctc_model_name} \\\n",
" input_manifest={test_nemo_manifest} \\\n",
" preds_output_folder={exp_dir} \\\n",
" decoder_type=\"ctc\" \\\n",
" acoustic_batch_size=64 \\\n",
" apply_context_biasing=true \\\n",
" context_file={cb_list_file} \\\n",
" beam_threshold=[7.0] \\\n",
" context_score=[3.0] \\\n",
" ctc_ali_token_weight=[0.5]"
]
},
{
"cell_type": "markdown",
"id": "faa1e73c",
"metadata": {},
"source": [
"Now, recognition results are much better:\n",
"\n",
"`Precision`: 1.0000 (21/21) fp:0 \n",
"`Recall`: 0.7000 (21/30) \n",
"`Fscore`: 0.8235 \n",
"`Greedy WER/CER` = 17.09%/4.43%\n",
"\n",
"But we are still able to recognize only 21 out of 30 specific words.\\\n",
"You can see that unrecognized words are mostly abbreviations (`dgxs`, `dlss`, `gpu`, `aws`, etc.) or compound words (`culitho`).\\\n",
"The ASR models tends to recognize such words as a sequence of characters (`\"aws\" -> \"a w s\"`) or subwords (`\"culitho\" -> \"cu litho\"`).\\\n",
"We can try to improve the recognition of such words by adding alternative transcriptions to the context-biasing list."
]
},
{
"cell_type": "markdown",
"id": "d72b6391",
"metadata": {},
"source": [
"### Alternative transcriptions\n",
"\n",
"wordninja is used to split compound words into simple words according to the default word dictionary."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "f7e00263",
"metadata": {},
"outputs": [],
"source": [
"!pip install wordninja"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "46fe91e9",
"metadata": {},
"outputs": [],
"source": [
"import wordninja\n",
"\n",
"cb_list_file_modified = cb_list_file + \".abbr_and_ninja\"\n",
"\n",
"with open(cb_list_file, \"r\", encoding=\"utf-8\") as fn1, \\\n",
" open(cb_list_file_modified, \"w\", encoding=\"utf-8\") as fn2:\n",
"\n",
" for line in fn1:\n",
" word = line.strip().split(\"_\")[0]\n",
" new_line = f\"{word}_{word}\"\n",
" # split all the short words into characters\n",
" if len(word) <= 4 and len(word.split()) == 1:\n",
" abbr = ' '.join(list(word))\n",
" new_line += f\"_{abbr}\"\n",
" # split the long words into the simple words (not for phrases)\n",
" new_segmentation = wordninja.split(word)\n",
" if word != new_segmentation[0]:\n",
" new_segmentation = ' '.join(new_segmentation)\n",
" new_line += f\"_{new_segmentation}\"\n",
" fn2.write(f\"{new_line}\\n\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "da69da45",
"metadata": {},
"outputs": [],
"source": [
"!cat {cb_list_file_modified}"
]
},
{
"cell_type": "markdown",
"id": "4a21cbf4",
"metadata": {},
"source": [
"Run context-biasing with modified context-biasing list:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "913a0f5e",
"metadata": {
"scrolled": true
},
"outputs": [],
"source": [
"# ctc models (with context biasing and modified cb list)\n",
"!python {NEMO_DIR_PATH}/scripts/asr_context_biasing/eval_greedy_decoding_with_context_biasing.py \\\n",
" nemo_model_file={ctc_model_name} \\\n",
" input_manifest={test_nemo_manifest} \\\n",
" preds_output_folder={exp_dir} \\\n",
" decoder_type=\"ctc\" \\\n",
" acoustic_batch_size=64 \\\n",
" apply_context_biasing=true \\\n",
" context_file={cb_list_file_modified} \\\n",
" beam_threshold=[7.0] \\\n",
" context_score=[3.0] \\\n",
" ctc_ali_token_weight=[0.5]"
]
},
{
"cell_type": "markdown",
"id": "654751ed",
"metadata": {},
"source": [
"Now, the recognition results are:\n",
"\n",
"`Precision`: 1.0000 (28/28) fp:1 \n",
"`Recall`: 0.9333 (28/30) \n",
"`Fscore`: 0.9655 \n",
"`Greedy WER/CER` = 7.04%/2.93%\n",
"\n",
"As you can see, that adding alternative transcriptions to the cb_list file improved the recognition accuracy of the context-biasing words. However, we still miss 2 words. Unfortunately, this algorithm is not a silver bullet.\n",
"\n",
"In some cases, you can improve results by adding alternative transcriptions manually based on the recognition errors of your ASR model for the specific words (for example, `\"nvidia\" -> \"n video\"`). "
]
},
{
"cell_type": "markdown",
"id": "b96c4023",
"metadata": {},
"source": [
"### Hybrid Transducer-CTC model\n",
"The CTC-WS context-biasing method for Transducer (RNN-T) models is supported only for Hybrid Transducer-CTC model. \n",
"To use Transducer head of the Hybrid Transducer-CTC model, we need to set `decoder_type=\"rnnt\"`. \n",
"Other parameters are the same as for the CTC model because the context-biasing is applied only on the CTC part of the model. Spotted context-biasing words will have been merged with greedy decoding results of the Transducer head.\n",
"\n",
"We can use already prepared context-biasing list because the CTC and Hybrid Transducer-CTC models have almost the same BPE tokenizer."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "456e47df",
"metadata": {
"scrolled": true
},
"outputs": [],
"source": [
"# Transducer model (no context-biasing)\n",
"!python {NEMO_DIR_PATH}/scripts/asr_context_biasing/eval_greedy_decoding_with_context_biasing.py \\\n",
" nemo_model_file={hybrid_ctc_rnnt_model_name} \\\n",
" input_manifest={test_nemo_manifest} \\\n",
" preds_output_folder={exp_dir} \\\n",
" decoder_type=\"rnnt\" \\\n",
" acoustic_batch_size=64 \\\n",
" apply_context_biasing=false \\\n",
" context_file={cb_list_file_modified} \\\n",
" beam_threshold=[7.0] \\\n",
" context_score=[3.0] \\\n",
" ctc_ali_token_weight=[0.5]"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "773e11f1",
"metadata": {
"scrolled": true
},
"outputs": [],
"source": [
"# Transducer model (with context-biasing)\n",
"!python {NEMO_DIR_PATH}/scripts/asr_context_biasing/eval_greedy_decoding_with_context_biasing.py \\\n",
" nemo_model_file={hybrid_ctc_rnnt_model_name} \\\n",
" input_manifest={test_nemo_manifest} \\\n",
" preds_output_folder={exp_dir} \\\n",
" decoder_type=\"rnnt\" \\\n",
" acoustic_batch_size=64 \\\n",
" apply_context_biasing=true \\\n",
" context_file={cb_list_file_modified} \\\n",
" beam_threshold=[7.0] \\\n",
" context_score=[3.0] \\\n",
" ctc_ali_token_weight=[0.5]"
]
},
{
"cell_type": "markdown",
"id": "45a91385",
"metadata": {},
"source": [
"CTC-WS context-biasing works for Transducer model as well as for CTC (`F-score improvenment: 0.3784 -> 0.9286`). Differences in the nature of offline and online models may cause differences in results (usually, online models have a tendency to predict tokens earlier what can affect the difference between the timestamps of CTC and RNN-T models). "
]
},
{
"cell_type": "markdown",
"id": "1968e7bc",
"metadata": {},
"source": [
"## Practical part 2 (advanced)\n",
"In this section, we will consider the context-biasing process more deeply:\n",
"- Visualization of the context-biasing graph\n",
"- Running CTC-WS with the context-biasing graph\n",
"- Merge the obtained spotted words with greedy decoding results\n",
"- Analysis of the results"
]
},
{
"cell_type": "markdown",
"id": "277104b5",
"metadata": {},
"source": [
"### Build a context graph (for visualization only)\n",
"The context graph consists of a composition of a prefix tree (Trie) with the CTC transition topology for words and phrases from the context-biasing list. We use a BPE tokenizer from the target ASR model for word segmentation."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "55a36a27-919c-4d64-9163-b0b2c9dca15e",
"metadata": {},
"outputs": [],
"source": [
"# install graphviz from source in case of local run (not Google Colab)\n",
"# this may take about 5-10 minutes\n",
"# make sure that env variables have been set\n",
"\n",
"if not IN_COLAB:\n",
"\n",
" os.environ['DEBIAN_FRONTEND'] = 'noninteractive'\n",
" os.environ['TZ'] = 'Etc/UTC'\n",
"\n",
" !echo $DEBIAN_FRONTEND\n",
" !echo $TZ\n",
"\n",
" !{NEMO_DIR_PATH}/scripts/installers/install_graphviz.sh"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "904ea41b",
"metadata": {},
"outputs": [],
"source": [
"from nemo.collections.asr.parts import context_biasing\n",
"\n",
"# get bpe tokenization\n",
"cb_words_small = [\"nvidia\", \"gpu\", \"nvlink\", \"numpy\"]\n",
"context_transcripts = []\n",
"for word in cb_words_small:\n",
" # use text_to_tokens method for viasualization only\n",
" word_tokenization = ctc_model.tokenizer.text_to_tokens(word)\n",
" print(f\"{word}: {word_tokenization}\")\n",
" context_transcripts.append([word, [word_tokenization]])\n",
"\n",
"# build context graph\n",
"context_graph = context_biasing.ContextGraphCTC(blank_id=\"⊘\")\n",
"context_graph.add_to_graph(context_transcripts)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "f7fab1e1",
"metadata": {},
"outputs": [],
"source": [
"context_graph.draw()"
]
},
{
"cell_type": "markdown",
"id": "04a6f4be",
"metadata": {},
"source": [
"### Build a real context graph (for decoding)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "2ba2d8a1",
"metadata": {},
"outputs": [],
"source": [
"# get bpe tokenization\n",
"context_transcripts = []\n",
"for word in cb_words:\n",
" word_tokenization = [ctc_model.tokenizer.text_to_ids(x) for x in word]\n",
" context_transcripts.append([word, word_tokenization])\n",
"\n",
"# build context graph\n",
"context_graph = context_biasing.ContextGraphCTC(blank_id=ctc_model.decoding.blank_id)\n",
"context_graph.add_to_graph(context_transcripts)"
]
},
{
"cell_type": "markdown",
"id": "71e0e86b",
"metadata": {},
"source": [
"### Run CTC-based Word Spotter\n",
"\n",
"The CTC-WS task is to search for words by decoding CTC log probabilities using the context graph. As a result, we obtain a list of detected words with exact start/end frames in the audio file and their overall scores. The relatively small size of the context graph and hypotheses pruning methods allow this algorithm to work very quickly."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "7c2c942e-e8df-4c09-a7de-87606ae453c9",
"metadata": {},
"outputs": [],
"source": [
"%pip install --upgrade jupyter ipywidgets"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "f2bc370b",
"metadata": {},
"outputs": [],
"source": [
"import torch\n",
"from tqdm import tqdm\n",
"\n",
"# get ctc logprobs\n",
"audio_file_paths = [item['audio_filepath'] for item in test_data]\n",
"\n",
"with torch.no_grad():\n",
" ctc_model.eval()\n",
" ctc_model.encoder.freeze()\n",
" device = next(ctc_model.parameters()).device\n",
" hyp_results = ctc_model.transcribe(audio_file_paths, batch_size=10, return_hypotheses=True)\n",
" ctc_logprobs = [hyp.alignments.cpu().numpy() for hyp in hyp_results]\n",
" blank_idx = ctc_model.decoding.blank_id\n",
" \n",
"# run ctc-based word spotter\n",
"ws_results = {}\n",
"for idx, logits in tqdm(\n",
" enumerate(ctc_logprobs), desc=f\"Eval CTC-based Word Spotter...\", total=len(ctc_logprobs)\n",
"):\n",
" ws_results[audio_file_paths[idx]] = context_biasing.run_word_spotter(\n",
" logits,\n",
" context_graph,\n",
" ctc_model,\n",
" blank_idx=blank_idx,\n",
" beam_threshold=7.0,\n",
" cb_weight=3.0,\n",
" ctc_ali_token_weight=0.5,\n",
" )"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "4bd6645c",
"metadata": {},
"outputs": [],
"source": [
"# print CTC-WS hypotheses for the first audio file\n",
"ws_results[audio_file_paths[0]]"
]
},
{
"cell_type": "markdown",
"id": "245a66f0",
"metadata": {},
"source": [
"### Merge CTC-WS words with greedy CTC decoding results\n",
"\n",
"Use `print_stats=True` to get more information about spotted words and greedy CTC word alignment."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "423b2b9e",
"metadata": {},
"outputs": [],
"source": [
"import numpy as np\n",
"\n",
"target_transcripts = [item['text'] for item in test_data]\n",
"\n",
"# merge spotted words with greedy results\n",
"for idx, logprobs in enumerate(ctc_logprobs):\n",
" greedy_predicts = np.argmax(logprobs, axis=1)\n",
" if ws_results[audio_file_paths[idx]]:\n",
" # make new text by mearging alignment with ctc-ws predictions:\n",
" print(\"\\n\" + \"********\" * 10)\n",
" print(f\"File name: {audio_file_paths[idx]}\")\n",
" pred_text, raw_text = context_biasing.merge_alignment_with_ws_hyps(\n",
" greedy_predicts,\n",
" ctc_model,\n",
" ws_results[audio_file_paths[idx]],\n",
" decoder_type=\"ctc\",\n",
" blank_idx=blank_idx,\n",
" print_stats=True,\n",
" )\n",
" print(f\"[raw text]: {raw_text}\")\n",
" print(f\"[hyp text]: {pred_text}\")\n",
" print(f\"[ref text]: {target_transcripts[idx]}\")\n",
" else:\n",
" # if no spotted words, use standard greedy predictions\n",
" pred_text = ctc_model.wer.decoding.ctc_decoder_predictions_tensor(greedy_predicts)[0].text"
]
},
{
"cell_type": "markdown",
"id": "fb8b5f51",
"metadata": {},
"source": [
"In these logs, you can find detailed context-biasing statistics about each audio file:\n",
"- Audio file name\n",
"- Greedy word alignment\n",
"- List of spotted words\n",
"- Text results:\n",
" - Greedy decoding (raw text)\n",
" - Text after applying context-biasing (hyp text)\n",
" - Ground truth transcription (ref text)\n",
" \n",
"These statistics can be helpful in case of problems with context-biasing word recognition. For example, Transducer models sometimes recognize tokens 1-2 frames earlier than CTC models. To solve this problem, you can shift the start frame of the detected word left in the CTC-WS sources."
]
},
{
"cell_type": "markdown",
"id": "11220db2",
"metadata": {},
"source": [
"## Summary\n",
"\n",
"This tutorial demonstrates how to use the CTC-WS context-biasing technique to improve the recognition accuracy of specific words in the case of CTC and Transducer (RNN-T) ASR models. The tutorial includes the methodology for creating the context-biasing list, improving recognition accuracy of abbreviations and compound words, visualization of the context-biasing process, and results analysis.\n"
]
}
],
"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.12"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
tutorials/asr/ASR_Example_CommonVoice_Finetuning.ipynb
+432
View File
@@ -0,0 +1,432 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "52580f1b",
"metadata": {},
"source": [
"# Example: Training Esperanto ASR model using Mozilla Common Voice Dataset\n",
"\n",
"NOTE: User is responsible for checking the content of datasets and the applicable licenses and determining if suitable for the intended use.\n",
"\n",
"\n",
"Training an ASR model for a new language can be challenging, especially for low-resource languages (see [example](https://github.com/NVIDIA/NeMo/blob/main/docs/source/asr/examples/kinyarwanda_asr.rst) for Kinyarwanda CommonVoice ASR model).\n",
"\n",
"This example describes all basic steps required to build ASR model for Esperanto:\n",
"\n",
"* Data preparation\n",
"* Tokenization\n",
"* Training hyper-parameters\n",
"* Training from scratch\n",
"* Fine-tuning from pretrained models on other languages (English, Spanish, Italian).\n",
"* Fine-tuning from pretrained English SSL ([Self-supervised learning](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/asr/ssl/intro.html?highlight=self%20supervised)) model\n",
"* Model evaluation"
]
},
{
"cell_type": "markdown",
"id": "ff84f5fc",
"metadata": {},
"source": [
"## Data Preparation\n",
"\n",
"Mozilla Common Voice provides 1400 hours of validated Esperanto speech (see [here](https://arxiv.org/abs/1912.0667)). However, the final training dataset consists only of 250 hours because the train, test, and development sets are bucketed such that any given speaker may appear in only one. This ensures that contributors seen at train time are not seen at test time, which would skew results. Additionally, repetitions of text sentences are removed from the train, test, and development sets of the corpus”. \n",
"\n",
"### Downloading the Data\n",
"\n",
"You can use the NeMo script to download MCV dataset from Hugging Face and get NeMo data manifests for Esperanto:\n",
"\n",
"# Setup\n",
"After installation of huggingface datasets (`pip install datasets`), CommonVoice requires authentication.\n",
"\n",
"Website steps:\n",
"- Visit https://huggingface.co/settings/profile\n",
"- Visit \"Access Tokens\" on list of items.\n",
"- Create new token - provide a name for the token and \"read\" access is sufficient.\n",
" - PRESERVE THAT TOKEN API KEY. You can copy that key for next step.\n",
"- Visit the HuggingFace Dataset page for Mozilla Common Voice\n",
" - There should be a section that asks you for your approval.\n",
" - Make sure you are logged in and then read that agreement.\n",
" - If and only if you agree to the text, then accept the terms.\n",
"\n",
"Code steps:\n",
"- Now on your machine, run `huggingface-cli login`\n",
"- Paste your preserved HF TOKEN API KEY (from above).\n",
"\n",
"Once the above is complete, to download the data:\n",
"\n",
"```bash\n",
"python ${NEMO_ROOT}/scripts/speech_recognition/convert_hf_dataset_to_nemo.py \\\n",
" output_dir=${OUTPUT_DIR} \\\n",
" path=\"mozilla-foundation/common_voice_11_0\" \\\n",
" name=\"eo\" \\\n",
" ensure_ascii=False \\\n",
" use_auth_token=True\n",
"```\n",
"You will get the next data structure:\n",
"\n",
"```bash\n",
" .\n",
" └── mozilla-foundation\n",
" └── common_voice_11_0\n",
" └── eo\n",
" ├── test\n",
" ├── train\n",
" └── validation\n",
"```\n",
"\n",
"### Dataset Preprocessing\n",
"\n",
"Next, we must clear the text data from punctuation and various “garbage” characters. In addition to deleting a standard set of elements (as in Kinyarwanda), you can compute the frequency of characters in the train set and add the rarest (occurring less than ten times) to the list for deletion. \n",
"\n",
"\n",
"```python\n",
"import json\n",
"from tqdm import tqdm\n",
"\n",
"dev_manifest = f\"{YOUR_DATA_ROOT}/validation/validation_mozilla-foundation_common_voice_11_0_manifest.json\"\n",
"test_manifest = f\"{YOUR_DATA_ROOT}/test/test_mozilla-foundation_common_voice_11_0_manifest.json\"\n",
"train_manifest = f\"{YOUR_DATA_ROOT}/train/train_mozilla-foundation_common_voice_11_0_manifest.json\"\n",
"\n",
"def compute_char_counts(manifest):\n",
" char_counts = {}\n",
" with open(manifest, 'r') as fn_in:\n",
" for line in tqdm(fn_in, desc=\"Compute counts..\"):\n",
" line = line.replace(\"\\n\", \"\")\n",
" data = json.loads(line)\n",
" text = data[\"text\"]\n",
" for word in text.split():\n",
" for char in word:\n",
" if char not in char_counts:\n",
" char_counts[char] = 1\n",
" else:\n",
" char_counts[char] += 1\n",
" return char_counts\n",
"\n",
"char_counts = compute_char_counts(train_manifest)\n",
"\n",
"threshold = 10\n",
"trash_char_list = []\n",
"\n",
"for char in char_counts:\n",
" if char_counts[char] <= threshold:\n",
" trash_char_list.append(char)\n",
"```\n",
"\n",
"Let's check:\n",
"\n",
"```python\n",
"print(trash_char_list)\n",
"['é', 'ǔ', 'á', '¨', 'Ŭ', 'fi', '=', 'y', '`', 'q', 'ü', '♫', '‑', 'x', '¸', 'ʼ', '‹', '›', 'ñ']\n",
"```\n",
"\n",
"Next we will check the data for anomalies in audio file (for example, audio file with noise only). For this end, we check character rate (number of chars per second). For example, If the char rate is too high (more than 15 chars per second), then something is wrong with the audio file. It is better to filter such data from the training dataset in advance. Other problematic files can be filtered out after receiving the first trained model. We will consider this method at the end of our example.\n",
"\n",
"```python\n",
"import re\n",
"import json\n",
"from tqdm import tqdm\n",
"\n",
"def clear_data_set(manifest, char_rate_threshold=None):\n",
"\n",
" chars_to_ignore_regex = \"[\\.\\,\\?\\:\\-!;()«»…\\]\\[/\\*–‽+&_\\\\½√>€™$•¼}{~—=“\\\"”″‟„]\"\n",
" addition_ignore_regex = f\"[{''.join(trash_char_list)}]\"\n",
"\n",
" manifest_clean = manifest + '.clean'\n",
" war_count = 0\n",
" with open(manifest, 'r') as fn_in, \\\n",
" open(manifest_clean, 'w', encoding='utf-8') as fn_out:\n",
" for line in tqdm(fn_in, desc=\"Cleaning manifest data\"):\n",
" line = line.replace(\"\\n\", \"\")\n",
" data = json.loads(line)\n",
" text = data[\"text\"]\n",
" if char_rate_threshold and len(text.replace(' ', '')) / float(data['duration']) > char_rate_threshold:\n",
" print(f\"[WARNING]: {data['audio_filepath']} has char rate > 15 per sec: {len(text)} chars, {data['duration']} duration\")\n",
" war_count += 1\n",
" continue\n",
" text = re.sub(chars_to_ignore_regex, \"\", text)\n",
" text = re.sub(addition_ignore_regex, \"\", text)\n",
" data[\"text\"] = text.lower()\n",
" data = json.dumps(data, ensure_ascii=False)\n",
" fn_out.write(f\"{data}\\n\")\n",
" print(f\"[INFO]: {war_count} files were removed from manifest\")\n",
"\n",
"clear_data_set(dev_manifest)\n",
"clear_data_set(test_manifest)\n",
"clear_data_set(train_manifest, char_rate_threshold=15)\n",
"```\n",
"\n",
"### Creating the Tarred Training Dataset\n",
"\n",
"The tarred dataset allows storing the dataset as large *.tar files instead of small separate audio files. It may speed up the training and minimizes the load when data is moved from storage to GPU nodes.\n",
"\n",
"The NeMo toolkit provides a [script]( https://github.com/NVIDIA/NeMo/blob/main/scripts/speech_recognition/convert_to_tarred_audio_dataset.py) to get tarred dataset.\n",
"\n",
"```bash\n",
"\n",
"TRAIN_MANIFEST=${YOUR_DATA_ROOT}/train/train_mozilla-foundation_common_voice_11_0_manifest.json.clean\n",
"\n",
"python ${NEMO_ROOT}/scripts/speech_recognition/convert_to_tarred_audio_dataset.py \\\n",
" --manifest_path=${TRAIN_MANIFEST} \\\n",
" --target_dir=${YOUR_DATA_ROOT}/train_tarred_1bk \\\n",
" --num_shards=1024 \\\n",
" --max_duration=15.0 \\\n",
" --min_duration=1.0 \\\n",
" --shuffle \\\n",
" --shuffle_seed=1 \\\n",
" --sort_in_shards \\\n",
" --workers=-1\n",
"```"
]
},
{
"cell_type": "markdown",
"id": "c15bc180",
"metadata": {},
"source": [
"## Text Tokenization\n",
"\n",
"We use the standard [Byte-pair](https://en.wikipedia.org/wiki/Byte_pair_encoding) encoding algorithm with 128, 512, and 1024 vocabulary size. We found that 128 works best for relatively small Esperanto dataset (~250 hours). For larger datasets, one can get better results with larger vocabulary size (512…1024 BPE tokens).\n",
"\n",
"```\n",
"VOCAB_SIZE=128\n",
"\n",
"python ${NEMO_ROOT}/scripts/tokenizers/process_asr_text_tokenizer.py \\\n",
" --manifest=${TRAIN_MANIFEST} \\\n",
" --vocab_size=${VOCAB_SIZE} \\\n",
" --data_root=${YOUR_DATA_ROOT}/esperanto/tokenizers \\\n",
" --tokenizer=\"spe\" \\\n",
" --spe_type=bpe \n",
"```"
]
},
{
"cell_type": "markdown",
"id": "106ec1dc",
"metadata": {},
"source": [
"## Training hyper-parameters\n",
"\n",
"The training parameters are defined in the [config file](https://github.com/NVIDIA/NeMo/blob/main/examples/asr/conf/conformer/conformer_ctc_bpe.yaml) (general description of the [ASR configuration file](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/asr/configs.html)). As an encoder, the [Conformer model](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/asr/models.html#conformer-ctc) is used here, the training parameters for which are already well configured based on the training English models. However, the set of optimal parameters may differ for a new language. In this section, we will look at the set of simple parameters that can improve recognition quality for a new language without digging into the details of the Conformer model too much.\n",
"\n",
"### Select Training Batch Size\n",
"\n",
"We trained model on server with 16 V100 GPUs with 32 GB. We use a local batch size = 32 per GPU V100), so global batch size is 32x16=512. In general, we observed, that global batch between 512 and 2048 works well for Conformer-CTC-Large model. One can use the [accumulate_grad_batches](https://github.com/NVIDIA/NeMo/blob/main/examples/asr/conf/conformer/conformer_ctc_bpe.yaml#L173) parameter to increase the size of the global batch, which is equal to *local_batch * num_gpu * accumulate_grad_batches*.\n",
"\n",
"### Selecting Optimizer and Learning Rate Scheduler\n",
"\n",
"The model was trained with AdamW optimizer and CosineAnealing Learning Rate (LR) scheduler. We use Learning Rate warmup when LR goes from 0 to maximum LR to stabilize initial phase of training. The number of warmup steps determines how quickly the scheduler will reach the peak learning rate during model training. The recommended number of steps is approximately 10-20% of total training duration. We used 8,000-10,000 warmup steps.\n",
"\n",
"Now we can plot our learning rate for CosineAnnealing schedule:\n",
"\n",
"```python\n",
"import nemo\n",
"import torch\n",
"import matplotlib.pyplot as plt\n",
"\n",
"# params:\n",
"train_files_num = 144000 # number of training audio_files\n",
"global_batch_size = 1024 # local_batch * gpu_num * accum_gradient\n",
"num_epoch = 300\n",
"warmup_steps = 10000\n",
"config_learning_rate = 1e-3\n",
"\n",
"steps_num = int(train_files_num / global_batch_size * num_epoch)\n",
"print(f\"steps number is: {steps_num}\")\n",
"\n",
"optimizer = torch.optim.SGD(model.parameters(), lr=config_learning_rate)\n",
"scheduler = nemo.core.optim.lr_scheduler.CosineAnnealing(optimizer,\n",
" max_steps=steps_num,\n",
" warmup_steps=warmup_steps,\n",
" min_lr=1e-6)\n",
"lrs = []\n",
"\n",
"for i in range(steps_num):\n",
" optimizer.step()\n",
" lr = optimizer.param_groups[0][\"lr\"]\n",
" lrs.append(lr)\n",
" scheduler.step()\n",
"\n",
"plt.plot(lrs)\n",
"```\n",
"\n",
"<img src=\"./images/CosineAnnealing_scheduler.png\" alt=\"NeMo CosineAnnealing scheduler\" style=\"width: 400px;\"/>\n",
"\n",
"### Numerical Precision\n",
"\n",
"By default, it is recommended to use half-precision float (FP16 for V100 and BF16 for A100 GPU) to speed up the training process. However, training with half-precision may affect the convergence of the model, for example training loss can explode. In this case, we recommend to decrease LR or switch to float32. "
]
},
{
"cell_type": "markdown",
"id": "124aefb8",
"metadata": {},
"source": [
"## Training\n",
"\n",
"We use three main scenarios to train Espearnto ASR model:\n",
"\n",
"* Training from scratch.\n",
"* Fine-tuning from ASR models for other languages (English, Spanish, Italian).\n",
"* Fine-tuning from an English SSL ([Self-supervised learning](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/asr/ssl/intro.html?highlight=self%20supervised)) model.\n",
"\n",
"For the training of the [Conformer-CTC](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/asr/models.html#conformer-ctc) model, we use [speech_to_text_ctc_bpe.py](https://github.com/NVIDIA/NeMo/tree/stable/examples/asr/asr_ctc/speech_to_text_ctc_bpe.py) with the default config [conformer_ctc_bpe.yaml](https://github.com/NVIDIA/NeMo/tree/stable/examples/asr/conf/conformer/conformer_ctc_bpe.yaml). Here you can see the example of how to run this training:\n",
"\n",
"```bash\n",
"TOKENIZER=${YOUR_DATA_ROOT}/esperanto/tokenizers/tokenizer_spe_bpe_v128\n",
"TRAIN_MANIFEST=${YOUR_DATA_ROOT}/train_tarred_1bk/tarred_audio_manifest.json\n",
"TARRED_AUDIO_FILEPATHS=${YOUR_DATA_ROOT}/train_tarred_1bk/audio__OP_0..1023_CL_.tar # \"_OP_0..1023_CL_\" is the range for the batch of files audio_0.tar, audio_1.tar, ..., audio_1023.tar\n",
"DEV_MANIFEST=${YOUR_DATA_ROOT}/validation/validation_mozilla-foundation_common_voice_11_0_manifest.json.clean\n",
"TEST_MANIFEST=${YOUR_DATA_ROOT}/test/test_mozilla-foundation_common_voice_11_0_manifest.json.clean\n",
"\n",
"python ${NEMO_ROOT}/examples/asr/asr_ctc/speech_to_text_ctc_bpe.py \\\n",
" --config-path=../conf/conformer/ \\\n",
" --config-name=conformer_ctc_bpe \\\n",
" exp_manager.name=\"Name of our experiment\" \\\n",
" exp_manager.resume_if_exists=true \\\n",
" exp_manager.resume_ignore_no_checkpoint=true \\\n",
" exp_manager.exp_dir=results/ \\\n",
" ++model.encoder.conv_norm_type=layer_norm \\\n",
" model.tokenizer.dir=$TOKENIZER \\\n",
" model.train_ds.is_tarred=true \\\n",
" model.train_ds.tarred_audio_filepaths=$TARRED_AUDIO_FILEPATHS \\\n",
" model.train_ds.manifest_filepath=$TRAIN_MANIFEST \\\n",
" model.validation_ds.manifest_filepath=$DEV_MANIFEST \\\n",
" model.test_ds.manifest_filepath=$TEST_MANIFEST\n",
"```\n",
"\n",
"Main training parameters:\n",
"\n",
"* Tokenization: BPE 128/512/1024\n",
"* Model: Conformer-CTC-large with Layer Normalization\n",
"* Optimizer: AdamW, weight_decay 1e-3, LR 1e-3\n",
"* Scheduler: CosineAnnealing, warmup_steps 10000, min_lr 1e-6\n",
"* Batch: 32 local, 1024 global (2 grad accumulation)\n",
"* Precision: FP16\n",
"* GPUs: 16 V100\n",
"\n",
"The following table provides the results for training Esperanto Conformer-CTC-large model from scratch with different BPE vocabulary size.\n",
"\n",
"BPE Size | DEV WER (%) | TEST WER (%)\n",
"-----|-----|----- \n",
"128|**3.96**|**6.48**\n",
"512|4.62|7.31\n",
"1024|5.81|8.56\n",
"\n",
"BPE vocabulary with 128 size provides the lowest WER since our training dataset is l (~250 hours) is insufficient to small to train models with larger BPE vocabulary sizes.\n",
"\n",
"For fine-tuning from already trained ASR models, we use three different models:\n",
"\n",
"* English [stt_en_conformer_ctc_large](https://huggingface.co/nvidia/stt_en_conformer_ctc_large) (several thousand hours of English speech).\n",
"* Spanish [stt_es_conformer_ctc_large](https://huggingface.co/nvidia/stt_es_conformer_ctc_large) (1340 hours of Spanish speech).\n",
"* Italian [stt_it_conformer_ctc_large](https://huggingface.co/nvidia/stt_it_conformer_ctc_large) (487 hours of Italian speech).\n",
"\n",
"To finetune a model with the same vocabulary size, just set the desired model via the *init_from_pretrained_model* parameter:\n",
"\n",
"```yaml\n",
"+init_from_pretrained_model=${PRETRAINED_MODEL_NAME}\n",
"```\n",
"\n",
"If the size of the vocabulary differs from the one presented in the pretrained model, you need to change the vocabulary manually as done in the [finetuning tutorial](https://github.com/NVIDIA/NeMo/blob/main/tutorials/asr/ASR_CTC_Language_Finetuning.ipynb).\n",
"\n",
"```python\n",
"model = nemo_asr.models.EncDecCTCModelBPE.from_pretrained(f\"nvidia/{PRETRAINED_MODEL_NAME}\", map_location='cpu')\n",
"model.change_vocabulary(new_tokenizer_dir=TOKENIZER, new_tokenizer_type=\"bpe\")\n",
"model.encoder.unfreeze()\n",
"model.save_to(f\"{save_path}\")\n",
"```\n",
"\n",
"There is no need to change anything for the SSL model, it will replace the vocabulary itself. However, you will need to first download this model and set it through another parameter *init_from_nemo_model*:\n",
"\n",
"```yaml\n",
"++init_from_nemo_model=${PRETRAINED_MODEL} \\\n",
"```\n",
"\n",
"As the SSL model, we use [ssl_en_conformer_large](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/nemo/models/ssl_en_conformer_large) which is trained using LibriLight corpus (~56k hrs of unlabeled English speech).\n",
"All models for fine-tuning are available on [Nvidia Hugging Face](https://huggingface.co/nvidia) or [NGC](https://catalog.ngc.nvidia.com/models) repo.\n",
"\n",
"The following table shows all results for fine-tuning from pretrained models for the Conformer-CTC-large model and compares them with the model that was obtained by training from scratch (here we use BPE size 128 for all the models because it gives the best results).\n",
"\n",
"\n",
"Training Mode | DEV WER (%) | TEST WER (%)\n",
"-----|-----|----- \n",
"From scratch|3.96|6.48\n",
"Finetuning (English)|3.45|5.45\n",
"Finetuning (Spanish)|3.40|5.52\n",
"Finetuning (Italian)|3.29|5.36\n",
"Finetuning (SSL English)|**2.90**|**4.76**\n",
"\n",
"We can also monitor test WER behavior during training process using wandb plots (X - global step, Y - test WER):\n",
"\n",
"![test WER](./images/test_wer_wandb.png)\n",
"\n",
"As you can see, the best way to get the Esperanto ASR model (the model can be found on [NGC](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/nemo/models/stt_eo_conformer_ctc_large) and [Hugging Face](https://huggingface.co/nvidia/stt_eo_conformer_ctc_large) is finetuning from the pretrained SSL model for English."
]
},
{
"cell_type": "markdown",
"id": "1e3abb4d",
"metadata": {},
"source": [
"## Decoding\n",
"\n",
"At the end of the training, several checkpoints (usually 5) and the best model (not always from the latest epoch) are stored in the model folder. Checkpoint averaging (script) can help to improve the final decoding accuracy. In our case, this did not improve the CTC models. However, it was possible to get an improvement in the range of 0.1-0.2% WER for some RNNT models. To make averaging, use the following command:\n",
"\n",
"```bash\n",
"python ${NEMO_ROOT}/scripts/checkpoint_averaging/checkpoint_averaging.py <your_trained_model.nemo>\n",
"```\n",
"\n",
"For decoding you can use:\n",
"\n",
"```bash\n",
"python ${NEMO_ROOT}/examples/asr/speech_to_text_eval.py \\\n",
" model_path=${MODEL} \\\n",
" pretrained_name=null \\\n",
" dataset_manifest=${TEST_MANIFEST} \\\n",
" batch_size=${BATCH_SIZE} \\\n",
" output_filename=${OUTPUT_MANIFEST} \\\n",
" amp=False \\\n",
" use_cer=False)\n",
"```\n",
"\n",
"You can use the Speech Data Explorer to analyze recognition errors, similar to the Kinyarwanda example.\n",
"We listened to files with an anomaly high WER (>50%) and found many problematic files. They have wrong transcriptions and cut or empty audio files in the dev and test sets.\n",
"\n",
"```bash\n",
"python ${NEMO_ROOT}/tools/speech_data_explorer/data_explorer.py <your_decoded_manifest_file>\n",
"```"
]
},
{
"cell_type": "markdown",
"id": "48e37bf3",
"metadata": {},
"source": [
"## Training data analysis\n",
"\n",
"For an additional analysis of the training dataset, you can decode it using an already trained model. Train examples with a high error rate (WER > 50%) are likely to be problematic files. Removing them from the training set is preferred because a model can train text even for almost empty audio. We do not want this behavior from the ASR model."
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.9.12"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
tutorials/asr/ASR_TTS_Tutorial.ipynb
+29
View File
@@ -0,0 +1,29 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "a3570803-9bfa-4e97-9891-5ae0759eb8ca",
"metadata": {},
"source": [
"# Hybrid ASR-TTS Tutorial (Removed)\n",
"\n",
"The Hybrid ASR-TTS model (``ASRWithTTSModel``) and the Spectrogram Enhancer model (``SpectrogramEnhancerModel``) have been removed from NeMo.\n",
"\n",
"This tutorial is no longer available. For ASR fine-tuning and training, please see other ASR tutorials and documentation."
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"name": "python",
"version": "3.10.0"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
tutorials/asr/ASR_for_telephony_speech.ipynb
+344
View File
@@ -0,0 +1,344 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "lJz6FDU1lRzc"
},
"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\nNOTE: User is responsible for checking the content of datasets and the applicable licenses and determining if suitable for the intended use.\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 matplotlib>=3.3.2\n",
"\n",
"## Install NeMo\n",
"BRANCH = 'main'\n",
"!python -m pip install \"nemo_toolkit[all] @ git+https://github.com/NVIDIA-NeMo/Speech.git@$BRANCH\"\n",
"\n",
"## Grab the config we'll use in this example\n",
"!mkdir configs\n",
"!wget -P configs/ https://raw.githubusercontent.com/NVIDIA/NeMo/$BRANCH/examples/asr/conf/config.yaml\n",
"\n",
"\"\"\"\n",
"Remember to restart the runtime for the kernel to pick up any upgraded packages (e.g. matplotlib)!\n",
"Alternatively, you can uncomment the exit() below to crash and restart the kernel, in the case\n",
"that you want to use the \"Run All Cells\" (or similar) option.\n",
"\"\"\"\n",
"# exit()"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "v1Jk9etFlRzf"
},
"source": [
"# Telephony speech (8 kHz)\n",
"This notebook covers general recommendations for using NeMo models with 8 kHz speech. All the pretrained models currently available through NeMo are trained with audio at 16 kHz. This means that if the original audio was sampled at a different rate, it's sampling rate was converted to 16 kHz through upsampling or downsampling. One of the common applications for ASR is to recognize telephony speech which typically consists of speech sampled at 8 kHz.\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Mixed sample rate\n",
"Most of the pretrained English models distributed with NeMo are trained with mixed sample rate data, i.e. the training data typically consists of data sampled at both 8 kHz and 16 kHz. As an example pretrained Citrinet model \"stt_en_citrinet_1024\" was trained with the following datasets. \n",
"* Librispeech 960 hours of English speech\n",
"* Fisher Corpus\n",
"* Switchboard-1 Dataset\n",
"* WSJ-0 and WSJ-1\n",
"* National Speech Corpus - 1\n",
"* Mozilla Common Voice\n",
"\n",
"Among these, Fisher and Switchboard datasets are conversational telephone speech datasets with audio sampled at 8 kHz while the other datasets were originally sampled at least 16 kHz. Before training, all audio files from Fisher and Switchboard datasets were upsampled to 16 kHz. Because of this mixed sample rate training, our models can be used to recognize both narrowband (8kHz) and wideband speech (16kHz)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Inference with NeMo\n",
"NeMo ASR currently supports inference of audio in .wav format. Internally, the audio file is resampled to 16 kHz before inference is called on the model, so there is no difference running inference on 8 kHz audio compared to say 16 kHz or any other sampling rate audio with NeMo. Let's look at an example for running inference on 8 kHz audio. "
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"# This is where the an4/ directory will be placed.\n",
"# Change this if you don't want the data to be extracted in the current directory.\n",
"data_dir = '.'"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import glob\n",
"import os\n",
"import subprocess\n",
"import tarfile\n",
"import wget\n",
"\n",
"# Download the dataset. This will take a few moments...\n",
"print(\"******\")\n",
"if not os.path.exists(data_dir + '/an4_sphere.tar.gz'):\n",
" an4_url = 'https://dldata-public.s3.us-east-2.amazonaws.com/an4_sphere.tar.gz'\n",
" an4_path = wget.download(an4_url, data_dir)\n",
" print(f\"Dataset downloaded at: {an4_path}\")\n",
"else:\n",
" print(\"Tarfile already exists.\")\n",
" an4_path = data_dir + '/an4_sphere.tar.gz'\n",
"\n",
"if not os.path.exists(data_dir + '/an4/'):\n",
" # Untar and convert .sph to .wav (using sox)\n",
" tar = tarfile.open(an4_path)\n",
" tar.extractall(path=data_dir)\n",
"\n",
" print(\"Converting .sph to .wav...\")\n",
" sph_list = glob.glob(data_dir + '/an4/**/*.sph', recursive=True)\n",
" for sph_path in sph_list:\n",
" wav_path = sph_path[:-4] + '.wav'\n",
" cmd = [\"sox\", sph_path, wav_path]\n",
" subprocess.run(cmd)\n",
"print(\"Finished conversion.\\n******\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Audio in an4 dataset is sampled at 22 kHz. Let's first downsample an audio file to 16 kHz."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import librosa\n",
"import IPython.display as ipd\n",
"import librosa.display\n",
"import matplotlib.pyplot as plt\n",
"\n",
"# Load and listen to the audio file\n",
"example_file = data_dir + '/an4/wav/an4_clstk/mgah/cen2-mgah-b.wav'\n",
"audio, sample_rate = librosa.load(example_file)\n",
"print(sample_rate)\n",
"audio_16kHz = librosa.core.resample(audio, orig_sr=sample_rate, target_sr=16000)\n",
"\n",
"import numpy as np\n",
"\n",
"# Get spectrogram using Librosa's Short-Time Fourier Transform (stft)\n",
"spec = np.abs(librosa.stft(audio_16kHz))\n",
"spec_db = librosa.amplitude_to_db(spec, ref=np.max) # Decibels\n",
"\n",
"# Use log scale to view frequencies\n",
"librosa.display.specshow(spec_db, y_axis='log', x_axis='time', sr=16000)\n",
"plt.colorbar()\n",
"plt.title('Audio Spectrogram');\n",
"plt.ylim([0, 8000])"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now, let's downsample the audio to 8 kHz"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"audio_8kHz = librosa.core.resample(audio, orig_sr=sample_rate, target_sr=8000)\n",
"spec = np.abs(librosa.stft(audio_8kHz))\n",
"spec_db = librosa.amplitude_to_db(spec, ref=np.max) # Decibels\n",
"\n",
"# Use log scale to view frequencies\n",
"librosa.display.specshow(spec_db, y_axis='log', x_axis='time', sr=8000)\n",
"plt.colorbar()\n",
"plt.title('Audio Spectrogram');\n",
"plt.ylim([0, 8000])"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import soundfile as sf\n",
"sf.write(data_dir + '/audio_16kHz.wav', audio_16kHz, 16000)\n",
"sample, sr = librosa.core.load(data_dir + '/audio_16kHz.wav')\n",
"ipd.Audio(sample, rate=sr)\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"sf.write(data_dir + '/audio_8kHz.wav', audio_8kHz, 8000)\n",
"sample, sr = librosa.core.load(data_dir + '/audio_8kHz.wav')\n",
"ipd.Audio(sample, rate=sr)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"\n",
"Let's look at inference results using one of the pretrained models on the original, 16 kHz and 8 kHz versions of the example file we chose above."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from nemo.collections.asr.models import ASRModel\n",
"import torch\n",
"if torch.cuda.is_available():\n",
" device = torch.device(f'cuda:0')\n",
"asr_model = ASRModel.from_pretrained(model_name='stt_en_citrinet_1024', map_location=device)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"As discussed above, there are no changes required for inference based on the sampling rate of audio and as we see below the pretrained Citrinet model gives accurate transcription even for audio downsampled to 8 Khz."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"print(asr_model.transcribe(audio=[example_file]))"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"print(asr_model.transcribe(audio=[data_dir + '/audio_16kHz.wav']))"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"print(asr_model.transcribe(audio=[data_dir + '/audio_8kHz.wav']))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Training / fine-tuning with 8 kHz data\n",
"For training a model with new 8 kHz data, one could take two approaches. The first approach, **which is recommended**, is to finetune a pretrained 16 kHz model by upsampling all the data to 16 kHz. Note that upsampling offline before training is not necessary but recommended as online upsampling during training is very time consuming and may slow down training significantly. The second approach is to train an 8 kHz model from scratch. **Note**: For the second approach, in our experiments we saw that loading the weights of a 16 kHz model as initialization helps the model to converge faster with better accuracy.\n",
"\n",
"To upsample your 8 kHz data to 16 kHz command line tools like sox or ffmpeg are very useful. Here is the command to upsample and audio file using sox:\n",
"```shell\n",
"sox input_8k.wav -r 16000 -o output_16k.wav\n",
"```\n",
"Now to finetune a pre-trained model with this upsampled data, you can just restore the model weights from the pre-trained model and call trainer with the upsampled data. As an example, here is how one would fine-tune a Citrinet model:\n",
"```python\n",
"python examples/asr/script_to_bpe.py \\\n",
" --config-path=\"examples/asr/conf/citrinet\" \\\n",
" --config-name=\"citrinet_512.yaml\" \\\n",
" model.train_ds.manifest_filepath=\"<path to manifest file with upsampled 16kHz data>\" \\\n",
" model.validation_ds.manifest_filepath=\"<path to manifest file>\" \\\n",
" trainer.devices=-1 \\\n",
" trainer.accelerator='gpu' \\\n",
" trainer.max_epochs=50 \\\n",
" +init_from_pretrained_model=\"stt_en_citrinet_512\"\n",
"```\n",
"\n",
"To train an 8 kHz model, just change the sample rate in the config to 8000 as follows:\n",
"\n",
"```python\n",
"python examples/asr/script_to_bpe.py \\\n",
" --config-path=\"examples/asr/conf/citrinet\" \\\n",
" --config-name=\"citrinet_512.yaml\" \\\n",
" model.sample_rate=8000 \\\n",
" model.train_ds.manifest_filepath=\"<path to manifest file with 8kHz data>\" \\\n",
" model.validation_ds.manifest_filepath=\"<path to manifest file>\" \\\n",
" trainer.devices=-1 \\\n",
" trainer.accelerator='gpu' \\\n",
" trainer.max_epochs=50 \\\n",
" +init_from_pretrained_model=\"stt_en_citrinet_512\"\n",
"```"
]
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"collapsed_sections": [],
"name": "ASR_with_NeMo.ipynb",
"provenance": [],
"toc_visible": true
},
"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.8.5"
},
"pycharm": {
"stem_cell": {
"cell_type": "raw",
"metadata": {
"collapsed": false
},
"source": []
}
}
},
"nbformat": 4,
"nbformat_minor": 4
}
tutorials/asr/ASR_with_NeMo.ipynb
+1889
View File
File diff suppressed because one or more lines are too long
tutorials/asr/ASR_with_Subword_Tokenization.ipynb
+1331
View File
File diff suppressed because it is too large Load Diff
tutorials/asr/ASR_with_Transducers.ipynb
+1345
View File
File diff suppressed because it is too large Load Diff
tutorials/asr/Buffered_Transducer_Inference.ipynb
+1628
View File
File diff suppressed because it is too large Load Diff
tutorials/asr/Buffered_Transducer_Inference_with_LCS_Merge.ipynb
+1295
View File
File diff suppressed because it is too large Load Diff
tutorials/asr/Canary_Multitask_Speech_Model.ipynb
+1546
View File
File diff suppressed because it is too large Load Diff
tutorials/asr/Intro_to_Transducers.ipynb
+777
View File
@@ -0,0 +1,777 @@
{
"nbformat": 4,
"nbformat_minor": 0,
"metadata": {
"colab": {
"name": "Intro-to-Transducers.ipynb",
"provenance": [],
"collapsed_sections": [],
"toc_visible": true
},
"kernelspec": {
"name": "python3",
"display_name": "Python 3"
},
"language_info": {
"name": "python"
}
},
"cells": [
{
"cell_type": "code",
"metadata": {
"id": "qI25LKhSNg02"
},
"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\nNOTE: User is responsible for checking the content of datasets and the applicable licenses and determining if suitable for the intended use.\n",
"\"\"\"\n",
"# If you're using Google Colab and not running locally, run this cell.\n",
"import os\n",
"\n",
"# Install dependencies\n",
"!pip install wget\n",
"!apt-get install sox libsndfile1 ffmpeg\n",
"!pip install text-unidecode\n",
"!pip install matplotlib>=3.3.2\n",
"\n",
"## Install NeMo\n",
"BRANCH = 'main'\n",
"!python -m pip install \"nemo_toolkit[all] @ git+https://github.com/NVIDIA-NeMo/Speech.git@$BRANCH\""
],
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"metadata": {
"id": "Ac55MjAM5cls"
},
"source": [
"# In a conda environment, you would use the following command\n",
"# Update Numba to > 0.53\n",
"# conda install -c conda-forge numba\n",
"# or\n",
"# conda update -c conda-forge numba\n",
"\n",
"# For pip based environments,\n",
"# Update Numba to > 0.54\n",
"!pip install --upgrade numba>=0.54.1"
],
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {
"id": "dqbpRwpnQ-1D"
},
"source": [
"# Intro to Transducers\n",
"\n",
"By following the earlier tutorials for Automatic Speech Recognition in NeMo, one would have probably noticed that we always end up using [Connectionist Temporal Classification (CTC) loss](https://distill.pub/2017/ctc/) in order to train the model. Speech Recognition can be formulated in many different ways, and CTC is a more popular approach because it is a monotonic loss - an acoustic feature at timestep $t_1$ and $t_2$ will correspond to a target token at timestep $u_1$ and only then $u_2$. This monotonic property significantly simplifies the training of ASR models and speeds up convergence. However, it has certain drawbacks that we will discuss below.\n",
"\n",
"In general, ASR can be described as a sequence-to-sequence prediction task - the original sequence is an audio sequence (often transformed into mel spectrograms). The target sequence is a sequence of characters (or subword tokens). Attention models are capable of the same sequence-to-sequence prediction tasks. They can even perform better than CTC due to their autoregressive decoding. However, they lack certain inductive biases that can be leveraged to stabilize and speed up training (such as the monotonicity exhibited by the CTC loss). Furthermore, by design, attention models require the entire sequence to be available to align the sequence to the output, thereby preventing their use for streaming inference.\n",
"\n",
"Then comes the [Transducer Loss](https://arxiv.org/abs/1211.3711). Proposed by Alex Graves, it aimed to resolve the issues in CTC loss while resolving the transcription accuracy issues by performing autoregressive decoding. \n",
"\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "OaPS4_xSRGNv"
},
"source": [
"## Drawbacks of Connectionist Temporal Classification (CTC)\n",
"\n",
"CTC is an excellent loss to train ASR models in a stable manner but comes with certain limitations on model design. If we presume speech recognition to be a sequence-to-sequence problem, let $T$ be the sequence length of the acoustic model's output, and let $U$ be the sequence length of the target text transcript (post tokenization, either as characters or subwords). \n",
"\n",
"-------\n",
"\n",
"1) CTC imposes the limitation : $T \\ge U$. Normally, this assumption is naturally valid because $T$ is generally a lot longer than the final text transcription. However, there are many cases where this assumption fails.\n",
"\n",
"- Acoustic model performs downsampling to such a degree that $T \\ge U$. Why would we want to perform so much downsampling? For convolutions, longer sequences take more stride steps and more memory. For Attention-based models (say Conformer), there's a quadratic memory cost of computing the attention step in proportion to $T$. So more downsampling significantly helps relieve the memory requirements. There are ways to bypass this limitation, as discussed in the `ASR_with_Subword_Tokenization` notebook, but even that has limits.\n",
"\n",
"- The target sequence is generally very long. Think of languages such as German, which have very long translations for short English words. In the task of ASR, if there is more than 2x downsampling and character tokenization is used, the model will often fail to learn due to this CTC limitation.\n",
"\n",
"2) Tokens predicted by models which are trained with just CTC loss are assumed to be *conditionally independent*. This means that, unlike language models where *h*-*e*-*l*-*l* as input would probably predict *o* to complete *hello*, for CTC trained models - any character from the English alphabet has equal likelihood for prediction. So CTC trained models often have misspellings or missing tokens when transcribing the audio segment to text. \n",
"\n",
"- Since we often use the Word Error Rate (WER) metric when evaluating models, even a single misspelling contributes significantly to the \"word\" being incorrect. \n",
"\n",
"- To alleviate this issue, we have to resort to Beam Search via an external language model. While this often works and significantly improves transcription accuracy, it is a slow process and involves large N-gram or Neural language models. "
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "5EVBcBDNf658"
},
"source": [
"--------\n",
"\n",
"Let's see CTC loss's limitation (1) in action:"
]
},
{
"cell_type": "code",
"metadata": {
"id": "4whMzIjYf4w8"
},
"source": [
"import torch\n",
"import torch.nn as nn"
],
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"metadata": {
"id": "aGdKAFe7gGY4"
},
"source": [
"T = 10 # acoustic sequence length\n",
"U = 16 # target sequence length\n",
"V = 28 # vocabulary size\n",
"\n",
"def get_sample(T, U, V, require_grad=True):\n",
" torch.manual_seed(0)\n",
"\n",
" acoustic_seq = torch.randn(1, T, V + 1, requires_grad=require_grad)\n",
" acoustic_seq_len = torch.tensor([T], dtype=torch.int32) # actual seq length in padded tensor (here no padding is done)\n",
"\n",
" target_seq = torch.randint(low=0, high=V, size=(1, U))\n",
" target_seq_len = torch.tensor([U], dtype=torch.int32)\n",
"\n",
" return acoustic_seq, acoustic_seq_len, target_seq, target_seq_len"
],
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"metadata": {
"id": "DTYIb-7ngo_L"
},
"source": [
"# First, we use CTC loss in the general sense.\n",
"loss = torch.nn.CTCLoss(blank=V, zero_infinity=False)\n",
"\n",
"acoustic_seq, acoustic_seq_len, target_seq, target_seq_len = get_sample(T, U, V)\n",
"\n",
"# CTC loss expects acoustic sequence to be in shape (T, B, V)\n",
"val = loss(acoustic_seq.transpose(1, 0), target_seq, acoustic_seq_len, target_seq_len)\n",
"print(\"CTC Loss :\", val)\n",
"\n",
"val.backward()\n",
"print(\"Grad of Acoustic model (over V):\", acoustic_seq.grad[0, 0, :])"
],
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"metadata": {
"id": "lBDvC2RykFC4"
},
"source": [
"# Next, we use CTC loss with `zero_infinity` flag set.\n",
"loss = torch.nn.CTCLoss(blank=V, zero_infinity=True)\n",
"\n",
"acoustic_seq, acoustic_seq_len, target_seq, target_seq_len = get_sample(T, U, V)\n",
"\n",
"# CTC loss expects acoustic sequence to be in shape (T, B, V)\n",
"val = loss(acoustic_seq.transpose(1, 0), target_seq, acoustic_seq_len, target_seq_len)\n",
"print(\"CTC Loss :\", val)\n",
"\n",
"val.backward()\n",
"print(\"Grad of Acoustic model (over V):\", acoustic_seq.grad[0, 0, :])"
],
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {
"id": "SQe6WYnWkSAZ"
},
"source": [
"-------\n",
"\n",
"As we saw, CTC loss in general case will not be able to compute the loss or the gradient when $T \\ge U$. In the PyTorch specific implementation of CTC Loss, we can specify a flag `zero_infinity`, which explicitly checks for such cases, zeroes out the loss and the gradient if such a case occurs. The flag allows us to train a batch of samples where some samples may accidentally violate this limitation, but training will not halt, and gradients will not become NAN."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "EGnc5HqnZ-GZ"
},
"source": [
"## What is the Transducer Loss ?"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "0W12xF_CqcVF"
},
"source": [
"![](https://github.com/NVIDIA/NeMo/blob/main/tutorials/asr/images/transducer.png?raw=true)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "RoOQJtIkqxbA"
},
"source": [
"A model that seeks to use the Transducer loss is composed of three models that interact with each other. They are:\n",
"\n",
"-------\n",
"\n",
"1) **Acoustic model** : This is nearly the same acoustic model used for CTC models. The output shape of these models is generally $(Batch, \\, T, \\, AM-Hidden)$. You will note that unlike for CTC, the output of the acoustic model is no longer passed through a decoder layer which would have the shape $(Batch, \\, T, \\, Vocabulary + 1)$.\n",
"\n",
"2) **Prediction / Decoder model** : The prediction model accepts a sequence of target tokens (in the case of ASR, text tokens) and is usually a causal auto-regressive model that is tasked with predicting some hidden feature dimension of shape $(Batch, \\, U, \\, Pred-Hidden)$.\n",
"\n",
"3) **Joint model** : This model accepts the outputs of the Acoustic model and the Prediction model and joins them to compute a joint probability distribution over the vocabulary space to compute the alignments from Acoustic sequence to Target sequence. The output of this model is of the shape $(Batch, \\, T, \\, U, \\, Vocabulary + 1)$.\n",
"\n",
"--------\n",
"\n",
"During training, the transducer loss is computed on the output of the joint model, which computes the joint probability distribution of a target vocabulary token $v_{t, u}$ (for all $v \\in V$) being predicted given the acoustic feature at timestep $t \\le T$ and the prediction network features at timestep $u \\le U$.\n",
"\n",
"--------\n",
"\n",
"During inference, we perform a single forward pass over the Acoustic Network to obtain the features of shape $(Batch, \\, T, \\, AM-Hidden)$, and autoregressively perform the forward passes of the Prediction Network and the Joint Network to decode several $u \\le U$ target tokens per acoustic timestep $t \\le T$. We will discuss decoding in the following sections.\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "yBxtxt2Ztuoo"
},
"source": [
"---------\n",
"\n",
"**Note**: For an excellent in-depth explanation of how Transducer loss works, how it computes the alignment, and how the gradient of this alignment is calculated, we highly encourage you to read this post about [Sequence-to-sequence learning with Transducers by Loren Lugosch](https://lorenlugosch.github.io/posts/2020/11/transducer/).\n",
"\n",
"---------"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "VgdYFkeyRGP-"
},
"source": [
"## Benefits of Transducer Loss\n",
"\n",
"Now that we understand what a Transducer model is comprised of and how it is trained, the next question that comes to mind is - What is the benefit of the Transducer loss?\n",
"\n",
"------\n",
"\n",
"1) It is a monotonic loss (similar to CTC). Monotonicity speeds up convergence and does not require auxiliary losses to stabilize training (which is required when using only attention-based loss for sequence-to-sequence training).\n",
"\n",
"2) Autoregressive decoding enables the model to implicitly have a dependency between predicted tokens (the conditional independence assumption of CTC trained models is corrected). As such, missing characters or incorrect spellings are less frequent (but still exist since no model is perfect).\n",
"\n",
"3) It no longer has the $T \\ge U$ limitation that CTC imposed. This is because the total joint probability distribution is calculated now - mapping every acoustic timestep $t \\le T$ to one or more target timestep $u \\le U$. This means that for each timestep $t$, the model has at most $U$ tokens that it can predict, and therefore in the extreme case, it can predict a total of $T \\times U$ tokens!"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "wisUfV8aRGSY"
},
"source": [
"## Drawbacks of Transducer Loss\n",
"\n",
"All of these benefits come with certain costs. As is (almost) always the case in machine learning, there is no free lunch. \n",
"\n",
"-------\n",
"\n",
"1) During training, the Joint model is required to compute a joint matrix of shape $(Batch, \\, T, \\, U, \\, Vocabulary + 1)$. If you consider the value of these constants for a general dataset like Librispeech, $T \\sim 1600$, $U \\sim 450$ (with character encoding) and vocabulary $V \\sim 28+1$. Considering a batch size of 32, that total memory cost comes out to roughly **2.7 GB** at float precision. The model would also need another **2.7 GB** for the gradients. Of course, the model needs more memory still for the actual Acoustic model + Prediction model + their gradients. Note, however - this issue can be *partially* resolved with some simple tricks, which are discussed in the next tutorial. Also, this memory cost is no longer an issue during inference!\n",
"\n",
"2) Autoregressive decoding is slow. Much slower than CTC models, which require just a simple argmax of the output tensor. So while we do get superior transcription quality, we sacrifice decoding speed."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "RuASpPlD2con"
},
"source": [
"--------\n",
"\n",
"Let's check that RNNT loss no longer shows the limitations of CTC loss - "
]
},
{
"cell_type": "code",
"metadata": {
"id": "XXodnve02c8h"
},
"source": [
"T = 10 # acoustic sequence length\n",
"U = 16 # target sequence length\n",
"V = 28 # vocabulary size\n",
"\n",
"def get_rnnt_sample(T, U, V, require_grad=True):\n",
" torch.manual_seed(0)\n",
"\n",
" joint_tensor = torch.randn(1, T, U + 1, V + 1, requires_grad=require_grad)\n",
" acoustic_seq_len = torch.tensor([T], dtype=torch.int32) # actual seq length in padded tensor (here no padding is done)\n",
"\n",
" target_seq = torch.randint(low=0, high=V, size=(1, U))\n",
" target_seq_len = torch.tensor([U], dtype=torch.int32)\n",
"\n",
" return joint_tensor, acoustic_seq_len, target_seq, target_seq_len"
],
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"metadata": {
"id": "w-9Qx01G21oK"
},
"source": [
"import nemo.collections.asr as nemo_asr"
],
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"metadata": {
"id": "6hb7q81f21qj"
},
"source": [
"joint_tensor, acoustic_seq_len, target_seq, target_seq_len = get_rnnt_sample(T, U, V)\n",
"\n",
"# RNNT loss expects joint tensor to be in shape (B, T, U, V)\n",
"loss = nemo_asr.losses.rnnt.RNNTLoss(num_classes=V)\n",
"\n",
"# Uncomment to check out the keyword arguments required to call the RNNT loss\n",
"print(\"Transducer loss input types :\", loss.input_types)\n",
"print()\n",
"\n",
"val = loss(log_probs=joint_tensor, targets=target_seq, input_lengths=acoustic_seq_len, target_lengths=target_seq_len)\n",
"print(\"Transducer Loss :\", val)\n",
"\n",
"val.backward()\n",
"print(\"Grad of Acoustic model (over V):\", joint_tensor.grad[0, 0, 0, :])"
],
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {
"id": "5pfrpy7wRGUc"
},
"source": [
"# Configure a Transducer Model\n",
"\n",
"We now understand a bit more about the transducer loss. Next, we will take a deep dive into how to set up the config for a transducer model.\n",
"\n",
"Transducer configs contain a fair bit more detail compared to CTC configs. However, the vast majority of the defaults can be copied and pasted into your configs to have a perfectly functioning transducer model!\n",
"\n",
"------\n",
"\n",
"Let us download one of the transducer configs already available in NeMo to analyze the components."
]
},
{
"cell_type": "code",
"metadata": {
"id": "cgJQXfwy7LO_"
},
"source": [
"import os\n",
"\n",
"if not os.path.exists(\"contextnet_rnnt.yaml\"):\n",
" !wget https://raw.githubusercontent.com/NVIDIA/NeMo/$BRANCH/examples/asr/conf/contextnet_rnnt/contextnet_rnnt.yaml"
],
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"metadata": {
"id": "CJ2-ORS17XbF"
},
"source": [
"from omegaconf import OmegaConf, open_dict\n",
"\n",
"cfg = OmegaConf.load('contextnet_rnnt.yaml')"
],
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {
"id": "h5TsAJQk6o4N"
},
"source": [
"## Model Defaults\n",
"\n",
"Since the transducer model is comprised of three separate models working in unison, it is practical to have some shared section of the config. That shared section is called `model.model_defaults`."
]
},
{
"cell_type": "code",
"metadata": {
"id": "N8tWZ9eb75Gx"
},
"source": [
"print(OmegaConf.to_yaml(cfg.model.model_defaults))"
],
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {
"id": "m8IxgJFj7_gc"
},
"source": [
"-------\n",
"\n",
"Of the many components shared here, the last three values are the primary components that a transducer model **must** possess. They are :\n",
"\n",
"1) `enc_hidden`: The hidden dimension of the final layer of the Encoder network.\n",
"\n",
"2) `pred_hidden`: The hidden dimension of the final layer of the Prediction network.\n",
"\n",
"3) `joint_hidden`: The hidden dimension of the intermediate layer of the Joint network.\n",
"\n",
"--------\n",
"\n",
"One can access these values inside the config by using OmegaConf interpolation as follows :\n",
"\n",
"```yaml\n",
"model:\n",
" ...\n",
" decoder:\n",
" ...\n",
" prednet:\n",
" pred_hidden: ${model.model_defaults.pred_hidden}\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "uRckIz_eRGWr"
},
"source": [
"## Acoustic Model\n",
"\n",
"As we discussed before, the transducer model is comprised of three models combined. One of these models is the Acoustic (encoder) model. We should be able to drop in any CTC Acoustic model config into this section of the transducer config.\n",
"\n",
"The only condition that needs to be met is that **the final layer of the acoustic model must have the dimension defined in `model_defaults.enc_hidden`**."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "o505IGX4RGYy"
},
"source": [
"## Decoder / Prediction Model\n",
"\n",
"The Prediction model is generally an autoregressive, causal model that consumes text tokens and returns embeddings that will be used by the Joint model. \n",
"\n",
"**This config can be dropped into any custom transducer model with no modification.**"
]
},
{
"cell_type": "code",
"metadata": {
"id": "A2a9Y5CCArLs"
},
"source": [
"print(OmegaConf.to_yaml(cfg.model.decoder))"
],
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {
"id": "ry5a_Z-zAvll"
},
"source": [
"------\n",
"\n",
"This config will build an LSTM based Transducer Decoder model. Let us discuss some of the important arguments:\n",
"\n",
"1) `blank_as_pad`: In ordinary transducer models, the embedding matrix does not acknowledge the `Transducer Blank` token (similar to CTC Blank). However, this causes the autoregressive loop to be more complicated and less efficient. Instead, this flag which is set by default, will add the `Transducer Blank` token to the embedding matrix - and use it as a pad value (zeros tensor). This enables more efficient inference without harming training.\n",
"\n",
"2) `prednet.pred_hidden`: The hidden dimension of the LSTM and the output dimension of the Prediction network.\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "FtdYE25cW1j_"
},
"source": [
"## Joint Model\n",
"\n",
"The Joint model is a simple feed-forward Multi-Layer Perceptron network. This MLP accepts the output of the Acoustic and Prediction models and computes a joint probability distribution over the entire vocabulary space.\n",
"\n",
"**This config can be dropped into any custom transducer model with no modification.**"
]
},
{
"cell_type": "code",
"metadata": {
"id": "pP8fL1bED3Dv"
},
"source": [
"print(OmegaConf.to_yaml(cfg.model.joint))"
],
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {
"id": "rA11eez_FYUl"
},
"source": [
"------\n",
"\n",
"The Joint model config has several essential components which we discuss below :\n",
"\n",
"1) `log_softmax`: Due to the cost of computing softmax on such large tensors, the Numba CUDA implementation of RNNT loss will implicitly compute the log softmax when called (so its inputs should be logits). The CPU version of the loss doesn't face such memory issues so it requires log-probabilities instead. Since the behaviour is different for CPU-GPU, the `None` value will automatically switch behaviour dependent on whether the input tensor is on a CPU or GPU device.\n",
"\n",
"2) `preserve_memory`: This flag will call `torch.cuda.empty_cache()` at certain critical sections when computing the Joint tensor. While this operation might allow us to preserve some memory, the empty_cache() operation is tremendously slow and will slow down training by an order of magnitude or more. It is available to use but not recommended.\n",
"\n",
"3) `fuse_loss_wer`: This flag performs \"batch splitting\" and then \"fused loss + metric\" calculation. It will be discussed in detail in the next tutorial that will train a Transducer model.\n",
"\n",
"4) `fused_batch_size`: When the above flag is set to True, the model will have two distinct \"batch sizes\". The batch size provided in the three data loader configs (`model.*_ds.batch_size`) will now be the `Acoustic model` batch size, whereas the `fused_batch_size` will be the batch size of the `Prediction model`, the `Joint model`, the `transducer loss` module and the `decoding` module.\n",
"\n",
"5) `jointnet.joint_hidden`: The hidden intermediate dimension of the joint network."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "cmIwDscCW1mP"
},
"source": [
"## Transducer Decoding\n",
"\n",
"Models which have been trained with CTC can transcribe text simply by performing a regular argmax over the output of their decoder.\n",
"\n",
"For transducer-based models, the three networks must operate in a synchronized manner in order to transcribe the acoustic features.\n",
"\n",
"The following section of the config describes how to change the decoding logic of the transducer model.\n",
"\n",
"\n",
"**This config can be dropped into any custom transducer model with no modification.**"
]
},
{
"cell_type": "code",
"metadata": {
"id": "LQjfXJsrIqFJ"
},
"source": [
"print(OmegaConf.to_yaml(cfg.model.decoding))"
],
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {
"id": "6FXtn41wIvu8"
},
"source": [
"-------\n",
"\n",
"The most important component at the top level is the `strategy`. It can take one of many values:\n",
"\n",
"1) `greedy`: This is sample-level greedy decoding. It is generally exceptionally slow as each sample in the batch will be decoded independently. For publications, this should be used alongside batch size of 1 for exact results.\n",
"\n",
"2) `greedy_batch`: This is the general default and should nearly match the `greedy` decoding scores (if the acoustic features are not affected by feature mixing in batch mode). Even for small batch sizes, this strategy is significantly faster than `greedy`.\n",
"\n",
"3) `beam`: Runs beam search with the implicit language model of the Prediction model. It will generally be quite slow, and might need some tuning of the beam size to get better transcriptions.\n",
"\n",
"4) `tsd`: Time synchronous decoding. Please refer to the paper: [Alignment-Length Synchronous Decoding for RNN Transducer](https://ieeexplore.ieee.org/document/9053040) for details on the algorithm implemented. Time synchronous decoding (TSD) execution time grows by the factor T * max_symmetric_expansions. For longer sequences, T is greater and can therefore take a long time for beams to obtain good results. TSD also requires more memory to execute.\n",
"\n",
"5) `alsd`: Alignment-length synchronous decoding. Please refer to the paper: [Alignment-Length Synchronous Decoding for RNN Transducer](https://ieeexplore.ieee.org/document/9053040) for details on the algorithm implemented. Alignment-length synchronous decoding (ALSD) execution time is faster than TSD, with a growth factor of T + U_max, where U_max is the maximum target length expected during execution. Generally, T + U_max < T * max_symmetric_expansions. However, ALSD beams are non-unique. Therefore it is required to use larger beam sizes to achieve the same (or close to the same) decoding accuracy as TSD. For a given decoding accuracy, it is possible to attain faster decoding via ALSD than TSD.\n",
"\n",
"-------\n",
"\n",
"Below, we discuss the various decoding strategies."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "PXzY7laMW1oo"
},
"source": [
"### Greedy Decoding\n",
"\n",
"When `strategy` is one of `greedy` or `greedy_batch`, an additional subconfig of `decoding.greedy` can be used to set an important decoding value."
]
},
{
"cell_type": "code",
"metadata": {
"id": "778R5oy6Ipha"
},
"source": [
"print(OmegaConf.to_yaml(cfg.model.decoding.greedy))"
],
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {
"id": "vItXzTbZKwyB"
},
"source": [
"-------\n",
"\n",
"This argument `max_symbols` is the maximum number of `target token` decoding steps $u \\le U$ per acoustic timestep $t \\le T$. Note that during training, this was implicitly constrained by the shape of the joint matrix (max_symbols = $U$). However, there is no such $U$ upper bound during inference (we don't have the ground truth $U$).\n",
"\n",
"So we explicitly set a heuristic upper bound on how many decoding steps can be performed per acoustic timestep. Generally a value of 5 and above is sufficient."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "ebFogfLvW1q9"
},
"source": [
"### Beam Decoding\n",
"\n",
"Next, we discuss the subconfig when `strategy` is one of `beam`, `tsd` or `alsd`."
]
},
{
"cell_type": "code",
"metadata": {
"id": "w073zT8ILtki"
},
"source": [
"print(OmegaConf.to_yaml(cfg.model.decoding.beam))"
],
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {
"id": "AvOeRhsULtrx"
},
"source": [
"------\n",
"\n",
"There are several important arguments in this section :\n",
"\n",
"1) `beam_size`: This determines the beam size for all types of beam decoding strategy. Since this is implemented in PyTorch, large beam sizes will take exorbitant amounts of time.\n",
"\n",
"2) `score_norm`: Whether to normalize scores prior to pruning the beam.\n",
"\n",
"3) `return_best_hypothesis`: If beam search is being performed, we can choose to return just the best hypothesis or all the hypotheses.\n",
"\n",
"4) `tsd_max_sym_exp`: The maximum symmetric expansions allowed per timestep during beam search. Larger values should be used to attempt decoding of longer sequences, but this in turn increases execution time and memory usage.\n",
"\n",
"5) `alsd_max_target_len`: The maximum expected target sequence length during beam search. Larger values allow decoding of longer sequences at the expense of execution time and memory.\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "MkoHp0dQW1tP"
},
"source": [
"## Transducer Loss\n",
"\n",
"Finally, we reach the Transducer loss config itself. This section configures the type of Transducer loss itself, along with possible sub-sections.\n",
"\n",
"**This config can be dropped into any custom transducer model with no modification.**"
]
},
{
"cell_type": "code",
"metadata": {
"id": "l3Uk11uHOa4O"
},
"source": [
"print(OmegaConf.to_yaml(cfg.model.loss))"
],
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {
"id": "1z_mYV9UOk_7"
},
"source": [
"---------\n",
"\n",
"The loss config is based on a resolver pattern and can be used as follows:\n",
"\n",
"1) `loss_name`: `default` is generally a good option. Will select one of the available resolved losses and match the kwargs from a sub-configs passed via explicit `{loss_name}_kwargs` sub-config. \n",
"\n",
"2) `{loss_name}_kwargs`: This sub-config is passed to the resolved loss above and can be used to configure the resolved loss."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "7w3Z3-IaRGaz"
},
"source": [
"### WarpRNNT Numba Loss\n",
"\n",
"The default transducer loss implemented in NeMo is a Numba port of the excellent CUDA implementation of Transducer Loss found in https://github.com/HawkAaron/warp-transducer.\n",
"\n",
"It should suffice for most use cases (CPU / GPU) transducer training."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "u9bPDu2-XYPB"
},
"source": [
"### FastEmit Regularization\n",
"\n",
"Recently proposed regularization approach - [FastEmit: Low-latency Streaming ASR with Sequence-level Emission Regularization](https://arxiv.org/abs/2010.11148) allows us near-direct control over the latency of transducer models.\n",
"\n",
"Refer to the above paper for results and recommendations of `fastemit_lambda`."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "iGG9UKTZXlme"
},
"source": [
"# Next Steps\n",
"\n",
"After that deep dive into how to configure Transducer models, the next tutorial will use one such config to build a transducer model and train it on a small dataset. We will then move on to exploring various decoding strategies and how to evaluate the model."
]
}
]
}
tutorials/asr/Multilang_ASR.ipynb
+5203
View File
File diff suppressed because it is too large Load Diff
tutorials/asr/Offline_ASR_with_VAD_for_CTC_models.ipynb
+417
View File
@@ -0,0 +1,417 @@
{
"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",
"\n",
"NOTE: User is responsible for checking the content of datasets and the applicable licenses and determining if suitable for the intended use.\n",
"\"\"\"\n",
"# If you're using Google Colab and not running locally, run this cell.\n",
"\n",
"## Install dependencies\n",
"!pip install wget\n",
"\n",
"## Install NeMo\n",
"BRANCH = 'main'\n",
"!python -m pip install \"nemo_toolkit[all] @ git+https://github.com/NVIDIA-NeMo/Speech.git@$BRANCH\"\n",
"\n",
"\"\"\"\n",
"Remember to restart the runtime for the kernel to pick up any upgraded packages (e.g. matplotlib)!\n",
"Alternatively, you can uncomment the exit() below to crash and restart the kernel, in the case\n",
"that you want to use the \"Run All Cells\" (or similar) option.\n",
"\"\"\"\n",
"# exit()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import json\n",
"import torch\n",
"import os\n",
"from nemo.collections.asr.metrics.wer import word_error_rate\n",
"from nemo.collections.asr.parts.utils.vad_utils import stitch_segmented_asr_output, construct_manifest_eval"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"# Offline ASR+VAD"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"In this tutorial, we will demonstrate how to use offline VAD to extract speech segments and transcribe the speech segments with CTC models. This will help to exclude some non_speech utterances and could save computation resources by removing unnecessary input to the ASR system. \n",
"\n",
"The pipeline includes the following steps.\n",
"\n",
"0. [Prepare data and script for demonstration](#Prepare-data-and-script-for-demonstration)\n",
"1. [Use offline VAD to extract speech segments](#Use-offline-VAD-to-extract-speech-segments)\n",
"2. [Transcribe speech segments with CTC models](#Transcribe-speech-segments-with-CTC-models)\n",
"3. [Stitch the prediction text of speech segments](#Stitch-the-prediction-text-of-speech-segments)\n",
"4. [Evaluate the performance of offline ASR with VAD ](#Evaluate-the-performance-of-offline-VAD-with-ASR)"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"## Prepare data and script for demonstration\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"scrolled": true
},
"outputs": [],
"source": [
"!mkdir -p data\n",
"!wget -P data/ https://nemo-public.s3.us-east-2.amazonaws.com/chris-sample01_02.wav\n",
"!wget -P data/ https://nemo-public.s3.us-east-2.amazonaws.com/chris-sample03.wav\n",
"!wget https://nemo-public.s3.us-east-2.amazonaws.com/chris_demo.json"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"input_manifest=\"chris_demo.json\"\n",
"vad_out_manifest_filepath=\"vad_out.json\"\n",
"vad_model=\"vad_multilingual_marblenet\" # here we use vad_multilingual_marblenet for example, you can choose other VAD models."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!head -n 10 $input_manifest"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# This cell is mainly for colab. \n",
"# You can ignore it if run locally but do make sure change the filepaths of scripts and config file in cells below.\n",
"!mkdir -p scripts\n",
"if not os.path.exists(\"scripts/vad_infer.py\"):\n",
" !wget -P scripts/ https://raw.githubusercontent.com/NVIDIA/NeMo/$BRANCH/examples/asr/speech_classification/vad_infer.py\n",
"if not os.path.exists(\"scripts/transcribe_speech.py\"):\n",
" !wget -P scripts/ https://raw.githubusercontent.com/NVIDIA/NeMo/$BRANCH/examples/asr/transcribe_speech.py\n",
" \n",
"!mkdir -p conf/vad\n",
"if not os.path.exists(\"conf/vad/vad_inference_postprocessing.yaml\"):\n",
" !wget -P conf/vad/ https://raw.githubusercontent.com/NVIDIA/NeMo/$BRANCH/examples/asr/conf/vad/vad_inference_postprocessing.yaml"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"## Use offline VAD to extract speech segments"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"Here we are using very simple parameters to demonstrate the process. \n",
"\n",
"Please choose or tune your own postprocessing parameters. \n",
"\n",
"You can find more details in \n",
"```python \n",
"<NeMo_git_root>/tutorials/asr/Online_Offline_Microphone_VAD_Demo.ipynb and \n",
"<NeMo_git_root>/scripts/voice_activity_detection/vad_tune_threshold.py\n",
"```"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"The <code>vad_infer.py</code> script will help you generate speech segments. See more details in the script below."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"scrolled": true
},
"outputs": [],
"source": [
"# if run locally, vad_infer.py is located in <NeMo_git_root>/examples/asr/speech_classification/vad_infer.py\n",
"%run -i scripts/vad_infer.py --config-path=\"../conf/vad\" --config-name=\"vad_inference_postprocessing.yaml\" \\\n",
"dataset=$input_manifest \\\n",
"vad.model_path=$vad_model \\\n",
"frame_out_dir=\"chris_demo\" \\\n",
"vad.parameters.window_length_in_sec=0.63 \\\n",
"vad.parameters.postprocessing.onset=0.7 \\\n",
"vad.parameters.postprocessing.offset=0.4 \\\n",
"vad.parameters.postprocessing.min_duration_on=1 \\\n",
"vad.parameters.postprocessing.min_duration_off=0.5 \\\n",
"out_manifest_filepath=$vad_out_manifest_filepath"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"Let's have a look at VAD output. If there are no speech segments in the sample. The sample will not appear in VAD output."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!head -n 10 $vad_out_manifest_filepath"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"## Transcribe speech segments with CTC models"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"segmented_output_manifest=\"asr_segmented_output_manifest.json\"\n",
"asr_model=\"stt_en_citrinet_1024_gamma_0_25\" # here we use citrinet for example, you can choose other CTC models."
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"The <code>transcribe_speech.py</code> script will help you transcribe each speech segment. See more details in the script below."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"scrolled": true
},
"outputs": [],
"source": [
"# if run locally, transcribe_speech.py is located in <NeMo_git_root>/examples/asr/transcribe_speech.py\n",
"%run -i scripts/transcribe_speech.py \\\n",
" pretrained_name=$asr_model \\\n",
" dataset_manifest=$vad_out_manifest_filepath \\\n",
" batch_size=32 \\\n",
" amp=True \\\n",
" output_filename=$segmented_output_manifest"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"Let's have a look at the segmented ASR transcript."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!head -n 5 $segmented_output_manifest"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"## Stitch the prediction text of speech segments"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"You can also evaluate the whole ASR output by stitching the segmented outputs together.\n",
"\n",
"Note, there would be a better method to stitch them together. Here, we just demonstrate the simplest method, concatenating."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"stitched_output_manifest=\"stitched_asr_output_manifest.json\""
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"stitched_output_manifest = stitch_segmented_asr_output(segmented_output_manifest)"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"Let's have a look at the stitched output and the stored speech segments of the first sample."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"scrolled": true
},
"outputs": [],
"source": [
"stitched_output = []\n",
"for line in open(stitched_output_manifest, 'r', encoding='utf-8'):\n",
" file = json.loads(line)\n",
" stitched_output.append(file)\n",
"\n",
"print(stitched_output[0])\n",
"print(f\"\\n The speech segments of above file are \\n {torch.load(stitched_output[0]['speech_segments_filepath'])}\")"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"# Evaluate the performance of offline VAD with ASR "
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"If we have ground-truth <code>'text'</code> in input_manifest, we can evaluate our performance of stitched output. Let's align the <code>'text'</code> in input manifest and <code>'pred_text'</code> in stitched segmented asr output first, since some samples from input_manifest might be pure noise and have been removed in VAD output and excluded for ASR inference. "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"aligned_vad_asr_output_manifest = construct_manifest_eval(input_manifest, stitched_output_manifest)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!head -n 10 $aligned_vad_asr_output_manifest"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"predicted_text, ground_truth_text = [], []\n",
"for line in open(aligned_vad_asr_output_manifest, 'r', encoding='utf-8'):\n",
" sample = json.loads(line)\n",
" predicted_text.append(sample['pred_text'])\n",
" ground_truth_text.append(sample['text'])"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"metric_value = word_error_rate(hypotheses=predicted_text, references=ground_truth_text, use_cer=False)\n",
"print(f\"WER is {metric_value}\")"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"# Further Reading\n",
"\n",
"There are two ways to incorporate VAD into ASR pipeline. The first strategy is to drop the frames that are predicted as `non-speech` by VAD, as already discussed in this tutorial. The second strategy is to keep all the frames and mask the `non-speech` frames with zero-signal values. Also, instead of using segment-VAD as shown in this tutorial, we can use frame-VAD model for faster inference and better accuracy. For more information, please refer to the script [speech_to_text_with_vad.py](https://github.com/NVIDIA/NeMo/blob/stable/examples/asr/asr_vad/speech_to_text_with_vad.py)."
]
}
],
"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.7.11"
}
},
"nbformat": 4,
"nbformat_minor": 4
}
tutorials/asr/Online_ASR_Microphone_Demo_Buffered_Streaming.ipynb
+522
View File
@@ -0,0 +1,522 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"\"\"\"\n",
"Please run notebook locally (if you have all the dependencies and a GPU). \n",
"Technically you can run this notebook on Google Colab but you need to set up microphone for 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. Set up microphone for Colab\n",
"\n\nNOTE: User is responsible for checking the content of datasets and the applicable licenses and determining if suitable for the intended use.\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 portaudio19-dev\n",
"!pip install text-unidecode\n",
"!pip install pyaudio\n",
"\n",
"# ## Install NeMo\n",
"BRANCH = 'main'\n",
"!python -m pip install \"nemo_toolkit[asr] @ git+https://github.com/NVIDIA-NeMo/Speech.git@$BRANCH\"\n",
"\n",
"## Grab the config we'll use in this example\n",
"!mkdir configs"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"This notebook demonstrates automatic speech recognition (ASR) from a microphone's stream in NeMo.\n",
"\n",
"It is **not a recommended** way to do inference in production workflows. And the incompatibility of components could lead to failure of running this notebook locally with container, we might deprecate this notebook and provide a better tutorial in soon releases. If you are interested in production-level inference using NeMo ASR models, please refer to NVIDIA RIVA: https://developer.nvidia.com/riva"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The notebook requires PyAudio library to get a signal from an audio device.\n",
"For Ubuntu, please run the following commands to install it:\n",
"```\n",
"sudo apt install python3-pyaudio\n",
"pip install pyaudio\n",
"```"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import numpy as np\n",
"import pyaudio as pa\n",
"import os, time"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import nemo\n",
"import nemo.collections.asr as nemo_asr"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# sample rate, Hz\n",
"SAMPLE_RATE = 16000"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Restore the model from NGC"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": "asr_model = nemo_asr.models.EncDecCTCModelBPE.from_pretrained('stt_en_fastconformer_ctc_large')"
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Observing the config of the model"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from omegaconf import OmegaConf\n",
"import copy"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"scrolled": true
},
"outputs": [],
"source": [
"# Preserve a copy of the full config\n",
"cfg = copy.deepcopy(asr_model._cfg)\n",
"print(OmegaConf.to_yaml(cfg))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Modify preprocessor parameters for inference"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Make config overwrite-able\n",
"OmegaConf.set_struct(cfg.preprocessor, False)\n",
"\n",
"# some changes for streaming scenario\n",
"cfg.preprocessor.dither = 0.0\n",
"cfg.preprocessor.pad_to = 0\n",
"\n",
"# spectrogram normalization constants\n",
"normalization = {}\n",
"normalization['fixed_mean'] = [\n",
" -14.95827016, -12.71798736, -11.76067913, -10.83311182,\n",
" -10.6746914, -10.15163465, -10.05378331, -9.53918999,\n",
" -9.41858904, -9.23382904, -9.46470918, -9.56037,\n",
" -9.57434245, -9.47498732, -9.7635205, -10.08113074,\n",
" -10.05454561, -9.81112681, -9.68673603, -9.83652977,\n",
" -9.90046248, -9.85404766, -9.92560366, -9.95440354,\n",
" -10.17162966, -9.90102482, -9.47471025, -9.54416855,\n",
" -10.07109475, -9.98249912, -9.74359465, -9.55632283,\n",
" -9.23399915, -9.36487649, -9.81791084, -9.56799225,\n",
" -9.70630899, -9.85148006, -9.8594418, -10.01378735,\n",
" -9.98505315, -9.62016094, -10.342285, -10.41070709,\n",
" -10.10687659, -10.14536695, -10.30828702, -10.23542833,\n",
" -10.88546868, -11.31723646, -11.46087382, -11.54877829,\n",
" -11.62400934, -11.92190509, -12.14063815, -11.65130117,\n",
" -11.58308531, -12.22214663, -12.42927197, -12.58039805,\n",
" -13.10098969, -13.14345864, -13.31835645, -14.47345634]\n",
"normalization['fixed_std'] = [\n",
" 3.81402054, 4.12647781, 4.05007065, 3.87790987,\n",
" 3.74721178, 3.68377423, 3.69344, 3.54001005,\n",
" 3.59530412, 3.63752368, 3.62826417, 3.56488469,\n",
" 3.53740577, 3.68313898, 3.67138151, 3.55707266,\n",
" 3.54919572, 3.55721289, 3.56723346, 3.46029304,\n",
" 3.44119672, 3.49030548, 3.39328435, 3.28244406,\n",
" 3.28001423, 3.26744937, 3.46692348, 3.35378948,\n",
" 2.96330901, 2.97663111, 3.04575148, 2.89717604,\n",
" 2.95659301, 2.90181116, 2.7111687, 2.93041291,\n",
" 2.86647897, 2.73473181, 2.71495654, 2.75543763,\n",
" 2.79174615, 2.96076456, 2.57376336, 2.68789782,\n",
" 2.90930817, 2.90412004, 2.76187531, 2.89905006,\n",
" 2.65896173, 2.81032176, 2.87769857, 2.84665271,\n",
" 2.80863137, 2.80707634, 2.83752184, 3.01914511,\n",
" 2.92046439, 2.78461139, 2.90034605, 2.94599508,\n",
" 2.99099718, 3.0167554, 3.04649716, 2.94116777]\n",
"\n",
"cfg.preprocessor.normalize = normalization\n",
"\n",
"# Disable config overwriting\n",
"OmegaConf.set_struct(cfg.preprocessor, True)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Setup preprocessor with these settings"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"asr_model.preprocessor = asr_model.from_config_dict(cfg.preprocessor)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Set model to inference mode\n",
"asr_model.eval();"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"asr_model = asr_model.to(asr_model.device)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Setting up data for Streaming Inference"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from nemo.core.classes import IterableDataset\n",
"from nemo.core.neural_types import NeuralType, AudioSignal, LengthsType\n",
"import torch\n",
"from torch.utils.data import DataLoader"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# simple data layer to pass audio signal\n",
"class AudioDataLayer(IterableDataset):\n",
" @property\n",
" def output_types(self):\n",
" return {\n",
" 'audio_signal': NeuralType(('B', 'T'), AudioSignal(freq=self._sample_rate)),\n",
" 'a_sig_length': NeuralType(tuple('B'), LengthsType()),\n",
" }\n",
"\n",
" def __init__(self, sample_rate):\n",
" super().__init__()\n",
" self._sample_rate = sample_rate\n",
" self.output = True\n",
" \n",
" def __iter__(self):\n",
" return self\n",
" \n",
" def __next__(self):\n",
" if not self.output:\n",
" raise StopIteration\n",
" self.output = False\n",
" return torch.as_tensor(self.signal, dtype=torch.float32), \\\n",
" torch.as_tensor(self.signal_shape, dtype=torch.int64)\n",
" \n",
" def set_signal(self, signal):\n",
" self.signal = signal.astype(np.float32)/32768.\n",
" self.signal_shape = self.signal.size\n",
" self.output = True\n",
"\n",
" def __len__(self):\n",
" return 1"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"data_layer = AudioDataLayer(sample_rate=cfg.preprocessor.sample_rate)\n",
"data_loader = DataLoader(data_layer, batch_size=1, collate_fn=data_layer.collate_fn)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# inference method for audio signal (single instance)\n",
"def infer_signal(model, signal):\n",
" data_layer.set_signal(signal)\n",
" batch = next(iter(data_loader))\n",
" audio_signal, audio_signal_len = batch\n",
" audio_signal, audio_signal_len = audio_signal.to(asr_model.device), audio_signal_len.to(asr_model.device)\n",
" log_probs, encoded_len, predictions = model.forward(\n",
" input_signal=audio_signal, input_signal_length=audio_signal_len\n",
" )\n",
" return log_probs"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# class for streaming frame-based ASR\n",
"# 1) use reset() method to reset FrameASR's state\n",
"# 2) call transcribe(frame) to do ASR on\n",
"# contiguous signal's frames\n",
"class FrameASR:\n",
" \n",
" def __init__(self, model_definition,\n",
" frame_len=2, frame_overlap=2.5, \n",
" offset=10):\n",
" '''\n",
" Args:\n",
" frame_len: frame's duration, seconds\n",
" frame_overlap: duration of overlaps before and after current frame, seconds\n",
" offset: number of symbols to drop for smooth streaming\n",
" '''\n",
" self.vocab = list(model_definition['labels'])\n",
" self.vocab.append('_')\n",
" \n",
" self.sr = model_definition['sample_rate']\n",
" self.frame_len = frame_len\n",
" self.n_frame_len = int(frame_len * self.sr)\n",
" self.frame_overlap = frame_overlap\n",
" self.n_frame_overlap = int(frame_overlap * self.sr)\n",
" timestep_duration = model_definition['AudioToMelSpectrogramPreprocessor']['window_stride']\n",
" for block in model_definition['JasperEncoder']['jasper']:\n",
" timestep_duration *= block['stride'][0] ** block['repeat']\n",
" self.n_timesteps_overlap = int(frame_overlap / timestep_duration) - 2\n",
" self.buffer = np.zeros(shape=2*self.n_frame_overlap + self.n_frame_len,\n",
" dtype=np.float32)\n",
" self.offset = offset\n",
" self.reset()\n",
" \n",
" def _decode(self, frame, offset=0):\n",
" assert len(frame)==self.n_frame_len\n",
" self.buffer[:-self.n_frame_len] = self.buffer[self.n_frame_len:]\n",
" self.buffer[-self.n_frame_len:] = frame\n",
" logits = infer_signal(asr_model, self.buffer).cpu().numpy()[0]\n",
" # print(logits.shape)\n",
" decoded = self._greedy_decoder(\n",
" logits[self.n_timesteps_overlap:-self.n_timesteps_overlap], \n",
" self.vocab\n",
" )\n",
" return decoded[:len(decoded)-offset]\n",
" \n",
" @torch.no_grad()\n",
" def transcribe(self, frame=None, merge=True):\n",
" if frame is None:\n",
" frame = np.zeros(shape=self.n_frame_len, dtype=np.float32)\n",
" if len(frame) < self.n_frame_len:\n",
" frame = np.pad(frame, [0, self.n_frame_len - len(frame)], 'constant')\n",
" unmerged = self._decode(frame, self.offset)\n",
" if not merge:\n",
" return unmerged\n",
" return self.greedy_merge(unmerged)\n",
" \n",
" def reset(self):\n",
" '''\n",
" Reset frame_history and decoder's state\n",
" '''\n",
" self.buffer=np.zeros(shape=self.buffer.shape, dtype=np.float32)\n",
" self.prev_char = ''\n",
"\n",
" @staticmethod\n",
" def _greedy_decoder(logits, vocab):\n",
" s = ''\n",
" for i in range(logits.shape[0]):\n",
" s += vocab[np.argmax(logits[i])]\n",
" return s\n",
"\n",
" def greedy_merge(self, s):\n",
" s_merged = ''\n",
" \n",
" for i in range(len(s)):\n",
" if s[i] != self.prev_char:\n",
" self.prev_char = s[i]\n",
" if self.prev_char != '_':\n",
" s_merged += self.prev_char\n",
" return s_merged"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Streaming Inference\n",
"\n",
"Streaming inference depends on a few factors, such as the frame length and buffer size. Experiment with a few values to see their effects in the below cells."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# duration of signal frame, seconds\n",
"FRAME_LEN = 1.0\n",
"# number of audio channels (expect mono signal)\n",
"CHANNELS = 1\n",
"\n",
"CHUNK_SIZE = int(FRAME_LEN*SAMPLE_RATE)\n",
"asr = FrameASR(model_definition = {\n",
" 'sample_rate': SAMPLE_RATE,\n",
" 'AudioToMelSpectrogramPreprocessor': cfg.preprocessor,\n",
" 'JasperEncoder': cfg.encoder,\n",
" 'labels': cfg.decoder.vocabulary\n",
" },\n",
" frame_len=FRAME_LEN, frame_overlap=2, \n",
" offset=4)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"asr.reset()\n",
"\n",
"p = pa.PyAudio()\n",
"print('Available audio input devices:')\n",
"input_devices = []\n",
"for i in range(p.get_device_count()):\n",
" dev = p.get_device_info_by_index(i)\n",
" if dev.get('maxInputChannels'):\n",
" input_devices.append(i)\n",
" print(i, dev.get('name'))\n",
"\n",
"if len(input_devices):\n",
" dev_idx = -2\n",
" while dev_idx not in input_devices:\n",
" print('Please type input device ID:')\n",
" dev_idx = int(input())\n",
"\n",
" empty_counter = 0\n",
"\n",
" def callback(in_data, frame_count, time_info, status):\n",
" global empty_counter\n",
" signal = np.frombuffer(in_data, dtype=np.int16)\n",
" text = asr.transcribe(signal)\n",
" if len(text):\n",
" print(text,end='')\n",
" empty_counter = asr.offset\n",
" elif empty_counter > 0:\n",
" empty_counter -= 1\n",
" if empty_counter == 0:\n",
" print(' ',end='')\n",
" return (in_data, pa.paContinue)\n",
"\n",
" stream = p.open(format=pa.paInt16,\n",
" channels=CHANNELS,\n",
" rate=SAMPLE_RATE,\n",
" input=True,\n",
" input_device_index=dev_idx,\n",
" stream_callback=callback,\n",
" frames_per_buffer=CHUNK_SIZE)\n",
"\n",
" print('Listening...')\n",
"\n",
" stream.start_stream()\n",
" \n",
" # Interrupt kernel and then speak for a few more words to exit the pyaudio loop !\n",
" try:\n",
" while stream.is_active():\n",
" time.sleep(0.1)\n",
" finally: \n",
" stream.stop_stream()\n",
" stream.close()\n",
" p.terminate()\n",
"\n",
" print()\n",
" print(\"PyAudio stopped\")\n",
" \n",
"else:\n",
" print('ERROR: No audio input device found.')"
]
}
],
"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.7.7"
}
},
"nbformat": 4,
"nbformat_minor": 4
}
tutorials/asr/Online_ASR_Microphone_Demo_Cache_Aware_Streaming.ipynb
+433
View File
@@ -0,0 +1,433 @@
{
"cells": [
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"## Introduction\n",
"\n",
"This notebook allows you to do real-time (\"streaming\") speech recognition using audio recorded from your microphone. This notebook shows how to use a NeMo chunk-aware FastConformer model with caching enabled.\n",
"\n",
"## Installation\n",
"\n",
"The notebook requires PyAudio library, which is used to capture an audio stream from your machine. This means that you need to run this notebook locally. This notebook will not be able to record your audio if you run it in Google Colab or in a Docker container.\n",
"\n",
"For Ubuntu, please run the following commands to install it:\n",
"\n",
"```\n",
"sudo apt install python3-pyaudio\n",
"pip install pyaudio\n",
"```\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"## Install dependencies\n",
"!pip install wget\n",
"!apt-get install sox libsndfile1 ffmpeg portaudio19-dev\n",
"!pip install text-unidecode\n",
"!pip install pyaudio"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# ## Uncomment this cell to install NeMo if it has not been installed\n",
"# BRANCH = 'main'\n",
"# !python -m pip install \"nemo_toolkit[asr] @ git+https://github.com/NVIDIA-NeMo/Speech.git@$BRANCH\""
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# import dependencies\n",
"import copy\n",
"import time\n",
"import pyaudio as pa\n",
"import numpy as np\n",
"import torch\n",
"\n",
"from omegaconf import OmegaConf, open_dict\n",
"\n",
"import nemo.collections.asr as nemo_asr\n",
"from nemo.collections.asr.models.ctc_bpe_models import EncDecCTCModelBPE\n",
"from nemo.collections.asr.parts.utils.streaming_utils import CacheAwareStreamingAudioBuffer\n",
"from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis\n",
"\n",
"# specify sample rate we will use for recording audio\n",
"SAMPLE_RATE = 16000 # Hz"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"## Cache-aware streaming Fastconformer\n",
"In this tutorial, we will do streaming transcription using NeMo models that were specially trained for use in streaming applications. These models are described in the paper released by the NeMo team: [*Noroozi et al.* \"Stateful FastConformer with Cache-based Inference for Streaming Automatic Speech Recognition](https://arxiv.org/abs/2312.17279)\" (accepted to ICASSP 2024).\n",
"\n",
"These models have the following features:\n",
"* They were trained such that at each timestep, the decoder (either RNNT or CTC) would receive a limited amount of context on the left and (most importantly) the right side. Keeping the right side context small means that in a real time streaming scenario, we do not need to keep recording for very long before we are able to compute the output token at that timestep - thus we are able to get transcriptions with a low latency.\n",
"* The model implementation has **caching** enabled, meaning we do not need to recalculate activations that were obtained in previous timesteps, thus reducing latency further.\n",
"\n",
"\n",
"## Model checkpoints\n",
"The following checkpoints of these models are currently available, and are compatible with this notebook. The meaning of \"lookahead\" and \"chunk size\" is described in the following section.\n",
"\n",
"1) [`stt_en_fastconformer_hybrid_large_streaming_80ms`](https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_fastconformer_hybrid_large_streaming_80ms) - 80ms lookahead / 160ms chunk size\n",
"\n",
"2) [`stt_en_fastconformer_hybrid_large_streaming_480ms`](https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_fastconformer_hybrid_large_streaming_480ms) - 480ms lookahead / 540ms chunk size\n",
"\n",
"3) [`stt_en_fastconformer_hybrid_large_streaming_1040ms`](https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_fastconformer_hybrid_large_streaming_1040ms) - 1040ms lookahead / 1120ms chunk size\n",
"\n",
"4) [`stt_en_fastconformer_hybrid_large_streaming_multi`](https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_fastconformer_hybrid_large_streaming_multi) - 0ms, 80ms, 480ms, 1040ms lookahead / 80ms, 160ms, 540ms, 1120ms chunk size\n",
"\n",
"## Model inference explanation\n",
"We run inference by continuously recording our audio in chunks, and feeding the chunks into the chosen ASR model. In this notebook we use `pyaudio` to open an audio input stream, and pass the audio to a `stream_callback` function every \"chunk-sized\" number of seconds. In the `stream_callback` function, we pass the audio signal to a `transcribe` function (which we will specify in this notebook), and print the resulting transcription.\n",
"\n",
"As mentioned, the \"chunk size\" is the duration of audio that we feed into the ASR model at a time (and we keep doing this continuously, to allow for real-time, streaming speech recognition).\n",
"\n",
"\"Lookahead\" size is the \"chunk size\" minus the duration of a single output timestep from the decoder. For FastConformer models, the duration of an output timestep is always 80ms, hence in this notebook always `lookahead size = chunk size - 80 ms`."
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"## Model selection\n",
"In the next cell, you can select which pretrained `model_name` and `lookahead_size` you would like to try.\n",
"\n",
"Additionally, note that all of the available models are [Hybrid RNNT-CTC models](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/models.html#hybrid-transducer-ctc). Inference is by default done using the RNNT decoder (which tends to produce a higher transcription accuracy), but you may choose to use the CTC decoder instead. For this, we also provide a `decoder_type` variable in the cell below."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# You may wish to try different values of model_name and lookahead_size\n",
"\n",
"# Choose a the name of a model to use.\n",
"# Currently available options:\n",
"# 1) \"stt_en_fastconformer_hybrid_large_streaming_multi\"\n",
"# 2) \"stt_en_fastconformer_hybrid_large_streaming_80ms\"\n",
"# 3) \"stt_en_fastconformer_hybrid_large_streaming_480ms\"\n",
"# 4) \"stt_en_fastconformer_hybrid_large_streaming_1040ms\"\n",
"\n",
"model_name = \"stt_en_fastconformer_hybrid_large_streaming_multi\"\n",
"\n",
"# Specify the lookahead_size.\n",
"# If model_name == \"stt_en_fastconformer_hybrid_large_streaming_multi\" then\n",
"# lookahead_size can be 0, 80, 480 or 1040 (ms)\n",
"# Else, lookahead_size should be whatever is written in the model_name:\n",
"# \"stt_en_fastconformer_hybrid_large_streaming_<lookahead_size>ms\"\n",
"\n",
"lookahead_size = 80 # in milliseconds\n",
"\n",
"# Specify the decoder to use.\n",
"# Can be \"rnnt\" or \"ctc\"\n",
"decoder_type = \"rnnt\""
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"## Model set-up\n",
"Next we:\n",
"* set up the `asr_model` according to the chosen `model_name` and `lookahead_size`\n",
"* make sure we use the specified `decoder_type`\n",
"* make sure the model's decoding strategy has suitable parameters\n",
"* instantiate a `CacheAwareStreamingAudioBuffer`\n",
"* get some parameters to use as the initial cache state"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"scrolled": true
},
"outputs": [],
"source": [
"# setting up model and validating the choice of model_name and lookahead size\n",
"asr_model = nemo_asr.models.ASRModel.from_pretrained(model_name=model_name)\n",
"\n",
"\n",
"# specify ENCODER_STEP_LENGTH (which is 80 ms for FastConformer models)\n",
"ENCODER_STEP_LENGTH = 80 # ms\n",
"\n",
"# update att_context_size if using multi-lookahead model\n",
"# (for single-lookahead models, the default context size will be used and the\n",
"# `lookahead_size` variable will be ignored)\n",
"if model_name == \"stt_en_fastconformer_hybrid_large_streaming_multi\":\n",
" # check that lookahead_size is one of the valid ones\n",
" if lookahead_size not in [0, 80, 480, 1040]:\n",
" raise ValueError(\n",
" f\"specified lookahead_size {lookahead_size} is not one of the \"\n",
" \"allowed lookaheads (can select 0, 80, 480 or 1040 ms)\"\n",
" )\n",
"\n",
" # update att_context_size\n",
" left_context_size = asr_model.encoder.att_context_size[0]\n",
" asr_model.encoder.set_default_att_context_size([left_context_size, int(lookahead_size / ENCODER_STEP_LENGTH)])\n",
"\n",
"\n",
"# make sure we use the specified decoder_type\n",
"asr_model.change_decoding_strategy(decoder_type=decoder_type)\n",
"\n",
"# make sure the model's decoding strategy is optimal\n",
"decoding_cfg = asr_model.cfg.decoding\n",
"with open_dict(decoding_cfg):\n",
" # save time by doing greedy decoding and not trying to record the alignments\n",
" decoding_cfg.strategy = \"greedy\"\n",
" decoding_cfg.preserve_alignments = False\n",
" if hasattr(asr_model, 'joint'): # if an RNNT model\n",
" # restrict max_symbols to make sure not stuck in infinite loop\n",
" decoding_cfg.greedy.max_symbols = 10\n",
" # sensible default parameter, but not necessary since batch size is 1\n",
" decoding_cfg.fused_batch_size = -1\n",
" asr_model.change_decoding_strategy(decoding_cfg)\n",
"\n",
"\n",
"# set model to eval mode\n",
"asr_model.eval()\n",
"\n",
"\n",
"# get parameters to use as the initial cache state\n",
"cache_last_channel, cache_last_time, cache_last_channel_len = asr_model.encoder.get_initial_cache_state(\n",
" batch_size=1\n",
")"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"## Transcribing a single chunk\n",
"In the following code block we specify the `transcribe_chunk` function that transcribes a single chunk."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# init params we will use for streaming\n",
"previous_hypotheses = None\n",
"pred_out_stream = None\n",
"step_num = 0\n",
"pre_encode_cache_size = asr_model.encoder.streaming_cfg.pre_encode_cache_size[1]\n",
"# cache-aware models require some small section of the previous processed_signal to\n",
"# be fed in at each timestep - we initialize this to a tensor filled with zeros\n",
"# so that we will do zero-padding for the very first chunk(s)\n",
"num_channels = asr_model.cfg.preprocessor.features\n",
"cache_pre_encode = torch.zeros((1, num_channels, pre_encode_cache_size), device=asr_model.device)\n",
"\n",
"\n",
"# helper function for extracting transcriptions\n",
"def extract_transcriptions(hyps):\n",
" \"\"\"\n",
" The transcribed_texts returned by CTC and RNNT models are different.\n",
" This method would extract and return the text section of the hypothesis.\n",
" \"\"\"\n",
" if isinstance(hyps[0], Hypothesis):\n",
" transcriptions = []\n",
" for hyp in hyps:\n",
" transcriptions.append(hyp.text)\n",
" else:\n",
" transcriptions = hyps\n",
" return transcriptions\n",
"\n",
"# define functions to init audio preprocessor and to\n",
"# preprocess the audio (ie obtain the mel-spectrogram)\n",
"def init_preprocessor(asr_model):\n",
" cfg = copy.deepcopy(asr_model._cfg)\n",
" OmegaConf.set_struct(cfg.preprocessor, False)\n",
"\n",
" # some changes for streaming scenario\n",
" cfg.preprocessor.dither = 0.0\n",
" cfg.preprocessor.pad_to = 0\n",
" cfg.preprocessor.normalize = \"None\"\n",
" \n",
" preprocessor = EncDecCTCModelBPE.from_config_dict(cfg.preprocessor)\n",
" preprocessor.to(asr_model.device)\n",
" \n",
" return preprocessor\n",
"\n",
"preprocessor = init_preprocessor(asr_model)\n",
"\n",
"def preprocess_audio(audio, asr_model):\n",
" device = asr_model.device\n",
"\n",
" # doing audio preprocessing\n",
" audio_signal = torch.from_numpy(audio).unsqueeze_(0).to(device)\n",
" audio_signal_len = torch.Tensor([audio.shape[0]]).to(device)\n",
" processed_signal, processed_signal_length = preprocessor(\n",
" input_signal=audio_signal, length=audio_signal_len\n",
" )\n",
" return processed_signal, processed_signal_length\n",
"\n",
"\n",
"def transcribe_chunk(new_chunk):\n",
" \n",
" global cache_last_channel, cache_last_time, cache_last_channel_len\n",
" global previous_hypotheses, pred_out_stream, step_num\n",
" global cache_pre_encode\n",
" \n",
" # new_chunk is provided as np.int16, so we convert it to np.float32\n",
" # as that is what our ASR models expect\n",
" audio_data = new_chunk.astype(np.float32)\n",
" audio_data = audio_data / 32768.0\n",
"\n",
" # get mel-spectrogram signal & length\n",
" processed_signal, processed_signal_length = preprocess_audio(audio_data, asr_model)\n",
" \n",
" # prepend with cache_pre_encode\n",
" processed_signal = torch.cat([cache_pre_encode, processed_signal], dim=-1)\n",
" processed_signal_length += cache_pre_encode.shape[1]\n",
" \n",
" # save cache for next time\n",
" cache_pre_encode = processed_signal[:, :, -pre_encode_cache_size:]\n",
" \n",
" with torch.no_grad():\n",
" (\n",
" pred_out_stream,\n",
" transcribed_texts,\n",
" cache_last_channel,\n",
" cache_last_time,\n",
" cache_last_channel_len,\n",
" previous_hypotheses,\n",
" ) = asr_model.conformer_stream_step(\n",
" processed_signal=processed_signal,\n",
" processed_signal_length=processed_signal_length,\n",
" cache_last_channel=cache_last_channel,\n",
" cache_last_time=cache_last_time,\n",
" cache_last_channel_len=cache_last_channel_len,\n",
" keep_all_outputs=False,\n",
" previous_hypotheses=previous_hypotheses,\n",
" previous_pred_out=pred_out_stream,\n",
" drop_extra_pre_encoded=None,\n",
" return_transcription=True,\n",
" )\n",
" \n",
" final_streaming_tran = extract_transcriptions(transcribed_texts)\n",
" step_num += 1\n",
" \n",
" return final_streaming_tran[0]"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"## Simple streaming with microphone\n",
"We use `pyaudio` to record audio from an input audio device on your local machine. We use a `stream_callback` which will be called every `frames_per_buffer` number of frames, and conduct the transcription, which will be printed in the output of the cell below."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# calculate chunk_size in milliseconds\n",
"chunk_size = lookahead_size + ENCODER_STEP_LENGTH\n",
"\n",
"p = pa.PyAudio()\n",
"print('Available audio input devices:')\n",
"input_devices = []\n",
"for i in range(p.get_device_count()):\n",
" dev = p.get_device_info_by_index(i)\n",
" if dev.get('maxInputChannels'):\n",
" input_devices.append(i)\n",
" print(i, dev.get('name'))\n",
"\n",
"if len(input_devices):\n",
" dev_idx = -2\n",
" while dev_idx not in input_devices:\n",
" print('Please type input device ID:')\n",
" dev_idx = int(input())\n",
"\n",
" def callback(in_data, frame_count, time_info, status):\n",
" signal = np.frombuffer(in_data, dtype=np.int16)\n",
" text = transcribe_chunk(signal)\n",
" print(text, end='\\r')\n",
" return (in_data, pa.paContinue)\n",
"\n",
" stream = p.open(format=pa.paInt16,\n",
" channels=1,\n",
" rate=SAMPLE_RATE,\n",
" input=True,\n",
" input_device_index=dev_idx,\n",
" stream_callback=callback,\n",
" frames_per_buffer=int(SAMPLE_RATE * chunk_size / 1000) - 1\n",
" )\n",
"\n",
" print('Listening...')\n",
"\n",
" stream.start_stream()\n",
" \n",
" # Interrupt kernel and then speak for a few more words to exit the pyaudio loop !\n",
" try:\n",
" while stream.is_active():\n",
" time.sleep(0.1)\n",
" finally: \n",
" stream.stop_stream()\n",
" stream.close()\n",
" p.terminate()\n",
"\n",
" print()\n",
" print(\"PyAudio stopped\")\n",
" \n",
"else:\n",
" print('ERROR: No audio input device found.')"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
}
],
"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/asr/Online_Noise_Augmentation.ipynb
+1343
View File
File diff suppressed because it is too large Load Diff
tutorials/asr/Online_Offline_Microphone_VAD_Demo.ipynb
+830
View File
@@ -0,0 +1,830 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"\"\"\"\n",
"Please run notebook locally (if you have all the dependencies and a GPU).\n",
"Technically you can run this notebook on Google Colab but you need to set up microphone for 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. Set up microphone for Colab\n",
"\n",
"\n",
"NOTE: User is responsible for checking the content of datasets and the applicable licenses and determining if suitable for the intended use.\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 portaudio19-dev\n",
"!pip install text-unidecode\n",
"!pip install pyaudio\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": [
"# Voice Activity Detection (VAD)\n",
"\n",
"\n",
"This notebook demonstrates how to perform\n",
"1. [offline streaming inference on audio files (offline VAD)](#Offline-streaming-inference);\n",
"2. [finetuning](#Finetune) and use [posterior](#Posterior);\n",
"3. [vad postprocessing and threshold tuning](#VAD-postprocessing-and-Tuning-threshold);\n",
"4. [online streaming inference](#Online-streaming-inference);\n",
"5. [online streaming inference from a microphone's stream](#Online-streaming-inference-through-microphone).\n",
"\n",
"Note the incompatibility of components could lead to failure of running this notebook locally with container, we might deprecate this notebook and provide a better tutorial in soon releases."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The notebook requires PyAudio library to get a signal from an audio device.\n",
"For Ubuntu, please run the following commands to install it:\n",
"```\n",
"sudo apt install python3-pyaudio\n",
"pip install pyaudio\n",
"```"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import numpy as np\n",
"import pyaudio as pa\n",
"import os, time\n",
"import librosa\n",
"import IPython.display as ipd\n",
"import matplotlib.pyplot as plt\n",
"%matplotlib inline\n",
"\n",
"import nemo\n",
"import nemo.collections.asr as nemo_asr"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# sample rate, Hz\n",
"SAMPLE_RATE = 16000"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Restore the model from NGC"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"scrolled": true
},
"outputs": [],
"source": [
"vad_model = nemo_asr.models.EncDecClassificationModel.from_pretrained('vad_marblenet')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Observing the config of the model"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from omegaconf import OmegaConf\n",
"import copy"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Preserve a copy of the full config\n",
"cfg = copy.deepcopy(vad_model._cfg)\n",
"print(OmegaConf.to_yaml(cfg))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Setup preprocessor with these settings"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"vad_model.preprocessor = vad_model.from_config_dict(cfg.preprocessor)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Set model to inference mode\n",
"vad_model.eval();"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"vad_model = vad_model.to(vad_model.device)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We demonstrate two methods for streaming inference:\n",
"1. [offline streaming inference (script)](#Offline-streaming-inference)\n",
"2. [online streaming inference (step-by-step)](#Online-streaming-inference)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Offline streaming inference\n",
"\n",
"VAD relies on shorter fixed-length segments for prediction. \n",
"\n",
"You can find all necessary steps about inference in \n",
"```python\n",
" Script: <NeMo_git_root>/examples/asr/speech_classification/vad_infer.py \n",
" Config: <NeMo_git_root>/examples/asr/conf/vad/vad_inference_postprocessing.yaml\n",
"```\n",
"Duration inference, we generate frame-level prediction by two approaches:\n",
"\n",
"1. shift the window of length `window_length_in_sec` (e.g. 0.63s) by `shift_length_in_sec` (e.g. 10ms) to generate the frame and use the prediction of the window to represent the label for the frame; Use \n",
"```python\n",
" <NeMo_git_root>/examples/asr/speech_classification/vad_infer.py\n",
"```\n",
"\n",
" This script will automatically split long audio file to avoid CUDA memory issue and performing **streaming** inside `AudioLabelDataset`."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Posterior\n",
"<img src=\"https://raw.githubusercontent.com/NVIDIA/NeMo/v1.0.2/tutorials/asr/images/vad_post_overlap_diagram.png\" width=\"500\">\n",
"\n",
"2. generate predictions with overlapping input segments. Then a smoothing filter is applied to decide the label for a frame spanned by multiple segments. Perform this step alongside with above step with flag **gen_overlap_seq=True** or use\n",
"```python\n",
"<NeMo_git_root>/scripts/voice_activity_detection/vad_overlap_posterior.py\n",
"```\n",
"if you already have frame level prediction. \n",
"\n",
"Have a look at [MarbleNet paper](https://arxiv.org/pdf/2010.13886.pdf) for choices about segment length, smoothing filter, etc. And play with those parameters with your data.\n",
"\n",
"You can also find posterior about converting frame level prediction to speech/no-speech segment in start and end times format in `vad_overlap_posterior.py` or use flag **gen_seg_table=True** alongside with `vad_infer.py`"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Finetune\n",
"You might need to finetune on your data for better performance. For finetuning/transfer learning, please refer to [**Transfer learning** part of ASR tutorial](https://github.com/NVIDIA/NeMo/blob/stable/tutorials/asr/ASR_with_NeMo.ipynb)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## VAD postprocessing and Tuning threshold"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We can use a single **threshold** (achieved by onset=offset=0.5) to binarize predictions or use typical VAD postprocessing including\n",
"\n",
"### Binarization:\n",
"1. **onset** and **offset** threshold for detecting the beginning and end of a speech;\n",
"2. padding durations before (**pad_onset**) and after (**pad_offset**) each speech segment.\n",
"\n",
"### Filtering:\n",
"1. threshold for short speech segment deletion (**min_duration_on**);\n",
"2. threshold for small silence deletion (**min_duration_off**);\n",
"3. Whether to perform short speech segment deletion first (**filter_speech_first**).\n",
"\n",
"\n",
"Of course you can do threshold tuning on frame level prediction. We also provide a script \n",
"```python\n",
"<NeMo_git_root>/scripts/voice_activity_detection/vad_tune_threshold.py\n",
"```\n",
"\n",
"to help you find best thresholds if you have ground truth label file in RTTM format. "
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Online streaming inference"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Setting up data for Streaming Inference"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from nemo.core.classes import IterableDataset\n",
"from nemo.core.neural_types import NeuralType, AudioSignal, LengthsType\n",
"import torch\n",
"from torch.utils.data import DataLoader"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# simple data layer to pass audio signal\n",
"class AudioDataLayer(IterableDataset):\n",
" @property\n",
" def output_types(self):\n",
" return {\n",
" 'audio_signal': NeuralType(('B', 'T'), AudioSignal(freq=self._sample_rate)),\n",
" 'a_sig_length': NeuralType(tuple('B'), LengthsType()),\n",
" }\n",
"\n",
" def __init__(self, sample_rate):\n",
" super().__init__()\n",
" self._sample_rate = sample_rate\n",
" self.output = True\n",
"\n",
" def __iter__(self):\n",
" return self\n",
"\n",
" def __next__(self):\n",
" if not self.output:\n",
" raise StopIteration\n",
" self.output = False\n",
" return torch.as_tensor(self.signal, dtype=torch.float32), \\\n",
" torch.as_tensor(self.signal_shape, dtype=torch.int64)\n",
"\n",
" def set_signal(self, signal):\n",
" self.signal = signal.astype(np.float32)/32768.\n",
" self.signal_shape = self.signal.size\n",
" self.output = True\n",
"\n",
" def __len__(self):\n",
" return 1"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"data_layer = AudioDataLayer(sample_rate=cfg.train_ds.sample_rate)\n",
"data_loader = DataLoader(data_layer, batch_size=1, collate_fn=data_layer.collate_fn)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# inference method for audio signal (single instance)\n",
"def infer_signal(model, signal):\n",
" data_layer.set_signal(signal)\n",
" batch = next(iter(data_loader))\n",
" audio_signal, audio_signal_len = batch\n",
" audio_signal, audio_signal_len = audio_signal.to(vad_model.device), audio_signal_len.to(vad_model.device)\n",
" logits = model.forward(input_signal=audio_signal, input_signal_length=audio_signal_len)\n",
" return logits"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# class for streaming frame-based VAD\n",
"# 1) use reset() method to reset FrameVAD's state\n",
"# 2) call transcribe(frame) to do VAD on\n",
"# contiguous signal's frames\n",
"# To simplify the flow, we use single threshold to binarize predictions.\n",
"class FrameVAD:\n",
"\n",
" def __init__(self, model_definition,\n",
" threshold=0.5,\n",
" frame_len=2, frame_overlap=2.5,\n",
" offset=10):\n",
" '''\n",
" Args:\n",
" threshold: If prob of speech is larger than threshold, classify the segment to be speech.\n",
" frame_len: frame's duration, seconds\n",
" frame_overlap: duration of overlaps before and after current frame, seconds\n",
" offset: number of symbols to drop for smooth streaming\n",
" '''\n",
" self.vocab = list(model_definition['labels'])\n",
" self.vocab.append('_')\n",
"\n",
" self.sr = model_definition['sample_rate']\n",
" self.threshold = threshold\n",
" self.frame_len = frame_len\n",
" self.n_frame_len = int(frame_len * self.sr)\n",
" self.frame_overlap = frame_overlap\n",
" self.n_frame_overlap = int(frame_overlap * self.sr)\n",
" timestep_duration = model_definition['AudioToMFCCPreprocessor']['window_stride']\n",
" for block in model_definition['JasperEncoder']['jasper']:\n",
" timestep_duration *= block['stride'][0] ** block['repeat']\n",
" self.buffer = np.zeros(shape=2*self.n_frame_overlap + self.n_frame_len,\n",
" dtype=np.float32)\n",
" self.offset = offset\n",
" self.reset()\n",
"\n",
" def _decode(self, frame, offset=0):\n",
" assert len(frame)==self.n_frame_len\n",
" self.buffer[:-self.n_frame_len] = self.buffer[self.n_frame_len:]\n",
" self.buffer[-self.n_frame_len:] = frame\n",
" logits = infer_signal(vad_model, self.buffer).cpu().numpy()[0]\n",
" decoded = self._greedy_decoder(\n",
" self.threshold,\n",
" logits,\n",
" self.vocab\n",
" )\n",
" return decoded\n",
"\n",
"\n",
" @torch.no_grad()\n",
" def transcribe(self, frame=None):\n",
" if frame is None:\n",
" frame = np.zeros(shape=self.n_frame_len, dtype=np.float32)\n",
" if len(frame) < self.n_frame_len:\n",
" frame = np.pad(frame, [0, self.n_frame_len - len(frame)], 'constant')\n",
" unmerged = self._decode(frame, self.offset)\n",
" return unmerged\n",
"\n",
" def reset(self):\n",
" '''\n",
" Reset frame_history and decoder's state\n",
" '''\n",
" self.buffer=np.zeros(shape=self.buffer.shape, dtype=np.float32)\n",
" self.prev_char = ''\n",
"\n",
" @staticmethod\n",
" def _greedy_decoder(threshold, logits, vocab):\n",
" s = []\n",
" if logits.shape[0]:\n",
" probs = torch.softmax(torch.as_tensor(logits), dim=-1)\n",
" probas, _ = torch.max(probs, dim=-1)\n",
" probas_s = probs[1].item()\n",
" preds = 1 if probas_s >= threshold else 0\n",
" s = [preds, str(vocab[preds]), probs[0].item(), probs[1].item(), str(logits)]\n",
" return s"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"\n",
"\n",
"Streaming inference depends on a few factors, such as the frame length (STEP) and buffer size (WINDOW SIZE). Experiment with a few values to see their effects in the below cells."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"STEP_LIST = [0.01,0.01]\n",
"WINDOW_SIZE_LIST = [0.31,0.15]"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import wave\n",
"\n",
"def offline_inference(wave_file, STEP = 0.025, WINDOW_SIZE = 0.5, threshold=0.5):\n",
"\n",
" FRAME_LEN = STEP # infer every STEP seconds\n",
" CHANNELS = 1 # number of audio channels (expect mono signal)\n",
" RATE = 16000 # sample rate, Hz\n",
"\n",
"\n",
" CHUNK_SIZE = int(FRAME_LEN*RATE)\n",
"\n",
" vad = FrameVAD(model_definition = {\n",
" 'sample_rate': SAMPLE_RATE,\n",
" 'AudioToMFCCPreprocessor': cfg.preprocessor,\n",
" 'JasperEncoder': cfg.encoder,\n",
" 'labels': cfg.labels\n",
" },\n",
" threshold=threshold,\n",
" frame_len=FRAME_LEN, frame_overlap = (WINDOW_SIZE-FRAME_LEN)/2,\n",
" offset=0)\n",
"\n",
" wf = wave.open(wave_file, 'rb')\n",
" p = pa.PyAudio()\n",
"\n",
" empty_counter = 0\n",
"\n",
" preds = []\n",
" proba_b = []\n",
" proba_s = []\n",
"\n",
" data = wf.readframes(CHUNK_SIZE)\n",
"\n",
" while len(data) > 0:\n",
"\n",
" data = wf.readframes(CHUNK_SIZE)\n",
" signal = np.frombuffer(data, dtype=np.int16)\n",
" result = vad.transcribe(signal)\n",
"\n",
" preds.append(result[0])\n",
" proba_b.append(result[2])\n",
" proba_s.append(result[3])\n",
"\n",
" if len(result):\n",
" print(result,end='\\n')\n",
" empty_counter = 3\n",
" elif empty_counter > 0:\n",
" empty_counter -= 1\n",
" if empty_counter == 0:\n",
" print(' ',end='')\n",
"\n",
" p.terminate()\n",
" vad.reset()\n",
"\n",
" return preds, proba_b, proba_s"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Here we show an example of online streaming inference\n",
"You can use your file or download the provided demo audio file. "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"scrolled": true
},
"outputs": [],
"source": [
"demo_wave = 'VAD_demo.wav'\n",
"if not os.path.exists(demo_wave):\n",
" !wget \"https://dldata-public.s3.us-east-2.amazonaws.com/VAD_demo.wav\""
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"wave_file = demo_wave\n",
"\n",
"CHANNELS = 1\n",
"RATE = 16000\n",
"audio, sample_rate = librosa.load(wave_file, sr=RATE)\n",
"dur = librosa.get_duration(y=audio, sr=sample_rate)\n",
"print(dur)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"scrolled": true
},
"outputs": [],
"source": [
"ipd.Audio(audio, rate=sample_rate)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"scrolled": true
},
"outputs": [],
"source": [
"threshold=0.4\n",
"\n",
"results = []\n",
"for STEP, WINDOW_SIZE in zip(STEP_LIST, WINDOW_SIZE_LIST, ):\n",
" print(f'====== STEP is {STEP}s, WINDOW_SIZE is {WINDOW_SIZE}s ====== ')\n",
" preds, proba_b, proba_s = offline_inference(wave_file, STEP, WINDOW_SIZE, threshold)\n",
" results.append([STEP, WINDOW_SIZE, preds, proba_b, proba_s])"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"To simplify the flow, the above prediction is based on single threshold and `threshold=0.4`.\n",
"\n",
"You can play with other [threshold](#VAD-postprocessing-and-Tuning-threshold) or use postprocessing and see how they would impact performance. \n",
"\n",
"**Note** if you want better performance, [finetune](#Finetune) on your data and use posteriors such as [overlapped prediction](#Posterior). \n",
"\n",
"Let's plot the prediction and melspectrogram"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import librosa.display\n",
"plt.figure(figsize=[20,10])\n",
"\n",
"num = len(results)\n",
"for i in range(num):\n",
" len_pred = len(results[i][2])\n",
" FRAME_LEN = results[i][0]\n",
" ax1 = plt.subplot(num+1,1,i+1)\n",
"\n",
" ax1.plot(np.arange(audio.size) / sample_rate, audio, 'b')\n",
" ax1.set_xlim([-0.01, int(dur)+1])\n",
" ax1.tick_params(axis='y', labelcolor= 'b')\n",
" ax1.set_ylabel('Signal')\n",
" ax1.set_ylim([-1, 1])\n",
"\n",
" proba_s = results[i][4]\n",
" pred = [1 if p > threshold else 0 for p in proba_s]\n",
" ax2 = ax1.twinx()\n",
" ax2.plot(np.arange(len_pred)/(1/results[i][0]), np.array(pred) , 'r', label='pred')\n",
" ax2.plot(np.arange(len_pred)/(1/results[i][0]), np.array(proba_s) , 'g--', label='speech prob')\n",
" ax2.tick_params(axis='y', labelcolor='r')\n",
" legend = ax2.legend(loc='lower right', shadow=True)\n",
" ax1.set_ylabel('prediction')\n",
"\n",
" ax2.set_title(f'step {results[i][0]}s, buffer size {results[i][1]}s')\n",
" ax2.set_ylabel('Preds and Probas')\n",
"\n",
"\n",
"ax = plt.subplot(num+1,1,num+1)\n",
"S = librosa.feature.melspectrogram(y=audio, sr=sample_rate, n_mels=64, fmax=8000)\n",
"S_dB = librosa.power_to_db(S, ref=np.max)\n",
"librosa.display.specshow(S_dB, x_axis='time', y_axis='mel', sr=sample_rate, fmax=8000)\n",
"ax.set_title('Mel-frequency spectrogram')\n",
"ax.grid()\n",
"plt.show()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Online streaming inference through microphone"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"**Please note the VAD model is not perfect for various microphone input and you might need to finetune on your input and play with different parameters.**"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"STEP = 0.01\n",
"WINDOW_SIZE = 0.31\n",
"CHANNELS = 1\n",
"RATE = 16000\n",
"FRAME_LEN = STEP\n",
"THRESHOLD = 0.5\n",
"\n",
"CHUNK_SIZE = int(STEP * RATE)\n",
"vad = FrameVAD(model_definition = {\n",
" 'sample_rate': SAMPLE_RATE,\n",
" 'AudioToMFCCPreprocessor': cfg.preprocessor,\n",
" 'JasperEncoder': cfg.encoder,\n",
" 'labels': cfg.labels\n",
" },\n",
" threshold=THRESHOLD,\n",
" frame_len=FRAME_LEN, frame_overlap=(WINDOW_SIZE - FRAME_LEN) / 2,\n",
" offset=0)\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"vad.reset()\n",
"\n",
"p = pa.PyAudio()\n",
"print('Available audio input devices:')\n",
"input_devices = []\n",
"for i in range(p.get_device_count()):\n",
" dev = p.get_device_info_by_index(i)\n",
" if dev.get('maxInputChannels'):\n",
" input_devices.append(i)\n",
" print(i, dev.get('name'))\n",
"\n",
"if len(input_devices):\n",
" dev_idx = -2\n",
" while dev_idx not in input_devices:\n",
" print('Please type input device ID:')\n",
" dev_idx = int(input())\n",
"\n",
" empty_counter = 0\n",
"\n",
" def callback(in_data, frame_count, time_info, status):\n",
" global empty_counter\n",
" signal = np.frombuffer(in_data, dtype=np.int16)\n",
" text = vad.transcribe(signal)\n",
" if len(text):\n",
" print(text,end='\\n')\n",
" empty_counter = vad.offset\n",
" elif empty_counter > 0:\n",
" empty_counter -= 1\n",
" if empty_counter == 0:\n",
" print(' ',end='\\n')\n",
" return (in_data, pa.paContinue)\n",
"\n",
" stream = p.open(format=pa.paInt16,\n",
" channels=CHANNELS,\n",
" rate=SAMPLE_RATE,\n",
" input=True,\n",
" input_device_index=dev_idx,\n",
" stream_callback=callback,\n",
" frames_per_buffer=CHUNK_SIZE)\n",
"\n",
" print('Listening...')\n",
"\n",
" stream.start_stream()\n",
"\n",
" # Interrupt kernel and then speak for a few more words to exit the pyaudio loop !\n",
" try:\n",
" while stream.is_active():\n",
" time.sleep(0.1)\n",
" finally:\n",
" stream.stop_stream()\n",
" stream.close()\n",
" p.terminate()\n",
"\n",
" print()\n",
" print(\"PyAudio stopped\")\n",
"\n",
"else:\n",
" print(\"ERROR: No audio input device found, please check if the jupyter notebook has access to your computer's microphone.\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"pycharm": {
"name": "#%% md\n"
}
},
"source": [
"## ONNX Deployment\n",
"You can also export the model to ONNX file and deploy it to TensorRT or MS ONNX Runtime inference engines. If you don't have one installed yet, please run:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!pip install --upgrade onnxruntime # for gpu, use onnxruntime-gpu\n",
"# !mkdir -p ort\n",
"# %cd ort\n",
"# !git clone --depth 1 --branch v1.8.0 https://github.com/microsoft/onnxruntime.git .\n",
"# !./build.sh --skip_tests --config Release --build_shared_lib --parallel --use_cuda --cuda_home /usr/local/cuda --cudnn_home /usr/lib/x86_64-linux-gnu --build_wheel\n",
"# !pip install ./build/Linux/Release/dist/onnxruntime*.whl\n",
"# %cd .."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Then just replace `infer_signal` implementation with this code:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"pycharm": {
"name": "#%%\n"
}
},
"outputs": [],
"source": [
"import onnxruntime\n",
"vad_model.export('vad.onnx')\n",
"ort_session = onnxruntime.InferenceSession('vad.onnx', providers=['CPUExecutionProvider'])\n",
"\n",
"def to_numpy(tensor):\n",
" return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()\n",
"\n",
"def infer_signal(signal):\n",
" data_layer.set_signal(signal)\n",
" batch = next(iter(data_loader))\n",
" audio_signal, audio_signal_len = batch\n",
" audio_signal, audio_signal_len = audio_signal.to(vad_model.device), audio_signal_len.to(vad_model.device)\n",
" processed_signal, processed_signal_len = vad_model.preprocessor(\n",
" input_signal=audio_signal, length=audio_signal_len,\n",
" )\n",
" ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(processed_signal), }\n",
" ologits = ort_session.run(None, ort_inputs)\n",
" alogits = np.asarray(ologits)\n",
" logits = torch.from_numpy(alogits[0])\n",
" return logits"
]
}
],
"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.7.7"
}
},
"nbformat": 4,
"nbformat_minor": 4
}
tutorials/asr/README.md
+67
View File
@@ -0,0 +1,67 @@
# Speech Recognition Tutorials
------------
In this repository, you will find several tutorials discussing what is Automatic Speech Recognition (ASR), general concepts, specific models and multiple sub-domains of ASR such as Speech Classification, Voice Activity Detection, Speaker Recognition, Speaker Identification and Speaker Diarization.
------------
# Automatic Speech Recognition
1) `ASR_with_NeMo`: Discussion of the task of ASR, handling of data, understanding the acoustic features, using an Acoustic Model and train on an ASR dataset, and finally evaluating the model's performance.
2) `ASR_with_Subword_Tokenization`: Modern ASR models benefit from several improvements in neural network design and data processing. In this tutorial we discuss how we can use Tokenizers (commonly found in NLP) to significantly improve the efficiency of ASR models without sacrificing any accuracy during transcription.
3) `ASR_CTC_Language_Finetuning`: Until now, we have discussed how to train ASR models from scratch. Once we get pretrained ASR models, we can then fine-tune them on domain specific use cases, or even other languages! This notebook discusses how to fine-tune an English ASR model onto another language, and discusses several methods to improve the efficiency of transfer learning.
4) `Online_ASR_Microphone_Demo`: A short notebook that enables us to speak into a microphone and transcribe speech in an online manner. Note that this is not the most efficient way to perform streaming ASR, and it is more of a demo.
5) `Online_ASR_Microphone_Demo_Cache_Aware_Streaming`: This notebook allows you to do real-time ("streaming") speech recognition on audio recorded from your microphone, using "cache-aware" NeMo ASR models specifically tuned for the streaming ASR usecase.
6) `ASR_for_telephony_speech`: Audio sources are not homogenous, nor are the ways to store large audio datasets. Here, we discuss our observations and recommendations when working with audio obtained from Telephony speech sources.
7) `Online_Noise_Augmentation`: While academic datasets are useful for training ASR model, there can often be cases where such datasets are pristine and don't really represent the use case in the real world. So we discuss how to make the model more noise robust with Online audio augmentation.
8) `Intro_to_Transducers`: Previous tutorials discuss ASR models in context of the Connectionist Temporal Classification Loss. In this tutorial, we introduce the Transducer loss, and the components of this loss function that are constructed in the config file. This tutorial is a prerequisite to the `ASR_with_Transducers` tutorial.
9) `ASR_with_Transducers`: In this tutorial, we take a deep dive into Transducer based ASR models, discussing the similarity of setup and config to CTC models and then train a small ContextNet model on the AN4 dataset. We then discuss how to change the decoding strategy of a trained Transducer from greedy search to beam search. Finally, we wrap up this tutorial by extraining the alignment matrix from a trained Transducer model.
10) `Self_Supervised_Pre_Training`: It can often be difficult to obtain labeled data for ASR training. In this tutorial, we demonstrate how to pre-train a speech model in an unsupervised manner, and then fine-tune with CTC loss.
11) `Offline_ASR_with_VAD_for_CTC_models`: In this tutorial, we will demonstrate how to use offline VAD to extract speech segments and transcribe the speech segments with CTC models. This will help to exclude some non_speech utterances and could save computation resources by removing unnecessary input to the ASR system.
12) `Multilang_ASR`: We will learn how to work with existing checkpoints of multilingual ASR models and how to train new ones. It is possible to create a multilingual version of any ASR model that uses tokenizers. This notebook shows how to create a multilingual version of the small monolingual Conformer Transducer model.
13) `ASR_Example_CommonVoice_Finetuning`: Learn how to fine-tune an ASR model using CommonVoice to a new alphabet, Esperanto. We walk through the data processing steps of MCV data using HuggingFace Datasets, preparation of the tokenizer, model and then setup fine-tuning.
14) `ASR_Context_Biasing`: This tutorial aims to show how to improve the recognition accuracy of specific words in NeMo Framework for CTC and Trasducer (RNN-T) ASR models by using the fast context-biasing method with CTC-based Word Spotter.
----------------
# Automatic Speech Recognition with Adapters
Please refer to the `asr_adapter` sub-folder which contains tutorials on the use of `Adapter` modules to perform domain adaptation on ASR models, as well as its sub-domains.
----------------
# Streaming / Buffered Automatic Speech Recognition
1) `Streaming_ASR`: Some ASR models cannot be used to evaluate very long audio segments due to their design. For example, self attention models consume quadratic memory with respect to sequence length. For such cases, this notebook shows how to perform streaming audio recognition in a buffered manner.
2) `Buffered_Transducer_Inference`: In this notebook, we explore a simple algorithm to perform streaming audio recognition in a buffered manner for Transducer models. This enables the use of transducers on very long speech segments, similar to CTC models.
3) `Buffered_Transducer_Inference_with_LCS_Merge`: This is an optional notebook, that discusses a different merge algorithm that can be utilized for streaming/buffered inference for Transducer models. It is not a required tutorial, but is useful for researchers who wish to analyse and improve buffered inference algorithms.
4) `Streaming_ASR_Pipelines`: This tutorial introduces the ASR Inference Pipeline API for building real-time streaming speech recognition systems. It covers the core abstractions (Frame, Stream, Pipeline, TranscribeStepOutput) and demonstrates how to construct pipelines for buffered inference (CTC, RNNT, TDT) and cache-aware streaming models (CTC, RNNT). Advanced features such as continuous batching with per-stream options, End-of-Utterance (EoU) detection, word timestamps, Inverse Text Normalization (ITN), and Neural Machine Translation (NMT) for streaming speech translation are also explored.
----------------
# Speech Command Recognition
1) `Speech_Commands`: Here, we study the task of speech classification - a subset of speech recognition that allows us to classify a spoken sentence into a single label. This allows to speak a command and the model can then recognize this command and perform an action.
2) `Online_Offline_Speech_Commands_Demo`: We perform a joint online-offline inference of speech command recognition. We utilize an online VAD model to detect speech segments (whether audio is in fact speech or background), and if speech is detected then a speech command recognition model classifies that speech in an offline manner. Note that this demo is a demonstration of a possible approach and is not meant for large scale use.
3) `Voice_Activity_Detection`: A special case of Speech Command Recognition - where the task is to classify whether some audio segment is speech or not. It is often a tiny model that is used prior to a large ASR model being used.
4) `Online_Offline_Microphone_VAD_Demo`: Similar to before, we demo an online-offline inference of voice activity detection. We discuss metrics for comparing the performance of streaming VAD models, and how one can try to perform streaming VAD inference with a microphone. Note that as always, this demo is a demonstration of a possible approach and is not meant for large scale use.
tutorials/asr/Self_Supervised_Pre_Training.ipynb
+730
View File
@@ -0,0 +1,730 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "lJz6FDU1lRzc"
},
"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",
"\n",
"NOTE: User is responsible for checking the content of datasets and the applicable licenses and determining if suitable for the intended use.\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 matplotlib>=3.3.2\n",
"\n",
"## Install NeMo\n",
"BRANCH = 'main'\n",
"!python -m pip install \"nemo_toolkit[all] @ git+https://github.com/NVIDIA-NeMo/Speech.git@$BRANCH\"\n",
"\n",
"\"\"\"\n",
"Remember to restart the runtime for the kernel to pick up any upgraded packages (e.g. matplotlib)!\n",
"Alternatively, you can uncomment the exit() below to crash and restart the kernel, in the case\n",
"that you want to use the \"Run All Cells\" (or similar) option.\n",
"\"\"\"\n",
"# exit()"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "FGIVHjS1YEPw"
},
"source": [
"# Self-Supervised pre-training for ASR\n",
"\n",
"This notebook is a basic tutorial for pre-training a model using the self-supervised approach. With this approach, we use a training objective that does not require our dataset to be labeled, which significantly reduces the difficulty of collecting data for pre-training the model. After pre-training our encoder in this way, we can use it in a CTC or RNNT ASR model.\n",
"\n",
"The approach we will use for pre-training our models is represented in the following diagram:\n",
"\n",
" ![SSL diagram](https://raw.githubusercontent.com/NVIDIA/NeMo/main/tutorials/asr/images/contrastive_ssl.png)\n",
"\n",
"We first mask parts of our input using SpecAugment. The model is then trained to solve a contrastive task of distinguishing the latent representation of the masked time steps from several sampled distractors. Since our encoders also contain stride blocks which reduce the length of the inputs, in order to obtain target representations we combine several consecutive time steps. They are then passed through a quantizer, which has been found to help with contrastive pre-training."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "i5XZWBnTf1pT"
},
"source": [
"# Preparing our data"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "38aYTCTIlRzh"
},
"source": [
"## Downloading dataset\n",
"\n",
"In order to demonstrate how to pre-train the model, we will use the AN4 dataset. Note: this is dataset is much smaller than one that we would typically want to use for self-supervised training, however it will suffice for this tutorial. This dataset also contains transcriptions, but they will be ignored for self-supervised pre-training.\n",
"\n",
"Before we get started, let's download and prepare the dataset. The utterances are available as `.sph` files, so we will need to convert them to `.wav` for later processing. If you are not using Google Colab, please make sure you have [Sox](http://sox.sourceforge.net/) installed for this step--see the \"Downloads\" section of the linked Sox homepage. (If you are using Google Colab, Sox should have already been installed in the setup cell at the beginning.)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "gAhsmi6HlRzh"
},
"outputs": [],
"source": [
"# This is where the an4/ directory will be placed.\n",
"# Change this if you don't want the data to be extracted in the current directory.\n",
"data_dir = '.'"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": true,
"id": "Yb4fuUvWlRzk",
"scrolled": true
},
"outputs": [],
"source": [
"import glob\n",
"import os\n",
"import subprocess\n",
"import tarfile\n",
"import wget\n",
"\n",
"# Download the dataset. This will take a few moments...\n",
"print(\"******\")\n",
"if not os.path.exists(data_dir + '/an4_sphere.tar.gz'):\n",
" an4_url = 'https://dldata-public.s3.us-east-2.amazonaws.com/an4_sphere.tar.gz'\n",
" an4_path = wget.download(an4_url, data_dir)\n",
" print(f\"Dataset downloaded at: {an4_path}\")\n",
"else:\n",
" print(\"Tarfile already exists.\")\n",
" an4_path = data_dir + '/an4_sphere.tar.gz'\n",
"\n",
"if not os.path.exists(data_dir + '/an4/'):\n",
" # Untar and convert .sph to .wav (using sox)\n",
" tar = tarfile.open(an4_path)\n",
" tar.extractall(path=data_dir)\n",
"\n",
" print(\"Converting .sph to .wav...\")\n",
" sph_list = glob.glob(data_dir + '/an4/**/*.sph', recursive=True)\n",
" for sph_path in sph_list:\n",
" wav_path = sph_path[:-4] + '.wav'\n",
" cmd = [\"sox\", sph_path, wav_path]\n",
" subprocess.run(cmd)\n",
"print(\"Finished conversion.\\n******\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "m_LFeM0elRzm"
},
"source": [
"You should now have a folder called `an4` that contains `etc/an4_train.transcription`, `etc/an4_test.transcription`, audio files in `wav/an4_clstk` and `wav/an4test_clstk`, along with some other files we will not need.\n",
"\n",
"Now we can load and take a look at the data. As an example, file `cen2-mgah-b.wav` is a 2.6 second-long audio recording of a man saying the letters \"G L E N N\" one-by-one. To confirm this, we can listen to the file:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "_M_bSs3MjQlz"
},
"outputs": [],
"source": [
"import librosa\n",
"import IPython.display as ipd\n",
"\n",
"# Load and listen to the audio file\n",
"example_file = data_dir + '/an4/wav/an4_clstk/mgah/cen2-mgah-b.wav'\n",
"audio, sample_rate = librosa.load(example_file)\n",
"\n",
"ipd.Audio(example_file, rate=sample_rate)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "RdNyw1b_zgtm"
},
"source": [
"## Creating Data Manifests\n",
"\n",
"The first thing we need to do now is to create manifests for our training and evaluation data, which will contain the metadata of our audio files. NeMo data sets take in a standardized manifest format where each line corresponds to one sample of audio, such that the number of lines in a manifest is equal to the number of samples that are represented by that manifest. \n",
"\n",
"\n",
"A line must contain the path to an audio file, and the duration of the audio sample. For labeled datasets it will also contain the corresponding transcript (or path to a transcript file). Even though for self-supervised pre-training the transcript is unnecessary, we will still add it to our manifest, since we will also be using this dataset to demonstrate fine-tuning later on.\n",
"\n",
"Here's an example of what one line in a NeMo-compatible manifest might look like:\n",
"```\n",
"{\"audio_filepath\": \"path/to/audio.wav\", \"duration\": 3.45, \"text\": \"this is a nemo tutorial\"}\n",
"```\n",
"\n",
"We can build our training and evaluation manifests using `an4/etc/an4_train.transcription` and `an4/etc/an4_test.transcription`, which have lines containing transcripts and their corresponding audio file IDs:\n",
"```\n",
"...\n",
"<s> P I T T S B U R G H </s> (cen5-fash-b)\n",
"<s> TWO SIX EIGHT FOUR FOUR ONE EIGHT </s> (cen7-fash-b)\n",
"...\n",
"```"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "lVB1sG1GlRzz"
},
"outputs": [],
"source": [
"# --- Building Manifest Files --- #\n",
"import json\n",
"\n",
"# Function to build a manifest\n",
"def build_manifest(transcripts_path, manifest_path, wav_path):\n",
" with open(transcripts_path, 'r') as fin:\n",
" with open(manifest_path, 'w') as fout:\n",
" for line in fin:\n",
" # Lines look like this:\n",
" # <s> transcript </s> (fileID)\n",
" transcript = line[: line.find('(')-1].lower()\n",
" transcript = transcript.replace('<s>', '').replace('</s>', '')\n",
" transcript = transcript.strip()\n",
"\n",
" file_id = line[line.find('(')+1 : -2] # e.g. \"cen4-fash-b\"\n",
" audio_path = os.path.join(\n",
" data_dir, wav_path,\n",
" file_id[file_id.find('-')+1 : file_id.rfind('-')],\n",
" file_id + '.wav')\n",
"\n",
" duration = librosa.core.get_duration(path=audio_path)\n",
"\n",
" # Write the metadata to the manifest\n",
" metadata = {\n",
" \"audio_filepath\": audio_path,\n",
" \"duration\": duration,\n",
" \"text\": transcript\n",
" }\n",
" json.dump(metadata, fout)\n",
" fout.write('\\n')\n",
" \n",
"# Building Manifests\n",
"print(\"******\")\n",
"train_transcripts = data_dir + '/an4/etc/an4_train.transcription'\n",
"train_manifest = data_dir + '/an4/train_manifest.json'\n",
"if not os.path.isfile(train_manifest):\n",
" build_manifest(train_transcripts, train_manifest, 'an4/wav/an4_clstk')\n",
" print(\"Training manifest created.\")\n",
"\n",
"test_transcripts = data_dir + '/an4/etc/an4_test.transcription'\n",
"test_manifest = data_dir + '/an4/test_manifest.json'\n",
"if not os.path.isfile(test_manifest):\n",
" build_manifest(test_transcripts, test_manifest, 'an4/wav/an4test_clstk')\n",
" print(\"Test manifest created.\")\n",
"print(\"***Done***\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "FLr1aJ-lf4_7"
},
"source": [
"# Self-supervised pre-training"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "sT_uewIDk8oA"
},
"source": [
"## Setting up the config\n",
"\n",
"First, let's download the default config for ssl pre-training of Citrinet-1024, and the corresponding supervised training config, which we will use for fine-tuning later."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "978wWK7AKIvU"
},
"outputs": [],
"source": [
"## Grab the configs we'll use in this example\n",
"!mkdir configs\n",
"!wget -P configs/ https://raw.githubusercontent.com/NVIDIA/NeMo/$BRANCH/examples/asr/conf/ssl/fastconformer/fast-conformer.yaml\n",
"!wget -P configs/ https://raw.githubusercontent.com/NVIDIA/NeMo/$BRANCH/examples/asr/conf/fastconformer/fast-conformer_ctc_bpe.yaml\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "RLzjCgmHuJ_j"
},
"source": [
"Since this config is for a very large model, we will modify it to make a much smaller version for the purpose of this tutorial by reducing the number of channels and the number of sub-blocks in each block."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Vmw9WbTPKHHA"
},
"outputs": [],
"source": [
"from omegaconf import OmegaConf\n",
"import torch\n",
"\n",
"config_path = data_dir + '/configs/fast-conformer.yaml'\n",
"\n",
"cfg = OmegaConf.load(config_path)\n",
"\n",
"cfg.model.encoder.n_layers = 8\n",
"cfg.model.encoder.d_model = 256\n",
"cfg.model.encoder.n_heads = 4\n",
"\n",
"cfg.model.train_ds.manifest_filepath = train_manifest\n",
"cfg.model.train_ds.batch_size = 16\n",
"\n",
"cfg.model.validation_ds.manifest_filepath = test_manifest\n",
"cfg.model.validation_ds.batch_size = 16\n",
"\n",
"if torch.cuda.is_available():\n",
" cfg.trainer.accelerator = 'gpu'\n",
" cfg.trainer.strategy = 'auto'\n",
" cfg.trainer.devices = 1\n",
"else:\n",
" cfg.trainer.accelerator = 'cpu'\n",
" cfg.trainer.strategy = 'auto'\n",
" cfg.trainer.devices = 0\n",
"\n",
"cfg.exp_manager.exp_dir = data_dir + \"/content/exp\"\n",
"cfg.exp_manager.name = \"pre_trained\"\n",
"cfg.exp_manager.use_datetime_version = False\n",
"cfg.exp_manager.create_tensorboard_logger = False\n",
"cfg.exp_manager.resume_if_exists = True\n",
"cfg.exp_manager.resume_ignore_no_checkpoint = True\n",
"cfg.exp_manager.checkpoint_callback_params.save_best_model = True\n",
"\n",
"cfg.trainer.check_val_every_n_epoch = 1\n",
"\n",
"cfg.model.optim.sched.name = \"CosineAnnealing\"\n",
"cfg.model.optim.sched.warmup_steps = 1000\n",
"cfg.model.optim.sched.max_steps = 2000\n",
"#in practice you will usually want a much larger amount of pre-training steps\n",
"cfg.model.optim.sched.min_lr = 0\n",
"cfg.model.optim.lr = 0.015\n",
"cfg.model.optim.weight_decay = 0\n",
"del cfg.model.optim.sched.d_model\n",
"\n",
"cfg.trainer.max_steps = cfg.model.optim.sched.max_steps\n",
"\n",
"cfg.model.spec_augment.patch_size = 16\n",
"cfg.model.spec_augment.mask_patches = 0.5\n",
"\n",
"cfg.model.train_ds.min_duration = 3.2\n",
"cfg.model.validation_ds.min_duration = 3.2\n",
"#with patch_size set to 16 and 10 patches, \n",
"#we need to be able to mask 160 time steps;\n",
"#at preprocessor stride 0.01 this means we need minimum duration of 1.6 seconds \n",
"#or more to sample only from masked steps in the same utterance\n",
"\n",
"cfg.model.loss_list.contrastive.loss.codebook_size = 32\n",
"cfg.model.loss_list.mlm.decoder.feat_out = 1024\n",
"\n",
"cfg.model.loss_list.contrastive.loss.num_negatives = 40\n",
"cfg.model.loss_list.contrastive.loss.quantizer_temp_decay = 0.999"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "TrfAb1DjWzpL"
},
"source": [
"The following parameters will be used for decoder, loss, and masking:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "DTcn93gMXQba"
},
"outputs": [],
"source": [
"print(OmegaConf.to_yaml(cfg.model.loss_list))\n",
"print(OmegaConf.to_yaml(cfg.model.spec_augment))"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "4JnepitBZ3ta"
},
"source": [
"Note that for this loss the outputs must match the inputs, so since we are using Citrinet architecture with 8x stride, we would need to either set \"combine_time_steps\" to 8, or put additional stride layers in the decoder. By default for Citrinet with 8x stride we use \"combine_time_steps=4\" and \"stride_layers=1\" to match the 8x stride.\n",
"\n",
"Since in MaskedPatchAugmentation we set mask_patches to 0.5 and our min_durations are set to 3.2, we are guaranteed to have 1.6 masked second per utterance, or 160 masked steps. Since combine_time_steps is set to 4, this means that 160 / 4 = 40 total negatives can be sampled, so we set num_negatives to 40 (unless you set sample_from_same_utterance_only to false or sample_from_non_masked to true, but this tends to make results worse).\n",
"\n",
"In the default configs we assume that min_duration for samples is higher (8 seconds by default), so there we can set patch_size to 48 for a total of 480 masked steps, and use 100 sampled negatives. If the min_duration of samples that you are training on allows, the amount of masked steps as well as negatives can be increased further (masking around 50% of the sample duration tends to work well)."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "yoUIMS7mgrUs"
},
"source": [
"Now we can create the config object."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "l0IqI0Yugqc1"
},
"outputs": [],
"source": [
"cfg = OmegaConf.to_container(cfg, resolve=True)\n",
"cfg = OmegaConf.create(cfg)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "OxxEwR05XatQ"
},
"source": [
"## Pre-training the model\n",
"\n",
"We can now create the model based on our config and start pre-training with pytorch lightning. In NeMo you can use the examples/asr/speech_pretraining/speech_pre_training.py script for this, which contains the following lines:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "WbnVpz03Jxto"
},
"outputs": [],
"source": [
"import lightning.pytorch as pl\n",
"from omegaconf import OmegaConf\n",
"\n",
"from nemo.collections.asr.models.ssl_models import SpeechEncDecSelfSupervisedModel\n",
"from nemo.core.config import hydra_runner\n",
"from nemo.utils import logging\n",
"from nemo.utils.exp_manager import exp_manager\n",
"\n",
"trainer = pl.Trainer(**cfg.trainer)\n",
"exp_manager(trainer, cfg.get(\"exp_manager\", None))\n",
"asr_model = SpeechEncDecSelfSupervisedModel(cfg=cfg.model, trainer=trainer)\n",
"\n",
"trainer.fit(asr_model)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "FylnQvwvKIy9"
},
"source": [
"# Fine-tuning with pre-trained encoder"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "RvhsJbsiKNj6"
},
"source": [
"Now that we have pre-trained our encoder, we can fine-tune our model with a supervised objective. Normally the dataset used for pre-training will be large and unlabeled, while the one used for fine-tuning will be a smaller labeled dataset. For the purpose of this tutorial we will just use the same dataset for fine-tuning."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Rk7uDmuHAD2e"
},
"source": [
"## Building tokenizer"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "9WsIHRVVADZu"
},
"outputs": [],
"source": [
"!mkdir scripts\n",
"!wget -P scripts/ https://raw.githubusercontent.com/NVIDIA/NeMo/$BRANCH/scripts/tokenizers/process_asr_text_tokenizer.py\n",
"\n",
"!python ./scripts/process_asr_text_tokenizer.py \\\n",
" --manifest=\"{data_dir}/an4/train_manifest.json\" \\\n",
" --data_root=\"{data_dir}/tokenizers/an4/\" \\\n",
" --vocab_size=128 \\\n",
" --tokenizer=\"spe\" \\\n",
" --no_lower_case \\\n",
" --spe_type=\"unigram\" \\\n",
" --log"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Na0iuLoQMfgj"
},
"source": [
"## Creating config for fine-tuning with CTC objective\n",
"\n",
"We will now create a config for a ctc-based model with a matching encoder to the model we pre-trained."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "tzF8uOWOoqeO"
},
"outputs": [],
"source": [
"config_path = data_dir + '/configs/fast-conformer_ctc_bpe.yaml'\n",
"\n",
"cfg = OmegaConf.load(config_path)\n",
"\n",
"cfg.model.encoder.n_layers = 8\n",
"cfg.model.encoder.d_model = 256\n",
"cfg.model.encoder.n_heads = 4\n",
"\n",
"cfg.model.spec_augment.freq_masks = 2\n",
"cfg.model.spec_augment.time_masks = 5\n",
"\n",
"cfg.model.train_ds.manifest_filepath = train_manifest\n",
"cfg.model.train_ds.batch_size = 16\n",
"\n",
"cfg.model.validation_ds.manifest_filepath = test_manifest\n",
"cfg.model.validation_ds.batch_size = 16\n",
"\n",
"cfg.model.log_prediction = False\n",
"\n",
"if torch.cuda.is_available():\n",
" cfg.trainer.accelerator = 'gpu'\n",
" cfg.trainer.strategy = 'auto'\n",
" cfg.trainer.devices = 1\n",
"else:\n",
" cfg.trainer.accelerator = 'cpu'\n",
" cfg.trainer.strategy = 'auto'\n",
" cfg.trainer.devices = 0\n",
"\n",
"cfg.model.tokenizer.dir = data_dir + \"/tokenizers/an4/tokenizer_spe_unigram_v128/\" # note this is a directory, not a path to a vocabulary file\n",
"cfg.model.tokenizer.type = \"bpe\"\n",
"\n",
"cfg.exp_manager.exp_dir = data_dir + \"/content/exp\"\n",
"cfg.exp_manager.name = \"fine_tuned\"\n",
"cfg.exp_manager.use_datetime_version = False\n",
"cfg.exp_manager.create_tensorboard_logger = False\n",
"cfg.exp_manager.resume_if_exists = True\n",
"cfg.exp_manager.resume_ignore_no_checkpoint = True\n",
"cfg.exp_manager.checkpoint_callback_params.save_best_model = True\n",
"\n",
"cfg.model.optim.sched.name = \"CosineAnnealing\"\n",
"cfg.model.optim.sched.warmup_steps = 500\n",
"cfg.model.optim.sched.max_steps = 1000\n",
"cfg.model.optim.sched.min_lr = 0\n",
"cfg.model.optim.lr = 0.015 #if encoder is frozen, lr can be much higher\n",
"cfg.model.optim.weight_decay = 0\n",
"\n",
"cfg.trainer.max_steps = cfg.model.optim.sched.max_steps"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "i3-_tFGDNPvZ"
},
"source": [
"In order to load our encoder, we add the following to the config. We need to set init_strict to False, since the decoder for this model will be different, in order to have the correct output for ctc."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "3GKv7ERxNPRc"
},
"outputs": [],
"source": [
"cfg.init_from_nemo_model=data_dir + \"/content/exp/pre_trained/checkpoints/pre_trained.nemo\"\n",
"cfg.init_strict = False"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "MOZmD9sjNgfX"
},
"outputs": [],
"source": [
"cfg = OmegaConf.to_container(cfg, resolve=True)\n",
"cfg = OmegaConf.create(cfg)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "InacQMmE8tPZ"
},
"source": [
"## Fine-tuning the model"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "l3eNBs_sNhjE"
},
"source": [
"Initializing our model from config:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "IEScq3t_vB6N"
},
"outputs": [],
"source": [
"from nemo.collections.asr.models.ctc_bpe_models import EncDecCTCModelBPE\n",
"\n",
"trainer = pl.Trainer(**cfg.trainer)\n",
"exp_manager(trainer, cfg.get(\"exp_manager\", None))\n",
"asr_model = EncDecCTCModelBPE(cfg=cfg.model, trainer=trainer)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "uif9DLOQ8d0b"
},
"source": [
"Loading our pre-trained checkpoint:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Jes4vqx4vLRC"
},
"outputs": [],
"source": [
"asr_model.maybe_init_from_pretrained_checkpoint(cfg)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "S5aVb2F8WuAR"
},
"source": [
"We can optionally freeze the encoder and only fine-tune the decoder during training. This can be done to lower the memory and time requirements of fine-tuning, but will likely result in a higher word error rate."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "LpF_YQUmXUR8"
},
"outputs": [],
"source": [
"#asr_model.encoder.freeze()"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "GdoIbsLn8gNC"
},
"source": [
"Now we can run fine-tuning."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "AACa0T6jJtMQ"
},
"outputs": [],
"source": [
"trainer.fit(asr_model)"
]
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"collapsed_sections": [],
"name": "Self_Supervised_Pre_Training.ipynb",
"provenance": []
},
"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": 0
}
tutorials/asr/Streaming_ASR.ipynb
+703
View File
@@ -0,0 +1,703 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "lJz6FDU1lRzc"
},
"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\nNOTE: User is responsible for checking the content of datasets and the applicable licenses and determining if suitable for the intended use.\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 matplotlib>=3.3.2\n",
"\n",
"## Install NeMo\n",
"BRANCH = 'main'\n",
"!python -m pip install \"nemo_toolkit[all] @ git+https://github.com/NVIDIA-NeMo/Speech.git@$BRANCH\"\n",
"\n",
"## Grab the config we'll use in this example\n",
"!mkdir configs\n",
"!wget -P configs/ https://raw.githubusercontent.com/NVIDIA/NeMo/$BRANCH/examples/asr/conf/config.yaml\n",
"\n",
"\"\"\"\n",
"Remember to restart the runtime for the kernel to pick up any upgraded packages (e.g. matplotlib)!\n",
"Alternatively, you can uncomment the exit() below to crash and restart the kernel, in the case\n",
"that you want to use the \"Run All Cells\" (or similar) option.\n",
"\"\"\"\n",
"# exit()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Streaming ASR\n",
"In this tutorial, we will look at one way to use one of NeMo's pretrained Conformer-CTC models for streaming inference. We will first look at some use cases where we may need streaming inference and then we will work towards developing a method for transcribing a long audio file using streaming."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "v1Jk9etFlRzf"
},
"source": [
"# Why Stream?\n",
"Streaming inference may be needed in one of the following scenarios:\n",
"* Real-time or close to real-time inference for live transcriptions\n",
"* Offline transcriptions of very long audio\n",
"\n",
"In this tutorial, we will mainly focus on streaming for handling long form audio and close to real-time inference with CTC based models. For training ASR models we usually use short segments of audio (<20s) that may be smaller chunks of a long audio that is aligned with the transcriptions and segmented into smaller chunks (see [tools/](https://github.com/NVIDIA/NeMo/tree/main/tools) for some great tools to do this). For running inference on long audio files we are restricted by the available GPU memory that dictates the maximum length of audio that can be transcribed in one inference call. We will take a look at one of the ways to overcome this restriction using NeMo's Conformer-CTC ASR model."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Conformer-CTC\n",
"Conformer-CTC models distributed with NeMo use a combination of self-attention and convolution modules to achieve the best of the two approaches, the self-attention layers can learn the global interaction while the convolutions efficiently capture the local correlations. Use of self-attention layers comes with a cost of increased memory usage at a quadratic rate with the sequence length. That means that transcribing long audio files with Conformer-CTC models needs streaming inference to break up the audio into smaller chunks. We will develop one method to do such inference through the course of this tutorial."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Data\n",
"To demonstrate transcribing a long audio file we will use utterances from the dev-clean set of the [mini Librispeech corpus](https://www.openslr.org/31/). "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# If something goes wrong during data processing, un-comment the following line to delete the cached dataset \n",
"# !rm -rf datasets/mini-dev-clean\n",
"!mkdir -p datasets/mini-dev-clean"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!python ../../scripts/dataset_processing/get_librispeech_data.py \\\n",
" --data_root \"datasets/mini-dev-clean/\" \\\n",
" --data_sets dev_clean_2"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"manifest = \"datasets/mini-dev-clean/dev_clean_2.json\""
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Let's create a long audio that is about 15 minutes long by concatenating audio from dev-clean and also create the corresponding concatenated transcript."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import json\n",
"def concat_audio(manifest_file, final_len=3600):\n",
" concat_len = 0\n",
" final_transcript = \"\"\n",
" with open(\"concat_file.txt\", \"w\") as cat_f:\n",
" while concat_len < final_len:\n",
" with open(manifest_file, \"r\") as mfst_f:\n",
" for l in mfst_f:\n",
" row = json.loads(l.strip())\n",
" if concat_len >= final_len:\n",
" break\n",
" cat_f.write(f\"file {row['audio_filepath']}\\n\")\n",
" final_transcript += (\" \" + row['text'])\n",
" concat_len += float(row['duration'])\n",
" return concat_len, final_transcript\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"new_duration, ref_transcript = concat_audio(manifest, 15*60)\n",
"\n",
"concat_audio_path = \"datasets/mini-dev-clean/concatenated_audio.wav\"\n",
"\n",
"!ffmpeg -t {new_duration} -safe 0 -f concat -i concat_file.txt -c copy -t {new_duration} {concat_audio_path} -y\n",
"print(\"Finished concatenating audio file!\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Streaming with CTC based models\n",
"Now let's try to transcribe the long audio file created above using a conformer-large model."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import torch\n",
"import nemo.collections.asr as nemo_asr\n",
"import contextlib\n",
"import gc"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
"device"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We are mainly concerned about decoding on the GPU in this tutorial. CPU decoding may be able to handle longer files but would also not be as fast as GPU decoding. Let's check if we can run transcribe() on the long audio file that we created above."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Clear up memory\n",
"torch.cuda.empty_cache()\n",
"gc.collect()\n",
"model = nemo_asr.models.EncDecCTCModelBPE.from_pretrained(\"stt_en_conformer_ctc_large\", map_location=device)\n",
"device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
"# device = 'cpu' # You can transcribe even longer samples on the CPU, though it will take much longer !\n",
"model = model.to(device)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Helper for torch amp autocast\n",
"if torch.cuda.is_available():\n",
" autocast = torch.cuda.amp.autocast\n",
"else:\n",
" @contextlib.contextmanager\n",
" def autocast():\n",
" print(\"AMP was not available, using FP32!\")\n",
" yield"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The call to transcribe() below should fail with a \"CUDA out of memory\" error when run on a GPU with 32 GB memory."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"with autocast():\n",
" transcript = model.transcribe([concat_audio_path], batch_size=1)[0]"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Clear up memory\n",
"torch.cuda.empty_cache()\n",
"gc.collect()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Buffer mechanism for streaming long audio files\n",
"One way to transcribe long audio with a Conformer-CTC model would be to split the audio into consecutive smaller chunks and running inference on each chunk. Care should be taken to have enough context for audio at either edges for accurate transcription. Let's introduce some terminology here to help us navigate the rest of this tutorial. \n",
"\n",
"* Buffer size is the length of audio on which inference is run\n",
"* Chunk size is the length of new audio that is added to the buffer.\n",
"\n",
"An audio buffer is made up of a chunk of audio with some padded audio from previous chunk. In order to make the best predictions with enough context for the beginning and end portions of the buffer, we only collect tokens for the middle portion of the buffer of length equal to the size of each chunk. \n",
"\n",
"Let's suppose that the maximum length of audio that can be transcribed with conformer-large model is 20s, then we can use 20s as the buffer size and use 15s (for example) as the chunk size, so one hour of audio is broken into 240 chunks of 15s each. Let's take a look at a few audio buffers that may be created for this audio."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# A simple iterator class to return successive chunks of samples\n",
"class AudioChunkIterator():\n",
" def __init__(self, samples, frame_len, sample_rate):\n",
" self._samples = samples\n",
" self._chunk_len = chunk_len_in_secs*sample_rate\n",
" self._start = 0\n",
" self.output=True\n",
" \n",
" def __iter__(self):\n",
" return self\n",
" \n",
" def __next__(self):\n",
" if not self.output:\n",
" raise StopIteration\n",
" last = int(self._start + self._chunk_len)\n",
" if last <= len(self._samples):\n",
" chunk = self._samples[self._start: last]\n",
" self._start = last\n",
" else:\n",
" chunk = np.zeros([int(self._chunk_len)], dtype='float32')\n",
" samp_len = len(self._samples) - self._start\n",
" chunk[0:samp_len] = self._samples[self._start:len(self._samples)]\n",
" self.output = False\n",
" \n",
" return chunk"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# a helper function for extracting samples as a numpy array from the audio file\n",
"import soundfile as sf\n",
"def get_samples(audio_file, target_sr=16000):\n",
" with sf.SoundFile(audio_file, 'r') as f:\n",
" sample_rate = f.samplerate\n",
" samples = f.read()\n",
" if sample_rate != target_sr:\n",
" samples = librosa.core.resample(samples, orig_sr=sample_rate, target_sr=target_sr)\n",
" samples = samples.transpose()\n",
" return samples\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Let's take a look at each chunk of speech that is used for decoding."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import matplotlib.pyplot as plt\n",
"samples = get_samples(concat_audio_path)\n",
"sample_rate = model.preprocessor._cfg['sample_rate'] \n",
"chunk_len_in_secs = 1 \n",
"chunk_reader = AudioChunkIterator(samples, chunk_len_in_secs, sample_rate)\n",
"count = 0\n",
"for chunk in chunk_reader:\n",
" count +=1\n",
" plt.plot(chunk)\n",
" plt.show()\n",
" if count >= 5:\n",
" break\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now, let's plot the actual buffers at each stage after a new chunk is added to the buffer. Audio buffer can be thought of as a fixed size queue with each incoming chunk added at the end of the buffer and the oldest samples removed from the beginning."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import numpy as np\n",
"context_len_in_secs = 1\n",
"\n",
"buffer_len_in_secs = chunk_len_in_secs + 2* context_len_in_secs\n",
"\n",
"buffer_len = sample_rate*buffer_len_in_secs\n",
"sampbuffer = np.zeros([buffer_len], dtype=np.float32)\n",
"\n",
"chunk_reader = AudioChunkIterator(samples, chunk_len_in_secs, sample_rate)\n",
"chunk_len = sample_rate*chunk_len_in_secs\n",
"count = 0\n",
"for chunk in chunk_reader:\n",
" count +=1\n",
" sampbuffer[:-chunk_len] = sampbuffer[chunk_len:]\n",
" sampbuffer[-chunk_len:] = chunk\n",
" plt.plot(sampbuffer)\n",
" plt.show()\n",
" if count >= 5:\n",
" break"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now that we have a method to split the long audio into smaller chunks, we can now work on transcribing the individual buffers and merging the outputs to get the transcription of the whole audio.\n",
"First, we implement some helper functions to help load the buffers into the data layer."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from nemo.core.classes import IterableDataset\n",
"\n",
"def speech_collate_fn(batch):\n",
" \"\"\"collate batch of audio sig, audio len\n",
" Args:\n",
" batch (FloatTensor, LongTensor): A tuple of tuples of signal, signal lengths.\n",
" This collate func assumes the signals are 1d torch tensors (i.e. mono audio).\n",
" \"\"\"\n",
"\n",
" _, audio_lengths = zip(*batch)\n",
"\n",
" max_audio_len = 0\n",
" has_audio = audio_lengths[0] is not None\n",
" if has_audio:\n",
" max_audio_len = max(audio_lengths).item()\n",
" \n",
" \n",
" audio_signal= []\n",
" for sig, sig_len in batch:\n",
" if has_audio:\n",
" sig_len = sig_len.item()\n",
" if sig_len < max_audio_len:\n",
" pad = (0, max_audio_len - sig_len)\n",
" sig = torch.nn.functional.pad(sig, pad)\n",
" audio_signal.append(sig)\n",
" \n",
" if has_audio:\n",
" audio_signal = torch.stack(audio_signal)\n",
" audio_lengths = torch.stack(audio_lengths)\n",
" else:\n",
" audio_signal, audio_lengths = None, None\n",
"\n",
" return audio_signal, audio_lengths\n",
"\n",
"# simple data layer to pass audio signal\n",
"class AudioBuffersDataLayer(IterableDataset):\n",
" \n",
"\n",
" def __init__(self):\n",
" super().__init__()\n",
"\n",
" \n",
" def __iter__(self):\n",
" return self\n",
" \n",
" def __next__(self):\n",
" if self._buf_count == len(self.signal) :\n",
" raise StopIteration\n",
" self._buf_count +=1\n",
" return torch.as_tensor(self.signal[self._buf_count-1], dtype=torch.float32), \\\n",
" torch.as_tensor(self.signal_shape[0], dtype=torch.int64)\n",
" \n",
" def set_signal(self, signals):\n",
" self.signal = signals\n",
" self.signal_shape = self.signal[0].shape\n",
" self._buf_count = 0\n",
"\n",
" def __len__(self):\n",
" return 1\n",
" "
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Next we implement a class that implements transcribing audio buffers and merging the tokens corresponding to a chunk of audio within each buffer. \n",
"\n",
"For each buffer, we pick tokens corresponding to one chunk length of audio. The chunk within each buffer is chosen such that there is equal left and right context available to the audio within the chunk.\n",
"\n",
"For example, if the chunk size is 1s and buffer size is 3s, we collect tokens corresponding to audio starting from 1s to 2s within each buffer. Conformer-CTC models have a model stride of 4, i.e., a token is produced for every 4 feature vectors in the time domain. MelSpectrogram features are generated once every 10 ms, so a token is produced for every 40 ms of audio.\n",
"\n",
"**Note:** The inherent assumption here is that the output tokens from the model are well aligned with corresponding audio segments. This may not always be true for models trained with CTC loss, so this method of streaming inference may not always work with CTC based models. "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from torch.utils.data import DataLoader\n",
"import math\n",
"class ChunkBufferDecoder:\n",
"\n",
" def __init__(self,asr_model, stride, chunk_len_in_secs=1, buffer_len_in_secs=3):\n",
" self.asr_model = asr_model\n",
" self.asr_model.eval()\n",
" self.data_layer = AudioBuffersDataLayer()\n",
" self.data_loader = DataLoader(self.data_layer, batch_size=1, collate_fn=speech_collate_fn)\n",
" self.buffers = []\n",
" self.all_preds = []\n",
" self.chunk_len = chunk_len_in_secs\n",
" self.buffer_len = buffer_len_in_secs\n",
" assert(chunk_len_in_secs<=buffer_len_in_secs)\n",
" \n",
" feature_stride = asr_model._cfg.preprocessor['window_stride']\n",
" self.model_stride_in_secs = feature_stride * stride\n",
" self.n_tokens_per_chunk = math.ceil(self.chunk_len / self.model_stride_in_secs)\n",
" self.blank_id = len(asr_model.decoder.vocabulary)\n",
" self.plot=False\n",
" \n",
" @torch.no_grad() \n",
" def transcribe_buffers(self, buffers, merge=True, plot=False):\n",
" self.plot = plot\n",
" self.buffers = buffers\n",
" self.data_layer.set_signal(buffers[:])\n",
" self._get_batch_preds() \n",
" return self.decode_final(merge)\n",
" \n",
" def _get_batch_preds(self):\n",
"\n",
" device = self.asr_model.device\n",
" for batch in iter(self.data_loader):\n",
"\n",
" audio_signal, audio_signal_len = batch\n",
"\n",
" audio_signal, audio_signal_len = audio_signal.to(device), audio_signal_len.to(device)\n",
" log_probs, encoded_len, predictions = self.asr_model(input_signal=audio_signal, input_signal_length=audio_signal_len)\n",
" preds = torch.unbind(predictions)\n",
" for pred in preds:\n",
" self.all_preds.append(pred.cpu().numpy())\n",
" \n",
" def decode_final(self, merge=True, extra=0):\n",
" self.unmerged = []\n",
" self.toks_unmerged = []\n",
" # index for the first token corresponding to a chunk of audio would be len(decoded) - 1 - delay\n",
" delay = math.ceil((self.chunk_len + (self.buffer_len - self.chunk_len) / 2) / self.model_stride_in_secs)\n",
"\n",
" decoded_frames = []\n",
" all_toks = []\n",
" for pred in self.all_preds:\n",
" ids, toks = self._greedy_decoder(pred, self.asr_model.tokenizer)\n",
" decoded_frames.append(ids)\n",
" all_toks.append(toks)\n",
"\n",
" for decoded in decoded_frames:\n",
" self.unmerged += decoded[len(decoded) - 1 - delay:len(decoded) - 1 - delay + self.n_tokens_per_chunk]\n",
" if self.plot:\n",
" for i, tok in enumerate(all_toks):\n",
" plt.plot(self.buffers[i])\n",
" plt.show()\n",
" print(\"\\nGreedy labels collected from this buffer\")\n",
" print(tok[len(tok) - 1 - delay:len(tok) - 1 - delay + self.n_tokens_per_chunk]) \n",
" self.toks_unmerged += tok[len(tok) - 1 - delay:len(tok) - 1 - delay + self.n_tokens_per_chunk]\n",
" print(\"\\nTokens collected from successive buffers before CTC merge\")\n",
" print(self.toks_unmerged)\n",
"\n",
"\n",
" if not merge:\n",
" return self.unmerged\n",
" return self.greedy_merge(self.unmerged)\n",
" \n",
" \n",
" def _greedy_decoder(self, preds, tokenizer):\n",
" s = []\n",
" ids = []\n",
" for i in range(preds.shape[0]):\n",
" if preds[i] == self.blank_id:\n",
" s.append(\"_\")\n",
" else:\n",
" pred = preds[i]\n",
" s.append(tokenizer.ids_to_tokens([pred.item()])[0])\n",
" ids.append(preds[i])\n",
" return ids, s\n",
" \n",
" def greedy_merge(self, preds):\n",
" decoded_prediction = []\n",
" previous = self.blank_id\n",
" for p in preds:\n",
" if (p != previous or previous == self.blank_id) and p != self.blank_id:\n",
" decoded_prediction.append(p.item())\n",
" previous = p\n",
" hypothesis = self.asr_model.tokenizer.ids_to_text(decoded_prediction)\n",
" return hypothesis\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"To see how this chunk based decoder comes together, let's call the decoder with a few buffers we create from our long audio file.\n",
"Some interesting experiments to try would be to see how changing sizes of the chunk and the context affects transcription accuracy. "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"\n",
"chunk_len_in_secs = 4\n",
"context_len_in_secs = 2\n",
"\n",
"buffer_len_in_secs = chunk_len_in_secs + 2* context_len_in_secs\n",
"\n",
"n_buffers = 5\n",
"\n",
"buffer_len = sample_rate*buffer_len_in_secs\n",
"sampbuffer = np.zeros([buffer_len], dtype=np.float32)\n",
"\n",
"chunk_reader = AudioChunkIterator(samples, chunk_len_in_secs, sample_rate)\n",
"chunk_len = sample_rate*chunk_len_in_secs\n",
"count = 0\n",
"buffer_list = []\n",
"for chunk in chunk_reader:\n",
" count +=1\n",
" sampbuffer[:-chunk_len] = sampbuffer[chunk_len:]\n",
" sampbuffer[-chunk_len:] = chunk\n",
" buffer_list.append(np.array(sampbuffer))\n",
" \n",
" if count >= n_buffers:\n",
" break\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"stride = 4 # 8 for Citrinet\n",
"asr_decoder = ChunkBufferDecoder(model, stride=stride, chunk_len_in_secs=chunk_len_in_secs, buffer_len_in_secs=buffer_len_in_secs )\n",
"transcription = asr_decoder.transcribe_buffers(buffer_list, plot=True)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Final transcription after CTC merge\n",
"print(transcription)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Time to evaluate our streaming inference on the whole long file that we created."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# WER calculation\n",
"from nemo.collections.asr.metrics.wer import word_error_rate\n",
"# Collect all buffers from the audio file\n",
"sampbuffer = np.zeros([buffer_len], dtype=np.float32)\n",
"\n",
"chunk_reader = AudioChunkIterator(samples, chunk_len_in_secs, sample_rate)\n",
"buffer_list = []\n",
"for chunk in chunk_reader:\n",
" sampbuffer[:-chunk_len] = sampbuffer[chunk_len:]\n",
" sampbuffer[-chunk_len:] = chunk\n",
" buffer_list.append(np.array(sampbuffer))\n",
"\n",
"asr_decoder = ChunkBufferDecoder(model, stride=stride, chunk_len_in_secs=chunk_len_in_secs, buffer_len_in_secs=buffer_len_in_secs )\n",
"transcription = asr_decoder.transcribe_buffers(buffer_list, plot=False)\n",
"wer = word_error_rate(hypotheses=[transcription], references=[ref_transcript])\n",
"\n",
"print(f\"WER: {round(wer*100,2)}%\")"
]
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"collapsed_sections": [],
"name": "ASR_with_NeMo.ipynb",
"provenance": [],
"toc_visible": true
},
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.9.7"
},
"pycharm": {
"stem_cell": {
"cell_type": "raw",
"metadata": {
"collapsed": false
},
"source": []
}
}
},
"nbformat": 4,
"nbformat_minor": 4
}
tutorials/asr/Streaming_ASR_Pipelines.ipynb
+1485
View File
File diff suppressed because it is too large Load Diff
tutorials/asr/Streaming_Multitalker_ASR.ipynb
+799
View File
@@ -0,0 +1,799 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "lJz6FDU1lRzc"
},
"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 Multitalker ASR"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "v1Jk9etFlRzf"
},
"source": [
"## Streaming Multitalker ASR with Self-Speaker Adaptation\n",
"\n",
"This tutorial shows you how to use NeMo's streaming multitalker ASR system based on the approach described in [(Wang et al., 2025)](https://arxiv.org/abs/2506.22646). This system transcribes each speaker separately in multispeaker audio using speaker activity information from a streaming diarization model.\n",
"\n",
"### How This Approach Works\n",
"\n",
"The streaming multitalker Parakeet model uses **self-speaker adaptation**, which means:\n",
"\n",
"1. **No Speaker Enrollment Required**: You only need speaker activity predictions from a diarization model (like Streaming Sortformer)\n",
"2. **Speaker Kernel Injection**: The model injects speaker-specific kernels into encoder layers to focus on each target speaker\n",
"3. **Multi-Instance Architecture**: You run one model instance per speaker, and each instance processes the same audio\n",
"4. **Handles Overlapping Speech**: Each instance focuses on one speaker, so it can transcribe overlapped speech segments\n",
"\n",
"### Cache-Aware Streaming\n",
"\n",
"The model uses stateful cache-based inference [(Noroozi et al., 2023)](https://arxiv.org/abs/2312.17279) for streaming:\n",
"- Left and right contexts in the encoder are constrained for low latency\n",
"- An activation caching mechanism enables the encoder to operate autoregressively during inference\n",
"- The model maintains consistent behavior between training and inference"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Multi-Instance Architecture Overview\n",
"\n",
"The streaming multitalker Parakeet model employs a **multi-instance approach** where one model instance is deployed per speaker:\n",
"\n",
"<img src=\"images/multi_instance.png\" alt=\"Multi-instance inference of Multitalker Parakeet model\" style=\"width: 800px;\"/>\n",
"\n",
"Each model instance:\n",
"- Receives the same mixed audio input\n",
"- Injects **speaker-specific kernels** generated from diarization-based speaker activity\n",
"- Produces transcription output specific to its target speaker\n",
"- Operates independently and can run in parallel with other instances\n",
"\n",
"### Speaker Kernel Injection Mechanism\n",
"\n",
"Learnable speaker kernels are injected into selected layers of the Fast-Conformer encoder:\n",
"\n",
"<img src=\"images/speaker_injection.png\" alt=\"Speaker Kernel Injection Mechanism\" style=\"width: 800px;\"/>\n",
"\n",
"The speaker kernels are generated through speaker supervision activations that detect speech activity for each target speaker from the streaming diarization output. This enables the encoder states to become more responsive to the targeted speaker's speech characteristics."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Setup and Data Preparation\n",
"\n",
"In this tutorial, we will demonstrate streaming multitalker ASR using:\n",
"1. **Streaming Sortformer** for real-time speaker diarization\n",
"2. **Streaming Multitalker Parakeet** for speaker-wise ASR\n",
"\n",
"Let's start by downloading a toy example audio file with multiple speakers. "
]
},
{
"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 visualize the waveform and listen to the audio. You'll notice that there are two speakers in this 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('Multispeaker Audio Waveform', 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()\n",
"plt.xticks(a, a/sr)\n",
"\n",
"IPython.display.Audio(signal, rate=sr)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Step 1: Streaming Speaker Diarization\n",
"\n",
"Now that we have a multispeaker audio file, the first step is to perform streaming speaker diarization using **Streaming Sortformer**. This will provide us with speaker activity information needed for self-speaker adaptation.\n",
"\n",
"### Download Streaming Sortformer Diarization Model\n",
"\n",
"To download the streaming Sortformer diarizer from [HuggingFace](https://huggingface.co/nvidia), you need a [HuggingFace Access Token](https://huggingface.co/docs/hub/en/security-tokens). \n",
"\n",
"Alternatively, you can download the `.nemo` file from [Streaming Sortformer HuggingFace model card](https://huggingface.co/nvidia/diar_streaming_sortformer_4spk-v2) and specify the file path."
]
},
{
"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 download \".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": [
"### Configure Streaming Parameters for Sortformer\n",
"\n",
"Set streaming parameters (all measured in 80ms frames):\n",
"- `chunk_len`: Number of frames in a processing chunk\n",
"- `chunk_right_context`: Right context length (determines latency with `chunk_len`)\n",
"- `fifo_len`: Number of previous frames from FIFO queue\n",
"- `spkcache_update_period`: Frames extracted from FIFO for speaker cache update\n",
"- `spkcache_len`: Total frames in speaker cache\n",
"\n",
"The input buffer latency is determined by `chunk_len` + `chunk_right_context`."
]
},
{
"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\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 and Streaming Diarization\n",
"\n",
"Extract log-mel features from the audio signal and prepare for streaming diarization:"
]
},
{
"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 Streaming Diarization Loop\n",
"\n",
"Initialize the streaming state and run the streaming diarization to get speaker activity predictions:"
]
},
{
"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",
"print(f\"Processing {num_chunks} chunks for diarization...\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Run streaming diarization\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 Diarization\",\n",
" disable=False,\n",
"):\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",
"\n",
"print(f\"Diarization complete! Total predictions shape: {total_preds.shape}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Visualize Diarization Results\n",
"\n",
"Let's visualize the speaker diarization predictions:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def plot_diarout(preds):\n",
" \"\"\"Visualize diarization predictions\"\"\"\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('Diarization Predictions (Speaker Activity)', 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",
" plt.show()\n",
"\n",
"plot_diarout(total_preds[0,:])"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Step 2: Streaming Multitalker ASR with Self-Speaker Adaptation\n",
"\n",
"Now that we have speaker activity information from diarization, we can load the streaming multitalker Parakeet model and perform speaker-wise ASR.\n",
"\n",
"### Download Streaming Multitalker Parakeet Model"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from nemo.collections.asr.models import ASRModel\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",
" asr_model = ASRModel.from_pretrained(\"nvidia/multitalker-parakeet-streaming-0.6b-v1\")\n",
"else:\n",
" # You can download \".nemo\" file from https://huggingface.co/nvidia/multitalker-parakeet-streaming-0.6b-v1 and specify the path.\n",
" asr_model = ASRModel.restore_from(restore_path=\"/path/to/multitalker-parakeet-streaming-0.6b-v1.nemo\", map_location=torch.device('cuda'))\n",
"diar_model.eval()\n",
"\n",
"asr_model.eval()\n",
"if torch.cuda.is_available():\n",
" asr_model.to(torch.device(\"cuda\"))\n",
"\n",
"print(\"ASR Model loaded successfully!\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Configure Cache-Aware Streaming Parameters\n",
"\n",
"Set the streaming attention context size for the ASR model. The latency is determined by the attention context configuration (measured in 80ms frames):\n",
"\n",
"- `[70, 0]`: Chunk size = 1 (1 * 80ms = 0.08s) \n",
"- `[70, 1]`: Chunk size = 2 (2 * 80ms = 0.16s) \n",
"- `[70, 6]`: Chunk size = 7 (7 * 80ms = 0.56s) \n",
"- `[70, 13]`: Chunk size = 14 (14 * 80ms = 1.12s)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Set streaming parameters for 1.12s latency (matching diarization)\n",
"att_context_size = [70, 13] # [left_context, right_context] in frames\n",
"\n",
"\n",
"print(f\"ASR streaming configured with attention context: {att_context_size}\")\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Determine Number of Active Speakers\n",
"\n",
"Analyze the diarization output to determine how many speakers are active in the audio:"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Step 3: Multi-Instance Streaming ASR\n",
"\n",
"Now we'll run streaming multitalker ASR using the multi-instance architecture. We'll create one model instance per detected speaker.\n",
"\n",
"### Step 3-1: Prepare the configurations\n",
"\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Since streaming processing of speech signals involves many complications in cache handling, we first need to set up a config dataclass that aggregates all the parameters in one place. You can access this class in the example multitalker streaming ASR script: [speech_to_text_multitalker_streaming_infer.py](https://raw.githubusercontent.com/NVIDIA-NeMo/NeMo/main/examples/asr/asr_cache_aware_streaming/speech_to_text_multitalker_streaming_infer.py\") "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from dataclasses import dataclass, field\n",
"from typing import List, Optional\n",
"\n",
"@dataclass\n",
"class MultitalkerTranscriptionConfig:\n",
" \"\"\"\n",
" Configuration for Multi-talker transcription with an ASR model and a diarization model.\n",
" \"\"\"\n",
" # Required configs\n",
" diar_model: Optional[str] = None\n",
" diar_pretrained_name: Optional[str] = None\n",
" max_num_of_spks: Optional[int] = 4\n",
" parallel_speaker_strategy: bool = True\n",
" masked_asr: bool = True\n",
" mask_preencode: bool = False\n",
" cache_gating: bool = True\n",
" cache_gating_buffer_size: int = 2\n",
" single_speaker_mode: bool = False\n",
" feat_len_sec: float = 0.01\n",
"\n",
" # General configs\n",
" session_len_sec: float = -1\n",
" num_workers: int = 8\n",
" random_seed: Optional[int] = None\n",
" log: bool = True\n",
"\n",
" # Streaming diarization configs\n",
" streaming_mode: bool = True\n",
" spkcache_len: int = 188\n",
" spkcache_refresh_rate: int = 0\n",
" fifo_len: int = 188\n",
" chunk_len: int = 0\n",
" chunk_left_context: int = 0\n",
" chunk_right_context: int = 0\n",
"\n",
" # If `cuda` is a negative number, inference will be on CPU only.\n",
" cuda: Optional[int] = None\n",
" allow_mps: bool = False\n",
" matmul_precision: str = \"highest\" # Literal[\"highest\", \"high\", \"medium\"]\n",
"\n",
" # ASR Configs\n",
" asr_model: Optional[str] = None\n",
" device: str = 'cuda'\n",
" audio_file: Optional[str] = None\n",
" manifest_file: Optional[str] = None\n",
" att_context_size: Optional[List[int]] = field(default_factory=lambda: [70, 13])\n",
" use_amp: bool = True\n",
" debug_mode: bool = False\n",
" deploy_mode: bool = False\n",
" batch_size: int = 32\n",
" chunk_size: int = -1\n",
" shift_size: int = -1\n",
" left_chunks: int = 2\n",
" online_normalization: bool = False\n",
" output_path: Optional[str] = None\n",
" pad_and_drop_preencoded: bool = False\n",
" set_decoder: Optional[str] = None # [\"ctc\", \"rnnt\"]\n",
" generate_realtime_scripts: bool = False\n",
" spk_supervision: str = \"diar\" # [\"diar\", \"rttm\"]\n",
" binary_diar_preds: bool = False\n",
"\n",
" # Multitalker transcription configs\n",
" verbose: bool = False\n",
" word_window: int = 50\n",
" sent_break_sec: float = 30.0\n",
" fix_prev_words_count: int = 5\n",
" update_prev_words_sentence: int = 5\n",
" left_frame_shift: int = -1\n",
" right_frame_shift: int = 0\n",
" min_sigmoid_val: float = 1e-2\n",
" discarded_frames: int = 8\n",
" print_time: bool = True\n",
" print_sample_indices: List[int] = field(default_factory=lambda: [0])\n",
" colored_text: bool = True\n",
" real_time_mode: bool = False\n",
" print_path: Optional[str] = None\n",
" ignored_initial_frame_steps: int = 5\n",
" finetune_realtime_ratio: float = 0.01"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Using this dataclass for configurations, assign the designated streaming speaker diarization parameters to a diarization model. This process ensures that the processing window size and cache sizes are synchronized."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from omegaconf import OmegaConf\n",
"# Create configuration object for multitalker transcription\n",
"cfg = MultitalkerTranscriptionConfig()\n",
"# Convert dataclass to OmegaConf DictConfig so it has .get() method\n",
"cfg = OmegaConf.structured(cfg)\n",
"\n",
"cfg.att_context_size = [70,13]\n",
"cfg.output_path = \"/path/to/output.json\"\n",
"cfg.audio_file = an4_audio\n",
"print(f\"ASR streaming configured with attention context: {att_context_size}\")\n",
"\n",
"for key in cfg:\n",
" cfg[key] = None if cfg[key] == 'None' else cfg[key]\n",
"\n",
"# Set streaming mode diar_model params (matching the diarization setup from lines 263-271 of reference file)\n",
"diar_model.streaming_mode = cfg.streaming_mode\n",
"diar_model.sortformer_modules.chunk_len = cfg.chunk_len if cfg.chunk_len > 0 else 6\n",
"diar_model.sortformer_modules.spkcache_len = cfg.spkcache_len\n",
"diar_model.sortformer_modules.chunk_left_context = cfg.chunk_left_context\n",
"diar_model.sortformer_modules.chunk_right_context = cfg.chunk_right_context if cfg.chunk_right_context > 0 else 7\n",
"diar_model.sortformer_modules.fifo_len = cfg.fifo_len\n",
"diar_model.sortformer_modules.log = cfg.log\n",
"diar_model.sortformer_modules.spkcache_refresh_rate = cfg.spkcache_refresh_rate\n",
"\n",
"# Set online normalization flag\n",
"online_normalization = cfg.online_normalization\n",
"\n",
"# Set pad_and_drop_preencoded flag\n",
"pad_and_drop_preencoded = cfg.pad_and_drop_preencoded\n",
"\n",
"print(f\"Configuration setup complete!\")\n",
"print(f\"Audio file: {cfg.audio_file}\")\n",
"print(f\"Streaming mode: {diar_model.streaming_mode}\")\n",
"print(f\"Diar model chunk_len: {diar_model.sortformer_modules.chunk_len}\")\n",
"print(f\"Diar model chunk_right_context: {diar_model.sortformer_modules.chunk_right_context}\")\n",
"print(f\"Online normalization: {online_normalization}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Run Multi-Instance Streaming ASR\n",
"\n",
"For each active speaker, we'll run a separate ASR model instance with speaker-specific kernel injection. In practice, these instances can run in parallel."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"For testing purposes, first set up a single-sample batch and feed it into the streaming buffer. The streaming buffer simulates an input audio stream so that we can run the multitalker ASR model in a streaming manner."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from nemo.collections.asr.parts.utils.streaming_utils import CacheAwareStreamingAudioBuffer\n",
"\n",
"samples = [\n",
" {\n",
" 'audio_filepath': cfg.audio_file,\n",
" }\n",
"]\n",
"streaming_buffer = CacheAwareStreamingAudioBuffer(\n",
" model=asr_model,\n",
" online_normalization=online_normalization,\n",
" pad_and_drop_preencoded=cfg.pad_and_drop_preencoded,\n",
")\n",
"streaming_buffer.append_audio_file(audio_filepath=cfg.audio_file, stream_id=-1)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Step 3-2: Run Multitalker ASR with the prepared configurations and streaming buffer\n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from nemo.collections.asr.parts.utils.multispk_transcribe_utils import SpeakerTaggedASR\n",
"from pprint import pprint\n",
"speaker_transcriptions = {}\n",
"\n",
"streaming_buffer_iter = iter(streaming_buffer)\n",
"multispk_asr_streamer = SpeakerTaggedASR(cfg, asr_model, diar_model)\n",
"feat_frame_count = 0\n",
"print(f\"length of streaming_buffer_iter: {len(streaming_buffer)} {time.time()}\")\n",
"\n",
"for step_num, (chunk_audio, chunk_lengths) in enumerate(streaming_buffer_iter):\n",
" drop_extra_pre_encoded = (\n",
" 0\n",
" if step_num == 0 and not pad_and_drop_preencoded\n",
" else asr_model.encoder.streaming_cfg.drop_extra_pre_encoded\n",
" )\n",
" loop_start_time = time.time()\n",
" with torch.inference_mode():\n",
" with autocast:\n",
" with torch.no_grad():\n",
" multispk_asr_streamer.perform_parallel_streaming_stt_spk(\n",
" step_num=step_num,\n",
" chunk_audio=chunk_audio,\n",
" chunk_lengths=chunk_lengths,\n",
" is_buffer_empty=streaming_buffer.is_buffer_empty(),\n",
" drop_extra_pre_encoded=drop_extra_pre_encoded,\n",
" )\n",
" pprint(multispk_asr_streamer.instance_manager.batch_asr_states[0].seglsts)\n",
"\n",
"seglst_dict_list = multispk_asr_streamer.generate_seglst_dicts_from_parallel_streaming(samples=samples)\n",
"\n",
"from pprint import pprint\n",
"print(f\"SegLST style multispeaker transcription\\n {seglst_dict_list}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Results and Analysis\n",
"\n",
"### Display Speaker-Wise Transcriptions\n",
"\n",
"Let's display the final transcriptions for each speaker: "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"print(\"\\n\" + \"=\"*80)\n",
"print(\"Final transcriptions with speaker tagging and timestamps\")\n",
"print(\"=\"*80 + \"\\n\")\n",
"\n",
"# Display SegLST dict list with timestamps and speaker-tagged transcriptions\n",
"# Format: {'speaker': 'speaker_0', 'start_time': 2.64, 'end_time': 5.6, 'words': 'eleven twenty seven fifty seven', 'session_id': 'an4_diarize_test'}\n",
"if seglst_dict_list:\n",
" for idx, seglst in enumerate(seglst_dict_list):\n",
" speaker = seglst.get('speaker', 'Unknown')\n",
" start_time = seglst.get('start_time', 0.0)\n",
" end_time = seglst.get('end_time', 0.0)\n",
" words = seglst.get('words', '')\n",
" session_id = seglst.get('session_id', '')\n",
"\n",
" print(f\"[{idx+1}] {speaker} ({start_time:.2f}s - {end_time:.2f}s): {words}\")\n",
"\n",
" print(f\"\\n{'-'*80}\")\n",
" print(f\"Total segments: {len(seglst_dict_list)}\")\n",
"else:\n",
" print(\"No transcriptions available in seglst_dict_list.\")\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"This tutorial demonstrated streaming multitalker ASR using self-speaker adaptation with NeMo:\n",
"\n",
"### Key Components\n",
"\n",
"1. **Streaming Sortformer Diarization**: Provides real-time speaker activity predictions using arrival-order speaker cache (AOSC)\n",
"2. **Cache-Aware Streaming ASR**: FastConformer-based model with stateful cache-based inference for low-latency transcription\n",
"3. **Self-Speaker Adaptation**: Speaker kernels injected into encoder layers based on diarization, enabling speaker-focused recognition without enrollment\n",
"4. **Multi-Instance Architecture**: One model instance per speaker, enabling parallel processing and handling of severe speech overlap\n",
"\n",
"### Advantages\n",
"\n",
"- **No enrollment required**: Only needs diarization output, no pre-recorded speaker audio\n",
"- **Handles overlap**: Each instance focuses on one speaker, even during fully overlapped speech\n",
"- **Streaming capable**: Real-time processing with configurable latency (0.08s to 1.12s+)\n",
"- **State-of-the-art performance**: Achieves strong results on challenging multitalker benchmarks\n",
"\n",
"### Configuration Summary\n",
"\n",
"In this tutorial, we used:\n",
"- **Diarization latency**: 1.04s (chunk_len=6, chunk_right_context=7)\n",
"- **ASR latency**: 1.12s (att_context_size=[70, 13])\n",
"- **Number of speakers**: Automatically detected from diarization output"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## References\n",
"\n",
"[1] [Speaker Targeting via Self-Speaker Adaptation for Multi-talker ASR](https://arxiv.org/abs/2506.22646) \n",
"\n",
"\n",
"[2] [Stateful Conformer with Cache-based Inference for Streaming Automatic Speech Recognition](https://arxiv.org/abs/2312.17279) \n",
"\n",
"[3] [Streaming Sortformer: Speaker Cache-Based Online Speaker Diarization with Arrival-Time Ordering](https://arxiv.org/abs/2507.18446)\n",
"\n",
"[4] [NEST: Self-supervised Fast Conformer as All-purpose Seasoning to Speech Processing Tasks](https://arxiv.org/abs/2408.13106)\n",
"\n",
"[5] [Fast Conformer with Linearly Scalable Attention for Efficient Speech Recognition](https://arxiv.org/abs/2305.05084)"
]
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"collapsed_sections": [],
"name": "ASR_with_NeMo.ipynb",
"provenance": [],
"toc_visible": true
},
"kernelspec": {
"display_name": "nemo093025",
"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/asr/Transducers_with_HF_Datasets.ipynb
+696
View File
@@ -0,0 +1,696 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "SUOXg71A3w78"
},
"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",
"\n",
"NOTE: User is responsible for checking the content of datasets and the applicable licenses and determining if suitable for the intended use.\n",
"\"\"\"\n",
"# If you're using Google Colab and not running locally, run this cell.\n",
"import os\n",
"\n",
"# Allow loading trusted checkpoints that hold non-tensor objects (torch>=2.6 weights_only default).\n",
"os.environ[\"TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD\"] = \"1\"\n",
"\n",
"# # Install dependencies\n",
"!pip install wget\n",
"!apt-get install sox libsndfile1 ffmpeg\n",
"!pip install text-unidecode\n",
"!pip install matplotlib>=3.3.2\n",
"\n",
"## Install NeMo\n",
"BRANCH = 'main'\n",
"!python -m pip install \"nemo_toolkit[all] @ git+https://github.com/NVIDIA-NeMo/Speech.git@$BRANCH\"\n",
"\n",
"!pip install --upgrade datasets==3.4.0 # downgrading to 3.4.0 because latest version (4.0.0) has some issues\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "RGNuJWr66C38"
},
"source": [
"# Automatic Speech Recognition with Transducer Models using HF Datasets\n",
"\n",
"We have discussed training various ASR models in NeMo using custom datasets, either for fine-tuning or for scratch-training. In this tutorial, we will showcase how to use Hugging Face datasets library in order to finetune a Transducer ASR model on a small dataset from for the Telugu language. \n",
"It includes discussions relevant to preparing datasets with HF and how to use them to finetune NeMo models. The same method applies to training from scratch. However, for training, we recommend using scripts directly from the `examples/asr/` folder.\n",
"\n",
"In this tutorial, we demonstrate the usage of HF datasets for the Telugu language, where we use the Fluers dataset for training, validation, and testing. However, the same procedure can be used for other languages or domains and finetuned for specific use cases accordingly. \n",
"\n",
"For scripts, refer to [speech_to_text_finetune.py]('https://github.com/NVIDIA/NeMo/blob/main/examples/asr/speech_to_text_finetune.py') for training from scratch. \n",
"\n",
"--------\n",
"\n",
"**Note**: It is assumed that the previous tutorial - \"Intro-to-Transducers\" has been reviewed, and there is some familiarity with the config components of transducer models."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "eAqCcJ-T6C6k"
},
"source": [
"# Preparing the dataset\n",
"\n",
"In this tutorial, we will be utilizing the `Telugu` language dataset part of fluers set. More details about dataset can be found at: https://huggingface.co/datasets/google/fleurs/viewer/te_in \n",
"\n",
"Dataset consists of 3 splits: `train`, `validation`, and `test`. We will use the `train` and `validation` splits for training and validation, respectively. The `test` split will be used for testing the model. Train set consists of 2,300 utterances and validation set consists of 311 examples. \n",
"\n",
"You may listen to some of the samples from the dataset using the following code:\n",
"\n",
"```python\n",
"\n",
"```"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import datasets\n",
"import IPython.display as ipd\n",
"# sample from the validation set\n",
"sample = datasets.load_dataset('google/fleurs', 'te_in', split='validation')[0]\n",
" \n",
"# Let's listen to the audio file and its transcription\n",
"ipd.Audio(sample['audio']['array'], rate=16000)\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Lets look at the transcription\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"print(\"Transcription:\", sample['transcription'])"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Audio samples are in the `audio` field and the corresponding transcriptions are in the `transcription` field. We will use these fields to prepare data and train our model.\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Preparing tokenizer"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Since we are finetuning Parakeet model, which is an English language model, we need to update the tokenizer and update the decoder to support the new language. \n",
"\n",
"First, we will extract text transcriptions from the dataset and use them to train a tokenizer. We will use the scripts from NeMo first to get the data from HF dataset using `get_hf_dataset.py` script. Next we use `process_asr_text_tokenizer.py` script to prepare the tokenizer from [scripts](https://github.com/NVIDIA/NeMo/tree/main/scripts/tokenizers) folder. \n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Download the `get_hf_text_data.py` script from the [scripts](https://github.com/NVIDIA/NeMo/blob/main/scripts/tokenizers) folder and run the following command to get the data from HF dataset. "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "i2hD4LkoJvrx"
},
"outputs": [],
"source": [
"if not os.path.exists(\"scripts/get_hf_text_data.py\"):\n",
" !wget -P scripts/ https://raw.githubusercontent.com/NVIDIA/NeMo/$BRANCH/scripts/tokenizers/get_hf_text_data.py"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Major difference from NeMo dataset configs to training using HF-datasets is for using hf-datasets, users need to mention the hf-dataset information through hf data config and pass to the script for downloading necessary data. Users can switch to another dataset by changing necessary fields in the hf data config. \n",
"Let's create that config here."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from omegaconf import OmegaConf, open_dict\n",
"train_split = {\n",
" 'path': 'google/fleurs',\n",
" 'name': 'te_in',\n",
" 'split': 'train', # we will update this accordingly based on the split\n",
" 'streaming': False}\n",
"print(OmegaConf.to_yaml(train_split))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Since we need clean data for training tokenizer and models, we need to filter the data based on how dataset was constructed and how we would like the ASR model output to be. \n",
"\n",
"Based on prior analysis of text transcripts of the current hf dataset, we skip all non-alphanumeric characters except `full-stop` using `normalize_text` option of the `get_hf_text_data.py` script, based on `huggingface_data_tokenizer.yaml` config file. "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"\n",
"if not os.path.exists('configs/huggingface_data_tokenizer.yaml'):\n",
" !wget -P configs/ https://raw.githubusercontent.com/NVIDIA/NeMo/$BRANCH/scripts/tokenizers/conf/huggingface_data_tokenizer.yaml\n",
"\n",
"\n",
"!export HYDRA_FULL_ERROR=1;python scripts/get_hf_text_data.py \\\n",
" --config-path=\"../configs\" \\\n",
" --config-name=\"huggingface_data_tokenizer\" \\\n",
" normalize_text=True \\\n",
" symbols_to_keep=[\".\"] \\\n",
" text_key=\"transcription\" \\\n",
" output_file='telugu_train_corpus.txt' \\\n",
" +hf_data_cfg='[{train_split}]'"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"From the above command, we were able to use the `huggingface_data_tokenizer.yaml` config file to download the data from HF dataset. The download data is saved to `telugu_train_corpus.txt` file, which we will use to train the tokenizer. Before that, let's look at some utterances from the normalized (preprocessed) text file."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"with open('telugu_train_corpus.txt', 'r') as f:\n",
" for i in range(5):\n",
" print(f.readline().strip())"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"As can be seen, there are characters from English and also numerical characters as well. For this tutorial we will retain those characters and prepare the tokenizer. But if you would like to filter out those characters, you can use your own function in the `get_hf_text_data.py` script to filter out the characters. "
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now, we will use the `process_asr_text_tokenizer.py` script to prepare the tokenizer, for finetuning we can limit the vocab size to 256 for telugu language."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"\n",
"if not os.path.exists(\"scripts/process_asr_text_tokenizer.py\"):\n",
" !wget -P scripts/ https://raw.githubusercontent.com/NVIDIA/NeMo/$BRANCH/scripts/tokenizers/process_asr_text_tokenizer.py\n",
"\n",
"# Now this downloads the text corpus of data to tokenizers script\n",
"VOCAB_SIZE = 256 # can be any value above 29\n",
"TOKENIZER_TYPE = \"spe\" # can be wpe or spe\n",
"SPE_TYPE = \"bpe\" # can be bpe or unigram\n",
"\n",
"# ------------------------------------------------------------------- #\n",
"!rm -r tokenizers/\n",
"\n",
"if not os.path.exists(\"tokenizers\"):\n",
" os.makedirs(\"tokenizers\")\n",
"\n",
"!python scripts/process_asr_text_tokenizer.py \\\n",
" --data_file='telugu_train_corpus.txt' \\\n",
" --data_root=\"tokenizers\" \\\n",
" --tokenizer=$TOKENIZER_TYPE \\\n",
" --spe_type=$SPE_TYPE \\\n",
" --no_lower_case \\\n",
" --log \\\n",
" --vocab_size=$VOCAB_SIZE"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "JHDZswN6LIBJ"
},
"outputs": [],
"source": [
"# Tokenizer path\n",
"if TOKENIZER_TYPE == 'spe':\n",
" TOKENIZER = os.path.join(\"tokenizers\", f\"tokenizer_spe_{SPE_TYPE}_v{VOCAB_SIZE}\")\n",
" TOKENIZER_TYPE_CFG = \"bpe\"\n",
"else:\n",
" TOKENIZER = os.path.join(\"tokenizers\", f\"tokenizer_wpe_v{VOCAB_SIZE}\")\n",
" TOKENIZER_TYPE_CFG = \"wpe\""
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"print(\"Tokenizers are saved at:\", TOKENIZER)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "j8tx2m2w6C87"
},
"source": [
"# Preparing a Transducer Model\n",
"\n",
"Now that we have the dataset and tokenizer prepared, let us begin by setting up the config of the Transducer model! In this tutorial, we will finetune [Parakeet RNNT 0.6B](https://huggingface.co/nvidia/parakeet-rnnt-0.6b) model, but the same procedure can be used for other RNNT models as well.\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "pQETBCSML0Us"
},
"source": [
"## Prepare the config \n",
"For finetuning the model, we need to update the config file to include the tokenizer and the dataset information. We will use the `parakeet-rnnt-6b` model and update the config file to include the tokenizer and the dataset information. For this we use `speech_to_text_hf_finetune.yaml` config file, and training script `speech_to_text_finetune.py` from the `examples/asr` folder.\n",
"\n",
"For this demo, we will replicate only a portion of the script - in order to show just the necessary components for training the model on single GPU. However, we recommend users to use the scripts directly from the `examples/asr` folder for training the model on multiple GPUs."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "6a_vedo0Lyo8"
},
"outputs": [],
"source": [
"## Grab the config we'll use in this example\n",
"!mkdir -p configs\n",
"if not os.path.exists('configs/speech_to_text_hf_finetune.yaml'):\n",
" !wget -P configs/ https://raw.githubusercontent.com/NVIDIA/NeMo/$BRANCH/examples/asr/conf/asr_finetune/speech_to_text_hf_finetune.yaml\n",
"\n",
"config = OmegaConf.load(\"configs/speech_to_text_hf_finetune.yaml\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Update tokenizer info"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"config.model.tokenizer.update_tokenizer = True\n",
"config.model.tokenizer.dir = TOKENIZER\n",
"config.model.tokenizer.type = TOKENIZER_TYPE_CFG\n",
"print(OmegaConf.to_yaml(config.model.tokenizer))"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "N-t7fRl6GS3A"
},
"source": [
"Need to mention the same data preprocessing scripts here as HF datasets are processed batchwise and the same preprocessing scripts are used for applying the same preprocessing to the dataset, as was used for training the tokenizer. "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "B9nY5JQaIhKz"
},
"outputs": [],
"source": [
"\n",
"config.model.train_ds.normalize_text=False\n",
"config.model.normalize_text = True # same as the normalize_text in the get_hf_text_data.py\n",
"config.model.symbols_to_keep=[\".\"] # same as the symbols_to_keep in the get_hf_text_data.py\n",
"config.model.data_path = 'google/fleurs'\n",
"config.model.data_name = 'te_in'\n",
"config.model.streaming = False\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "d0tyOaqNLq-4"
},
"source": [
"Setting up dataset config for validation and test datasets is similar to the training dataset."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "mJbqMEkjMTfM"
},
"outputs": [],
"source": [
"config.model.train_ds.hf_data_cfg=train_split\n",
"config.model.train_ds.text_key='transcription'\n",
"config.model.train_ds.batch_size=16 # change this based on your GPU memory.\n",
"config.model.train_ds.normalize_text=True\n",
"\n",
"config.model.validation_ds.hf_data_cfg=train_split\n",
"config.model.validation_ds.hf_data_cfg.split='validation' # updated this based on the split\n",
"config.model.validation_ds.text_key='transcription'\n",
"config.model.validation_ds.normalize_text=True\n",
"\n",
"config.model.test_ds.hf_data_cfg=train_split\n",
"config.model.test_ds.hf_data_cfg.split='test' # updated this based on the split\n",
"config.model.test_ds.text_key='transcription'\n",
"config.model.test_ds.normalize_text=True\n",
"\n",
"print(OmegaConf.to_yaml(config.model))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Update the trainer config to include the necessary parameters for training the model."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Training the model\n",
"\n",
"For finetuning the model, basic parameters we need to primarily focus on are optimizer, scheduler, learning rate and the number of steps to finetune.\n",
"\n",
"For this example we will stick to AdamW optimizer and CosineAnnealing scheduler, for 2500 steps. \n",
"These parameters can be updated based on the dataset and the model being used."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"print(OmegaConf.to_yaml(config.model.optim))"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Current optim looks fine except for the warmup_steps\n",
"config.model.optim.sched.warmup_steps = 500 # 10% of the total steps\n",
"config.model.optim.lr = 3e-4\n",
"\n",
"del config.model.spec_augment #For this example, we are not using SpecAugment"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Let's load the model and also the dataloader for training the model"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Unv8-GvOWhad"
},
"outputs": [],
"source": [
"from nemo.collections.asr.models import ASRModel\n",
"asr_model=ASRModel.from_pretrained('nvidia/parakeet-rnnt-0.6b')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"After loading the pretrained model, decoder of the model needs to updated to support the new language."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"\n",
"decoder=asr_model.decoder.state_dict()\n",
"joint_state = asr_model.joint.state_dict()\n",
"prev_vocab_size = asr_model.tokenizer.vocab_size\n",
"asr_model.change_vocabulary(new_tokenizer_dir=TOKENIZER, new_tokenizer_type=TOKENIZER_TYPE_CFG)\n",
"if asr_model.tokenizer.vocab_size == prev_vocab_size: # checking new tokenizer vocab size\n",
" asr_model.decoder.load_state_dict(decoder)\n",
" asr_model.joint.load_state_dict(joint_state)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Loading dataloaders"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "VsunP99saF5c"
},
"outputs": [],
"source": [
"from nemo.utils import logging, model_utils\n",
"cfg = model_utils.convert_model_config_to_dict_config(config)\n",
"asr_model.setup_training_data(cfg.model.train_ds)\n",
"asr_model.setup_validation_data(cfg.model.validation_ds)\n",
"if hasattr(cfg.model, 'test_ds') and cfg.model.test_ds.manifest_filepath is not None:\n",
" asr_model.setup_test_data(cfg.model.test_ds)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "yT1vJH9OkS0u"
},
"outputs": [],
"source": [
"# We will also reduce the hidden dimension of the joint and the prediction networks to preserve some memory\n",
"asr_model.setup_optimization(cfg.model.optim)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "SH4ZfXHOdGhX"
},
"source": [
"------\n",
"\n",
"Setup a Pytorch Lightning Trainer:"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We can reduce the number of steps to 5000 for this small dataset.\n",
"and also update precision to float16 for faster training. # Change this bf16 training on Ampere based GPUs. "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Fmf0iSY-a6LC"
},
"outputs": [],
"source": [
"import torch\n",
"from lightning.pytorch import Trainer\n",
"\n",
"if torch.cuda.is_available():\n",
" accelerator = 'gpu'\n",
"else:\n",
" accelerator = 'gpu'\n",
"\n",
"MAX_STEPS = 5000\n",
"\n",
"# Initialize a Trainer for the Transducer model\n",
"trainer = Trainer(devices=1, accelerator=accelerator, max_epochs=-1, max_steps=MAX_STEPS,\n",
" enable_checkpointing=False, logger=False,\n",
" log_every_n_steps=100, check_val_every_n_epoch=10, precision='bf16')"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "RheLsmA1cRz0"
},
"outputs": [],
"source": [
"# Build the model\n",
"trainer.fit(asr_model)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Once trained users can save the model and use it for inference."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"asr_model.save_to(\"telugu_asr_model.nemo\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Running inference on a sample from the test dataset"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"del asr_model\n",
"asr_model = ASRModel.restore_from(\"telugu_asr_model.nemo\")\n",
"\n",
"!rm -r telugu_asr_model.nemo # remove the model to save space\n",
"\n",
"# Let's test the model on a sample from the test set\n",
"sample = datasets.load_dataset('google/fleurs', 'te_in', split='test')[0]\n",
"transcription = sample['transcription']\n",
"\n",
"# Let's listen to the audio file and its transcription\n",
"ipd.Audio(sample['audio']['array'], rate=16000)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"asr_model.eval()\n",
"# save audio file to disk\n",
"audio_path = \"sample.wav\"\n",
"import soundfile as sf\n",
"sf.write(audio_path, sample['audio']['array'], 16000)\n",
"transcription = asr_model.transcribe([audio_path])"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"\n",
"print(\"Predicted:\", transcription[0])\n",
"print(\"Target:\", sample['transcription'])"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"os.remove(audio_path) # remove the audio file to save space"
]
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"collapsed_sections": [
"V5sMoFHmVvhg"
],
"name": "ASR-with-Transducers.ipynb",
"provenance": [],
"toc_visible": true
},
"kernelspec": {
"display_name": "base",
"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.12.3"
}
},
"nbformat": 4,
"nbformat_minor": 4
}
tutorials/asr/Voice_Activity_Detection.ipynb
+1211
View File
File diff suppressed because it is too large Load Diff
tutorials/asr/asr_adapters/ASR_with_Adapters.ipynb
+1311
View File
File diff suppressed because it is too large Load Diff
tutorials/asr/asr_adapters/Multi_Task_Adapters.ipynb
+1715
View File
File diff suppressed because it is too large Load Diff
tutorials/asr/asr_adapters/README.md
+15
View File
@@ -0,0 +1,15 @@
# Speech Recognition Tutorials With Adapters
------------
In this repository, you will find several tutorials discussing how to utilize Adapter modules for Automatic Speech Recognition (ASR). We also discuss the general concepts and specific ways to apply Adapters for multiple sub-domains of ASR such as domain adaptation.
------------
# Automatic Speech Recognition
1) `ASR_with_Adapters`: An introduction of adapters and their use case with ASR models. Dives into domain adaptation of a pre-trained model with adapter modules, general advantages and disadvantages of adapters and finally trains a model to adapt on a toy dataset.
2) `Multi_Task_Adapters`: An introduction of how to customize multi-task models with adapters. We will train a model on two tasks, one being ASR and the other being a downstream task. We will discuss how to use adapters to finetune a model for Speech Recognition and Speech Translation task on a toy dataset, and dive into construction of custom datasets and prompt formatters for Multi Task Models.
------------
tutorials/asr/images/CosineAnnealing_scheduler.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 86 KiB

tutorials/asr/images/alignment_example.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 23 KiB

tutorials/asr/images/buffered_pipeline.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 74 KiB

tutorials/asr/images/cache_aware_pipeline.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 110 KiB

tutorials/asr/images/canary2_timestamps.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 131 KiB

tutorials/asr/images/contrastive_ssl.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 56 KiB

tutorials/asr/images/multi_instance.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 819 KiB

tutorials/asr/images/multilang_asr_inference.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 327 KiB

tutorials/asr/images/multilang_asr_train.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 381 KiB

tutorials/asr/images/promptformat.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 111 KiB

tutorials/asr/images/speaker_injection.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 328 KiB

tutorials/asr/images/test_wer_wandb.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 224 KiB

tutorials/asr/images/tokenizer.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 259 KiB

tutorials/asr/images/transducer.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 96 KiB

tutorials/asr/images/vad_post_overlap_diagram.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 20 KiB

tutorials/asr/speech_to_text_multitalker_streaming_infer.py
+442
View File
@@ -0,0 +1,442 @@
# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import time
from dataclasses import dataclass, field, is_dataclass
from typing import List, Optional, Union
import pytorch_lightning as pl
import torch
from omegaconf import OmegaConf, open_dict
import nemo.collections.asr as nemo_asr
from nemo.collections.asr.models.sortformer_diar_models import SortformerEncLabelModel
from nemo.collections.asr.parts.utils.multispk_transcribe_utils import (
SpeakerTaggedASR,
add_delay_for_real_time,
get_multi_talker_samples_from_manifest,
write_seglst_file,
)
from nemo.collections.asr.parts.utils.streaming_utils import CacheAwareStreamingAudioBuffer
from nemo.core.config import hydra_runner
from nemo.utils import logging
@dataclass
class MultitalkerTranscriptionConfig:
"""
Configuration for Multi-talker transcription with an ASR model and a diarization model.
"""
# Required configs
diar_model: Optional[str] = None # Path to a .nemo file
diar_pretrained_name: Optional[str] = None # Name of a pretrained model
max_num_of_spks: Optional[int] = 4 # maximum number of speakers
parallel_speaker_strategy: bool = True # whether to use parallel speaker strategy
masked_asr: bool = True # whether to use masked ASR
mask_preencode: bool = False # whether to mask preencode or mask features
cache_gating: bool = True # whether to use cache gating
cache_gating_buffer_size: int = 2 # buffer size for cache gating
single_speaker_mode: bool = False # whether to use single speaker mode
# General configs
session_len_sec: float = -1 # End-to-end diarization session length in seconds
num_workers: int = 8
random_seed: Optional[int] = None # seed number going to be used in seed_everything()
log: bool = True # If True, log will be printed
# Streaming diarization configs
streaming_mode: bool = True # If True, streaming diarization will be used.
spkcache_len: int = 188
spkcache_refresh_rate: int = 0
fifo_len: int = 188
chunk_len: int = 0
chunk_left_context: int = 0
chunk_right_context: int = 0
# If `cuda` is a negative number, inference will be on CPU only.
cuda: Optional[int] = None
allow_mps: bool = False # allow to select MPS device (Apple Silicon M-series GPU)
matmul_precision: str = "highest" # Literal["highest", "high", "medium"]
# ASR Configs
asr_model: Optional[str] = None
device: str = 'cuda'
audio_file: Optional[str] = None
manifest_file: Optional[str] = None
use_amp: bool = True
debug_mode: bool = False
batch_size: int = 32
chunk_size: int = -1
shift_size: int = -1
left_chunks: int = 2
online_normalization: bool = False
output_path: Optional[str] = None
pad_and_drop_preencoded: bool = False
set_decoder: Optional[str] = None # ["ctc", "rnnt"]
att_context_size: Optional[list] = None
generate_realtime_scripts: bool = True
word_window: int = 50
sent_break_sec: float = 30.0
fix_prev_words_count: int = 5
update_prev_words_sentence: int = 5
left_frame_shift: int = -1
right_frame_shift: int = 0
min_sigmoid_val: float = 1e-2
discarded_frames: int = 8
print_time: bool = True
print_sample_indices: List[int] = field(default_factory=lambda: [0])
colored_text: bool = True
real_time_mode: bool = False
print_path: str = "./"
ignored_initial_frame_steps: int = 5
verbose: bool = False
feat_len_sec: float = 0.01
finetune_realtime_ratio: float = 0.01
spk_supervision: str = "diar" # ["diar", "rttm"]
binary_diar_preds: bool = False
def launch_serial_streaming(
cfg,
asr_model,
diar_model,
streaming_buffer,
pad_and_drop_preencoded=False,
):
"""
Launch the serial streaming inference with ASR model and diarization model.
Args:
cfg (Any): The configuration object containing the parameters for the streaming inference.
asr_model (Any): The ASR model loaded from the path provided in MultitalkerTranscriptionConfig.
diar_model (Any): The diarization model loadded from the path provided in MultitalkerTranscriptionConfig.
streaming_buffer: An iterator that yields chunks of audio data and their lengths.
pad_and_drop_preencoded: A boolean flag indicating whether to pad and drop the extra pre-encoded tokens.
"""
streaming_buffer_iter = iter(streaming_buffer)
multispk_asr_streamer = SpeakerTaggedASR(cfg, asr_model, diar_model)
feat_frame_count = 0
session_start_time = time.time()
for step_num, (chunk_audio, chunk_lengths) in enumerate(streaming_buffer_iter):
drop_extra_pre_encoded = (
0
if step_num == 0 and not pad_and_drop_preencoded
else asr_model.encoder.streaming_cfg.drop_extra_pre_encoded
)
loop_start_time = time.time()
with torch.inference_mode():
with autocast:
with torch.no_grad():
multispk_asr_streamer.perform_serial_streaming_stt_spk(
step_num=step_num,
chunk_audio=chunk_audio,
chunk_lengths=chunk_lengths,
is_buffer_empty=streaming_buffer.is_buffer_empty(),
drop_extra_pre_encoded=drop_extra_pre_encoded,
)
if cfg.real_time_mode:
add_delay_for_real_time(
cfg=cfg,
chunk_audio=chunk_audio,
session_start_time=session_start_time,
feat_frame_count=feat_frame_count,
loop_end_time=time.time(),
loop_start_time=loop_start_time,
)
feat_frame_count += chunk_audio.shape[-1] - cfg.discarded_frames
return multispk_asr_streamer
def launch_parallel_streaming(
cfg,
asr_model,
diar_model,
streaming_buffer,
pad_and_drop_preencoded=False,
):
streaming_buffer_iter = iter(streaming_buffer)
multispk_asr_streamer = SpeakerTaggedASR(cfg, asr_model, diar_model)
feat_frame_count = 0
session_start_time = time.time()
for step_num, (chunk_audio, chunk_lengths) in enumerate(streaming_buffer_iter):
drop_extra_pre_encoded = (
0
if step_num == 0 and not pad_and_drop_preencoded
else asr_model.encoder.streaming_cfg.drop_extra_pre_encoded
)
loop_start_time = time.time()
with torch.inference_mode():
with autocast:
with torch.no_grad():
multispk_asr_streamer.perform_parallel_streaming_stt_spk(
step_num=step_num,
chunk_audio=chunk_audio,
chunk_lengths=chunk_lengths,
is_buffer_empty=streaming_buffer.is_buffer_empty(),
drop_extra_pre_encoded=drop_extra_pre_encoded,
)
if cfg.real_time_mode:
add_delay_for_real_time(
cfg=cfg,
chunk_audio=chunk_audio,
session_start_time=session_start_time,
feat_frame_count=feat_frame_count,
loop_end_time=time.time(),
loop_start_time=loop_start_time,
)
feat_frame_count += chunk_audio.shape[-1] - cfg.discarded_frames
return multispk_asr_streamer
@hydra_runner(config_name="MultitalkerTranscriptionConfig", schema=MultitalkerTranscriptionConfig)
def main(cfg: MultitalkerTranscriptionConfig) -> Union[MultitalkerTranscriptionConfig]:
for key in cfg:
cfg[key] = None if cfg[key] == 'None' else cfg[key]
if is_dataclass(cfg):
cfg = OmegaConf.structured(cfg)
if cfg.random_seed:
pl.seed_everything(cfg.random_seed)
if cfg.diar_model is None and cfg.diar_pretrained_name is None:
raise ValueError("Both cfg.diar_model and cfg.pretrained_name cannot be None!")
if cfg.audio_file is None and cfg.manifest_file is None:
raise ValueError("Both cfg.audio_file and cfg.manifest_file cannot be None!")
# setup GPU
torch.set_float32_matmul_precision(cfg.matmul_precision)
if cfg.cuda is None:
if torch.cuda.is_available():
device = [0] # use 0th CUDA device
accelerator = 'gpu'
map_location = torch.device('cuda:0')
elif cfg.allow_mps and hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
device = [0]
accelerator = 'mps'
map_location = torch.device('mps')
else:
device = 1
accelerator = 'cpu'
map_location = torch.device('cpu')
else:
device = [cfg.cuda]
accelerator = 'gpu'
map_location = torch.device(f'cuda:{cfg.cuda}')
if cfg.diar_model.endswith(".ckpt"):
diar_model = SortformerEncLabelModel.load_from_checkpoint(
checkpoint_path=cfg.diar_model, map_location=map_location, strict=False
)
elif cfg.diar_model.endswith(".nemo"):
diar_model = SortformerEncLabelModel.restore_from(restore_path=cfg.diar_model, map_location=map_location)
else:
raise ValueError("cfg.diar_model must end with.ckpt or.nemo!")
# Model setup for inference
trainer = pl.Trainer(devices=device, accelerator=accelerator)
diar_model.set_trainer(trainer)
diar_model._cfg.test_ds.session_len_sec = cfg.session_len_sec
diar_model._cfg.test_ds.manifest_filepath = cfg.manifest_file
diar_model._cfg.test_ds.batch_size = cfg.batch_size
diar_model._cfg.test_ds.num_workers = cfg.num_workers
diar_model.setup_test_data(test_data_config=diar_model._cfg.test_ds)
diar_model = diar_model.eval()
# Steaming mode setup
diar_model.streaming_mode = cfg.streaming_mode
diar_model.sortformer_modules.chunk_len = cfg.chunk_len
diar_model.sortformer_modules.spkcache_len = cfg.spkcache_len
diar_model.sortformer_modules.chunk_left_context = cfg.chunk_left_context
diar_model.sortformer_modules.chunk_right_context = cfg.chunk_right_context
diar_model.sortformer_modules.fifo_len = cfg.fifo_len
diar_model.sortformer_modules.log = cfg.log
diar_model.sortformer_modules.spkcache_refresh_rate = cfg.spkcache_refresh_rate
if cfg.audio_file is not None and cfg.manifest_file is not None:
logging.warning("Both audio_file and manifest_file are specified. Audio_file will be used with top priority.")
elif cfg.audio_file is not None:
logging.info("audio_file is specified. Using audio_file as input.")
elif cfg.manifest_file is not None:
logging.info("manifest_file is specified. Using manifest_file as input.")
else:
raise ValueError("One of audio_file or manifest_file must be specified!")
if cfg.asr_model.endswith('.nemo'):
logging.info(f"Using local ASR model from {cfg.asr_model}")
asr_model = nemo_asr.models.ASRModel.restore_from(restore_path=cfg.asr_model)
else:
logging.info(f"Using NGC cloud ASR model {cfg.asr_model}")
asr_model = nemo_asr.models.ASRModel.from_pretrained(model_name=cfg.asr_model)
logging.info(asr_model.encoder.streaming_cfg)
if cfg.set_decoder is not None:
if hasattr(asr_model, "cur_decoder"):
asr_model.change_decoding_strategy(decoder_type=cfg.set_decoder)
else:
raise ValueError("Decoder cannot get changed for non-Hybrid ASR models.")
if cfg.att_context_size is not None:
if hasattr(asr_model.encoder, "set_default_att_context_size"):
asr_model.encoder.set_default_att_context_size(att_context_size=cfg.att_context_size)
else:
raise ValueError("Model does not support multiple lookaheads.")
global autocast
autocast = torch.amp.autocast(asr_model.device.type, enabled=cfg.use_amp)
# Initialize to avoid "possibly used before assignment" error
multispk_asr_streamer = None
# configure the decoding config
decoding_cfg = asr_model.cfg.decoding
with open_dict(decoding_cfg):
decoding_cfg.strategy = "greedy"
decoding_cfg.preserve_alignments = False
if hasattr(asr_model, 'joint'): # if an RNNT model
decoding_cfg.greedy.max_symbols = 10
decoding_cfg.fused_batch_size = -1
asr_model.change_decoding_strategy(decoding_cfg)
asr_model = asr_model.to(cfg.device)
asr_model.eval()
# chunk_size is set automatically for models trained for streaming.
# For models trained for offline mode with full context, we need to pass the chunk_size explicitly.
if cfg.chunk_size > 0:
if cfg.shift_size < 0:
shift_size = cfg.chunk_size
else:
shift_size = cfg.shift_size
asr_model.encoder.setup_streaming_params(
chunk_size=cfg.chunk_size, left_chunks=cfg.left_chunks, shift_size=shift_size
)
# In streaming, offline normalization is not feasible as we don't have access to the
# whole audio at the beginning When online_normalization is enabled, the normalization
# of the input features (mel-spectrograms) are done per step It is suggested to train
# the streaming models without any normalization in the input features.
if cfg.online_normalization:
if asr_model.cfg.preprocessor.normalize not in ["per_feature", "all_feature"]:
logging.warning(
"online_normalization is enabled but the model has"
"no normalization in the feature extration part, so it is ignored."
)
online_normalization = False
else:
online_normalization = True
else:
online_normalization = False
seglst_dict_list = []
if cfg.audio_file is not None:
# Stream a single audio file
samples = [
{
'audio_filepath': cfg.audio_file,
}
]
streaming_buffer = CacheAwareStreamingAudioBuffer(
model=asr_model,
online_normalization=online_normalization,
pad_and_drop_preencoded=cfg.pad_and_drop_preencoded,
)
cfg.batch_size = len(samples)
streaming_buffer.append_audio_file(audio_filepath=cfg.audio_file, stream_id=-1)
if cfg.parallel_speaker_strategy:
multispk_asr_streamer = launch_parallel_streaming(
cfg=cfg,
asr_model=asr_model,
diar_model=diar_model,
streaming_buffer=streaming_buffer,
pad_and_drop_preencoded=cfg.pad_and_drop_preencoded,
)
multispk_asr_streamer.generate_seglst_dicts_from_parallel_streaming(samples=samples)
else:
multispk_asr_streamer = launch_serial_streaming(
cfg=cfg,
asr_model=asr_model,
diar_model=diar_model,
streaming_buffer=streaming_buffer,
)
multispk_asr_streamer.generate_seglst_dicts_from_serial_streaming(samples=samples)
seglst_dict_list.extend(multispk_asr_streamer.instance_manager.seglst_dict_list)
else:
# Stream audio files in a manifest file in batched mode
feat_per_sec = round(asr_model.cfg.preprocessor.window_stride * asr_model.cfg.encoder.subsampling_factor, 2)
samples, rttms_mask_mats = get_multi_talker_samples_from_manifest(
cfg, manifest_file=cfg.manifest_file, feat_per_sec=feat_per_sec, max_spks=cfg.max_num_of_spks
)
# Note: rttms_mask_mats contains PyTorch tensors, so we pass it directly instead of storing in config
if cfg.spk_supervision == "rttm":
diar_model.add_rttms_mask_mats(rttms_mask_mats, device=asr_model.device)
logging.info(f"Loaded {len(samples)} from the manifest at {cfg.manifest_file}.")
streaming_buffer = CacheAwareStreamingAudioBuffer(
model=asr_model,
online_normalization=online_normalization,
pad_and_drop_preencoded=cfg.pad_and_drop_preencoded,
)
batch_samples = []
for sample_idx, sample in enumerate(samples):
batch_samples.append(sample)
streaming_buffer.append_audio_file(sample['audio_filepath'], stream_id=-1)
logging.info(f'Added this sample to the buffer: {sample["audio_filepath"]}')
if (sample_idx + 1) % cfg.batch_size == 0 or sample_idx == len(samples) - 1:
logging.info(f"Starting to stream samples {sample_idx - len(streaming_buffer) + 1} to {sample_idx}...")
if cfg.parallel_speaker_strategy:
multispk_asr_streamer = launch_parallel_streaming(
cfg=cfg,
asr_model=asr_model,
diar_model=diar_model,
streaming_buffer=streaming_buffer,
pad_and_drop_preencoded=cfg.pad_and_drop_preencoded,
)
multispk_asr_streamer.generate_seglst_dicts_from_parallel_streaming(samples=batch_samples)
else:
multispk_asr_streamer = launch_serial_streaming(
cfg=cfg,
asr_model=asr_model,
diar_model=diar_model,
streaming_buffer=streaming_buffer,
)
multispk_asr_streamer.generate_seglst_dicts_from_serial_streaming(samples=batch_samples)
seglst_dict_list.extend(multispk_asr_streamer.instance_manager.seglst_dict_list)
streaming_buffer.reset_buffer()
batch_samples = []
if len(seglst_dict_list) == 0:
logging.warning("No segmentation list dictionary found.")
return
if cfg.output_path is not None and multispk_asr_streamer is not None:
if cfg.parallel_speaker_strategy:
write_seglst_file(seglst_dict_list=seglst_dict_list, output_path=cfg.output_path)
else:
write_seglst_file(seglst_dict_list=seglst_dict_list, output_path=cfg.output_path)
if __name__ == '__main__':
main()
tutorials/audio/README.md
+6
View File
@@ -0,0 +1,6 @@
# NeMo Audio Processing Tutorials
An quick way to get started with audio processing in NeMo.
## Speech Enhancement
* [Quickstart: Speech Enhancement with NeMo](speech_enhancement/Speech_Enhancement_with_NeMo.ipynb)
tutorials/audio/speech_enhancement/BNR_Speech_enhancement_with_NeMo.ipynb
+880
View File
@@ -0,0 +1,880 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "2fb43f21",
"metadata": {
"id": "2fb43f21"
},
"source": [
"# Introduction\n",
"\n",
"The goal of this tutorial is to demonstrate the basic steps required to setup and train the Maxine Background Noise removal network [1] in NeMo.\n",
"\n",
"This notebook covers the following steps:\n",
"\n",
"* Download speech and noise data\n",
"* Prepare the training data by mixing speech and noise\n",
"* Configure and train a simple single-output model\n",
"\n",
"Note that this tutorial is only for demonstration purposes.\n",
"\n",
"*Disclaimer:*\n",
"User is responsible for checking the content of datasets and the applicable licenses and determining if suitable for the intended use."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "50525ece",
"metadata": {
"id": "50525ece"
},
"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",
"\n",
"GIT_USER = 'NVIDIA-NeMo'\n",
"BRANCH = 'main'\n",
"\n",
"if 'google.colab' in str(get_ipython()):\n",
"\n",
" # Install dependencies\n",
" !pip install wget\n",
" !apt-get install sox libsndfile1 ffmpeg\n",
" !pip install text-unidecode\n",
" !pip install matplotlib>=3.3.2\n",
"\n",
" ## Install NeMo\n",
" !python -m pip install \"nemo_toolkit[all] @ git+https://github.com/{GIT_USER}/Speech.git@{BRANCH}\"\n",
"\n",
" ## Install TorchAudio\n",
" !pip install torchaudio>=0.13.0 -f https://download.pytorch.org/whl/torch_stable.html"
]
},
{
"cell_type": "markdown",
"id": "765e8b8a",
"metadata": {
"id": "765e8b8a"
},
"source": [
"The following cell will take care of the necessary imports and prepare utility functions used throughout the notebook."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "9248dd5f",
"metadata": {
"id": "9248dd5f"
},
"outputs": [],
"source": [
"import glob\n",
"import librosa\n",
"import os\n",
"import torch\n",
"import tqdm\n",
"\n",
"import IPython.display as ipd\n",
"import matplotlib.pyplot as plt\n",
"import numpy as np\n",
"import lightning.pytorch as pl\n",
"import soundfile as sf\n",
"\n",
"from omegaconf import OmegaConf, open_dict\n",
"from pathlib import Path\n",
"from torchmetrics.functional.audio import signal_distortion_ratio, scale_invariant_signal_distortion_ratio\n",
"\n",
"from nemo.utils.notebook_utils import download_an4\n",
"from nemo.collections.asr.parts.preprocessing.segment import AudioSegment\n",
"from nemo.collections.asr.parts.utils.manifest_utils import read_manifest, write_manifest\n",
"\n",
"\n",
"# Utility functions for displaying signals and metrics\n",
"def show_signal(signal: np.ndarray, sample_rate: int = 16000, tag: str = 'Signal'):\n",
" \"\"\"Show the time-domain signal and its spectrogram.\n",
" \"\"\"\n",
" fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(10, 2.5))\n",
"\n",
" # show waveform\n",
" t = np.arange(0, len(signal)) / sample_rate\n",
"\n",
" ax[0].plot(t, signal)\n",
" ax[0].set_xlim(0, t.max())\n",
" ax[0].grid()\n",
" ax[0].set_xlabel('time / s')\n",
" ax[0].set_ylabel('amplitude')\n",
" ax[0].set_title(tag)\n",
"\n",
" n_fft = 1024\n",
" hop_length = 256\n",
"\n",
" D = librosa.amplitude_to_db(np.abs(librosa.stft(signal, n_fft=n_fft, hop_length=hop_length)), ref=np.max)\n",
" img = librosa.display.specshow(D, y_axis='linear', x_axis='time', sr=sample_rate, n_fft=n_fft, hop_length=hop_length, ax=ax[1])\n",
" ax[1].set_title(tag)\n",
" \n",
" plt.tight_layout()\n",
" plt.colorbar(img, format=\"%+2.f dB\", ax=ax)\n",
"\n",
"def show_metrics(signal: np.ndarray, reference: np.ndarray, sample_rate: int = 16000, tag: str = 'Signal'):\n",
" \"\"\"Show metrics for the time-domain signal and the reference signal.\n",
" \"\"\"\n",
" sdr = signal_distortion_ratio(preds=torch.tensor(signal), target=torch.tensor(reference))\n",
" sisdr = scale_invariant_signal_distortion_ratio(preds=torch.tensor(signal), target=torch.tensor(reference))\n",
" print(tag)\n",
" print('\\tsdr: ', sdr.item())\n",
" print('\\tsisdr:', sisdr.item())"
]
},
{
"cell_type": "markdown",
"id": "bfa2199e",
"metadata": {
"id": "bfa2199e"
},
"source": [
"### Data preparation"
]
},
{
"cell_type": "markdown",
"id": "0aa05a36",
"metadata": {
"id": "0aa05a36"
},
"source": [
"In this notebook, it is assumed that all audio will be resampled to 16kHz and the data and configuration will be stored under `root_dir` as defined below."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "8e790630",
"metadata": {
"id": "8e790630"
},
"outputs": [],
"source": [
"# sample rate used throughout the notebook\n",
"sample_rate = 16000\n",
"\n",
"# root directory for data preparation, configurations, etc\n",
"root_dir = Path('./')\n",
"\n",
"# data directory\n",
"data_dir = root_dir / 'data'\n",
"data_dir.mkdir(exist_ok=True)\n",
"\n",
"# scripts directory\n",
"scripts_dir = root_dir / 'scripts'\n",
"scripts_dir.mkdir(exist_ok=True)"
]
},
{
"cell_type": "markdown",
"id": "9af717ad",
"metadata": {
"id": "9af717ad"
},
"source": [
"Clean speech data is used to prepare datasets used for training a simple speech enhancement model.\n",
"\n",
"In this tutorial, a subset of LibriSpeech dataset [2] will be downloaded and used as the speech material.\n",
"The following cell will download and prepare the speech data."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e1c5a60a",
"metadata": {
"id": "e1c5a60a"
},
"outputs": [],
"source": [
"speech_dir = data_dir / 'speech'\n",
"speech_data_set = 'mini'\n",
"\n",
"# Copy script\n",
"get_librispeech_script = os.path.join(scripts_dir, 'get_librispeech_data.py')\n",
"if not os.path.exists(get_librispeech_script):\n",
" !wget -P $scripts_dir https://raw.githubusercontent.com/{GIT_USER}/NeMo/{BRANCH}/scripts/dataset_processing/get_librispeech_data.py\n",
"\n",
"# Download the data\n",
"if not speech_dir.is_dir():\n",
" speech_dir.mkdir(exist_ok=True)\n",
" !python {get_librispeech_script} --data_root={speech_dir} --data_set={speech_data_set}\n",
"else:\n",
" print('Speech dataset already exists in:', speech_dir)\n",
"\n",
"# Reduce the size of test dataset for this tutorial to 1000 clean utterances for train and 100 clean utterances for test\n",
"train_metadata = read_manifest(speech_dir / 'train_clean_5.json')\n",
"write_manifest(speech_dir / 'train.json', train_metadata[:1000])\n",
"\n",
"test_metadata = read_manifest(speech_dir / 'dev_clean_2.json')\n",
"write_manifest(speech_dir / 'test.json', test_metadata[:100])\n",
"\n",
"# Speech manifests\n",
"speech_manifest = {\n",
" 'train': speech_dir / 'train.json',\n",
" 'test': speech_dir / 'test.json',\n",
"}"
]
},
{
"cell_type": "markdown",
"id": "06f07374",
"metadata": {
"id": "06f07374"
},
"source": [
"Noise data will be mixed with the downloaded speech data to prepare a noisy dataset.\n",
"\n",
"The following cell will download and prepare the noise data using a subset of the DEMAND dataset [3] will be downloaded and used as the noise data."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "b1de1089",
"metadata": {
"id": "b1de1089"
},
"outputs": [],
"source": [
"noise_dir = data_dir / 'noise'\n",
"noise_data_set = 'STRAFFIC,PSTATION'\n",
"\n",
"# Copy script\n",
"get_demand_script = os.path.join(scripts_dir, 'get_demand_data.py')\n",
"if not os.path.exists(get_demand_script):\n",
" !wget -P $scripts_dir https://raw.githubusercontent.com/{GIT_USER}/NeMo/{BRANCH}/scripts/dataset_processing/get_demand_data.py\n",
"\n",
"if not noise_dir.is_dir():\n",
" noise_dir.mkdir(exist_ok=True)\n",
" !python {get_demand_script} --data_root={noise_dir} --data_sets={noise_data_set}\n",
"else:\n",
" print('Noise directory already exists in:', noise_dir)\n",
"\n",
"\n",
"def create_noise_manifest(base_dir, subset, offset=0, duration=None):\n",
" \"\"\"Split the noise data set into train and test subsets.\n",
" \"\"\"\n",
" complete_noise_manifests = glob.glob(str(base_dir / 'manifests' / '*.json'))\n",
" subset_noise_manifest = base_dir / f'{subset}_manifest.json'\n",
" \n",
" subset_metadata = []\n",
"\n",
" for noise_manifest in complete_noise_manifests:\n",
" complete_metadata = read_manifest(noise_manifest)\n",
" \n",
" for item in complete_metadata:\n",
" new_item = item.copy()\n",
" new_item['offset'] = offset\n",
" new_item['duration'] = duration\n",
" subset_metadata.append(new_item)\n",
"\n",
" write_manifest(subset_noise_manifest.as_posix(), subset_metadata)\n",
"\n",
" return subset_noise_manifest\n",
"\n",
"noise_manifest = {\n",
" 'train': create_noise_manifest(noise_dir, 'train', offset=0, duration=200),\n",
" 'test': create_noise_manifest(noise_dir, 'test', offset=200, duration=100),\n",
"}"
]
},
{
"cell_type": "markdown",
"id": "1378d78a",
"metadata": {
"id": "1378d78a"
},
"source": [
"For this tutorial, a single-channel noisy dataset is constructed by adding speech and noise.\n",
"\n",
"The following block will add speech and noise and save the noisy data. The noisy data is created by mixing speech and noise at a few pre-defined signal-to-noise ratios (SNRs). Note that a separate manifest will be created for each SNR."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "d6c678a0",
"metadata": {
"id": "d6c678a0"
},
"outputs": [],
"source": [
"%%capture\n",
"# Suppress output of this cell, since the script used below is relatively verbose.\n",
"\n",
"# Copy script\n",
"add_noise_script = os.path.join(scripts_dir, 'add_noise.py')\n",
"if not os.path.exists(add_noise_script):\n",
" !wget -P $scripts_dir https://raw.githubusercontent.com/{GIT_USER}/NeMo/{BRANCH}/scripts/dataset_processing/add_noise.py\n",
"\n",
"# Generate noisy datasets and save the noise component as well.\n",
"noisy_dir = data_dir / 'noisy'\n",
"noisy_dir.mkdir(exist_ok=True)\n",
"\n",
"for subset in ['train', 'test']:\n",
" noisy_subset_dir = noisy_dir / subset\n",
"\n",
" if not noisy_subset_dir.is_dir():\n",
" noisy_subset_dir.mkdir(exist_ok=True)\n",
" !python {add_noise_script} --input_manifest={speech_manifest[subset]} --noise_manifest={noise_manifest[subset]} --out_dir={noisy_subset_dir} --snrs 0 5 10 15 20 --num_workers 4 --save_noise"
]
},
{
"cell_type": "markdown",
"id": "c4cd1426",
"metadata": {
"id": "c4cd1426"
},
"source": [
"Training a model requires an input dataset which includes information about the noisy input signal and the desired (target) output signal.\n",
"\n",
"In this tutorial, train and test manifests are created by combining the information from the speech manifests and each noisy manifest generated in the previous step. Note that the final manifests include `noisy_filepath`, `speech_filepath` and `noise_filepath`. These keys can be used to define the input signal and the output signal for the model."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "fcfeb639",
"metadata": {
"id": "fcfeb639"
},
"outputs": [],
"source": [
"dataset_manifest = {\n",
" 'train': data_dir / 'dataset_train.json',\n",
" 'test': data_dir / 'dataset_test.json',\n",
"}\n",
"\n",
"for subset in ['train', 'test']:\n",
" # Load clean manifest\n",
" speech_metadata = read_manifest(speech_manifest[subset])\n",
"\n",
" # Load noisy manifests\n",
" noisy_manifests = glob.glob(str(noisy_dir / subset / 'manifests/*.json'))\n",
" noisy_manifests.sort()\n",
"\n",
" subset_metadata = []\n",
"\n",
" for noisy_manifest in noisy_manifests:\n",
" noisy_metadata = read_manifest(noisy_manifest)\n",
"\n",
" for speech_item, noisy_item in tqdm.tqdm(zip(speech_metadata, noisy_metadata), total=len(noisy_metadata)):\n",
" # Check that the file matches\n",
" assert os.path.basename(speech_item['audio_filepath']) == os.path.basename(noisy_item['audio_filepath']), f'Speech: {speech_item}. Noisy: {noisy_item}'\n",
"\n",
" # Create a new item for the subset manifest\n",
" subset_item = {\n",
" 'noisy_filepath': noisy_item['audio_filepath'],\n",
" 'speech_filepath': speech_item['audio_filepath'],\n",
" 'noise_filepath': noisy_item['noise_filepath'],\n",
" 'duration': noisy_item['duration'],\n",
" 'offset': noisy_item.get('offset', 0)\n",
" }\n",
"\n",
" subset_metadata.append(subset_item)\n",
"\n",
" # Save the subset manifest\n",
" write_manifest(dataset_manifest[subset].as_posix(), subset_metadata)"
]
},
{
"cell_type": "markdown",
"id": "aa4248ab",
"metadata": {
"id": "aa4248ab"
},
"source": [
"### Model configuration\n",
"\n",
"We use the SEASR model architecture for training the dataset.",
"\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "002b917c",
"metadata": {
"id": "002b917c"
},
"outputs": [],
"source": [
"config_dir = root_dir / 'conf'\n",
"config_dir.mkdir(exist_ok=True)\n",
"\n",
"config_path = config_dir / 'maxine_bnr.yaml'\n",
"\n",
"if not config_path.is_file():\n",
" !wget https://raw.githubusercontent.com/{GIT_USER}/NeMo/{BRANCH}/examples/audio/conf/maxine_bnr.yaml -P {config_dir.as_posix()}\n",
"\n",
"config = OmegaConf.load(config_path)\n",
"config = OmegaConf.to_container(config, resolve=True)\n",
"config = OmegaConf.create(config)\n",
"\n",
"print('Loaded config')\n",
"print(OmegaConf.to_yaml(config))"
]
},
{
"cell_type": "markdown",
"id": "65cab3ad",
"metadata": {
"id": "65cab3ad"
},
"source": [
"Training dataset is configured with the following parameters\n",
"* `manifest_filepath` points to a manifest file, with each line containing a dictionary corresponding to a single example\n",
"* `input_key` is the key corresponding to the input audio signal in the example dictionary\n",
"* `target_key` is the key corresponding to the desired output (target) audio signal in the example dictionary\n",
"* `min_duration` can be used to filter out short examples"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "91717b15",
"metadata": {
"id": "91717b15"
},
"outputs": [],
"source": [
"# Setup training dataset\n",
"config.model.train_ds.manifest_filepath = dataset_manifest['train'].as_posix()\n",
"config.model.train_ds.input_key = 'noisy_filepath'\n",
"config.model.train_ds.target_key = 'speech_filepath'\n",
"config.model.train_ds.min_duration = 0 # load all audio files, without filtering short ones\n",
"\n",
"print(\"Train dataset config:\")\n",
"print(OmegaConf.to_yaml(config.model.train_ds))"
]
},
{
"cell_type": "markdown",
"id": "c11592e4",
"metadata": {
"id": "c11592e4"
},
"source": [
"Validation and test datasets can be configured in the same way as the training dataset. Here, we use the same dataset for validation and testing purposes for simplicity."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "d8bc0a0a",
"metadata": {
"id": "d8bc0a0a"
},
"outputs": [],
"source": [
"# Use test manifest for validation and test sets\n",
"config.model.validation_ds.manifest_filepath = dataset_manifest['test'].as_posix()\n",
"config.model.validation_ds.input_key = 'noisy_filepath'\n",
"config.model.validation_ds.target_key = 'speech_filepath'\n",
"\n",
"config.model.test_ds.manifest_filepath = dataset_manifest['test'].as_posix()\n",
"config.model.test_ds.input_key = 'noisy_filepath'\n",
"config.model.test_ds.target_key = 'speech_filepath'\n",
"\n",
"print(\"Validation dataset config:\")\n",
"print(OmegaConf.to_yaml(config.model.validation_ds))\n",
"\n",
"print(\"Test dataset config:\")\n",
"print(OmegaConf.to_yaml(config.model.test_ds))"
]
},
{
"cell_type": "markdown",
"id": "1bc0f8f3",
"metadata": {
"id": "1bc0f8f3"
},
"source": [
"Metrics for validation and test set are configured in the following cell.\n",
"\n",
"In this tutorial, signal-to-distortion ratio (SDR) and scale-invariant SDR from torch metrics are used [4]."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "82de1dbb",
"metadata": {
"id": "82de1dbb"
},
"outputs": [],
"source": [
"# Setup metrics to compute on validation and test sets\n",
"metrics = OmegaConf.create({\n",
" 'sisdr': {\n",
" '_target_': 'torchmetrics.audio.ScaleInvariantSignalDistortionRatio',\n",
" },\n",
" 'sdr': {\n",
" '_target_': 'torchmetrics.audio.SignalDistortionRatio',\n",
" }\n",
"})\n",
"config.model.metrics.val = metrics\n",
"config.model.metrics.test = metrics\n",
"\n",
"print(\"Metrics config:\")\n",
"print(OmegaConf.to_yaml(config.model.metrics))"
]
},
{
"cell_type": "markdown",
"id": "4eda2d65",
"metadata": {
"id": "4eda2d65"
},
"source": [
"### Trainer configuration\n",
"NeMo models are primarily PyTorch Lightning modules and therefore are entirely compatible with the PyTorch Lightning ecosystem."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e53ba6ef",
"metadata": {
"id": "e53ba6ef"
},
"outputs": [],
"source": [
"print(\"Trainer config:\")\n",
"print(OmegaConf.to_yaml(config.trainer))"
]
},
{
"cell_type": "markdown",
"id": "a9201689",
"metadata": {
"id": "a9201689"
},
"source": [
"We can modify some trainer configs for this tutorial.\n",
"Most importantly, the number of epochs is set to a small value, to limit the runtime for the purpose of this example."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "17b7499f",
"metadata": {
"id": "17b7499f"
},
"outputs": [],
"source": [
"# Checks if we have GPU available and uses it\n",
"accelerator = 'gpu' if torch.cuda.is_available() else 'cpu'\n",
"config.trainer.devices = 1\n",
"config.trainer.accelerator = accelerator\n",
"\n",
"# Reduces maximum number of epochs for quick demonstration\n",
"config.trainer.max_epochs = 10\n",
"\n",
"# Remove distributed training flags\n",
"config.trainer.strategy = 'auto'\n",
"\n",
"# Instantiate the trainer\n",
"trainer = pl.Trainer(**config.trainer)"
]
},
{
"cell_type": "markdown",
"id": "053192a4",
"metadata": {
"id": "053192a4"
},
"source": [
"### Experiment manager\n",
"\n",
"NeMo has an experiment manager that handles logging and checkpointing."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "ccb8948d",
"metadata": {
"id": "ccb8948d"
},
"outputs": [],
"source": [
"from nemo.utils.exp_manager import exp_manager\n",
"\n",
"exp_dir = exp_manager(trainer, config.get(\"exp_manager\", None))\n",
"# The exp_dir provides a path to the current experiment for easy access\n",
"\n",
"print(\"Experiment directory:\")\n",
"print(exp_dir)"
]
},
{
"cell_type": "markdown",
"id": "8bb583f8",
"metadata": {
"id": "8bb583f8"
},
"source": [
"### Model instantiation"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "6905fea2",
"metadata": {
"id": "6905fea2"
},
"outputs": [],
"source": [
"from nemo.collections.audio.models.maxine import BNR2",
"\n",
"enhancement_model = BNR2(cfg=config.model, trainer=trainer)"
]
},
{
"cell_type": "markdown",
"id": "6ee1a6d2",
"metadata": {
"id": "6ee1a6d2"
},
"source": [
"### Training\n",
"Create a Tensorboard visualization to monitor progress"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "dc9e7639",
"metadata": {
"id": "dc9e7639"
},
"outputs": [],
"source": [
"try:\n",
" from google import colab\n",
" COLAB_ENV = True\n",
"except (ImportError, ModuleNotFoundError):\n",
" COLAB_ENV = False\n",
"\n",
"# Load the TensorBoard notebook extension\n",
"if COLAB_ENV:\n",
" %load_ext tensorboard\n",
" %tensorboard --logdir {exp_dir}\n",
"else:\n",
" print(\"To use tensorboard, please use this notebook in a Google Colab environment.\")"
]
},
{
"cell_type": "markdown",
"id": "5e962b9b",
"metadata": {
"id": "5e962b9b"
},
"source": [
"Training can be started using `trainer.fit`:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "0b00b3f4",
"metadata": {
"id": "0b00b3f4"
},
"outputs": [],
"source": [
"trainer.fit(enhancement_model)"
]
},
{
"cell_type": "markdown",
"id": "5b5c1e78",
"metadata": {
"id": "5b5c1e78"
},
"source": [
"After the training is completed, the configured metrics can be easily computed on the test set as follows:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "5216195d",
"metadata": {
"id": "5216195d"
},
"outputs": [],
"source": [
"trainer.test(enhancement_model, ckpt_path=None)"
]
},
{
"cell_type": "markdown",
"id": "bf5b73e6",
"metadata": {
"id": "bf5b73e6"
},
"source": [
"### Inference\n",
"\n",
"The following cell provides an example of inference on an single audio file.\n",
"For simplicity, the audio file information is taken from the test dataset."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "6b5cf923",
"metadata": {
"id": "6b5cf923"
},
"outputs": [],
"source": [
"# Load a single audio example from the test set\n",
"test_metadata = read_manifest(dataset_manifest['test'].as_posix())\n",
"\n",
"# Path to audio files\n",
"noisy_filepath = test_metadata[-1]['noisy_filepath'] # noisy audio\n",
"speech_filepath = test_metadata[-1]['speech_filepath'] # clean speech\n",
"noise_filepath = test_metadata[-1]['noise_filepath'] # corresponding noise\n",
"\n",
"# Load audio\n",
"noisy_signal = AudioSegment.from_file(noisy_filepath, target_sr=sample_rate).samples\n",
"speech_signal = AudioSegment.from_file(speech_filepath, target_sr=sample_rate).samples\n",
"\n",
"# Move to device\n",
"device = 'cuda' if accelerator == 'gpu' else 'cpu'\n",
"enhancement_model = enhancement_model.to(device)\n",
"\n",
"# Process using the model\n",
"noisy_tensor = torch.tensor(noisy_signal).reshape(1, 1, -1).to(device) # (batch, channel, time)\n",
"with torch.no_grad():\n",
" output_tensor = enhancement_model(input_signal=noisy_tensor)\n",
"output_signal = output_tensor[0][0].detach().cpu().numpy()"
]
},
{
"cell_type": "markdown",
"id": "90f6c0d0",
"metadata": {
"id": "90f6c0d0"
},
"source": [
"Signals can be easily plotted and signal metrics can be calculated for the given example."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "4aec5c64",
"metadata": {
"id": "4aec5c64"
},
"outputs": [],
"source": [
"# Show noisy and clean signals\n",
"show_metrics(signal=noisy_signal, reference=speech_signal, tag='Noisy signal', sample_rate=sample_rate)\n",
"show_metrics(signal=output_signal, reference=speech_signal, tag='Output signal', sample_rate=sample_rate)\n",
"\n",
"# Show signals\n",
"show_signal(speech_signal, tag='Speech signal')\n",
"show_signal(noisy_signal, tag='Noisy signal')\n",
"show_signal(output_signal, tag='Output signal')\n",
"\n",
"# Play audio\n",
"print('Speech signal')\n",
"ipd.display(ipd.Audio(speech_signal, rate=sample_rate))\n",
"\n",
"print('Noisy signal')\n",
"ipd.display(ipd.Audio(noisy_signal, rate=sample_rate))\n",
"\n",
"print('Output signal')\n",
"ipd.display(ipd.Audio(output_signal, rate=sample_rate))"
]
},
{
"cell_type": "markdown",
"id": "cd6da1f0",
"metadata": {
"id": "cd6da1f0"
},
"source": [
"## Next steps\n",
"This is a simple tutorial which can serve as a starting point for prototyping and experimentation with audio-to-audio models.\n",
"A processed audio output can be used, for example, for ASR or TTS.\n",
"\n",
"For more details about NeMo models and applications in in ASR and TTS, we recommend you checkout other tutorials next:\n",
"\n",
"* [NeMo fundamentals](https://colab.research.google.com/github/NVIDIA/NeMo/blob/stable/tutorials/00_NeMo_Primer.ipynb)\n",
"* [NeMo models](https://colab.research.google.com/github/NVIDIA/NeMo/blob/stable/tutorials/01_NeMo_Models.ipynb)\n",
"* [Speech Recognition](https://colab.research.google.com/github/NVIDIA/NeMo/blob/stable/tutorials/asr/ASR_with_NeMo.ipynb)\n",
"* [Speech Synthesis](https://colab.research.google.com/github/NVIDIA/NeMo/blob/stable/tutorials/tts/Inference_ModelSelect.ipynb)"
]
},
{
"cell_type": "markdown",
"id": "46a855e3",
"metadata": {
"id": "46a855e3"
},
"source": [
"## References\n",
"\n",
"[1] M. Remane, R. R. Nalla and A. Dantrey, \"SEASR: Speech Enhancement for Automatic Speech Recognition Systems using Convolution Recurrent Neural Network with Residual Connections,\" 2024 IEEE 5th Women in Technology Conference (WINTECHCON), Bengaluru, India, 2024, pp. 1-5, doi: 10.1109/Wintechcon61988.2024.10837982\n",
"\n",
"[2] V. Panayotov, G. Chen, D. Povery, S. Khudanpur, \"LibriSpeech: An ASR corpus based on public domain audio books,\" ICASSP 2015\n",
"\n",
"[3] J. Thieman, N. Ito, V. Emmanuel, \"DEMAND: collection of multi-channel recordings of acoustic noise in diverse environments,\" ICA 2013\n",
"\n",
"[4] https://github.com/Lightning-AI/torchmetrics\n"
]
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"gpuType": "T4",
"provenance": []
},
"gpuClass": "standard",
"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.10"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
tutorials/audio/speech_enhancement/Speech_Enhancement_with_NeMo.ipynb
+1318
View File
File diff suppressed because it is too large Load Diff
tutorials/audio/speech_enhancement/Speech_Enhancement_with_Online_Augmentation.ipynb
+984
View File
@@ -0,0 +1,984 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "3b53ae0a-4cb5-44fd-bcde-5ce7abce79e9",
"metadata": {
"id": "3b53ae0a-4cb5-44fd-bcde-5ce7abce79e9"
},
"source": [
"# Introduction\n",
"\n",
"The goal of this tutorial is to demonstrate the basic steps required to setup and train a simple single-channel speech enhancement model in NeMo using online augmentation with noise and room impulse responce (RIR). Online augmentation is performed using a dataloader based on Lhotse speech data processing toolkit [1].\n",
"\n",
"\n",
"This notebook covers the following steps:\n",
"\n",
"* Download speech, noise and RIR data\n",
"* Prepare Lhotse manifests for speech, noise and RIR data\n",
"* Prepare fixed validation set by mixing speech, noise and RIR data\n",
"* Configure and train a simple single-output model\n",
"\n",
"Note that this tutorial is only for demonstration purposes.\n",
"To achieve best performance for a particular use case, carefully prepared data and more advanced models should be used.\n",
"\n",
"*Disclaimer:*\n",
"User is responsible for checking the content of datasets and the applicable licenses and determining if suitable for the intended use."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "3d603b38-b63c-42a6-92b3-575925149bde",
"metadata": {
"id": "3d603b38-b63c-42a6-92b3-575925149bde"
},
"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",
"\n",
"GIT_USER = 'NVIDIA-NeMo'\n",
"BRANCH = 'main'\n",
"\n",
"if 'google.colab' in str(get_ipython()):\n",
"\n",
" # Install dependencies\n",
" !pip install wget\n",
" !apt-get install sox libsndfile1 ffmpeg\n",
" !pip install text-unidecode\n",
" !pip install matplotlib>=3.3.2\n",
"\n",
" ## Install NeMo\n",
" !python -m pip install \"nemo_toolkit[all] @ git+https://github.com/{GIT_USER}/Speech.git@{BRANCH}\"\n",
"\n",
" ## Install TorchAudio\n",
" !pip install torchaudio>=0.13.0 -f https://download.pytorch.org/whl/torch_stable.html"
]
},
{
"cell_type": "markdown",
"id": "bdac6ac8-a21f-4ea8-8484-41a6d187a2fe",
"metadata": {
"id": "bdac6ac8-a21f-4ea8-8484-41a6d187a2fe"
},
"source": [
"The following cell will take care of the necessary imports and prepare utility functions used throughout the notebook."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "baca13c6-b0ed-429b-93b6-77249fdf4710",
"metadata": {
"id": "baca13c6-b0ed-429b-93b6-77249fdf4710"
},
"outputs": [],
"source": [
"import glob\n",
"import librosa\n",
"import os\n",
"import torch\n",
"import tqdm\n",
"from itertools import islice\n",
"\n",
"import IPython.display as ipd\n",
"import matplotlib.pyplot as plt\n",
"import numpy as np\n",
"import lightning.pytorch as pl\n",
"import soundfile as sf\n",
"from pathlib import Path\n",
"from omegaconf import OmegaConf, open_dict\n",
"from sklearn.model_selection import train_test_split\n",
"from torchmetrics.functional.audio import signal_distortion_ratio, scale_invariant_signal_distortion_ratio\n",
"from lhotse import CutSet, RecordingSet, Recording, MonoCut\n",
"from lhotse.recipes import (\n",
" download_rir_noise,\n",
" prepare_rir_noise,\n",
" download_librispeech,\n",
" prepare_librispeech\n",
")\n",
"\n",
"from nemo.collections.common.data.lhotse import get_lhotse_dataloader_from_config\n",
"from nemo.collections.audio.data.audio_to_audio_lhotse import LhotseAudioToTargetDataset"
]
},
{
"cell_type": "markdown",
"id": "fc0fafa4-bc65-4066-8e75-d740e2d15259",
"metadata": {
"id": "fc0fafa4-bc65-4066-8e75-d740e2d15259"
},
"source": [
"Utility functions for displaying signals and metrics"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "49720c06-3b4d-45b6-a054-73800c3bddea",
"metadata": {
"id": "49720c06-3b4d-45b6-a054-73800c3bddea"
},
"outputs": [],
"source": [
"def show_signal(signal: np.ndarray, sample_rate: int = 16000, tag: str = 'Signal'):\n",
" \"\"\"Show the time-domain signal and its spectrogram.\n",
" \"\"\"\n",
" fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(10, 2.5))\n",
"\n",
" # show waveform\n",
" t = np.arange(0, len(signal)) / sample_rate\n",
"\n",
" ax[0].plot(t, signal)\n",
" ax[0].set_xlim(0, t.max())\n",
" ax[0].grid()\n",
" ax[0].set_xlabel('time / s')\n",
" ax[0].set_ylabel('amplitude')\n",
" ax[0].set_title(tag)\n",
"\n",
" n_fft = 1024\n",
" hop_length = 256\n",
"\n",
" D = librosa.amplitude_to_db(np.abs(librosa.stft(signal, n_fft=n_fft, hop_length=hop_length)), ref=np.max)\n",
" img = librosa.display.specshow(D, y_axis='linear', x_axis='time', sr=sample_rate, n_fft=n_fft, hop_length=hop_length, ax=ax[1])\n",
" ax[1].set_title(tag)\n",
"\n",
" plt.tight_layout()\n",
" plt.colorbar(img, format=\"%+2.f dB\", ax=ax)\n",
"\n",
"def show_metrics(signal: np.ndarray, reference: np.ndarray, sample_rate: int = 16000, tag: str = 'Signal'):\n",
" \"\"\"Show metrics for the time-domain signal and the reference signal.\n",
" \"\"\"\n",
" sdr = signal_distortion_ratio(preds=torch.tensor(signal), target=torch.tensor(reference))\n",
" sisdr = scale_invariant_signal_distortion_ratio(preds=torch.tensor(signal), target=torch.tensor(reference))\n",
" print(tag)\n",
" print('\\tsdr: ', sdr.item())\n",
" print('\\tsisdr:', sisdr.item())"
]
},
{
"cell_type": "markdown",
"id": "653eb22f-b09e-4421-8028-05d123fc47a5",
"metadata": {
"id": "653eb22f-b09e-4421-8028-05d123fc47a5"
},
"source": [
"### Data preparation"
]
},
{
"cell_type": "markdown",
"id": "42deb721-f734-43ea-bfef-931733d6379b",
"metadata": {
"id": "42deb721-f734-43ea-bfef-931733d6379b"
},
"source": [
"In this notebook, it is assumed that all audio will be resampled to 16kHz and the data and configuration will be stored under `root_dir` as defined below."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "577d4be2-10f8-4050-8f40-b2566ef9dfef",
"metadata": {
"id": "577d4be2-10f8-4050-8f40-b2566ef9dfef"
},
"outputs": [],
"source": [
"# sample rate used throughout the notebook\n",
"sample_rate = 16000\n",
"\n",
"# root directory for data preparation, configurations, etc\n",
"root_dir = Path('./')\n",
"\n",
"# data directory\n",
"data_dir = root_dir / 'data'\n",
"data_dir.mkdir(exist_ok=True)\n",
"\n",
"# scripts directory\n",
"scripts_dir = root_dir / 'scripts'\n",
"scripts_dir.mkdir(exist_ok=True)"
]
},
{
"cell_type": "markdown",
"id": "9f6dc925-04b1-4118-8092-c3107a734f4f",
"metadata": {
"id": "9f6dc925-04b1-4118-8092-c3107a734f4f"
},
"source": [
"Create dictionary with paths for all of the manifests files which will be stored under `data_dir`."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "da232926-a529-4409-b983-8a921c578c91",
"metadata": {
"id": "da232926-a529-4409-b983-8a921c578c91"
},
"outputs": [],
"source": [
"dataset_manifest = {\n",
" 'speech_train': data_dir / 'libri_cuts_train.jsonl.gz',\n",
" 'speech_val': data_dir / 'libri_cuts_val.jsonl.gz',\n",
" 'noise_train': data_dir / 'demand_cuts_train.jsonl.gz',\n",
" 'noise_val': data_dir / 'demand_cuts_val.jsonl.gz',\n",
" 'rir_train': data_dir / 'rir_recordings_train.jsonl.gz',\n",
" 'rir_val': data_dir / 'rir_recordings_val.jsonl.gz',\n",
" 'noisy_val': data_dir / 'noisy_cuts_val.jsonl.gz'\n",
"}"
]
},
{
"cell_type": "markdown",
"id": "d0e19997-b151-47f5-b19b-97d8c7e65e2e",
"metadata": {
"id": "d0e19997-b151-47f5-b19b-97d8c7e65e2e"
},
"source": [
"In this tutorial, a subset of LibriSpeech dataset [2] will be downloaded and used as the speech material.\n",
"\n",
"To use a dataset with the Lhotse dataloader, we need to create manifest files from Lhotse cuts (refer to [3] for the details). In this cell, we first download and prepare the LibriSpeech dataset in a Lhotse format and then save it as manifest files for training and validation sets. Note that the target recording in the speech enhancement task is the original (unchanged) clean speech signal, which is defined under the custom field \"target_recording\" in the cuts."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "d4156a68-1906-4a4c-9ecf-cab5d79eb13b",
"metadata": {
"id": "d4156a68-1906-4a4c-9ecf-cab5d79eb13b"
},
"outputs": [],
"source": [
"libri_variant = 'mini_librispeech'\n",
"speech_dir = data_dir / 'speech'\n",
"\n",
"libri_root = download_librispeech(speech_dir, dataset_parts=libri_variant)\n",
"\n",
"# Use script from Lhotse to prepate Librispeech dataset to Lhotse format\n",
"libri = prepare_librispeech(\n",
" libri_root, dataset_parts=libri_variant,\n",
")\n",
"cuts_train = CutSet.from_manifests(**libri[\"train-clean-5\"]).trim_to_supervisions()\n",
"cuts_val = CutSet.from_manifests(**libri[\"dev-clean-2\"]).trim_to_supervisions()\n",
"\n",
"# Save the manifest with a custom \"target_recording\"\n",
"with CutSet.open_writer(dataset_manifest['speech_train']) as writer:\n",
" for cut in cuts_train:\n",
" cut.target_recording = cut.recording\n",
" writer.write(cut)\n",
"\n",
"with CutSet.open_writer(dataset_manifest['speech_val']) as writer:\n",
" for cut in cuts_val:\n",
" cut.target_recording = cut.recording\n",
" writer.write(cut)"
]
},
{
"cell_type": "markdown",
"id": "6f5a0eb7-b913-4323-a8ce-1da77521730b",
"metadata": {
"id": "6f5a0eb7-b913-4323-a8ce-1da77521730b"
},
"source": [
"During the training phase, noise data will be used for online augmentation by mixing it with the downloaded speech on-the-fly. During the validation and test phases, the noise will be used to create fixed sets.\n",
"\n",
"The following cell will download and prepare the noise data using a subset of the DEMAND dataset [4]."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "ba4dcfe4-5614-438d-9c2b-73c0fbf3bc97",
"metadata": {
"id": "ba4dcfe4-5614-438d-9c2b-73c0fbf3bc97"
},
"outputs": [],
"source": [
"noise_dir = data_dir / 'noise'\n",
"noise_data_set = 'STRAFFIC,PSTATION'\n",
"\n",
"# Copy script\n",
"get_demand_script = os.path.join(scripts_dir, 'get_demand_data.py')\n",
"if not os.path.exists(get_demand_script):\n",
" !wget -P $scripts_dir https://raw.githubusercontent.com/{GIT_USER}/NeMo/{BRANCH}/scripts/dataset_processing/get_demand_data.py\n",
"\n",
"if not noise_dir.is_dir():\n",
" noise_dir.mkdir(exist_ok=True)\n",
" !python {get_demand_script} --data_root={noise_dir} --data_sets={noise_data_set}\n",
"else:\n",
" print('Noise directory already exists in:', noise_dir)\n",
"\n",
"noise_dir = data_dir / 'noise'\n",
"demand_recordings = RecordingSet.from_dir(noise_dir, pattern='*.wav')\n",
"\n",
"demand_cuts = CutSet.from_manifests(recordings=demand_recordings)\n",
"shuffled_demand_cuts = demand_cuts.shuffle()\n",
"\n",
"demand_cuts_train = shuffled_demand_cuts.subset(last=len(shuffled_demand_cuts)-3)\n",
"demand_cuts_val = shuffled_demand_cuts.subset(first=3)\n",
"\n",
"demand_cuts_train.to_file(dataset_manifest['noise_train'])\n",
"demand_cuts_val.to_file(dataset_manifest['noise_val'])"
]
},
{
"cell_type": "markdown",
"id": "f9429dd5-fedc-486d-92f1-5a0255d600e6",
"metadata": {
"id": "f9429dd5-fedc-486d-92f1-5a0255d600e6"
},
"source": [
"The following cell will download and prepare a simulated subset from room impulse responses dataset, described in the following paper [5]."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "3e83f5d2-26cd-44cb-9efe-4eafcd709bfe",
"metadata": {
"id": "3e83f5d2-26cd-44cb-9efe-4eafcd709bfe"
},
"outputs": [],
"source": [
"rir_recordings = RecordingSet()\n",
"rir_raw_dir = download_rir_noise(data_dir)\n",
"rirs = prepare_rir_noise(rir_raw_dir, parts=[\"sim_rir\"])\n",
"rir_recordings = rirs[\"sim_rir\"][\"recordings\"]\n",
"shuffled_rir_recordings = rir_recordings.shuffle()\n",
"\n",
"rir_val_part = int(len(rir_recordings) * 0.1)\n",
"rir_train_part = len(rir_recordings) - rir_val_part\n",
"\n",
"rir_recordings_train = shuffled_rir_recordings.subset(last=rir_train_part)\n",
"rir_recordings_val = shuffled_rir_recordings.subset(first=rir_val_part)\n",
"\n",
"rir_recordings_train.to_file(dataset_manifest['rir_train'])\n",
"rir_recordings_val.to_file(dataset_manifest['rir_val'])"
]
},
{
"cell_type": "markdown",
"id": "b568a03e-57a1-412f-a3a2-85af3a1e3a20",
"metadata": {
"id": "b568a03e-57a1-412f-a3a2-85af3a1e3a20"
},
"source": [
"For this tutorial, a single-channel noisy validation set is constructed by adding speech and noise.\n",
"\n",
"The following block will use based on Lhotse data loader from NeMo to create fixed noisy validation set and save it do `data/val` folder."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "62dc5755-ccc5-4eb8-80d4-d59b418406a3",
"metadata": {
"id": "62dc5755-ccc5-4eb8-80d4-d59b418406a3"
},
"outputs": [],
"source": [
"# Create the cofing for the Lhotse data loader\n",
"val_noise_config = {\n",
"'cuts_path': dataset_manifest['speech_val'].as_posix(), # path to Lhotse cuts manifest with speech signals for augmentation\n",
"'sample_rate': 16000,\n",
"'batch_size': 1,\n",
"'rir_enabled': True, # enable room impulse response augmentation\n",
"'rir_path': dataset_manifest['rir_val'].as_posix(), # path to Lhotse recordings manifest with room impulse response signals\n",
"'rir_prob': 1.0, # probability of applying RIR augmentation\n",
"'noise_path': dataset_manifest['noise_val'].as_posix(), # path to Lhotse cuts manifest with noise signals\n",
"'noise_mix_prob': 1.0, # probability of applying noise augmentation\n",
"'noise_snr': (0, 20), # range of speech-to-noise ratio for the noise augmentation\n",
"'shuffle': False\n",
"}\n",
"\n",
"# Instantiate the data loader\n",
"dl = get_lhotse_dataloader_from_config(\n",
"OmegaConf.create(val_noise_config), global_rank=0, world_size=1, dataset=LhotseAudioToTargetDataset()\n",
")\n",
"\n",
"# Define number of samples for the validation set\n",
"num_examples = 100\n",
"print(f'Get {num_examples} samples for the validation set')\n",
"samples = [sample for sample in islice(dl, num_examples)]\n",
"\n",
"\n",
"# Create folders for saving noisy (input) and clean (target) samples\n",
"val_dir = data_dir / 'val'\n",
"val_noisy_dir = val_dir / 'noisy'\n",
"val_clean_dir = val_dir / 'clean'\n",
"\n",
"val_dir.mkdir(exist_ok=True)\n",
"val_noisy_dir.mkdir(exist_ok =True)\n",
"val_clean_dir.mkdir(exist_ok=True)\n",
"\n",
"val_noisy_basename = 'val_noisy_fileid'\n",
"val_clean_basename = 'val_clean_fileid'\n",
"\n",
"with CutSet.open_writer(dataset_manifest['noisy_val']) as writer:\n",
" for n, sample in enumerate(samples):\n",
" noisy, clean = sample['input_signal'].numpy()[0], sample['target_signal'].numpy()[0]\n",
" #Save\n",
" sf.write(val_noisy_dir / f'{val_noisy_basename}_{str(n)}.wav', noisy, samplerate=val_noise_config['sample_rate'])\n",
" sf.write(val_clean_dir / f'{val_clean_basename}_{str(n)}.wav', clean, samplerate=val_noise_config['sample_rate'])\n",
" noisy_rec = Recording.from_file(val_noisy_dir / f'{val_noisy_basename}_{str(n)}.wav')\n",
" clean_rec = Recording.from_file(val_clean_dir / f'{val_clean_basename}_{str(n)}.wav')\n",
"\n",
" val_cut = MonoCut(id=noisy_rec.id,\n",
" start=0,\n",
" duration=noisy_rec.duration,\n",
" channel=0,\n",
" recording=noisy_rec)\n",
" val_cut.target_recording = clean_rec\n",
" writer.write(val_cut)"
]
},
{
"cell_type": "markdown",
"id": "f108c5d2-c87b-47a7-8c2d-d363a9234abb",
"metadata": {
"id": "f108c5d2-c87b-47a7-8c2d-d363a9234abb"
},
"source": [
"### Model configuration\n",
"\n",
"Here, a simple encoder-mask-decoder model will be used to process the noisy input signal and produce an enhanced output signal.\n",
"\n",
"In general, an encoder-mask-decoder model can be configured using `EncMaskDecAudioToAudioModel` class, which is depicted in the following block diagram."
]
},
{
"cell_type": "markdown",
"id": "c0ff2bff-1637-4a17-85b2-c92307a4d8d7",
"metadata": {
"id": "c0ff2bff-1637-4a17-85b2-c92307a4d8d7"
},
"source": [
"<img src=\"https://github.com/NVIDIA/NeMo/releases/download/v1.18.0/encmaskdecoder_model.png\" alt=\"encmaskdecoder_model\" style=\"width: 800px;\"/>"
]
},
{
"cell_type": "markdown",
"id": "470b0a50-fe6c-4ded-9d27-64db20f83cad",
"metadata": {
"id": "470b0a50-fe6c-4ded-9d27-64db20f83cad"
},
"source": [
"The model structure can briefly be described as follows:\n",
"* Input to the model is a time-domain signal.\n",
"* Encoder transforms the input signal to the analysis domain.\n",
"* Mask estimator estimates a mask used to generate the output signal.\n",
"* Mask processor combines the estimated mask and the encoded input to produce the encoded output.\n",
"* Decoder transforms the encoded output into a time-domain signal.\n",
"* Output is a time-domain signal."
]
},
{
"cell_type": "markdown",
"id": "4cf1cfde-913f-42da-a9e0-fdb5ae8e50c5",
"metadata": {
"id": "4cf1cfde-913f-42da-a9e0-fdb5ae8e50c5"
},
"source": [
"For this example, the model will be configured to use a fixed short-time Fourier transform-based encoder and decoder, and the mask will be estimated using a recurrent neural network. The model used here is depicted in the following block diagram."
]
},
{
"cell_type": "markdown",
"id": "8ed37321-2e8b-4d8e-90bf-efc0897517e5",
"metadata": {
"id": "8ed37321-2e8b-4d8e-90bf-efc0897517e5"
},
"source": [
"<img src=\"https://github.com/NVIDIA/NeMo/releases/download/v1.18.0/single_output_example_model.png\" alt=\"single_output_example_model\" style=\"width: 1000px;\"/>"
]
},
{
"cell_type": "markdown",
"id": "36596c01-db7d-402f-a5df-0fb9b4641b08",
"metadata": {
"id": "36596c01-db7d-402f-a5df-0fb9b4641b08"
},
"source": [
"In this particular configuration, the model structure can be described as follows:\n",
"* `AudioToSpectrogram` implements the analysis STFT transform.\n",
"* `MaskEstimatorRNN` is a mask estimator using RNNs.\n",
"* `MaskReferenceChannel` is a simple processor which applies the estimated mask on the reference channel. In this tutorial, the input signal has only a single channel, so the reference channel will be set to `0`.\n",
"* `SpectrogramToAudio` implements the synthesis STFT transform."
]
},
{
"cell_type": "markdown",
"id": "03546307-1c4e-4775-a754-276c9a69be5c",
"metadata": {
"id": "03546307-1c4e-4775-a754-276c9a69be5c"
},
"source": [
"The following cell will load and show the default configuration for the model depicted above."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "03952acf-1792-4470-9120-dddf20a9f646",
"metadata": {
"id": "03952acf-1792-4470-9120-dddf20a9f646",
"scrolled": true
},
"outputs": [],
"source": [
"config_dir = root_dir / 'conf'\n",
"config_dir.mkdir(exist_ok=True)\n",
"\n",
"config_path = config_dir / 'masking_with_online_augmentation.yaml'\n",
"\n",
"if not config_path.is_file():\n",
" !wget https://raw.githubusercontent.com/{GIT_USER}/NeMo/{BRANCH}/examples/audio/conf/masking_with_online_augmentation.yaml -P {config_dir.as_posix()}\n",
"\n",
"config = OmegaConf.load(config_path)\n",
"config = OmegaConf.to_container(config, resolve=True)\n",
"config = OmegaConf.create(config)\n",
"\n",
"print('Loaded config')\n",
"print(OmegaConf.to_yaml(config))"
]
},
{
"cell_type": "markdown",
"id": "1c872345-3bfd-4933-8c51-dbb62fce790a",
"metadata": {
"id": "1c872345-3bfd-4933-8c51-dbb62fce790a"
},
"source": [
"Training dataset is configured with the following parameters\n",
"* `cuts_path` points to a Lhotse manifest file, containing speech samples\n",
"* `noise_path` poins to a Lhotse manifest file, containing noise samples\n",
"* `noise_mix_prob` defines the probabilty with which noise will be added during training\n",
"* `noise_snr` defines an SNR range for mixing noise samples\n",
"* `rir_enabled` enables room impulse response agmentation\n",
"* `rir_path` points to a Lhotse manifest file, containing RIR samples\n",
"* `rir_prob` defines the probabilty with which RIR will be added during training\n",
" \n",
"For the validation and test sets only `cuts_path` parameter is used since the `val` manifest already contains noisy and clean samples"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e19c4209-b10b-4924-a3de-e792de57c313",
"metadata": {
"id": "e19c4209-b10b-4924-a3de-e792de57c313"
},
"outputs": [],
"source": [
"# Setup training dataset\n",
"config.model.train_ds.cuts_path = dataset_manifest['speech_train'].as_posix() # path to Lhotse cuts manifest with speech signals for augmentation\n",
"config.model.train_ds.noise_path = dataset_manifest['noise_train'].as_posix() # path to Lhotse cuts manifest with noise signals\n",
"config.model.train_ds.noise_mix_prob = 1.0 # probability of applying noise augmentation\n",
"config.model.train_ds.noise_snr = (0, 20) # range of speech-to-noise ratio for the noise augmentation\n",
"config.model.train_ds.rir_enabled = True # enable room impulse response augmentation\n",
"config.model.train_ds.rir_path = dataset_manifest['rir_val'].as_posix() # path to Lhotse recordings manifest with room impulse response signals\n",
"config.model.train_ds.rir_prob = 1.0 # probability of applying RIR augmentation\n",
"\n",
"config.model.validation_ds.cuts_path = dataset_manifest['noisy_val'].as_posix() # fixed noisy validation set\n",
"\n",
"config.model.test_ds.cuts_path = dataset_manifest['noisy_val'].as_posix() # fixed noisy test set\n",
"\n",
"\n",
"print(\"Train dataset config:\")\n",
"print(OmegaConf.to_yaml(config.model.train_ds))"
]
},
{
"cell_type": "markdown",
"id": "534aef42-d3a7-4477-8f89-8fe780f2b5f7",
"metadata": {
"id": "534aef42-d3a7-4477-8f89-8fe780f2b5f7"
},
"source": [
"Metrics for validation and test set are configured in the following cell.\n",
"\n",
"In this tutorial, signal-to-distortion ratio (SDR) and scale-invariant SDR from torch metrics are used [5]."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "522f1065-fe23-4101-b6a9-0706cc8db389",
"metadata": {
"id": "522f1065-fe23-4101-b6a9-0706cc8db389"
},
"outputs": [],
"source": [
"# Setup metrics to compute on validation and test sets\n",
"metrics = OmegaConf.create({\n",
" 'sisdr': {\n",
" '_target_': 'torchmetrics.audio.ScaleInvariantSignalDistortionRatio',\n",
" },\n",
" 'sdr': {\n",
" '_target_': 'torchmetrics.audio.SignalDistortionRatio',\n",
" }\n",
"})\n",
"config.model.metrics.val = metrics\n",
"config.model.metrics.test = metrics\n",
"\n",
"print(\"Metrics config:\")\n",
"print(OmegaConf.to_yaml(config.model.metrics))"
]
},
{
"cell_type": "markdown",
"id": "4771d02f-3b01-481c-8894-bb43644a941a",
"metadata": {
"id": "4771d02f-3b01-481c-8894-bb43644a941a"
},
"source": [
"### Trainer configuration\n",
"NeMo models are primarily PyTorch Lightning modules and therefore are entirely compatible with the PyTorch Lightning ecosystem."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "2208b407-07e2-4bb0-bfad-2ef488acf217",
"metadata": {
"id": "2208b407-07e2-4bb0-bfad-2ef488acf217"
},
"outputs": [],
"source": [
"print(\"Trainer config:\")\n",
"print(OmegaConf.to_yaml(config.trainer))"
]
},
{
"cell_type": "markdown",
"id": "f2dac35a-b487-4037-907d-326f321564ca",
"metadata": {
"id": "f2dac35a-b487-4037-907d-326f321564ca"
},
"source": [
"We can modify some trainer configs for this tutorial.\n",
"Most importantly, the number of epochs is set to a small value, to limit the runtime for the purpose of this example."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "2c18dbd2-ad22-4eae-912a-56d13aa74a6b",
"metadata": {
"id": "2c18dbd2-ad22-4eae-912a-56d13aa74a6b"
},
"outputs": [],
"source": [
"# Checks if we have GPU available and uses it\n",
"accelerator = 'gpu' if torch.cuda.is_available() else 'cpu'\n",
"config.trainer.devices = 1\n",
"config.trainer.accelerator = accelerator\n",
"\n",
"# Reduces maximum number of epochs for quick demonstration\n",
"config.trainer.max_epochs = 30\n",
"\n",
"# Remove distributed training flags\n",
"config.trainer.strategy = 'auto'\n",
"\n",
"# Instantiate the trainer\n",
"trainer = pl.Trainer(**config.trainer)"
]
},
{
"cell_type": "markdown",
"id": "c1198508-7fb5-4ba7-a4e5-ad15153037a1",
"metadata": {
"id": "c1198508-7fb5-4ba7-a4e5-ad15153037a1"
},
"source": [
"### Experiment manager\n",
"\n",
"NeMo has an experiment manager that handles logging and checkpointing."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "72642c8c-97f5-4b6c-9e63-030c7a07810c",
"metadata": {
"id": "72642c8c-97f5-4b6c-9e63-030c7a07810c",
"scrolled": true
},
"outputs": [],
"source": [
"from nemo.utils.exp_manager import exp_manager\n",
"\n",
"exp_dir = exp_manager(trainer, config.get(\"exp_manager\", None))\n",
"# The exp_dir provides a path to the current experiment for easy access\n",
"\n",
"print(\"Experiment directory:\")\n",
"print(exp_dir)"
]
},
{
"cell_type": "markdown",
"id": "00562a21-0c0f-4c53-b89e-344431e75a42",
"metadata": {
"id": "00562a21-0c0f-4c53-b89e-344431e75a42"
},
"source": [
"### Model instantiation"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "0b73c3de-6ded-491d-8026-8a093f7ffc16",
"metadata": {
"id": "0b73c3de-6ded-491d-8026-8a093f7ffc16",
"scrolled": true
},
"outputs": [],
"source": [
"from nemo.collections.audio.models.enhancement import EncMaskDecAudioToAudioModel\n",
"\n",
"enhancement_model = EncMaskDecAudioToAudioModel(cfg=config.model, trainer=trainer)"
]
},
{
"cell_type": "markdown",
"id": "a8f4373e-6aca-41d3-a69d-5bef53a88f20",
"metadata": {
"id": "a8f4373e-6aca-41d3-a69d-5bef53a88f20"
},
"source": [
"### Training\n",
"Create a Tensorboard visualization to monitor progress"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "a253ce7c-6cca-441f-818e-1c8a2f00a9f5",
"metadata": {
"id": "a253ce7c-6cca-441f-818e-1c8a2f00a9f5"
},
"outputs": [],
"source": [
"try:\n",
" from google import colab\n",
" COLAB_ENV = True\n",
"except (ImportError, ModuleNotFoundError):\n",
" COLAB_ENV = False\n",
"\n",
"# Load the TensorBoard notebook extension\n",
"if COLAB_ENV:\n",
" %load_ext tensorboard\n",
" %tensorboard --logdir {exp_dir}\n",
"else:\n",
" print(\"To use tensorboard, please use this notebook in a Google Colab environment.\")"
]
},
{
"cell_type": "markdown",
"id": "9183a35d-8332-49f6-96b3-3d14dd70ff4e",
"metadata": {
"id": "9183a35d-8332-49f6-96b3-3d14dd70ff4e"
},
"source": [
"Training can be started using `trainer.fit`:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "0eab4e7e-3004-49fa-a09b-10dfd12ba294",
"metadata": {
"id": "0eab4e7e-3004-49fa-a09b-10dfd12ba294",
"scrolled": true
},
"outputs": [],
"source": [
"trainer.fit(enhancement_model)"
]
},
{
"cell_type": "markdown",
"id": "af906f16-626c-4ca4-a465-fb0cca8d514a",
"metadata": {
"id": "af906f16-626c-4ca4-a465-fb0cca8d514a"
},
"source": [
"After the training is completed, the configured metrics can be easily computed on the test set as follows:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "589af425-18b2-4445-8c88-d3161e9928c5",
"metadata": {
"id": "589af425-18b2-4445-8c88-d3161e9928c5"
},
"outputs": [],
"source": [
"trainer.test(enhancement_model, ckpt_path=None)"
]
},
{
"cell_type": "markdown",
"id": "cce96446-a969-4c6f-a865-968ee61119a8",
"metadata": {
"id": "cce96446-a969-4c6f-a865-968ee61119a8"
},
"source": [
"### Inference\n",
"\n",
"The following cell provides an example of inference on an single audio file.\n",
"For simplicity, the audio file information is taken from the test dataset."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "9f29ade5-424e-4a7a-81b7-dc3c5cc90bf7",
"metadata": {
"id": "9f29ade5-424e-4a7a-81b7-dc3c5cc90bf7"
},
"outputs": [],
"source": [
"# Load 10 samples from test_dataloader\n",
"samples = [sample for sample in islice(enhancement_model.test_dataloader(), 10)]\n",
"\n",
"# Different sample can be used via list index\n",
"sample = samples[0]\n",
"\n",
"noisy_tensor = sample['input_signal']\n",
"speech_tensor = sample['target_signal']\n",
"\n",
"# Get the one-dimentional numpy signals for the plotting audio files and metrics calculation\n",
"noisy_signal = noisy_tensor.squeeze(0).numpy()\n",
"speech_signal = speech_tensor.squeeze(0).numpy()\n",
"\n",
"\n",
"# Move to device\n",
"device = 'cuda' if accelerator == 'gpu' else 'cpu'\n",
"enhancement_model = enhancement_model.to(device)\n",
"\n",
"# Process using the model\n",
"with torch.no_grad():\n",
" output_tensor, _ = enhancement_model(input_signal=noisy_tensor.unsqueeze(1).cuda())\n",
"output_signal = output_tensor[0][0].detach().cpu().numpy()"
]
},
{
"cell_type": "markdown",
"id": "5110ee6c-3939-4d70-a4d6-6422de31da51",
"metadata": {
"id": "5110ee6c-3939-4d70-a4d6-6422de31da51"
},
"source": [
"Signals can be easily plotted and signal metrics can be calculated for the given example."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "0275a032-969c-4e31-b8b5-8a61cb1f82ee",
"metadata": {
"id": "0275a032-969c-4e31-b8b5-8a61cb1f82ee"
},
"outputs": [],
"source": [
"# Show noisy and clean signals\n",
"show_metrics(signal=noisy_signal, reference=speech_signal, tag='Noisy signal', sample_rate=sample_rate)\n",
"show_metrics(signal=output_signal, reference=speech_signal, tag='Output signal', sample_rate=sample_rate)\n",
"\n",
"# Show signals\n",
"show_signal(speech_signal, tag='Speech signal')\n",
"show_signal(noisy_signal, tag='Noisy signal')\n",
"show_signal(output_signal, tag='Output signal')\n",
"\n",
"# Play audio\n",
"print('Speech signal')\n",
"ipd.display(ipd.Audio(speech_signal, rate=sample_rate))\n",
"\n",
"print('Noisy signal')\n",
"ipd.display(ipd.Audio(noisy_signal, rate=sample_rate))\n",
"\n",
"print('Output signal')\n",
"ipd.display(ipd.Audio(output_signal, rate=sample_rate))"
]
},
{
"cell_type": "markdown",
"id": "89eb3478-bf82-4fbd-aa7d-0a55edf57f3a",
"metadata": {
"id": "89eb3478-bf82-4fbd-aa7d-0a55edf57f3a"
},
"source": [
"## Next steps\n",
"This is a simple tutorial which can serve as a starting point for prototyping and experimentation with audio-to-audio models.\n",
"A processed audio output can be used, for example, for ASR or TTS.\n",
"\n",
"For more details about NeMo models and applications in in ASR and TTS, we recommend you checkout other tutorials next:\n",
"\n",
"* [NeMo fundamentals](https://colab.research.google.com/github/NVIDIA/NeMo/blob/stable/tutorials/00_NeMo_Primer.ipynb)\n",
"* [NeMo models](https://colab.research.google.com/github/NVIDIA/NeMo/blob/stable/tutorials/01_NeMo_Models.ipynb)\n",
"* [Speech Recognition](https://colab.research.google.com/github/NVIDIA/NeMo/blob/stable/tutorials/asr/ASR_with_NeMo.ipynb)\n",
"* [Speech Synthesis](https://colab.research.google.com/github/NVIDIA/NeMo/blob/stable/tutorials/tts/Inference_ModelSelect.ipynb)"
]
},
{
"cell_type": "markdown",
"id": "31a162ae-6b9e-40e4-a468-fbbad3360b11",
"metadata": {
"id": "31a162ae-6b9e-40e4-a468-fbbad3360b11"
},
"source": [
"## References\n",
"\n",
"[1] Żelasko, P., Povey, D., Trmal, J., & Khudanpur, S. (2021). Lhotse: a speech data representation library for the modern deep learning ecosystem. https://arxiv.org/abs/2110.12561\n",
"\n",
"[2] V. Panayotov, G. Chen, D. Povery, S. Khudanpur, \"LibriSpeech: An ASR corpus based on public domain audio books,\" ICASSP 2015\n",
"\n",
"[3] Lhotse documentation, https://lhotse.readthedocs.io/\n",
"\n",
"[4] J. Thieman, N. Ito, V. Emmanuel, \"DEMAND: collection of multi-channel recordings of acoustic noise in diverse environments,\" ICA 2013\n",
"\n",
"[5] T. Ko, V. Peddinti, D. Povey, M. L. Seltzer and S. Khudanpur, \"A study on data augmentation of reverberant speech for robust speech recognition,\" 2017 ICASSP\n",
"\n",
"[6] https://github.com/Lightning-AI/torchmetrics"
]
}
],
"metadata": {
"colab": {
"provenance": []
},
"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.12"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
tutorials/cloud/README.md
+7
View File
@@ -0,0 +1,7 @@
# NeMo Cloud Tutorials
The best way to get started with NeMo in the cloud.
## ASR
* [Quickstart: Training an ASR model on AWS SageMaker](aws/SageMaker_ASR_Training.ipynb)
* [Multi-Node Scaling: ASR Fine-Tuning on AWS SageMaker](aws/ASR_Finetuning_at_Scale_with_AWS_SageMaker.ipynb)
tutorials/cloud/aws/ASR_Finetuning_at_Scale_with_AWS_SageMaker.ipynb
+764
View File
@@ -0,0 +1,764 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "2e86a2b3",
"metadata": {},
"source": [
"# ASR Fine-Tuning at Scale using AWS SageMaker"
]
},
{
"cell_type": "markdown",
"id": "215e3d3c",
"metadata": {},
"source": [
"In this tutorial we show how you can train a NeMo ASR Model using [Amazon SageMaker](https://docs.aws.amazon.com/sagemaker/latest/dg/whatis.html) using the CommonVoice Esperanto dataset (~250hrs), and scale onto multiple GPUs and nodes with near-linear scaling.\n",
"\n",
"AWS SageMaker is useful for practitioners/researchers who are familiar with training locally or on a remote instance (via SSH). SageMaker also supports multi-GPU & Multi-node."
]
},
{
"cell_type": "markdown",
"id": "5e6d484f",
"metadata": {},
"source": [
"<img src=\"images/sagemaker_benchmark.png\" alt=\"sagemaker_benchmark\" style=\"width: 700px;\"/>"
]
},
{
"cell_type": "markdown",
"id": "81d7883d",
"metadata": {},
"source": [
"We fine-tune the [Conformer SSL Large (En)](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/nemo/models/ssl_en_conformer_large) model using the CommonVoice esperanto subset on a remote instance with a 8 GPUs (ml.p3.16xlarge). We use S3 to store the data and our checkpoints/logs.\n",
"\n",
"The overall steps are:\n",
"\n",
"1. Setup your AWS Credentials to access SageMaker\n",
"2. Download the source code we'll be running\n",
"3. Setup the CommonVoice Esperanto dataset\n",
"4. Upload data to S3\n",
"5. Configure the training job\n",
"6. Run training job on SageMaker\n",
"7. Download model, (Optional) Tensorboard Logs"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "ac621da0",
"metadata": {},
"outputs": [],
"source": [
"\"\"\"\n",
"You can run either this notebook locally (if you have all the dependencies) 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",
"## Install dependencies\n",
"!pip install wget\n",
"!apt-get install sox libsndfile1 ffmpeg\n",
"!pip install text-unidecode\n",
"!pip install matplotlib>=3.3.2\n",
"\n",
"## Install NeMo\n",
"BRANCH = 'main'\n",
"!python -m pip install \"nemo_toolkit[all] @ git+https://github.com/NVIDIA-NeMo/Speech.git@$BRANCH\"\n",
"\n",
"\"\"\"\n",
"Remember to restart the runtime for the kernel to pick up any upgraded packages (e.g. matplotlib)!\n",
"Alternatively, you can uncomment the exit() below to crash and restart the kernel, in the case\n",
"that you want to use the \"Run All Cells\" (or similar) option.\n",
"\"\"\"\n",
"# exit()"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "61c4fbe2",
"metadata": {},
"outputs": [],
"source": [
"pip install sagemaker awscli"
]
},
{
"cell_type": "markdown",
"id": "876f553d",
"metadata": {},
"source": [
"### 1. Setup SageMaker with AWS Credentials\n",
"\n",
"If you haven't setup your AWS credentials, you can setup using the AWS CLI.\n",
"You will need your access and Secret key, with permissions to use SageMaker and S3."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "1328482d",
"metadata": {},
"outputs": [],
"source": [
"!aws configure"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "01477d55",
"metadata": {},
"outputs": [],
"source": [
"%load_ext tensorboard\n",
"from pathlib import Path\n",
"import os\n",
"import sagemaker\n",
"import wget\n",
"import json\n",
"import re\n",
"\n",
"from tqdm import tqdm\n",
"from omegaconf import OmegaConf\n",
"from sagemaker import get_execution_role\n",
"from sagemaker.pytorch import PyTorch"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "405806f9",
"metadata": {},
"outputs": [],
"source": [
"sess = sagemaker.Session()"
]
},
{
"cell_type": "markdown",
"id": "7d099a96",
"metadata": {},
"source": [
"### 2. Download the NeMo source code\n",
"\n",
"SageMaker allows you to pass in your own source code, with an entrypoint script.\n",
"\n",
"Below we download the AWS NeMo `config.yaml` which contains our configuration, and the `speech_to_text_ctc_bpe.py` script to run training.\n",
"\n",
"Our folder structure will look like this:\n",
"\n",
" code/\n",
" speech_to_text_ctc_bpe.py\n",
" conf/\n",
" config.yaml\n",
" data/\n",
" manifests/\n",
" raw_data/"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "0b456c57",
"metadata": {},
"outputs": [],
"source": [
"root_dir = Path('./')\n",
"code_dir = root_dir / 'code/'\n",
"config_dir = code_dir / 'conf/'\n",
"data_dir = root_dir / 'data/'\n",
"\n",
"root_dir.mkdir(exist_ok=True)\n",
"code_dir.mkdir(exist_ok=True)\n",
"config_dir.mkdir(exist_ok=True)\n",
"data_dir.mkdir(exist_ok=True)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "721c7b7d",
"metadata": {},
"outputs": [],
"source": [
"config_path = str(config_dir / \"config.yaml\")\n",
"\n",
"# download scripts to format the data source.\n",
"wget.download(\"https://raw.githubusercontent.com/NVIDIA/NeMo/main/scripts/speech_recognition/convert_hf_dataset_to_nemo.py\", str(code_dir))\n",
"wget.download(\"https://raw.githubusercontent.com/NVIDIA/NeMo/main/scripts/speech_recognition/convert_to_tarred_audio_dataset.py\",\n",
" str(code_dir))\n",
"\n",
"# download scripts to run training\n",
"wget.download(\"https://raw.githubusercontent.com/NVIDIA/NeMo/main/examples/asr/conf/conformer/conformer_ctc_bpe.yaml\", config_path)\n",
"wget.download(\"https://raw.githubusercontent.com/NVIDIA/NeMo/main/examples/asr/asr_ctc/speech_to_text_ctc_bpe.py\",\n",
" str(code_dir))\n",
"\n",
"# download script to create tokenizer\n",
"wget.download(\"https://raw.githubusercontent.com/NVIDIA/NeMo/main/scripts/tokenizers/process_asr_text_tokenizer.py\",\n",
" str(code_dir))"
]
},
{
"cell_type": "markdown",
"id": "7934baab",
"metadata": {},
"source": [
"We also create a `requirements.txt` file within our source code to install NeMo."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "397d8eef",
"metadata": {},
"outputs": [],
"source": [
"with open(code_dir / 'requirements.txt', 'w') as f:\n",
" f.write(\"nemo_toolkit[all]\")"
]
},
{
"cell_type": "markdown",
"id": "a7bc7f51",
"metadata": {},
"source": [
"### 2.1 Initialize SageMaker within Training Script\n",
"\n",
"We provide a helper function that we require to be imported and run at the top of the training script.\n",
"\n",
"This installs and setups DDP for you. It also alleviates having to import a custom container, and can leverage all of the SageMaker containers. Rather than running this cell, you could also manually do this in your script."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e17535a7",
"metadata": {},
"outputs": [],
"source": [
"line = \"from nemo.utils.cloud import initialize_sagemaker; initialize_sagemaker()\"\n",
"with open(code_dir / \"speech_to_text_ctc.py\", 'r+') as f:\n",
" content = f.read()\n",
" f.seek(0, 0)\n",
" f.write(line.rstrip('\\r\\n') + '\\n' + content)"
]
},
{
"cell_type": "markdown",
"id": "ce3e5ae8",
"metadata": {},
"source": [
"### 3. Setup the CommonVoice Esperanto Dataset\n",
"\n",
"Mozilla Common Voice requires you to use huggingface datasets (`pip install datasets`) and register as a user to generate an API key.\n",
"\n",
"#### Authenticated Setup Steps\n",
"\n",
"Website steps:\n",
"- Visit https://huggingface.co/settings/profile\n",
"- Visit \"Access Tokens\" on list of items.\n",
"- Create new token - provide a name for the token and \"read\" access is sufficient.\n",
" - PRESERVE THAT TOKEN API KEY. You can copy that key for next step.\n",
"- Visit the HuggingFace Dataset page for Mozilla Common Voice\n",
" - There should be a section that asks you for your approval.\n",
" - Make sure you are logged in and then read that agreement.\n",
" - If and only if you agree to the text, then accept the terms."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "5cbca268",
"metadata": {},
"outputs": [],
"source": [
"# Paste your preserved HF TOKEN API KEY (from above) when asked.\n",
"!huggingface-cli login"
]
},
{
"cell_type": "markdown",
"id": "9b964276",
"metadata": {},
"source": [
"Now you should be logged in. When running the script, dont forget to set `use_auth_token=True`!"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e89e4ee8",
"metadata": {},
"outputs": [],
"source": [
"!python ./code/convert_hf_dataset_to_nemo.py \\\n",
" output_dir=./data/ \\\n",
" path=\"mozilla-foundation/common_voice_11_0\" \\\n",
" name=\"eo\" \\\n",
" ensure_ascii=False \\\n",
" use_auth_token=True"
]
},
{
"cell_type": "markdown",
"id": "2e34a392",
"metadata": {},
"source": [
"We apply filtering/data-processing to the dataset. We've skipped explanations as to why we're filtering, as extensive information can be found [here](https://github.com/andrusenkoau/NeMo/blob/esperanto_example/docs/source/asr/examples/esperanto_asr/esperanto_asr.rst)."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "3ca04023",
"metadata": {},
"outputs": [],
"source": [
"dev_manifest = f\"{data_dir}/validation/validation_mozilla-foundation_common_voice_11_0_manifest.json\"\n",
"test_manifest = f\"{data_dir}/test/test_mozilla-foundation_common_voice_11_0_manifest.json\"\n",
"train_manifest = f\"{data_dir}/train/train_mozilla-foundation_common_voice_11_0_manifest.json\"\n",
"\n",
"def compute_char_counts(manifest):\n",
" char_counts = {}\n",
" with open(manifest, 'r') as fn_in:\n",
" for line in tqdm(fn_in, desc=\"Compute counts..\"):\n",
" line = line.replace(\"\\n\", \"\")\n",
" data = json.loads(line)\n",
" text = data[\"text\"]\n",
" for word in text.split():\n",
" for char in word:\n",
" if char not in char_counts:\n",
" char_counts[char] = 1\n",
" else:\n",
" char_counts[char] += 1\n",
" return char_counts\n",
"\n",
"char_counts = compute_char_counts(train_manifest)\n",
"\n",
"threshold = 10\n",
"trash_char_list = []\n",
"\n",
"for char in char_counts:\n",
" if char_counts[char] <= threshold:\n",
" trash_char_list.append(char)\n",
"\n",
"def clear_data_set(manifest, char_rate_threshold=None):\n",
"\n",
" chars_to_ignore_regex = \"[\\.\\,\\?\\:\\-!;()«»…\\]\\[/\\*–‽+&_\\\\½√>€™$•¼}{~—=“\\\"”″‟„]\"\n",
" addition_ignore_regex = f\"[{''.join(trash_char_list)}]\"\n",
"\n",
" manifest_clean = manifest + '.clean'\n",
" war_count = 0\n",
" with open(manifest, 'r') as fn_in, \\\n",
" open(manifest_clean, 'w', encoding='utf-8') as fn_out:\n",
" for line in tqdm(fn_in, desc=\"Cleaning manifest data\"):\n",
" line = line.replace(\"\\n\", \"\")\n",
" data = json.loads(line)\n",
" text = data[\"text\"]\n",
" if char_rate_threshold and len(text.replace(' ', '')) / float(data['duration']) > char_rate_threshold:\n",
" print(f\"[WARNING]: {data['audio_filepath']} has char rate > 15 per sec: {len(text)} chars, {data['duration']} duration\")\n",
" war_count += 1\n",
" continue\n",
" text = re.sub(chars_to_ignore_regex, \"\", text)\n",
" text = re.sub(addition_ignore_regex, \"\", text)\n",
" data[\"text\"] = text\n",
" data = json.dumps(data, ensure_ascii=False)\n",
" fn_out.write(f\"{data}\\n\")\n",
" print(f\"[INFO]: {war_count} files were removed from manifest\")\n",
"\n",
"clear_data_set(dev_manifest)\n",
"clear_data_set(test_manifest)\n",
"clear_data_set(train_manifest, char_rate_threshold=15)"
]
},
{
"cell_type": "markdown",
"id": "4b50de19",
"metadata": {},
"source": [
"When the dataset is mounted to the SageMaker instance, the filepaths change. To handle this, we change the root path of the manifest entries to point to the correct mounted directory."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e31d9e91",
"metadata": {},
"outputs": [],
"source": [
"dev_manifest = dev_manifest + '.clean'\n",
"\n",
"\n",
"def replace_real_dir_with_mount_dir(manifest_path, data_root, mount_dir):\n",
" with open(manifest_path) as f:\n",
" entries = [json.loads(x) for x in f]\n",
" for entry in entries:\n",
" entry['audio_filepath'] = entry['audio_filepath'].replace(str(data_root), str(mount_dir))\n",
" with open(manifest_path + '.mount', 'w') as f:\n",
" f.write('\\n'.join([json.dumps(x) for x in entries]))\n",
"\n",
"replace_real_dir_with_mount_dir(dev_manifest, data_dir.absolute(), \"/opt/ml/input/data/testing/\")"
]
},
{
"cell_type": "markdown",
"id": "35d71747",
"metadata": {},
"source": [
"To improve performance, we also tar the training dataset to improve throughput when reading/writing from S3 directly within the remote instance."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "29d858a2",
"metadata": {},
"outputs": [],
"source": [
"train_manifest = train_manifest + '.clean'"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "bc495ca7",
"metadata": {},
"outputs": [],
"source": [
"!python ./code/convert_to_tarred_audio_dataset.py \\\n",
" --manifest_path={train_manifest} \\\n",
" --target_dir=./data/train_tarred_1bk \\\n",
" --num_shards=512 \\\n",
" --max_duration=15.0 \\\n",
" --min_duration=1.0 \\\n",
" --shuffle \\\n",
" --shuffle_seed=1 \\\n",
" --sort_in_shards \\\n",
" --workers=-1"
]
},
{
"cell_type": "markdown",
"id": "c850634e",
"metadata": {},
"source": [
"Finally, we generate the tokenizer that we will be using for training using the training manifest. Further information about parameters when choosing vocab sizes can be found [here](https://github.com/andrusenkoau/NeMo/blob/esperanto_example/docs/source/asr/examples/esperanto_asr/esperanto_asr.rst). "
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "ca394bcd",
"metadata": {},
"outputs": [],
"source": [
"vocab_size=128\n",
"\n",
"!python ./code/process_asr_text_tokenizer.py \\\n",
" --manifest={train_manifest} \\\n",
" --vocab_size={vocab_size} \\\n",
" --data_root=./data/tokenizers \\\n",
" --tokenizer=\"spe\" \\\n",
" --spe_type=bpe \\\n",
" --no_lower_case"
]
},
{
"cell_type": "markdown",
"id": "bfa2199e",
"metadata": {},
"source": [
"### 5. Upload the Dataset to S3\n",
"\n",
"This can take some time depending on your upload speed. We are uploading roughly 111GB of data to S3."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e1c5a60a",
"metadata": {},
"outputs": [],
"source": [
"# optional, move the raw training dataset out of the data folder to reduce data transfer\n",
"shutil.move(f\"{dataset_dir}/train/\", root_dir / \"train_raw/\")\n",
"\n",
"prefix = \"cv_esperanto\"\n",
"bucket = sess.default_bucket()\n",
"\n",
"loc = sess.upload_data(path=str(data_dir), bucket=bucket, key_prefix=prefix)\n",
"loc"
]
},
{
"cell_type": "markdown",
"id": "6321e3a9",
"metadata": {},
"source": [
"### 4. Configure the training job\n",
"\n",
"Now we configure the training job, by modifying the `config.yaml` file that is stored in our source code directory.\n",
"We pass relative directory paths for the data based on the SageMaker mount directory on the remote instance."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "ac709398",
"metadata": {},
"outputs": [],
"source": [
"# the output path for the model checkpoints/logs\n",
"output_path = \"s3://\" + sess.default_bucket() + \"/nemo-output/\""
]
},
{
"cell_type": "markdown",
"id": "8186678f",
"metadata": {},
"source": [
"We setup the training configuration based on the Esperanto experiments ran [here](https://github.com/andrusenkoau/NeMo/blob/esperanto_example/docs/source/asr/examples/esperanto_asr/esperanto_asr.rst).\n",
"\n",
"If you change the number of GPUs on the instance, remember to change the values below."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "a4b67faa",
"metadata": {},
"outputs": [],
"source": [
"global_batch_size = 1024\n",
"train_batch_size = 8\n",
"n_gpus = 8\n",
"num_nodes = 1 # set this to the number of nodes you'd like to train on\n",
"accumulate_grad_batches = int(global_batch_size / (n_gpus * num_nodes * train_batch_size))"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "4bb61640",
"metadata": {},
"outputs": [],
"source": [
"conf = OmegaConf.load(config_path)\n",
"\n",
"# data config setup\n",
"conf.model.train_ds.batch_size=train_batch_size\n",
"conf.model.train_ds.is_tarred = True\n",
"conf.model.train_ds.tarred_audio_filepaths = \"/opt/ml/input/data/training/train_tarred_1bk/audio__OP_0..511_CL_.tar\"\n",
"conf.model.train_ds.manifest_filepath = \"/opt/ml/input/data/training/train_tarred_1bk/tarred_audio_manifest.json\"\n",
"conf.model.validation_ds.manifest_filepath = \"/opt/ml/input/data/testing/validation/validation_mozilla-foundation_common_voice_11_0_manifest.json.clean.mount\"\n",
"\n",
"# optimization setup\n",
"conf.model.optim.lr = 1e-3\n",
"conf.model.optim.sched.name = \"CosineAnnealing\"\n",
"\n",
"# remove variable which is not needed when using CosineAnnealing\n",
"conf.model.optim.sched.pop('d_model')\n",
"\n",
"# logging/output setup\n",
"conf.exp_manager.exp_dir = \"/opt/ml/model/\"\n",
"\n",
"# the BPE tokenizer location based on the remote instance mounted directory\n",
"conf.model.tokenizer.dir = \"/opt/ml/input/data/training/tokenizers/tokenizer_spe_bpe_v128\"\n",
"\n",
"# the pre-trained model we want to fine-tune\n",
"conf.init_from_nemo_model = \"ssl_en_conformer_large.nemo\"\n",
"\n",
"# training setup\n",
"conf.trainer.accelerator = \"gpu\"\n",
"# enable SageMaker DDP\n",
"conf.trainer.strategy = \"smddp\"\n",
"conf.trainer.max_epochs = 300\n",
"conf.trainer.accumulate_grad_batches = accumulate_grad_batches\n",
"conf.trainer.precision = 16 # if using Ampere or above (i.e ml.p4d.24xlarge), you can use bf16 (BFloat16) \n",
"\n",
"# resume flags if crashes occur\n",
"conf.exp_manager.resume_if_exists=True\n",
"conf.exp_manager.resume_ignore_no_checkpoint=True"
]
},
{
"cell_type": "markdown",
"id": "1e287a26",
"metadata": {},
"source": [
"Uncomment the below cell you want to use W&B, you'll need to pass your W&B token at a later step."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "3ba11ab8",
"metadata": {},
"outputs": [],
"source": [
"# conf.exp_manager.wandb_logger_kwargs.name = \"common_voice_esperanto\"\n",
"# conf.exp_manager.create_wandb_logger = True\n",
"# conf.exp_manager.wandb_logger_kwargs.resume=True"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e4b27920",
"metadata": {},
"outputs": [],
"source": [
"# save the updated config\n",
"OmegaConf.save(conf, config_dir / 'config.yaml')"
]
},
{
"cell_type": "markdown",
"id": "1d8cd4cb",
"metadata": {},
"source": [
"Download the pre-trained SSL model from NGC. We'll upload this model with our code to SageMaker."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "8d11ff03",
"metadata": {},
"outputs": [],
"source": [
"!wget --content-disposition https://api.ngc.nvidia.com/v2/models/nvidia/nemo/ssl_en_conformer_large/versions/1.10.1/zip -O {code_dir}/ssl_en_conformer_large_1.10.1.zip\n",
"!unzip {code_dir}/ssl_en_conformer_large_1.10.1.zip -d {code_dir}"
]
},
{
"cell_type": "markdown",
"id": "959da702",
"metadata": {},
"source": [
"### 5. Run training on SageMaker\n",
"\n",
"Pass the path of the training and validation data on S3 + the output directory on S3 to the PyTorch estimator, and call fit with the appropriate bucket locations for the training and testing data."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "1d2e44e3",
"metadata": {},
"outputs": [],
"source": [
"channels = {\"training\": loc, \"testing\": loc}\n",
"\n",
"role = get_execution_role()\n",
"\n",
"output_path = \"s3://\" + sess.default_bucket() + \"/nemo-output/\"\n",
"\n",
"local_mode = False\n",
"\n",
"if local_mode:\n",
" instance_type = \"local_gpu\"\n",
"else:\n",
" instance_type = \"ml.p3.16xlarge\"\n",
" \n",
"instance_count = num_nodes\n",
"\n",
"est = PyTorch(\n",
" entry_point=\"speech_to_text_ctc_bpe.py\",\n",
" source_dir=\"code\", # directory of your training script\n",
" role=role,\n",
" instance_type=instance_type,\n",
" instance_count=instance_count,\n",
" framework_version=\"1.12.0\",\n",
" py_version=\"py38\",\n",
" volume_size=125,\n",
" output_path=output_path,\n",
" hyperparameters={'config-path': 'conf', 'config-name': 'config'},\n",
" # uncomment if you're enabling WANDB, and pass your own API key.\n",
" # environment={\"WANDB_API_KEY\": \"<WANDB_API_KEY>\"},\n",
" # enable SageMaker Optimized DDP for multi-node scaling\n",
" distribution={\"smdistributed\":{\"dataparallel\":{\"enabled\":True, \"fp16\": True}}},\n",
" input_mode='FastFile' # enables reading from S3 directly, no need to download the data\n",
")\n",
"\n",
"est.fit(inputs=channels)"
]
},
{
"cell_type": "markdown",
"id": "2be67b8e",
"metadata": {},
"source": [
"### 6. Download model, (Optional) Tensorboard Logs\n",
"\n",
"SageMaker stores our models/logs within a tar file after training has finished. These can be obtained from S3 like below.\n",
"\n",
"We also visualize the training logs. We suggest using an external logger (such as W&B) to track training progress during the run."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "81cd58d3",
"metadata": {},
"outputs": [],
"source": [
"key = est.model_data.replace(\"s3://\" + sess.default_bucket() + '/', '')\n",
"\n",
"sess.boto_session.client(\"s3\", region_name=sess.boto_region_name).download_file(\n",
" Bucket=sess.default_bucket(), Key=key, Filename='model.tar.gz',\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "656d53f3",
"metadata": {},
"outputs": [],
"source": [
"!tar -xvzf model.tar.gz"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "62018b6e",
"metadata": {},
"outputs": [],
"source": [
"%tensorboard --logdir ./ --host 0.0.0.0"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.9.12"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
tutorials/cloud/aws/SageMaker_ASR_Training.ipynb
+418
View File
@@ -0,0 +1,418 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "2e86a2b3",
"metadata": {},
"source": [
"# NeMo ASR Training Using AWS SageMaker"
]
},
{
"cell_type": "markdown",
"id": "215e3d3c",
"metadata": {},
"source": [
"In this tutorial we show how you can train a NeMo ASR Model using [Amazon SageMaker](https://docs.aws.amazon.com/sagemaker/latest/dg/whatis.html). This is meant to be a minimalistic example of how to use SageMaker with NeMo.\n",
"\n",
"AWS SageMaker is useful for practitioners/researchers who are familiar with training locally or on a remote instance (via SSH). SageMaker also supports multi-GPU & Multi-node.\n",
"\n",
"Using AWS SageMaker we train a simple Conformer CTC model using the AN4 dataset on a remote instance with a GPU (p3.2xlarge). We use S3 to store the data and our checkpoints/logs.\n",
"\n",
"The overall steps are:\n",
"\n",
"1. Setup your AWS Credentials to access SageMaker\n",
"2. Download the source code we'll be running\n",
"3. Setup AN4 dataset, upload data to S3\n",
"4. Configure the training job\n",
"5. Run training job on SageMaker\n",
"6. Download model, (Optional) Tensorboard Logs"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "ac621da0",
"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",
"## Install dependencies\n",
"!pip install wget\n",
"!apt-get install sox libsndfile1 ffmpeg\n",
"!pip install text-unidecode\n",
"!pip install matplotlib>=3.3.2\n",
"\n",
"## Install NeMo\n",
"BRANCH = 'main'\n",
"!python -m pip install \"nemo_toolkit[all] @ git+https://github.com/NVIDIA-NeMo/Speech.git@$BRANCH\"\n",
"\n",
"\"\"\"\n",
"Remember to restart the runtime for the kernel to pick up any upgraded packages (e.g. matplotlib)!\n",
"Alternatively, you can uncomment the exit() below to crash and restart the kernel, in the case\n",
"that you want to use the \"Run All Cells\" (or similar) option.\n",
"\"\"\"\n",
"# exit()"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "61c4fbe2",
"metadata": {},
"outputs": [],
"source": [
"pip install sagemaker awscli"
]
},
{
"cell_type": "markdown",
"id": "876f553d",
"metadata": {},
"source": [
"### 1. Setup SageMaker with AWS Credentials\n",
"\n",
"If you haven't setup your AWS credentials, you can setup using the AWS CLI.\n",
"You will need your access and Secret key, with permissions to use SageMaker and S3."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "1328482d",
"metadata": {},
"outputs": [],
"source": [
"!aws configure"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "01477d55",
"metadata": {},
"outputs": [],
"source": [
"%load_ext tensorboard\n",
"from pathlib import Path\n",
"import os\n",
"import sagemaker\n",
"import wget\n",
"from omegaconf import OmegaConf\n",
"from sagemaker import get_execution_role\n",
"from sagemaker.pytorch import PyTorch\n",
"\n",
"from nemo.utils.notebook_utils import download_an4"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "405806f9",
"metadata": {},
"outputs": [],
"source": [
"sess = sagemaker.Session()"
]
},
{
"cell_type": "markdown",
"id": "7d099a96",
"metadata": {},
"source": [
"### 2. Download the NeMo source code\n",
"\n",
"SageMaker allows you to pass in your own source code, with an entrypoint script.\n",
"\n",
"Below we download the AWS NeMo `config.yaml` which contains our configuration, and the `speech_to_text_ctc.py` script to run training.\n",
"\n",
"Our folder structure will look like this:\n",
"\n",
" code/\n",
" speech_to_text_ctc.py\n",
" conf/\n",
" config.yaml\n",
" data/\n",
" an4/"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "0b456c57",
"metadata": {},
"outputs": [],
"source": [
"root_dir = Path('./an4_nemo_sagemaker/')\n",
"code_dir = root_dir / 'code/'\n",
"config_dir = code_dir / 'conf/'\n",
"data_dir = root_dir / 'data/'\n",
"\n",
"root_dir.mkdir(exist_ok=True)\n",
"code_dir.mkdir(exist_ok=True)\n",
"config_dir.mkdir(exist_ok=True)\n",
"data_dir.mkdir(exist_ok=True)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "721c7b7d",
"metadata": {},
"outputs": [],
"source": [
"config_path = str(config_dir / \"config.yaml\")\n",
"wget.download(\"https://raw.githubusercontent.com/NVIDIA/NeMo/main/examples/asr/conf/conformer/conformer_ctc_char.yaml\", config_path)\n",
"wget.download(\"https://raw.githubusercontent.com/NVIDIA/NeMo/main/examples/asr/asr_ctc/speech_to_text_ctc.py\", str(code_dir))"
]
},
{
"cell_type": "markdown",
"id": "7934baab",
"metadata": {},
"source": [
"We also create a `requirements.txt` file within our source code to install NeMo."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "397d8eef",
"metadata": {},
"outputs": [],
"source": [
"with open(code_dir / 'requirements.txt', 'w') as f:\n",
" f.write(\"nemo_toolkit[all]\")"
]
},
{
"cell_type": "markdown",
"id": "a7bc7f51",
"metadata": {},
"source": [
"### 2.1 Initialize SageMaker within Training Script\n",
"\n",
"We provide a helper function that we require to be imported and run at the top of the training script.\n",
"\n",
"This installs and setups DDP for you. It also alleviates having to import a custom container, and can leverage all of the SageMaker containers. Rather than running this cell, you could also manually do this in your script."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e17535a7",
"metadata": {},
"outputs": [],
"source": [
"line = \"from nemo.utils.cloud import initialize_sagemaker; initialize_sagemaker()\"\n",
"with open(code_dir / \"speech_to_text_ctc.py\", 'r+') as f:\n",
" content = f.read()\n",
" f.seek(0, 0)\n",
" f.write(line.rstrip('\\r\\n') + '\\n' + content)"
]
},
{
"cell_type": "markdown",
"id": "bfa2199e",
"metadata": {},
"source": [
"### 3. Setup the AN4 Dataset, upload data to S3\n",
"\n",
"We now download our training and validation data, uploading to S3 so that SageMaker can mount our data to the instance at runtime."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e1c5a60a",
"metadata": {},
"outputs": [],
"source": [
"# within the SageMaker container, mount_dir will be where our data is stored.\n",
"download_an4(\n",
" data_dir=str(data_dir),\n",
" train_mount_dir=\"/opt/ml/input/data/training/\",\n",
" test_mount_dir=\"/opt/ml/input/data/testing/\",\n",
")\n",
"\n",
"# Upload to the default bucket\n",
"prefix = \"an4\"\n",
"bucket = sess.default_bucket()\n",
"loc = sess.upload_data(path=str(data_dir), bucket=bucket, key_prefix=prefix)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "b1de1089",
"metadata": {},
"outputs": [],
"source": [
"output_path = \"s3://\" + sess.default_bucket() + \"/nemo-output/\""
]
},
{
"cell_type": "markdown",
"id": "6321e3a9",
"metadata": {},
"source": [
"### 4. Configure the training job\n",
"\n",
"Now we configure the training job, by modifying the `config.yaml` file that is stored in our source code directory.\n",
"We pass relative directory paths for the data based on the SageMaker mount directory on the remote instance."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "4bb61640",
"metadata": {},
"outputs": [],
"source": [
"conf = OmegaConf.load(config_path)\n",
"\n",
"# Set Data Locations based on the mounted directory in the SageMaker instance\n",
"conf.model.train_ds.manifest_filepath = \"/opt/ml/input/data/training/an4/train_manifest.json\"\n",
"conf.model.validation_ds.manifest_filepath = \"/opt/ml/input/data/testing/an4/test_manifest.json\"\n",
"conf.trainer.accelerator = \"gpu\"\n",
"conf.trainer.max_epochs = 150\n",
"\n",
"# Output directory for our experiment within the SageMaker instance\n",
"conf.exp_manager.exp_dir=\"/opt/ml/model/\"\n",
"\n",
"# Create a Small Variant of the Conformer Model\n",
"conf.model.encoder.n_layers = 8\n",
"conf.model.n_heads = 4\n",
"conf.model.spec_augment.time_masks = 5\n",
"\n",
"# Set Optimizer parameters\n",
"conf.model.optim.lr = 2.0 # by default we using Noam scheduling, the LR is a multiplier \n",
"\n",
"OmegaConf.save(conf, config_dir / 'config.yaml')"
]
},
{
"cell_type": "markdown",
"id": "959da702",
"metadata": {},
"source": [
"### 5. Run training on SageMaker\n",
"\n",
"Pass the path of the training and validation data on S3 + the output directory on S3 to the PyTorch estimator, and call fit with the appropriate bucket locations for the training and testing data."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "1d2e44e3",
"metadata": {},
"outputs": [],
"source": [
"channels = {\"training\": loc, \"testing\": loc}\n",
"\n",
"role = get_execution_role()\n",
"\n",
"# Set to True to enable SageMaker to run locally\n",
"local_mode = False\n",
"\n",
"if local_mode:\n",
" instance_type = \"local_gpu\"\n",
"else:\n",
" instance_type = \"ml.p3.2xlarge\"\n",
"\n",
"est = PyTorch(\n",
" entry_point=\"speech_to_text_ctc.py\", # the script we want to run\n",
" source_dir=str(code_dir), # where our conf/script is\n",
" role=role,\n",
" instance_type=instance_type,\n",
" instance_count=1,\n",
" framework_version=\"1.12.0\", # version of PyTorch\n",
" py_version=\"py38\",\n",
" volume_size=250,\n",
" output_path=output_path,\n",
" hyperparameters={'config-path': 'conf'},\n",
")\n",
"\n",
"est.fit(inputs=channels)"
]
},
{
"cell_type": "markdown",
"id": "2be67b8e",
"metadata": {},
"source": [
"### 6. Download model, (Optional) Tensorboard Logs\n",
"\n",
"SageMaker stores our models/logs within a tar file after training has finished. These can be obtained from S3 like below.\n",
"\n",
"We also visualize the training logs. We suggest using an external logger (such as W&B) to track training progress during the run."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "81cd58d3",
"metadata": {},
"outputs": [],
"source": [
"key = est.model_data.replace(\"s3://\" + sess.default_bucket() + '/', '')\n",
"\n",
"sess.boto_session.client(\"s3\", region_name=sess.boto_region_name).download_file(\n",
" Bucket=sess.default_bucket(), Key=key, Filename='model.tar.gz',\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "656d53f3",
"metadata": {},
"outputs": [],
"source": [
"!tar -xvzf model.tar.gz"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "62018b6e",
"metadata": {},
"outputs": [],
"source": [
"%tensorboard --logdir ./ --host 0.0.0.0"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.9.12"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
tutorials/cloud/aws/images/sagemaker_benchmark.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 32 KiB

tutorials/helper_files/__init__.py
+15
View File
@@ -0,0 +1,15 @@
# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Helper files for NeMo tutorials."""
tutorials/helper_files/gpt_components.py
+187
View File
@@ -0,0 +1,187 @@
# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
GPT components for the NeMo Models tutorial.
This module contains the neural network components used in the tutorial 01_NeMo_Models.ipynb
"""
import math
from typing import Optional
import torch
import torch.nn as nn
from torch.nn import functional as F
from nemo.core import NeuralModule, typecheck
from nemo.core.neural_types import EmbeddedTextType, EncodedRepresentation, Index, LogitsType, NeuralType
from nemo.core.neural_types.elements import *
# Custom Element Types
class AttentionType(EncodedRepresentation):
"""Basic Attention Element Type"""
class SelfAttentionType(AttentionType):
"""Self Attention Element Type"""
class CausalSelfAttentionType(SelfAttentionType):
"""Causal Self Attention Element Type"""
# Neural Network Modules (not NeMo neural modules)
class CausalSelfAttention(nn.Module):
"""
A vanilla multi-head masked self-attention layer with a projection at the end.
It is possible to use torch.nn.MultiheadAttention here but I am including an
explicit implementation here to show that there is nothing too scary here.
"""
def __init__(self, n_embd, block_size, n_head, attn_pdrop, resid_pdrop):
super().__init__()
assert n_embd % n_head == 0
self.n_head = n_head
# key, query, value projections for all heads
self.key = nn.Linear(n_embd, n_embd)
self.query = nn.Linear(n_embd, n_embd)
self.value = nn.Linear(n_embd, n_embd)
# regularization
self.attn_drop = nn.Dropout(attn_pdrop)
self.resid_drop = nn.Dropout(resid_pdrop)
# output projection
self.proj = nn.Linear(n_embd, n_embd)
# causal mask to ensure that attention is only applied to the left in the input sequence
self.register_buffer("mask", torch.tril(torch.ones(block_size, block_size)).view(1, 1, block_size, block_size))
def forward(self, x, layer_past=None):
B, T, C = x.size()
# calculate query, key, values for all heads in batch and move head forward to be the batch dim
k = self.key(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
q = self.query(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
v = self.value(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
# causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
att = att.masked_fill(self.mask[:, :, :T, :T] == 0, float('-inf'))
att = F.softmax(att, dim=-1)
att = self.attn_drop(att)
y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side
# output projection
y = self.resid_drop(self.proj(y))
return y
class Block(nn.Module):
"""an unassuming Transformer block"""
def __init__(self, n_embd, block_size, n_head, attn_pdrop, resid_pdrop):
super().__init__()
self.ln1 = nn.LayerNorm(n_embd)
self.ln2 = nn.LayerNorm(n_embd)
self.attn = CausalSelfAttention(n_embd, block_size, n_head, attn_pdrop, resid_pdrop)
self.mlp = nn.Sequential(
nn.Linear(n_embd, 4 * n_embd),
nn.GELU(),
nn.Linear(4 * n_embd, n_embd),
nn.Dropout(resid_pdrop),
)
def forward(self, x):
x = x + self.attn(self.ln1(x))
x = x + self.mlp(self.ln2(x))
return x
# NeMo Neural Modules
class GPTEmbedding(NeuralModule):
def __init__(self, vocab_size: int, n_embd: int, block_size: int, embd_pdrop: float = 0.0):
super().__init__()
# input embedding stem: drop(content + position)
self.tok_emb = nn.Embedding(vocab_size, n_embd)
self.pos_emb = nn.Parameter(torch.zeros(1, block_size, n_embd))
self.drop = nn.Dropout(embd_pdrop)
@typecheck()
def forward(self, idx):
b, t = idx.size()
# forward the GPT model
token_embeddings = self.tok_emb(idx) # each index maps to a (learnable) vector
position_embeddings = self.pos_emb[:, :t, :] # each position maps to a (learnable) vector
x = self.drop(token_embeddings + position_embeddings)
return x
@property
def input_types(self):
return {'idx': NeuralType(('B', 'T'), Index())}
@property
def output_types(self):
return {'embeddings': NeuralType(('B', 'T', 'C'), EmbeddedTextType())}
class GPTTransformerEncoder(NeuralModule):
def __init__(
self,
n_embd: int,
block_size: int,
n_head: int,
n_layer: int,
attn_pdrop: float = 0.0,
resid_pdrop: float = 0.0,
):
super().__init__()
self.blocks = nn.Sequential(
*[Block(n_embd, block_size, n_head, attn_pdrop, resid_pdrop) for _ in range(n_layer)]
)
@typecheck()
def forward(self, embed):
return self.blocks(embed)
@property
def input_types(self):
return {'embed': NeuralType(('B', 'T', 'C'), EmbeddedTextType())}
@property
def output_types(self):
return {'encoding': NeuralType(('B', 'T', 'C'), CausalSelfAttentionType())}
class GPTDecoder(NeuralModule):
def __init__(self, n_embd: int, vocab_size: int):
super().__init__()
self.ln_f = nn.LayerNorm(n_embd)
self.head = nn.Linear(n_embd, vocab_size, bias=False) # no need for extra bias due to one in ln_f
@typecheck()
def forward(self, encoding):
x = self.ln_f(encoding)
logits = self.head(x)
return logits
@property
def input_types(self):
return {'encoding': NeuralType(('B', 'T', 'C'), EncodedRepresentation())}
@property
def output_types(self):
return {'logits': NeuralType(('B', 'T', 'C'), LogitsType())}
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

tutorials/speechlm2/SpeechLM_With_NeMo_Automodel.ipynb
+538
View File
@@ -0,0 +1,538 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Training a Speech LLM with NeMo Automodel and MoE\n",
"\n",
"This tutorial walks through the full **SALMAutomodel** pipeline: data preparation,\n",
"training, checkpoint conversion, and evaluation.\n",
"\n",
"## What is SALMAutomodel?\n",
"\n",
"SALMAutomodel is a Speech-Augmented Language Model that uses\n",
"[NeMo Automodel](https://github.com/NVIDIA-NeMo/Automodel) as the LLM backend.\n",
"It connects a pretrained ASR encoder (e.g., Canary) to a pretrained LLM via a\n",
"learnable linear projection layer.\n",
"\n",
"## Why MoE?\n",
"\n",
"We use **NVIDIA Nemotron Nano V3** as the LLM backbone — a Mixture-of-Experts\n",
"model with 30B total parameters but only 3B active per token. NeMo Automodel\n",
"provides two key MoE optimizations:\n",
"\n",
"- **Grouped GEMM**: Fuses expert computations into a single batched matrix multiply.\n",
"- **DeepEP**: Efficient all-to-all expert routing across GPUs.\n",
"\n",
"## EP and FSDP2: Same Axis\n",
"\n",
"A key point: **Expert Parallelism (EP) reuses the FSDP2 data-parallel axis\n",
"(`dp_size`)**. Dense layers are sharded via FSDP2, while MoE expert layers\n",
"use EP for all-to-all routing — both on the same set of GPUs. Setting\n",
"`ep_size` does *not* add a separate dimension.\n",
"\n",
"## What We Cover\n",
"\n",
"1. Download Mini LibriSpeech with Lhotse\n",
"2. Train SALMAutomodel (only `perception.proj` is trained; LLM and ASR encoder are frozen)\n",
"3. Convert the distributed checkpoint to HuggingFace format\n",
"4. Evaluate with distributed inference and compute WER\n",
"\n",
"## Prerequisites\n",
"\n",
"- The tutorial was tested on **2x RTX 6000 Pro (Blackwell) GPUs**; for a smaller setup, you might need to use a smaller LLM backbone\n",
"- `pip install nemo-toolkit[speechlm2]`"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!pip install nemo-toolkit[speechlm2]"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 1. Download Mini LibriSpeech\n",
"\n",
"We use [Lhotse](https://github.com/lhotse-speech/lhotse) to download and prepare\n",
"Mini LibriSpeech — a small subset of LibriSpeech with two splits:\n",
"- `train-clean-5` (~5 hours)\n",
"- `dev-clean-2` (~2 hours)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"scrolled": true
},
"outputs": [],
"source": [
"import os\n",
"from pathlib import Path\n",
"from lhotse.recipes import download_librispeech, prepare_librispeech\n",
"from lhotse import CutSet\n",
"\n",
"def find_nemo_root(start: Path) -> Path:\n",
" start = start.resolve()\n",
" for candidate in (start, *start.parents):\n",
" if (candidate / \"examples\" / \"speechlm2\").is_dir() and (candidate / \"tutorials\" / \"speechlm2\").is_dir():\n",
" return candidate\n",
" raise RuntimeError(\n",
" f\"Could not locate the NeMo source tree from {start}. Open this notebook from inside the repo checkout.\"\n",
" )\n",
"\n",
"NEMO_ROOT = find_nemo_root(Path.cwd())\n",
"TUTORIAL_ROOT = NEMO_ROOT / \"tutorials\" / \"speechlm2\"\n",
"DATA_ROOT = TUTORIAL_ROOT / \"data\"\n",
"SPEECHLM2_EXAMPLES_ROOT = NEMO_ROOT / \"examples\" / \"speechlm2\"\n",
"RESULTS_ROOT = TUTORIAL_ROOT / \"salm_tutorial_results\"\n",
"HF_OUTPUT_DIR = TUTORIAL_ROOT / \"salm_tutorial_hf\"\n",
"GENERATIONS_PATH = TUTORIAL_ROOT / \"tutorial_generations.jsonl\"\n",
"\n",
"existing_pythonpath = os.environ.get(\"PYTHONPATH\")\n",
"REPO_PYTHONPATH = str(NEMO_ROOT)\n",
"if existing_pythonpath:\n",
" REPO_PYTHONPATH = os.pathsep.join([REPO_PYTHONPATH, existing_pythonpath])\n",
"SCRIPT_ENV = dict(os.environ, PYTHONPATH=REPO_PYTHONPATH)\n",
"\n",
"print(f\"Using NeMo root: {NEMO_ROOT}\")\n",
"print(f\"Using SpeechLM2 examples: {SPEECHLM2_EXAMPLES_ROOT}\")\n",
"\n",
"CORPUS_DIR = download_librispeech(DATA_ROOT, dataset_parts=\"mini_librispeech\")\n",
"manifests = prepare_librispeech(\n",
" CORPUS_DIR, dataset_parts=\"mini_librispeech\", output_dir=DATA_ROOT / \"manifests\"\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 2. Create Lhotse CutSets\n",
"\n",
"Convert the raw manifests into CutSets and save as `.jsonl` files. LibriSpeech has all uppercase transcripts which we normalize: the first letter is capital, and each sentence ends with a full stop, to match LLM training data more closely."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"for split, parts in manifests.items():\n",
" cuts = CutSet.from_manifests(\n",
" recordings=parts[\"recordings\"], supervisions=parts[\"supervisions\"]\n",
" ).with_recording_path_prefix(DATA_ROOT.absolute())\n",
" path = DATA_ROOT / f\"cuts_{split}.jsonl\"\n",
" cuts = cuts.transform_text(lambda txt: txt.capitalize() + \".\")\n",
" cuts.to_file(path)\n",
" print(\n",
" f\"{split}: {len(cuts)} cuts, \"\n",
" f\"total duration: {sum(c.duration for c in cuts):.1f}s\"\n",
" )"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 3. Inspect the Data\n",
"\n",
"Each `Cut` holds a pointer to an audio recording, its duration, and one or more\n",
"supervisions (transcripts with timing information)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"cuts = CutSet.from_file(DATA_ROOT / \"cuts_train-clean-5.jsonl\")\n",
"cut = cuts[0]\n",
"print(f\"ID: {cut.id}\")\n",
"print(f\"Duration: {cut.duration:.2f}s\")\n",
"print(f\"Text: {cut.supervisions[0].text}\")\n",
"cut.play_audio()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 4. Write Training Config\n",
"\n",
"We write a YAML config for SALMAutomodel with:\n",
"\n",
"- **Nemotron Nano V3** as the LLM backbone\n",
"- **EP=2** on 2 GPUs (dense layers → FSDP2, MoE layers → EP, same axis)\n",
"- Only `perception.proj` (Linear 1024→4096) is trainable; LLM, ASR encoder,\n",
" preprocessor, and modality adapter are frozen"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"train_cuts = str((DATA_ROOT / \"cuts_train-clean-5.jsonl\").resolve())\n",
"val_cuts = str((DATA_ROOT / \"cuts_dev-clean-2.jsonl\").resolve())\n",
"\n",
"results_dir = str(RESULTS_ROOT.resolve())\n",
"\n",
"config_yaml = f\"\"\"\\\n",
"model:\n",
" pretrained_llm: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16\n",
" pretrained_asr: nvidia/canary-1b-flash\n",
" pretrained_weights: true\n",
" use_nemo_automodel: true\n",
" trust_remote_code: true # needed for Nemotron-3-Nano\n",
" automodel_backend:\n",
" dispatcher: torch # set to \"torch\" if your GPUs don't have NVLINK/NVSHMEM\n",
"\n",
" prompt_format: nemotron-nano-v3\n",
" audio_locator_tag: \"<|audio|>\"\n",
"\n",
" freeze_params:\n",
" - \"^llm\\\\\\\\..+$\"\n",
" - \"^perception\\\\\\\\.preprocessor\\\\\\\\..+$\"\n",
" - \"^perception\\\\\\\\.encoder\\\\\\\\..+$\"\n",
" prevent_freeze_params: []\n",
"\n",
" # Uncomment to enable LoRA on the LLM:\n",
" # lora:\n",
" # dim: 128\n",
" # alpha: 256\n",
" # dropout: 0.01\n",
" # target_modules: [\"q_proj\", \"v_proj\"]\n",
"\n",
" perception:\n",
" target: nemo.collections.speechlm2.modules.perception.AudioPerceptionModule\n",
" output_dim: 4096\n",
" modality_adapter:\n",
" _target_: nemo.collections.speechlm2.modules.perception.IdentityConnector\n",
" d_model: 1024\n",
" spec_augment:\n",
" _target_: nemo.collections.asr.modules.SpectrogramAugmentation\n",
" freq_masks: 2 # set to zero to disable it\n",
" time_masks: 10 # set to zero to disable it\n",
" freq_width: 27\n",
" time_width: 5 # 5 frames = 50ms\n",
"\n",
" optimizer:\n",
" _target_: torch.optim.AdamW\n",
" lr: 3e-4\n",
" betas: [0.9, 0.98]\n",
" weight_decay: 1e-3\n",
" foreach: true\n",
"\n",
" lr_scheduler:\n",
" _target_: nemo.core.optim.lr_scheduler.CosineAnnealing\n",
" warmup_steps: 50\n",
" min_lr: 1e-6\n",
" max_steps: ${{trainer.max_steps}}\n",
"\n",
"trainer:\n",
" devices: 2\n",
" accelerator: gpu\n",
" num_nodes: 1\n",
" precision: bf16-true\n",
" logger: false\n",
" enable_checkpointing: false\n",
" use_distributed_sampler: false\n",
" max_steps: 500\n",
" val_check_interval: 1.0\n",
" limit_val_batches: 5\n",
" log_every_n_steps: 10\n",
" num_sanity_val_steps: 0\n",
" gradient_clip_val: 1.0\n",
" accumulate_grad_batches: 1\n",
" strategy:\n",
" _target_: nemo.collections.speechlm2.parts.parallel.AutomodelParallelStrategy\n",
" dp_size: null\n",
" dp_replicate_size: 1\n",
" tp_size: 1\n",
" pp_size: 1\n",
" cp_size: 1\n",
" ep_size: 2\n",
"\n",
"data:\n",
" train_ds:\n",
" sample_rate: 16000\n",
" prompt_format: ${{model.prompt_format}}\n",
" token_equivalent_duration: 0.08\n",
" input_cfg:\n",
" - type: lhotse_as_conversation\n",
" cuts_path: {train_cuts}\n",
" audio_locator_tag: ${{model.audio_locator_tag}}\n",
" tags:\n",
" context: \"\" # a fixed prompt can be added here since, or example-specific prompt can be attached in the manfiest directly\n",
" seed: 42\n",
" shuffle: true\n",
" shard_seed: randomized\n",
" num_workers: 1\n",
" batch_size: 16\n",
"\n",
" validation_ds:\n",
" prompt_format: ${{model.prompt_format}}\n",
" token_equivalent_duration: 0.08\n",
" datasets:\n",
" dev_clean_2:\n",
" input_cfg:\n",
" - type: lhotse_as_conversation\n",
" cuts_path: {val_cuts}\n",
" audio_locator_tag: ${{model.audio_locator_tag}}\n",
" tags:\n",
" context: \"\" # a fixed prompt can be added here since, or example-specific prompt can be attached in the manfiest directly\n",
" sample_rate: 16000\n",
" batch_size: 1\n",
" seed: 42\n",
" shard_seed: randomized\n",
"\n",
"exp_manager:\n",
" exp_dir: null\n",
" explicit_log_dir: {results_dir}\n",
" name: salm\n",
" create_tensorboard_logger: false\n",
" create_checkpoint_callback: true\n",
" use_datetime_version: true\n",
" resume_if_exists: true\n",
" resume_ignore_no_checkpoint: true\n",
" create_wandb_logger: false\n",
" checkpoint_callback_params:\n",
" filename: \"{{step}}\"\n",
" monitor: val_loss\n",
" mode: min\n",
" every_n_epochs: 1\n",
" always_save_nemo: false\n",
" save_top_k: 1\n",
" save_nemo_on_train_end: false\n",
"\"\"\"\n",
"\n",
"config_path = DATA_ROOT / \"salm_automodel_tutorial.yaml\"\n",
"config_path.write_text(config_yaml)\n",
"print(f\"Config written to {config_path.resolve()}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 5. Config Explained\n",
"\n",
"Key settings in the config above:\n",
"\n",
"| Setting | Meaning |\n",
"|---|---|\n",
"| `use_nemo_automodel: true` | Selects `SALMAutomodel` (NeMo Automodel backend) |\n",
"| `ep_size: 2` | Expert Parallelism on the FSDP data-parallel axis. Dense layers are sharded via FSDP2, MoE expert layers use EP — both on the same 2 GPUs |\n",
"| `perception.output_dim: 4096` | Nemotron Nano V3 has `hidden_size=4096` |\n",
"| `freeze_params` | Freezes the LLM, ASR encoder, preprocessor, and modality adapter. Only `perception.proj` (Linear 1024→4096) is trainable |\n",
"| `lora` (commented) | Can be enabled for parameter-efficient LLM fine-tuning using Automodel-native LoRA |"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 6. Launch Training\n",
"\n",
"We use `torchrun` with 2 GPUs. The training script (`salm_train.py`) reads\n",
"`use_nemo_automodel: true` and instantiates `SALMAutomodel`."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"scrolled": true
},
"outputs": [],
"source": [
"import subprocess\n",
"\n",
"TRAIN_SCRIPT = str((SPEECHLM2_EXAMPLES_ROOT / \"salm_train.py\").resolve())\n",
"\n",
"cmd = [\n",
" \"torchrun\", \"--nproc_per_node=2\",\n",
" TRAIN_SCRIPT,\n",
" f\"--config-path={str(DATA_ROOT.resolve())}\",\n",
" \"--config-name=salm_automodel_tutorial\",\n",
"]\n",
"print(\"Running:\", \" \".join(cmd))\n",
"subprocess.run(cmd, check=True, cwd=str(NEMO_ROOT), env=SCRIPT_ENV)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 7. Locate Checkpoint"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"ckpt_paths = sorted((RESULTS_ROOT / \"checkpoints\").glob(\"*\"))\n",
"ckpt_paths = [str(path.resolve()) for path in ckpt_paths]\n",
"print(\"Checkpoints found:\", ckpt_paths)\n",
"CKPT_PATH = ckpt_paths[-1]\n",
"CONFIG_PATH = str((RESULTS_ROOT / \"exp_config.yaml\").resolve())\n",
"print(f\"Using checkpoint: {CKPT_PATH}\")\n",
"print(f\"Using config: {CONFIG_PATH}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 8. Convert Checkpoint to HuggingFace Format\n",
"\n",
"Training with EP=2 produces **distributed checkpoints** (a directory with\n",
"per-rank shards). The `to_hf.py` script consolidates DTensors into regular\n",
"tensors and saves them as `config.json` + `model.safetensors`.\n",
"\n",
"We launch with `torchrun --nproc_per_node=2` to match the training topology."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"scrolled": true
},
"outputs": [],
"source": [
"TO_HF_SCRIPT = str((SPEECHLM2_EXAMPLES_ROOT / \"to_hf.py\").resolve())\n",
"\n",
"cmd = [\n",
" \"torchrun\", \"--nproc_per_node=2\",\n",
" TO_HF_SCRIPT,\n",
" \"class_path=nemo.collections.speechlm2.models.SALMAutomodel\",\n",
" f\"ckpt_path='{CKPT_PATH}'\",\n",
" f\"ckpt_config='{CONFIG_PATH}'\",\n",
" f\"output_dir='{HF_OUTPUT_DIR.resolve()}'\",\n",
"]\n",
"print(\"Running:\", \" \".join(cmd))\n",
"subprocess.run(cmd, check=True, cwd=str(NEMO_ROOT), env=SCRIPT_ENV)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 9. Evaluate with WER\n",
"\n",
"Run distributed inference on the dev set and compute Word Error Rate.\n",
"The `salm_eval.py` script prints per-batch and overall WER."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"scrolled": true
},
"outputs": [],
"source": [
"DEV_CUTS = str((DATA_ROOT / \"cuts_dev-clean-2.jsonl\").resolve())\n",
"EVAL_SCRIPT = str((SPEECHLM2_EXAMPLES_ROOT / \"salm_eval.py\").resolve())\n",
"\n",
"cmd = [\n",
" \"torchrun\", \"--nproc_per_node=2\",\n",
" EVAL_SCRIPT,\n",
" f\"pretrained_name='{HF_OUTPUT_DIR.resolve()}'\",\n",
" f\"inputs={DEV_CUTS}\",\n",
" # f\"inputs={train_cuts}\",\n",
" \"batch_size=128\",\n",
" \"max_new_tokens=128\",\n",
" \"ep_size=2\",\n",
" f\"output_manifest='{GENERATIONS_PATH.resolve()}'\",\n",
" \"user_prompt=Transcribe the following:\",\n",
" \"enable_thinking=False\",\n",
"]\n",
"print(\"Running:\", \" \".join(cmd))\n",
"subprocess.run(cmd, check=True, cwd=str(NEMO_ROOT), env=SCRIPT_ENV)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 10. Inspect Results"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import json\n",
"\n",
"with open(GENERATIONS_PATH) as f:\n",
" results = [json.loads(line) for line in f]\n",
"\n",
"for r in results[:5]:\n",
" print(f\"REF: {r['text']}\")\n",
" print(f\"HYP: {r['pred_text']}\")\n",
" print()\n",
"\n",
"print(f\"Total examples: {len(results)}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Summary\n",
"\n",
"In this tutorial we:\n",
"\n",
"1. Downloaded Mini LibriSpeech with Lhotse\n",
"2. Trained SALMAutomodel with Nemotron Nano V3 MoE backbone (EP=2 on the FSDP data-parallel axis)\n",
"3. Converted the distributed checkpoint to HuggingFace format\n",
"4. Evaluated with distributed inference and computed WER\n",
"\n",
"### Next Steps\n",
"\n",
"- **Scale up**: more GPUs, larger datasets\n",
"- **Enable LoRA**: uncomment the `lora:` block for parameter-efficient LLM fine-tuning\n",
"- **Try other parallelism combos**: TP + EP, HSDP (`dp_replicate_size > 1`)\n",
"- See `docs/source/speechlm2/` for full documentation"
]
}
],
"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.13.12"
}
},
"nbformat": 4,
"nbformat_minor": 4
}
tutorials/tools/CTC_Segmentation_Tutorial.ipynb
+762
View File
@@ -0,0 +1,762 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "d4KCUoxSpdoZ"
},
"outputs": [],
"source": [
"BRANCH = 'main'\n",
"\n",
"\"\"\"\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",
"\"\"\""
]
},
{
"cell_type": "code",
"source": [
"import os\n",
"# either provide a path to local NeMo repository with NeMo already installed or git clone\n",
"\n",
"# option #1: local path to NeMo repo with NeMo already installed\n",
"NEMO_DIR_PATH = \"NeMo\"\n",
"\n",
"# option #2: download NeMo repo\n",
"if 'google.colab' in str(get_ipython()) or not os.path.exists(NEMO_DIR_PATH):\n",
" ! git clone -b $BRANCH https://github.com/NVIDIA-NeMo/Speech\n",
" ! cd NeMo\n",
" ! python -m pip install \"nemo_toolkit[all] @ git+https://github.com/NVIDIA-NeMo/Speech.git@$BRANCH\""
],
"metadata": {
"id": "6DGWYSp62hs1"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"source": [
"# install reqs for this notebook\n",
"!python -m pip install --no-cache-dir -r \"{NEMO_DIR_PATH}/tools/ctc_segmentation/requirements.txt\""
],
"metadata": {
"id": "73jaO-HZ4b_5"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"source": [
"import json\n",
"import os\n",
"import wget\n",
"\n",
"from IPython.display import Audio\n",
"import numpy as np\n",
"import scipy\n",
"import soundfile as sf\n",
"\n",
"\n",
"from plotly import graph_objects as go\n"
],
"metadata": {
"id": "WFvmiWv02jr2"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"source": [
"!apt-get install sox libsox-fmt-mp3"
],
"metadata": {
"id": "JbdAPFb99Mff"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {
"id": "xXRARM8XtK_g"
},
"source": [
"# 1. Introduction\n",
"End-to-end Automatic Speech Recognition (ASR) systems surpassed traditional systems in performance but require large amounts of labeled data for training.\n",
"\n",
"This tutorial will show how to use a pre-trained with Connectionist Temporal Classification (CTC) ASR model, such as [QuartzNet Model](https://arxiv.org/abs/1910.10261) or [Citrinet](https://arxiv.org/abs/2104.01721) to split long audio files and the corresponding transcripts into shorter fragments that are suitable for an ASR model training.\n",
"\n",
"We're going to use [ctc-segmentation](https://github.com/lumaku/ctc-segmentation) Python package based on the algorithm described in [CTC-Segmentation of Large Corpora for German End-to-end Speech Recognition](https://arxiv.org/pdf/2007.09127.pdf)."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "S1DZk-inQGTI"
},
"source": [
"`TOOLS_DIR` contains scripts that we are going to need during the next steps, all necessary scripts could be found [here](https://github.com/NVIDIA/NeMo/tree/main/tools/ctc_segmentation/scripts)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "1C9DdMfvRFM-"
},
"outputs": [],
"source": [
"if 'google.colab' in str(get_ipython()):\n",
" NEMO_DIR_PATH = \"/content/NeMo\"\n",
"elif not os.path.exists(NEMO_DIR_PATH):\n",
" NEMO_DIR_PATH = \"NeMo\"\n",
"\n",
"TOOLS_DIR = f'{NEMO_DIR_PATH}/tools/ctc_segmentation/scripts'\n",
"print(TOOLS_DIR)\n",
"! ls -l $TOOLS_DIR"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "XUEncnqTIzF6"
},
"source": [
"# 2. Data Download\n",
"First, let's download audio and text data (data source: [https://librivox.org/](https://librivox.org/) and [http://www.gutenberg.org](http://www.gutenberg.org)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "bkeKX2I_tIgV"
},
"outputs": [],
"source": [
"## create data directory and download an audio file\n",
"WORK_DIR = 'WORK_DIR'\n",
"DATA_DIR = WORK_DIR + '/DATA'\n",
"os.makedirs(DATA_DIR, exist_ok=True)\n",
"\n",
"print('downloading audio samples...')\n",
"wget.download(\"https://multilangaudiosamples.s3.us-east-2.amazonaws.com/audio_samples.zip\", DATA_DIR)\n",
"! unzip -o $DATA_DIR/audio_samples.zip -d $DATA_DIR\n",
"! rm $DATA_DIR/audio_samples.zip\n",
"\n",
"DATA_DIR = os.path.join(DATA_DIR, \"audio_samples\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "2JAv7ePmpdok"
},
"source": [
"We downloaded audio and text samples in `English` and `Spanish`:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Y6VYVk9mpdol"
},
"outputs": [],
"source": [
"! ls $DATA_DIR"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "-_XE9MkKuAA7"
},
"source": [
"Data folder for each language contains both audio and text files. Note, the text file and the audio file share the same base name. For example, an audio file `example.wav` should have a corresponding text file called `example.txt`."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "IGhijb-Bpdol"
},
"outputs": [],
"source": [
"! ls $DATA_DIR/es/audio/ $DATA_DIR/es/text/"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "FWqlbSryw_WL"
},
"source": [
"# 3. Segmentation of a single file (Spanish sample)\n",
"\n",
"Let's listen to our Spanish audio sample:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "ulkPrqwipdom"
},
"outputs": [],
"source": [
"base_name_es = \"el19demarzoyel2demayo_03_perezgaldos\"\n",
"Audio(f\"{DATA_DIR}/es/audio/{base_name_es}.wav\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "RzlLsmMXpdom"
},
"source": [
"Let's take a look at the ground truth text:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "9Qfp10Xnpdom"
},
"outputs": [],
"source": [
"text = f\"{DATA_DIR}/es/text/{base_name_es}.txt\"\n",
"! cat $text"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "RMT5lkPYzZHK"
},
"source": [
"As one probably noticed, the audio file contains a prologue and an epilogue that are missing in the corresponding text. The segmentation algorithm could handle extra audio fragments at the end and the beginning of the audio, but prolonged untranscribed audio segments in the middle of the file could deteriorate segmentation results. That is why, it is recommended to normalize text, so that transcripts contain spoken equivalents of abbreviations and numbers.\n",
"\n",
"## 3.1. Prepare Text and Audio\n",
"\n",
"We're going to use `prepare_data.py` script to prepare both text and audio data for segmentation.\n",
"\n",
"### Text preprocessing:\n",
"* the text will be roughly split into sentences and stored under '$OUTPUT_DIR/processed/*.txt' where each sentence is going to start with a new line (we're going to find alignments for these sentences in the next steps)\n",
"* to change the lengths of the final sentences/fragments, specify additional punctuation marks to split the text into fragments, use `--additional_split_symbols` argument. Use `|` as a separator between symbols, for example: `--additional_split_symbols=;|:`\n",
"* `max_length` argument - max number of words in a segment for alignment (used only if there are no punctuation marks present in the original text. Long non-speech segments are better for segments split and are more likely to co-occur with punctuation marks. Random text split could deteriorate the quality of the alignment.\n",
"* out-of-vocabulary words will be removed based on pre-trained ASR model vocabulary, and the text will be changed to lowercase\n",
"* sentences for alignment with the original punctuation and capitalization will be stored under `$OUTPUT_DIR/processed/*_with_punct.txt`\n",
"* numbers will be converted from written to their spoken form with `num2words` package. For English, it's recommended to use NeMo normalization tool use `--use_nemo_normalization` argument (not supported if running this segmentation tutorial in Colab, see the text normalization tutorial: [`https://github.com/NVIDIA/NeMo-text-processing/blob/main/tutorials/Text_(Inverse)_Normalization.ipynb`](https://colab.research.google.com/github/NVIDIA/NeMo-text-processing/blob/main/tutorials/Text_(Inverse)_Normalization.ipynb) for more details). Even `num2words` normalization is usually enough for proper segmentation. However, it does not take audio into account. NeMo supports audio-based normalization for English, German and Russian languages that can be applied to the segmented data as a post-processing step. Audio-based normalization produces multiple normalization options. For example, `901` could be normalized as `nine zero one` or `nine hundred and one`. The audio-based normalization chooses the best match among the possible normalization options and the transcript based on the character error rate. See [https://github.com/NVIDIA/NeMo-text-processing/blob/main/nemo_text_processing/text_normalization/normalize_with_audio.py](https://github.com/NVIDIA/NeMo-text-processing/blob/main/nemo_text_processing/text_normalization/normalize_with_audio.py) for more details.\n",
"\n",
"### Audio preprocessing:\n",
"* non '.wav' audio files will be converted to `.wav` format\n",
"* audio files will be resampled to 16kHz (sampling rate used during training NeMo ASR models)\n",
"* stereo tracks will be converted to mono\n",
"* In some cases, if an audio contains a very long untranscribed prologue, increasing `--cut_prefix` value might help improve segmentation quality.\n",
"\n",
"\n",
"The `prepare_data.py` will preprocess all `.txt` files found in the `--in_text=$DATA_DIR` and all audio files located at `--audio_dir=$DATA_DIR`.\n",
"\n",
"We are going to use [Spanish Citrinet-512 model](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/nemo/models/stt_es_citrinet_512)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "u4zjeVVv-UXR"
},
"outputs": [],
"source": [
"import sys\n",
"MODEL = \"stt_es_citrinet_512\"\n",
"OUTPUT_DIR = WORK_DIR + \"/es_output\"\n",
"\n",
"! rm -rf $OUTPUT_DIR\n",
"\n",
"!{sys.executable} $TOOLS_DIR/prepare_data.py \\\n",
"--in_text=$DATA_DIR/es/text \\\n",
"--output_dir=$OUTPUT_DIR/processed/ \\\n",
"--language='en' \\\n",
"--model=$MODEL \\\n",
"--audio_dir=$DATA_DIR/es/audio"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "kmDTCuTLH7pm"
},
"source": [
"The following four files should be generated and stored at the `$OUTPUT_DIR/processed` folder:\n",
"\n",
"* el19demarzoyel2demayo_03_perezgaldos.txt (lower cased and normalized text with punctuation removed, each line represents an utterance for alignment)\n",
"* el19demarzoyel2demayo_03_perezgaldos.wav (.wav mono file, 16kHz)\n",
"* el19demarzoyel2demayo_03_perezgaldos_with_punct.txt (raw utterances for alignment with punctuation and case preserved)\n",
"* el19demarzoyel2demayo_03_perezgaldos_with_punct_normalized.txt (normalized utterances for alignment utterance for alignment with punctuation and case preserved)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "6R7OKAsYH9p0"
},
"outputs": [],
"source": [
"! ls $OUTPUT_DIR/processed"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "bIvKBwRcH_9W"
},
"source": [
"The `.txt` file without punctuation contains preprocessed text phrases that we're going to align within the audio file. Here, we split the text into sentences. Each line should contain a text snippet for alignment."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "74GLpMgoICmk"
},
"outputs": [],
"source": [
"! head $OUTPUT_DIR/processed/el19demarzoyel2demayo_03_perezgaldos.txt"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "QrvZAjeoR9U1"
},
"source": [
"## 3.2. Run CTC-Segmentation\n",
"\n",
"In this step, we're going to use the [`ctc-segmentation`](https://github.com/lumaku/ctc-segmentation) to find the start and end time stamps for the segments we created during the previous step.\n",
"\n",
"\n",
"As described in the [CTC-Segmentation of Large Corpora for German End-to-end Speech Recognition](https://arxiv.org/pdf/2007.09127.pdf), the algorithm is relying on a CTC-based ASR model to extract utterance segments with exact time-wise alignments."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "xyKtaqAd-Tvk"
},
"outputs": [],
"source": [
"import sys\n",
"WINDOW = 8000\n",
"\n",
"!{sys.executable} $TOOLS_DIR/run_ctc_segmentation.py \\\n",
"--output_dir=$OUTPUT_DIR \\\n",
"--data=$OUTPUT_DIR/processed \\\n",
"--model=$MODEL \\\n",
"--window_len=$WINDOW"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "wY27__e3HmhH"
},
"source": [
"`WINDOW` parameter might need to be adjusted depending on the length of the utterance one wants to align, the default value should work in most cases. By default, if the alignment is not found for the initial `WINDOW` size, the window size will be doubled a few times to re-try backtracing.\n",
"\n",
"Let's take a look at the generated alignments."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "ktBAsfJRVCwI"
},
"outputs": [],
"source": [
"alignment_file = f\"{WINDOW}_{base_name_es}_segments.txt\"\n",
"! head -n 3 $OUTPUT_DIR/segments/$alignment_file"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "xCwEFefHZz1C"
},
"source": [
"The expected output for our audio sample looks like this:\n",
"\n",
"```\n",
"<PATH_TO>/processed/el19demarzoyel2demayo_03_perezgaldos.wav\n",
"11.518331881862931 25.916246734191596 -1.2629864680237006 | entraron en la habitación donde estábamos y al punto que d mauro vio a su sobrina dirigiose a ella con los brazos abiertos y al estrecharla en ellos exclamó endulzando la voz ¡inés de mi alma inocente hija de mi prima juana | Entraron en la habitación donde estábamos, y al punto que D. Mauro vio a su sobrina dirigiose a ella con los brazos abiertos, y al estrecharla en ellos, exclamó endulzando la voz: -¡Inés de mi alma, inocente hija de mi prima Juana! | Entraron en la habitación donde estábamos, y al punto que D. Mauro vio a su sobrina dirigiose a ella con los brazos abiertos, y al estrecharla en ellos, exclamó endulzando la voz: - Inés de mi alma, inocente hija de mi prima Juana!\n",
"25.756269902499053 28.155922377887165 -0.0003830735786323203 | al fin al fin te veo | Al fin, al fin te veo. | Al fin, al fin te veo.\n",
"...\n",
"```\n",
"\n",
"**Details of the file content**:\n",
"- the first line of the file contains the path to the original audio file\n",
"- all subsequent lines contain:\n",
" * the first number is the start of the segment (in seconds)\n",
" * the second one is the end of the segment (in seconds)\n",
" * the third value - alignment confidence score (in log space)\n",
" * text fragments corresponding to the timestamps\n",
" * original text without pre-processing\n",
" * normalized text\n",
"\n",
"Finally, we're going to split the original audio file into segments based on the found alignments. We're going to save only segments with alignment score above the threshold value (default threshold=-2:\n",
"* high scored clips (segments with the segmentation score above the threshold value)\n",
"* low scored clips (segments with the segmentation score below the threshold)\n",
"* deleted segments (segments that were excluded during the alignment. For example, in our sample audio file, the prologue and epilogue that don't have the corresponding transcript were excluded. Oftentimes, deleted files also contain such things as clapping, music, or hard breathing.\n",
"\n",
"The alignment score values depend on the pre-trained model quality and the dataset.\n",
"\n",
"Also note, that the `OFFSET` parameter is something one might want to experiment with since timestamps have a delay (offset) depending on the model.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "6YM64RPlitPL"
},
"outputs": [],
"source": [
"import sys\n",
"OFFSET = 0\n",
"THRESHOLD = -2\n",
"\n",
"!{sys.executable} $TOOLS_DIR/cut_audio.py \\\n",
"--output_dir=$OUTPUT_DIR \\\n",
"--alignment=$OUTPUT_DIR/segments/ \\\n",
"--threshold=$THRESHOLD \\\n",
"--offset=$OFFSET"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "QoyS0T8AZxcx"
},
"source": [
"## 3.3. Transcribe segmented audio\n",
"\n",
"The transcripts will be saved in a new manifest file in `pred_text` field."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "1UaSIflBZwaV"
},
"outputs": [],
"source": [
"import sys\n",
"wget.download(f'https://raw.githubusercontent.com/NVIDIA/NeMo/{BRANCH}/examples/asr/transcribe_speech.py')\n",
"\n",
"!{sys.executable} transcribe_speech.py \\\n",
"pretrained_name=$MODEL \\\n",
"dataset_manifest=$OUTPUT_DIR/manifests/manifest.json \\\n",
"output_filename=$OUTPUT_DIR/manifests/manifest_transcribed.json"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "F-nPT8z_IVD-"
},
"outputs": [],
"source": [
"import soundfile as sf\n",
"\n",
"def plot_signal(signal, sample_rate):\n",
" \"\"\" Plot the signal in time domain \"\"\"\n",
" fig_signal = go.Figure(\n",
" go.Scatter(x=np.arange(signal.shape[0])/sample_rate,\n",
" y=signal, line={'color': 'green'},\n",
" name='Waveform',\n",
" hovertemplate='Time: %{x:.2f} s<br>Amplitude: %{y:.2f}<br><extra></extra>'),\n",
" layout={\n",
" 'height': 200,\n",
" 'xaxis': {'title': 'Time, s'},\n",
" 'yaxis': {'title': 'Amplitude'},\n",
" 'title': 'Audio Signal',\n",
" 'margin': dict(l=0, r=0, t=40, b=0, pad=0),\n",
" }\n",
" )\n",
" fig_signal.show()\n",
"\n",
"def display_samples(manifest):\n",
" \"\"\" Display audio and reference text.\"\"\"\n",
" with open(manifest, 'r') as f:\n",
" for line in f:\n",
" sample = json.loads(line)\n",
"\n",
"\n",
" signal, sample_rate = sf.read(sample['audio_filepath'])\n",
"\n",
" plot_signal(signal, sample_rate)\n",
" display(Audio(sample['audio_filepath']))\n",
" display('Reference text: ' + sample['text_no_preprocessing'])\n",
" if 'pred_text' in sample:\n",
" display('ASR transcript: ' + sample['pred_text'])\n",
" print(f\"Score: {sample['score']}\")\n",
" print('\\n' + '-' * 110)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "S69UFA30ZvxV"
},
"source": [
"Let's examine the high scored segments we obtained.\n",
"\n",
"The `Reference text` in the next cell represents the original text without pre-processing, while `ASR transcript` is an ASR model prediction with greedy decoding. Also notice, that `ASR transcript` in some cases contains errors that could decrease the alignment score, but usually it doesn’t hurt the quality of the aligned segments.\n",
"\n",
"Displaying audio in Jupyter Notebook could be slow, it's recommended to use [Speech Data Explorer](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/tools/speech_data_explorer.html) to analyze speech data."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Q45uBtsHIaAD"
},
"outputs": [],
"source": [
"# let's examine only a few first samples\n",
"! head -n 2 $OUTPUT_DIR/manifests/manifest_transcribed.json > $OUTPUT_DIR/manifests/samples.json\n",
"\n",
"display_samples(f\"{OUTPUT_DIR}/manifests/samples.json\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "yivXpD25T4Ir"
},
"source": [
"# 4. Processing of multiple files (English samples)\n",
"\n",
"Up until now, we were processing only one file at a time, but to create a large dataset processing of multiple files simultaneously could help speed up things considerably.\n",
"\n",
"Our English data folder contains 2 audio files and corresponding text files:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "KRc9yMjPXPgj"
},
"outputs": [],
"source": [
"! ls $DATA_DIR/en/audio $DATA_DIR/en/text"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "3ftilXu-5tzT"
},
"source": [
"We are going to use `run_segmentation.sh` to perform all the above steps starting from the text and audio preprocessing to segmentation and manifest creation in a single step:"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "2I4w34Hepdor"
},
"source": [
"`run_segmentation.sh` script takes `DATA_DIR` argument and assumes that it contains folders `text` and `audio`.\n",
"An example of the `DATA_DIR` folder structure:\n",
"\n",
"\n",
"--DATA_DIR\n",
"\n",
" |----audio\n",
" |---1.mp3\n",
" |---2.mp3\n",
" \n",
" |-----text\n",
" |---1.txt\n",
" |---2.txt"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "nYXNvBDsHMEu"
},
"source": [
"`run_segmentation.sh` could use multiple `WINDOW` sizes for segmentation, and then adds segments that were similarly aligned with various window sizes to `verified_segments` folder. This could be useful to reduce the amount of manual work while checking the alignment quality."
]
},
{
"cell_type": "code",
"source": [
"!pip install sox"
],
"metadata": {
"id": "TijPmstX9x4W"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "hRFAl0gO92bp"
},
"outputs": [],
"source": [
"if 'google.colab' in str(get_ipython()):\n",
" !sudo ln -s /usr/local/bin/python /usr/local/bin/ctc_segmentation #for comatability with run_segmentation.sh\n",
"\n",
"MODEL = \"stt_en_citrinet_512\" # QuartzNet15x5Base-En was removed from NeMo (PR #15507); use a registered CTC model\n",
"OUTPUT_DIR_2 = WORK_DIR + \"/en_output\"\n",
"\n",
"! rm -rf $OUTPUT_DIR_2\n",
"\n",
"! bash $TOOLS_DIR/../run_segmentation.sh \\\n",
"--MODEL_NAME_OR_PATH=$MODEL \\\n",
"--DATA_DIR=$DATA_DIR/en \\\n",
"--OUTPUT_DIR=$OUTPUT_DIR_2 \\\n",
"--SCRIPTS_DIR=$TOOLS_DIR \\\n",
"--MIN_SCORE=$THRESHOLD \\\n",
"--USE_NEMO_NORMALIZATION=False"
]
},
{
"cell_type": "markdown",
"source": [],
"metadata": {
"id": "9qgq9m6Fyp-1"
}
},
{
"cell_type": "markdown",
"metadata": {
"id": "zzJTwKq2Kl9U"
},
"source": [
"Manifest file with segments with alignment score above the threshold values are saved in `en/output/manifests/manifest.json`.\n",
"\n",
"Next, we are going to run `run_filter.sh`. The script does the following:\n",
"* adds ASR transcripts to the manifest\n",
"* calculates and saves metrics such as Word Error Rate (WER), Character Error Rate (CER), CER at the tails of the audio file, word difference between reference and transcript, mean absolute values at the tails of the audio.\n",
"* filters out samples that do not satisfy threshold values and saves selected segments in `manifest_transcribed_metrics_filtered.json`.\n",
"\n",
"Note, it's better to analyze the manifest with metrics in Speech Data Explorer to decide on what thresholds should be used for final sample selection."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "xsm89hYlpdor"
},
"outputs": [],
"source": [
"! bash $TOOLS_DIR/../run_filter.sh \\\n",
"--SCRIPTS_DIR=$TOOLS_DIR \\\n",
"--MODEL_NAME_OR_PATH=stt_en_conformer_ctc_large \\\n",
"--MANIFEST=$OUTPUT_DIR_2/manifests/manifest.json \\\n",
"--INPUT_AUDIO_DIR=$DATA_DIR/en/audio/"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "nacE_iQ2_85L"
},
"outputs": [],
"source": [
"# let's examine only a few first samples\n",
"! head -n 2 $OUTPUT_DIR_2/manifests/manifest_transcribed_metrics_filtered.json > $OUTPUT_DIR_2/manifests/samples.json\n",
"\n",
"display_samples(f\"{OUTPUT_DIR_2}/manifests/samples.json\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "lcvT3P2lQ_GS"
},
"source": [
"# Next Steps\n",
"\n",
"- Check out [NeMo Speech Data Explorer tool](https://github.com/NVIDIA/NeMo/tree/main/tools/speech_data_explorer#speech-data-explorer) to interactively evaluate the aligned segments.\n",
"- Try Audio-based normalization tool."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "GYylwvTX2VSF"
},
"source": [
"# References\n",
"Kürzinger, Ludwig, et al. [\"CTC-Segmentation of Large Corpora for German End-to-End Speech Recognition.\"](https://arxiv.org/abs/2007.09127) International Conference on Speech and Computer. Springer, Cham, 2020."
]
},
{
"cell_type": "code",
"source": [],
"metadata": {
"id": "i_dk6NJA3Vre"
},
"execution_count": null,
"outputs": []
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"private_outputs": true,
"provenance": []
},
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.9.7"
}
},
"nbformat": 4,
"nbformat_minor": 0
}
tutorials/tools/DefinedCrowd_x_NeMo_ASR_Training_Tutorial.ipynb
+1853
View File
File diff suppressed because one or more lines are too long
tutorials/tools/Multispeaker_Simulator.ipynb
+396
View File
@@ -0,0 +1,396 @@
{
"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",
"\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[all] @ git+https://github.com/NVIDIA-NeMo/Speech.git@$BRANCH\"\n",
"%cd .."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Install dependencies\n",
"!apt-get install sox libsndfile1 ffmpeg\n",
"!pip install wget\n",
"!pip install text-unidecode\n",
"!pip install \"matplotlib>=3.3.2\"\n",
"!pip install seaborn"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Multispeaker Simulator\n",
"\n",
"This tutorial shows how to use the speech data simulator to generate synthetic multispeaker audio sessions that can be used to train or evaluate models for multispeaker ASR or speaker diarization. This tool aims to address the lack of labelled multispeaker training data and to help models deal with overlapping speech."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Step 1: Download Required Datasets\n",
"\n",
"The LibriSpeech dataset and corresponding word alignments are required for generating synthetic multispeaker audio sessions. For simplicity, only the dev-clean dataset is used for generating synthetic sessions."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import os\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"
]
},
{
"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": [
"The LibriSpeech forced word alignments are from [this repository](https://github.com/CorentinJ/librispeech-alignments). You can access to the whole LibriSpeech splits at this google drive [link](https://drive.google.com/file/d/1WYfgr31T-PPwMcxuAq09XZfHQO5Mw8fE/view?usp=sharing). We will download the dev-clean part for demo purpose."
]
},
{
"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": [
"# Step 2: Produce Manifest File with Forced Alignments\n",
"\n",
"The LibriSpeech manifest file and LibriSpeech forced alignments will now be merged into one manifest file for ease of use when generating synthetic data. \n",
"\n",
"Here, the input LibriSpeech alignments are first converted to CTM files, and the CTM files are then combined with the base manifest in order to produce the manifest file with word alignments. When using another dataset, the --use_ctm argument can be used to generate the manifest file using alignments in CTM files."
]
},
{
"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": [
"# Step 3: Download Background Noise Data\n",
"\n",
"The background noise is constructed from [here](https://www.openslr.org/resources/28/rirs_noises.zip) (although it can be constructed from other background noise datasets instead)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!wget -nc https://www.openslr.org/resources/28/rirs_noises.zip\n",
"!unzip -o rirs_noises.zip\n",
"!rm -f rirs_noises.zip"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Data Simulator Overview\n",
"\n",
"The simulator creates a speech session using utterances from a desired number of speakers. The simulator first selects the LibriSpeech speaker IDs that will be used for the current session, and sets speaker dominance values for each speaker to determine how often each speaker will talk in the session. The session is then constructed by iterating through the following steps:\n",
"\n",
"* The next speaker is selected (which could be the same speaker again with some probability, and which accounts for the speaker dominance values).\n",
"* The sentence length is determined using a probability distribution, and an utterance of the desired length is then constructed by concatenating together (or truncating) LibriSpeech sentences corresponding to the desired speaker. Individual word alignments are used to truncate the last LibriSpeech sentence such that the entire utterance has the desired length.\n",
"* Next, either the current utterance is overlapped with a previous utterance or silence is introduced before inserting the current utterance. \n",
"\n",
"The simulator includes a multi-microphone far-field mode that incorporates synthetic room impulse response generation in order to simulate multi-microphone multispeaker sessions. When using RIR generation, the RIR is computed once per batch of sessions, and then each constructed utterance is convolved with the RIR in order to get the sound recorded by each microphone before adding the utterance to the audio session. In this tutorial, only near field sessions will be generated.\n",
"\n",
"The simulator also has a speaker enforcement mode which ensures that the correct number of speakers appear in each session, since it is possible that fewer than the desired number may be present since speaker turns are probabilistic. In speaker enforcement mode, the length of the session or speaker probabilities may be adjusted to ensure all speakers are included before the session finishes."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Parameters"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from omegaconf import OmegaConf\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": [
"# Step 4: Create Background Noise Manifest"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!find RIRS_NOISES/real_rirs_isotropic_noises/*.wav > bg_noise.list \n",
"\n",
"# this function can also be run using the pathfiles_to_diarize_manifest.py script\n",
"from nemo.collections.asr.parts.utils.manifest_utils import create_manifest\n",
"create_manifest('bg_noise.list', 'bg_noise.json')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Step 5: Generate Simulated Audio Session\n",
"\n",
"A single 4-speaker session of 30 seconds is generated as an example. "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from nemo.collections.asr.data.data_simulation import MultiSpeakerSimulator\n",
"\n",
"ROOT = os.getcwd()\n",
"data_dir = os.path.join(ROOT,'simulated_data')\n",
"config.data_simulator.random_seed=42\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=1\n",
"config.data_simulator.session_config.session_length=30\n",
"config.data_simulator.background_noise.add_bg=True\n",
"config.data_simulator.background_noise.background_manifest=\"./bg_noise.json\"\n",
"config.data_simulator.session_params.mean_silence=0.2\n",
"config.data_simulator.session_params.mean_silence_var=0.02\n",
"config.data_simulator.session_params.mean_overlap=0.15\n",
"\n",
"lg = MultiSpeakerSimulator(cfg=config)\n",
"lg.generate_sessions()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Step 6: 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_data')\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",
"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(audio)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The visualization is useful for seeing both the distribution of utterance lengths, the differing speaker dominance values, and the amount of overlap in the session."
]
},
{
"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"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"# Step 7: Get Simulated Data Statistics "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import wget\n",
"if not os.path.exists(\"multispeaker_data_analysis.py\"):\n",
" !wget https://raw.githubusercontent.com/NVIDIA/NeMo/main/scripts/speaker_tasks/multispeaker_data_analysis.py\n",
"\n",
"from multispeaker_data_analysis import run_multispeaker_data_analysis\n",
"\n",
"session_dur = 30\n",
"silence_mean = 0.2\n",
"silence_var = 0.1\n",
"overlap_mean = 0.15\n",
"run_multispeaker_data_analysis(data_dir, session_dur=session_dur, silence_var=silence_var, overlap_mean=overlap_mean, precise=True);"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Files Produced\n",
"\n",
"The following files are produced by the 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",
"* json files (one per audio session) - the output manifest file for the corresponding audio session (containing text transcriptions, utterance durations, full paths to audio files, words, and word alignments)\n",
"* ctm files (one per audio session) - contains word-by-word alignments, speaker ID, and word\n",
"* txt files (one per audio session) - contains the full text transcription for a given session\n",
"* list files (one per file type per batch of sessions) - a list of generated files of the given type (wav, rttm, json, ctm, or txt), used primarily for manifest creation\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!ls simulated_data"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "nemo093025",
"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/tools/NeMo_Forced_Aligner_Tutorial.ipynb
+560
View File
@@ -0,0 +1,560 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 88
},
"id": "d4KCUoxSpdoZ",
"outputId": "51c8a307-54e4-4056-a677-2baee15081b8"
},
"outputs": [],
"source": [
"BRANCH = 'main'\n",
"\n",
"\"\"\"\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",
"\"\"\""
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "JDk9zxC6pdod",
"outputId": "3ba8e4ef-65b6-4731-ad27-d34d15adc3d1"
},
"outputs": [],
"source": [
"import os\n",
"# either provide a path to local NeMo repository with NeMo already installed or git clone\n",
"\n",
"# option #1: local path to NeMo repo with NeMo already installed\n",
"NEMO_DIR_PATH = \"NeMo\"\n",
"\n",
"# option #2: download NeMo repo\n",
"if 'google.colab' in str(get_ipython()) or not os.path.exists(NEMO_DIR_PATH):\n",
" !git clone -b $BRANCH https://github.com/NVIDIA-NeMo/Speech\n",
" %cd NeMo\n",
" !python -m pip install \"nemo_toolkit[all] @ git+https://github.com/NVIDIA-NeMo/Speech.git@$BRANCH\"\n",
" %cd -"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "7ZO_421L7bOH"
},
"source": [
"This tutorial contains external links. Each user is responsible for checking the content and the applicable licenses and determining if suitable for the intended use."
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {
"id": "A4NE9GhNn8f-"
},
"source": [
"In this tutorial, we will use [NeMo Forced Aligner](https://github.com/NVIDIA/NeMo/tree/main/tools/nemo_forced_aligner) to generate token and word alignments for a video of Neil Armstrong's first steps on the moon. We will use the ASS-format subtitle files generated by NFA to add subtitles with token-by-token and word-by-word highlighting to the video.\n",
"\n",
"\n",
"We will use the video at this [link](https://www.nasa.gov/wp-content/uploads/static/history/alsj/a11/a11.v1092338.mov), which is in the public domain and was obtained from the NASA website [here](https://history.nasa.gov/alsj/a11/video11.html#Step). The transcript for the video is obtained from the transcript of the mission [here](https://history.nasa.gov/alsj/a11/a11transcript_tec.pdf). As referenced on this [page](https://history.nasa.gov/alsj/a11/a11trans.html), this is a raw transcript with no copyright asserted.\n",
"\n",
"The alignment process is shown below. To better understand the 'Viterbi decoding' process, you can refer to our tutorial [here](https://nvidia.github.io/NeMo/blogs/2023/2023-08-forced-alignment/).\n",
"\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "PGr_hMTCcm3J"
},
"source": [
"![NFA forced alignment pipeline](https://github.com/NVIDIA/NeMo/releases/download/v1.20.0/nfa_forced_alignment_pipeline.png)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "gQOQ-EK1aIBN"
},
"source": [
"\n",
"\n",
"We will do the following:\n",
"\n",
"1. Download the source video.\n",
"2. Prepare a manifest to input to NFA.\n",
"3. Generate alignments using NFA.\n",
"4. Generate a video with subtitles with token-by-token or word-by-word highlighting."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "CH7yR7cSwPKr"
},
"outputs": [],
"source": [
"import json\n",
"import os\n",
"\n",
"from IPython.display import HTML\n",
"from base64 import b64encode"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "40XnfPNop_7V"
},
"source": [
"# 1. Video download\n",
"We will download the video from the provided link and convert it to a higher-resolution format - the video footage will not get any clearer, but it will allow the subtitles to look sharper when we burn them in at the end of the tutorial."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "sl34bUsUp_EV",
"outputId": "a1170a38-7fa0-439f-8ca8-2b79d9d0406d"
},
"outputs": [],
"source": [
"# first make a directory WORK_DIR that we will save all our new files in\n",
"WORK_DIR=\"WORK_DIR\"\n",
"!mkdir $WORK_DIR\n",
"!wget https://www.nasa.gov/wp-content/uploads/static/history/alsj/a11/a11.v1092338.mov --directory-prefix=$WORK_DIR"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "GuFjl8-mrJub",
"outputId": "a216ef90-3d6e-49e5-8c23-88f60b2fc11f"
},
"outputs": [],
"source": [
"# scale up the number of pixels in video so that the subtitles will look sharp when we burn them in later\n",
"!/usr/bin/ffmpeg -loglevel warning -y -i $WORK_DIR/a11.v1092338.mov -vf \"scale=-1:480\" $WORK_DIR/one_small_step.mp4\n",
"# also save the audio as a separate file, so that we can use this as input to NFA\n",
"!/usr/bin/ffmpeg -loglevel warning -y -i $WORK_DIR/a11.v1092338.mov $WORK_DIR/one_small_step.wav"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "laCH__s6KYSF"
},
"source": [
"A note on FFMPEG commands:\n",
"\n",
"*In this tutorial, instead of just calling the `ffmpeg` command, we will specifically call `/usr/bin/ffmpeg`. If we just call `ffmpeg`, there is a chance the conda version of `ffmpeg` will be called (depending on how you installed the dependencies for this tutorial). We have observed that the conda version of `ffmpeg` may generate incorrectly formatted MP4 files with the commands used, and may not be able to generate the subtitle videos which we will create at the end of this tutorial.*\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 322
},
"id": "5In-mlX7pPqq",
"outputId": "47651480-9e64-48e8-d3bb-08bae51e6f50"
},
"outputs": [],
"source": [
"# display video so we know what we will be working with\n",
"mp4 = open(f'{WORK_DIR}/one_small_step.mp4','rb').read()\n",
"data_url = \"data:video/mp4;base64,\" + b64encode(mp4).decode()\n",
"HTML(\"\"\"\n",
"<video width=400 controls>\n",
" <source src=\"%s\" type=\"video/mp4\">\n",
"</video>\n",
"\"\"\" % data_url)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "dNSOrcDyscvV"
},
"source": [
"The audio and video quality of the video isn't great - it was taken on the Moon in 1969 after all. Using NFA, we will still be able to obtain good token & word alignments despite the poor audio quality."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "VGsMdjOsdBn6"
},
"source": [
"# 2. Prepare a manifest to input to NFA.\n",
"NFA requires a manifest as input, in the form shown in the diagram below. As we are only aligning a single audio file, our manifest will only contain one line.\n",
"\n",
"**Note**: the text field is optional, but if you omit it, you need to specify `align_using_pred_text=true` in the config you feed into NFA.\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "i84J0AyiW6X7"
},
"source": [
"![NFA usage pipeline](https://github.com/NVIDIA/NeMo/releases/download/v1.20.0/nfa_run.png)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "lmxvl-tcsyz6"
},
"source": [
"We prepare the text obtained from [here](https://history.nasa.gov/alsj/a11/a11transcript_tec.pdf).\n",
"\n",
"We will specify in the NFA config `additional_segment_grouping_separator=['.', '?', '!', '...']` to separate the text into logically defined segments that are finished expressions. This will make sure NFA produces timestamps for each segment separated by `additional_segment_grouping_separator`. It also means that, by default, the ASS subtitle files will display those segments together.\n",
"\n",
"**Extra info**: Alternatively, you can specify in the NFA config that `ass_file_config.resegment_text_to_fill_space=true`, which will cause automatic grouping of words in the ASS files. You can also specify `ass_file_config.max_lines_per_segment` to set the maximum number of lines of text that will appear at a time."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "IS-ef4o5s_tv"
},
"outputs": [],
"source": [
"text = \"\"\"\n",
"I'm at the foot of the ladder. The LM footpads are only depressed in the \n",
"surface about 1 or 2 inches, although the surface appears to be very, very \n",
"fine grained, as you get close to it. It's almost like a powder. \n",
"Down there, it's very fine. I'm going to step off the LM now. That's one \n",
"small step for man, one giant leap for mankind.\n",
"\"\"\""
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "OHq0trf2h24g"
},
"outputs": [],
"source": [
"manifest_filepath = f\"{WORK_DIR}/manifest.json\"\n",
"manifest_data = {\n",
" \"audio_filepath\": f\"{WORK_DIR}/one_small_step.wav\",\n",
" \"text\": text\n",
"}\n",
"with open(manifest_filepath, 'w') as f:\n",
" line = json.dumps(manifest_data)\n",
" f.write(line + \"\\n\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "RvTrg8C3wHYw",
"outputId": "b257a766-4dfc-4f69-f75e-4b3bbf84ef60"
},
"outputs": [],
"source": [
"!cat $manifest_filepath"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "QKnTJG5Ig4lP"
},
"source": [
"# 3. Run NFA"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {
"id": "izC7x71-T0Il"
},
"source": [
"We run NFA by calling the `align.py` function and passing in config parameters. In our case these want to set:\n",
"\n",
"**Compulsory parameters**\n",
"\n",
"* `pretrained_name` (either this or a `model_path` is required) - this is the name of the ASR model that we will use to do alignment. You can use any of NeMo's CTC models for alignment.\n",
"* `manifest_filepath` - the path to the manifest specifying the audio that you want to align and its reference text.\n",
"* `output_dir` - the path to the directory where your output files will be saved.\n",
"\n",
"**Optional parameters**\n",
"* `additional_segment_grouping_separator` - a list of strings that will be used to separate the text into segments. In our case we use default values of `['.', '?', '!', '...']`.\n",
"* `ass_file_config.vertical_alignment` - by default this is set to `\"center\"`, meaning the subtitles will appear in the center of the screen. We want them to appear at the bottom of our video, so we set this parameter to `\"bottom\"`.\n",
"* `ass_file_config.text_already_spoken_rgb`, `ass_file_config.text_being_spoken_rgb`, `ass_file_config.text_not_yet_spoken_rgb` - these parameters define the RGB values of the colors of the text. The default colors do not show up very clearly against our video, so we set them to some different values, which will generate the colors teal, yellow, and a very pale blue."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "QDcyfW_RWrqB",
"outputId": "9b76c513-30bd-44d6-ef70-7daba8f604f0"
},
"outputs": [],
"source": [
"!python $NEMO_DIR_PATH/tools/nemo_forced_aligner/align.py \\\n",
" pretrained_name=\"stt_en_fastconformer_hybrid_large_pc\" \\\n",
" manifest_filepath=$manifest_filepath \\\n",
" output_dir=$WORK_DIR/nfa_output/ \\\n",
" additional_segment_grouping_separator='[\".\",\"?\",\"!\",\"...\"]' \\\n",
" ass_file_config.vertical_alignment=\"bottom\" \\\n",
" ass_file_config.text_already_spoken_rgb=[66,245,212] \\\n",
" ass_file_config.text_being_spoken_rgb=[242,222,44] \\\n",
" ass_file_config.text_not_yet_spoken_rgb=[223,242,239]"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "dHU-YmALUvVf"
},
"source": [
"The alignment process should have finished successfully, let's look at some of the output files."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "VSyM9zp4wQSJ",
"outputId": "3d77e04d-cc56-425a-92bc-d70a34545f00"
},
"outputs": [],
"source": [
"!head $WORK_DIR/nfa_output/ctm/*/*.ctm\n",
"\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "-LESm1cudqZ0"
},
"source": [
"The token timestamps are produced directly from the Viterbi alignment process, and the word and segment timestamps are obtained from the token timestamps using a simple grouping process shown below:"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "OXUL-KyUdVpL"
},
"source": [
"![How NFA generates word and segment alignments from token alignments](https://github.com/NVIDIA/NeMo/releases/download/v1.20.0/nfa_word_segment_alignments.png)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "3ZZuKU5YwXBh"
},
"source": [
"# 4. Make subtitled video"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "tlYr1ymnU4Du"
},
"source": [
"To make the subtitled video, we will use the ASS-format files that NFA produces."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "RauJ91RgzCvp",
"outputId": "b0b34085-847a-4af0-d937-b67de12304e5"
},
"outputs": [],
"source": [
"!head -n20 $WORK_DIR/nfa_output/ass/words/*.ass"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "OaYjR7mDdVKF",
"outputId": "2fedc7ae-b931-409f-a338-5a36e6c457c4"
},
"outputs": [],
"source": [
"!head -n20 $WORK_DIR/nfa_output/ass/tokens/*.ass"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "ipIlViegVDgQ"
},
"source": [
"We burn the subtitles into the video using the below commands. We generate 2 videos, one with token-by-token highlighting, and one with word-by-word highlighting."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "Av7L4Qzojykv",
"outputId": "dd1dc2af-3631-4776-e79d-157b61afa98e"
},
"outputs": [],
"source": [
"!/usr/bin/ffmpeg -loglevel warning -y -i \"$WORK_DIR/one_small_step.mp4\" \\\n",
" -vf \"ass=$WORK_DIR/nfa_output/ass/tokens/one_small_step.ass\" \\\n",
" $WORK_DIR/one_small_step_tokens_aligned.mp4\n",
"\n",
"!/usr/bin/ffmpeg -loglevel warning -y -i \"$WORK_DIR/one_small_step.mp4\" \\\n",
" -vf \"ass=$WORK_DIR/nfa_output/ass/words/one_small_step.ass\" \\\n",
" $WORK_DIR/one_small_step_words_aligned.mp4"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "gfhQxkbPVPSj"
},
"source": [
"Let's look at the resulting videos below.\n",
"\n",
"You can see that the token timestamps (in the first video) are very accurate, even despite the poor audio quality of the video.\n",
"\n",
"The word timestamps (in the second video) are also very good. The only noticeable mistakes are when the word has punctuation at the end (or beginning). This is because punctuation that is not separated from a word by a space is considered to be part of that word. If the alignment for the punctuation is in a region of non-speech, then the word alignment will also contain that region of non-speech."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 322
},
"id": "ik6FcCgIMYsa",
"outputId": "4781a087-a8a2-44a7-9e0f-ea57e05cde59"
},
"outputs": [],
"source": [
"mp4 = open(f'{WORK_DIR}/one_small_step_tokens_aligned.mp4','rb').read()\n",
"data_url = \"data:video/mp4;base64,\" + b64encode(mp4).decode()\n",
"HTML(\"\"\"\n",
"<video width=400 controls>\n",
" <source src=\"%s\" type=\"video/mp4\">\n",
"</video>\n",
"\"\"\" % data_url)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 322
},
"id": "mgFUhdlsdXLs",
"outputId": "94f8ae2f-8752-4671-aeca-a19f03db7fd6"
},
"outputs": [],
"source": [
"mp4 = open(f'{WORK_DIR}/one_small_step_words_aligned.mp4','rb').read()\n",
"data_url = \"data:video/mp4;base64,\" + b64encode(mp4).decode()\n",
"HTML(\"\"\"\n",
"<video width=400 controls>\n",
" <source src=\"%s\" type=\"video/mp4\">\n",
"</video>\n",
"\"\"\" % data_url)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Gr70zxMoeLqR"
},
"outputs": [],
"source": []
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"provenance": []
},
"gpuClass": "standard",
"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/tools/SDE_HowTo_v2.ipynb
+1116
View File
File diff suppressed because one or more lines are too long
tutorials/tools/label-studio/setup-asr-preannotations.ipynb
+249
View File
@@ -0,0 +1,249 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "exact-strand",
"metadata": {},
"source": [
"## Install NVidia NeMo environment"
]
},
{
"cell_type": "markdown",
"id": "compound-found",
"metadata": {},
"source": [
"You can locally install NeMo environment by following [installation guide](https://github.com/heartexlabs/NeMo#installation), or quickstart it from the prebuilt Docker container:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "pacific-pepper",
"metadata": {},
"outputs": [],
"source": [
"!docker run --gpus all -it --rm --shm-size=8g \\\n",
"-p 8888:8888 -p 6006:6006 -p 8080:8080 --ulimit memlock=-1 --ulimit \\\n",
"stack=67108864 --device=/dev/snd nvcr.io/nvidia/nemo:1.0.1"
]
},
{
"cell_type": "markdown",
"id": "strong-therapist",
"metadata": {},
"source": [
"Note that the default Label Studio port 8080 is exposed from Docker."
]
},
{
"cell_type": "markdown",
"id": "everyday-depth",
"metadata": {},
"source": [
"## Install Label Studio"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "spoken-venice",
"metadata": {},
"outputs": [],
"source": [
"!pip install label-studio"
]
},
{
"cell_type": "markdown",
"id": "integral-introduction",
"metadata": {},
"source": [
"## Create ML backend with NeMo model"
]
},
{
"cell_type": "markdown",
"id": "dimensional-playing",
"metadata": {},
"source": [
"Let's create a simple script `asr.py` that wraps NeMo inference call and converts its output to annotation format expected by Label Studio"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "tribal-minority",
"metadata": {},
"outputs": [],
"source": [
"import nemo\n",
"import nemo.collections.asr as nemo_asr\n",
"from label_studio_ml.model import LabelStudioMLBase\n",
"\n",
"\n",
"class NemoASR(LabelStudioMLBase):\n",
"\n",
" def __init__(self, model_name='QuartzNet15x5Base-En', **kwargs):\n",
" super(NemoASR, self).__init__(**kwargs)\n",
"\n",
" # Find TextArea control tag and bind ASR model to it\n",
" self.from_name, self.to_name, self.value = self._bind_to_textarea()\n",
"\n",
" # This line will download pre-trained QuartzNet15x5 model from NVIDIA's NGC cloud and instantiate it for you\n",
" self.model = nemo_asr.models.EncDecCTCModel.from_pretrained(model_name=model_name)\n",
"\n",
" def predict(self, tasks, **kwargs):\n",
" \"\"\"Returns NeMo ASR predictions given audio files in Label Studio's tasks\"\"\"\n",
" audio_path = self.get_local_path(tasks[0]['data'][self.value])\n",
" transcription = self.model.transcribe(paths2audio_files=[audio_path])[0]\n",
" return [{\n",
" 'result': [{\n",
" 'from_name': self.from_name,\n",
" 'to_name': self.to_name,\n",
" 'type': 'textarea',\n",
" 'value': {\n",
" 'text': [transcription]\n",
" }\n",
" }],\n",
" 'score': 1.0\n",
" }]\n",
"\n",
" def _bind_to_textarea(self):\n",
" \"\"\"Helper to bind inference output to annotation format expected by Label Studio\"\"\"\n",
" from_name, to_name, value = None, None, None\n",
" for tag_name, tag_info in self.parsed_label_config.items():\n",
" if tag_info['type'] == 'TextArea':\n",
" from_name = tag_name\n",
" if len(tag_info['inputs']) > 1:\n",
" logger.warning(\n",
" 'ASR model works with single Audio or AudioPlus input, '\n",
" 'but {0} found: {1}. We\\'ll use only the first one'.format(\n",
" len(tag_info['inputs']), ', '.join(tag_info['to_name'])))\n",
" if tag_info['inputs'][0]['type'] not in ('Audio', 'AudioPlus'):\n",
" raise ValueError('{0} tag expected to be of type Audio or AudioPlus, but type {1} found'.format(\n",
" tag_info['to_name'][0], tag_info['inputs'][0]['type']))\n",
" to_name = tag_info['to_name'][0]\n",
" value = tag_info['inputs'][0]['value']\n",
" if from_name is None:\n",
" raise ValueError('ASR model expects <TextArea> tag to be presented in a label config.')\n",
" return from_name, to_name, value"
]
},
{
"cell_type": "markdown",
"id": "induced-pacific",
"metadata": {},
"source": [
"## Run ML backend"
]
},
{
"cell_type": "markdown",
"id": "fuzzy-malta",
"metadata": {},
"source": [
"The following initializes ML backend by creating a directory `./nemo-ml-backend` and copying everything needed to run, including `asr.py` script."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "outstanding-russell",
"metadata": {},
"outputs": [],
"source": [
"!label-studio-ml init nemo-ml-backend --from asr.py"
]
},
{
"cell_type": "markdown",
"id": "hybrid-thread",
"metadata": {},
"source": [
"Then launch ML backend serving on default `http://localhost:9090`"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "black-hazard",
"metadata": {},
"outputs": [],
"source": [
"!label-studio-ml start nemo-ml-backend"
]
},
{
"cell_type": "markdown",
"id": "private-recommendation",
"metadata": {},
"source": [
"## Connect ML backend to Label Studio"
]
},
{
"cell_type": "markdown",
"id": "aerial-circulation",
"metadata": {},
"source": [
"Launch Label Studio web application running on `http://localhost:8080`"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "skilled-giant",
"metadata": {},
"outputs": [],
"source": [
"!label-studio start annotation-with-nemo --init"
]
},
{
"cell_type": "markdown",
"id": "afraid-revision",
"metadata": {},
"source": [
"In Label Studio, upload audio files either by drag-and-drop, or by importing a text file with one URL referencing an audio file per line. Then, go to the **Settings** page and select the **Speech Transcription** template. Click **Save**."
]
},
{
"cell_type": "markdown",
"id": "suffering-respect",
"metadata": {},
"source": [
"On the **Model** page, add the ML backend URL `http://localhost:9090`. If it connects successfully, you see \"Connected\" status in green."
]
},
{
"cell_type": "markdown",
"id": "alike-realtor",
"metadata": {},
"source": [
"Then you can start to annotate your audio files by correcting the text areas prepopulated by NeMo ASR's output. After you finish labeling, you can export results in the `ASR_MANIFEST` format ready to use for [training a NeMo ASR model](https://colab.research.google.com/github/NVIDIA/NeMo/blob/stable/tutorials/asr/ASR_with_NeMo.ipynb)"
]
}
],
"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.8.2"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
tutorials/tts/Audio_Codec_Inference.ipynb
+478
View File
@@ -0,0 +1,478 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "7X-TwhdTGmlc"
},
"source": [
"# License"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "fCQUeZRPGnoe"
},
"source": [
"> Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n",
">\n",
"> Licensed under the Apache License, Version 2.0 (the \"License\"); you may not use this file except in compliance with the License. You may obtain a copy of the License at\n",
">\n",
"> http://www.apache.org/licenses/LICENSE-2.0\n",
">\n",
"> Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "rtBDkKqVGZJ8"
},
"source": [
"# Introduction"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "pZ2QSsXuGbMe"
},
"source": [
"In this tutorial we show how use NeMo **neural audio codecs** at inference time. To learn more about training and finetuning neural audio codecs in NeMo, check the [Audio Codec Training tutorial](https://github.com/NVIDIA/NeMo/blob/main/tutorials/tts/Audio_Codec_Training.ipynb).\n",
"\n",
"An audio codec typically consists of an encoder, a quantizer and a decoder, with a typical architecture depicted in the figure below.\n",
"An audio codec can be used to encode an input audio signal into a sequence of discrete values.\n",
"In this tutorial, the discrete values will be referred to as **audio tokens**.\n",
"The obtained audio tokens can be decoded into an output audio signal.\n",
"\n",
"Audio tokens can be used to represent the input audio for an automatic speech recognition (ASR) model [[1](https://arxiv.org/abs/2309.10922), [2](https://arxiv.org/pdf/2407.03495)], or to represent the output audio of a text-to-speech (TTS) system [[3](https://arxiv.org/abs/2406.05298), [4](https://arxiv.org/pdf/2406.17957)].\n",
"\n",
"NeMo provides several neural audio codec models, inlcuding audio codecs and mel codecs at different sampling rates.\n",
"The list of the available models can be found [here](https://docs.nvidia.com/nemo-framework/user-guide/latest/nemotoolkit/tts/checkpoints.html#codec-models).\n",
"\n",
"<div>\n",
"<img src=\"https://github.com/NVIDIA/NeMo/releases/download/v1.22.0/nemo_audio_codec.png\" width=\"800\", height=\"400\"/>\n",
"</div>"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "3OZassNG5xff"
},
"source": [
"# Setup"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "WZvQvPkIhRi3"
},
"outputs": [],
"source": [
"BRANCH = 'main'\n",
"# Install NeMo library. If you are running locally (rather than on Google Colab), follow the instructions at https://github.com/NVIDIA/NeMo#Installation\n",
"\n",
"if 'google.colab' in str(get_ipython()):\n",
" !python -m pip install \"nemo_toolkit[all] @ git+https://github.com/NVIDIA-NeMo/Speech.git@$BRANCH\""
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "v8NGOM0EzK8W"
},
"outputs": [],
"source": [
"import math\n",
"import wget\n",
"import os\n",
"import librosa\n",
"import torch\n",
"import numpy as np\n",
"import IPython.display as ipd\n",
"import matplotlib.pyplot as plt\n",
"from pathlib import Path\n",
"\n",
"\n",
"# Utility for displaying signals and metrics\n",
"def show_signal(signal: np.ndarray, sample_rate: int = 16000, tag: str = 'Signal'):\n",
" \"\"\"Show the time-domain signal and its spectrogram.\n",
" \"\"\"\n",
" fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(12, 2.5))\n",
"\n",
" # show waveform\n",
" t = np.arange(0, len(signal)) / sample_rate\n",
"\n",
" ax[0].plot(t, signal)\n",
" ax[0].set_xlim(0, t.max())\n",
" ax[0].grid()\n",
" ax[0].set_xlabel('time / s')\n",
" ax[0].set_ylabel('amplitude')\n",
" ax[0].set_title(tag)\n",
"\n",
" n_fft = 1024\n",
" hop_length = 256\n",
"\n",
" D = librosa.amplitude_to_db(np.abs(librosa.stft(signal, n_fft=n_fft, hop_length=hop_length)), ref=np.max)\n",
" img = librosa.display.specshow(D, y_axis='linear', x_axis='time', sr=sample_rate, n_fft=n_fft, hop_length=hop_length, ax=ax[1])\n",
" ax[1].set_title(tag)\n",
"\n",
" plt.tight_layout()\n",
" plt.colorbar(img, format=\"%+2.f dB\", ax=ax)\n",
"\n",
"\n",
"# Utility for displaying a latent representation\n",
"def show_latent(latent: np.ndarray, tag: str):\n",
" plt.figure(figsize = (16, 3))\n",
" img = plt.imshow(latent, aspect='equal')\n",
" plt.colorbar(img, ax=plt.gca())\n",
" plt.title(tag)\n",
" plt.xlabel('Time frame')\n",
" plt.ylabel('Latent vector index')\n",
" plt.tight_layout()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "8ZKDMTwsEY1K"
},
"outputs": [],
"source": [
"# Working directory\n",
"ROOT_DIR = Path().absolute() / 'codec_tutorial'\n",
"\n",
"# Create dataset directory\n",
"DATA_DIR = ROOT_DIR / 'data'\n",
"DATA_DIR.mkdir(parents=True, exist_ok=True)\n",
"\n",
"audio_path = DATA_DIR / 'LJ023-0089.wav'\n",
"audio_url = \"https://multilangaudiosamples.s3.us-east-2.amazonaws.com/LJ023-0089.wav\"\n",
"\n",
"if not os.path.exists(audio_path):\n",
" wget.download(audio_url, audio_path.as_posix())"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "KAbH7N427FdT"
},
"source": [
"# Load a model from NGC"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "ODgdGgsAAUku"
},
"source": [
"Any of the [pretrained checkpoints](https://docs.nvidia.com/nemo-framework/user-guide/latest/nemotoolkit/tts/checkpoints.html#codec-models) could be used for inference.\n",
"Here, we use `mel_codec_22khz_fullband_medium`, which works for 22.05 kHz audio signals.\n",
"\n",
"The model can be easily restored from NGC:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "XqAYWR65aKTx"
},
"outputs": [],
"source": [
"from nemo.collections.tts.models.audio_codec import AudioCodecModel\n",
"\n",
"# Optionally specify a pretrained model to fine-tune from. To train from scratch, set this to 'None'.\n",
"model_name = 'mel_codec_22khz_fullband_medium'\n",
"codec_model = AudioCodecModel.from_pretrained(model_name)\n",
"codec_model.freeze()"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "ZnnjL28pEY1L"
},
"source": [
"Show information about the loaded model:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "4xsfeHVyEY1L"
},
"outputs": [],
"source": [
"print(f'Loaded model from NeMo:')\n",
"print(f'\\tmodel name : {model_name}')\n",
"print(f'\\tsample rate : {codec_model.sample_rate} Hz')\n",
"print(f'\\tlatent dimension : {codec_model.vector_quantizer.codebook_dim}')\n",
"\n",
"print('\\n\\nModel summary:')\n",
"print(codec_model.summarize())"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "fM4QPsLTnzK7"
},
"source": [
"# Inference"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "tkZC6Dl7KRl6"
},
"source": [
"## Processing audio\n",
"\n",
"Here we use the codec model to process the input audio by applying the complete model. The input signal is encoded, quantized, dequantized and decoded. Finally, a reconstructed signal is obtained."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "sYzvAYr2vo1K"
},
"outputs": [],
"source": [
"input_audio, sr = librosa.load(audio_path, sr=codec_model.sample_rate)\n",
"\n",
"# Shape (batch, time)\n",
"input_audio_tensor = torch.from_numpy(input_audio).unsqueeze(dim=0).to(codec_model.device)\n",
"\n",
"# Shape (batch,)\n",
"input_audio_len = torch.tensor([input_audio_tensor.size(-1)]).to(codec_model.device)\n",
"\n",
"# Process audio using the codec model\n",
"output_audio_tensor, _ = codec_model(audio=input_audio_tensor, audio_len=input_audio_len)\n",
"\n",
"# Output audio\n",
"output_audio = output_audio_tensor.squeeze().cpu().numpy()\n",
"\n",
"# Show signals\n",
"show_signal(input_audio, tag='Input audio', sample_rate=codec_model.sample_rate)\n",
"show_signal(output_audio, tag='Output audio', sample_rate=codec_model.sample_rate)\n",
"\n",
"# Play audio\n",
"print('Input audio')\n",
"ipd.display(ipd.Audio(input_audio, rate=codec_model.sample_rate))\n",
"\n",
"print('Output audio')\n",
"ipd.display(ipd.Audio(output_audio, rate=codec_model.sample_rate))"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "rynZYwg2VP5d"
},
"source": [
"## Audio tokens\n",
"\n",
"Audio tokens can be easily computed by using the `encode` method of the `AudioCodec` model."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "ASKM_jKVEY1L"
},
"outputs": [],
"source": [
"# Convert audio to tokens\n",
"tokens, tokens_len = codec_model.encode(audio=input_audio_tensor, audio_len=input_audio_len)\n",
"\n",
"print('tokens information:')\n",
"print(f'\\tshape (batch, codebook, time frame) : {tokens.size()}')\n",
"print(f'\\tdtype : {tokens.dtype}')\n",
"print(f'\\tmin : {tokens.min()}')\n",
"print(f'\\tmax : {tokens.max()}')\n",
"\n",
"# Number of codebooks should match the number of codebooks/groups\n",
"if hasattr(codec_model.vector_quantizer, 'num_groups'):\n",
" # Group FSQ\n",
" assert tokens.size(1) == codec_model.vector_quantizer.num_groups\n",
" print(f'\\tnum_groups : {tokens.size(1)}')\n",
"elif hasattr(codec_model.vector_quantizer, 'codebooks'):\n",
" # RVQ\n",
" assert tokens.size(1) == len(codec_model.vector_quantizer.codebooks)\n",
" print(f'\\tnum_codebooks : {tokens.size(1)}')"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "CmliPMnDEY1L"
},
"source": [
"Similarly, audio can be easily reconstructed from audio tokens using the `decode` method of the `AudioCodec` models."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "RTQ1M9PMEY1L"
},
"outputs": [],
"source": [
"# Convert tokens back to audio\n",
"output_audio_from_tokens_tensor, _ = codec_model.decode(tokens=tokens, tokens_len=tokens_len)\n",
"output_audio_from_tokens = output_audio_from_tokens_tensor.squeeze().cpu().numpy()\n",
"\n",
"# Show signals\n",
"show_signal(output_audio_from_tokens, tag='Output audio from tokens', sample_rate=codec_model.sample_rate)\n",
"show_signal(output_audio_from_tokens - output_audio, tag='Difference compared to forward pass', sample_rate=codec_model.sample_rate)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "kGqotZkqEY1M"
},
"source": [
"## Latent representation\n",
"\n",
"Continuous (non-discrete) latent representation at the output of the encoder can be easily computed using the `encode_audio` method of the `AudioCodec` model."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "r-89-gG3EY1M"
},
"outputs": [],
"source": [
"# Convert audio to the encoded representation\n",
"encoded, encoded_len = codec_model.encode_audio(audio=input_audio_tensor, audio_len=input_audio_len)\n",
"\n",
"print('encoded information:')\n",
"print(f'\\tshape (batch, codebook, time frame) : {encoded.size()}')\n",
"print(f'\\tdtype : {encoded.dtype}')\n",
"print(f'\\tmin : {encoded.min()}')\n",
"print(f'\\tmax : {encoded.max()}')\n",
"\n",
"\n",
"# Show the encoded representation\n",
"show_latent(encoded.squeeze().cpu().numpy(), tag='Encoder output')"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "3Ory1U1uEY1M"
},
"source": [
"The encoded representation can be easily converted to tokens, dequantized into a continuous latent representation and decoded back to audio."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "btmqUWNkEY1M"
},
"outputs": [],
"source": [
"# Encoder output to tokens\n",
"tokens = codec_model.quantize(encoded=encoded, encoded_len=encoded_len)\n",
"\n",
"# Tokens back to a continuous vector\n",
"dequantized = codec_model.dequantize(tokens=tokens, tokens_len=encoded_len)\n",
"\n",
"# Reconstruct audio\n",
"output_audio_from_latent_tensor, _ = codec_model.decode_audio(inputs=dequantized, input_len=encoded_len)\n",
"output_audio_from_latent = output_audio_from_latent_tensor.squeeze().cpu().numpy()\n",
"\n",
"# Show dequantized latent representation\n",
"show_latent(dequantized.squeeze().cpu().numpy(), tag='Decoder input')\n",
"\n",
"# Show signals\n",
"show_signal(output_audio_from_latent, tag='Output audio from latent', sample_rate=codec_model.sample_rate)\n",
"show_signal(output_audio_from_latent - output_audio, tag='Difference compared to forward pass', sample_rate=codec_model.sample_rate)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "cMvU0WxlEY1M"
},
"source": [
"# Related information"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "_LtyHHuLkNDv"
},
"source": [
"To learn more about audio codec models in NeMo, look at our [documentation](https://docs.nvidia.com/nemo-framework/user-guide/latest/nemotoolkit/tts/models.html#codecs).\n",
"\n",
"For more information on training and finetuning neural audio codecs in NeMo, check the [Audio Codec Training tutorial](https://github.com/NVIDIA/NeMo/blob/main/tutorials/tts/Audio_Codec_Training.ipynb)."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "LeqV3VvJVOb-"
},
"source": [
"# References"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Rvu4w2x_3RSY"
},
"source": [
"1. [Discrete Audio Representation as an Alternative to Mel-Spectrograms for Speaker and Speech Recognition](https://arxiv.org/abs/2309.10922)\n",
"2. [Codec-ASR: Training Performant Automatic Speech Recognition Systems with Discrete Speech Representations](https://arxiv.org/pdf/2407.03495)\n",
"3. [Spectral Codecs: Spectrogram-Based Audio Codecs for High Quality Speech Synthesis](https://arxiv.org/abs/2406.05298)\n",
"4. [Improving Robustness of LLM-based Speech Synthesis by Learning Monotonic Alignment](https://arxiv.org/pdf/2406.17957)"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"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"
},
"colab": {
"provenance": [],
"toc_visible": true
}
},
"nbformat": 4,
"nbformat_minor": 0
}
tutorials/tts/Audio_Codec_Training.ipynb
+796
View File
@@ -0,0 +1,796 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "7X-TwhdTGmlc"
},
"source": [
"# License"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "fCQUeZRPGnoe"
},
"source": [
"> Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n",
">\n",
"> Licensed under the Apache License, Version 2.0 (the \"License\"); you may not use this file except in compliance with the License. You may obtain a copy of the License at\n",
">\n",
"> http://www.apache.org/licenses/LICENSE-2.0\n",
">\n",
"> Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "rtBDkKqVGZJ8"
},
"source": [
"# Introduction"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "pZ2QSsXuGbMe"
},
"source": [
"In this tutorial we show how to use NeMo to train and fine-tune **neural audio codecs**.\n",
"\n",
"Neural audio codecs are deep learning models that compress audio into a low bitrate representation. The compact embedding space created by these models can be useful for various speech tasks, such as TTS and ASR.\n",
"\n",
"<div>\n",
"<img src=\"https://github.com/NVIDIA/NeMo/releases/download/v1.22.0/nemo_audio_codec.png\" width=\"800\", height=\"400\"/>\n",
"</div>\n",
"\n",
"Audio codec models typically have an *encoder-quantizer-decoder* structure. The **encoder** takes an input audio signal and encodes it into a sequence of embeddings. The **quantizer** discretizes the embeddings to create a lookup table known as a **codebook**. The embeddings saved in the codebook are referred to as **audio codes**. The **decoder** takes the audio codes as input and attempts to reconstruct the original audio signal.\n",
"\n",
"To store compressed audio we only need to save the codebook index for each embedding in an audio sequence. This is how audio codec models achieve low bitrates. The codebook indices for an audio are referred to **audio tokens**. It is becoming common for speech generation models to synthesize speech by predicting audio tokens.\n",
"\n",
"In NeMo we have implementations of the [SEANet encoder and decoder](https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/tts/modules/encodec_modules.py#L146) used by [EnCodec](https://github.com/facebookresearch/encodec). As well as a [ResNet encoder](https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/tts/modules/audio_codec_modules.py#L1035) and [HiFi-GAN decoder](https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/tts/modules/audio_codec_modules.py#L875). For quantizers we support [Residual Vector Quantizer](https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/tts/modules/encodec_modules.py#L694) (**RVQ**) and [Finite Scalar Quantizer](https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/tts/modules/audio_codec_modules.py#L409) (**FSQ**).\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "3OZassNG5xff"
},
"source": [
"# Install"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "WZvQvPkIhRi3"
},
"outputs": [],
"source": [
"BRANCH = 'main'\n",
"# Install NeMo library. If you are running locally (rather than on Google Colab), comment out the below line\n",
"# and instead follow the instructions at https://github.com/NVIDIA/NeMo#Installation\n",
"!python -m pip install \"nemo_toolkit[tts] @ git+https://github.com/NVIDIA-NeMo/Speech.git@$BRANCH\""
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "v8NGOM0EzK8W"
},
"outputs": [],
"source": [
"from pathlib import Path"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "tvsgWO_WhV3M"
},
"outputs": [],
"source": [
"# Directory where tutorialscripts will run and outputs will be saved.\n",
"ROOT_DIR = Path().absolute() / \"codec_tutorial\"\n",
"\n",
"# Nemo code paths\n",
"NEMO_DIR = ROOT_DIR / \"nemo\"\n",
"NEMO_SCRIPT_DIR = NEMO_DIR / \"scripts\" / \"dataset_processing\" / \"tts\"\n",
"NEMO_EXAMPLES_DIR = NEMO_DIR / \"examples\" / \"tts\"\n",
"NEMO_CONFIG_DIR = NEMO_EXAMPLES_DIR / \"conf\"\n",
"\n",
"nemo_download_dir = str(NEMO_DIR)\n",
"# Download local version of NeMo scripts. If you are running locally and want to use your own local NeMo code,\n",
"# comment out the below line and set NEMO_ROOT_DIR to your local path.\n",
"!git clone -b $BRANCH https://github.com/NVIDIA-NeMo/Speech.git $nemo_download_dir"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "KAbH7N427FdT"
},
"source": [
"# Configuration"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "ODgdGgsAAUku"
},
"source": [
"Predefined model configurations are available in https://github.com/NVIDIA/NeMo/tree/main/examples/tts/conf/audio_codec.\n",
"\n",
"Configurations available include:\n",
"\n",
"* **audio_codec_*.yaml**: Audio codec configurations optimized for various sampling rates.\n",
"* **mel_codec_*.yaml**: A mel-spectrogram based codec designed to maximize the performance of speech synthesis models.\n",
"* **encodec_*.yaml**: A reproduction of the original [EnCodec](https://arxiv.org/abs/2210.13438) model setup.\n",
"* **audio_codec_low_frame_rate_22050.yaml**: The [Low Frame-rate Speech Codec](https://arxiv.org/abs/2409.12117): an audio codec that achieves high-quality audio compression with a 1.89 kbps bitrate and 21.5 frames per second.\n",
"\n",
"This tutorial can be run with any of our predefined configs. As a default we have selected `audio_codec_16000.yaml`, which works for 16kHz audio."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "SPtjS2LkzE9Q"
},
"outputs": [],
"source": [
"from omegaconf import OmegaConf"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "iCPJFKg63Dsv"
},
"outputs": [],
"source": [
"CONFIG_FILENAME = \"audio_codec_22050.yaml\"\n",
"CONFIG_DIR = NEMO_CONFIG_DIR / \"audio_codec\"\n",
"\n",
"config_filepath = CONFIG_DIR / CONFIG_FILENAME\n",
"\n",
"if not config_filepath.exists():\n",
" raise ValueError(f\"Config file does not exist {config_filepath}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "QE0HYh7FjAR3"
},
"outputs": [],
"source": [
"# Read model name and sample rate from model configuration\n",
"omega_conf = OmegaConf.load(config_filepath)\n",
"MODEL_NAME = omega_conf.name\n",
"SAMPLE_RATE = omega_conf.sample_rate\n",
"print(f\"Training {MODEL_NAME} with sample rate {SAMPLE_RATE}\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "W7F--_0maLh5"
},
"source": [
"We provide pretrained model checkpoints for fine-tuning.\n",
"\n",
"A list of models available on NGC can be found [here](https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/tts/models/audio_codec.py#L645).\n",
"\n",
"A list of models available on Hugging Face can be found [here](https://huggingface.co/collections/nvidia/nemo-audio-codecs-674f57ab6cb1324f997b5d5b). To use a checkpoint from hugging face, add \"nvidia/\" before the model name."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "cADIAIDUcGWd"
},
"outputs": [],
"source": [
"from nemo.collections.tts.models.audio_codec import AudioCodecModel\n",
"\n",
"pretrained_model_name = \"nvidia/audio-codec-22khz\"\n",
"\n",
"if pretrained_model_name is None:\n",
" MODEL_CHECKPOINT_PATH = None\n",
"else:\n",
" MODEL_CHECKPOINT_PATH = AudioCodecModel.from_pretrained(model_name=pretrained_model_name, return_model_file=True)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "fM4QPsLTnzK7"
},
"source": [
"# Dataset Preparation"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "tkZC6Dl7KRl6"
},
"source": [
"For our tutorial, we use a subset of [VCTK](https://datashare.ed.ac.uk/handle/10283/2950) dataset with 5 speakers (p225-p229)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "sYzvAYr2vo1K"
},
"outputs": [],
"source": [
"import tarfile\n",
"import wget\n",
"\n",
"from nemo.collections.asr.parts.utils.manifest_utils import read_manifest, write_manifest"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "aoxN1QsUzX-k"
},
"outputs": [],
"source": [
"# Create dataset directory\n",
"DATA_DIR = ROOT_DIR / \"data\"\n",
"\n",
"DATA_DIR.mkdir(parents=True, exist_ok=True)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "mArlQd5Hk36b"
},
"outputs": [],
"source": [
"# Download the dataset\n",
"dataset_url = \"https://vctk-subset.s3.amazonaws.com/vctk_subset_multispeaker.tar.gz\"\n",
"dataset_tar_filepath = DATA_DIR / \"vctk.tar.gz\"\n",
"\n",
"if not dataset_tar_filepath.exists():\n",
" wget.download(dataset_url, out=str(dataset_tar_filepath))"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "p987cjtOy9C7"
},
"outputs": [],
"source": [
"# Extract the dataset\n",
"with tarfile.open(dataset_tar_filepath) as tar_f:\n",
" tar_f.extractall(DATA_DIR)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Ko6dxYJW0i3G"
},
"outputs": [],
"source": [
"DATASET_DIR = DATA_DIR / \"vctk_subset_multispeaker\""
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "We5FHYQt5BeO"
},
"outputs": [],
"source": [
"# Visualize the raw dataset\n",
"train_raw_filepath = DATASET_DIR / \"train.json\"\n",
"!head $train_raw_filepath"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "i3jsk2HCMSU5"
},
"source": [
"## Manifest Processing"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "N8WuAGJsMHRn"
},
"source": [
"The downloaded manifest is formatted for TTS training, which contains metadata such as text and speaker.\n",
"\n",
"For codec training we need `audio_filepath`. The `audio_filepath` field can either be an *absolute path*, or a *relative path* with the root directory provided as an argument to each script. Here we use relative paths.\n",
"\n",
"If you include `duration` the training script will automatically calculate the total size of each dataset used, and can be useful for filtering based on utterance length."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "zoCRrKQ20VZP"
},
"outputs": [],
"source": [
"def update_manifest(data_type):\n",
" input_filepath = DATASET_DIR / f\"{data_type}.json\"\n",
" output_filepath = DATASET_DIR / f\"{data_type}_raw.json\"\n",
"\n",
" entries = read_manifest(input_filepath)\n",
" new_entries = []\n",
" for entry in entries:\n",
" # Provide relative path instead of absolute path\n",
" audio_filepath = entry[\"audio_filepath\"].replace(\"audio/\", \"\")\n",
" duration = round(entry[\"duration\"], 2)\n",
" new_entry = {\n",
" \"audio_filepath\": audio_filepath,\n",
" \"duration\": duration\n",
" }\n",
" new_entries.append(new_entry)\n",
"\n",
" write_manifest(output_path=output_filepath, target_manifest=new_entries, ensure_ascii=False)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "PaCc3GCG1UbH"
},
"outputs": [],
"source": [
"update_manifest(\"dev\")\n",
"update_manifest(\"train\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "bVLIB3Ip1Aqn"
},
"outputs": [],
"source": [
"# Visualize updated 'audio_filepath' field.\n",
"train_filepath = DATASET_DIR / \"train_raw.json\"\n",
"!head $train_filepath"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "alrRDWio41qi"
},
"source": [
"## Audio Preprocessing"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "4WfEaMwpUsFt"
},
"source": [
"Next we process the audio data using [preprocess_audio.py](https://github.com/NVIDIA/NeMo/blob/main/scripts/dataset_processing/tts/preprocess_audio.py).\n",
"\n",
"During this step we can apply the following transformations:\n",
"\n",
"1. Resample the audio from 48khz to the target sample rate for codec training.\n",
"2. Remove long silence from the beginning and end of each audio file. This can be done using an *energy* based approach which will work on clean audio, or using *voice activity detection (VAD)* which is slower but also works on audio with background or static noise (eg. from a microphone). Here we suggest VAD because some audio in VCTK has background noise."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "WEvIefjnd7AG"
},
"outputs": [],
"source": [
"import IPython.display as ipd"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "-qEuCH8S4vFP"
},
"outputs": [],
"source": [
"# Python wrapper to invoke the given bash script with the given input args\n",
"def run_script(script, args):\n",
" args = ' \\\\'.join(args)\n",
" cmd = f\"python {script} \\\\{args}\"\n",
"\n",
" print(cmd.replace(\" \\\\\", \"\\n\"))\n",
" print()\n",
" !$cmd"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "0kQ1UDnGfdX6"
},
"outputs": [],
"source": [
"audio_preprocessing_script = NEMO_SCRIPT_DIR / \"preprocess_audio.py\"\n",
"\n",
"# Directory with raw audio data\n",
"input_audio_dir = DATASET_DIR / \"audio\"\n",
"# Directory to write preprocessed audio to\n",
"output_audio_dir = DATASET_DIR / \"audio_preprocessed\"\n",
"# Whether to overwrite existing audio, if it exists in the output directory\n",
"overwrite_audio = True\n",
"# Whether to overwrite output manifest, if it exists\n",
"overwrite_manifest = True\n",
"# Number of threads to parallelize audio processing across\n",
"num_workers = 4\n",
"# Format of output audio files. Use \"flac\" to compress to a smaller file size.\n",
"output_format = \"flac\"\n",
"# Method for silence trimming. Can use \"energy.yaml\" or \"vad.yaml\".\n",
"trim_config_path = NEMO_CONFIG_DIR / \"trim\" / \"vad.yaml\"\n",
"\n",
"def preprocess_audio(data_type):\n",
" input_filepath = DATASET_DIR / f\"{data_type}_raw.json\"\n",
" output_filepath = DATASET_DIR / f\"{data_type}_manifest.json\"\n",
"\n",
" args = [\n",
" f\"--input_manifest={input_filepath}\",\n",
" f\"--output_manifest={output_filepath}\",\n",
" f\"--input_audio_dir={input_audio_dir}\",\n",
" f\"--output_audio_dir={output_audio_dir}\",\n",
" f\"--num_workers={num_workers}\",\n",
" f\"--output_sample_rate={SAMPLE_RATE}\",\n",
" f\"--output_format={output_format}\",\n",
" f\"--trim_config_path={trim_config_path}\"\n",
" ]\n",
" if overwrite_manifest:\n",
" args.append(\"--overwrite_manifest\")\n",
" if overwrite_audio:\n",
" args.append(\"--overwrite_audio\")\n",
"\n",
" run_script(audio_preprocessing_script, args)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "ai0zbXSOriuY"
},
"outputs": [],
"source": [
"preprocess_audio(\"dev\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "NUKnidQYfgDo"
},
"outputs": [],
"source": [
"preprocess_audio(\"train\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "x2yhJtsj2lDR"
},
"source": [
"Before we proceed, it is important to verify that the audio processing works as expected. Let's listen to an audio file before and after processing.\n",
"\n",
"Note that the processed audio is shorter because we trimmed the leading and trailing silence."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "AfdHUHAWuF-G"
},
"outputs": [],
"source": [
"!ls $processed_audio_filepath"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "_fM3GwJxkjOA"
},
"outputs": [],
"source": [
"audio_file = \"p228_009.wav\"\n",
"audio_filepath = input_audio_dir / audio_file\n",
"processed_audio_filepath = output_audio_dir / audio_file.replace(\".wav\", \".flac\")\n",
"\n",
"print(\"Original audio.\")\n",
"ipd.display(ipd.Audio(audio_filepath, rate=SAMPLE_RATE))\n",
"\n",
"print(\"Processed audio.\")\n",
"ipd.display(ipd.Audio(processed_audio_filepath, rate=SAMPLE_RATE))"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "oRO842MUyODC"
},
"source": [
"# Audio Codec Training"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "E4wUKYOfH8ax"
},
"source": [
"Here we show how to train an audio codec model from scratch. Instructions and checkpoints for fine-tuning will be provided later.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "pqfl9jAYMJob"
},
"outputs": [],
"source": [
"import os\n",
"import torch\n",
"from omegaconf import OmegaConf"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "jK2rr-Kr6Qg8"
},
"outputs": [],
"source": [
"dataset_name = \"vctk\"\n",
"audio_dir = DATASET_DIR / \"audio_preprocessed\"\n",
"train_manifest_filepath = DATASET_DIR / \"train_manifest.json\"\n",
"dev_manifest_filepath = DATASET_DIR / \"dev_manifest.json\""
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Vr4D-NB-yQx8"
},
"outputs": [],
"source": [
"audio_codec_training_script = NEMO_EXAMPLES_DIR / \"audio_codec.py\"\n",
"\n",
"# The total number of training steps will be (epochs * steps_per_epoch)\n",
"epochs = 10\n",
"steps_per_epoch = 10\n",
"\n",
"# Name of the experiment that will determine where it is saved locally and in TensorBoard and WandB\n",
"run_id = \"test_run\"\n",
"exp_dir = ROOT_DIR / \"exps\"\n",
"codec_exp_output_dir = exp_dir / MODEL_NAME / run_id\n",
"# Directory where predicted audio will be stored periodically throughout training\n",
"codec_log_dir = codec_exp_output_dir / \"logs\"\n",
"# Optionally log visualization of learned codes.\n",
"log_dequantized = True\n",
"# Optionally log predicted audio and other artifacts to WandB\n",
"log_to_wandb = False\n",
"# Optionally log predicted audio and other artifacts to Tensorboard\n",
"log_to_tensorboard = False\n",
"\n",
"if torch.cuda.is_available():\n",
" accelerator=\"gpu\"\n",
" batch_size = 4\n",
" devices = -1\n",
"else:\n",
" import multiprocessing\n",
" accelerator=\"cpu\"\n",
" batch_size = 2\n",
" devices = multiprocessing.cpu_count()\n",
"\n",
"args = [\n",
" f\"--config-path={CONFIG_DIR}\",\n",
" f\"--config-name={CONFIG_FILENAME}\",\n",
" f\"max_epochs={epochs}\",\n",
" f\"weighted_sampling_steps_per_epoch={steps_per_epoch}\",\n",
" f\"batch_size={batch_size}\",\n",
" f\"log_dir={codec_log_dir}\",\n",
" f\"exp_manager.exp_dir={exp_dir}\",\n",
" f\"+exp_manager.version={run_id}\",\n",
" f\"model.log_config.log_wandb={log_to_wandb}\",\n",
" f\"model.log_config.log_tensorboard={log_to_tensorboard}\",\n",
" f\"model.log_config.generators.0.log_dequantized={log_dequantized}\",\n",
" f\"trainer.accelerator={accelerator}\",\n",
" f\"trainer.devices={devices}\",\n",
" f\"+train_ds_meta.{dataset_name}.manifest_path={train_manifest_filepath}\",\n",
" f\"+train_ds_meta.{dataset_name}.audio_dir={audio_dir}\",\n",
" f\"+val_ds_meta.{dataset_name}.manifest_path={dev_manifest_filepath}\",\n",
" f\"+val_ds_meta.{dataset_name}.audio_dir={audio_dir}\",\n",
" f\"+log_ds_meta.{dataset_name}.manifest_path={dev_manifest_filepath}\",\n",
" f\"+log_ds_meta.{dataset_name}.audio_dir={audio_dir}\"\n",
"]\n",
"\n",
"if MODEL_CHECKPOINT_PATH is not None:\n",
" args.append(f\"+init_from_nemo_model={MODEL_CHECKPOINT_PATH}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Bn8lQG0PxWGi"
},
"outputs": [],
"source": [
"# If an error occurs, log the entire stacktrace.\n",
"os.environ[\"HYDRA_FULL_ERROR\"] = \"1\""
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "yUxFCNrE3Ywi"
},
"outputs": [],
"source": [
"# Do the model training. For some configurations this step might hang when using CPU.\n",
"run_script(audio_codec_training_script, args)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "BBPIpS-lL6z9"
},
"source": [
"During training, the model will automatically save predictions for all audio files specified in the `log_ds_meta` manifest."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "rSFOm1Sg46Lh"
},
"outputs": [],
"source": [
"codec_log_epoch_dir = codec_log_dir / \"epoch_10\" / dataset_name\n",
"!ls $codec_log_epoch_dir"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "oCJs7oCLMIjD"
},
"source": [
"This makes it easy to listen to the audio to determine how well the model is performing. We can decide to stop training when either:\n",
"\n",
"* The predicted audio sounds almost identical to the original audio.\n",
"* The predicted audio stops improving in between epochs.\n",
"\n",
"**Note that when training from scratch, the dataset in this tutorial is too small to get good audio quality.**"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "G6k4ymzfJ5Y6"
},
"outputs": [],
"source": [
"audio_filepath_ground_truth = output_audio_dir / \"p228_009.flac\"\n",
"audio_filepath_reconstructed = codec_log_epoch_dir / \"p228_009_audio_out.wav\"\n",
"\n",
"print(\"Ground truth audio.\")\n",
"ipd.display(ipd.Audio(audio_filepath_ground_truth, rate=SAMPLE_RATE))\n",
"\n",
"print(\"Reconstructed audio.\")\n",
"ipd.display(ipd.Audio(audio_filepath_reconstructed, rate=SAMPLE_RATE))\n",
"\n",
"dequantized_filepath = codec_log_epoch_dir / \"p228_009_dequantized.png\"\n",
"ipd.Image(dequantized_filepath)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "rynZYwg2VP5d"
},
"source": [
"# Related Information"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "_LtyHHuLkNDv"
},
"source": [
"To learn more about audio codec models in NeMo, look at our [documentation](https://docs.nvidia.com/nemo-framework/user-guide/latest/nemotoolkit/tts/models.html#codecs).\n",
"\n",
"For more information on how to download and run pre-trained audio codec models, visit [NGC](https://catalog.ngc.nvidia.com/models?filters=&orderBy=scoreDESC&query=codec)."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "LeqV3VvJVOb-"
},
"source": [
"# References"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Rvu4w2x_3RSY"
},
"source": [
"1. [EnCodec](https://arxiv.org/abs/2210.13438)\n",
"2. [Finite Scalar Quantization (FSQ)](https://arxiv.org/abs/2309.15505)\n",
"3. [HiFi-GAN](https://arxiv.org/abs/2010.05646)\n",
"4. [SEANet](https://arxiv.org/abs/2009.02095)\n",
"5. [Spectral Codecs](https://arxiv.org/abs/2406.05298)\n",
"6. [Low Frame-rate Speech Codec](https://arxiv.org/abs/2409.12117)"
]
}
],
"metadata": {
"colab": {
"provenance": []
},
"kernelspec": {
"display_name": "Python 3",
"name": "python3"
},
"language_info": {
"name": "python"
}
},
"nbformat": 4,
"nbformat_minor": 0
}
tutorials/tts/Evaluation_MelCepstralDistortion.ipynb
+631
View File
@@ -0,0 +1,631 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "f0408502",
"metadata": {},
"source": [
"# Model Evaluation: Mel Cepstral Distortion with Dynamic Time Warping"
]
},
{
"cell_type": "markdown",
"id": "8a749b10",
"metadata": {},
"source": [
"In this tutorial, we will learn how to calculate **mel cepstral distortion (MCD)** with **dynamic time warping (DTW)** between a synthesized mel spec and a reference mel spec. MCD DTW can be useful for comparing models trained on the same training data.\n",
"\n",
"MCD is an objective measure of speech quality that is calculated between pairs of TTS-generated mel spectrograms and ground truth mel spectrograms. Two mel spectrograms are similar if the MCD value between them is low; as you might expect, the MCD of a mel spec with itself is 0.\n",
"\n",
"MCD DTW is a modification of MCD that works with non-aligned mels by using a dynamic time warping cost matrix. (Vanilla MCD can only be measured for two mels that are both the same length, and which are assumed to be aligned.) The scale depends on factors such as the mel extractor and reduction algorithm (mean of DTW cost or min DTW path cost)."
]
},
{
"cell_type": "markdown",
"id": "8801b8cb",
"metadata": {},
"source": [
"## License\n",
"\n",
"> Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n",
">\n",
"> Licensed under the Apache License, Version 2.0 (the \"License\");\n",
"> you may not use this file except in compliance with the License.\n",
"> You may obtain a copy of the License at\n",
">\n",
"> http://www.apache.org/licenses/LICENSE-2.0\n",
">\n",
"> Unless required by applicable law or agreed to in writing, software\n",
"> distributed under the License is distributed on an \"AS IS\" BASIS,\n",
"> WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
"> See the License for the specific language governing permissions and\n",
"> limitations under the License."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "09cff1ab",
"metadata": {},
"outputs": [],
"source": [
"\"\"\"\n",
"You can either run this notebook locally (if you have all the dependencies and a GPU) or on Google Colab.\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",
"BRANCH = 'main'\n",
"# If you're using Google Colab and not running locally, uncomment and run this cell.\n",
"# !pip install librosa numpy matplotlib"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "8b9b6e1c",
"metadata": {},
"outputs": [],
"source": [
"## Import libraries\n",
"import IPython.display as ipd\n",
"import librosa\n",
"import numpy as np\n",
"import math\n",
"import matplotlib.pyplot as plt\n",
"import os"
]
},
{
"cell_type": "markdown",
"id": "398b01e0",
"metadata": {},
"source": [
"## Defining Parameters for Spectrogram Generation\n",
"\n",
"In this section we define parameters required to generate our mel spectrograms. More information about these parameters can found in the Librosa documentation for [stft](https://librosa.org/doc/main/generated/librosa.stft.html), [mels](https://librosa.org/doc/main/generated/librosa.filters.mel.html) and [mfcc](https://librosa.org/doc/main/generated/librosa.feature.mfcc.html) generation."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "191012de",
"metadata": {},
"outputs": [],
"source": [
"## Mel spec params\n",
"n_fft=1024\n",
"hop_length=256\n",
"win_length=None\n",
"window='hann'\n",
"n_mels = 80\n",
"sr = 22050\n",
"\n",
"## Mfcc params\n",
"n_mfcc=34"
]
},
{
"cell_type": "markdown",
"id": "54830c8c",
"metadata": {},
"source": [
"A little bit about these parameters: \n",
" - `n_fft`: Number of fft_components or length of windowed signal after padding in STFT. \n",
" - `hop_length`: Number of audio samples between adjacent STFT columns. \n",
" - `window`: Window to use in STFT. \n",
" - `win_length`: Length of the window to be used. \n",
" - `n_mels`: Number of number of mel bands to generate. \n",
" - `sr`: Sample rate of the samples. \n",
" - `n_mfcc`: Number of MFCCs to generate."
]
},
{
"cell_type": "markdown",
"id": "4cb9e37c",
"metadata": {},
"source": [
"## Loading and Visualizing Data\n",
"\n",
"Lets apply the algorithm on a synthesized mel and ground truth mel pair for understanding.\n",
"\n",
"First, we need a function to generate mel spectrograms from audio files. Mel spectrograms are generated using the [librosa mel extractor](https://librosa.org/doc/main/generated/librosa.filters.mel.html)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "042ceb7f",
"metadata": {},
"outputs": [],
"source": [
"def wav2mel(filename):\n",
" \"\"\"\n",
" Function to load an audio file and generate/return mel specs:\n",
" Args:\n",
" filename: Full path of the audio file.\n",
" Returns:\n",
" mels: Corresponding mel spectrogram.\n",
" \"\"\"\n",
" wav_, _ = librosa.load(filename, sr=sr) # load() returns an (audio data, sample rate) tuple\n",
" mels = librosa.feature.melspectrogram(\n",
" y=wav_,\n",
" sr=sr,\n",
" n_fft=n_fft,\n",
" hop_length=hop_length,\n",
" win_length=win_length,\n",
" window=window,\n",
" n_mels=n_mels\n",
" )\n",
" mels = librosa.power_to_db(mels, ref=np.max)\n",
" return mels"
]
},
{
"cell_type": "markdown",
"id": "8546edb5",
"metadata": {},
"source": [
"We also need a function to convert mels to MFCC, we will use Librosa's [mfcc](https://librosa.org/doc/main/generated/librosa.feature.mfcc.html) generation:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "c1a0663d",
"metadata": {},
"outputs": [],
"source": [
"def mel2mfcc(mels):\n",
" mfcc = librosa.feature.mfcc(S=mels, n_mfcc=n_mfcc)\n",
" return mfcc"
]
},
{
"cell_type": "markdown",
"id": "6efac554",
"metadata": {},
"source": [
"For this tutorial, we have already generated mels from trained [FastPitch](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/nemo/models/tts_en_fastpitch) and [RAD-TTS](https://github.com/NVIDIA/radtts) models, and we have the ground truth sample that they correspond to. Let's first download the tarball with the data and expand it."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "156290a9",
"metadata": {},
"outputs": [],
"source": [
"# Untar the example files.\n",
"!wget https://tts-tutorial-data.s3.us-east-2.amazonaws.com/MCD_DTW_examples.tar\n",
"!tar -xvf MCD_DTW_examples.tar"
]
},
{
"cell_type": "markdown",
"id": "7aeabe75",
"metadata": {},
"source": [
"Now we can generate a mel spectrogram for the ground truth audio, load mel specs for generated audio and generate MFCC(Mel frequency cepstral coefficient) for both."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "d1a71f88",
"metadata": {
"scrolled": false
},
"outputs": [],
"source": [
"## Generate spectrograms\n",
"gt_mels = wav2mel(\"MCD_DTW/gt/sample_0.wav\")\n",
"synt_mels = np.load(\"MCD_DTW/fastpitch/mels/mels_0.npy\")\n",
"\n",
"## Generate MFCCs\n",
"gt_mfcc = mel2mfcc(gt_mels)\n",
"synt_mfcc = mel2mfcc(synt_mels)"
]
},
{
"cell_type": "markdown",
"id": "faa5df64",
"metadata": {},
"source": [
"### Visualizing the Audio\n",
"\n",
"Let's first listen to the ground truth and synthesized audio."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "08f08b89",
"metadata": {},
"outputs": [],
"source": [
"print(\"Ground truth audio:\")\n",
"audio_gt, sr_gt = librosa.load(\"MCD_DTW/gt/sample_0.wav\")\n",
"ipd.display(ipd.Audio(audio_gt, rate=sr_gt))\n",
"\n",
"print(\"Synthesized audio:\")\n",
"audio_fp, sr_fp = librosa.load(\"MCD_DTW/fastpitch/audio/sample_0.wav\")\n",
"ipd.display(ipd.Audio(audio_fp, rate=sr_fp))"
]
},
{
"cell_type": "markdown",
"id": "21b5c780",
"metadata": {},
"source": [
"Now, let's take a look at the mel spectrograms."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "550d9b12",
"metadata": {},
"outputs": [],
"source": [
"# Visualize mel spectrograms\n",
"fig, ax = plt.subplots(2, 1, figsize=(12,10))\n",
"\n",
"_ = ax[0].pcolormesh(gt_mels, cmap='viridis')\n",
"_ = ax[0].set_title(\"Ground Truth Mel Spectrogram\")\n",
"\n",
"_ = ax[1].pcolormesh(synt_mels, cmap='viridis')\n",
"_ = ax[1].set_title(\"FastPitch Synthesized Mel Spectrogram\")"
]
},
{
"cell_type": "markdown",
"id": "93c4c301",
"metadata": {},
"source": [
"## Calculating the DTW matrix\n",
"\n",
"Dynamic time warping finds the most optimal path to align two sequences of different length, and in doing so measures the similarity between the two time series that are not in sync.\n",
"\n",
"DTW uses dynamic programming to **calculate the cost of every alignment path and chooses the path with least accumulated cost**. The value of accumulated cost matrix $D$ at index $(x_a, y_b)$ is the minimum distance between the points, where $x$ and $y$ are two time series. Formally, the cost matrix can be defined as:\n",
" \n",
"$D(a,b) = min(D(a-1, b), D(a, b-1), D(a-1, b-1)) + c(x_{a}, y_{b})$ \n",
"$D(1, b) = \\sum(c(1, y_{b}))$ \n",
"$D(a, 1) = \\sum(c(x_{a}, 1))$ \n",
" \n",
"Where ***x*** and ***y*** are the audio signals and ***c*** is the log cost function we have defined.\n",
"\n",
"We will use [the DTW function from Librosa](https://librosa.org/doc/main/generated/librosa.sequence.dtw.html) on MFCC. It returns the DTW accumulated cost matrix and DTW optimum path.\n",
"\n",
"In the following function, we define the cost function for DTW to pass in."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "8a4adae2",
"metadata": {},
"outputs": [],
"source": [
"## Define the cost function for calculating DTW\n",
"def log_spec_dB_dist(x, y):\n",
" log_spec_dB_const = 10.0 / math.log(10.0) * math.sqrt(2.0)\n",
" diff = x - y\n",
" return log_spec_dB_const * math.sqrt(np.inner(diff, diff))"
]
},
{
"cell_type": "markdown",
"id": "ca8c6974",
"metadata": {},
"source": [
"Now we can run Librosa's `dtw()`:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "0c4c5338",
"metadata": {},
"outputs": [],
"source": [
"# dtw_cost: Shape (N, M), where N and M are the lengths of gt_mfcc, synt_mfcc respectively. Cost matrix.\n",
"# dtw_min_path: Shape (N, 2). Pairs of coordinates with of the min-cost path from bottom-right to top-left.\n",
"dtw_cost, dtw_min_path = librosa.sequence.dtw(gt_mfcc, synt_mfcc, metric=log_spec_dB_dist)"
]
},
{
"cell_type": "markdown",
"id": "913b6548",
"metadata": {},
"source": [
"Reduction of the DTW matrix can be done by either taking the mean of the entire cost matrix or averaging the DTW cost for the minimum cost path per frame. We will use the DTW cost along the min cost path."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "6c85ab38",
"metadata": {},
"outputs": [],
"source": [
"# First, sum up the costs over the path\n",
"path_cost_matrix = dtw_cost[dtw_min_path[:, 0], dtw_min_path[:, 1]]\n",
"path_cost = np.sum(path_cost_matrix)\n",
"\n",
"# Average over path length\n",
"path_length = dtw_min_path.shape[0]\n",
"reduced_dtw_cost = path_cost/path_length\n",
"\n",
"# Average over number of frames\n",
"frames = synt_mels.shape[1]\n",
"mcd = reduced_dtw_cost/frames\n",
"\n",
"print(f\"MCD_DTW is: {mcd}\")"
]
},
{
"cell_type": "markdown",
"id": "24801d61",
"metadata": {},
"source": [
"## MCD DTW Usecase\n",
"\n",
"MCD is a very useful metric to **compare the convergence of two models**. Therefore in this section, we will calculate the average MCD for audio files generated by two different models: FastPitch and RAD-TTS.\n",
"\n",
"But first, let's put the calculations we have just performed in functions for better readability and reusability.\n",
"\n",
"Here are the helper functions for getting the average cost of DTW along a path."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e4181d60",
"metadata": {},
"outputs": [],
"source": [
"def extract_path_cost(D, wp):\n",
" \"\"\"\n",
" Get the path cost from D(cost matrix), wp (warped path)\n",
" Returns the sum of the path cost \n",
" \"\"\"\n",
" path_cost = D[wp[:, 0], wp[:, 1]]\n",
" return np.sum(path_cost)\n",
"\n",
"def extract_frame_avg_path_cost(D, wp):\n",
" \"\"\"\n",
" Get the average path cost over the length of the given path\n",
" \"\"\"\n",
" path_cost = extract_path_cost(D, wp)\n",
" path_length = wp.shape[0]\n",
" frame_avg_path_cost = path_cost / float(path_length)\n",
" return frame_avg_path_cost"
]
},
{
"cell_type": "markdown",
"id": "08ca0f47",
"metadata": {},
"source": [
"And a function for calculating MCD for a synthetic mel spectrogram:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "99800a97",
"metadata": {},
"outputs": [],
"source": [
"def cal_mcd(gt_mfcc, synt_mfcc, cost_function, dtw_type='path_cost'):\n",
" \"\"\"\n",
" Calculates MCD between a ground truth and synthetic mel.\n",
" \"\"\"\n",
" frames = synt_mfcc.shape[1]\n",
" path_cost = 0\n",
" \n",
" # dynamic time warping for MCD\n",
" dtw_cost, dtw_min_path = librosa.sequence.dtw(gt_mfcc, synt_mfcc, metric=cost_function)\n",
" if dtw_type == 'mean':\n",
" path_cost = np.mean(dtw_cost)\n",
" else:\n",
" path_cost = extract_frame_avg_path_cost(dtw_cost, dtw_min_path)\n",
" \n",
" mcd = path_cost / frames\n",
" \n",
" return mcd, frames"
]
},
{
"cell_type": "markdown",
"id": "f1489ecf",
"metadata": {},
"source": [
"This last function will calculate the MCDs for all the synthetic mels in a directory.\n",
"\n",
"We'll get a better comparison if we compute the MCDs of (much) more than one synthesized spectrogram per model! This function streamlines the process by letting us pass in (1) a directory with all the synthesized mel spectrograms from a TTS model, and (2) a directory with all the ground-truth audio.\n",
"\n",
"The function assumes that, after sorting the contents of each directory, wav *#i* in the ground truth directory will correspond to mel *#i* in the synthesized mel directory. In our case, the directories for ground truth and synthesized FastPitch mels look like this:\n",
"\n",
"```bash\n",
"% ls gt/\n",
"sample_0.wav sample_1.wav sample_2.wav sample_3.wav sample_4.wav\n",
"% ls fastpitch/mels/\n",
"mels_0.npy mels_1.npy mels_2.npy mels_3.npy mels_4.npy\n",
"```\n",
"\n",
"We only have ten files each to compare as a toy example, for a full evaluation you will likely want more!"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "07bb51df",
"metadata": {},
"outputs": [],
"source": [
"def cal_mcd_dir(synt_dir, gt_dir):\n",
" \"\"\"\n",
" Calculate MCDs for pairs of synthetic and ground truth audio files.\n",
" \n",
" Args:\n",
" synt_dir: Path to the directory that contains all the synthetic mels.\n",
" gt_dir: Path to the directory that contains all the ground truth wavs.\n",
" These should correspond 1:1 with the mels in synt_dir.\n",
" \n",
" Returns:\n",
" List of MCDs (where length is the number of files in each directory)\n",
" \"\"\"\n",
" mcds = []\n",
" \n",
" synt_filenames = os.listdir(synt_dir)\n",
" synt_filepaths = [os.path.join(synt_dir, filename) for filename in synt_filenames]\n",
" synt_filepaths.sort()\n",
" gt_filenames = os.listdir(gt_dir)\n",
" gt_filepaths = [os.path.join(gt_dir, filename) for filename in gt_filenames]\n",
" gt_filepaths.sort()\n",
"\n",
" for synt_melname, gt_audio in zip(synt_filepaths, gt_filepaths):\n",
" synt_mels = np.load(synt_melname)\n",
" synt_mfcc = mel2mfcc(synt_mels)\n",
" \n",
" gt_mels = wav2mel(gt_audio)\n",
" gt_mfcc = mel2mfcc(gt_mels)\n",
" \n",
" mcd, _ = cal_mcd(gt_mfcc, synt_mfcc, log_spec_dB_dist)\n",
" mcds.append(mcd)\n",
" return mcds"
]
},
{
"cell_type": "markdown",
"id": "fc295283",
"metadata": {
"scrolled": true
},
"source": [
"### Calculate MCD DTW on Synthesized Files from Each Model\n",
"\n",
"We can now calculate the MCD DTW for the synthesized FastPitch mels (compared to the ground truth) as well as for the synthesized RAD-TTS mels (ditto), then compare them. This will take a few seconds."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "2bc8da46",
"metadata": {
"scrolled": false
},
"outputs": [],
"source": [
"%%capture --no-display\n",
"mels_dir_m1 = \"MCD_DTW/fastpitch/mels/\"\n",
"mels_dir_m2 = \"MCD_DTW/radtts/mels/\"\n",
"mels_dir_gt = \"MCD_DTW/gt/\"\n",
"\n",
"mcds_m1 = cal_mcd_dir(mels_dir_m1, mels_dir_gt) # FastPitch\n",
"mcds_m2 = cal_mcd_dir(mels_dir_m2, mels_dir_gt) # RAD-TTS"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "6bd1ef33",
"metadata": {},
"outputs": [],
"source": [
"print(f\"Average MCD for Model 1 (FastPitch) is: {sum(mcds_m1)/len(mcds_m1):.2f}\")\n",
"print(f\"Average MCD for Model 2 (RAD-TTS) is: {sum(mcds_m2)/len(mcds_m2):.2f}\")"
]
},
{
"cell_type": "markdown",
"id": "f9ef5363",
"metadata": {},
"source": [
"We're measuring divergence from the ground truth for each of these, so **lower is better**!"
]
},
{
"cell_type": "markdown",
"id": "71e998a9",
"metadata": {},
"source": [
"### Plotting MCD"
]
},
{
"cell_type": "markdown",
"id": "7f81c5b5",
"metadata": {},
"source": [
"We can also plot the MCD DTW values for both models, with the model with a lower MCD DTW value being closer to ground truth."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "3f8d71bf",
"metadata": {},
"outputs": [],
"source": [
"x_axis = np.linspace(1, len(mcds_m1), len(mcds_m1)) ## Define an x axis for for plotting in matplotlib\n",
"plt.plot(x_axis, mcds_m1, label=\"FastPitch\")\n",
"plt.plot(x_axis, mcds_m2, label=\"RAD-TTS\")\n",
"\n",
"plt.title(\"MCD DTW value for each file\")\n",
"plt.ylabel(\"MCD_DTW value\")\n",
"\n",
"plt.grid()\n",
"plt.legend()\n",
"plt.show()"
]
},
{
"cell_type": "markdown",
"id": "fd4eb0a6",
"metadata": {},
"source": [
"From the graph above, we can see that in general the value of MCD is greater for RAD-TTS mel spectrograms than for the FastPitch ones. This is also reflected in the average MCD value for both models that we computed before.\n",
"\n",
"Therefore, we can conclude that FastPitch has better convergence than RAD-TTS. However, we cannot evaluate the quality of audio generated by these models using MCD alone! **MCD is a great tool for testing model convergence, but generated audio may have pronunciation and quality artifacts.** Therefore MCD evaluation should be followed by a MOS (Mean Opinion Score) and CMOS (Comparative Mean Opinion Score) evaluation."
]
},
{
"cell_type": "markdown",
"id": "d66a52d5",
"metadata": {},
"source": [
"## Additional NeMo Resources\n",
"\n",
"If you are unsure where to begin for training a TTS model, you may want to start with the [FastPitch and Mixer-TTS Training notebook](https://github.com/NVIDIA/NeMo/blob/main/tutorials/tts/FastPitch_MixerTTS_Training.ipynb) or the [NeMo TTS Primer notebook](https://github.com/NVIDIA/NeMo/blob/main/tutorials/tts/NeMo_TTS_Primer.ipynb). For fine-tuning, there is also the [FastPitch Fine-Tuning notebook](https://github.com/NVIDIA/NeMo/blob/main/tutorials/tts/FastPitch_Finetuning.ipynb).\n",
"\n",
"For some guidance on how to load a trained model and perform inference to generate mels or waveforms, check out how it's done in the [Inference notebook](https://github.com/NVIDIA/NeMo/blob/main/tutorials/tts/Inference_ModelSelect.ipynb). Important functions to know are include `from_pretrained()` (if loading from an NGC model) and `restore_from()` (if loading a `.nemo` file). See the [NeMo Primer notebook](https://github.com/NVIDIA/NeMo/blob/stable/tutorials/00_NeMo_Primer.ipynb) for more general information about model training, saving, and loading."
]
}
],
"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.9.15"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
tutorials/tts/NeMo_TTS_Primer.ipynb
+2069
View File
File diff suppressed because it is too large Load Diff
tutorials/tts/Pronunciation_customization.ipynb
+306
View File
@@ -0,0 +1,306 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"<img src=\"http://developer.download.nvidia.com/notebooks/dlsw-notebooks/riva_tts_tts-python-basics/nvidia_logo.png\" style=\"width: 90px; float: right;\">\n",
"\n",
"# How do I customize TTS pronunciations?\n",
"\n",
"This tutorial walks you through the basics of NeMo TTS pronunciation customization. "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"\"\"\"\n",
"You can either run this notebook locally (if you have all the dependencies and a GPU) or on Google Colab.\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",
"\n",
"BRANCH = 'main'\n",
"# # If you're using Google Colab and not running locally, uncomment and run this cell.\n",
"# !apt-get install sox libsndfile1 ffmpeg\n",
"# !pip install wget text-unidecode \n",
"# !python -m pip install \"nemo_toolkit[all] @ git+https://github.com/NVIDIA-NeMo/Speech.git@$BRANCH\"\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Grapheme-to-phoneme (G2P) Overview\n",
"\n",
"Modern **text-to-speech** (TTS) models can learn pronunciations from raw text input and its corresponding audio data.\n",
"Sometimes, however, it is desirable to customize pronunciations, for example, for domain-specific terms. As a result, many TTS systems use grapheme and phonetic input during training to directly access and correct pronunciations at inference time.\n",
"\n",
"\n",
"[The International Phonetic Alphabet (IPA)](https://en.wikipedia.org/wiki/International_Phonetic_Alphabet) and [ARPABET](https://en.wikipedia.org/wiki/ARPABET) are the most common phonetic alphabets. \n",
"\n",
"There are two ways to customize pronunciations:\n",
"\n",
"1. pass phonemes as an input to the TTS model, note that the request-time overrides are best suited for one-off adjustments\n",
"2. configure TTS model with the desired domain-specific terms using custom phonetic dictionary\n",
"\n",
"Both methods require users to convert graphemes into phonemes (G2P). \n",
"\n",
"#### All words for G2P purposes could be divided into the following groups:\n",
"* *known* words - words that are present in the model's phonetic dictionary\n",
"* *out-of-vocabulary (OOV)* words - words that are missing from the model's phonetic dictionary. \n",
"* *[heteronyms](https://en.wikipedia.org/wiki/Heteronym_&#40;linguistics&#41;)* - words with the same spelling but different pronunciations and/or meanings, e.g., *bass* (the fish) and *bass* (the musical instrument).\n",
"\n",
"#### Important NeMo flags:\n",
"* `your_spec_generator_model.vocab.g2p.phoneme_dict` - phoneme dictionary that maps words to their phonetic transcriptions, e.g., [ARPABET-based CMU Dictionary](https://raw.githubusercontent.com/NVIDIA/NeMo/stable/scripts/tts_dataset_files/cmudict-0.7b_nv22.10) or [IPA-based CMU Dictionary](https://github.com/NVIDIA/NeMo/blob/stable/scripts/tts_dataset_files/ipa_cmudict-0.7b_nv23.01.txt)\n",
"* `your_spec_generator_model.vocab.g2p.heteronyms` - list of the model's heteronyms, grapheme form of these words will be used even if the word is present in the phoneme dictionary.\n",
"* `your_spec_generator_model.vocab.g2p.ignore_ambiguous_words`: if is set to **True**, words with more than one phonetic representation in the pronunciation dictionary are ignored. This flag is relevant to the words with multiple valid phonetic transcriptions in the dictionary that are not in `your_spec_generator_model.vocab.g2p.heteronyms` list.\n",
"* `your_spec_generator_model.vocab.phoneme_probability` - phoneme probability flag in the Tokenizer and the same from in the G2P module: `your_spec_generator_model.vocab.g2p.phoneme_probability` ([0, 1]). If a word is present in the phoneme dictionary, we still want our TTS model to see graphemes and phonemes during training to handle OOV words during inference. The `phoneme_probability` determines the probability of an unambiguous dictionary word appearing in phonetic form during model training, `(1 - phoneme_probability)` is the probability of the graphemes. This flag is set to `1` in the parse() method during inference.\n",
"\n",
"To ensure the desired pronunciation, we need to add a new entry to the model's phonetic dictionary. If the target word is already in the dictionary, we need to remove the default pronunciation so that only the target pronunciation is present. "
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Default G2P\n",
"\n",
"Below we show how to analyze default G2P output of the NeMo models"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"import nemo.collections.tts as nemo_tts\n",
"import IPython.display as ipd\n",
"\n",
"# Load mel spectrogram generator\n",
"spec_generator = nemo_tts.models.FastPitchModel.from_pretrained(\"tts_en_fastpitch_ipa\").eval()\n",
"# Load vocoder\n",
"vocoder = nemo_tts.models.HifiGanModel.from_pretrained(model_name=\"tts_en_hifigan\").eval()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# helper functions\n",
"\n",
"def generate_audio(input_text):\n",
" # parse() sets phoneme probability to 1, i.e. dictionary phoneme transcriptions are used for known words\n",
" parsed = spec_generator.parse(input_text)\n",
" spectrogram = spec_generator.generate_spectrogram(tokens=parsed)\n",
" audio = vocoder.convert_spectrogram_to_audio(spec=spectrogram)\n",
" display(ipd.Audio(audio.detach().to('cpu').numpy(), rate=22050))\n",
" \n",
"def display_postprocessed_text(text):\n",
" # to use dictionary entries for known words, not needed for generate_audio() as parse() handles this\n",
" spec_generator.vocab.phoneme_probability = 1\n",
" spec_generator.vocab.g2p.phoneme_probability = 1\n",
" print(f\"Input before tokenization: {' '.join(spec_generator.vocab.g2p(text))}\\n\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"text = \"paracetamol can help reduce fever.\"\n",
"generate_audio(text)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### Expected results if you run the tutorial:\n",
"<audio controls src=\"https://github.com/NVIDIA/NeMo/raw/main/tutorials/tts/audio_samples/default_ipa.wav\" type=\"audio/ogg\"></audio> \n",
"\n",
"\n",
"During preprocessing, unambiguous dictionary words are converted to phonemes, while OOV and words with multiple entries are kept as graphemes. For example, **paracetamol** is missing from the phoneme dictionary, and **can** has 2 forms."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"display_postprocessed_text(text)\n",
"for word in [\"paracetamol\", \"can\"]:\n",
" word = word.upper()\n",
" phoneme_forms = spec_generator.vocab.g2p.phoneme_dict[word]\n",
" print(f\"Number of phoneme forms for wordPhoneme forms for '{word}': {len(phoneme_forms)} -- {phoneme_forms}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Input customization\n",
"\n",
"One can pass phonemes directly as input to the model to customize pronunciation.\n",
"\n",
"Let's replace the word **paracetamol** with the desired phonemic transcription in our example by adding vertical bars around each phone, e.g., `ˌpæɹəˈsitəmɔl` -> `|ˌ||p||æ||ɹ||ə||ˈ||s||i||t||ə||m||ɔ||l|`."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"print(f\"Original text: {text}\")\n",
"\n",
"new_pronunciation = \"ˌpæɹəˈsitəmɔl\"\n",
"phoneme_form = \"\".join([f\"|{s}|\" for s in new_pronunciation])\n",
"text_with_phonemes = text.replace(\"paracetamol\", phoneme_form)\n",
"print(f\"Text with phonemes: {text_with_phonemes}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"generate_audio(text_with_phonemes)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### Expected results if you run the tutorial:\n",
"<audio controls src=\"https://github.com/NVIDIA/NeMo/raw/main/tutorials/tts/audio_samples/phonemes_as_input.wav\" type=\"audio/ogg\"></audio> \n",
"\n",
"\n",
"## Dictionary customization\n",
"\n",
"Below we show how to customize phonetic dictionary for NeMo TTS models. "
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Let's add a new entry to the dictionary for the word **paracetamol**. "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# we download IPA-based CMU Dictionary and add a custom entry for the target word\n",
"ipa_cmu_dict = \"ipa_cmudict-0.7b_nv23.01.txt\"\n",
"if os.path.exists(ipa_cmu_dict):\n",
" ! rm $ipa_cmu_dict\n",
"\n",
"! wget https://raw.githubusercontent.com/NVIDIA/NeMo/main/scripts/tts_dataset_files/$ipa_cmu_dict\n",
"\n",
"with open(ipa_cmu_dict, \"a\") as f:\n",
" f.write(f\"PARACETAMOL {new_pronunciation}\\n\")\n",
" \n",
"! tail $ipa_cmu_dict"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# let's now use our updated dictionary as the model's phonetic dictionary\n",
"spec_generator.vocab.g2p.replace_dict(ipa_cmu_dict)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"**Paracetamol** is no longer an OOV, and the model uses the phonetic form we provided:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"display_postprocessed_text(text)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Finally, let's use the new phoneme dictionary for synthesis."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"generate_audio(text)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### Expected results if you run the tutorial:\n",
"<audio controls src=\"https://github.com/NVIDIA/NeMo/raw/main/tutorials/tts/audio_samples/new_dict_entry.wav\" type=\"audio/ogg\"></audio> "
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Resources\n",
"* [TTS pipeline customization](https://docs.nvidia.com/deeplearning/riva/user-guide/docs/tts/tts-custom.html#tts-pipeline-configuration)\n",
"* [Overview of TTS in NeMo](https://github.com/NVIDIA/NeMo/blob/main/tutorials/tts/NeMo_TTS_Primer.ipynb)\n",
"* [G2P models in NeMo](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/tts/g2p.html)\n",
"* [Riva TTS documentation](https://docs.nvidia.com/deeplearning/riva/user-guide/docs/tts/tts-overview.html)"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "38",
"language": "python",
"name": "38"
},
"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.8.13"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
tutorials/tts/audio_samples/default_ipa.wav
BIN
View File
Binary file not shown.
tutorials/tts/audio_samples/new_dict_entry.wav
BIN
View File
Binary file not shown.
tutorials/tts/audio_samples/phonemes_as_input.wav
BIN
View File
Binary file not shown.
tutorials/tts/images/fastpitch-pitch.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 82 KiB

tutorials/tts/images/fastpitch_diagram.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 21 KiB

tutorials/tts/images/hifi_gan_discriminator_diagram.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 23 KiB

tutorials/tts/images/sampling_rate.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 3.1 KiB

tutorials/tts/images/tacotron2_diagram.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 197 KiB

tutorials/tts/images/tts_pipeline.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 75 KiB

tutorials/tts/images/wavenet_diagram.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 19 KiB