week04/Seminar_04.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "## Dependencies\n",
        "\n",
        "approximately 5-6 minutes"
      ],
      "metadata": {
        "id": "z5a_3aA-9VB8"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6AGxp-OREthL"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "%%bash\n",
        "\n",
        "git clone https://github.com/NVIDIA/NeMo\n",
        "cd NeMo\n",
        "git checkout v1.20.0\n",
        "for f in $(ls requirements/requirements*.txt); do pip3 install --disable-pip-version-check --no-cache-dir -r $f; done\n",
        "pip install -e .\n",
        "pip install huggingface_hub==0.23.2\n",
        "pip install gdown\n",
        "pip install jiwer\n",
        "\n",
        "# restart session"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "\n",
        "---\n",
        "\n",
        "## Hybrid CTC & RNN-t model\n",
        "\n",
        "\n",
        "* Hybrid RNNT-CTC models is a group of models with both the RNNT and CTC decoders. Training a unified model would speedup the convergence for the CTC models and would enable the user to use a single model which works as both a CTC and RNNT model. This category can be used with any of the ASR models.[1]\n",
        "\n",
        "* So we can get speed from a CTC decoder and quality from a RNN-t decoder. This is extremely useful for production systems where you need to make partial predictions to show on screen while people are talking, and then make a final prediction. The first requests are usually handled by a fast CTC decoder, and the final prediction is done by RNN-t decoder.\n",
        "\n",
        "\\\n",
        "\n",
        "<img alt=\"hybrid\" src=\"https://drive.google.com/uc?id=1e8oe4CfBf8UmvWdm--DK_q126EK9A9tg\" width=400>\n",
        "\n",
        "\\\n",
        "\n",
        "* More about hybrid models:\n",
        "  * [[1] NeMo docs](https://docs.nvidia.com/nemo-framework/user-guide/latest/nemotoolkit/asr/models.html#hybrid-transducer-ctc)\n",
        "  * [[2] RNNT + LAS](https://arxiv.org/pdf/1908.10992)\n",
        "  * [[3] CTC + LAS](https://arxiv.org/pdf/1609.06773)\n",
        "  * [[4] Hybrid Rescoring 1](https://arxiv.org/pdf/2008.13093)\n",
        "  * [[5] Hybrid Rescoring 2](https://arxiv.org/pdf/2101.11577)\n"
      ],
      "metadata": {
        "id": "qGN9af_7wGKY"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "\n",
        "---\n",
        "\n"
      ],
      "metadata": {
        "id": "lLAKyHpMRoLn"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import re\n",
        "import typing as tp\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torchaudio\n",
        "import soundfile as sf\n",
        "from jiwer import wer\n",
        "from tqdm import tqdm\n",
        "import IPython.display as dsp\n",
        "from sentencepiece import SentencePieceProcessor\n",
        "\n",
        "from nemo.collections.asr.models import EncDecHybridRNNTCTCBPEModel\n",
        "\n",
        "BLANK_IND: int = 1024\n",
        "\n",
        "\n",
        "def clear(text: str):\n",
        "  return re.sub(r'[^A-Za-z +]', '', text.lower())"
      ],
      "metadata": {
        "id": "3znauxkwFPiM"
      },
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "\n",
        "model = EncDecHybridRNNTCTCBPEModel.from_pretrained(\n",
        "    model_name=\"stt_en_fastconformer_hybrid_large_pc\"\n",
        ").to(device).eval()\n",
        "\n",
        "TOKENIZER: SentencePieceProcessor = model.tokenizer.tokenizer\n",
        "\n",
        "dsp.clear_output()"
      ],
      "metadata": {
        "id": "ePqZU1OeGeHw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Audio"
      ],
      "metadata": {
        "id": "8uakB4rxP-da"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "! gdown https://drive.google.com/uc?id=1eNQt0R7Dm71utkLuRhc9wjTjdWoYrwsk -O best-song-ever.wav"
      ],
      "metadata": {
        "id": "_XGaz4jyeOZG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "dsp.display(dsp.Audio('best-song-ever.wav'))"
      ],
      "metadata": {
        "id": "AQQmAC5Ardqk"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "transcription = clear((\n",
        "    \"Never gonna give you up, never gonna let you down \"\n",
        "    \"Never gonna run around and desert you \"\n",
        "    \"Never gonna make you cry, never gonna say goodbye \"\n",
        "    \"Never gonna tell a lie and hurt you\"\n",
        "))\n",
        "transcription"
      ],
      "metadata": {
        "id": "e1ef738zybeF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### RNN-t inference"
      ],
      "metadata": {
        "id": "HKw7QkuFnzNC"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "**RNN-t** modules:\n",
        "\n",
        "<img alt=\"rnnt\" src=\"https://www.mdpi.com/symmetry/symmetry-11-01018/article_deploy/html/images/symmetry-11-01018-g004.png\" width=400>\n",
        "\n",
        "\n",
        "The encoder can be arbitrary, like RNN, DeepSpeech 2 encoder or Сonformer encoder, it can be streamable or non streamable, then whole model will be streamable or non streamable respectively.\n",
        "\n",
        "Inference stage looks like:\n",
        "\n",
        "<img alt=\"rnnt\" src=\"https://drive.google.com/uc?id=1EoSRLSSIg2fSge0yVKakKnbcgWUlCVJJ\" width=700>\n",
        "\n",
        "The prediction network consists of two required parts: embedder and RNN.\n",
        "\n",
        "<img alt=\"rnnt\" src=\"https://drive.google.com/uc?id=1SaMiv5F3bDRngNS6ot-TBi12Xo6IFOsX\" width=700>\n",
        "\n",
        "And the joint network can have arbitrary complexity and architecture, but in a simple case, it is a simple DNN.\n",
        "\n",
        "<img alt=\"rnnt\" src=\"https://drive.google.com/uc?id=11qccpDLBuAEXvsdkOIB9UZbVXqwD4zJC\" width=700>\n",
        "\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "neezeJDMIJkn"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Read wav\n",
        "wav, sr = torchaudio.load('best-song-ever.wav')\n",
        "wav = wav.to(device)\n",
        "assert sr == 16_000, sr\n",
        "\n",
        "# Get mel spectrogram\n",
        "input_signal_length = torch.tensor([wav.size(-1)], dtype=torch.int32, device=device)\n",
        "spectrogram, spec_length = model.preprocessor.forward(\n",
        "    input_signal=wav,\n",
        "    length=input_signal_length,\n",
        ")\n",
        "\n",
        "# Get encoded acoustic embeddings\n",
        "acoustic_embs, acoustic_embs_length = model.encoder.forward(\n",
        "    audio_signal=spectrogram, length=spec_length\n",
        ")"
      ],
      "metadata": {
        "id": "cD0g1vOFgqJv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "acoustic_embs.size()"
      ],
      "metadata": {
        "id": "YWtVnmbvR-QI"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### CTC Inference\n",
        "\n",
        "Let's use the `ctc_decode` function from the previous seminar and make a prediction by argmax."
      ],
      "metadata": {
        "id": "AOP9KlkF1H0G"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def ctc_decode(inds: list):\n",
        "    # your code here\n",
        "    raise NotImplementedError()"
      ],
      "metadata": {
        "id": "ofqQE_1C1RoF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "logits = model.ctc_decoder.forward(encoder_output=acoustic_embs)\n",
        "\n",
        "inds = logits.argmax(-1).tolist()[0]\n",
        "inds = ctc_decode(inds)\n",
        "\n",
        "ctc_hypothesis = model.tokenizer.tokenizer.decode_ids(inds)\n",
        "ctc_hypothesis = clear(ctc_hypothesis)\n",
        "ctc_hypothesis"
      ],
      "metadata": {
        "id": "I5vPy9m3rBkC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "wer(transcription, ctc_hypothesis)"
      ],
      "metadata": {
        "id": "ySbfl-oo19BR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### RNN-t inference\n",
        "\n",
        "Use `PredictionNetwork` and `JointNetwork` modules for RNN-t decoding. Sometimes it is useful to limit the number of tokens that will be emitted per frame, try to use this in your code with the `MAX_SYMBOLS_PER_FRAME: int` variable."
      ],
      "metadata": {
        "id": "EuiKbZoc19TN"
      }
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "KlDoE3uhZzGN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class PredictionNetwork(nn.Module):\n",
        "  def __init__(\n",
        "      self,\n",
        "      input_size: int,\n",
        "      hidden_size: int,\n",
        "      num_layers: int,\n",
        "      dropout: float,\n",
        "      num_embeddings: int,\n",
        "      embedding_dim: int,\n",
        "      padding_idx=None,\n",
        "    ):\n",
        "    super().__init__()\n",
        "\n",
        "    # https://pytorch.org/docs/stable/generated/torch.nn.Embedding.html\n",
        "    self.embed = nn.modules.sparse.Embedding(\n",
        "        num_embeddings=num_embeddings,\n",
        "        embedding_dim=embedding_dim,\n",
        "        padding_idx=padding_idx,\n",
        "    )\n",
        "    # https://pytorch.org/docs/stable/generated/torch.nn.LSTM.html\n",
        "    self.lstm = nn.LSTM(\n",
        "        input_size=input_size,\n",
        "        hidden_size=hidden_size,\n",
        "        num_layers=num_layers,\n",
        "        dropout=dropout,\n",
        "    )\n",
        "    self.dropout = nn.Dropout(dropout) if dropout else None\n",
        "\n",
        "  def forward(\n",
        "      self,\n",
        "      y: torch.tensor,\n",
        "      state: tp.Optional[tp.Tuple[torch.tensor, ...]] = None,\n",
        "    ) -> tp.Tuple[torch.tensor, tp.Tuple[torch.tensor, ...]]:\n",
        "    \"\"\"\n",
        "    input:\n",
        "      y_labels (bs, seq_len): ids from tokenizer of labels\n",
        "      state: lstm state, can be None in the first moment, (see torch docs)\n",
        "    output:\n",
        "      g (bs, seq_len, hid_dim): language context\n",
        "      state: lstm state\n",
        "    \"\"\"\n",
        "    # Get embeddings for labels\n",
        "    # your code goes here\n",
        "\n",
        "    # Proccess it with LSTM\n",
        "    # your code goes here\n",
        "    raise NotImplementedError()\n"
      ],
      "metadata": {
        "id": "m97sv-J1S-tU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "ZoC-Ui1lblZQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "prediction_network = PredictionNetwork(\n",
        "    input_size=640,\n",
        "    hidden_size=640,\n",
        "    num_layers=1,\n",
        "    dropout=0.2,\n",
        "    num_embeddings=1025,\n",
        "    embedding_dim=640,\n",
        "    padding_idx=BLANK_IND,\n",
        ").to(device).eval()\n",
        "\n",
        "prediction_network.embed.load_state_dict(\n",
        "    model.decoder.prediction.embed.state_dict()\n",
        ")\n",
        "prediction_network.lstm.load_state_dict(\n",
        "    model.decoder.prediction.dec_rnn.lstm.state_dict()\n",
        ")"
      ],
      "metadata": {
        "id": "lu2nsyMybgtv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "LqZBklVUkM_n"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class JointNetwork(nn.Module):\n",
        "  def __init__(\n",
        "      self,\n",
        "      pred_emb_size: int,\n",
        "      enc_emb_size: int,\n",
        "      hidden_size: int,\n",
        "      dropout: float,\n",
        "      vocab_size: int,\n",
        "    ):\n",
        "    super().__init__()\n",
        "\n",
        "    self.pred_proj = nn.Linear(\n",
        "        pred_emb_size, hidden_size\n",
        "    )\n",
        "    self.enc_proj = nn.Linear(\n",
        "        enc_emb_size, hidden_size\n",
        "    )\n",
        "    self.joint_net = nn.Sequential(\n",
        "        nn.ReLU(),\n",
        "        nn.Dropout(dropout),\n",
        "        nn.Linear(hidden_size, vocab_size + 1)\n",
        "    )\n",
        "\n",
        "  def forward(\n",
        "      self,\n",
        "      encoder_outputs: torch.tensor,\n",
        "      decoder_outputs: torch.tensor,\n",
        "    ) -> torch.Tensor:\n",
        "    \"\"\"\n",
        "    input:\n",
        "      encoder outputs (B, H1, T): acoustic context\n",
        "      decoder outputs (B, H2, U): language context\n",
        "    output:\n",
        "      joint activation (B, T, U, V+1)\n",
        "    \"\"\"\n",
        "    # Project the output of the encoder/decoder into the latent space and concatenate them\n",
        "    # your code goes here\n",
        "\n",
        "    # Project the following state into the vocab distribution space\n",
        "    # your code goes here\n",
        "    raise NotImplementedError()\n"
      ],
      "metadata": {
        "id": "M2pqe5WIS-vi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "joint_network = JointNetwork(\n",
        "    pred_emb_size=640,\n",
        "    enc_emb_size=512,\n",
        "    hidden_size=640,\n",
        "    dropout=0.2,\n",
        "    vocab_size=1024,\n",
        ").to(device).eval()\n",
        "\n",
        "joint_network.pred_proj.load_state_dict(\n",
        "    model.joint.pred.state_dict()\n",
        ")\n",
        "joint_network.enc_proj.load_state_dict(\n",
        "    model.joint.enc.state_dict()\n",
        ")\n",
        "joint_network.joint_net.load_state_dict(\n",
        "    model.joint.joint_net.state_dict()\n",
        ")"
      ],
      "metadata": {
        "id": "JQ8bRGTDTLZZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "OaBFJfMEh32c"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "write `rnnt_decoder_inference` function:\n",
        "\n",
        "<img alt=\"rnnt\" src=\"https://drive.google.com/uc?id=1EoSRLSSIg2fSge0yVKakKnbcgWUlCVJJ\" width=700>\n"
      ],
      "metadata": {
        "id": "2H9DyVN9pYBp"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "MAX_SYMBOLS_PER_FRAME: int = 100"
      ],
      "metadata": {
        "id": "6lOrJ8svwj_D"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "bZphhxkJzMlb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "@torch.inference_mode()\n",
        "def rnnt_decoder_inference(\n",
        "    prediction_network: nn.Module,\n",
        "    joint_network: nn.Module,\n",
        "    f: torch.Tensor,  # acoustic context\n",
        ") -> tp.List[int]:\n",
        "    \"\"\"\n",
        "    f - torch.tensor (B, H1, T): acoustic context\n",
        "    \"\"\"\n",
        "    bs, _, T = f.size()\n",
        "    assert bs == 1, bs\n",
        "\n",
        "    y_cur = torch.tensor([[BLANK_IND]], dtype=torch.long, device=device)\n",
        "    prediction_network_state = None\n",
        "\n",
        "    for time_step in tqdm(range(T)):\n",
        "        # your code here\n",
        "        pass\n",
        "\n",
        "    raise NotImplementedError()\n"
      ],
      "metadata": {
        "id": "brpo15JzMswo"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "decoded_outut = rnnt_decoder_inference(\n",
        "    prediction_network=prediction_network,\n",
        "    joint_network=joint_network,\n",
        "    f=acoustic_embs,\n",
        ")"
      ],
      "metadata": {
        "id": "Ilw5dI-ZMsyo"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "rnnt_hypothesis = clear(TOKENIZER.decode_ids(decoded_outut))\n",
        "rnnt_hypothesis"
      ],
      "metadata": {
        "id": "h7Ir2hClII_A"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "transcription"
      ],
      "metadata": {
        "id": "4ThGp_8--IF-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "wer(transcription, rnnt_hypothesis)"
      ],
      "metadata": {
        "id": "j1EH84Rp3Mgp"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### RNN-t training step"
      ],
      "metadata": {
        "id": "lUHg2ctd3NGf"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "transcription"
      ],
      "metadata": {
        "id": "rFmdqeCk3Mi7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "transcription_ids = TOKENIZER.encode(transcription)\n",
        "transcription_ids = torch.tensor(transcription_ids, dtype=torch.long, device=device).unsqueeze(0)\n",
        "transcription_ids"
      ],
      "metadata": {
        "id": "iOgNELeF3Pi0"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Read wav\n",
        "wav, sr = torchaudio.load('best-song-ever.wav')\n",
        "wav = wav.to(device)\n",
        "assert sr == 16_000, sr\n",
        "\n",
        "# Get mel spectrogram\n",
        "input_signal_length = torch.tensor([wav.size(-1)], dtype=torch.int32, device=device)\n",
        "spectrogram, spec_length = model.preprocessor.forward(\n",
        "    input_signal=wav,\n",
        "    length=input_signal_length,\n",
        ")\n",
        "\n",
        "# Get encoded acoustic embeddings\n",
        "acoustic_embs, acoustic_embs_length = model.encoder.forward(\n",
        "    audio_signal=spectrogram, length=spec_length\n",
        ")"
      ],
      "metadata": {
        "id": "IJ9O2Lj5GnDA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "cur_token_emb, hidden_state = prediction_network(\n",
        "    y=transcription_ids,\n",
        "    state=None,\n",
        ")\n",
        "inp, vocab_distributon = joint_network(\n",
        "    encoder_outputs=cur_token_emb,\n",
        "    decoder_outputs=acoustic_embs.mT,\n",
        "    return_inp=True,\n",
        ")"
      ],
      "metadata": {
        "id": "Jv_v_6G33Pk_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "inp.size()"
      ],
      "metadata": {
        "id": "rc3g-i0DII2-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "vocab_distributon.size()"
      ],
      "metadata": {
        "id": "WmrzGH3JFjRF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "FsG5wvtFYEKm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "\n",
        "---\n",
        "\n",
        "\n",
        "\n",
        "for further reading:\n",
        "  * [Sequence-to-sequence learning with Transducers](https://lorenlugosch.github.io/posts/2020/11/transducer/)\n",
        "  * RNN-t optimizations:\n",
        "    * [Multi-Blank Transducers for Speech Recognition, Hainan Xu et al., NVIDIA, 2024](https://arxiv.org/pdf/2211.03541v2)\n",
        "    * [Efficient Sequence Transduction by Jointly Predicting Tokens and Durations, Hainan Xu et al., NVIDIA, 2023](https://arxiv.org/abs/2304.06795)\n",
        "    * [FastEmit: Low-latency Streaming ASR with Sequence-level Emission Regularization, Jiahui Yu et al., Google, 2021](https://arxiv.org/abs/2010.11148)\n",
        "    * [Fast Conformer with Linearly Scalable Attention for Efficient Speech Recognition, Dima Rekesh et al., NVIDIA, 2023](https://arxiv.org/abs/2305.05084)\n",
        "    * [Rnn-Transducer with Stateless Prediction Network, Mohammadreza Ghodsi et al., 2020](https://ieeexplore.ieee.org/document/9054419)\n"
      ],
      "metadata": {
        "id": "ioHHAJ9TkKhR"
      }
    }
  ]
}