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 model with [OpenVINO Generate API](https://github.com/openvinotoolkit/openvino.genai) for [Whisper automatic speech recognition scenarios](https://github.com/openvinotoolkit/openvino.genai/blob/master/samples/python/whisper_speech_recognition/README.md) 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": [
    "#### 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": 3,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import platform\n",
    "\n",
    "%pip install -q \"nncf>=2.14.0\"\n",
    "%pip install -q -U \"openvino>=2024.5.0\" \"openvino-tokenizers>=2024.5.0\" \"openvino-genai>=2024.5.0\"\n",
    "%pip install -q \"python-ffmpeg<=1.0.16\" \"ffmpeg\" \"moviepy\" \"transformers>=4.45\" \"git+https://github.com/huggingface/optimum-intel.git\" \"torch>=2.1\" --extra-index-url https://download.pytorch.org/whl/cpu\n",
    "%pip install -q -U \"yt_dlp>=2024.8.6\" soundfile librosa jiwer\n",
    "%pip install -q  \"gradio>=4.19\" \"typing_extensions>=4.9\"\n",
    "\n",
    "if platform.system() == \"Darwin\":\n",
    "    %pip install -q \"numpy<2.0\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "from pathlib import Path\n",
    "\n",
    "if not Path(\"notebook_utils.py\").exists():\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",
    "if not Path(\"cmd_helper.py\").exists():\n",
    "    r = requests.get(\n",
    "        url=\"https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/cmd_helper.py\",\n",
    "    )\n",
    "    open(\"cmd_helper.py\", \"w\").write(r.text)"
   ]
  },
  {
   "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": "f12e1281744d4a769015824e4382c3c1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dropdown(description='Model:', index=7, options=('openai/whisper-large-v3-turbo', 'openai/whisper-large-v3', '…"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import ipywidgets as widgets\n",
    "\n",
    "MODELS = [\n",
    "    \"openai/whisper-large-v3-turbo\",\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",
    "Listed Whisper model are available for downloading via the [HuggingFace hub](https://huggingface.co/openai). We will use optimum-cli interface for exporting it into OpenVINO Intermediate Representation (IR) format.\n",
    "\n",
    "Optimum CLI interface for converting models supports export to OpenVINO (supported starting optimum-intel 1.12 version).\n",
    "General command format:\n",
    "\n",
    "```bash\n",
    "optimum-cli export openvino --model <model_id_or_path> --task <task> <output_dir>\n",
    "```\n",
    "\n",
    "where `--model` argument is model id from HuggingFace Hub or local directory with model (saved using `.save_pretrained` method), `--task ` is one of [supported task](https://huggingface.co/docs/optimum/exporters/task_manager) that exported model should solve. For LLMs it will be `automatic-speech-recognition-with-past`. If model initialization requires to use remote code, `--trust-remote-code` flag additionally should be passed. Full list of supported arguments available via `--help` For more details and examples of usage, please check [optimum documentation](https://huggingface.co/docs/optimum/intel/inference#export).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from cmd_helper import optimum_cli\n",
    "\n",
    "model_dir = model_id.value.split(\"/\")[-1]\n",
    "\n",
    "if not Path(model_dir).exists():\n",
    "    optimum_cli(model_id.value, model_dir)"
   ]
  },
  {
   "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",
    "\n",
    "To simplify user experience we will use [OpenVINO Generate API](https://github.com/openvinotoolkit/openvino.genai/blob/master/samples/python/whisper_speech_recognition/README.md). Firstly we will create pipeline with `WhisperPipeline`. You can construct it straight away from the folder with the converted model. It will automatically load the `model`, `tokenizer`, `detokenizer` and default `generation configuration`. "
   ]
  },
  {
   "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": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4edc715b0ede40168416a4514ddfeae0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dropdown(description='Device:', options=('CPU', 'AUTO'), value='CPU')"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from notebook_utils import device_widget\n",
    "\n",
    "device = device_widget(default=\"CPU\", exclude=[\"NPU\"])\n",
    "\n",
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openvino_genai as ov_genai\n",
    "\n",
    "ov_pipe = ov_genai.WhisperPipeline(str(model_dir), device=device.value)"
   ]
  },
  {
   "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. Let's load the video first."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "'downloaded_video.mp4' already exists.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "PosixPath('/home/labuser/work/notebook/openvino_notebooks/notebooks/whisper-subtitles-generation/downloaded_video.mp4')"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from notebook_utils import download_file\n",
    "\n",
    "output_file = Path(\"downloaded_video.mp4\")\n",
    "\n",
    "download_file(\n",
    "    \"https://storage.openvinotoolkit.org/repositories/openvino_notebooks/data/data/video/Sheldon%20Cooper%20Jim%20Parsons%20at%20Intels%20Lab.mp4\",\n",
    "    filename=output_file.name,\n",
    ")"
   ]
  },
  {
   "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": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "de3b9c7002264661a855fa84c3411de7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Select(description='Select task:', index=1, options=('transcribe', 'translate'), value='translate')"
      ]
     },
     "execution_count": 7,
     "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": 8,
   "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, 16000)\n",
    "    return {\n",
    "        \"raw\": audio,\n",
    "        \"sampling_rate\": 16000,\n",
    "    }, duration"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's run generation method. We will put input data as `np array`. Also we will specify `task` and `return_timestamps=True` options. If task is `translate`, you can place `language` option, for example `<|fr|>` for French or it would be detect automatically. We can set up generation parameters in different ways. We can get default config with `get_generation_config()`, setup parameters and put config directly to `generate()`. It's also possible to specify the needed options just as inputs in the `generate()` method and we will use this way. Then we just run `generate` method and get the output in text format.\n",
    "\n",
    "`generate` method with `return_timestamps` set to `True` will return `chunks`, which contain attributes: `text`, `start_ts` and `end_ts`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs, duration = get_audio(output_file)\n",
    "\n",
    "transcription = ov_pipe.generate(inputs[\"raw\"], task=task.value, return_timestamps=True).chunks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "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",
    "        timestamp = (segment.start_ts, segment.end_ts)\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.end_ts == -1:\n",
    "            timestamp[1] = filter_duration\n",
    "\n",
    "        if filter_duration is not None and (timestamp[0] >= math.floor(filter_duration) or timestamp[1] > math.ceil(filter_duration) + 1):\n",
    "            break\n",
    "        segment_lines.append(str(idx + 1) + \"\\n\")\n",
    "        time_start = format_timestamp(timestamp[0])\n",
    "        time_end = format_timestamp(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": 11,
   "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": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b27923d770014dde84ce789a0a6042ec",
       "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": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "widgets.Video.from_file(output_file, loop=False, width=800, height=800)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "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": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "14d51e7bb1104de1862ccb72adf305ce",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Checkbox(value=True, description='Quantization')"
      ]
     },
     "execution_count": 14,
     "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": 15,
   "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",
    "quantized_ov_pipe = None\n",
    "\n",
    "%load_ext skip_kernel_extension"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's load converted OpenVINO model format using Optimum-Intel to easily quantize it.\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) or local folder to 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",
    "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`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%skip not $to_quantize.value\n",
    "\n",
    "from transformers import AutoProcessor\n",
    "from optimum.intel.openvino import OVModelForSpeechSeq2Seq\n",
    "\n",
    "ov_model = OVModelForSpeechSeq2Seq.from_pretrained(model_dir, device=device.value)\n",
    "processor = AutoProcessor.from_pretrained(model_dir)"
   ]
  },
  {
   "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": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%skip not $to_quantize.value\n",
    "\n",
    "from itertools import islice\n",
    "from tqdm.notebook import tqdm\n",
    "from datasets import load_dataset\n",
    "from transformers import pipeline\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": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%skip not $to_quantize.value\n",
    "\n",
    "import gc\n",
    "import shutil\n",
    "import nncf\n",
    "import openvino as ov\n",
    "\n",
    "\n",
    "CALIBRATION_DATASET_SIZE = 30\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(ov_model, calibration_dataset_size)\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.80),\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",
    "        shutil.copy(model_path / \"openvino_tokenizer.xml\", quantized_model_path / \"openvino_tokenizer.xml\")\n",
    "        shutil.copy(model_path / \"openvino_tokenizer.bin\", quantized_model_path / \"openvino_tokenizer.bin\")\n",
    "        shutil.copy(model_path / \"openvino_detokenizer.xml\", quantized_model_path / \"openvino_detokenizer.xml\")\n",
    "        shutil.copy(model_path / \"openvino_detokenizer.bin\", quantized_model_path / \"openvino_detokenizer.bin\")\n",
    "        shutil.copy(model_path / \"tokenizer_config.json\", quantized_model_path / \"tokenizer_config.json\")\n",
    "        shutil.copy(model_path / \"tokenizer.json\", quantized_model_path / \"tokenizer.json\")\n",
    "        shutil.copy(model_path / \"vocab.json\", quantized_model_path / \"vocab.json\")\n",
    "        shutil.copy(model_path / \"preprocessor_config.json\", quantized_model_path / \"preprocessor_config.json\")\n",
    "        shutil.copy(model_path / \"special_tokens_map.json\", quantized_model_path / \"special_tokens_map.json\")\n",
    "        shutil.copy(model_path / \"normalizer.json\", quantized_model_path / \"normalizer.json\")\n",
    "        shutil.copy(model_path / \"merges.txt\", quantized_model_path / \"merges.txt\")\n",
    "        shutil.copy(model_path / \"added_tokens.json\", quantized_model_path / \"added_tokens.json\")\n",
    "\n",
    "    quantized_ov_pipe = ov_genai.WhisperPipeline(str(quantized_model_path), device=device.value)\n",
    "    return quantized_ov_pipe\n",
    "\n",
    "\n",
    "quantized_ov_pipe = 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": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "if ov_quantized_model is not None:\n",
    "    inputs, duration = get_audio(output_file)\n",
    "    transcription = quantized_ov_pipe.generate(inputs[\"raw\"], 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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d748cc7a3f394026907c0ef3780d58a8",
       "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": "c2377a7e2dd840328913bc77e6f51ca9",
       "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": [
      "Whole pipeline performance speedup: 1.452\n",
      "Whisper transcription word accuracy. Original model: 81.77%. Quantized model: 82.97%.\n",
      "Accuracy drop: -1.20%.\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",
    "TEST_DATASET_SIZE = 50\n",
    "\n",
    "def calculate_transcription_time_and_accuracy(ov_model, test_samples):\n",
    "    whole_infer_times = []\n",
    "\n",
    "    ground_truths = []\n",
    "    predictions = []\n",
    "    for data_item in tqdm(test_samples, desc=\"Measuring performance and accuracy\"):\n",
    "        start_time = time.perf_counter()\n",
    "        transcription = ov_model.generate(data_item[\"audio\"][\"array\"], return_timestamps=True)\n",
    "        end_time = time.perf_counter()\n",
    "        whole_infer_times.append(end_time - start_time)\n",
    "\n",
    "        ground_truths.append(data_item[\"text\"])\n",
    "        predictions.append(transcription.texts[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",
    "    return word_accuracy, mean_whole_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_pipe, test_samples)\n",
    "accuracy_quantized, times_quantized = calculate_transcription_time_and_accuracy(quantized_ov_pipe, test_samples)\n",
    "print(f\"Whole pipeline performance speedup: {times_original / times_quantized:.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": {
    "scrolled": true,
    "test_replace": {
     "    demo.launch(debug=True)": "    demo.launch()",
     "    demo.launch(share=True, debug=True)": "    demo.launch(share=True)"
    }
   },
   "outputs": [],
   "source": [
    "def_config = ov_pipe.get_generation_config()\n",
    "\n",
    "\n",
    "def transcribe(video_path, task, use_int8):\n",
    "    data_path = Path(video_path)\n",
    "    inputs, duration = get_audio(data_path)\n",
    "    m_pipe = quantized_ov_pipe if use_int8 else ov_pipe\n",
    "\n",
    "    frame_num = len(inputs[\"raw\"]) / 16000\n",
    "    if frame_num > 30:\n",
    "        config = ov_pipe.get_generation_config()\n",
    "        chink_num = math.ceil(frame_num / 30)\n",
    "        config.max_length = chink_num * def_config.max_length\n",
    "        m_pipe.set_generation_config(config)\n",
    "\n",
    "    transcription = m_pipe.generate(inputs[\"raw\"], task=task.lower(), return_timestamps=True).chunks\n",
    "    srt_lines = prepare_srt(transcription, duration)\n",
    "    with data_path.with_suffix(\".srt\").open(\"w\") as f:\n",
    "        f.writelines(srt_lines)\n",
    "    return [str(data_path), str(data_path.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, sample_path=output_file)\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 (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.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
}