tutorials/nlp/SpellMapper_English_ASR_Customization.ipynb

{
  "cells": [
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "PiRuohn_FQco"
      },
      "source": [
        "# Overview\n",
        "This tutorial demonstrates how to run inference with [SpellMapper](https://arxiv.org/abs/2306.02317) - a model for Spellchecking ASR (Automatic Speech Recognition) Customization.\n",
        "\n",
        "Estimated time: 10-15 min.\n",
        "\n",
        "SpellMapper is a non-autoregressive (NAR) model based on transformer architecture ([BERT](https://arxiv.org/pdf/1810.04805.pdf) with multiple separators).\n",
        "It gets as input a single ASR hypothesis (text) and a **custom vocabulary** and predicts which fragments in the ASR hypothesis should be replaced by which custom words/phrases if any.\n",
        "\n",
        "This model is an alternative to word boosting/shallow fusion approaches:\n",
        "  - does not require retraining ASR model;\n",
        "  - does not require beam-search/language model(LM);\n",
        "  - can be applied on top of any English ASR model output;"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qm5wmxVEGXgH"
      },
      "source": [
        "## What is custom vocabulary?\n",
        "**Custom vocabulary** is a list of words/phrases that are important for a particular user. For example, user's contact names, playlist, selected terminology and so on. The size of the custom vocabulary can vary from several hundreds to **several thousand entries** - but this is not an equivalent to ngram language model.\n",
        "\n",
        "![Scope of customization with user vocabulary](images/spellmapper_customization_vocabulary.png)\n",
        "\n",
        "Note that unlike traditional spellchecking approaches, which aim to correct known words using language models, the goal of contextual spelling correction is to correct highly specific user terms, most of which can be 1) out-of-vocabulary (OOV) words, 2) spelling variations (e.g., \"John Koehn\", \"Jon Cohen\") and language models cannot help much with that."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D5_XwuXDOKho"
      },
      "source": [
        "## Tutorial Plan\n",
        "\n",
        "1.   Create a sample custom vocabulary using some medical terminology.\n",
        "2.   Study what customization does - a detailed analysis of a small example.\n",
        "3.   Run a bigger example:\n",
        "   *  Create sample ASR results by running TTS (text-to-speech synthesis) + ASR on some medical paper abstracts.\n",
        "   *  Run SpellMapper inference and show how it can improve ASR results using custom vocabulary.\n",
        "\n",
        "TL;DR We reduce WER from `14.3%` to `11.4%` by correcting medical terms, e.g.\n",
        "* `puramesin` => `puromycin`\n",
        "* `parromsin` => `puromycin`\n",
        "* `and hydrod` => `anhydride`\n",
        "* `lesh night and` => `lesch-nyhan`\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "agz8B2CxXBBG"
      },
      "source": [
        "# Preparation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "koRPpYISNPuH"
      },
      "source": [
        "## Installing NeMo"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HCnnz3cgVc4Q"
      },
      "outputs": [],
      "source": [
        "# Install NeMo library. If you are running locally (rather than on Google Colab), comment out the below lines\n",
        "# and instead follow the instructions at https://github.com/NVIDIA/NeMo#Installation\n",
        "GITHUB_ACCOUNT = \"NVIDIA\"\n",
        "BRANCH = 'main'\n",
        "!python -m pip install git+https://github.com/{GITHUB_ACCOUNT}/NeMo.git@{BRANCH}#egg=nemo_toolkit[all]\n",
        "\n",
        "# Download local version of NeMo scripts. If you are running locally and want to use your own local NeMo code,\n",
        "# comment out the below lines and set NEMO_DIR to your local path.\n",
        "NEMO_DIR = 'nemo'\n",
        "!git clone -b {BRANCH} https://github.com/{GITHUB_ACCOUNT}/NeMo.git $NEMO_DIR"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_M92gCn_NW1_"
      },
      "source": [
        "## Additional installs\n",
        "We will use `sentence_splitter` to split abstracts to sentences."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ddyJA3NtGl9C"
      },
      "outputs": [],
      "source": [
        "!pip install sentence_splitter"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qVa91rGkeFje"
      },
      "source": [
        "Clone the SpellMapper model from HuggingFace.\n",
        "Note that we will need not only the checkpoint itself, but also the ngram mapping vocabulary `replacement_vocab_filt.txt` from the same folder."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JiI9dkEm5cpW"
      },
      "outputs": [],
      "source": [
        "!git clone https://huggingface.co/bene-ges/spellmapper_asr_customization_en"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8saqFOePVfFf"
      },
      "source": [
        "## Imports\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tAJyiYn_VnrF"
      },
      "outputs": [],
      "source": [
        "import IPython.display as ipd\n",
        "import json\n",
        "import random\n",
        "import re\n",
        "import soundfile as sf\n",
        "import torch\n",
        "\n",
        "from collections import Counter, defaultdict\n",
        "from difflib import SequenceMatcher\n",
        "from matplotlib.pyplot import imshow\n",
        "from matplotlib import pyplot as plt\n",
        "from sentence_splitter import SentenceSplitter\n",
        "from typing import List, Set, Tuple\n",
        "\n",
        "from nemo.collections.tts.models import FastPitchModel\n",
        "from nemo.collections.tts.models import HifiGanModel\n",
        "\n",
        "from nemo.collections.asr.parts.utils.manifest_utils import read_manifest\n",
        "\n",
        "from nemo.collections.nlp.data.spellchecking_asr_customization.utils import (\n",
        "    get_all_candidates_coverage,\n",
        "    get_index,\n",
        "    load_ngram_mappings,\n",
        "    search_in_index,\n",
        "    get_candidates,\n",
        "    read_spellmapper_predictions,\n",
        "    apply_replacements_to_text,\n",
        "    load_ngram_mappings_for_dp,\n",
        "    get_alignment_by_dp,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mfAaOdAWUGUV"
      },
      "source": [
        "Use seed to get a reproducible behaviour."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UlGnNKTuT_6A"
      },
      "outputs": [],
      "source": [
        "random.seed(0)\n",
        "torch.manual_seed(0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RPPHI7Zd_fDz"
      },
      "source": [
        "## Download data\n",
        "\n",
        "File `pubmed24n0009.xml` taken from public ftp server of https://www.ncbi.nlm.nih.gov/pmc/ contains information about 5593 medical papers, from which we extract only their abstracts. We will feed sentences from there to TTS + ASR to get initial ASR results.\n",
        "\n",
        "File `wordlist.txt` contains 100k **single-word** medical terms.\n",
        "\n",
        "File `valid_adam.txt` contains 24k medical abbreviations with their full forms. We will use those full forms as examples of **multi-word** medical terms.\n",
        "\n",
        "File `count_1w.txt` contains 330k single words with their frequencies from Google Ngrams corpus. We will use this file to filter out frequent words from our custom vocabulary.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mX6cvE8xw2n1"
      },
      "outputs": [],
      "source": [
        "!wget https://ftp.ncbi.nlm.nih.gov/pubmed/baseline/pubmed24n0009.xml.gz\n",
        "!gunzip pubmed24n0009.xml.gz\n",
        "!grep \"AbstractText\" pubmed24n0009.xml > abstract.txt\n",
        "\n",
        "!wget https://raw.githubusercontent.com/McGill-NLP/medal/master/toy_data/valid_adam.txt\n",
        "!wget https://raw.githubusercontent.com/glutanimate/wordlist-medicalterms-en/master/wordlist.txt\n",
        "!wget https://norvig.com/ngrams/count_1w.txt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mBm9BeqNaRlC"
      },
      "source": [
        "## Auxiliary functions\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kVUKhSh48Ypi"
      },
      "outputs": [],
      "source": [
        "CHARS_TO_IGNORE_REGEX = re.compile(r\"[\\.\\,\\?\\:!;()«»…\\]\\[/\\*–‽+&_\\\\½√>€™$•¼}{~—=“\\\"”″‟„]\")\n",
        "\n",
        "\n",
        "def get_medical_vocabulary() -> Tuple[Set[str], Set[str]]:\n",
        "    \"\"\"This function builds a vocabulary of medical terms using downloaded sources:\n",
        "        wordlist.txt - 100k single-word medical terms.\n",
        "        valid_adam.txt - 24k medical abbreviations with their full forms. We use those full forms as examples of multi-word medical terms.\n",
        "        count_1w.txt - 330k single words with their frequencies from Google Ngrams corpus. We will use this file to filter out frequent words from our custom vocabulary.\n",
        "    \"\"\"\n",
        "    common_words  = set()\n",
        "    with open(\"count_1w.txt\", \"r\", encoding=\"utf-8\") as f:\n",
        "        for line in f:\n",
        "            word, freq = line.strip().casefold().split(\"\\t\")\n",
        "            if int(freq) < 500000:\n",
        "                break\n",
        "            common_words.add(word)\n",
        "    print(\"Size of common words vocabulary:\", len(common_words))\n",
        "\n",
        "    abbreviations = defaultdict(set)\n",
        "    medical_vocabulary = set()\n",
        "    with open(\"valid_adam.txt\", \"r\", encoding=\"utf-8\") as f:\n",
        "        lines = f.readlines()\n",
        "        # first line is header\n",
        "        for line in lines[1:]:\n",
        "            abbrev, _, phrase = line.strip().split(\"\\t\")\n",
        "            # skip phrases longer than 3 words because some of them are long explanations\n",
        "            if phrase.count(\" \") > 2:\n",
        "                continue\n",
        "            if phrase in common_words:\n",
        "                continue\n",
        "            medical_vocabulary.add(phrase)\n",
        "            abbrev = abbrev.lower()\n",
        "            abbreviations[abbrev].add(phrase)\n",
        "\n",
        "    with open(\"wordlist.txt\", \"r\", encoding=\"utf-8\") as f:\n",
        "        for line in f:\n",
        "            word = line.strip().casefold()\n",
        "            # skip words containing digits\n",
        "            if re.match(r\".*\\d.*\", word):\n",
        "                continue\n",
        "            if re.match(r\".*[\\[\\]\\(\\)\\+\\,\\.].*\", word):\n",
        "                continue\n",
        "            if word in common_words:\n",
        "                continue\n",
        "            medical_vocabulary.add(word)\n",
        "\n",
        "    print(\"Size of medical vocabulary:\", len(medical_vocabulary))\n",
        "    print(\"Size of abbreviation vocabulary:\", len(abbreviations))\n",
        "    return medical_vocabulary, abbreviations\n",
        "\n",
        "\n",
        "def read_abstracts(medical_vocabulary: Set[str]) -> Tuple[List[str], Set[str], Set[str]]:\n",
        "    \"\"\"This function reads the downloaded medical abstracts, and extracts sentences containing any word/phrase from the medical vocabulary.\n",
        "    Args:\n",
        "        medical_vocabulary: set of known medical words or phrases\n",
        "    Returns:\n",
        "        sentences: list of extracted sentences\n",
        "        all_found_singleword: set of single words from medical vocabulary that occurred at least in one sentence\n",
        "        all_found_multiword: set of multi-word phrases from medical vocabulary that occurred at least in one sentence\n",
        "    \"\"\"\n",
        "    splitter = SentenceSplitter(language='en')\n",
        "\n",
        "    all_sentences = []\n",
        "    all_found_singleword = set()\n",
        "    all_found_multiword = set()\n",
        "    with open(\"abstract.txt\", \"r\", encoding=\"utf-8\") as f:\n",
        "        for line in f:\n",
        "            text = line.strip().replace(\"<AbstractText>\", \"\").replace(\"</AbstractText>\", \"\")\n",
        "            sents = splitter.split(text)\n",
        "            found_singleword = set()\n",
        "            found_multiword = set()\n",
        "            for sent in sents:\n",
        "                # remove anything in brackets from text\n",
        "                sent = re.sub(r\"\\(.+\\)\", r\"\", sent)\n",
        "                # remove quotes from text\n",
        "                sent = sent.replace(\"\\\"\", \"\")\n",
        "                # skip sentences containing digits because normalization is out of scope of this tutorial\n",
        "                if re.match(r\".*\\d.*\", sent):\n",
        "                    continue\n",
        "                # skip sentences containing abbreviations with period inside the sentence (for the same reason)\n",
        "                if \". \" in sent:\n",
        "                    continue\n",
        "                # skip long sentences as they may cause OOM issues\n",
        "                if len(sent) > 150:\n",
        "                    continue\n",
        "                # replace all punctuation to space and convert to lowercase\n",
        "                sent_clean = CHARS_TO_IGNORE_REGEX.sub(\" \", sent).lower()\n",
        "                sent_clean = \" \".join(sent_clean.split(\" \"))\n",
        "                words = sent_clean.split(\" \")\n",
        "\n",
        "                found_phrases = set()\n",
        "                for begin in range(len(words)):\n",
        "                    for end in range(begin + 1, min(begin + 4, len(words))):\n",
        "                        phrase = \" \".join(words[begin:end])\n",
        "                        if phrase in medical_vocabulary:\n",
        "                            found_phrases.add(phrase)\n",
        "                            if end - begin == 1:\n",
        "                                found_singleword.add(phrase)\n",
        "                            else:\n",
        "                                found_multiword.add(phrase)\n",
        "                if len(found_phrases) > 0:\n",
        "                    all_sentences.append((sent, \";\".join(found_phrases)))\n",
        "            all_found_singleword = all_found_singleword.union(found_singleword)\n",
        "            all_found_multiword = all_found_multiword.union(found_multiword)\n",
        "\n",
        "    print(\"Sentences:\", len(all_sentences))\n",
        "    print(\"Unique single-word terms found:\", len(all_found_singleword))\n",
        "    print(\"Unique multi-word terms found:\", len(all_found_multiword))\n",
        "    print(\"Examples of multi-word terms\", str(list(all_found_multiword)[0:10]))\n",
        "    \n",
        "    return all_sentences, all_found_singleword, all_found_multiword"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XU3xeCBVpWOL"
      },
      "outputs": [],
      "source": [
        "def get_fragments(i_words: List[str], j_words: List[str]) -> List[Tuple[str, str, str, int, int, int, int]]:\n",
        "    \"\"\"This function is used to compare two word sequences to find minimal fragments that differ.\n",
        "    Args:\n",
        "        i_words: list of words in first sequence\n",
        "        j_words: list of words in second sequence\n",
        "    Returns:\n",
        "        list of tuples (difference_type, fragment1, fragment2, begin_of_fragment1, end_of_fragment1, begin_of_fragment2, end_of_fragment2)\n",
        "    \"\"\"\n",
        "    s = SequenceMatcher(None, i_words, j_words)\n",
        "    result = []\n",
        "    for tag, i1, i2, j1, j2 in s.get_opcodes():\n",
        "        result.append((tag, \" \".join(i_words[i1:i2]), \" \".join(j_words[j1:j2]), i1, i2, j1, j2))\n",
        "    result = sorted(result, key=lambda x: x[3])\n",
        "    return result"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2ydXp_pFYmYu"
      },
      "source": [
        "## Read medical data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WAeauax0SV1-"
      },
      "outputs": [],
      "source": [
        "medical_vocabulary, _ = get_medical_vocabulary()\n",
        "sentences, found_singleword, found_multiword = read_abstracts(medical_vocabulary)\n",
        "# in case if we need random candidates from a big sample - we will use full medical vocabulary for that purpose.\n",
        "big_sample = list(medical_vocabulary)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FRli7-Kx7sOO"
      },
      "outputs": [],
      "source": [
        "for sent, phrases in sentences[0:10]:\n",
        "    print(sent, \"\\t\", phrases)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rL1VqH2_dk93"
      },
      "source": [
        "# SpellMapper ASR Customization\n",
        "\n",
        "SpellMapper model relies on two offline preparation steps:\n",
        "1. Collecting n-gram mappings from a large corpus (this mappings vocabulary had been collected once on a large corpus and is supplied with the model).\n",
        "2. Indexing of user vocabulary by n-grams.\n",
        "\n",
        "![Offline data preparation](images/spellmapper_data_preparation.png)\n",
        "\n",
        "At inference time we take as input an ASR hypothesis and an n-gram-indexed user vocabulary and perform following steps:\n",
        "1. Retrieve the top 10 candidate phrases from the user vocabulary that are likely to be contained in the given ASR-hypothesis, possibly in a misspelled form.\n",
        "2. Run the neural model that tags the input characters with correct candidate labels or 0 if no match is found.\n",
        "3. Do post-processing to combine results.\n",
        "\n",
        "![Inference pipeline](images/spellmapper_inference_pipeline.png)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OeJpsMwslmrd"
      },
      "source": [
        "## N-gram mappings\n",
        "Note that n-gram mappings vocabulary had been collected from a large corpus and is supplied with the model. It is supposed to be \"universal\" for English language.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uH6p0mOd12pi"
      },
      "source": [
        "Let's see what n-gram mappings are like, for example, for an n-gram `l u c`.\n",
        "Note that n-grams in `replacement_vocab_filt.txt` preserve one-to-one correspondence between original letters and misspelled fragments (this additional markup is handled during loading). \n",
        "* `+` means that adjacent letters are concatenated and correspond to a single source letter. \n",
        "* `<DELETE>` means that the original letter is deleted. \n",
        "This auxiliary markup will be removed automatically during loading.\n",
        "\n",
        "`_` is used instead of real space symbol.\n",
        "\n",
        "Last three columns are:\n",
        "* joint frequency\n",
        "* frequency of original n-gram\n",
        "* frequency of misspelled n-gram\n",
        "\n",
        "$$\\frac{JointFrequency}{SourceFrequency}=TranslationProbability$$\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qul163dB1sKp"
      },
      "outputs": [],
      "source": [
        "!awk 'BEGIN {FS=\"\\t\"} ($1==\"l u c\"){print $0}' < spellmapper_asr_customization_en/replacement_vocab_filt.txt | sort -t$'\\t' -k3nr"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eWxcrVWZ3Pfq"
      },
      "source": [
        "Now we read n-gram mappings from the file. Parameter `max_misspelled_freq` controls maximum frequency of misspelled n-grams. N-grams more frequent than that are put in the list of banned n-grams and won't be used in indexing."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WHKhE945-N7o"
      },
      "outputs": [],
      "source": [
        "print(\"load n-gram mappings...\")\n",
        "ngram_mapping_vocab, ban_ngram = load_ngram_mappings(\"spellmapper_asr_customization_en/replacement_vocab_filt.txt\", max_misspelled_freq=125000)\n",
        "# CAUTION: entries in ban_ngram end with a space and can contain \"+\" \"=\"\n",
        "print(\"Size of ngram mapping vocabulary:\", len(ngram_mapping_vocab))\n",
        "print(\"Size of banned ngrams:\", len(ban_ngram))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "49IcMBfllvXN"
      },
      "source": [
        "## Indexing of custom vocabulary"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b1K6paeee2Iu"
      },
      "source": [
        "As we mentioned earlier, this model pipeline is intended to work with custom vocabularies up to several thousand entries. Since the whole medical vocabulary contains 110k entries, we restrict our custom vocabulary to 5000+ terms that occurred in given corpus of abstracts.\n",
        "\n",
        "The goal of indexing our custom vocabulary is to build an index where key is a letter n-gram and value is the whole phrase. The keys are n-grams in the given user phrase and their misspelled variants taken from our collection of n-\n",
        "gram mappings (see Index of custom vocabulary in Fig. 1)\n",
        "\n",
        "*Though it is possible to index and search the whole 110k vocabulary, it will require additional optimizations and is beyond the scope of this tutorial.*"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xWb0jGqw6Woi"
      },
      "outputs": [],
      "source": [
        "custom_phrases = []\n",
        "for phrase in medical_vocabulary:\n",
        "    if phrase not in found_singleword and phrase not in found_multiword:\n",
        "        continue\n",
        "    custom_phrases.append(\" \".join(list(phrase.replace(\" \", \"_\"))))\n",
        "print(\"Size of customization vocabulary:\", len(custom_phrases))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UHWor5pD2Eyb"
      },
      "source": [
        "Now we build the index for our custom phrases.\n",
        "\n",
        "Parameter `min_log_prob` controls minimum log probability, after which we stop growing this n-gram.\n",
        "\n",
        "Parameter `max_phrases_per_ngram` controls maximum number of phrases that can be indexed by one ngram. N-grams exceeding this limit are also banned and not used in indexing.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hs4RDXj0-xW9"
      },
      "outputs": [],
      "source": [
        "phrases, ngram2phrases = get_index(custom_phrases, ngram_mapping_vocab, ban_ngram, min_log_prob=-4.0, max_phrases_per_ngram=600)\n",
        "print(\"Size of phrases:\", len(phrases))\n",
        "print(\"Size of ngram2phrases:\", len(ngram2phrases))\n",
        "\n",
        "# Save index to file - later we will use it in other script\n",
        "with open(\"index.txt\", \"w\", encoding=\"utf-8\") as out:\n",
        "    for ngram in ngram2phrases:\n",
        "        for phrase_id, begin, size, logprob in ngram2phrases[ngram]:\n",
        "            phrase = phrases[phrase_id]\n",
        "            out.write(ngram + \"\\t\" + phrase + \"\\t\" + str(begin) + \"\\t\" + str(size) + \"\\t\" + str(logprob) + \"\\n\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RV1sdQ9rvar8"
      },
      "source": [
        "## Small detailed example\n",
        "\n",
        "Let's consider, for example, one custom phrase `thoracic aorta` and an incorrect ASR-hypothesis `the tarasic oorda is a part of the aorta located in the thorax`, containing a misspelled phrase `tarasic_oorda`. \n",
        "\n",
        "We will see \n",
        "1. How this custom phrase is indexed.\n",
        "2. How candidate retrieval works, given ASR-hypothesis.\n",
        "3. How inference and post-processing work.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kGBTTJXixnrG"
      },
      "source": [
        "### N-grams in index"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ryfUlqNMl4vQ"
      },
      "source": [
        "Let's look, for example, by what n-grams a custom phrase `thoracic aorta` is indexed. \n",
        "Columns: \n",
        "1. n-gram\n",
        "2. beginning position in the phrase\n",
        "3. length\n",
        "4. log probability\n",
        "\n",
        "Note that many n-grams are not from n-gram mappings file. Those are derived by  growing previous n-grams with new replacements. In this case log probabilities are summed up. Growing stops, when minimum log prob is exceeded.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "x0ZVsXGBo8pt"
      },
      "outputs": [],
      "source": [
        "for ngram in ngram2phrases:\n",
        "    for phrase_id, b, length, lprob in ngram2phrases[ngram]:\n",
        "        if phrases[phrase_id] == \"t h o r a c i c _ a o r t a\":\n",
        "            print(ngram.ljust(16) + \"\\t\" + str(b).rjust(4) + \"\\t\" + str(length).rjust(4) + \"\\t\" + str(lprob))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "20ov23ze4xeQ"
      },
      "source": [
        "### Candidate retrieval\n",
        "Candidate retrieval tasks are:\n",
        " - Given an input sentence and an index of custom vocabulary find all n-grams from the index matching the sentence. \n",
        " - Find which sentence fragments and which custom phrases have most \"hits\" - potential candidates.\n",
        " - Find approximate starting position for each candidate phrase. \n",
        "\n",
        "\n",
        "Let's look at the hits, that phrase \"thoracic aorta\" gets by searching all ngrams in the input text. We can see some hits in different part of the sentence, but a moving window can find a fragment with most hits."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "t_rhKQ3Xqa8A"
      },
      "outputs": [],
      "source": [
        "sent = \"the_tarasic_oorda_is_a_part_of_the_aorta_located_in_the_thorax\"\n",
        "phrases2positions, position2ngrams = search_in_index(ngram2phrases, phrases, sent)\n",
        "print(\" \".join(list(sent)))\n",
        "print(\" \".join(list(map(str, phrases2positions[phrases.index(\"t h o r a c i c _ a o r t a\")].astype(int)))))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "orkRapbjF4aZ"
      },
      "source": [
        "`phrases2positions` is a matrix of size (len(phrases), len(ASR_hypothesis)).\n",
        "It is filled with 1.0 (hits) on intersection of letter n-grams and phrases that are indexed by these n-grams, 0.0 - elsewhere.\n",
        "It is used to find phrases with many hits within a contiguous window - potential matching candidates.\n",
        "\n",
        "`position2ngrams` is a list of sets of ngrams. List index is the starting position in the ASR-hypothesis.\n",
        "It is used later to check how well each found candidate is covered by n-grams (to avoid cases where some repeating n-gram gives many hits to a phrase, but the phrase itself is not well covered)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JF7u4_iiHLyI"
      },
      "outputs": [],
      "source": [
        "candidate2coverage, candidate2position = get_all_candidates_coverage(phrases, phrases2positions)\n",
        "print(\"Coverage=\", candidate2coverage[phrases.index(\"t h o r a c i c _ a o r t a\")])\n",
        "print(\"Starting position=\", candidate2position[phrases.index(\"t h o r a c i c _ a o r t a\")])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "45mvKg8ZyNbr"
      },
      "source": [
        "`candidate2coverage` is a list of size len(phrases) containing coverage (0.0 to 1.0) in best window.\n",
        "Coverage is a smoothed percentage of hits in the window of size of the given phrase.\n",
        "\n",
        "`candidate2position` is a list of size len(phrases) containing starting position of best window.\n",
        "\n",
        "Starting position is approximate, it's ok. If it is not at the beginning of some word, SpellMapper will try to adjust it later. In this particular example we get 5 as starting position instead of 4, missing the first letter."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Sjyn9I98udL9"
      },
      "source": [
        "### Inference\n",
        "\n",
        "Now let's generate input for SpellMapper inference. \n",
        "An input line should consist of 4 tab-separated columns:\n",
        "  - text of ASR-hypothesis\n",
        "  - texts of 10 candidates separated by semicolon\n",
        "  - 1-based ids of non-dummy candidates\n",
        "  - approximate start/end coordinates of non-dummy candidates (correspond to ids)\n",
        "Note that candidate retrieval is done inside the function `get_candidates`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cJnusVfBRhRX"
      },
      "outputs": [],
      "source": [
        "out = open(\"spellmapper_input.txt\", \"w\", encoding=\"utf-8\")\n",
        "letters = list(sent)\n",
        "candidates = get_candidates(ngram2phrases, phrases, letters, big_sample)\n",
        "# We add two columns with targets and span_info. \n",
        "# They have same format as during training, but start and end positions are APPROXIMATE, they will be adjusted when constructing BertExample.\n",
        "targets = []\n",
        "span_info = []\n",
        "for idx, c in enumerate(candidates):\n",
        "    if c[1] == -1:\n",
        "        continue\n",
        "    targets.append(str(idx + 1))  # targets are 1-based\n",
        "    start = c[1]\n",
        "    end = min(c[1] + c[2], len(letters))  # ensure that end is not outside sentence length (it can happen because c[2] is candidate length used as approximation)\n",
        "    span_info.append(\"CUSTOM \" + str(start) + \" \" + str(end))\n",
        "\n",
        "out.write(\" \".join(letters) + \"\\t\" + \";\".join([x[0] for x in candidates])  + \"\\t\" + \" \".join(targets) + \"\\t\" + \";\".join(span_info) + \"\\n\")\n",
        "out.close()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Qpei5o89SmaU"
      },
      "outputs": [],
      "source": [
        "!cat spellmapper_input.txt"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9rAmO15SS6go"
      },
      "outputs": [],
      "source": [
        "!python nemo/examples/nlp/spellchecking_asr_customization/spellchecking_asr_customization_infer.py \\\n",
        "      pretrained_model=spellmapper_asr_customization_en/training_10m_5ep.nemo \\\n",
        "      model.max_sequence_len=512 \\\n",
        "      inference.from_file=spellmapper_input.txt \\\n",
        "      inference.out_file=spellmapper_output.txt \\\n",
        "      inference.batch_size=16 \\\n",
        "      lang=en\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wd2aq4T1N5cs"
      },
      "source": [
        "Each line in SpellMapper output is tab-separated and consists of 4 columns:\n",
        "1. ASR-hypothesis (same as in input)\n",
        "2. 10 candidates separated with semicolon (same as in input)\n",
        "3. fragment predictions, separated with semicolon, each prediction is a tuple (start, end, candidate_id, probability)\n",
        "4. letter predictions - candidate_id predicted for each letter (this is only for debug purposes)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ravgEX8cTFty"
      },
      "outputs": [],
      "source": [
        "!cat spellmapper_output.txt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "az26364-PHb2"
      },
      "source": [
        "We can use some utility functions to apply found replacements and get actual corrected text."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lPtFa_EhK8pb"
      },
      "outputs": [],
      "source": [
        "spellmapper_results = read_spellmapper_predictions(\"spellmapper_output.txt\")\n",
        "text, replacements, _ = spellmapper_results[0]\n",
        "corrected_text = apply_replacements_to_text(text, replacements, replace_hyphen_to_space=False)\n",
        "print(\"Text before correction:\\n\", text)\n",
        "print(\"Text after correction:\\n\", corrected_text)\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "efF7O-D91FLX"
      },
      "source": [
        "# Bigger customization example\n",
        "\n",
        "Let's test customization on more data. The plan is\n",
        "   *  Get baseline ASR transcriptions by running TTS + ASR on some medical paper abstracts.\n",
        "   *  Run SpellMapper inference and show how it can improve ASR results using custom vocabulary.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r_EFPnyDcXZt"
      },
      "source": [
        "## Run TTS"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "i9F5SBhmr8rk"
      },
      "outputs": [],
      "source": [
        "# create a folder for wav files (TTS output)\n",
        "!rm -r audio\n",
        "!mkdir audio"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JMbkNVt7YBAO"
      },
      "outputs": [],
      "source": [
        "if torch.cuda.is_available():\n",
        "  device = \"cuda\"\n",
        "else:\n",
        "  device = \"cpu\"\n",
        "\n",
        "# Load FastPitch from HuggingFace\n",
        "spectrogram_generator = FastPitchModel.from_pretrained(\"nvidia/tts_en_fastpitch\").eval().to(device)\n",
        "# Load HifiGan vocoder from HuggingFace\n",
        "vocoder = HifiGanModel.from_pretrained(model_name=\"nvidia/tts_hifigan\").eval().to(device)\n",
        "\n",
        "# Write sentences that we want to feed to TTS\n",
        "with open(\"tts_input.txt\", \"w\", encoding=\"utf-8\") as out:\n",
        "    for sent, _ in sentences[0:100]:\n",
        "        out.write(sent + \"\\n\")\n",
        "\n",
        "out_manifest = open(\"manifest.json\", \"w\", encoding=\"utf-8\")\n",
        "i = 0\n",
        "with open(\"tts_input.txt\", \"r\", encoding=\"utf-8\") as inp:\n",
        "    for line in inp:\n",
        "        text = line.strip()\n",
        "        text_clean = CHARS_TO_IGNORE_REGEX.sub(\" \", text).lower()  #replace all punctuation to space and convert to lowercase\n",
        "        text_clean = \" \".join(text_clean.split())\n",
        "\n",
        "        parsed = spectrogram_generator.parse(text, normalize=True)\n",
        "\n",
        "        spectrogram = spectrogram_generator.generate_spectrogram(tokens=parsed)\n",
        "        audio = vocoder.convert_spectrogram_to_audio(spec=spectrogram)\n",
        "\n",
        "        # Note that vocoder return a batch of audio. In this example, we just take the first and only sample.\n",
        "        filename = \"audio/\" + str(i) + \".wav\"\n",
        "        sf.write(filename, audio.to('cpu').detach().numpy()[0], 16000)\n",
        "        out_manifest.write(\n",
        "            \"{\\\"audio_filepath\\\": \\\"\" + filename + \"\\\", \\\"text\\\": \\\"\" + text_clean + \"\\\", \\\"orig_text\\\": \\\"\" + text + \"\\\"}\\n\"\n",
        "        )\n",
        "        i += 1\n",
        "\n",
        "        # display some examples\n",
        "        if i < 10:\n",
        "            print(f'\"{text}\"\\n')\n",
        "            ipd.display(ipd.Audio(audio.to('cpu').detach(), rate=22050))\n",
        "\n",
        "out_manifest.close()\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9T3CZcCAmxCz"
      },
      "source": [
        "Now we have a folder with generated audios `audio/*.wav` and a nemo manifest with json records like `{\"audio_filepath\": \"audio/0.wav\", \"text\": \"no renal auditory or vestibular toxicity was observed\", \"orig_text\": \"No renal, auditory, or vestibular toxicity was observed.\"}`.",
        "\n",
        "Note that TTS model may mispronounce some unknown words, for example, abbreviations like `tRNAs`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pR_T1HnttVjm"
      },
      "outputs": [],
      "source": [
        "lines = []\n",
        "with open(\"manifest.json\", \"r\", encoding=\"utf-8\") as f:\n",
        "    lines = f.readlines()\n",
        "\n",
        "for line in lines:\n",
        "    try:\n",
        "        data = json.loads(line.strip())\n",
        "    except:\n",
        "        print(line)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bt2TMLLvdUHm"
      },
      "source": [
        "Free GPU memory to avoid OOM."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZwEpAOCaRH7s"
      },
      "outputs": [],
      "source": [
        "del spectrogram_generator\n",
        "del vocoder\n",
        "torch.cuda.empty_cache()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HrensakWdLkt"
      },
      "source": [
        "## Run baseline ASR"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IQNIo2M_mqJc"
      },
      "source": [
        "Next we transcribe our .wav files with a general domain [ASR model](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/nemo/models/stt_en_conformer_ctc_large). It will generate an output file `ctc_baseline_transcript.json` where the predicted transcriptions are stored in the field `pred_text` of each record.\n",
        "\n",
        "Note that this ASR model was not trained or fine-tuned on medical domain, so we expect it to make mistakes on medical terms."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NMN63ux1mJiG"
      },
      "outputs": [],
      "source": [
        "!python nemo/examples/asr/transcribe_speech.py \\\n",
        "      pretrained_name=\"stt_en_conformer_ctc_large\" \\\n",
        "      dataset_manifest=manifest.json \\\n",
        "      output_filename=ctc_baseline_transcript_tmp.json \\\n",
        "      batch_size=2"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L3swQ8uqqgnp"
      },
      "source": [
        "ATTENTION: SpellMapper relies on words to be separated by _single_ space\n",
        "\n",
        "There is a bug with multiple space, observed in ASR results produced by Conformer-CTC, probably connected to this issue: https://github.com/NVIDIA/NeMo/issues/4034.\n",
        "\n",
        "So we need to correct the manifests to ensure that all spaces are single."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "z17sxkmXrXpJ"
      },
      "outputs": [],
      "source": [
        "test_data = read_manifest(\"ctc_baseline_transcript_tmp.json\")\n",
        "\n",
        "for i in range(len(test_data)):\n",
        "    # if there are multiple spaces in the string they will be merged to one\n",
        "    test_data[i][\"pred_text\"] = \" \".join(test_data[i][\"pred_text\"].split())\n",
        "\n",
        "with open(\"ctc_baseline_transcript.json\", \"w\", encoding=\"utf-8\") as out:\n",
        "    for d in test_data:\n",
        "        line = json.dumps(d)\n",
        "        out.write(line + \"\\n\")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PuKtfhbVkVJY"
      },
      "outputs": [],
      "source": [
        "!head -n 4 ctc_baseline_transcript.json"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aCJw9NEXqRg8"
      },
      "source": [
        "### Calculating WER of baseline transcript\n",
        "We use the standard script from NeMo to calculate WER and CER of our baseline transcript. Internally it compares the text in `pred_text` (predicted transcript) to `text` (reference transcript). "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZmNEGVWQsGo2"
      },
      "outputs": [],
      "source": [
        "!python nemo/examples/asr/speech_to_text_eval.py \\\n",
        "  dataset_manifest=ctc_baseline_transcript.json \\\n",
        "  only_score_manifest=True\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AvPwJr0ZqdkN"
      },
      "source": [
        "### See fragments that differ\n",
        "We use SequenceMatcher to see fragments that differ. (Another option is to use a more powerful analytics tool [Speech Data Explorer](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/tools/speech_data_explorer.html))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RAeaVCpMv78y"
      },
      "outputs": [],
      "source": [
        "test_data = read_manifest(\"ctc_baseline_transcript.json\")\n",
        "pred_text = [data['pred_text'] for data in test_data]\n",
        "ref_text = [data['text'] for data in test_data]\n",
        "audio_filepath = [data['audio_filepath'] for data in test_data]\n",
        "\n",
        "diff_vocab = Counter()\n",
        "\n",
        "for i in range(len(test_data)):\n",
        "    ref_sent = \" \" + ref_text[i] + \" \"\n",
        "    pred_sent = \" \" + pred_text[i] + \" \"\n",
        "\n",
        "    pred_words = pred_sent.strip().split()\n",
        "    ref_words = ref_sent.strip().split()\n",
        "\n",
        "    for tag, hyp_fragment, ref_fragment, i1, i2, j1, j2 in get_fragments(pred_words, ref_words):\n",
        "        if tag != \"equal\":\n",
        "            diff_vocab[(tag, hyp_fragment, ref_fragment)] += 1\n",
        "\n",
        "sum_ = 0\n",
        "print(\"PRED vs REF\")\n",
        "for k, v in diff_vocab.most_common(1000000):\n",
        "    sum_ += v\n",
        "    print(k, v, \"sum=\", sum_)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dUSOF7iD1w_9"
      },
      "source": [
        "## Run SpellMapper"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "x39BQhYB6_Fr"
      },
      "source": [
        "Now we run retrieval on our input manifest and prepare input for SpellMapper inference. Note that we use index of custom vocabulary (file `index.txt` that we saved earlier)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "y8x-yT5WqfFz"
      },
      "outputs": [],
      "source": [
        "!python nemo/examples/nlp/spellchecking_asr_customization/prepare_input_from_manifest.py \\\n",
        "  --manifest ctc_baseline_transcript.json \\\n",
        "  --custom_vocab_index index.txt \\\n",
        "  --big_sample spellmapper_asr_customization_en/big_sample.txt \\\n",
        "  --short2full_name short2full.txt \\\n",
        "  --output_name spellmapper_input.txt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ueq_JAPWGs_Y"
      },
      "source": [
        "Run the inference."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zgkqiiZtJjcB"
      },
      "outputs": [],
      "source": [
        "!python nemo/examples/nlp/spellchecking_asr_customization/spellchecking_asr_customization_infer.py \\\n",
        "      pretrained_model=spellmapper_asr_customization_en/training_10m_5ep.nemo \\\n",
        "      model.max_sequence_len=512 \\\n",
        "      inference.from_file=spellmapper_input.txt \\\n",
        "      inference.out_file=spellmapper_output.txt \\\n",
        "      inference.batch_size=16 \\\n",
        "      lang=en\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RPQWJX8dFLfX"
      },
      "source": [
        "Now we postprocess SpellMapper output and create output corrected manifest."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3eFU515yKvXP"
      },
      "outputs": [],
      "source": [
        "!python nemo/examples/nlp/spellchecking_asr_customization/postprocess_and_update_manifest.py \\\n",
        "  --input_manifest ctc_baseline_transcript.json \\\n",
        "  --short2full_name short2full.txt \\\n",
        "  --output_manifest ctc_corrected_transcript.json \\\n",
        "  --spellmapper_result spellmapper_output.txt \\\n",
        "  --replace_hyphen_to_space \\\n",
        "  --field_name pred_text \\\n",
        "  --ngram_mappings \"\"\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hRoIhhGh17tp"
      },
      "source": [
        "### Calculating WER of corrected transcript."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qIT957bGo9AY"
      },
      "outputs": [],
      "source": [
        "!python nemo/examples/asr/speech_to_text_eval.py \\\n",
        "  dataset_manifest=ctc_corrected_transcript.json \\\n",
        "  only_score_manifest=True\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NYXIPusupqOQ"
      },
      "outputs": [],
      "source": [
        "test_data = read_manifest(\"ctc_corrected_transcript.json\")\n",
        "pred_text = [data['pred_text'] for data in test_data]\n",
        "ref_text = [data['pred_text_before_correction'] for data in test_data]\n",
        "\n",
        "diff_vocab = Counter()\n",
        "\n",
        "for i in range(len(test_data)):\n",
        "    ref_sent = \" \" + ref_text[i] + \" \"\n",
        "    pred_sent = \" \" + pred_text[i] + \" \"\n",
        "\n",
        "    pred_words = pred_sent.strip().split()\n",
        "    ref_words = ref_sent.strip().split()\n",
        "\n",
        "    for tag, hyp_fragment, ref_fragment, i1, i2, j1, j2 in get_fragments(pred_words, ref_words):\n",
        "        if tag != \"equal\":\n",
        "            diff_vocab[(tag, hyp_fragment, ref_fragment)] += 1\n",
        "\n",
        "sum_ = 0\n",
        "print(\"Corrected vs baseline\")\n",
        "for k, v in diff_vocab.most_common(1000000):\n",
        "    sum_ += v\n",
        "    print(k, v, \"sum=\", sum_)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DJtXlqXbTD6M"
      },
      "source": [
        "### Filtering by Dynamic Programming(DP) score\n",
        "\n",
        "What else can be done?\n",
        "Given a fragment and its potential replacement, we can apply **dynamic programming** to find the most probable \"translation\" path between them. We will use the same n-gram mapping vocabulary, because its frequencies give us \"translation probability\" of each n-gram pair. The final path score can be calculated as maximum sum of log probabilities of matching n-grams along this path.\n",
        "Let's look at an example. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "05Qf9wgHU_UR"
      },
      "outputs": [],
      "source": [
        "joint_vocab, orig_vocab, misspelled_vocab, max_len = load_ngram_mappings_for_dp(\"spellmapper_asr_customization_en/replacement_vocab_filt.txt\")\n",
        "\n",
        "fragment = \"and hydrod\"\n",
        "replacement = \"anhydride\"\n",
        "fragment_spaced = \" \".join(list(fragment.replace(\" \", \"_\")))\n",
        "replacement_spaced = \" \".join(list(replacement.replace(\" \", \"_\")))\n",
        "path = get_alignment_by_dp(\n",
        "    replacement_spaced,\n",
        "    fragment_spaced,\n",
        "    dp_data=(joint_vocab, orig_vocab, misspelled_vocab, max_len)\n",
        ")\n",
        "print(\"Dynamic Programming path:\")\n",
        "for fragment_ngram, replacement_ngram, score, sum_score, joint_freq, orig_freq, misspelled_freq in path:\n",
        "    print(\n",
        "        \"\\t\",\n",
        "        \"frag=\",\n",
        "        fragment_ngram,\n",
        "        \"; repl=\",\n",
        "        replacement_ngram,\n",
        "        \"; score=\",\n",
        "        score,\n",
        "        \"; sum_score=\",\n",
        "        sum_score,\n",
        "        \"; joint_freq=\",\n",
        "        joint_freq,\n",
        "        \"; orig_freq=\",\n",
        "        orig_freq,\n",
        "        \"; misspelled_freq=\",\n",
        "        misspelled_freq,\n",
        "    )\n",
        "\n",
        "print(\"Final path score is in path[-1][3]: \", path[-1][3])\n",
        "print(\"Dynamic programming(DP) score per symbol is final score divided by len(fragment): \", path[-1][3] / (len(fragment)))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hgfKPKckaLnc"
      },
      "source": [
        "The idea is that we can skip replacements whose average DP score per symbol is below some predefined minimum, say -1.5.\n",
        "Note that dynamic programming works slow because of quadratic complexity, but it allows to get rid of some false positives. Let's apply it on the same test set."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UhSXh7ht_JRn"
      },
      "outputs": [],
      "source": [
        "!python nemo/examples/nlp/spellchecking_asr_customization/postprocess_and_update_manifest.py \\\n",
        "  --input_manifest ctc_baseline_transcript.json \\\n",
        "  --short2full_name short2full.txt \\\n",
        "  --output_manifest ctc_corrected_transcript_dp.json \\\n",
        "  --spellmapper_result spellmapper_output.txt \\\n",
        "  --replace_hyphen_to_space \\\n",
        "  --field_name pred_text \\\n",
        "  --use_dp \\\n",
        "  --ngram_mappings spellmapper_asr_customization_en/replacement_vocab_filt.txt \\\n",
        "  --min_dp_score_per_symbol -1.5"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "u8R5YHB3vPC8"
      },
      "outputs": [],
      "source": [
        "!python nemo/examples/asr/speech_to_text_eval.py \\\n",
        "  dataset_manifest=ctc_corrected_transcript_dp.json \\\n",
        "  only_score_manifest=True"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "upvTbkFAeYtR"
      },
      "source": [
        "# Final notes\n",
        "1. Bash-script with example of inference pipeline [run_infer.sh](https://github.com/NVIDIA/NeMo/blob/main/examples/nlp/spellchecking_asr_customization/run_infer.sh)\n",
        "\n",
        "2. Check our paper: [SpellMapper: A non-autoregressive neural spellchecker for ASR customization with candidate retrieval based on n-gram mappings](https://arxiv.org/abs/2306.02317)\n",
        "\n",
        "3. To reproduce evaluation experiments from this paper see these scripts:\n",
        " - [test_on_kensho.sh](https://github.com/bene-ges/nemo_compatible/blob/main/scripts/nlp/en_spellmapper/evaluation/test_on_kensho.sh)\n",
        " - [test_on_userlibri.sh](https://github.com/bene-ges/nemo_compatible/blob/main/scripts/nlp/en_spellmapper/evaluation/test_on_kensho.sh)\n",
        " - [test_on_spoken_wikipedia.sh](https://github.com/bene-ges/nemo_compatible/blob/main/scripts/nlp/en_spellmapper/evaluation/test_on_kensho.sh)\n",
        "\n",
        "4. To reproduce creation of training data see [README.md](https://github.com/bene-ges/nemo_compatible/blob/main/scripts/nlp/en_spellmapper/README.md)\n",
        "\n",
        "5. To run training see [run_training.sh](https://github.com/NVIDIA/NeMo/blob/main/examples/nlp/spellchecking_asr_customization/run_training.sh)\n",
        "\n",
        "6. Promising future research directions would be:\n",
        "  - add a simple trainable classifier on top of SpellMapper predictions instead of using multiple thresholds\n",
        "  - retrain with adding more various false positives to the training data"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}