notebooks/227-whisper-subtitles-generation/227-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.\n",
    "Notebook contains the following steps:\n",
    "1. Download the model.\n",
    "2. Instantiate the PyTorch model pipeline.\n",
    "3. Export the ONNX model and convert it to OpenVINO IR, using the Model Optimizer tool.\n",
    "4. Run the Whisper pipeline with OpenVINO models."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prerequisites"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Clone and install the model repository."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip available: \u001b[0m\u001b[31;49m22.3.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m23.0.1\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "!pip install -q \"python-ffmpeg<=1.0.16\" moviepy "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/ea/work/openvino_notebooks/notebooks/227-whisper-subtitles-generation/whisper\n",
      "/home/ea/work/notebooks_env/lib/python3.8/site-packages/setuptools/_distutils/dist.py:265: UserWarning: Unknown distribution option: 'readme'\n",
      "  warnings.warn(msg)\n",
      "running develop\n",
      "/home/ea/work/notebooks_env/lib/python3.8/site-packages/setuptools/command/easy_install.py:144: EasyInstallDeprecationWarning: easy_install command is deprecated. Use build and pip and other standards-based tools.\n",
      "  warnings.warn(\n",
      "/home/ea/work/notebooks_env/lib/python3.8/site-packages/setuptools/command/install.py:34: SetuptoolsDeprecationWarning: setup.py install is deprecated. Use build and pip and other standards-based tools.\n",
      "  warnings.warn(\n",
      "running egg_info\n",
      "writing openai_whisper.egg-info/PKG-INFO\n",
      "writing dependency_links to openai_whisper.egg-info/dependency_links.txt\n",
      "writing entry points to openai_whisper.egg-info/entry_points.txt\n",
      "writing requirements to openai_whisper.egg-info/requires.txt\n",
      "writing top-level names to openai_whisper.egg-info/top_level.txt\n",
      "file whisper.py (for module whisper) not found\n",
      "reading manifest file 'openai_whisper.egg-info/SOURCES.txt'\n",
      "reading manifest template 'MANIFEST.in'\n",
      "adding license file 'LICENSE'\n",
      "writing manifest file 'openai_whisper.egg-info/SOURCES.txt'\n",
      "running build_ext\n",
      "Creating /home/ea/work/notebooks_env/lib/python3.8/site-packages/openai-whisper.egg-link (link to .)\n",
      "openai-whisper 20230124 is already the active version in easy-install.pth\n",
      "Installing whisper script to /home/ea/work/notebooks_env/bin\n",
      "\n",
      "Installed /home/ea/work/openvino_notebooks/notebooks/227-whisper-subtitles-generation/whisper\n",
      "Processing dependencies for openai-whisper==20230124\n",
      "Searching for ffmpeg-python==0.2.0\n",
      "Best match: ffmpeg-python 0.2.0\n",
      "Processing ffmpeg_python-0.2.0-py3.8.egg\n",
      "ffmpeg-python 0.2.0 is already the active version in easy-install.pth\n",
      "\n",
      "Using /home/ea/work/notebooks_env/lib/python3.8/site-packages/ffmpeg_python-0.2.0-py3.8.egg\n",
      "Searching for transformers==4.28.0.dev0\n",
      "Best match: transformers 4.28.0.dev0\n",
      "transformers 4.28.0.dev0 is already the active version in easy-install.pth\n",
      "Installing transformers-cli script to /home/ea/work/notebooks_env/bin\n",
      "\n",
      "Using /home/ea/work/transformers/src\n",
      "Searching for more-itertools==9.0.0\n",
      "Best match: more-itertools 9.0.0\n",
      "Processing more_itertools-9.0.0-py3.8.egg\n",
      "more-itertools 9.0.0 is already the active version in easy-install.pth\n",
      "\n",
      "Using /home/ea/work/notebooks_env/lib/python3.8/site-packages/more_itertools-9.0.0-py3.8.egg\n",
      "Searching for tqdm==4.65.0\n",
      "Best match: tqdm 4.65.0\n",
      "Adding tqdm 4.65.0 to easy-install.pth file\n",
      "Installing tqdm script to /home/ea/work/notebooks_env/bin\n",
      "\n",
      "Using /home/ea/work/notebooks_env/lib/python3.8/site-packages\n",
      "Searching for torch==1.13.1+cpu\n",
      "Best match: torch 1.13.1+cpu\n",
      "Adding torch 1.13.1+cpu to easy-install.pth file\n",
      "Installing convert-caffe2-to-onnx script to /home/ea/work/notebooks_env/bin\n",
      "Installing convert-onnx-to-caffe2 script to /home/ea/work/notebooks_env/bin\n",
      "Installing torchrun script to /home/ea/work/notebooks_env/bin\n",
      "\n",
      "Using /home/ea/work/notebooks_env/lib/python3.8/site-packages\n",
      "Searching for numpy==1.23.1\n",
      "Best match: numpy 1.23.1\n",
      "Adding numpy 1.23.1 to easy-install.pth file\n",
      "Installing f2py script to /home/ea/work/notebooks_env/bin\n",
      "Installing f2py3 script to /home/ea/work/notebooks_env/bin\n",
      "Installing f2py3.8 script to /home/ea/work/notebooks_env/bin\n",
      "\n",
      "Using /home/ea/work/notebooks_env/lib/python3.8/site-packages\n",
      "Searching for future==0.18.3\n",
      "Best match: future 0.18.3\n",
      "Adding future 0.18.3 to easy-install.pth file\n",
      "Installing futurize script to /home/ea/work/notebooks_env/bin\n",
      "Installing pasteurize script to /home/ea/work/notebooks_env/bin\n",
      "\n",
      "Using /home/ea/work/notebooks_env/lib/python3.8/site-packages\n",
      "Searching for tokenizers==0.13.2\n",
      "Best match: tokenizers 0.13.2\n",
      "Adding tokenizers 0.13.2 to easy-install.pth file\n",
      "\n",
      "Using /home/ea/work/notebooks_env/lib/python3.8/site-packages\n",
      "Searching for requests==2.28.2\n",
      "Best match: requests 2.28.2\n",
      "Adding requests 2.28.2 to easy-install.pth file\n",
      "\n",
      "Using /home/ea/work/notebooks_env/lib/python3.8/site-packages\n",
      "Searching for regex==2022.10.31\n",
      "Best match: regex 2022.10.31\n",
      "Adding regex 2022.10.31 to easy-install.pth file\n",
      "\n",
      "Using /home/ea/work/notebooks_env/lib/python3.8/site-packages\n",
      "Searching for PyYAML==6.0\n",
      "Best match: PyYAML 6.0\n",
      "Adding PyYAML 6.0 to easy-install.pth file\n",
      "\n",
      "Using /home/ea/work/notebooks_env/lib/python3.8/site-packages\n",
      "Searching for packaging==23.0\n",
      "Best match: packaging 23.0\n",
      "Adding packaging 23.0 to easy-install.pth file\n",
      "\n",
      "Using /home/ea/work/notebooks_env/lib/python3.8/site-packages\n",
      "Searching for huggingface-hub==0.12.1\n",
      "Best match: huggingface-hub 0.12.1\n",
      "Adding huggingface-hub 0.12.1 to easy-install.pth file\n",
      "Installing huggingface-cli script to /home/ea/work/notebooks_env/bin\n",
      "\n",
      "Using /home/ea/work/notebooks_env/lib/python3.8/site-packages\n",
      "Searching for filelock==3.9.0\n",
      "Best match: filelock 3.9.0\n",
      "Adding filelock 3.9.0 to easy-install.pth file\n",
      "\n",
      "Using /home/ea/work/notebooks_env/lib/python3.8/site-packages\n",
      "Searching for typing-extensions==4.5.0\n",
      "Best match: typing-extensions 4.5.0\n",
      "Adding typing-extensions 4.5.0 to easy-install.pth file\n",
      "\n",
      "Using /home/ea/work/notebooks_env/lib/python3.8/site-packages\n",
      "Searching for certifi==2022.12.7\n",
      "Best match: certifi 2022.12.7\n",
      "Adding certifi 2022.12.7 to easy-install.pth file\n",
      "\n",
      "Using /home/ea/work/notebooks_env/lib/python3.8/site-packages\n",
      "Searching for urllib3==1.26.14\n",
      "Best match: urllib3 1.26.14\n",
      "Adding urllib3 1.26.14 to easy-install.pth file\n",
      "\n",
      "Using /home/ea/work/notebooks_env/lib/python3.8/site-packages\n",
      "Searching for idna==3.4\n",
      "Best match: idna 3.4\n",
      "Adding idna 3.4 to easy-install.pth file\n",
      "\n",
      "Using /home/ea/work/notebooks_env/lib/python3.8/site-packages\n",
      "Searching for charset-normalizer==2.0.12\n",
      "Best match: charset-normalizer 2.0.12\n",
      "Adding charset-normalizer 2.0.12 to easy-install.pth file\n",
      "Installing normalizer script to /home/ea/work/notebooks_env/bin\n",
      "\n",
      "Using /home/ea/work/notebooks_env/lib/python3.8/site-packages\n",
      "Finished processing dependencies for openai-whisper==20230124\n"
     ]
    }
   ],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "REPO_DIR = Path(\"whisper\")\n",
    "if not REPO_DIR.exists():\n",
    "    !git clone https://github.com/openai/whisper.git -b v20230124\n",
    "%cd whisper\n",
    "!python setup.py develop"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Instantiate model\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 `base` model, but the same actions are also applicable to other models from Whisper family."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import whisper\n",
    "\n",
    "model = whisper.load_model(\"base\")\n",
    "model.to(\"cpu\")\n",
    "model.eval()\n",
    "pass"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Convert model to OpenVINO Intermediate Representation (IR) format.\n",
    "\n",
    "For best results with OpenVINO, it is recommended to convert the model to OpenVINO IR format. OpenVINO supports PyTorch via ONNX conversion. We will use `torch.onnx.export` for exporting the ONNX model from PyTorch. We need to provide initialized model object and example of inputs for shape inference. We will use `mo.convert_model` functionality to convert the ONNX models. The `mo.convert_model` Python function returns an OpenVINO model ready to load on device and start making predictions. We can save it on disk for next usage with `openvino.runtime.serialize`.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Convert Whisper Encoder to OpenVINO IR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ea/work/openvino_notebooks/notebooks/227-whisper-subtitles-generation/whisper/whisper/model.py:153: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
      "  assert x.shape[1:] == self.positional_embedding.shape, \"incorrect audio shape\"\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from openvino.tools import mo\n",
    "from openvino.runtime import serialize\n",
    "\n",
    "mel = torch.zeros((1, 80, 3000))\n",
    "audio_features = model.encoder(mel)\n",
    "torch.onnx.export(\n",
    "    model.encoder, \n",
    "    mel, \n",
    "    \"whisper_encoder.onnx\",\n",
    "    input_names=[\"mel\"], \n",
    "    output_names=[\"output_features\"]\n",
    ")\n",
    "encoder_model = mo.convert_model(\"whisper_encoder.onnx\", compress_to_fp16=True)\n",
    "serialize(encoder_model, xml_path=\"whisper_encoder.xml\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Convert Whisper decoder to OpenVINO IR\n",
    "\n",
    "To reduce computational complexity, the decoder uses cached key/value projections in attention modules from the previous steps. We need to modify this process for correct tracing to ONNX."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from typing import Optional, Union, List, Dict\n",
    "from functools import partial\n",
    "\n",
    "positional_embeddings_size = model.decoder.positional_embedding.shape[0]\n",
    "\n",
    "\n",
    "def save_to_cache(cache: Dict[str, torch.Tensor], module: str, output: torch.Tensor):\n",
    "    \"\"\"\n",
    "    Saving cached attention hidden states for previous tokens.\n",
    "    Parameters:\n",
    "      cache: dictionary with cache.\n",
    "      module: current attention module name.\n",
    "      output: predicted hidden state.\n",
    "    Returns:\n",
    "      output: cached attention hidden state for specified attention module.\n",
    "    \"\"\"\n",
    "    if module not in cache or output.shape[1] > positional_embeddings_size:\n",
    "        # save as-is, for the first token or cross attention\n",
    "        cache[module] = output\n",
    "    else:\n",
    "        cache[module] = torch.cat([cache[module], output], dim=1).detach()\n",
    "    return cache[module]\n",
    "\n",
    "\n",
    "def attention_forward(\n",
    "        attention_module,\n",
    "        x: torch.Tensor,\n",
    "        xa: Optional[torch.Tensor] = None,\n",
    "        mask: Optional[torch.Tensor] = None,\n",
    "        kv_cache: Optional[dict] = None,\n",
    "        idx: int = 0\n",
    "):\n",
    "    \"\"\"\n",
    "    Override for forward method of decoder attention module with storing cache values explicitly.\n",
    "    Parameters:\n",
    "      attention_module: current attention module\n",
    "      x: input token ids.\n",
    "      xa: input audio features (Optional).\n",
    "      mask: mask for applying attention (Optional).\n",
    "      kv_cache: dictionary with cached key values for attention modules.\n",
    "      idx: idx for search in kv_cache.\n",
    "    Returns:\n",
    "      attention module output tensor\n",
    "      updated kv_cache\n",
    "    \"\"\"\n",
    "    q = attention_module.query(x)\n",
    "\n",
    "    if kv_cache is None or xa is None:\n",
    "        # hooks, if installed (i.e. kv_cache is not None), will prepend the cached kv tensors;\n",
    "        # otherwise, perform key/value projections for self- or cross-attention as usual.\n",
    "        k = attention_module.key(x if xa is None else xa)\n",
    "        v = attention_module.value(x if xa is None else xa)\n",
    "        if kv_cache is not None:\n",
    "            k = save_to_cache(kv_cache, f'k_{idx}', k)\n",
    "            v = save_to_cache(kv_cache, f'v_{idx}', v)\n",
    "    else:\n",
    "        # for cross-attention, calculate keys and values once and reuse in subsequent calls.\n",
    "        k = kv_cache.get(f'k_{idx}', save_to_cache(\n",
    "            kv_cache, f'k_{idx}', attention_module.key(xa)))\n",
    "        v = kv_cache.get(f'v_{idx}', save_to_cache(\n",
    "            kv_cache, f'v_{idx}', attention_module.value(xa)))\n",
    "\n",
    "    wv, qk = attention_module.qkv_attention(q, k, v, mask)\n",
    "    return attention_module.out(wv), kv_cache\n",
    "\n",
    "\n",
    "def block_forward(\n",
    "    residual_block,\n",
    "    x: torch.Tensor,\n",
    "    xa: Optional[torch.Tensor] = None,\n",
    "    mask: Optional[torch.Tensor] = None,\n",
    "    kv_cache: Optional[dict] = None,\n",
    "    idx: int = 0\n",
    "):\n",
    "    \"\"\"\n",
    "    Override for residual block forward method for providing kv_cache to attention module.\n",
    "      Parameters:\n",
    "        residual_block: current residual block.\n",
    "        x: input token_ids.\n",
    "        xa: input audio features (Optional).\n",
    "        mask: attention mask (Optional).\n",
    "        kv_cache: cache for storing attention key values.\n",
    "        idx: index of current residual block for search in kv_cache.\n",
    "      Returns:\n",
    "        x: residual block output\n",
    "        kv_cache: updated kv_cache\n",
    "\n",
    "    \"\"\"\n",
    "    x0, kv_cache = residual_block.attn(residual_block.attn_ln(\n",
    "        x), mask=mask, kv_cache=kv_cache, idx=f'{idx}a')\n",
    "    x = x + x0\n",
    "    if residual_block.cross_attn:\n",
    "        x1, kv_cache = residual_block.cross_attn(\n",
    "            residual_block.cross_attn_ln(x), xa, kv_cache=kv_cache, idx=f'{idx}c')\n",
    "        x = x + x1\n",
    "    x = x + residual_block.mlp(residual_block.mlp_ln(x))\n",
    "    return x, kv_cache\n",
    "\n",
    "\n",
    "# update forward functions\n",
    "for idx, block in enumerate(model.decoder.blocks):\n",
    "    block.forward = partial(block_forward, block, idx=idx)\n",
    "    block.attn.forward = partial(attention_forward, block.attn)\n",
    "    if block.cross_attn:\n",
    "        block.cross_attn.forward = partial(attention_forward, block.cross_attn)\n",
    "\n",
    "\n",
    "def decoder_forward(decoder, x: torch.Tensor, xa: torch.Tensor, kv_cache: Optional[dict] = None):\n",
    "    \"\"\"\n",
    "    Override for decoder forward method.\n",
    "    Parameters:\n",
    "      x: torch.LongTensor, shape = (batch_size, <= n_ctx) the text tokens\n",
    "      xa: torch.Tensor, shape = (batch_size, n_mels, n_audio_ctx)\n",
    "           the encoded audio features to be attended on\n",
    "      kv_cache: Dict[str, torch.Tensor], attention modules hidden states cache from previous steps \n",
    "    \"\"\"\n",
    "    offset = next(iter(kv_cache.values())).shape[1] if kv_cache else 0\n",
    "    x = decoder.token_embedding(\n",
    "        x) + decoder.positional_embedding[offset: offset + x.shape[-1]]\n",
    "    x = x.to(xa.dtype)\n",
    "\n",
    "    for block in decoder.blocks:\n",
    "        x, kv_cache = block(x, xa, mask=decoder.mask, kv_cache=kv_cache)\n",
    "\n",
    "    x = decoder.ln(x)\n",
    "    logits = (\n",
    "        x @ torch.transpose(decoder.token_embedding.weight.to(x.dtype), 1, 0)).float()\n",
    "\n",
    "    return logits, kv_cache\n",
    "\n",
    "\n",
    "# override decoder forward\n",
    "model.decoder.forward = partial(decoder_forward, model.decoder)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokens = torch.ones((5, 3), dtype=torch.int64)\n",
    "\n",
    "logits, kv_cache = model.decoder(tokens, audio_features, kv_cache={})\n",
    "kv_cache = {k: v for k, v in kv_cache.items()}\n",
    "tokens = torch.ones((5, 1), dtype=torch.int64)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_71955/1737529362.py:18: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
      "  if module not in cache or output.shape[1] > positional_embeddings_size:\n"
     ]
    }
   ],
   "source": [
    "outputs = [f\"out_{k}\" for k in kv_cache.keys()]\n",
    "inputs = [f\"in_{k}\" for k in kv_cache.keys()]\n",
    "dynamic_axes = {\n",
    "    \"tokens\": {0: \"beam_size\", 1: \"seq_len\"},\n",
    "    \"audio_features\": {0: \"beam_size\"},\n",
    "    \"logits\": {0: \"beam_size\", 1: \"seq_len\"}}\n",
    "dynamic_outs = {o: {0: \"beam_size\", 1: \"prev_seq_len\"} for o in outputs}\n",
    "dynamic_inp = {i: {0: \"beam_size\", 1: \"prev_seq_len\"} for i in inputs}\n",
    "dynamic_axes.update(dynamic_outs)\n",
    "dynamic_axes.update(dynamic_inp)\n",
    "torch.onnx.export(\n",
    "    model.decoder, {'x': tokens, 'xa': audio_features, 'kv_cache': kv_cache},\n",
    "    'whisper_decoder.onnx',\n",
    "    input_names=[\"tokens\", \"audio_features\"] + inputs,\n",
    "    output_names=[\"logits\"] + outputs,\n",
    "    dynamic_axes=dynamic_axes\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The decoder model autoregressively predicts the next token guided by encoder hidden states and previously predicted sequence. This means that the shape of inputs which depends on the previous step (inputs for tokens and attention hidden states from previous step) are dynamic. For efficient utilization of memory, you define an upper bound for dynamic input shapes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_shapes = \"tokens[1..5 1..224],audio_features[1..5 1500 512]\"\n",
    "for k, v in kv_cache.items():\n",
    "    if k.endswith('a'):\n",
    "        input_shapes += f\",in_{k}[1..5 0..224 512]\"\n",
    "decoder_model = mo.convert_model(\n",
    "    input_model=\"whisper_decoder.onnx\",\n",
    "    compress_to_fp16=True,\n",
    "    input=input_shapes)\n",
    "serialize(decoder_model, \"whisper_decoder.xml\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare inference pipeline\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",
    "To run the PyTorch Whisper model, we just need to call the `model.transcribe(audio, **parameters)` function. We will try to reuse original model pipeline for audio transcribing after replacing the original models with OpenVINO IR versions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class OpenVINOAudioEncoder(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    Helper for inference Whisper encoder model with OpenVINO\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, core, model_path, device='CPU'):\n",
    "        super().__init__()\n",
    "        self.model = core.read_model(model_path)\n",
    "        self.compiled_model = core.compile_model(self.model, device)\n",
    "        self.output_blob = self.compiled_model.output(0)\n",
    "\n",
    "    def forward(self, mel: torch.Tensor):\n",
    "        \"\"\"\n",
    "        Inference OpenVINO whisper encoder model.\n",
    "\n",
    "        Parameters:\n",
    "          mel: input audio fragment mel spectrogram.\n",
    "        Returns:\n",
    "          audio_features: torch tensor with encoded audio features.\n",
    "        \"\"\"\n",
    "        return torch.from_numpy(self.compiled_model(mel)[self.output_blob])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openvino.runtime import Core, Tensor\n",
    "\n",
    "\n",
    "class OpenVINOTextDecoder(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    Helper for inference OpenVINO decoder model\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, core: Core, model_path: Path, device: str = 'CPU'):\n",
    "        super().__init__()\n",
    "        self._core = core\n",
    "        self.model = core.read_model(model_path)\n",
    "        self._input_names = [inp.any_name for inp in self.model.inputs]\n",
    "        self.compiled_model = core.compile_model(self.model, device)\n",
    "        self.device = device\n",
    "\n",
    "    def init_past_inputs(self, feed_dict):\n",
    "        \"\"\"\n",
    "        Initialize cache input for first step.\n",
    "\n",
    "        Parameters:\n",
    "          feed_dict: Dictonary with inputs for inference\n",
    "        Returns:\n",
    "          feed_dict: updated feed_dict\n",
    "        \"\"\"\n",
    "        beam_size = feed_dict['tokens'].shape[0]\n",
    "        audio_len = feed_dict['audio_features'].shape[2]\n",
    "        previous_seq_len = 0\n",
    "        for name in self._input_names:\n",
    "            if name in ['tokens', 'audio_features']:\n",
    "                continue\n",
    "            feed_dict[name] = Tensor(np.zeros(\n",
    "                (beam_size, previous_seq_len, audio_len), dtype=np.float32))\n",
    "        return feed_dict\n",
    "\n",
    "    def preprocess_kv_cache_inputs(self, feed_dict, kv_cache):\n",
    "        \"\"\"\n",
    "        Transform kv_cache to inputs\n",
    "\n",
    "        Parameters:\n",
    "          feed_dict: dictionary with inputs for inference\n",
    "          kv_cache: dictionary with cached attention hidden states from previous step\n",
    "        Returns:\n",
    "          feed_dict: updated feed dictionary with additional inputs\n",
    "        \"\"\"\n",
    "        if not kv_cache:\n",
    "            return self.init_past_inputs(feed_dict)\n",
    "        for k, v in kv_cache.items():\n",
    "            new_k = f'in_{k}'\n",
    "            if new_k in self._input_names:\n",
    "                feed_dict[new_k] = Tensor(v.numpy())\n",
    "        return feed_dict\n",
    "\n",
    "    def postprocess_outputs(self, outputs):\n",
    "        \"\"\"\n",
    "        Transform model output to format expected by the pipeline\n",
    "\n",
    "        Parameters:\n",
    "          outputs: outputs: raw inference results.\n",
    "        Returns:\n",
    "          logits: decoder predicted token logits\n",
    "          kv_cache: cached attention hidden states\n",
    "        \"\"\"\n",
    "        logits = None\n",
    "        kv_cache = {}\n",
    "        for output_t, out in outputs.items():\n",
    "            if 'logits' in output_t.get_names():\n",
    "                logits = torch.from_numpy(out)\n",
    "            else:\n",
    "                tensor_name = output_t.any_name\n",
    "                kv_cache[tensor_name.replace(\n",
    "                    'out_', '')] = torch.from_numpy(out)\n",
    "        return logits, kv_cache\n",
    "\n",
    "    def forward(self, x: torch.Tensor, xa: torch.Tensor, kv_cache: Optional[dict] = None):\n",
    "        \"\"\"\n",
    "        Inference decoder model.\n",
    "\n",
    "        Parameters:\n",
    "          x: torch.LongTensor, shape = (batch_size, <= n_ctx) the text tokens\n",
    "          xa: torch.Tensor, shape = (batch_size, n_mels, n_audio_ctx)\n",
    "             the encoded audio features to be attended on\n",
    "          kv_cache: Dict[str, torch.Tensor], attention modules hidden states cache from previous steps\n",
    "        Returns:\n",
    "          logits: decoder predicted logits\n",
    "          kv_cache: updated kv_cache with current step hidden states\n",
    "        \"\"\"\n",
    "        feed_dict = {'tokens': Tensor(x.numpy()), 'audio_features': Tensor(xa.numpy())}\n",
    "        feed_dict = (self.preprocess_kv_cache_inputs(feed_dict, kv_cache))\n",
    "        res = self.compiled_model(feed_dict)\n",
    "        return self.postprocess_outputs(res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "from whisper.decoding import DecodingTask, Inference, DecodingOptions, DecodingResult\n",
    "\n",
    "\n",
    "class OpenVINOInference(Inference):\n",
    "    \"\"\"\n",
    "    Wrapper for inference interface\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, model: \"Whisper\", initial_token_length: int):\n",
    "        self.model: \"Whisper\" = model\n",
    "        self.initial_token_length = initial_token_length\n",
    "        self.kv_cache = {}\n",
    "\n",
    "    def logits(self, tokens: torch.Tensor, audio_features: torch.Tensor) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        getting logits for given tokens sequence and audio features and save kv_cache\n",
    "\n",
    "        Parameters:\n",
    "          tokens: input tokens\n",
    "          audio_features: input audio features\n",
    "        Returns:\n",
    "          logits: predicted by decoder logits\n",
    "        \"\"\"\n",
    "        if tokens.shape[-1] > self.initial_token_length:\n",
    "            # only need to use the last token except in the first forward pass\n",
    "            tokens = tokens[:, -1:]\n",
    "        logits, self.kv_cache = self.model.decoder(\n",
    "            tokens, audio_features, kv_cache=self.kv_cache)\n",
    "        return logits\n",
    "\n",
    "    def cleanup_caching(self):\n",
    "        \"\"\"\n",
    "        Reset kv_cache to initial state\n",
    "        \"\"\"\n",
    "        self.kv_cache = {}\n",
    "\n",
    "    def rearrange_kv_cache(self, source_indices):\n",
    "        \"\"\"\n",
    "        Update hidden states cache for selected sequences\n",
    "        Parameters:\n",
    "          source_indicies: sequences indicies\n",
    "        Returns:\n",
    "          None\n",
    "        \"\"\"\n",
    "        for module, tensor in self.kv_cache.items():\n",
    "            # update the key/value cache to contain the selected sequences\n",
    "            self.kv_cache[module] = tensor[source_indices]\n",
    "\n",
    "\n",
    "class OpenVINODecodingTask(DecodingTask):\n",
    "    \"\"\"\n",
    "    Class for decoding using OpenVINO\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, model: \"Whisper\", options: DecodingOptions):\n",
    "        super().__init__(model, options)\n",
    "        self.inference = OpenVINOInference(model, len(self.initial_tokens))\n",
    "\n",
    "\n",
    "@torch.no_grad()\n",
    "def decode(model: \"Whisper\", mel: torch.Tensor, options: DecodingOptions = DecodingOptions()) -> Union[DecodingResult, List[DecodingResult]]:\n",
    "    \"\"\"\n",
    "    Performs decoding of 30-second audio segment(s), provided as Mel spectrogram(s).\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    model: Whisper\n",
    "        the Whisper model instance\n",
    "\n",
    "    mel: torch.Tensor, shape = (80, 3000) or (*, 80, 3000)\n",
    "        A tensor containing the Mel spectrogram(s)\n",
    "\n",
    "    options: DecodingOptions\n",
    "        A dataclass that contains all necessary options for decoding 30-second segments\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    result: Union[DecodingResult, List[DecodingResult]]\n",
    "        The result(s) of decoding contained in `DecodingResult` dataclass instance(s)\n",
    "    \"\"\"\n",
    "    single = mel.ndim == 2\n",
    "    if single:\n",
    "        mel = mel.unsqueeze(0)\n",
    "\n",
    "    result = OpenVINODecodingTask(model, options).run(mel)\n",
    "\n",
    "    if single:\n",
    "        result = result[0]\n",
    "\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "del model.decoder\n",
    "del model.encoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import namedtuple\n",
    "\n",
    "Parameter = namedtuple('Parameter', ['device'])\n",
    "\n",
    "core = Core()\n",
    "\n",
    "model.encoder = OpenVINOAudioEncoder(core, 'whisper_encoder.xml')\n",
    "model.decoder = OpenVINOTextDecoder(core, 'whisper_decoder.xml')\n",
    "model.decode = partial(decode, model)\n",
    "\n",
    "\n",
    "def parameters():\n",
    "    return iter([Parameter(torch.device('cpu'))])\n",
    "\n",
    "\n",
    "model.parameters = parameters\n",
    "\n",
    "\n",
    "def logits(model, tokens: torch.Tensor, audio_features: torch.Tensor):\n",
    "    \"\"\"\n",
    "    Override for logits extraction method\n",
    "    Parameters:\n",
    "      toekns: input tokens\n",
    "      audio_features: input audio features\n",
    "    Returns:\n",
    "      logits: decoder predicted logits\n",
    "    \"\"\"\n",
    "    return model.decoder(tokens, audio_features, None)[0]\n",
    "\n",
    "\n",
    "model.logits = partial(logits, model)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Define audio preprocessing\n",
    "\n",
    "The model expects mono-channel audio with a 16000 Hz sample rate, represented in floating point range. When the audio from the input video does not meet these requirements, we will need to apply preprocessing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "import io\n",
    "from pathlib import Path\n",
    "import numpy as np\n",
    "from scipy.io import wavfile\n",
    "from pytube import YouTube\n",
    "from moviepy.editor import VideoFileClip\n",
    "\n",
    "\n",
    "def resample(audio, src_sample_rate, dst_sample_rate):\n",
    "    \"\"\"\n",
    "    Resample audio to specific sample rate\n",
    "\n",
    "    Parameters:\n",
    "      audio: input audio signal\n",
    "      src_sample_rate: source audio sample rate\n",
    "      dst_sample_rate: destination audio sample rate\n",
    "    Returns:\n",
    "      resampled_audio: input audio signal resampled with dst_sample_rate\n",
    "    \"\"\"\n",
    "    if src_sample_rate == dst_sample_rate:\n",
    "        return audio\n",
    "    duration = audio.shape[0] / src_sample_rate\n",
    "    resampled_data = np.zeros(shape=(int(duration * dst_sample_rate)), dtype=np.float32)\n",
    "    x_old = np.linspace(0, duration, audio.shape[0], dtype=np.float32)\n",
    "    x_new = np.linspace(0, duration, resampled_data.shape[0], dtype=np.float32)\n",
    "    resampled_audio = np.interp(x_new, x_old, audio)\n",
    "    return resampled_audio.astype(np.float32)\n",
    "\n",
    "\n",
    "def audio_to_float(audio):\n",
    "    \"\"\"\n",
    "    convert audio signal to floating point format\n",
    "    \"\"\"\n",
    "    return audio.astype(np.float32) / np.iinfo(audio.dtype).max\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",
    "    \"\"\"\n",
    "    input_video = VideoFileClip(str(video_file))\n",
    "    input_video.audio.write_audiofile(video_file.stem + '.wav', verbose=False, logger=None)\n",
    "    input_audio_file = video_file.stem + '.wav'\n",
    "    sample_rate, audio = wavfile.read(\n",
    "        io.BytesIO(open(input_audio_file, 'rb').read()))\n",
    "    audio = audio_to_float(audio)\n",
    "    if audio.ndim == 2:\n",
    "        audio = audio.mean(axis=1)\n",
    "    resampled_audio = resample(audio, sample_rate, 16000)\n",
    "    return resampled_audio"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run video transcription pipeline\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": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e7e96dd24d5746eb91b73cd5d75a1fee",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Text(value='https://www.youtube.com/watch?v=kgL5LBM-hFI', description='Video:', placeholder='Type link for vid…"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import ipywidgets as widgets\n",
    "VIDEO_LINK = \"https://www.youtube.com/watch?v=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": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading video https://www.youtube.com/watch?v=kgL5LBM-hFI started\n",
      "Video saved to downloaded_video.mp4\n"
     ]
    }
   ],
   "source": [
    "print(f\"Downloading video {link.value} started\")\n",
    "\n",
    "output_file = Path(\"downloaded_video.mp4\")\n",
    "yt = YouTube(link.value)\n",
    "yt.streams.get_highest_resolution().download(filename=output_file)\n",
    "print(f\"Video saved to {output_file}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "audio = get_audio(output_file)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Select the task for the model:\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": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cfbad1b86e134c99ae5f4daaee3f8615",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Select(description='Select task:', index=1, options=('transcribe', 'translate'), value='translate')"
      ]
     },
     "execution_count": 18,
     "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": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "transcription = model.transcribe(audio, beam_size=5, best_of=5, task=task.value)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "def format_timestamp(seconds: float):\n",
    "    \"\"\"\n",
    "    format time in srt-file excpected 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):\n",
    "    \"\"\"\n",
    "    Format transcription into srt file format\n",
    "    \"\"\"\n",
    "    segment_lines = []\n",
    "    for segment in transcription[\"segments\"]:\n",
    "        segment_lines.append(str(segment[\"id\"] + 1) + \"\\n\")\n",
    "        time_start = format_timestamp(segment[\"start\"])\n",
    "        time_end = format_timestamp(segment[\"end\"])\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"
   ]
  },
  {
   "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": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "srt_lines = prepare_srt(transcription)\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": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8d3f89b9ea764abcb2c1d1b3e807e46b",
       "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": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "widgets.Video.from_file(output_file, loop=False, width=800, height=800)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "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:09,000\n",
      " Oh, wow.\n",
      "\n",
      "3\n",
      "00:00:09,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:18,000\n",
      " Focus on the guard.\n",
      "\n",
      "6\n",
      "00:00:18,000 --> 00:00:23,000\n",
      " Don't tell anyone what you've seen in here.\n",
      "\n",
      "7\n",
      "00:00:23,000 --> 00:00:30,000\n",
      " They have intel, what's\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(\"\".join(srt_lines))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  },
  "vscode": {
   "interpreter": {
    "hash": "916dbcbb3f70747c44a77c7bcd40155683ae19c65e1c03b4aa3499c5328201f1"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}