asr-finetune-conformer-ctc-nemo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4NpvEOMs_mKp"
   },
   "source": [
    "<img src=\"http://developer.download.nvidia.com/notebooks/dlsw-notebooks/rivaasrasr-finetuning-conformer-ctc-nemo/nvidia_logo.png\" style=\"width: 90px; float: right;\">\n",
    "\n",
    "# How to Fine-Tune a Riva ASR Acoustic Model with NVIDIA NeMo\n",
    "This tutorial walks you through how to fine-tune an NVIDIA Riva ASR acoustic model with NVIDIA NeMo.\n",
    "\n",
    "**Important**: If you plan to fine-tune an ASR acoustic model using the same tokenizer with which the model was trained, skip this tutorial and refer to the \"Sub-word Encoding CTC Model\" section (starting with the \"Load pre-trained model\" subsection) of the [NeMo ASR Language Finetuning tutorial](https://github.com/NVIDIA/NeMo/blob/main/tutorials/asr/ASR_CTC_Language_Finetuning.ipynb)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6h7SXF6G_mKw"
   },
   "source": [
    "## NVIDIA Riva Overview\n",
    "\n",
    "NVIDIA Riva is a GPU-accelerated SDK for building speech AI applications that are customized for your use case and deliver real-time performance. <br/>\n",
    "Riva offers a rich set of speech and natural language understanding (NLU) services such as:\n",
    "\n",
    "- Automated speech recognition (ASR). \n",
    "- Text-to-Speech synthesis (TTS). \n",
    "- A collection of natural language processing (NLP) services, such as named entity recognition (NER), punctuation, and intent classification.\n",
    "\n",
    "In this tutorial, we will fine-tune a Riva ASR acoustic model with NeMo. <br> \n",
    "To understand the basics of Riva ASR APIs, refer to [Getting started with Riva ASR in Python](https://github.com/nvidia-riva/tutorials/blob/stable/asr-python-basics.ipynb). <br>\n",
    "\n",
    "For more information about Riva, refer to the [Riva developer documentation](https://developer.nvidia.com/riva)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Rv47EmBp_mK2"
   },
   "source": [
    "## NeMo (Neural Modules)\n",
    "[NVIDIA NeMo](https://developer.nvidia.com/nvidia-nemo) is an open-source framework for building, training, and fine-tuning GPU-accelerated speech AI and NLU models with a simple Python interface. For information about how to set up NeMo, refer to the [NeMo GitHub](https://github.com/NVIDIA/NeMo) instructions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "1O2bGCMAQDs8",
    "outputId": "a8731f17-b61c-4457-ae39-7079b8c4edd3"
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "You can run either this tutorial locally (if you have all the dependencies and a GPU) or on Google Colab.\n",
    "\n",
    "Perform the following steps to setup in Google Colab:\n",
    "1. Open a new Python 3 notebook.\n",
    "2. Import this notebook from GitHub.\n",
    "   a. Click **File** > **Upload Notebook** > **GITHUB** tab > copy/paste the GitHub URL.\n",
    "3. Connect to an instance with a GPU.\n",
    "   a. Click **Runtime** > Change the runtime type > select **GPU** for the hardware accelerator.\n",
    "4. Run this cell to set up the dependencies.\n",
    "5. Restart the runtime.\n",
    "   a. Click **Runtime** > **Restart Runtime** for any upgraded packages to take effect.\n",
    "\"\"\"\n",
    "\n",
    "# Install Dependencies\n",
    "!pip install wget\n",
    "!apt-get install sox libsndfile1 ffmpeg libsox-fmt-mp3\n",
    "!pip install text-unidecode\n",
    "!pip install matplotlib>=3.3.2\n",
    "!pip install Cython\n",
    "\n",
    "## Install NeMo\n",
    "BRANCH = 'main'\n",
    "!python -m pip install git+https://github.com/NVIDIA/NeMo.git@$BRANCH#egg=nemo_toolkit[all]\n",
    "\n",
    "\"\"\"\n",
    "Remember to restart the runtime for the kernel to pick up any upgraded packages (e.g. matplotlib)!\n",
    "Alternatively, in the case where you want to use the \"Run All Cells\" (or similar) option, \n",
    "uncomment `exit()` below to crash and restart the kernel.\n",
    "\"\"\"\n",
    "# exit()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vo27Unex_mLG"
   },
   "source": [
    "---\n",
    "## Fine-Tuning an ASR model with NeMo"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QfRpENJJ_mLy"
   },
   "source": [
    "### Download the Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AKrmQQuK_mLz"
   },
   "source": [
    "In this tutorial, we will use the popular AN4 dataset. Let's download it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "0eFF0r7z_mLz",
    "outputId": "c9f3772d-76ea-4e9d-b61e-64e32a6281df"
   },
   "outputs": [],
   "source": [
    "! wget https://dldata-public.s3.us-east-2.amazonaws.com/an4_sphere.tar.gz  # for the original source, please visit http://www.speech.cs.cmu.edu/databases/an4/an4_sphere.tar.gz"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qAsvlh53_mL0"
   },
   "source": [
    "After downloading, untar the dataset and move it to the correct directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "OipV7YVq_mL1",
    "outputId": "eaeacde1-a279-444a-af72-31d244b78162"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "DATA_DIR = os.getcwd()\n",
    "os.environ[\"DATA_DIR\"] = DATA_DIR\n",
    "! tar -xvf an4_sphere.tar.gz \n",
    "! mv an4 $DATA_DIR"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2Tilg2Eg_mL2"
   },
   "source": [
    "### Pre-Processing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jKud8fQ8_mL3"
   },
   "source": [
    "This step converts the `.mp3` files into `.wav` files and splits the data into training and testing sets. It also generates a \"meta-data\" file to be consumed by the data-loader for training and testing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "54Pff6vr_mL3",
    "outputId": "221abc41-b0f4-4e06-8138-57d8d47ec1db"
   },
   "outputs": [],
   "source": [
    "import json, librosa, os, glob\n",
    "import subprocess\n",
    "\n",
    "\n",
    "source_data_dir = f\"{DATA_DIR}/an4\"\n",
    "target_data_dir = f\"{DATA_DIR}/an4_converted\"\n",
    "\n",
    "def an4_build_manifest(transcripts_path, manifest_path, target_wavs_dir):\n",
    "    \"\"\"Build an AN4 manifest from a given transcript file.\"\"\"\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(target_wavs_dir, file_id + '.wav')\n",
    "\n",
    "                duration = librosa.core.get_duration(filename=audio_path)\n",
    "\n",
    "                # Write the metadata to the manifest\n",
    "                metadata = {\"audio_filepath\": audio_path, \"duration\": duration, \"text\": transcript}\n",
    "                json.dump(metadata, fout)\n",
    "                fout.write('\\n')\n",
    "\n",
    "\"\"\"Process AN4 dataset.\"\"\"\n",
    "if not os.path.exists(source_data_dir):\n",
    "    link = 'http://www.speech.cs.cmu.edu/databases/an4/an4_sphere.tar.gz'\n",
    "    raise ValueError(\n",
    "        f\"Data not found at `{source_data_dir}`. Please download the AN4 dataset from `{link}` \"\n",
    "        f\"and extract it into the folder specified by the `source_data_dir` argument.\"\n",
    "    )\n",
    "\n",
    "# Convert SPH files to WAV files\n",
    "sph_list = glob.glob(os.path.join(source_data_dir, '**/*.sph'), recursive=True)\n",
    "target_wavs_dir = os.path.join(target_data_dir, 'wavs')\n",
    "if not os.path.exists(target_wavs_dir):\n",
    "    print(f\"Creating directories for {target_wavs_dir}.\")\n",
    "    os.makedirs(os.path.join(target_data_dir, 'wavs'))\n",
    "\n",
    "for sph_path in sph_list:\n",
    "    wav_path = os.path.join(target_wavs_dir, os.path.splitext(os.path.basename(sph_path))[0] + '.wav')\n",
    "    cmd = [\"sox\", sph_path, wav_path]\n",
    "    subprocess.run(cmd, check=True)\n",
    "\n",
    "# Build AN4 manifests\n",
    "train_transcripts = os.path.join(source_data_dir, 'etc/an4_train.transcription')\n",
    "train_manifest = os.path.join(target_data_dir, 'train_manifest.json')\n",
    "an4_build_manifest(train_transcripts, train_manifest, target_wavs_dir)\n",
    "\n",
    "test_transcripts = os.path.join(source_data_dir, 'etc/an4_test.transcription')\n",
    "test_manifest = os.path.join(target_data_dir, 'test_manifest.json')\n",
    "an4_build_manifest(test_transcripts, test_manifest, target_wavs_dir)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2RAtW4wg_mL4"
   },
   "source": [
    "Let's listen to a sample audio file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "bGH3btfn_mL6"
   },
   "outputs": [],
   "source": [
    "# change path of the file here\n",
    "import os\n",
    "import IPython.display as ipd\n",
    "path = os.environ[\"DATA_DIR\"] + '/an4_converted/wavs/an268-mbmg-b.wav'\n",
    "ipd.Audio(path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FubCSAin_mMC"
   },
   "source": [
    "### Training "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fjHrIfDj_mMD"
   },
   "source": [
    "#### Create Tokenizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4zBKAO36_mME"
   },
   "source": [
    "Before we can do the actual training, we need to create a tokenizer as this ASR model uses word-piece encoding. Character based models don't need the tokenizer creation as only single characters are regarded as elements in the vocabulary in their cases. We can use NeMo's `process_asr_text_tokenizer.py` script to create the tokenizer that generates the subword vocabulary for us for use in training. The size of the vocabulary (`vocab_size`) should be the same as the vocabulary size in the ASR model. We will clone the NeMo GitHub repository to use the scripts and examples available there.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "maoasPkL_mMF",
    "outputId": "f468df62-ce9f-4828-ef6a-06e890de8b0e"
   },
   "outputs": [],
   "source": [
    "# clone NeMo locally\n",
    "NEMO_DIR = 'FIX_ME/path/to/NeMo'\n",
    "! git clone https://github.com/NVIDIA/NeMo $NEMO_DIR\n",
    "\n",
    "# create the tokenizer\n",
    "!python $NEMO_DIR/scripts/tokenizers/process_asr_text_tokenizer.py \\\n",
    "         --manifest=$DATA_DIR/an4_converted/train_manifest.json \\\n",
    "         --data_root=$DATA_DIR/an4 \\\n",
    "         --vocab_size=128 \\\n",
    "         --tokenizer=spe \\\n",
    "         --spe_type=unigram"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6-yDMTMO_mMJ"
   },
   "source": [
    "#### Training Conformer-CTC\n",
    "\n",
    "NeMo uses `.yml` files to configure the training parameters. You may update them directly by editing the configuration file or from the command-line interface. For example, if the number of epochs needs to be modified, along with a change in the learning rate, you can add `trainer.max_epochs=100` and `optim.lr=0.02` and train the model. \n",
    "\n",
    "The following sample command uses the `speech_to_text_ctc_bpe.py` script in the `examples` folder to train/fine-tune a Conformer-CTC ASR model for 1 epoch. For other ASR models like Citrinet, you may find the appropriate config files in the NeMo GitHub repo under [examples/asr/conf/](https://github.com/NVIDIA/NeMo/tree/main/examples/asr/conf).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "-bkkgxxv_mMK",
    "outputId": "9aeca99e-c3e0-4784-97f4-13cf997c2247"
   },
   "outputs": [],
   "source": [
    "# To fully train the model from scratch, you'll need to increase trainer.max_epochs from 1.\n",
    "# Empirical evidence suggests that around 200 epochs should suffice.\n",
    "!python $NEMO_DIR/examples/asr/asr_ctc/speech_to_text_ctc_bpe.py \\\n",
    "    --config-path=../conf/conformer/ --config-name=conformer_ctc_bpe \\\n",
    "    +init_from_pretrained_model=stt_en_conformer_ctc_large \\\n",
    "    model.train_ds.manifest_filepath=$DATA_DIR/an4_converted/train_manifest.json \\\n",
    "    model.validation_ds.manifest_filepath=$DATA_DIR/an4_converted/test_manifest.json \\\n",
    "    model.tokenizer.dir=$DATA_DIR/an4/tokenizer_spe_unigram_v128 \\\n",
    "    trainer.devices=1 \\\n",
    "    trainer.max_epochs=1 \\\n",
    "    model.optim.name=\"adamw\" \\\n",
    "    model.optim.lr=1.0 \\\n",
    "    model.optim.weight_decay=0.001 \\\n",
    "    model.optim.sched.warmup_steps=2000 \\\n",
    "    ++exp_manager.exp_dir=$DATA_DIR/checkpoints \\\n",
    "    ++exp_manager.version=test \\\n",
    "    ++exp_manager.use_datetime_version=False\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "nhGF0z_-rl8o",
    "outputId": "060b0b4c-224d-4242-bdef-72e0e346e137"
   },
   "outputs": [],
   "source": [
    "!ls ./Conformer-CTC-BPE/test/checkpoints/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nemo_file_path = os.path.join(DATA_DIR, 'checkpoints/Conformer-CTC-BPE/test/checkpoints/Conformer-CTC-BPE.nemo')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CKDipfZF_mMK"
   },
   "source": [
    "### ASR Evaluation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "f3svjWpR_mMN"
   },
   "source": [
    "Now that we have a model trained, we need to check how well it performs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "D4ATNCSe_mMN",
    "outputId": "b9039c6c-8840-4ce2-9089-689c29a8a162"
   },
   "outputs": [],
   "source": [
    "!python $NEMO_DIR/examples/asr/speech_to_text_eval.py \\\n",
    "    model_path=$nemo_file_path \\\n",
    "    dataset_manifest=$DATA_DIR/an4_converted/test_manifest.json \\\n",
    "    output_filename=./test_manifest_predictions.json \\\n",
    "    batch_size=32 \\\n",
    "    amp=True\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "M9ygVAEC_mMS"
   },
   "source": [
    "### ASR Model Export"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ap5CjiV4_mMT"
   },
   "source": [
    "With NeMo, you can also export your model in a format that can be deployed using NVIDIA Riva: a highly performant application framework for multi-modal conversational AI services using GPUs. The same command for exporting to ONNX can be used here. The only small variation is the configuration for `export_format` in the spec file."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "myI104s2pY13"
   },
   "source": [
    "#### Install the Packages\n",
    "\n",
    "We will now install the NeMo and `nemo2riva` packages. `nemo2riva` is available on [NVIDIA NGC](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/riva/resources/riva_quickstart/files?version=2.8.1). Make sure you install NGC CLI first before running the following commands."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Pd0Y3Gys_mMU",
    "outputId": "31cb1a3a-e539-4d5a-e395-862525f3cb99"
   },
   "outputs": [],
   "source": [
    "from version import __riva_version__\n",
    "print(__riva_version__)\n",
    "!pip install nvidia-pyindex\n",
    "!ngc registry resource download-version \"nvidia/riva/riva_quickstart:\"$__riva_version__\n",
    "!pip install nemo2riva\n",
    "!pip install protobuf==3.20.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8XNolXIPptAr"
   },
   "source": [
    "#### Convert to Riva\n",
    "\n",
    "Convert the downloaded model to the `.riva` format. We will set the encryption key with `--key=nemotoriva`. Choose a different encryption key value when generating `.riva` models for production."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "PlUNmVHKp1ft",
    "outputId": "f612c5e8-604e-4eee-8806-4eb463d1e810"
   },
   "outputs": [],
   "source": [
    "riva_file_path = nemo_file_path[:-5]+\".riva\"\n",
    "!nemo2riva --out {riva_file_path} --key=nemotoriva {nemo_file_path}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IAtoNe9krC2z"
   },
   "source": [
    "## More Resources\n",
    "You can find more information about working with NeMo's ASR models in the [ASR section](https://github.com/NVIDIA/NeMo/tree/main/tutorials/asr) of the NeMo tutorials."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Lv7ZRPoc_mMa"
   },
   "source": [
    "## What's Next?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SJ7hWpRT_mMb"
   },
   "source": [
    "You can use NeMo to build custom models for your own applications, and deploy them with NVIDIA Riva! Refer to the [Conformer-CTC deployment tutorial](https://github.com/nvidia-riva/tutorials/blob/main/asr-deployment-conformer-ctc.ipynb)."
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "provenance": []
  },
  "gpuClass": "standard",
  "kernelspec": {
   "display_name": "virtualenv-riva-tutorials-py38",
   "language": "python",
   "name": "virtualenv-riva-tutorials-py38"
  },
  "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}