{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "EXA4lgxDIzwa"
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "You can run either this notebook locally (if you have all the dependencies and a GPU) or on Google Colab.\n",
    "\n",
    "Instructions for setting up Colab are as follows:\n",
    "1. Open a new Python 3 notebook.\n",
    "2. Import this notebook from GitHub (File -> Upload Notebook -> \"GITHUB\" tab -> copy/paste GitHub URL)\n",
    "3. Connect to an instance with a GPU (Runtime -> Change runtime type -> select \"GPU\" for hardware accelerator)\n",
    "4. Run this cell to set up dependencies.\n",
    "5. Restart the runtime (Runtime -> Restart Runtime) for any upgraded packages to take effect\n",
    "\"\"\"\n",
    "# If you're using Google Colab and not running locally, run this cell.\n",
    "\n",
    "## Install dependencies\n",
    "!pip install wget\n",
    "!apt-get install sox libsndfile1 ffmpeg\n",
    "!pip install text-unidecode\n",
    "!pip install matplotlib>=3.3.2\n",
    "\n",
    "## Install NeMo\n",
    "BRANCH = 'r1.17.0'\n",
    "!python -m pip install git+https://github.com/NVIDIA/NeMo.git@$BRANCH#egg=nemo_toolkit[all]\n",
    "\n",
    "# Update numba and restart (this is required to update internal numba version of Colab)\n",
    "\n",
    "# In a conda environment, you would use the following command\n",
    "# Update Numba to > 0.54\n",
    "# conda install -c conda-forge numba>=0.54\n",
    "# or\n",
    "# conda update -c conda-forge numba>=0.54\n",
    "\n",
    "# For pip based environments,\n",
    "# Update Numba to > 0.54\n",
    "import os\n",
    "import signal\n",
    "\n",
    "!pip install --upgrade numba\n",
    "\n",
    "# This will kill the kernel, click next cell to import the latest numba\n",
    "os.kill(os.getpid(), signal.SIGKILL)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_W1joUSQKvAd"
   },
   "source": [
    "# Buffered Transducer ASR\n",
    "\n",
    "There are many approaches to perform streaming/buffered inference for causal CTC / Transducer models. However, it is often observed that causal models sacrifice accuracy to perform streaming evaluation. \n",
    "\n",
    "In this notebook, similar to the CTC tutorial for [Streaming ASR](https://github.com/NVIDIA/NeMo/blob/stable/tutorials/asr/Streaming_ASR.ipynb), we will tackle the challenge of buffered ASR for long-form speech recognition, but this time we will use Transducer models as the basis for ASR. \n",
    "\n",
    "You may use this script [ASR Chunked Streaming Inference](https://github.com/NVIDIA/NeMo/blob/stable/examples/asr/asr_chunked_inference/rnnt/speech_to_text_buffered_infer_rnnt.py) to transcribe long audio files with Transducer models. \n",
    "\n",
    "**Note**: It is highly recommended to review the ``Streaming ASR`` tutorial for a good overview of how streaming/buffered inference works for CTC models and the underlying motivation of streaming ASR itself.\n",
    "\n",
    "------"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nJadC6xmQydl"
   },
   "source": [
    "Transducers surpass CTC models in speech recognition accuracy when greedy decoding with no LM is used. While CTC models can give better accuracy with beam search decoding and LM, large external language models are required to reach or surpass the accuracy of transducers with greedy decoding.\n",
    "\n",
    "Moreover, the challenging autoregressive strategy of transducer decoding imposes particular challenges, which we will tackle as a topic in this tutorial."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ymU6NlnVLMji"
   },
   "source": [
    "# Prepare the dataset\n",
    "\n",
    "We will continue to use the Librispeech dev-clean subset of [Mini Librispeech](https://www.openslr.org/31/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Bgck1OSaQ63d"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "if not os.path.exists(\"scripts/get_librispeech_data.py\"):\n",
    "  !wget -P scripts/ https://raw.githubusercontent.com/NVIDIA/NeMo/stable/scripts/dataset_processing/get_librispeech_data.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kZB99Ul0LJF3"
   },
   "outputs": [],
   "source": [
    "# If something goes wrong during data processing, un-comment the following line to delete the cached dataset \n",
    "# !rm -rf datasets/mini-dev-clean\n",
    "!mkdir -p datasets/mini-dev-clean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zCCC_ssiLTeK"
   },
   "outputs": [],
   "source": [
    "!python scripts/get_librispeech_data.py \\\n",
    "  --data_root \"datasets/mini-dev-clean/\" \\\n",
    "  --data_sets dev_clean_2 \\\n",
    "  --num_workers=10 \\\n",
    "  --log"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vssfT5M1LfPz"
   },
   "outputs": [],
   "source": [
    "manifest = os.path.join(os.getcwd(), \"datasets/mini-dev-clean/dev_clean_2.json\")\n",
    "print(\"Manifest path :\", manifest)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3GcHmUraL8RW"
   },
   "source": [
    "Let's create a long audio that is about 15 minutes long by concatenating audio from dev-clean and also create the corresponding concatenated transcript."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lt1Ne9SlL4h9"
   },
   "outputs": [],
   "source": [
    "import json\n",
    "def concat_audio(manifest_file, final_len=3600):\n",
    "    concat_len = 0\n",
    "    final_transcript = \"\"\n",
    "    with open(\"concat_file.txt\", \"w\") as cat_f:\n",
    "        while concat_len < final_len:\n",
    "            with open(manifest_file, \"r\") as mfst_f:\n",
    "                for l in mfst_f:\n",
    "                    row = json.loads(l.strip())\n",
    "                    if concat_len >= final_len:\n",
    "                        break\n",
    "                    cat_f.write(f\"file {row['audio_filepath']}\\n\")\n",
    "                    final_transcript += (\" \" + row['text'])\n",
    "                    concat_len += float(row['duration'])\n",
    "    return concat_len, final_transcript"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-G5l7mZ8L_2X"
   },
   "outputs": [],
   "source": [
    "new_duration, ref_transcript = concat_audio(manifest, 15*60)\n",
    "\n",
    "concat_audio_path = os.path.join(os.getcwd(), \"datasets/mini-dev-clean/concatenated_audio.wav\")\n",
    "\n",
    "!ffmpeg -t {new_duration} -safe 0 -f concat -i concat_file.txt -c copy -t {new_duration} {concat_audio_path} -y\n",
    "print(\"Finished concatenating audio file!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RvJBpVdsM9wy"
   },
   "source": [
    "# Buffered Transducer\n",
    "\n",
    "We will now prepare a Conformer Transducer model to set the stage for buffered inference. Conformers possess self-attention layers, which require quadratic cost in terms of memory and compute for a given audio sequence length. Self-attention naturally imposes a limit of 2-5 minute long audio clips, even on 32 GB of GPU memory. Therefore buffered inference is a prime candidate to resolve the issue of Conformer memory consumption.\n",
    "\n",
    "**Note**: While we primarily discuss buffered ASR here, the primary difference between buffered and streaming ASR is the size of the chunk (which determines the latency of prediction). Many of the techniques here can be tested with smaller chunk and buffer sizes, therefore significantly improving latency and approach  \"streaming\" mode inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "b-Szou9bMCwa"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import nemo.collections.asr as nemo_asr\n",
    "import contextlib\n",
    "import gc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "anMwz7MdRoYp"
   },
   "outputs": [],
   "source": [
    "pretrained_model_name = \"stt_en_conformer_transducer_large\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "W9KbfhycNCPf"
   },
   "outputs": [],
   "source": [
    "# Clear up memory\n",
    "torch.cuda.empty_cache()\n",
    "gc.collect()\n",
    "model = nemo_asr.models.EncDecRNNTBPEModel.from_pretrained(pretrained_model_name)\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "# device = 'cpu'  # You can transcribe even longer samples on the CPU, though it will take much longer !\n",
    "model = model.to(device)\n",
    "model.freeze()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "J79e_RSQQ7m1"
   },
   "outputs": [],
   "source": [
    "# Helper for torch amp autocast\n",
    "if torch.cuda.is_available():\n",
    "    autocast = torch.cuda.amp.autocast\n",
    "else:\n",
    "    @contextlib.contextmanager\n",
    "    def autocast():\n",
    "        print(\"AMP was not available, using FP32!\")\n",
    "        yield"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xX-w1eFuQxJp"
   },
   "source": [
    "The call to transcribe() below should fail with a \"CUDA out of memory\" error when run on a GPU with 32 GB memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3zdEXln4NKj0"
   },
   "outputs": [],
   "source": [
    "with autocast():\n",
    "  if torch.cuda.is_available():\n",
    "    transcript = model.transcribe([concat_audio_path], batch_size=1)[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xDWxuk65RJV9"
   },
   "source": [
    "## Offline Baseline\n",
    "\n",
    "Let us check the offline score of this model (on the individual segmented audio files) so that we have a baseline. This will evaluate if the buffered inference significantly sacrifices recognition accuracy.\n",
    "\n",
    "Note that it is often the case that such clean audio segments will not be available (unless it is a preprocessed dataset). Still, we are lucky since we are using Librispeech, which has been nearly perfectly segmented for our use case. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4a2U660DRCpl"
   },
   "source": [
    "------\n",
    "\n",
    "Let's download some scripts from the NeMo repo to easily score our model on this dataset in an offline manner.\n",
    "\n",
    "**Note**: It may take a few minutes to transcribe all the files due to network I/O on Colab. You may choose to uncomment and run the offline evaluation or continue on to the next cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zD2_otE_Q4AZ"
   },
   "outputs": [],
   "source": [
    "if not os.path.exists(\"scripts/transcribe_speech.py\"):\n",
    "  !wget -P scripts/ https://raw.githubusercontent.com/NVIDIA/NeMo/stable/examples/asr/transcribe_speech.py\n",
    "\n",
    "if not os.path.exists(\"scripts/speech_to_text_eval.py\"):\n",
    "  !wget -P scripts/ https://raw.githubusercontent.com/NVIDIA/NeMo/stable/examples/asr/speech_to_text_eval.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7TDTO3KLRfYb"
   },
   "outputs": [],
   "source": [
    "# Uncomment if you want to run the evaluation in offline mode\n",
    "# if torch.cuda.is_available():\n",
    "#   !python scripts/speech_to_text_eval.py \\\n",
    "#     pretrained_name={pretrained_model_name} \\\n",
    "#     dataset_manifest={manifest} \\\n",
    "#     batch_size=32 \\\n",
    "#     amp=True \\\n",
    "#     use_cer=False\n",
    "# else:\n",
    "#   print(\"CUDA not available, decoding full dataset would take too long.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "76uo4Bg-UA2v"
   },
   "outputs": [],
   "source": [
    "# Clear up memory\n",
    "torch.cuda.empty_cache()\n",
    "gc.collect()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DX-Arg9kUEfe"
   },
   "source": [
    "## Buffer mechanism for streaming long audio files\n",
    "\n",
    "As you will note below, audio chunking and buffering are identical steps for CTC and Transducer models. As such, we will perform the setup steps in the next cell without significant elaboration.\n",
    "\n",
    "**Note**: For detailed information on how audio is chunked and evaluated, you should refer to the ``Streaming ASR`` tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "X2tKBEXIRzMl"
   },
   "outputs": [],
   "source": [
    "#@title Setup Audio Chunk Iterator\n",
    "# A simple iterator class to return successive chunks of samples\n",
    "class AudioChunkIterator():\n",
    "    def __init__(self, samples, frame_len, sample_rate):\n",
    "        self._samples = samples\n",
    "        self._chunk_len = chunk_len_in_secs*sample_rate\n",
    "        self._start = 0\n",
    "        self.output=True\n",
    "   \n",
    "    def __iter__(self):\n",
    "        return self\n",
    "    \n",
    "    def __next__(self):\n",
    "        if not self.output:\n",
    "            raise StopIteration\n",
    "        last = int(self._start + self._chunk_len)\n",
    "        if last <= len(self._samples):\n",
    "            chunk = self._samples[self._start: last]\n",
    "            self._start = last\n",
    "        else:\n",
    "            chunk = np.zeros([int(self._chunk_len)], dtype='float32')\n",
    "            samp_len = len(self._samples) - self._start\n",
    "            chunk[0:samp_len] = self._samples[self._start:len(self._samples)]\n",
    "            self.output = False\n",
    "   \n",
    "        return chunk\n",
    "\n",
    "# a helper function for extracting samples as a numpy array from the audio file\n",
    "import soundfile as sf\n",
    "def get_samples(audio_file, target_sr=16000):\n",
    "    with sf.SoundFile(audio_file, 'r') as f:\n",
    "        sample_rate = f.samplerate\n",
    "        samples = f.read()\n",
    "        if sample_rate != target_sr:\n",
    "            samples = librosa.core.resample(samples, orig_sr=sample_rate, target_sr=target_sr)\n",
    "        samples = samples.transpose()\n",
    "        return samples\n",
    "  \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "l8j_F3GLUPrV"
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from IPython.display import display, Audio\n",
    "import numpy as np\n",
    "\n",
    "samples = get_samples(concat_audio_path)\n",
    "sample_rate  = model.preprocessor._cfg['sample_rate'] "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hrTq2XUcWBLo"
   },
   "source": [
    "## Batched Chunk Processor\n",
    "\n",
    "First, we write a batched variant of the ``FeatureFrameBufferer`` that was written implicitly as part of the ``Streaming ASR`` tutorial.\n",
    "\n",
    "The difference between the two versions is - the ``FeatureFrameBufferer`` will buffer across frames of a single sample and then loop for each sample in the dataset. The `BatchedFeatureFrameBufferer` will buffer across the dependent frames of the independent batch of samples. This significantly improves the efficiency of buffered transducer inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "p8r96TdbVYBO"
   },
   "outputs": [],
   "source": [
    "#@title Setup batched feature chunk bufferer\n",
    "\n",
    "from nemo.collections.asr.parts.utils import streaming_utils\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "class BatchedFeatureFrameBufferer(streaming_utils.BatchedFeatureFrameBufferer):\n",
    "    \"\"\"\n",
    "    Batched variant of FeatureFrameBufferer where batch dimension is the independent audio samples.\n",
    "    \"\"\"\n",
    "\n",
    "    def reset(self):\n",
    "        '''\n",
    "        Reset frame_history and decoder's state\n",
    "        '''\n",
    "        super().reset()\n",
    "        self.limit_frames = [None for _ in range(self.batch_size)]\n",
    "\n",
    "    def get_batch_frames(self):\n",
    "        # Exit if all buffers of all samples have been processed\n",
    "        if all(self.signal_end):\n",
    "            return []\n",
    "\n",
    "        # Otherwise sequentially process frames of each sample one by one.\n",
    "        batch_frames = []\n",
    "        for idx, frame_reader in enumerate(self.all_frame_reader):\n",
    "\n",
    "            limit_frames = self.limit_frames[idx]\n",
    "            try:\n",
    "                if limit_frames is not None and self.buffer_number >= limit_frames:\n",
    "                  raise StopIteration()\n",
    "\n",
    "                frame = next(frame_reader)\n",
    "                frame = np.copy(frame)\n",
    "\n",
    "                batch_frames.append(frame)\n",
    "            except StopIteration:\n",
    "                # If this sample has finished all of its buffers\n",
    "                # Set its signal_end flag, and assign it the id of which buffer index\n",
    "                # did it finish the sample (if not previously set)\n",
    "                # This will let the alignment module know which sample in the batch finished\n",
    "                # at which index.\n",
    "                batch_frames.append(None)\n",
    "                self.signal_end[idx] = True\n",
    "\n",
    "                if self.signal_end_index[idx] is None:\n",
    "                    self.signal_end_index[idx] = self.buffer_number\n",
    "\n",
    "        self.buffer_number += 1\n",
    "        return batch_frames\n",
    "\n",
    "    def set_frame_reader(self, frame_reader, idx, limit_frames=None):\n",
    "        self.all_frame_reader[idx] = frame_reader\n",
    "        self.signal_end[idx] = False\n",
    "        self.signal_end_index[idx] = None\n",
    "        self.limit_frames[idx] = limit_frames"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "YEJEaKBQXlHG"
   },
   "source": [
    "## Batched and Buffered ASR Transducer\n",
    "\n",
    "Next, we will build the actual buffered transducer evaluation class. \n",
    "\n",
    "\n",
    "Similar to Streaming CTC models, we pick tokens corresponding to one chunk length of audio for each buffer. The chunk within each buffer is chosen such that there is equal left and right context available to the audio within the chunk.\n",
    "\n",
    "\n",
    "Since this is a batched variant of the ``Streaming ASR`` tutorial, we will subclass the required method and override the parts that we need to support batching across independent samples and buffering across dependent frames per sample in the batch.\n",
    "\n",
    "----\n",
    "\n",
    "Due to the complexity of the code, we will hide the cell below then explain the essential sections of the code as sub-sections. If at any point you would like to review the code itself, click `Show code` below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "0z_Iff9HXjCv"
   },
   "outputs": [],
   "source": [
    "#@title Setup batched and buffered transducer ASR code \n",
    "\n",
    "import librosa\n",
    "\n",
    "def inplace_buffer_merge(buffer, data, timesteps, model):\n",
    "    \"\"\"\n",
    "    Merges the new text from the current frame with the previous text contained in the buffer.\n",
    "\n",
    "    The alignment is based on a Longest Common Subsequence algorithm, with some additional heuristics leveraging\n",
    "    the notion that the chunk size is >= the context window. In case this assumptio is violated, the results of the merge\n",
    "    will be incorrect (or at least obtain worse WER overall).\n",
    "    \"\"\"\n",
    "    # If delay timesteps is 0, that means no future context was used. Simply concatenate the buffer with new data.\n",
    "    if timesteps < 1:\n",
    "        buffer += data\n",
    "        return buffer\n",
    "\n",
    "    # If buffer is empty, simply concatenate the buffer and data.\n",
    "    if len(buffer) == 0:\n",
    "        buffer += data\n",
    "        return buffer\n",
    "\n",
    "    # Concat data to buffer\n",
    "    buffer += data\n",
    "    return buffer\n",
    "\n",
    "\n",
    "class BatchedFrameASRRNNT(streaming_utils.FrameBatchASR):\n",
    "    \"\"\"\n",
    "    Batched implementation of FrameBatchASR for RNNT models, where the batch dimension is independent audio samples.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, asr_model, frame_len=1.6, total_buffer=4.0,\n",
    "        batch_size=32, max_steps_per_timestep: int = 5, stateful_decoding: bool = False):\n",
    "        '''\n",
    "        Args:\n",
    "            asr_model: An RNNT model.\n",
    "            frame_len: frame's duration, seconds.\n",
    "            total_buffer: duration of total audio chunk size, in seconds.\n",
    "            batch_size: Number of independent audio samples to process at each step.\n",
    "            max_steps_per_timestep: Maximum number of tokens (u) to process per acoustic timestep (t).\n",
    "            stateful_decoding: Boolean whether to enable stateful decoding for preservation of state across buffers.\n",
    "        '''\n",
    "        super().__init__(asr_model, frame_len=frame_len, total_buffer=total_buffer, batch_size=batch_size)\n",
    "\n",
    "        # OVERRIDES OF THE BASE CLASS\n",
    "        self.max_steps_per_timestep = max_steps_per_timestep\n",
    "        self.stateful_decoding = stateful_decoding\n",
    "\n",
    "        self.all_alignments = [[] for _ in range(self.batch_size)]\n",
    "        self.all_preds = [[] for _ in range(self.batch_size)]\n",
    "        self.previous_hypotheses = None\n",
    "        self.batch_index_map = {\n",
    "            idx: idx for idx in range(self.batch_size)\n",
    "        }  # pointer from global batch id : local sub-batch id\n",
    "\n",
    "        try:\n",
    "            self.eos_id = self.asr_model.tokenizer.eos_id\n",
    "        except Exception:\n",
    "            self.eos_id = -1\n",
    "\n",
    "        print(\"Performing Stateful decoding :\", self.stateful_decoding)\n",
    "\n",
    "        # OVERRIDES\n",
    "        self.frame_bufferer = BatchedFeatureFrameBufferer(\n",
    "            asr_model=asr_model, frame_len=frame_len, batch_size=batch_size, total_buffer=total_buffer\n",
    "        )\n",
    "\n",
    "        self.reset()\n",
    "\n",
    "    def reset(self):\n",
    "        \"\"\"\n",
    "        Reset frame_history and decoder's state\n",
    "        \"\"\"\n",
    "        super().reset()\n",
    "\n",
    "        self.all_alignments = [[] for _ in range(self.batch_size)]\n",
    "        self.all_preds = [[] for _ in range(self.batch_size)]\n",
    "        self.previous_hypotheses = None\n",
    "        self.batch_index_map = {idx: idx for idx in range(self.batch_size)}\n",
    "\n",
    "        self.data_layer = [streaming_utils.AudioBuffersDataLayer() for _ in range(self.batch_size)]\n",
    "        self.data_loader = [\n",
    "            DataLoader(self.data_layer[idx], batch_size=1, collate_fn=streaming_utils.speech_collate_fn)\n",
    "            for idx in range(self.batch_size)\n",
    "        ]\n",
    "\n",
    "        self.buffers = []\n",
    "\n",
    "    def read_audio_file(self, audio_filepath: list, delay, model_stride_in_secs):\n",
    "        assert len(audio_filepath) == self.batch_size\n",
    "\n",
    "        # Read in a batch of audio files, one by one\n",
    "        for idx in range(self.batch_size):\n",
    "            samples = get_samples(audio_filepath[idx])\n",
    "            samples = np.pad(samples, (0, int(delay * model_stride_in_secs * self.asr_model._cfg.sample_rate)))\n",
    "            frame_reader = streaming_utils.AudioFeatureIterator(samples, self.frame_len, self.raw_preprocessor, self.asr_model.device)\n",
    "            self.set_frame_reader(frame_reader, idx)\n",
    "\n",
    "    def set_frame_reader(self, frame_reader, idx, limit_frames = None):\n",
    "        self.frame_bufferer.set_frame_reader(frame_reader, idx, limit_frames)\n",
    "\n",
    "    @torch.no_grad()\n",
    "    def infer_logits(self):\n",
    "        frame_buffers = self.frame_bufferer.get_buffers_batch()\n",
    "\n",
    "        while len(frame_buffers) > 0:\n",
    "            # While at least 1 sample has a buffer left to process\n",
    "            self.frame_buffers += frame_buffers[:]\n",
    "\n",
    "            for idx, buffer in enumerate(frame_buffers):\n",
    "                if self.plot:\n",
    "                  self.buffers.append(buffer[:][0])\n",
    "                self.data_layer[idx].set_signal(buffer[:])\n",
    "\n",
    "            self._get_batch_preds()\n",
    "            frame_buffers = self.frame_bufferer.get_buffers_batch()\n",
    "\n",
    "    @torch.no_grad()\n",
    "    def _get_batch_preds(self):\n",
    "        \"\"\"\n",
    "        Perform dynamic batch size decoding of frame buffers of all samples.\n",
    "\n",
    "        Steps:\n",
    "            -   Load all data loaders of every sample\n",
    "            -   For all samples, determine if signal has finished.\n",
    "                -   If so, skip calculation of mel-specs.\n",
    "                -   If not, compute mel spec and length\n",
    "            -   Perform Encoder forward over this sub-batch of samples. Maintain the indices of samples that were processed.\n",
    "            -   If performing stateful decoding, prior to decoder forward, remove the states of samples that were not processed.\n",
    "            -   Perform Decoder + Joint forward for samples that were processed.\n",
    "            -   For all output RNNT alignment matrix of the joint do:\n",
    "                -   If signal has ended previously (this was last buffer of padding), skip alignment\n",
    "                -   Otherwise, recalculate global index of this sample from the sub-batch index, and preserve alignment.\n",
    "            -   Same for preds\n",
    "            -   Update indices of sub-batch with global index map.\n",
    "            - Redo steps until all samples were processed (sub-batch size == 0).\n",
    "        \"\"\"\n",
    "        device = self.asr_model.device\n",
    "\n",
    "        data_iters = [iter(data_loader) for data_loader in self.data_loader]\n",
    "\n",
    "        feat_signals = []\n",
    "        feat_signal_lens = []\n",
    "\n",
    "        new_batch_keys = []\n",
    "        for idx in range(self.batch_size):\n",
    "            if self.frame_bufferer.signal_end[idx]:\n",
    "                continue\n",
    "\n",
    "            batch = next(data_iters[idx])\n",
    "            feat_signal, feat_signal_len = batch\n",
    "            feat_signal, feat_signal_len = feat_signal.to(device), feat_signal_len.to(device)\n",
    "\n",
    "            feat_signals.append(feat_signal)\n",
    "            feat_signal_lens.append(feat_signal_len)\n",
    "\n",
    "            # preserve batch indices\n",
    "            new_batch_keys.append(idx)\n",
    "\n",
    "        if len(feat_signals) == 0:\n",
    "            return\n",
    "\n",
    "        feat_signal = torch.cat(feat_signals, 0)\n",
    "        feat_signal_len = torch.cat(feat_signal_lens, 0)\n",
    "\n",
    "        del feat_signals, feat_signal_lens\n",
    "\n",
    "        encoded, encoded_len = self.asr_model(processed_signal=feat_signal, processed_signal_length=feat_signal_len)\n",
    "\n",
    "        # filter out partial hypotheses from older batch subset\n",
    "        if self.stateful_decoding and self.previous_hypotheses is not None:\n",
    "            new_prev_hypothesis = []\n",
    "            for new_batch_idx, global_index_key in enumerate(new_batch_keys):\n",
    "                old_pos = self.batch_index_map[global_index_key]\n",
    "                new_prev_hypothesis.append(self.previous_hypotheses[old_pos])\n",
    "            self.previous_hypotheses = new_prev_hypothesis\n",
    "\n",
    "        best_hyp, _ = self.asr_model.decoding.rnnt_decoder_predictions_tensor(\n",
    "            encoded, encoded_len, return_hypotheses=True, partial_hypotheses=self.previous_hypotheses\n",
    "        )\n",
    "\n",
    "        if self.stateful_decoding:\n",
    "            # preserve last state from hypothesis of new batch indices\n",
    "            self.previous_hypotheses = best_hyp\n",
    "\n",
    "        for idx, hyp in enumerate(best_hyp):\n",
    "            global_index_key = new_batch_keys[idx]  # get index of this sample in the global batch\n",
    "\n",
    "            has_signal_ended = self.frame_bufferer.signal_end[global_index_key]\n",
    "\n",
    "            if not has_signal_ended:\n",
    "                self.all_alignments[global_index_key].append(hyp.alignments)\n",
    "\n",
    "        preds = [hyp.y_sequence for hyp in best_hyp]\n",
    "        for idx, pred in enumerate(preds):\n",
    "            global_index_key = new_batch_keys[idx]  # get index of this sample in the global batch\n",
    "\n",
    "            has_signal_ended = self.frame_bufferer.signal_end[global_index_key]\n",
    "            if not has_signal_ended:\n",
    "                self.all_preds[global_index_key].append(pred.cpu().numpy())\n",
    "\n",
    "        if self.stateful_decoding:\n",
    "            # State resetting is being done on sub-batch only, global index information is not being updated\n",
    "            reset_states = self.asr_model.decoder.initialize_state(encoded)\n",
    "\n",
    "            for idx, pred in enumerate(preds):\n",
    "                if len(pred) > 0 and pred[-1] == self.eos_id:\n",
    "                    # reset states :\n",
    "                    self.previous_hypotheses[idx].y_sequence = self.previous_hypotheses[idx].y_sequence[:-1]\n",
    "                    self.previous_hypotheses[idx].dec_state = self.asr_model.decoder.batch_select_state(\n",
    "                        reset_states, idx\n",
    "                    )\n",
    "\n",
    "        # Position map update\n",
    "        if len(new_batch_keys) != len(self.batch_index_map):\n",
    "            for new_batch_idx, global_index_key in enumerate(new_batch_keys):\n",
    "                self.batch_index_map[global_index_key] = new_batch_idx  # let index point from global pos -> local pos\n",
    "\n",
    "        del encoded, encoded_len\n",
    "        del best_hyp, pred\n",
    "\n",
    "    def transcribe(\n",
    "        self, tokens_per_chunk: int, delay: int, plot=False,\n",
    "    ):\n",
    "        \"\"\"\n",
    "        Performs \"middle token\" alignment prediction using the buffered audio chunk.\n",
    "        \"\"\"\n",
    "        self.plot = plot\n",
    "        self.infer_logits()\n",
    "\n",
    "        self.unmerged = [[] for _ in range(self.batch_size)]\n",
    "        for idx, alignments in enumerate(self.all_alignments):\n",
    "\n",
    "            signal_end_idx = self.frame_bufferer.signal_end_index[idx]\n",
    "            if signal_end_idx is None:\n",
    "                raise ValueError(\"Signal did not end\")\n",
    "\n",
    "            all_toks = []\n",
    "\n",
    "            for a_idx, alignment in enumerate(alignments):\n",
    "                alignment = alignment[len(alignment) - 1 - delay : len(alignment) - 1 - delay + tokens_per_chunk]\n",
    "\n",
    "                ids, toks = self._alignment_decoder(alignment, self.asr_model.tokenizer, self.blank_id)\n",
    "                all_toks.append(toks)\n",
    "\n",
    "                if len(ids) > 0 and a_idx < signal_end_idx:\n",
    "                    self.unmerged[idx] = inplace_buffer_merge(self.unmerged[idx], ids, delay, model=self.asr_model,)\n",
    "\n",
    "            if plot:\n",
    "              for i, tok in enumerate(all_toks):\n",
    "                  print(\"\\nGreedy labels collected from this buffer\")\n",
    "                  print(tok[len(tok) - 1 - delay:len(tok) - 1 - delay + tokens_per_chunk])                \n",
    "                  self.toks_unmerged += tok[len(tok) - 1 - delay:len(tok) - 1 - delay + tokens_per_chunk]\n",
    "              print(\"\\nTokens collected from succesive buffers before RNNT merge\")\n",
    "              print(self.toks_unmerged)\n",
    "\n",
    "        output = []\n",
    "        for idx in range(self.batch_size):\n",
    "            output.append(self.greedy_merge(self.unmerged[idx]))\n",
    "        return output\n",
    "\n",
    "    def _alignment_decoder(self, alignments, tokenizer, blank_id):\n",
    "        s = []\n",
    "        ids = []\n",
    "\n",
    "        for t in range(len(alignments)):\n",
    "            for u in range(len(alignments[t])):\n",
    "                token_id = int(alignments[t][u][1])\n",
    "                if token_id != blank_id:\n",
    "                    token = tokenizer.ids_to_tokens([token_id])[0]\n",
    "                    s.append(token)\n",
    "                    ids.append(token_id)\n",
    "\n",
    "                else:\n",
    "                    # blank token\n",
    "                    pass\n",
    "\n",
    "        return ids, s\n",
    "\n",
    "    def greedy_merge(self, preds):\n",
    "        decoded_prediction = [p for p in preds]\n",
    "        hypothesis = self.asr_model.tokenizer.ids_to_text(decoded_prediction)\n",
    "        return hypothesis"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1xB919nvhcet"
   },
   "source": [
    "## Code Breakdown\n",
    "\n",
    "The following section is optional and describes the sub-sections of the code snippet above. It can improve understanding of how the code above works."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2T0RVdzSXYF7"
   },
   "source": [
    "### Code: `__init__`\n",
    "\n",
    "Transducers will operate on a batch of samples at once and then process the chunks of each of these samples independently with a single forward pass of the Encoder and then multiple autoregressive calls to the Prediction Network + Joint Network."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "U42JNBycYNg-"
   },
   "source": [
    "### Code: `_alignment_decoder(alignments, tokenizer, blank_id)`\n",
    "\n",
    "Since the models we are evaluating are trained with sub-word encoding, we will need to decode the tokens to a text format from the 2-dimensional dangling array, which represents the alignments of the transducer's prediction.\n",
    "\n",
    "**Note**: The alignment is a 2-dimensional dangling array with the shape `Ti x Uj`; there can be any number of `Uj` per `Ti`. The alignment also contains the id for the `Transducer Blank` token - which we need to remove during decoding to prevent the tokenizer from trying to decode an invalid id. An example of a transducer alignment will be presented at the end of the notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Lg_2yywrZWry"
   },
   "source": [
    "### Code: `_get_batch_preds()`\n",
    "\n",
    "The core of the transducer model's decoding step per chunk of provided audio for all independent audio samples. We batch together the independent acoustic segments through the encoder and then batch process the prediction net + joint net to improve the GPU efficiency of decoding.\n",
    "\n",
    "To further improve efficiency, we will perform adaptive batching during evaluation. In adaptive batching, once a sample has finished processing its audio sequence, it will be removed from the global set of all samples that should be processed. After each chunk is processed of each sample, the completed samples are removed from the next round of decoding.\n",
    "\n",
    "------\n",
    "\n",
    "Due to this additional complexity, we break down a few steps of this process below - "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EAa0Ko7aa8jx"
   },
   "source": [
    "#### Select the subset of samples that need to finish processing\n",
    "\n",
    "We will loop through all samples, checking if the sample has finished processing or not. If not, it will be added to the pool of samples that must be processed. These samples are passed through the encoder.\n",
    "\n",
    "```python\n",
    "new_batch_keys = []\n",
    "for idx in range(self.batch_size):\n",
    "    if self.frame_bufferer.signal_end[idx]:\n",
    "        continue\n",
    "    batch = next(data_iters[idx])\n",
    "    ...\n",
    "    new_batch_keys.append(idx)\n",
    "encoded, encoded_len = self.asr_model(processed_signal=feat_signal, processed_signal_length=feat_signal_len)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-QzUgpIsaxCt"
   },
   "source": [
    "#### Update partial states of the model and decode the prediction + joint steps\n",
    "\n",
    "If stateful decoding is being performed, update the states partially. In this step, we select the indices of states that existed in this sub-batch only.\n",
    "\n",
    "After this, we perform regular transducer decoding of the Prediction Network + Joint Network. Since it is being done on a subset of samples, it is much faster than padded decoding.\n",
    "\n",
    "```python\n",
    "best_hyp, _ = self.asr_model.decoding.rnnt_decoder_predictions_tensor(\n",
    "    encoded, encoded_len, return_hypotheses=True, partial_hypotheses=self.previous_hypotheses\n",
    ")\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "69tNXxtfcYBq"
   },
   "source": [
    "#### Preserve the alignments\n",
    "\n",
    "Now that we have the model's hypotheses, we need to preserve the alignments in the correct global index. Remember, we originally had a batch size (say B), but now we performed an inference step over just a sub-batch (say $B'$; $B' ≤ B$), so we need to de-reference the ids of this sub-batch $B'$ to the actual ids in $B$. \n",
    "\n",
    "We utilize `global_index_key = new_batch_keys[idx]` which we built when we were sub-sampling the chunks themselves.\n",
    "\n",
    "```python\n",
    " for idx, hyp in enumerate(best_hyp):\n",
    "    global_index_key = new_batch_keys[idx]  # get index of this sample in the global batch\n",
    "    has_signal_ended = self.frame_bufferer.signal_end[global_index_key]\n",
    "    if not has_signal_ended:\n",
    "        self.all_alignments[global_index_key].append(hyp.alignments)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BmhrYC9XdoeB"
   },
   "source": [
    "#### Update the index mapping from local sub-batch to global batch\n",
    "\n",
    "If, in the current step, the sub-batch $B'$ was smaller than the original batch size $B$, then we need to update the index that is tracked by `self.batch_index_map`.\n",
    "\n",
    "`self.batch_index_map` is a mapping from the global batch index to the current local batch index. \n",
    "\n",
    "```python\n",
    "# Position map update\n",
    "if len(new_batch_keys) != len(self.batch_index_map):\n",
    "    for new_batch_idx, global_index_key in enumerate(new_batch_keys):\n",
    "        self.batch_index_map[global_index_key] = new_batch_idx  # let index point from global pos -> local pos\n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VIfDE19CeXE5"
   },
   "source": [
    "------\n",
    "You may wonder, why do we need to preserve this information? \n",
    "\n",
    "Consider the case where you have five samples in the original global batch size $B$. Let their ids be - `[0, 1, 2, 3, 4]`.\n",
    "\n",
    "After a few steps, the ids `2` and `3` finished processing and are no longer part of sub-batch $B'$. So the new sub-batch is `[0, 1, 4]`.\n",
    "\n",
    "Now - index `sub_batch[2]` no longer corresponds to sample `[2]` but instead to sample `[4]`. Therefore, this information is preserved in `self.batch_index_map` where the key is the global index id (0-5) and the value is the index of this sample in the current sub-batch (0,1,4).\n",
    "\n",
    "-----\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ndf5as2Zg1tP"
   },
   "source": [
    "### Code: `transcribe(tokens_per_chunk, delay)`\n",
    "\n",
    "The method that actually performs transcriptions on chunks of audio segments. It loops two layers - the samples per batch and the alignments per chunk in each of these samples.\n",
    "\n",
    "\n",
    "```python\n",
    "self.unmerged = [[] for _ in range(self.batch_size)]\n",
    "for idx, alignments in enumerate(self.all_alignments):\n",
    "    signal_end_idx = self.frame_bufferer.signal_end_index[idx]\n",
    "    for a_idx, alignment in enumerate(alignments):\n",
    "        # The core of the \"middle token\" algorithm for buffered ASR.\n",
    "        alignment = alignment[len(alignment) - 1 - delay : len(alignment) - 1 - delay + tokens_per_chunk]\n",
    "        ids, toks = self._alignment_decoder(alignment, self.asr_model.tokenizer, self.blank_id)\n",
    "        if len(ids) > 0 and a_idx < signal_end_idx:\n",
    "            self.unmerged[idx] = inplace_buffer_merge(self.unmerged[idx], ids, delay, model=self.asr_model)\n",
    "...\n",
    "output = []\n",
    "for idx in range(self.batch_size):\n",
    "    output.append(self.greedy_merge(self.unmerged[idx]))\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "G_oL_mwnfnkj"
   },
   "source": [
    "## Evaluation\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1UT0MigeYshQ"
   },
   "source": [
    "Let's call the decoder with a few buffers we create from our long audio file to see how this chunk-based decoder comes together. Some interesting experiments to try would be to see how changing sizes of the chunk and the context affect transcription accuracy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0vr1dg-MXK_4"
   },
   "outputs": [],
   "source": [
    "chunk_len_in_secs: float = 8.0\n",
    "context_len_in_secs: float = 1.0\n",
    "\n",
    "max_steps_per_timestep: int = 5\n",
    "stateful_decoding: bool = False\n",
    "        \n",
    "\n",
    "##########################################################################\n",
    "buffer_len_in_secs = chunk_len_in_secs + 2* context_len_in_secs\n",
    "\n",
    "n_buffers = 5\n",
    "\n",
    "buffer_len = int(sample_rate*buffer_len_in_secs)\n",
    "sampbuffer = np.zeros([buffer_len], dtype=np.float32)\n",
    "\n",
    "chunk_reader = AudioChunkIterator(samples, chunk_len_in_secs, sample_rate)\n",
    "chunk_len = int(sample_rate*chunk_len_in_secs)\n",
    "count = 0\n",
    "buffer_list = []\n",
    "for chunk in chunk_reader:\n",
    "    count +=1\n",
    "    sampbuffer[:-chunk_len] = sampbuffer[chunk_len:]\n",
    "    sampbuffer[-chunk_len:] = chunk\n",
    "    buffer_list.append(np.array(sampbuffer))\n",
    "\n",
    "    plt.plot(buffer_list[-1])\n",
    "    plt.show()\n",
    "   \n",
    "    display(Audio(sampbuffer, rate=16000))\n",
    "    if count >= n_buffers:\n",
    "        break\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZZItpaiergK1"
   },
   "source": [
    "## Change Decoding Strategy for Transducer Model\n",
    "\n",
    "Below, we will change the decoding strategy for transducer models to preserve the alignments during autoregressive predictions. This will enable us to easily compute the \"middle token\" during decoding."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Ngrw2L2Krfbc"
   },
   "outputs": [],
   "source": [
    "# Change Decoding Config\n",
    "from omegaconf import OmegaConf, open_dict\n",
    "\n",
    "decoding_cfg = model.cfg.decoding\n",
    "with open_dict(decoding_cfg):\n",
    "    if stateful_decoding:  # Very slow procedure, avoid unless really needed\n",
    "        decoding_cfg.strategy = \"greedy\"\n",
    "    else:\n",
    "        decoding_cfg.strategy = \"greedy_batch\"\n",
    "\n",
    "    decoding_cfg.preserve_alignments = True  # required to compute the middle token for transducers.\n",
    "    decoding_cfg.fused_batch_size = -1  # temporarily stop fused batch during inference.\n",
    "\n",
    "model.change_decoding_strategy(decoding_cfg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6CRTw38bY35k"
   },
   "outputs": [],
   "source": [
    "stride = 4 # 8 for ContextNet\n",
    "asr_decoder = BatchedFrameASRRNNT(model, frame_len=chunk_len_in_secs, total_buffer=buffer_len_in_secs, \n",
    "                                  batch_size=1, \n",
    "                                  max_steps_per_timestep=max_steps_per_timestep, \n",
    "                                  stateful_decoding=stateful_decoding)\n",
    "\n",
    "samples = get_samples(concat_audio_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "VIH12cFfbzn8"
   },
   "outputs": [],
   "source": [
    "import tqdm\n",
    "import math\n",
    "\n",
    "def transcribe_buffers(asr_decoder, samples, num_frames, chunk_len_in_secs, buffer_len_in_secs, model_stride, plot=False):\n",
    "\n",
    "  model.freeze()\n",
    "  model_stride_in_secs = asr_decoder.asr_model.cfg.preprocessor.window_stride * model_stride\n",
    "  tokens_per_chunk = math.ceil(chunk_len_in_secs / model_stride_in_secs)\n",
    "  mid_delay = math.ceil((chunk_len_in_secs + (buffer_len_in_secs - chunk_len_in_secs) / 2) / model_stride_in_secs)\n",
    "\n",
    "  batch_size = asr_decoder.batch_size  # Since only one sample buffers are available, batch size = 1\n",
    "\n",
    "  assert batch_size == 1\n",
    "\n",
    "  with torch.inference_mode():\n",
    "    with torch.cuda.amp.autocast():\n",
    "      asr_decoder.reset()\n",
    "      asr_decoder.sample_offset = 0\n",
    "\n",
    "      frame_reader = streaming_utils.AudioFeatureIterator(samples.copy(), asr_decoder.frame_len, asr_decoder.raw_preprocessor, asr_decoder.asr_model.device)\n",
    "      asr_decoder.set_frame_reader(frame_reader, idx=0, limit_frames=num_frames if num_frames is not None else None)\n",
    "\n",
    "      transcription = asr_decoder.transcribe(tokens_per_chunk, mid_delay, plot=plot)\n",
    "  \n",
    "  return transcription"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ObAXQRQjjTUt"
   },
   "outputs": [],
   "source": [
    "transcription = transcribe_buffers(asr_decoder, samples, n_buffers, chunk_len_in_secs, buffer_len_in_secs, stride, plot=True)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "dy6YyKoowckV"
   },
   "outputs": [],
   "source": [
    "print()\n",
    "print(\"Transcription :\")\n",
    "print(transcription)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MHUN6q-NwhMK"
   },
   "source": [
    "# Transcribe the entire concatenated audio\n",
    "\n",
    "Finally, we will decode the entire 15-minute audio clip with the settings chosen above. It should take just a few seconds to transcribe the entire clip with large chunk sizes, but it increases significantly for shorter chunk sizes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ucFHfwyKk5oN"
   },
   "outputs": [],
   "source": [
    "# WER calculation\n",
    "from nemo.collections.asr.metrics.wer import word_error_rate\n",
    "# Collect all buffers from the audio file\n",
    "sampbuffer = np.zeros([buffer_len], dtype=np.float32)\n",
    "\n",
    "chunk_reader = AudioChunkIterator(samples, chunk_len_in_secs, sample_rate)\n",
    "buffer_list = []\n",
    "for chunk in chunk_reader:\n",
    "    sampbuffer[:-chunk_len] = sampbuffer[chunk_len:]\n",
    "    sampbuffer[-chunk_len:] = chunk\n",
    "    buffer_list.append(np.array(sampbuffer))\n",
    "\n",
    "asr_decoder = BatchedFrameASRRNNT(model, frame_len=chunk_len_in_secs, total_buffer=buffer_len_in_secs, \n",
    "                                  batch_size=1, \n",
    "                                  max_steps_per_timestep=max_steps_per_timestep, \n",
    "                                  stateful_decoding=stateful_decoding)\n",
    "\n",
    "transcription = transcribe_buffers(asr_decoder, samples, None, chunk_len_in_secs, buffer_len_in_secs, stride, plot=False)[0]\n",
    "wer = word_error_rate(hypotheses=[transcription], references=[ref_transcript], use_cer=False)\n",
    "\n",
    "print(f\"WER: {round(wer*100,2)}%\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lLyqdZf_0bhr"
   },
   "source": [
    "# Find the differences in the transcript\n",
    "\n",
    "Word Error Rate is a great tool to measure the performance of the model, but we can go further and debug exactly where mistakes were made. This will further help us determine if the transcript was incorrect due to the merge algorithm rather than the model making any mistake during transcription."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "-ZI2_j7xxVNu"
   },
   "outputs": [],
   "source": [
    "#@title Run to setup Text Diff\n",
    "import difflib\n",
    "from typing import List, Any, Callable, Tuple, Union\n",
    "from itertools import zip_longest\n",
    "import html\n",
    "import re\n",
    "\n",
    "Token = str\n",
    "TokenList = List[Token]\n",
    "\n",
    "whitespace = re.compile('\\s+')\n",
    "end_sentence = re.compile('[.]\\s+')\n",
    "\n",
    "def tokenize(s:str) -> TokenList:\n",
    "    '''Split a string into tokens'''\n",
    "    return whitespace.split(s)\n",
    "\n",
    "def untokenize(ts:TokenList) -> str:\n",
    "    '''Join a list of tokens into a string'''\n",
    "    return ' '.join(ts)\n",
    "\n",
    "def sentencize(s:str) -> TokenList:\n",
    "    '''Split a string into a list of sentences'''\n",
    "    return end_sentence.split(s)\n",
    "\n",
    "def unsentencise(ts:TokenList) -> str:\n",
    "    '''Join a list of sentences into a string'''\n",
    "    return '. '.join(ts)\n",
    "\n",
    "def html_unsentencise(ts:TokenList) -> str:\n",
    "    '''Joing a list of sentences into HTML for display'''\n",
    "    return ''.join(f'<p>{t}</p>' for t in ts)\n",
    "\n",
    "def mark_text(text:str) -> str:\n",
    "    return f'<span style=\"color: red;\">{text}</span>'\n",
    "    \n",
    "def mark_span(text:TokenList) -> TokenList:\n",
    "    if len(text) > 0:\n",
    "        text[0] = '<span style=\"background: #69E2FB;\">' + text[0]\n",
    "        text[-1] += '</span>'\n",
    "    return text\n",
    "\n",
    "def markup_diff(a:TokenList, b:TokenList,\n",
    "                mark=mark_span,\n",
    "                default_mark = lambda x: x,\n",
    "                isjunk=None) -> Tuple[TokenList, TokenList]:\n",
    "    \"\"\"Returns a and b with any differences processed by mark\n",
    "\n",
    "    Junk is ignored by the differ\n",
    "    \"\"\"\n",
    "    seqmatcher = difflib.SequenceMatcher(isjunk=isjunk, a=a, b=b, autojunk=False)\n",
    "    out_a, out_b = [], []\n",
    "    for tag, a0, a1, b0, b1 in seqmatcher.get_opcodes():\n",
    "        markup = default_mark if tag == 'equal' else mark\n",
    "        out_a += markup(a[a0:a1])\n",
    "        out_b += markup(b[b0:b1])\n",
    "    assert len(out_a) == len(a)\n",
    "    assert len(out_b) == len(b)\n",
    "    return out_a, out_b\n",
    "\n",
    "\n",
    "def align_seqs(a: TokenList, b: TokenList, fill:Token='') -> Tuple[TokenList, TokenList]:\n",
    "    out_a, out_b = [], []\n",
    "    seqmatcher = difflib.SequenceMatcher(a=a, b=b, autojunk=False)\n",
    "    for tag, a0, a1, b0, b1 in seqmatcher.get_opcodes():\n",
    "        delta = (a1 - a0) - (b1 - b0)\n",
    "        out_a += a[a0:a1] + [fill] * max(-delta, 0)\n",
    "        out_b += b[b0:b1] + [fill] * max(delta, 0)\n",
    "    assert len(out_a) == len(out_b)\n",
    "    return out_a, out_b\n",
    "\n",
    "\n",
    "def html_sidebyside(a, b):\n",
    "    # Set the panel display\n",
    "    out = '<div style=\"display: grid;grid-template-columns: 1fr 1fr;grid-gap: 20px;\">'\n",
    "    # There's some CSS in Jupyter notebooks that makes the first pair unalign. This is a workaround\n",
    "    out += '<p></p><p></p>'\n",
    "    for left, right in zip_longest(a, b, fillvalue=''):\n",
    "        out += f'<p>{left}</p>'\n",
    "        out += f'<p>{right}</p>'\n",
    "        out += '</div>'\n",
    "    return out\n",
    "\n",
    "def html_diffs(a, b):\n",
    "    a = html.escape(a)\n",
    "    b = html.escape(b)\n",
    "\n",
    "    out_a, out_b = [], []\n",
    "    for sent_a, sent_b in zip(*align_seqs(sentencize(a), sentencize(b))):\n",
    "        mark_a, mark_b = markup_diff(tokenize(sent_a), tokenize(sent_b))\n",
    "        out_a.append(untokenize(mark_a))\n",
    "        out_b.append(untokenize(mark_b))\n",
    "\n",
    "    return html_sidebyside(out_a, out_b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "P9jRGFIw0nDp"
   },
   "outputs": [],
   "source": [
    "from IPython.display import HTML\n",
    "\n",
    "def show_diffs(a, b):\n",
    "    display(HTML(html_diffs(a,b)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "CcKJrq3-024_"
   },
   "outputs": [],
   "source": [
    "show_diffs(ref_transcript,  transcription)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "NBSX6HS44pBR"
   },
   "source": [
    "# Extra: Middle token alignment in Transducers\n",
    "\n",
    "The middle token algorithm is a general algorithm that works even for non-causal offline models relatively well. On the other hand, it requires larger chunk sizes and larger buffer sizes (which causes some latency).\n",
    "\n",
    "For CTC, it is straightforward to decode the text from the log probabilities - take argmax. However, transducers have a much more complicated decoding technique, and taking middle token alignment is not so straightforward.\n",
    "\n",
    "So below, we will take a deep dive into what transducer alignments are and how we can use them to perform middle token selection."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wBxrDbyq5cv3"
   },
   "source": [
    "\n",
    "\n",
    "```python\n",
    "def transcribe(self, tokens_per_chunk: int, delay: int, plot=False):\n",
    "    ...\n",
    "        all_toks = []\n",
    "        for a_idx, alignment in enumerate(alignments):\n",
    "            alignment = alignment[len(alignment) - 1 - delay : len(alignment) - 1 - delay + tokens_per_chunk]\n",
    "            ids, toks = self._alignment_decoder(alignment, self.asr_model.tokenizer, self.blank_id)\n",
    "            all_toks.append(toks)\n",
    "    ...\n",
    "\n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "pN-HFw8i7wos"
   },
   "outputs": [],
   "source": [
    "from nemo.collections.asr.parts.utils.manifest_utils import read_manifest\n",
    "\n",
    "manifest_data = read_manifest(manifest)\n",
    "print(f\"Read {len(manifest_data)} samples from manifest {manifest}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "YKTYXU_38D__"
   },
   "outputs": [],
   "source": [
    "sample_idx = 0\n",
    "audio_filepath = manifest_data[sample_idx]['audio_filepath']\n",
    "text = manifest_data[sample_idx]['text']\n",
    "\n",
    "print(\"Audio filepath   :\", audio_filepath)\n",
    "print(\"Ground truth text:\", text)\n",
    "print(\"Duration         :\", manifest_data[sample_idx][\"duration\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MxRSKKX9qsby"
   },
   "source": [
    "------\n",
    "\n",
    "To track the alignment array, we need to let the Transducer Decoding strategy know to preserve the alignments. Since it is done during greedy / beam decoding, it can be simply done by setting a flag and updating the decoding strategy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "WFUGpH2W8Oke"
   },
   "outputs": [],
   "source": [
    "decoding_cfg = model.cfg.decoding\n",
    "\n",
    "with open_dict(decoding_cfg):\n",
    "  decoding_cfg.preserve_alignments = True  # required to compute the middle token for transducers.\n",
    "  decoding_cfg.fused_batch_size = -1  # temporarily stop fused batch during inference.\n",
    "\n",
    "model.change_decoding_strategy(decoding_cfg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "GO86ej_V8gbM"
   },
   "outputs": [],
   "source": [
    "greedy_hypotheses, beam_hypotheses = model.transcribe([audio_filepath], return_hypotheses=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "CYPg3T3K8ffV"
   },
   "outputs": [],
   "source": [
    "hyp = greedy_hypotheses[0]  # only one sample in batch\n",
    "print(\"Hypothesis text  :\", hyp.text)\n",
    "print(\"Ground truth text:\", text)\n",
    "\n",
    "wer = word_error_rate(hypotheses=[hyp.text], references=[text], use_cer=False)\n",
    "print(\"Word error rate of audio :\", wer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fgDkjlDX-hym"
   },
   "source": [
    "# Print alignments\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "SjjQ9ODz-No6"
   },
   "outputs": [],
   "source": [
    "alignments = hyp.alignments"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1gqFfQwZ-xpk"
   },
   "source": [
    "Let's calculate the alignment grid. We will de-tokenize the sub-word token if it is a valid index in the vocabulary and use '' as a placeholder for the Transducer Blank token.\n",
    "\n",
    "Note that each timestep here is (roughly) 40 milli-seconds timestamp (since the window stride is 10 ms, and Conformer has 4x stride). The resolution of the model differs based on the stride of the model - QuartzNet has 2x stride (20 ms), Conformer has 4x stride (40 ms), and ContextNet has 8x stride (80 ms).\n",
    "\n",
    "Note: You can modify the value of config.model.loss.warprnnt_numba_kwargs.fastemit_lambda before training and see an impact on final alignment latency! For a tutorial to train your Transducer models, refer to [ASR with Transducers in NeMo](https://colab.research.google.com/github/NVIDIA/NeMo/blob/stable/tutorials/asr/ASR_with_Transducers.ipynb)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Zvyh0ATNrLhQ"
   },
   "source": [
    "------\n",
    "\n",
    "First, let's listen to the audio clip itself. We can then determine if the text transcription is being accurately transcribed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9njFbF4P-qNL"
   },
   "outputs": [],
   "source": [
    "display(Audio(filename=audio_filepath, rate=model.cfg.sample_rate))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "82NLN1ZC-wCM"
   },
   "outputs": [],
   "source": [
    "model_stride = 4  # 4 for Conformers; 8 for ContextNet / Citrinet;\n",
    "window_stride = model.cfg.preprocessor.window_stride\n",
    "\n",
    "time_ms = 0.0  # time in ms\n",
    "\n",
    "# Compute the alignment grid\n",
    "for ti in range(len(alignments)):\n",
    "  t_u = []\n",
    "  for uj in range(len(alignments[ti])):\n",
    "    logprobs, token = alignments[ti][uj]\n",
    "    token = token.to('cpu').numpy().tolist()\n",
    "    decoded_token = model.decoding.decode_ids_to_tokens([token])[0] if token != model.decoding.blank_id else ''  # token at index len(vocab) == RNNT blank token\n",
    "    t_u.append(decoded_token)\n",
    "  \n",
    "  time_ms += model_stride * window_stride\n",
    "  print(f\"Tokens at timestep {ti} (time={time_ms:0.2f} s) = {t_u}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mEArX3Khrbuu"
   },
   "source": [
    "-------\n",
    "\n",
    "You can see that there were roughly as many timesteps as the original duration of the audio. Since the transducer alignment is characterized as a Ti x Uj dangling array, we can treat Ti as the equivalent to the timestep of the acoustic signal itself. Therefore, if we take \"middle\" slices over the Ti dimension during decoding, we can effectively compute the middle token alignment even for transducer models!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "y6ZhP-ENsMHH"
   },
   "source": [
    "# Final Notes\n",
    "\n",
    "Now, anyone can perform long audio transcription using any NeMo transducer model. You could even try to modify the chunk and buffer sizes to try to stream these models.\n",
    "\n",
    "For further references on training your own transducer models, please refer to [ASR with Transducers](https://github.com/NVIDIA/NeMo/blob/stable/tutorials/asr/ASR_with_Transducers.ipynb) tutorial."
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "Buffered_Transducer_Inference.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}