asr-customize-vocabulary-and-lexicon.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"http://developer.download.nvidia.com/notebooks/dlsw-notebooks/riva_asr_asr-python-advanced-customize-vocabulary-and-lexicon/nvidia_logo.png\" style=\"width: 90px; float: right;\">\n",
    "\n",
    "# How to Customize Riva ASR Vocabulary and Pronunciation with Lexicon Mapping\n",
    "\n",
    "This notebook walks you through the process of customizing Riva ASR vocabulary and lexicon, in order to improve Riva vocabulary coverage and recognition of difficult words, such as acronyms.\n",
    "\n",
    "## Overview\n",
    "\n",
    "The Flashlight decoder, deployed by default in Riva, is a lexicon-based decoder and only emits words that are present in the provided lexicon file. That means, uncommon and new words, such as domain specific terminologies, that are not present in the lexicon file, will have no chance of being generated.\n",
    "\n",
    "On the other hand, the greedy decoder (available as an option during the `riva-build` process with the flag `--decoder_type=greedy`) is not lexicon-based and hence can virtually produce any word or character sequence.\n",
    "\n",
    "### Prerequisite\n",
    "\n",
    "This notebook assumes that the user is familiar with manually deploying a Riva ASR pipeline using the Riva ServiceMaker tool, `riva-build` and `riva-deploy` commands. See Riva [documentation](https://docs.nvidia.com/deeplearning/riva/user-guide/docs/model-overview.html#).\n",
    "\n",
    "### Terminologies\n",
    "- **Vocabulary file**: The vocabulary file is a flat text file containing a list of vocabulary words, each on its own line. For example:\n",
    "```\n",
    "the\n",
    "i\n",
    "to\n",
    "and\n",
    "a\n",
    "you\n",
    "of\n",
    "that\n",
    "...\n",
    "```\n",
    "\n",
    "This file is used by the `riva-build` process to generate the lexicon file. \n",
    "\n",
    "- **Lexicon file**: The lexicon file is a flat text file that contains the mapping of each vocabulary word to its tokenized form, e.g, sentencepiece tokens, separated by a `tab`. Below is an example:\n",
    "\n",
    "```\n",
    "with    ▁with\n",
    "not     ▁not\n",
    "this    ▁this\n",
    "just    ▁just\n",
    "my      ▁my\n",
    "as      ▁as\n",
    "don't   ▁don ' t\n",
    "```\n",
    "\n",
    "*Note: Ultimately, the Riva decoder makes use only of the lexicon file directly at run time (but not the vocabulary file).*\n",
    "\n",
    "Riva ServiceMaker automatically tokenizes the words in the vocabulary file to generate the lexicon file. It uses the correct tokenizer model that is packaged together with the acoustic model in the `.riva` file. By default, Riva generates 1 tokenized form for each word in the vocabulary file. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What can be customized?\n",
    "\n",
    "Both the vocabulary and the lexicon files can be customized.\n",
    "\n",
    "- Extending the vocabulary enriches the Riva default vocabulary, providing additional coverage for out-of-vocabulary words, terminologies, and abbreviations.\n",
    "\n",
    "- Customizing the lexicon file can further enrich the Riva knowledge base by providing one or more explicit pronunciations, in the form of tokenized sequences."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extending the vocabulary\n",
    "\n",
    "Extending the vocabulary must be done at Riva **build** time.\n",
    "\n",
    "When building a Riva ASR pipeline, pass the [extended vocabulary file](#modify_vocab) to the `--decoding_vocab=<vocabulary_file>` parameter of the build command. For example, the build command for the Citrinet model:\n",
    "\n",
    "```\n",
    "    riva-build speech_recognition \\\n",
    "   <rmir_filename>:<key> <riva_filename>:<key> \\\n",
    "   --name=citrinet-1024-english-asr-streaming \\\n",
    "   --decoding_language_model_binary=<lm_binary> \\\n",
    "   --decoding_vocab=<vocabulary_file> \\\n",
    "   --language_code=en-US \\\n",
    "   <other_parameters>...\n",
    "```\n",
    "\n",
    "Refer to Riva [documentation](https://docs.nvidia.com/deeplearning/riva/user-guide/docs/asr/asr-pipeline-configuration.html) for build commands for supported models.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='modify_vocab'></a>\n",
    "### How to modify vocabulary file\n",
    "\n",
    "You can either provide your own vocabulary file, or extend Riva's default vocabulary file.\n",
    "\n",
    "- BYO vocabulary file: provide a flat text file containing a list of vocabulary words, each on its own line. Note that this file must not only contain a small list of \"difficult words\", but must contains all the words that you want the ASR pipeline to be able to generate, that is, including all common words.\n",
    "\n",
    "- Modifying an existing one: This is the recommended approach. Out-of-the-box vocabulary files  for Riva supported languages can be found either:\n",
    "    - On NGC, for example, for English, the vocabulary file named `flashlight_decoder_vocab.txt` can be found at this [link](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/riva/models/speechtotext_en_us_lm/files).\n",
    "    - Or in a local Riva deployment: The actual physical location of Riva assets depends on the value of the `riva_model_loc` variable in the `config.sh` file under the Riva quickstart folder. The vocabulary file is bundled with the Flashlight decoder. \n",
    "        - By default, `riva_model_loc` is set to `riva-model-repo`, which is a docker volume. You can inspect this docker volume and copy the vocabulary file from within the docker volume to the host file system with commands such as:\n",
    "        \n",
    "        ```bash\n",
    "        # Inspect the Riva model docker volume\n",
    "        docker inspect riva-model-repo\n",
    "        \n",
    "        # Inspect the content of the Riva model docker volume\n",
    "        docker run --rm -v riva-model-repo:/riva-model-repo alpine ls /riva-model-repo\n",
    "        \n",
    "        # Copy the vocabulary file from the docker volume to the current directory\n",
    "        docker run --rm -v $PWD:/dest -v riva-model-repo:/riva-model-repo alpine cp  /riva-model-repo/models/citrinet-1024-en-US-asr-offline-ctc-decoder-cpu-offline/1/dict_vocab.txt /dest\n",
    "        ```\n",
    "        \n",
    "        - If you modify `riva_model_loc` to an absolute path pointing to a folder, then the specified folder in the local file system will be used to store Riva assets instead. Assuming `<RIVA_REPO_DIR>` is the directory where Riva assets are stored, then the vocabulary file can similarly be found under, for example, `<RIVA_REPO_DIR>/models/citrinet-1024-en-US-asr-offline-ctc-decoder-cpu-offline/1/dict_vocab.txt`.\n",
    "\n",
    "You can make a copy, then extend this default vocabulary file with the words of interest.\n",
    "\n",
    "Once modified, you'll have to redeploy the Riva ASR pipeline with `riva-build` while passing the flag `--decoding_vocab=<modified_vocabulary_file>`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Customizing pronunciation with lexicon mapping\n",
    "\n",
    "The lexicon file that is used by the Flashlight decoder can be found in the Riva assets directory, as specified by the value of the `riva_model_loc` variable in the `config.sh` file under the Riva quickstart folder (see above).\n",
    "\n",
    "- If `riva_model_loc` points to a docker volume (by default), you can find and copy the lexicon file with:\n",
    "```bash\n",
    "        # Copy the lexicon file from the docker volume to the current directory\n",
    "        docker run --rm -v $PWD:/dest -v riva-model-repo:/riva-model-repo alpine cp  /riva-model-repo/models/citrinet-1024-en-US-asr-offline-ctc-decoder-cpu-offline/1/lexicon.txt /dest\n",
    "```      \n",
    "\n",
    "- If you modify `riva_model_loc` to an absolute path pointing to a folder, then the specified folder in the local file system will be used to store Riva assets instead. Assuming `<RIVA_REPO_DIR>` is the directory where Riva assets are stored, then the vocabulary file can similarly be found under, for example, `<RIVA_REPO_DIR>/models/citrinet-1024-en-US-asr-offline-ctc-decoder-cpu-offline/1/lexicon.txt`.\n",
    "\n",
    "### How to modify the lexicon file\n",
    "\n",
    "First, locate and make a copy of the lexicon file. For example:\n",
    "```\n",
    "cp <RIVA_REPO_DIR>/models/citrinet-1024-en-US-asr-offline-ctc-decoder-cpu-offline/1/lexicon.txt modified_lexicon.txt\n",
    "```\n",
    "\n",
    "Next, modify it to add the sentencepiece tokenizations for the words of interest. For example, one could add:\n",
    "```\n",
    "manu ▁ma n u\n",
    "manu ▁man n n ew\n",
    "manu ▁man n ew\n",
    "```\n",
    "which are 3 different pronunciations/tokenizations of the word `manu`.  If the acoustic model predicts those tokens, they will be decoded as `manu`.\n",
    "\n",
    "Finally, once this is done, regenerate the model repository using that new decoding lexicon tokenization by passing `--decoding_lexicon=modified_lexicon.txt` to `riva-build` instead of `--decoding_vocab=decoding_vocab.txt`.\n",
    "\n",
    "### How to generate the correct tokenized form\n",
    "\n",
    "When modifying the lexicon file, ensure that:\n",
    "\n",
    "- The new lines follow the indentation/space pattern like the rest of the file and that the tokens used are part of the tokenizer model. \n",
    "\n",
    "- The tokens are valid tokens as determined by the tokenizer model (packaged with the Riva acoustic model).\n",
    "\n",
    "The latter ensures that you use only tokens that the acoustic model has been trained on. To do this, you’ll need the tokenizer model and the `sentencepiece` Python package (`pip install sentencepiece`). You can get the tokenizer model for the deployed pipeline from the model repository `ctc-decoder-...` directory for your model. It will be named `<hash>_tokenizer.model`. For example:\n",
    "\n",
    "`<RIVA_REPO_DIR>/models/citrinet-1024-en-US-asr-offline-ctc-decoder-cpu-offline/1/498056ba420d4bb3831ad557fba06032_tokenizer.model`\n",
    "\n",
    "When using a docker volume to store Riva assets (by default), you can copy the tokenizer model to the local directory with a command such as:\n",
    "\n",
    "```bash\n",
    "        # Copy the tokenizer model file from the docker volume to the current directory\n",
    "        docker run --rm -v $PWD:/dest -v riva-model-repo:/riva-model-repo alpine cp  /riva-model-repo/models/citrinet-1024-en-US-asr-offline-ctc-decoder-cpu-offline/1/498056ba420d4bb3831ad557fba06032_tokenizer.model /dest\n",
    "```  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install sentencepiece"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can then generate new lexicon entries, for example:\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "BRAF\t▁b ▁ra f\n",
      "BRAF\t▁b ▁ r a f\n",
      "BRAF\t▁b ▁ra f\n",
      "BRAF\t▁b ▁ra f\n",
      "BRAF\t▁ b ▁ra f\n"
     ]
    }
   ],
   "source": [
    "TOKEN=\"BRAF\"\n",
    "PRONUNCIATION=\"b raf\"\n",
    "\n",
    "import sentencepiece as spm\n",
    "s = spm.SentencePieceProcessor(model_file='tokenizer.model')\n",
    "for n in range(5):\n",
    "    print(TOKEN + '\\t' + ' '.join(s.encode(PRONUNCIATION, out_type=str, enable_sampling=True, alpha=0.1, nbest_size=-1)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note: `TOKEN` represents the desired written form of the word, while `PRONUNCIATION` is what the word should sound like.\n",
    "\n",
    "Other examples:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WhatsApp\t▁what ' s ▁app\n",
      "WhatsApp\t▁w h at ' s ▁app\n",
      "WhatsApp\t▁w h at ' s ▁app\n",
      "WhatsApp\t▁what ' s ▁a pp\n",
      "WhatsApp\t▁ w ha t ' s ▁app\n"
     ]
    }
   ],
   "source": [
    "TOKEN=\"WhatsApp\"\n",
    "PRONUNCIATION=\"what's app\"\n",
    "\n",
    "import sentencepiece as spm\n",
    "s = spm.SentencePieceProcessor(model_file='tokenizer.model')\n",
    "for n in range(5):\n",
    "    print(TOKEN + '\\t' + ' '.join(s.encode(PRONUNCIATION, out_type=str, enable_sampling=True, alpha=0.1, nbest_size=-1)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cya\t▁ S e e ▁y a\n",
      "Cya\t▁ S e e ▁y a\n",
      "Cya\t▁ S e e ▁y a\n",
      "Cya\t▁ S e e ▁y a\n",
      "Cya\t▁ S e e ▁ y a\n"
     ]
    }
   ],
   "source": [
    "TOKEN=\"Cya\"\n",
    "PRONUNCIATION=\"See ya\"\n",
    "\n",
    "import sentencepiece as spm\n",
    "s = spm.SentencePieceProcessor(model_file='tokenizer.model')\n",
    "for n in range(5):\n",
    "    print(TOKEN + '\\t' + ' '.join(s.encode(PRONUNCIATION, out_type=str, enable_sampling=True, alpha=0.1, nbest_size=-1)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Go deeper into Riva capabilities\n",
    "\n",
    "\n",
    "### Additional Riva Tutorials\n",
    "\n",
    "Checkout more Riva tutorials [here](https://github.com/nvidia-riva/tutorials) to understand how to use some of the advanced features of Riva ASR, including customizing ASR for your specific needs.\n",
    "\n",
    "\n",
    "### Sample Applications\n",
    "\n",
    "Riva comes with various sample applications. They demonstrate how to use the APIs to build various applications. Refer to [Riva Sample Apps](https://docs.nvidia.com/deeplearning/riva/user-guide/docs/samples/index.html) for more information.  \n",
    "\n",
    "\n",
    "### Additional Resources\n",
    "\n",
    "For more information about each of the Riva APIs and their functionalities, refer to the [documentation](https://docs.nvidia.com/deeplearning/riva/user-guide/docs/reference/protos/protos.html)."
   ]
  }
 ],
 "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.8.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}