notebooks/227-whisper-subtitles-generation/227-whisper-nncf-quantize.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Post-Training Quantization of OpenAI Whisper model with NNCF\n",
    "\n",
    "The goal of this tutorial is to demonstrate how to speed up the model by applying 8-bit post-training quantization from [NNCF](https://github.com/openvinotoolkit/nncf/) (Neural Network Compression Framework) and infer quantized model via OpenVINO™ Toolkit. The optimization process contains the following steps:\n",
    "\n",
    "1. Quantize the converted OpenVINO model from [227-whisper-convert notebook](227-whisper-convert.ipynb) with NNCF.\n",
    "2. Check model result for the demo video.\n",
    "3. Compare model size, performance and accuracy of FP32 and quantized INT8 models.\n",
    "\n",
    "> **NOTE**: you should run [227-whisper-convert](227-whisper-convert.ipynb) notebook first to generate OpenVINO IR model that is used for quantization.\n",
    "\n",
    "Table of content:\n",
    "- [Prerequisites](#Prerequisites-Uparrow)\n",
    "- [Create and initialize quantization](#Create-and-initialize-quantization-Uparrow)\n",
    "    - [Prepare calibration datasets](#Prepare-calibration-datasets-Uparrow)\n",
    "    - [Quantize Whisper encoder and decoder models](#Quantize-Whisper-encoder-and-decoder-models-Uparrow)\n",
    "- [Transcribe video with quantized OpenVINO model](#Transcribe-video-with-quantized-OpenVINO-model-Uparrow)\n",
    "- [Compare performance and accuracy of the FP32 and INT8 IRs](#Compare-performance-and-accuracy-of-the-FP32-and-INT8-IRs-Uparrow)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prerequisites [$\\Uparrow$](#Table-of-content:)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Install dependencies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-08-24T07:22:38.138055200Z",
     "start_time": "2023-08-24T07:22:38.000326500Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "!pip install -q \"openvino==2023.1.0.dev20230811\"\n",
    "!pip install -q \"git+https://github.com/openvinotoolkit/nncf.git@release_v260\"\n",
    "!pip install -q datasets librosa soundfile\n",
    "!pip install -q evaluate jiwer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Select device from dropdown list for running inference using OpenVINO."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-08-24T07:22:38.357268200Z",
     "start_time": "2023-08-24T07:22:38.250550900Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6cb79741c37f49e4ad4e8b3e8ab073cc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dropdown(description='Device:', index=4, options=('CPU', 'GPU.0', 'GPU.1', 'GPU.2', 'AUTO'), value='AUTO')"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import ipywidgets as widgets\n",
    "\n",
    "from openvino import Core\n",
    "core = Core()\n",
    "\n",
    "device = widgets.Dropdown(\n",
    "    options=core.available_devices + [\"AUTO\"],\n",
    "    value='AUTO',\n",
    "    description='Device:',\n",
    "    disabled=False,\n",
    ")\n",
    "\n",
    "device"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Select the task for the model:\n",
    "\n",
    "* **transcribe** - generate audio transcription in the source language (automatically detected).\n",
    "* **translate** - generate audio transcription with translation to English language."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-08-24T07:22:38.359102800Z",
     "start_time": "2023-08-24T07:22:38.357540500Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7d1d9a6f461546288d227179653954ad",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Select(description='Select task:', index=1, options=('transcribe', 'translate'), value='translate')"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "task = widgets.Select(\n",
    "    options=[\"transcribe\", \"translate\"],\n",
    "    value=\"translate\",\n",
    "    description=\"Select task:\",\n",
    "    disabled=False\n",
    ")\n",
    "task"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "<a id=\"2\"></a>\n",
    "## Create and initialize quantization [&#8657;](#0)\n",
    "\n",
    "[NNCF](https://github.com/openvinotoolkit/nncf/) enables post-training quantization by adding the quantization layers into the model graph and then using a subset of the training dataset to initialize the parameters of these additional quantization layers. The framework is designed so that modifications to your original training code are minor. Quantization is the simplest scenario and requires a few modifications.\n",
    "\n",
    "The optimization process contains the following steps:\n",
    "\n",
    "1. Create a calibration dataset for quantization.\n",
    "2. Run `nncf.quantize` to obtain quantized models.\n",
    "3. Serialize the `INT8` model using `openvino.runtime.serialize` function."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Set paths to the model converted in [227-whisper-convert](227-whisper-convert.ipynb) notebook and the paths where quantized models will be saved."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "WHISPER_ENCODER_OV = Path(\"whisper_encoder.xml\")\n",
    "WHISPER_DECODER_OV = Path(\"whisper_decoder.xml\")\n",
    "\n",
    "WHISPER_ENCODER_OV_INT8 = Path(\"whisper_encoder_int8.xml\")\n",
    "WHISPER_DECODER_OV_INT8 = Path(\"whisper_decoder_int8.xml\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Load FP32 model IR."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-08-24T07:22:39.574746400Z",
     "start_time": "2023-08-24T07:22:38.358318Z"
    }
   },
   "outputs": [],
   "source": [
    "import whisper\n",
    "from utils import patch_whisper_for_ov_inference, OpenVINOAudioEncoder, OpenVINOTextDecoder\n",
    "\n",
    "model_fp32 = whisper.load_model(\"base\").to(\"cpu\").eval()\n",
    "patch_whisper_for_ov_inference(model_fp32)\n",
    "\n",
    "model_fp32.encoder = OpenVINOAudioEncoder(core, WHISPER_ENCODER_OV, device=device.value)\n",
    "model_fp32.decoder = OpenVINOTextDecoder(core, WHISPER_DECODER_OV, device=device.value)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "### Prepare calibration datasets [$\\Uparrow$](#Table-of-content:)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Whisper consists of an encoder and a decoder models. We need to collect calibration data for both of them.\n",
    "\n",
    "Below we overwrite encoder/decoder forward methods in order to collect calibration samples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-08-24T07:22:39.623947800Z",
     "start_time": "2023-08-24T07:22:39.575286700Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "from contextlib import contextmanager\n",
    "from functools import partial\n",
    "import openvino as ov\n",
    "from typing import Optional\n",
    "import torch\n",
    "\n",
    "COLLECT_CALIBRATION_DATA = False\n",
    "encoder_calibration_data = []\n",
    "decoder_calibration_data = []\n",
    "\n",
    "@contextmanager\n",
    "def calibration_data_collection():\n",
    "    global COLLECT_CALIBRATION_DATA\n",
    "    try:\n",
    "        COLLECT_CALIBRATION_DATA = True\n",
    "        yield\n",
    "    finally:\n",
    "        COLLECT_CALIBRATION_DATA = False\n",
    "\n",
    "\n",
    "def encoder_forward(self, mel: torch.Tensor):\n",
    "    if COLLECT_CALIBRATION_DATA:\n",
    "        encoder_calibration_data.append(mel)\n",
    "    return torch.from_numpy(self.compiled_model(mel)[self.output_blob])\n",
    "\n",
    "def decoder_forward(self, x: torch.Tensor, xa: torch.Tensor, kv_cache: Optional[dict] = None):\n",
    "    feed_dict = {'x': ov.Tensor(x.numpy()), 'xa': ov.Tensor(xa.numpy())}\n",
    "    feed_dict = (self.preprocess_kv_cache_inputs(feed_dict, kv_cache))\n",
    "    if COLLECT_CALIBRATION_DATA:\n",
    "        decoder_calibration_data.append(feed_dict)\n",
    "    res = self.compiled_model(feed_dict)\n",
    "    return self.postprocess_outputs(res)\n",
    "\n",
    "model_fp32.encoder.forward = partial(encoder_forward, model_fp32.encoder)\n",
    "model_fp32.decoder.forward = partial(decoder_forward, model_fp32.decoder)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "We use a portion of validation [librispeech_asr](https://huggingface.co/datasets/librispeech_asr) dataset from Hugging Face as calibration data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-08-24T07:23:05.269312200Z",
     "start_time": "2023-08-24T07:22:39.623947800Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "test_replace": {
     "CALIBRATION_DATASET_SIZE = 30": "CALIBRATION_DATASET_SIZE = 1"
    }
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "af02ea88a96f4042bc1a4b2d7159be44",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Collecting calibration data:   0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "CALIBRATION_DATASET_SIZE = 30\n",
    "\n",
    "calibration_dataset = load_dataset(\"librispeech_asr\", \"clean\", split=\"validation\", streaming=True).take(CALIBRATION_DATASET_SIZE)\n",
    "\n",
    "with calibration_data_collection():\n",
    "    for data_item in tqdm(calibration_dataset, desc=\"Collecting calibration data\", total=CALIBRATION_DATASET_SIZE):\n",
    "        model_fp32.transcribe(data_item[\"audio\"][\"array\"].astype(\"float32\"), task=task.value)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "### Quantize Whisper encoder and decoder models [$\\Uparrow$](#Table-of-content:)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Quantize both encoder and decoder models using `nncf.quantize()` API and save the quantized IRs after that."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-08-24T07:25:31.614190800Z",
     "start_time": "2023-08-24T07:23:05.269312200Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:nncf:NNCF initialized successfully. Supported frameworks detected: torch, tensorflow, onnx, openvino\n",
      "Quantizing encoder...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-08-30 19:38:10.314501: I tensorflow/core/util/port.cc:110] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n",
      "2023-08-30 19:38:10.347770: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2023-08-30 19:38:10.917857: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n",
      "Statistics collection: 100%|████████████████████████████████████████████████████████████████████████████████████████████████| 60/60 [00:04<00:00, 12.26it/s]\n",
      "Applying Smooth Quant: 100%|████████████████████████████████████████████████████████████████████████████████████████████████| 24/24 [00:00<00:00, 60.29it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:nncf:18 ignored nodes was found by name in the NNCFGraph\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Statistics collection: 100%|████████████████████████████████████████████████████████████████████████████████████████████████| 60/60 [00:14<00:00,  4.14it/s]\n",
      "Applying Fast Bias correction: 100%|████████████████████████████████████████████████████████████████████████████████████████| 32/32 [00:06<00:00,  5.22it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saved quantized encoder at ./whisper_encoder_int8.xml\n",
      "Quantizing decoder...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Statistics collection: 100%|██████████████████████████████████████████████████████████████████████████████████████████████| 664/664 [00:12<00:00, 54.92it/s]\n",
      "Applying Smooth Quant: 100%|████████████████████████████████████████████████████████████████████████████████████████████████| 38/38 [00:00<00:00, 39.37it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:nncf:36 ignored nodes was found by name in the NNCFGraph\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Statistics collection: 100%|██████████████████████████████████████████████████████████████████████████████████████████████| 664/664 [00:34<00:00, 19.20it/s]\n",
      "Applying Fast Bias correction: 100%|████████████████████████████████████████████████████████████████████████████████████████| 48/48 [00:07<00:00,  6.30it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saved quantized decoder at ./whisper_decoder_int8.xml\n"
     ]
    }
   ],
   "source": [
    "import nncf\n",
    "from nncf.quantization.advanced_parameters import OverflowFix\n",
    "from openvino.runtime import serialize\n",
    "\n",
    "print(\"Quantizing encoder...\")\n",
    "quantized_encoder = nncf.quantize(\n",
    "    model=model_fp32.encoder.model,\n",
    "    calibration_dataset=nncf.Dataset(encoder_calibration_data),\n",
    "    subset_size=len(encoder_calibration_data),\n",
    "    model_type=nncf.ModelType.TRANSFORMER,\n",
    "    advanced_parameters=nncf.AdvancedQuantizationParameters(\n",
    "        overflow_fix=OverflowFix.DISABLE,   # disable overflow fix (can lead to accuracy drop on legacy platforms w/o DL Boost)\n",
    "        smooth_quant_alpha=0.5      # Smooth Quant algorithm reduces activation quantization error; optimal alpha value was obtained through grid search\n",
    "    )\n",
    ")\n",
    "serialize(quantized_encoder, WHISPER_ENCODER_OV_INT8)\n",
    "print(f\"Saved quantized encoder at ./{WHISPER_ENCODER_OV_INT8}\")\n",
    "\n",
    "print(\"Quantizing decoder...\")\n",
    "quantized_decoder = nncf.quantize(\n",
    "    model=model_fp32.decoder.model,\n",
    "    calibration_dataset=nncf.Dataset(decoder_calibration_data),\n",
    "    subset_size=len(decoder_calibration_data),\n",
    "    model_type=nncf.ModelType.TRANSFORMER,\n",
    "    advanced_parameters=nncf.AdvancedQuantizationParameters(\n",
    "        overflow_fix=OverflowFix.DISABLE,   # disable overflow fix (can lead to accuracy drop on legacy platforms w/o DL Boost)\n",
    "        smooth_quant_alpha=0.95     # Smooth Quant algorithm reduces activation quantization error; optimal alpha value was obtained through grid search\n",
    "    )\n",
    ")\n",
    "serialize(quantized_decoder, WHISPER_DECODER_OV_INT8)\n",
    "print(f\"Saved quantized decoder at ./{WHISPER_DECODER_OV_INT8}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Transcribe video with quantized OpenVINO model [$\\Uparrow$](#Table-of-content:)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Load `INT8` models saved above into a new instance of Whisper model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "model_int8 = whisper.load_model(\"base\").to(\"cpu\").eval()\n",
    "patch_whisper_for_ov_inference(model_int8)\n",
    "\n",
    "model_int8.encoder = OpenVINOAudioEncoder(core, WHISPER_ENCODER_OV_INT8, device=device.value)\n",
    "model_int8.decoder = OpenVINOTextDecoder(core, WHISPER_DECODER_OV_INT8, device=device.value)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Select a video for transcription as in [227-whisper-convert](227-whisper-convert.ipynb) notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-08-24T07:25:33.101022100Z",
     "start_time": "2023-08-24T07:25:33.090494400Z"
    }
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f6ddf61a225a4e8b9f5b191c88094a01",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Text(value='https://youtu.be/kgL5LBM-hFI', description='Video:', placeholder='Type link for video')"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "VIDEO_LINK = \"https://youtu.be/kgL5LBM-hFI\"\n",
    "link = widgets.Text(\n",
    "    value=VIDEO_LINK,\n",
    "    placeholder=\"Type link for video\",\n",
    "    description=\"Video:\",\n",
    "    disabled=False\n",
    ")\n",
    "link"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-08-24T07:25:33.108346600Z",
     "start_time": "2023-08-24T07:25:33.105926700Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading video https://youtu.be/kgL5LBM-hFI started\n",
      "Video saved to downloaded_video.mp4\n"
     ]
    }
   ],
   "source": [
    "from pytube import YouTube\n",
    "\n",
    "print(f\"Downloading video {link.value} started\")\n",
    "\n",
    "output_file = Path(\"downloaded_video.mp4\")\n",
    "yt = YouTube(link.value)\n",
    "yt.streams.get_highest_resolution().download(filename=output_file)\n",
    "print(f\"Video saved to {output_file}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-08-24T07:25:36.876565700Z",
     "start_time": "2023-08-24T07:25:33.120502500Z"
    }
   },
   "outputs": [],
   "source": [
    "from utils import get_audio\n",
    "\n",
    "audio = get_audio(output_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Run transcription by the quantized model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-08-24T07:25:39.288568400Z",
     "start_time": "2023-08-24T07:25:36.875642800Z"
    }
   },
   "outputs": [],
   "source": [
    "transcription = model_int8.transcribe(audio, task=task.value)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-08-24T07:25:39.311444800Z",
     "start_time": "2023-08-24T07:25:39.311444800Z"
    }
   },
   "outputs": [],
   "source": [
    "from utils import prepare_srt\n",
    "\n",
    "srt_lines = prepare_srt(transcription)\n",
    "# save transcription\n",
    "with output_file.with_suffix(\".srt\").open(\"w\") as f:\n",
    "    f.writelines(srt_lines)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let us see the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-08-24T07:25:39.327567Z",
     "start_time": "2023-08-24T07:25:39.311444800Z"
    }
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b24a5b170d204fa5b613990fdcba74a1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Video(value=b'\\x00\\x00\\x00\\x18ftypmp42\\x00\\x00\\x00\\x00isommp42\\x00\\x00Aimoov\\x00\\x00\\x00lmvhd...', height='800…"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "widgets.Video.from_file(output_file, loop=False, width=800, height=800)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-08-24T07:25:39.327567Z",
     "start_time": "2023-08-24T07:25:39.327567Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1\n",
      "00:00:00,000 --> 00:00:07,000\n",
      " What's that? Oh, wow.\n",
      "\n",
      "2\n",
      "00:00:09,000 --> 00:00:11,000\n",
      " Hello humans.\n",
      "\n",
      "3\n",
      "00:00:14,000 --> 00:00:15,000\n",
      " Focus on me.\n",
      "\n",
      "4\n",
      "00:00:15,000 --> 00:00:16,000\n",
      " Focus on the guard.\n",
      "\n",
      "5\n",
      "00:00:18,000 --> 00:00:20,000\n",
      " Don't tell anyone what you've seen in here.\n",
      "\n",
      "6\n",
      "00:00:22,000 --> 00:00:24,000\n",
      " Have you seen what's in there?\n",
      "\n",
      "7\n",
      "00:00:24,000 --> 00:00:25,000\n",
      " They have intel.\n",
      "\n",
      "8\n",
      "00:00:25,000 --> 00:00:27,000\n",
      " This is where it all changes.\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(\"\".join(srt_lines))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "As you can see the result is almost the same."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Compare performance and accuracy of the FP32 and INT8 IRs [$\\Uparrow$](#Table-of-content:)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Compare model file size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-08-24T07:25:39.370126700Z",
     "start_time": "2023-08-24T07:25:39.336253900Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: whisper_encoder\n",
      "    * FP32 IR model size: 40216.07 KB\n",
      "    * INT8 IR model size: 21092.37 KB\n",
      "    * Model compression rate: 1.907\n",
      "Model: whisper_decoder\n",
      "    * FP32 IR model size: 101961.09 KB\n",
      "    * INT8 IR model size: 78058.77 KB\n",
      "    * Model compression rate: 1.306\n"
     ]
    }
   ],
   "source": [
    "def calculate_compression_rate(model_path_ov, model_path_ov_int8):\n",
    "    model_size_fp32 = model_path_ov.with_suffix(\".bin\").stat().st_size / 1024\n",
    "    model_size_int8 = model_path_ov_int8.with_suffix(\".bin\").stat().st_size / 1024\n",
    "    print(f\"Model: {model_path_ov.stem}\")\n",
    "    print(f\"    * FP32 IR model size: {model_size_fp32:.2f} KB\")\n",
    "    print(f\"    * INT8 IR model size: {model_size_int8:.2f} KB\")\n",
    "    print(f\"    * Model compression rate: {model_size_fp32 / model_size_int8:.3f}\")\n",
    "\n",
    "calculate_compression_rate(WHISPER_ENCODER_OV, WHISPER_ENCODER_OV_INT8)\n",
    "calculate_compression_rate(WHISPER_DECODER_OV, WHISPER_DECODER_OV_INT8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "To measure the inference performance of the `FP32` and `INT8` encoder/decoder models, we use median inference time on calibration dataset.\n",
    "So we can approximately estimate the speed-up of the dynamic quantized models.\n",
    "\n",
    "> **NOTE**: For the most accurate performance estimation, it is recommended to run `benchmark_app` with static shapes in a terminal/command prompt after closing other applications."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-08-24T07:26:39.712460600Z",
     "start_time": "2023-08-24T07:25:39.370126700Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6ef3f614e9d54c2294a6c6f4a16af701",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Measuring performance:   0%|          | 0/60 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f4e33272f850404bb36f674c82113b1b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Measuring performance:   0%|          | 0/60 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Encoder performance speedup: 1.325\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "792f16b4e5294affbed469f751f660c9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Measuring performance:   0%|          | 0/100 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cc31cba3799f4cb6a073465b665be924",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Measuring performance:   0%|          | 0/100 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Decoder performance speedup: 1.609\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "import numpy as np\n",
    "\n",
    "def calculate_call_inference_time(model, dataset):\n",
    "    inference_time = []\n",
    "    for data_item in tqdm(dataset[:100], desc=\"Measuring performance\"):\n",
    "        start = time.perf_counter()\n",
    "        model(data_item)\n",
    "        end = time.perf_counter()\n",
    "        delta = end - start\n",
    "        inference_time.append(delta)\n",
    "    return np.median(inference_time)\n",
    "\n",
    "\n",
    "encoder_time_fp32 = calculate_call_inference_time(model_fp32.encoder.compiled_model, encoder_calibration_data)\n",
    "encoder_time_int8 = calculate_call_inference_time(model_int8.encoder.compiled_model, encoder_calibration_data)\n",
    "print(f\"Encoder performance speedup: {encoder_time_fp32 / encoder_time_int8:.3f}\")\n",
    "\n",
    "decoder_time_fp32 = calculate_call_inference_time(model_fp32.decoder.compiled_model, decoder_calibration_data)\n",
    "decoder_time_int8 = calculate_call_inference_time(model_int8.decoder.compiled_model, decoder_calibration_data)\n",
    "print(f\"Decoder performance speedup: {decoder_time_fp32 / decoder_time_int8:.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "We measure the whole transcription performance separately, because a single Whisper `transcribe()` call triggers multiple encoder and decoder inference calls.\n",
    "And the number of these calls is dynamic depending on the model accuracy.\n",
    "In this experiment we use the mean time instead of the median because the model transcription time is less uniform.\n",
    "\n",
    "We also compare accuracy values of the `FP32` and `INT8` models on a subset of [librispeech_asr](https://huggingface.co/datasets/librispeech_asr) test dataset.\n",
    "We rely on the Word Error Rate (WER) metric and compute accuracy as `(1 - WER)`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "test_replace": {
     "TEST_DATASET_SIZE = 100": "TEST_DATASET_SIZE = 1"
    }
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "36c2da357eca40058256cc2a5289caf8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Measuring performance and accuracy:   0%|          | 0/100 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2545fb94537247eeaa5fe6e2f6c2d1f4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Measuring performance and accuracy:   0%|          | 0/100 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Whisper transcription performance speedup: 1.446\n",
      "Whisper transcription word accuracy. FP32: 95.61%. INT8: 94.23%. Accuracy drop :1.38%.\n"
     ]
    }
   ],
   "source": [
    "from evaluate import load\n",
    "from transformers import WhisperProcessor\n",
    "\n",
    "wer = load(\"wer\")\n",
    "\n",
    "TEST_DATASET_SIZE = 100\n",
    "test_dataset = load_dataset(\"librispeech_asr\", \"clean\", split=\"test\", streaming=True).take(TEST_DATASET_SIZE)\n",
    "\n",
    "def calculate_transcription_time_and_accuracy(model, dataset):\n",
    "    processor = WhisperProcessor.from_pretrained(\"openai/whisper-large\")\n",
    "\n",
    "    ground_truths = []\n",
    "    predictions = []\n",
    "    inference_time = []\n",
    "    for data_item in tqdm(dataset, desc=\"Measuring performance and accuracy\", total=TEST_DATASET_SIZE):\n",
    "        audio = data_item[\"audio\"][\"array\"].astype(\"float32\")\n",
    "\n",
    "        start_time = time.perf_counter()\n",
    "        transcription = model.transcribe(audio, task=task.value)\n",
    "        end_time = time.perf_counter()\n",
    "        delta_time = end_time - start_time\n",
    "\n",
    "        reference = processor.tokenizer._normalize(data_item[\"text\"])\n",
    "        prediction = processor.tokenizer._normalize(transcription[\"text\"])\n",
    "        ground_truths.append(reference)\n",
    "        predictions.append(prediction)\n",
    "        inference_time.append(delta_time)\n",
    "\n",
    "    word_accuracy = (1 - wer.compute(references=ground_truths, predictions=predictions)) * 100\n",
    "    mean_inference_time = np.mean(inference_time)\n",
    "    return mean_inference_time, word_accuracy\n",
    "\n",
    "transcription_time_fp32, accuracy_fp32 = calculate_transcription_time_and_accuracy(model_fp32, test_dataset)\n",
    "transcription_time_int8, accuracy_int8 = calculate_transcription_time_and_accuracy(model_int8, test_dataset)\n",
    "print(f\"Whisper transcription performance speedup: {transcription_time_fp32 / transcription_time_int8:.3f}\")\n",
    "print(f\"Whisper transcription word accuracy. FP32: {accuracy_fp32:.2f}%. INT8: {accuracy_int8:.2f}%. Accuracy drop :{accuracy_fp32 - accuracy_int8:.2f}%.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "> **NOTE**: Accuracy drop can generally be improved by increasing calibration dataset size."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  },
  "vscode": {
   "interpreter": {
    "hash": "916dbcbb3f70747c44a77c7bcd40155683ae19c65e1c03b4aa3499c5328201f1"
   }
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}