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

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

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:28:58 +08:00
commit ba4be087d5
2316 changed files with 2668701 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tutorials/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()