notebooks/whisper-subtitles-generation/whisper-subtitles-generation.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Video Subtitle Generation using Whisper and OpenVINO™\n",
    "\n",
    "[Whisper](https://openai.com/blog/whisper/) is an automatic speech recognition (ASR) system trained on 680,000 hours of multilingual and multitask supervised data collected from the web. It is a multi-task model that can perform multilingual speech recognition as well as speech translation and language identification.\n",
    "\n",
    "![asr-training-data-desktop.svg](https://user-images.githubusercontent.com/29454499/204536347-28976978-9a07-416c-acff-fc1214bbfbe0.svg)\n",
    "\n",
    "You can find more information about this model in the [research paper](https://cdn.openai.com/papers/whisper.pdf), [OpenAI blog](https://openai.com/blog/whisper/), [model card](https://github.com/openai/whisper/blob/main/model-card.md) and GitHub [repository](https://github.com/openai/whisper).\n",
    "\n",
    "In this notebook, we will use Whisper with OpenVINO to generate subtitles in a sample video. Additionally, we will use [NNCF](https://github.com/openvinotoolkit/nncf) improving model performance by INT8 quantization.\n",
    "Notebook contains the following steps:\n",
    "1. Download the model.\n",
    "2. Instantiate the PyTorch model pipeline.\n",
    "3. Convert model to OpenVINO IR, using model conversion API.\n",
    "4. Run the Whisper pipeline with OpenVINO models.\n",
    "5. Quantize the OpenVINO model with NNCF.\n",
    "6. Check quantized model result for the demo video.\n",
    "7. Compare model size, performance and accuracy of FP32 and quantized INT8 models.\n",
    "8. Launch Interactive demo for video subtitles generation.\n",
    "\n",
    "\n",
    "<img referrerpolicy=\"no-referrer-when-downgrade\" src=\"https://static.scarf.sh/a.png?x-pxid=5b5a4db0-7875-4bfb-bdbd-01698b5b1a77&file=notebooks/whisper-subtitles-generation/whisper-subtitles-generation.ipynb\" />\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "\n",
    "\n",
    "#### Table of contents:\n",
    "\n",
    "- [Prerequisites](#Prerequisites)\n",
    "- [Instantiate model](#Instantiate-model)\n",
    "    - [Convert model to OpenVINO Intermediate Representation (IR) format.](#Convert-model-to-OpenVINO-Intermediate-Representation-(IR)-format.)\n",
    "- [Prepare inference pipeline](#Prepare-inference-pipeline)\n",
    "    - [Select inference device](#Select-inference-device)\n",
    "- [Run video transcription pipeline](#Run-video-transcription-pipeline)\n",
    "- [Quantization](#Quantization)\n",
    "    - [Prepare calibration datasets](#Prepare-calibration-datasets)\n",
    "    - [Quantize Whisper encoder and decoder models](#Quantize-Whisper-encoder-and-decoder-models)\n",
    "    - [Run quantized model inference](#Run-quantized-model-inference)\n",
    "    - [Compare performance and accuracy of the original and quantized models](#Compare-performance-and-accuracy-of-the-original-and-quantized-models)\n",
    "- [Interactive demo](#Interactive-demo)\n",
    "\n",
    "\n",
    "### Installation Instructions\n",
    "\n",
    "This is a self-contained example that relies solely on its own code.\n",
    "\n",
    "We recommend  running the notebook in a virtual environment. You only need a Jupyter server to start.\n",
    "For details, please refer to [Installation Guide](https://github.com/openvinotoolkit/openvino_notebooks/blob/latest/README.md#-installation-guide)."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "[back to top ⬆️](#Table-of-contents:)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Install dependencies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%pip install -q \"openvino>=2024.1.0\" \"nncf>=2.10.0\"\n",
    "%pip install -q \"python-ffmpeg<=1.0.16\" moviepy transformers onnx \"git+https://github.com/huggingface/optimum-intel.git\" \"peft==0.6.2\" \"torch>=2.1,<2.4\" \"torchvision<0.19.0\" --extra-index-url https://download.pytorch.org/whl/cpu\n",
    "%pip install -q \"yt_dlp>=2024.8.6\" soundfile librosa jiwer\n",
    "%pip install -q  \"gradio>=4.19\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Instantiate model\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "Whisper is a Transformer based encoder-decoder model, also referred to as a sequence-to-sequence model. It maps a sequence of audio spectrogram features to a sequence of text tokens. First, the raw audio inputs are converted to a log-Mel spectrogram by action of the feature extractor. Then, the Transformer encoder encodes the spectrogram to form a sequence of encoder hidden states. Finally, the decoder autoregressively predicts text tokens, conditional on both the previous tokens and the encoder hidden states.\n",
    "\n",
    "You can see the model architecture in the diagram below:\n",
    "\n",
    "![whisper_architecture.svg](https://user-images.githubusercontent.com/29454499/204536571-8f6d8d77-5fbd-4c6d-8e29-14e734837860.svg)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are several models of different sizes and capabilities trained by the authors of the model. In this tutorial, we will use the `tiny` model, but the same actions are also applicable to other models from Whisper family."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6f1b0c699e444e7fa375cf5cc59d9a7c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dropdown(description='Model:', index=6, options=('openai/whisper-large-v3', 'openai/whisper-large-v2', 'openai…"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import ipywidgets as widgets\n",
    "\n",
    "MODELS = [\n",
    "    \"openai/whisper-large-v3\",\n",
    "    \"openai/whisper-large-v2\",\n",
    "    \"openai/whisper-large\",\n",
    "    \"openai/whisper-medium\",\n",
    "    \"openai/whisper-small\",\n",
    "    \"openai/whisper-base\",\n",
    "    \"openai/whisper-tiny\",\n",
    "]\n",
    "\n",
    "model_id = widgets.Dropdown(\n",
    "    options=list(MODELS),\n",
    "    value=\"openai/whisper-tiny\",\n",
    "    description=\"Model:\",\n",
    "    disabled=False,\n",
    ")\n",
    "\n",
    "model_id"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Convert model to OpenVINO Intermediate Representation (IR) format using Optimum-Intel.\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "The Hugging Face Optimum API is a high-level API that enables us to convert and quantize models from the Hugging Face Transformers library to the OpenVINO™ IR format. For more details, refer to the [Hugging Face Optimum documentation](https://huggingface.co/docs/optimum/intel/inference).\n",
    "\n",
    "Optimum Intel can be used to load optimized models from the [Hugging Face Hub](https://huggingface.co/docs/optimum/intel/hf.co/models) and create pipelines to run an inference with OpenVINO Runtime using Hugging Face APIs. The Optimum Inference models are API compatible with Hugging Face Transformers models.  This means we just need to replace the `AutoModelForXxx` class with the corresponding `OVModelForXxx` class.\n",
    "\n",
    "Below is an example of the whisper-tiny model\n",
    "\n",
    "```diff\n",
    "-from transformers import AutoModelForSpeechSeq2Seq\n",
    "+from optimum.intel.openvino import OVModelForSpeechSeq2Seq\n",
    "from transformers import AutoTokenizer, pipeline\n",
    "\n",
    "model_id = \"openai/whisper-tiny\"\n",
    "-model = AutoModelForSpeechSeq2Seq.from_pretrained(model_id)\n",
    "+model = OVModelForSpeechSeq2Seq.from_pretrained(model_id, export=True)\n",
    "```\n",
    "\n",
    "Model class initialization starts with calling the `from_pretrained` method. When downloading and converting the Transformers model, the parameter `export=True` should be added. We can save the converted model for the next usage with the `save_pretrained` method. Alternatively, model conversion can be performed using Optimum-CLI interface. You can find more details about Optimum-Intel and Optimum CLI usage in this [tutorial](../hugging-face-hub/hugging-face-hub.ipynb). The command bellow illustrates how to convert whisper using optimum cli.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "model_dir = model_id.value.split(\"/\")[-1]\n",
    "\n",
    "if not Path(model_dir).exists():\n",
    "    !optimum-cli export openvino -m {model_id.value} {model_dir} --weight-format fp16"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare inference pipeline\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "The image below illustrates the pipeline of video transcribing using the Whisper model.\n",
    "\n",
    "![whisper_pipeline.png](https://user-images.githubusercontent.com/29454499/204536733-1f4342f7-2328-476a-a431-cb596df69854.png)\n",
    "\n",
    "\n",
    "Preprocessing and post-processing are important in this model use. `transformers.AutoProcessor` class used for initialization `WhisperProcessor` is responsible for preparing audio input data for the PyTorch model, converting it to Mel-spectrogram and decoding predicted output token_ids into string using tokenizer. Tokenizers and Processors are distributed with models also compatible with the OpenVINO model.\n",
    "\n",
    "Like the original PyTorch model, the OpenVINO model is also compatible with HuggingFace [pipeline](https://huggingface.co/docs/transformers/main_classes/pipelines#transformers.AutomaticSpeechRecognitionPipeline) interface for `automatic-speech-recognition`. \n",
    "Pipeline can be used for long audio transcription. Distil-Whisper uses a chunked algorithm to transcribe long-form audio files. In practice, this chunked long-form algorithm is 9x faster than the sequential algorithm proposed by OpenAI in the Whisper paper. To enable chunking, pass the chunk_length_s parameter to the pipeline. For Distil-Whisper, a chunk length of 15 seconds is optimal. To activate batching, pass the argument batch_size."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Select inference device\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "select device from dropdown list for running inference using OpenVINO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "import openvino as ov\n",
    "\n",
    "core = ov.Core()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a37643c49958440285805210af50b2c2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dropdown(description='Device:', index=3, options=('CPU', 'GPU.0', 'GPU.1', 'AUTO'), value='AUTO')"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import requests\n",
    "\n",
    "r = requests.get(\n",
    "    url=\"https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/notebook_utils.py\",\n",
    ")\n",
    "open(\"notebook_utils.py\", \"w\").write(r.text)\n",
    "\n",
    "from notebook_utils import device_widget\n",
    "\n",
    "device = device_widget()\n",
    "\n",
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from optimum.intel.openvino import OVModelForSpeechSeq2Seq\n",
    "from transformers import AutoProcessor, pipeline\n",
    "\n",
    "ov_model = OVModelForSpeechSeq2Seq.from_pretrained(model_dir, device=device.value)\n",
    "\n",
    "processor = AutoProcessor.from_pretrained(model_dir)\n",
    "\n",
    "pipe = pipeline(\n",
    "    \"automatic-speech-recognition\",\n",
    "    model=ov_model,\n",
    "    chunk_length_s=30,\n",
    "    tokenizer=processor.tokenizer,\n",
    "    feature_extractor=processor.feature_extractor,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run video transcription pipeline\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "Now, we are ready to start transcription. We select a video from YouTube that we want to transcribe. Be patient, as downloading the video may take some time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f493dd18def948e5a50f5950b5e5c4ca",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Text(value='https://youtu.be/kgL5LBM-hFI', description='Video:', placeholder='Type link for video')"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import ipywidgets as widgets\n",
    "\n",
    "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",
    "\n",
    "link"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "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 pathlib import Path\n",
    "import yt_dlp\n",
    "\n",
    "print(f\"Downloading video {link.value} started\")\n",
    "\n",
    "output_file = Path(\"downloaded_video.mp4\")\n",
    "ydl_ops = {\"format\": \"best[ext=mp4]\", \"outtmpl\": output_file.as_posix()}\n",
    "with yt_dlp.YoutubeDL(ydl_ops) as ydl:\n",
    "    ydl.download(link.value)\n",
    "\n",
    "print(f\"Video saved to {output_file}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "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": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "692ad019edab4ffc81a9cf2808e50d36",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Select(description='Select task:', index=1, options=('transcribe', 'translate'), value='translate')"
      ]
     },
     "execution_count": 9,
     "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": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "from moviepy.editor import VideoFileClip\n",
    "from transformers.pipelines.audio_utils import ffmpeg_read\n",
    "\n",
    "\n",
    "def get_audio(video_file):\n",
    "    \"\"\"\n",
    "    Extract audio signal from a given video file, then convert it to float,\n",
    "    then mono-channel format and resample it to the expected sample rate\n",
    "\n",
    "    Parameters:\n",
    "        video_file: path to input video file\n",
    "    Returns:\n",
    "      resampled_audio: mono-channel float audio signal with 16000 Hz sample rate\n",
    "                       extracted from video\n",
    "      duration: duration of video fragment in seconds\n",
    "    \"\"\"\n",
    "    input_video = VideoFileClip(str(video_file))\n",
    "    duration = input_video.duration\n",
    "    audio_file = video_file.stem + \".wav\"\n",
    "    input_video.audio.write_audiofile(audio_file, verbose=False, logger=None)\n",
    "    with open(audio_file, \"rb\") as f:\n",
    "        inputs = f.read()\n",
    "    audio = ffmpeg_read(inputs, pipe.feature_extractor.sampling_rate)\n",
    "    return {\n",
    "        \"raw\": audio,\n",
    "        \"sampling_rate\": pipe.feature_extractor.sampling_rate,\n",
    "    }, duration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs, duration = get_audio(output_file)\n",
    "\n",
    "transcription = pipe(inputs, generate_kwargs={\"task\": task.value}, return_timestamps=True)[\"chunks\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "\n",
    "\n",
    "def format_timestamp(seconds: float):\n",
    "    \"\"\"\n",
    "    format time in srt-file expected format\n",
    "    \"\"\"\n",
    "    assert seconds >= 0, \"non-negative timestamp expected\"\n",
    "    milliseconds = round(seconds * 1000.0)\n",
    "\n",
    "    hours = milliseconds // 3_600_000\n",
    "    milliseconds -= hours * 3_600_000\n",
    "\n",
    "    minutes = milliseconds // 60_000\n",
    "    milliseconds -= minutes * 60_000\n",
    "\n",
    "    seconds = milliseconds // 1_000\n",
    "    milliseconds -= seconds * 1_000\n",
    "\n",
    "    return (f\"{hours}:\" if hours > 0 else \"00:\") + f\"{minutes:02d}:{seconds:02d},{milliseconds:03d}\"\n",
    "\n",
    "\n",
    "def prepare_srt(transcription, filter_duration=None):\n",
    "    \"\"\"\n",
    "    Format transcription into srt file format\n",
    "    \"\"\"\n",
    "    segment_lines = []\n",
    "    for idx, segment in enumerate(transcription):\n",
    "        # for the case where the model could not predict an ending timestamp, which can happen if audio is cut off in the middle of a word.\n",
    "        if segment[\"timestamp\"][1] is None:\n",
    "            segment[\"timestamp\"] = (segment[\"timestamp\"][0], filter_duration)\n",
    "\n",
    "        if filter_duration is not None and (segment[\"timestamp\"][0] >= math.floor(filter_duration) or segment[\"timestamp\"][1] > math.ceil(filter_duration) + 1):\n",
    "            break\n",
    "        segment_lines.append(str(idx + 1) + \"\\n\")\n",
    "        time_start = format_timestamp(segment[\"timestamp\"][0])\n",
    "        time_end = format_timestamp(segment[\"timestamp\"][1])\n",
    "        time_str = f\"{time_start} --> {time_end}\\n\"\n",
    "        segment_lines.append(time_str)\n",
    "        segment_lines.append(segment[\"text\"] + \"\\n\\n\")\n",
    "    return segment_lines"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\"The results will be saved in the `downloaded_video.srt` file. SRT is one of the most popular formats for storing subtitles and is compatible with many modern video players. This file can be used to embed transcription into videos during playback or by injecting them directly into video files using `ffmpeg`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "srt_lines = prepare_srt(transcription, filter_duration=duration)\n",
    "# save transcription\n",
    "with output_file.with_suffix(\".srt\").open(\"w\") as f:\n",
    "    f.writelines(srt_lines)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let us see the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8bcb6af81fc14c9a9e3e003b9a2a6e0f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Video(value=b\"\\x00\\x00\\x00\\x18ftypmp42\\x00\\x00\\x00\\x00isommp42\\x00\\x00:'moov\\x00\\x00\\x00lmvhd...\", height='800…"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "widgets.Video.from_file(output_file, loop=False, width=800, height=800)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1\n",
      "00:00:00,000 --> 00:00:05,000\n",
      " Oh, what's that?\n",
      "\n",
      "2\n",
      "00:00:05,000 --> 00:00:08,000\n",
      " Oh, wow.\n",
      "\n",
      "3\n",
      "00:00:08,000 --> 00:00:10,000\n",
      " Hello, humans.\n",
      "\n",
      "4\n",
      "00:00:13,000 --> 00:00:15,000\n",
      " Focus on me.\n",
      "\n",
      "5\n",
      "00:00:15,000 --> 00:00:17,000\n",
      " Focus on the guard.\n",
      "\n",
      "6\n",
      "00:00:17,000 --> 00:00:20,000\n",
      " Don't tell anyone what you're seeing in here.\n",
      "\n",
      "7\n",
      "00:00:22,000 --> 00:00:24,000\n",
      " Have you seen what's in there?\n",
      "\n",
      "8\n",
      "00:00:24,000 --> 00:00:25,000\n",
      " They have intel.\n",
      "\n",
      "9\n",
      "00:00:25,000 --> 00:00:27,000\n",
      " This is where it all changes.\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(\"\".join(srt_lines))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Quantization\n",
    "[back to top ⬆️](#Table-of-contents:)\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.\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 encoder and decoder models.\n",
    "3. Serialize the `INT8` model using `openvino.save_model` function.\n",
    "\n",
    ">**Note**: Quantization is time and memory consuming operation. Running quantization code below may take some time.\n",
    "\n",
    "Please select below whether you would like to run Whisper quantization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a21bfe2e9278413a9d1adaf94ac7388a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Checkbox(value=True, description='Quantization')"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "to_quantize = widgets.Checkbox(\n",
    "    value=True,\n",
    "    description=\"Quantization\",\n",
    "    disabled=False,\n",
    ")\n",
    "\n",
    "to_quantize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fetch `skip_kernel_extension` module\n",
    "import requests\n",
    "\n",
    "r = requests.get(\n",
    "    url=\"https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/skip_kernel_extension.py\",\n",
    ")\n",
    "open(\"skip_kernel_extension.py\", \"w\").write(r.text)\n",
    "\n",
    "ov_quantized_model = None\n",
    "\n",
    "%load_ext skip_kernel_extension"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prepare calibration datasets\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "First step is to prepare calibration datasets for quantization. Since we quantize whisper encoder and decoder separately, we need to prepare a calibration dataset for each of the models. We import an `InferRequestWrapper` class that will intercept model inputs and collect them to a list. Then we run model inference on some small amount of audio samples. Generally, increasing the calibration dataset size improves quantization quality."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%skip not $to_quantize.value\n",
    "\n",
    "from itertools import islice\n",
    "from optimum.intel.openvino.quantization import InferRequestWrapper\n",
    "\n",
    "\n",
    "def collect_calibration_dataset(ov_model: OVModelForSpeechSeq2Seq, calibration_dataset_size: int):\n",
    "    # Overwrite model request properties, saving the original ones for restoring later\n",
    "    encoder_calibration_data = []\n",
    "    decoder_calibration_data = []\n",
    "    ov_model.encoder.request = InferRequestWrapper(ov_model.encoder.request, encoder_calibration_data, apply_caching=True)\n",
    "    ov_model.decoder_with_past.request = InferRequestWrapper(ov_model.decoder_with_past.request,\n",
    "                                                             decoder_calibration_data,\n",
    "                                                             apply_caching=True)\n",
    "\n",
    "    pipe = pipeline(\n",
    "      \"automatic-speech-recognition\",\n",
    "      model=ov_model,\n",
    "      chunk_length_s=30,\n",
    "      tokenizer=processor.tokenizer,\n",
    "      feature_extractor=processor.feature_extractor)\n",
    "    try:\n",
    "        calibration_dataset = dataset = load_dataset(\"openslr/librispeech_asr\", \"clean\", split=\"validation\", streaming=True, trust_remote_code=True)\n",
    "        for sample in tqdm(islice(calibration_dataset, calibration_dataset_size), desc=\"Collecting calibration data\",\n",
    "                           total=calibration_dataset_size):\n",
    "            pipe(sample[\"audio\"], generate_kwargs={\"task\": task.value}, return_timestamps=True)\n",
    "    finally:\n",
    "        ov_model.encoder.request = ov_model.encoder.request.request\n",
    "        ov_model.decoder_with_past.request = ov_model.decoder_with_past.request.request\n",
    "\n",
    "    return encoder_calibration_data, decoder_calibration_data"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Quantize Whisper encoder and decoder models\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "Below we run the `quantize` function which calls `nncf.quantize` on Whisper encoder and decoder-with-past models. We don't quantize first-step-decoder because its share in whole inference time is negligible."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "148b4131e11f4363bcebfc0c78ed13df",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Collecting calibration data:   0%|          | 0/50 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e77ad19c904c464db194b5df82c7bfd6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Quantizing encoder\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"></pre>\n"
      ],
      "text/plain": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f30e8952c50e4f6ca374dc6972b95fca",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"></pre>\n"
      ],
      "text/plain": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:nncf:12 ignored nodes were found by name in the NNCFGraph\n",
      "INFO:nncf:16 ignored nodes were found by name in the NNCFGraph\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "520c7840eb5440859d2b9ba2123049a7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"></pre>\n"
      ],
      "text/plain": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5ddc5b2c750d4a4cbe67f0b8f7be4faf",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"></pre>\n"
      ],
      "text/plain": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "69846e3229834f7992738d50a33b354b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Quantizing decoder with past\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"></pre>\n"
      ],
      "text/plain": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4c48f03790324ee99afdb4031429a09a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"></pre>\n"
      ],
      "text/plain": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:nncf:24 ignored nodes were found by name in the NNCFGraph\n",
      "INFO:nncf:24 ignored nodes were found by name in the NNCFGraph\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "03249855c07f4b83bfb4289608bca05b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"></pre>\n"
      ],
      "text/plain": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fea73d3c378a442bacb43bf1ab11b4ec",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"></pre>\n"
      ],
      "text/plain": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Compiling the encoder to AUTO ...\n",
      "Compiling the decoder to AUTO ...\n",
      "Compiling the decoder to AUTO ...\n"
     ]
    }
   ],
   "source": [
    "%%skip not $to_quantize.value\n",
    "\n",
    "import gc\n",
    "import shutil\n",
    "import nncf\n",
    "from datasets import load_dataset\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "def extract_input_features(sample):\n",
    "    input_features = processor(\n",
    "        sample[\"audio\"][\"array\"],\n",
    "        sampling_rate=sample[\"audio\"][\"sampling_rate\"],\n",
    "        return_tensors=\"pt\",\n",
    "    ).input_features\n",
    "    return input_features\n",
    "\n",
    "\n",
    "\n",
    "CALIBRATION_DATASET_SIZE = 50\n",
    "quantized_model_path = Path(f\"{model_dir}_quantized\")\n",
    "\n",
    "\n",
    "def quantize(ov_model: OVModelForSpeechSeq2Seq, calibration_dataset_size: int):\n",
    "    if not quantized_model_path.exists():\n",
    "        encoder_calibration_data, decoder_calibration_data = collect_calibration_dataset(\n",
    "            ov_model, calibration_dataset_size\n",
    "        )\n",
    "        print(\"Quantizing encoder\")\n",
    "        quantized_encoder = nncf.quantize(\n",
    "            ov_model.encoder.model,\n",
    "            nncf.Dataset(encoder_calibration_data),\n",
    "            subset_size=len(encoder_calibration_data),\n",
    "            model_type=nncf.ModelType.TRANSFORMER,\n",
    "            # Smooth Quant algorithm reduces activation quantization error; optimal alpha value was obtained through grid search\n",
    "            advanced_parameters=nncf.AdvancedQuantizationParameters(smooth_quant_alpha=0.50)\n",
    "        )\n",
    "        ov.save_model(quantized_encoder, quantized_model_path / \"openvino_encoder_model.xml\")\n",
    "        del quantized_encoder\n",
    "        del encoder_calibration_data\n",
    "        gc.collect()\n",
    "\n",
    "        print(\"Quantizing decoder with past\")\n",
    "        quantized_decoder_with_past = nncf.quantize(\n",
    "            ov_model.decoder_with_past.model,\n",
    "            nncf.Dataset(decoder_calibration_data),\n",
    "            subset_size=len(decoder_calibration_data),\n",
    "            model_type=nncf.ModelType.TRANSFORMER,\n",
    "            # Smooth Quant algorithm reduces activation quantization error; optimal alpha value was obtained through grid search\n",
    "            advanced_parameters=nncf.AdvancedQuantizationParameters(smooth_quant_alpha=0.96)\n",
    "        )\n",
    "        ov.save_model(quantized_decoder_with_past, quantized_model_path / \"openvino_decoder_with_past_model.xml\")\n",
    "        del quantized_decoder_with_past\n",
    "        del decoder_calibration_data\n",
    "        gc.collect()\n",
    "\n",
    "        # Copy the config file and the first-step-decoder manually\n",
    "        model_path = Path(model_dir)\n",
    "        shutil.copy(model_path / \"config.json\", quantized_model_path / \"config.json\")\n",
    "        shutil.copy(model_path / \"generation_config.json\", quantized_model_path / \"generation_config.json\")\n",
    "        shutil.copy(model_path / \"openvino_decoder_model.xml\", quantized_model_path / \"openvino_decoder_model.xml\")\n",
    "        shutil.copy(model_path / \"openvino_decoder_model.bin\", quantized_model_path / \"openvino_decoder_model.bin\")\n",
    "\n",
    "    quantized_ov_model = OVModelForSpeechSeq2Seq.from_pretrained(quantized_model_path, compile=False)\n",
    "    quantized_ov_model.to(device.value)\n",
    "    quantized_ov_model.compile()\n",
    "    return quantized_ov_model\n",
    "\n",
    "\n",
    "ov_quantized_model = quantize(ov_model, CALIBRATION_DATASET_SIZE)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run quantized model inference\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "Let's compare the transcription results for original and quantized models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1\n",
      "00:00:00,000 --> 00:00:05,000\n",
      " What's that?\n",
      "\n",
      "2\n",
      "00:00:05,000 --> 00:00:07,000\n",
      " Oh, wow.\n",
      "\n",
      "3\n",
      "00:00:09,000 --> 00:00:11,000\n",
      " Hello humans.\n",
      "\n",
      "4\n",
      "00:00:14,000 --> 00:00:15,000\n",
      " Focus on me.\n",
      "\n",
      "5\n",
      "00:00:15,000 --> 00:00:16,000\n",
      " Focus on the guard.\n",
      "\n",
      "6\n",
      "00:00:18,000 --> 00:00:20,000\n",
      " Don't tell anyone what you're seen in here.\n",
      "\n",
      "7\n",
      "00:00:22,000 --> 00:00:24,000\n",
      " Have you seen what's in there?\n",
      "\n",
      "8\n",
      "00:00:24,000 --> 00:00:25,000\n",
      " They have intel.\n",
      "\n",
      "9\n",
      "00:00:25,000 --> 00:00:27,000\n",
      " This is where it all changes.\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "if ov_quantized_model is not None:\n",
    "    int8_pipe = pipeline(\n",
    "        \"automatic-speech-recognition\",\n",
    "        model=ov_quantized_model,\n",
    "        chunk_length_s=30,\n",
    "        tokenizer=processor.tokenizer,\n",
    "        feature_extractor=processor.feature_extractor,\n",
    "    )\n",
    "    inputs, duration = get_audio(output_file)\n",
    "    transcription = int8_pipe(inputs, generate_kwargs={\"task\": task.value}, return_timestamps=True)[\"chunks\"]\n",
    "    srt_lines = prepare_srt(transcription, filter_duration=duration)\n",
    "    print(\"\".join(srt_lines))\n",
    "    widgets.Video.from_file(output_file, loop=False, width=800, height=800)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compare performance and accuracy of the original and quantized models\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "Finally, we compare original and quantized Whisper models from accuracy and performance stand-points.\n",
    "\n",
    "To measure accuracy, we use `1 - WER` as a metric, where WER stands for Word Error Rate.\n",
    "\n",
    "When measuring inference time, we do it separately for encoder and decoder-with-past model forwards, and for the whole model inference too."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a2822b0dcd584fd2aa28e01c607926d0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Measuring performance and accuracy:   0%|          | 0/50 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3dc0e232e81c4de0ad82737f98f69d2a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Measuring performance and accuracy:   0%|          | 0/50 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Encoder performance speedup: 1.352\n",
      "Decoder with past performance speedup: 1.342\n",
      "Whole pipeline performance speedup: 1.350\n",
      "Whisper transcription word accuracy. Original model: 81.67%. Quantized model: 83.67%.\n",
      "Accuracy drop: -1.99%.\n"
     ]
    }
   ],
   "source": [
    "%%skip not $to_quantize.value\n",
    "\n",
    "import time\n",
    "from contextlib import contextmanager\n",
    "from jiwer import wer, wer_standardize\n",
    "\n",
    "\n",
    "TEST_DATASET_SIZE = 50\n",
    "MEASURE_TIME = False\n",
    "\n",
    "@contextmanager\n",
    "def time_measurement():\n",
    "    global MEASURE_TIME\n",
    "    try:\n",
    "        MEASURE_TIME = True\n",
    "        yield\n",
    "    finally:\n",
    "        MEASURE_TIME = False\n",
    "\n",
    "def time_fn(obj, fn_name, time_list):\n",
    "    original_fn = getattr(obj, fn_name)\n",
    "\n",
    "    def wrapper(*args, **kwargs):\n",
    "        if not MEASURE_TIME:\n",
    "            return original_fn(*args, **kwargs)\n",
    "        start_time = time.perf_counter()\n",
    "        result = original_fn(*args, **kwargs)\n",
    "        end_time = time.perf_counter()\n",
    "        time_list.append(end_time - start_time)\n",
    "        return result\n",
    "\n",
    "    setattr(obj, fn_name, wrapper)\n",
    "\n",
    "def calculate_transcription_time_and_accuracy(ov_model, test_samples):\n",
    "    encoder_infer_times = []\n",
    "    decoder_with_past_infer_times = []\n",
    "    whole_infer_times = []\n",
    "    time_fn(ov_model, \"generate\", whole_infer_times)\n",
    "    time_fn(ov_model.encoder, \"forward\", encoder_infer_times)\n",
    "    time_fn(ov_model.decoder_with_past, \"forward\", decoder_with_past_infer_times)\n",
    "\n",
    "    ground_truths = []\n",
    "    predictions = []\n",
    "    for data_item in tqdm(test_samples, desc=\"Measuring performance and accuracy\"):\n",
    "        input_features = extract_input_features(data_item)\n",
    "\n",
    "        with time_measurement():\n",
    "            predicted_ids = ov_model.generate(input_features)\n",
    "        transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)\n",
    "\n",
    "        ground_truths.append(data_item[\"text\"])\n",
    "        predictions.append(transcription[0])\n",
    "\n",
    "    word_accuracy = (1 - wer(ground_truths, predictions, reference_transform=wer_standardize,\n",
    "                             hypothesis_transform=wer_standardize)) * 100\n",
    "    mean_whole_infer_time = sum(whole_infer_times)\n",
    "    mean_encoder_infer_time = sum(encoder_infer_times)\n",
    "    mean_decoder_with_time_infer_time = sum(decoder_with_past_infer_times)\n",
    "    return word_accuracy, (mean_whole_infer_time, mean_encoder_infer_time, mean_decoder_with_time_infer_time)\n",
    "\n",
    "test_dataset = load_dataset(\"openslr/librispeech_asr\", \"clean\", split=\"validation\", streaming=True, trust_remote_code=True)\n",
    "test_dataset = test_dataset.shuffle(seed=42).take(TEST_DATASET_SIZE)\n",
    "test_samples = [sample for sample in test_dataset]\n",
    "\n",
    "accuracy_original, times_original = calculate_transcription_time_and_accuracy(ov_model, test_samples)\n",
    "accuracy_quantized, times_quantized = calculate_transcription_time_and_accuracy(ov_quantized_model, test_samples)\n",
    "print(f\"Encoder performance speedup: {times_original[1] / times_quantized[1]:.3f}\")\n",
    "print(f\"Decoder with past performance speedup: {times_original[2] / times_quantized[2]:.3f}\")\n",
    "print(f\"Whole pipeline performance speedup: {times_original[0] / times_quantized[0]:.3f}\")\n",
    "print(f\"Whisper transcription word accuracy. Original model: {accuracy_original:.2f}%. Quantized model: {accuracy_quantized:.2f}%.\")\n",
    "print(f\"Accuracy drop: {accuracy_original - accuracy_quantized:.2f}%.\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Interactive demo\n",
    "[back to top ⬆️](#Table-of-contents:)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "test_replace": {
     "    demo.launch(debug=True)": "    demo.launch()",
     "    demo.launch(share=True, debug=True)": "    demo.launch(share=True)"
    }
   },
   "outputs": [],
   "source": [
    "def transcribe(url, task, use_int8):\n",
    "    output_file = Path(\"downloaded_video.mp4\")\n",
    "    ydl_ops = {\"format\": \"best[ext=mp4]\", \"outtmpl\": output_file.as_posix()}\n",
    "    with yt_dlp.YoutubeDL(ydl_ops) as ydl:\n",
    "        ydl.download(link.value)\n",
    "    inputs, duration = get_audio(output_file)\n",
    "    m_pipe = int8_pipe if use_int8 else pipe\n",
    "    transcription = m_pipe(inputs, generate_kwargs={\"task\": task.lower()}, return_timestamps=True)[\"chunks\"]\n",
    "    srt_lines = prepare_srt(transcription, duration)\n",
    "    with output_file.with_suffix(\".srt\").open(\"w\") as f:\n",
    "        f.writelines(srt_lines)\n",
    "    return [str(output_file), str(output_file.with_suffix(\".srt\"))]\n",
    "\n",
    "\n",
    "if not Path(\"gradio_helper.py\").exists():\n",
    "    r = requests.get(url=\"https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/notebooks/whisper-subtitles-generation/gradio_helper.py\")\n",
    "    open(\"gradio_helper.py\", \"w\").write(r.text)\n",
    "\n",
    "from gradio_helper import make_demo\n",
    "\n",
    "demo = make_demo(fn=transcribe, quantized=ov_quantized_model is not None)\n",
    "\n",
    "try:\n",
    "    demo.launch(debug=True)\n",
    "except Exception:\n",
    "    demo.launch(share=True, debug=True)\n",
    "# if you are launching remotely, specify server_name and server_port\n",
    "# demo.launch(server_name='your server name', server_port='server port in int')\n",
    "# Read more in the docs: https://gradio.app/docs/"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.11.4"
  },
  "openvino_notebooks": {
   "imageUrl": "https://user-images.githubusercontent.com/29454499/204548693-1304ef33-c790-490d-8a8b-d5766acb6254.png",
   "tags": {
    "categories": [
     "Model Demos",
     "AI Trends"
    ],
    "libraries": [],
    "other": [],
    "tasks": [
     "Speech Recognition"
    ]
   }
  },
  "vscode": {
   "interpreter": {
    "hash": "916dbcbb3f70747c44a77c7bcd40155683ae19c65e1c03b4aa3499c5328201f1"
   }
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}