Video Subtitle Generation using Whisper and OpenVINO™

This tutorial is also available as a Jupyter notebook that can be cloned directly from GitHub. See the installation guide for instructions to run this tutorial locally on Windows, Linux or macOS.

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.

asr-training-data-desktop.svg

You can find more information about this model in the research paper, OpenAI blog, model card and GitHub repository.

In this notebook we will use Whisper with OpenVINO to generate subtitles in a sample video. Notebook contains following steps: 1. Download model 2. Instantiate PyTorch model pipeline 3. Export ONNX model and convert to OpenVINO IR using the Model Optimizer tool 4. Run Whisper pipeline with OpenVINO models.

Prerequisites

Clone and install model repository

!pip install "python-ffmpeg<=1.0.16" moviepy

Collecting python-ffmpeg<=1.0.16
  Using cached python_ffmpeg-1.0.16-py3-none-any.whl (6.2 kB)
Collecting moviepy
  Using cached moviepy-1.0.3-py3-none-any.whl
Collecting pyee
  Using cached pyee-9.0.4-py2.py3-none-any.whl (14 kB)
Requirement already satisfied: imageio<3.0,>=2.5 in /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (from moviepy) (2.9.0)
Requirement already satisfied: requests<3.0,>=2.8.1 in /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (from moviepy) (2.28.1)
Requirement already satisfied: numpy>=1.17.3 in /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (from moviepy) (1.23.4)
Collecting decorator<5.0,>=4.0.2
  Using cached decorator-4.4.2-py2.py3-none-any.whl (9.2 kB)
Requirement already satisfied: tqdm<5.0,>=4.11.2 in /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (from moviepy) (4.65.0)
Collecting proglog<=1.0.0
  Using cached proglog-0.1.10-py3-none-any.whl (6.1 kB)
Requirement already satisfied: imageio-ffmpeg>=0.2.0 in /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (from moviepy) (0.4.8)
Requirement already satisfied: pillow in /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (from imageio<3.0,>=2.5->moviepy) (9.4.0)
Requirement already satisfied: idna<4,>=2.5 in /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (from requests<3.0,>=2.8.1->moviepy) (3.4)
Requirement already satisfied: urllib3<1.27,>=1.21.1 in /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (from requests<3.0,>=2.8.1->moviepy) (1.26.14)
Requirement already satisfied: certifi>=2017.4.17 in /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (from requests<3.0,>=2.8.1->moviepy) (2022.12.7)
Requirement already satisfied: charset-normalizer<3,>=2 in /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (from requests<3.0,>=2.8.1->moviepy) (2.1.1)
Requirement already satisfied: typing-extensions in /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages (from pyee->python-ffmpeg<=1.0.16) (4.5.0)
Installing collected packages: pyee, proglog, decorator, python-ffmpeg, moviepy
  Attempting uninstall: decorator
    Found existing installation: decorator 5.1.1
    Uninstalling decorator-5.1.1:
      Successfully uninstalled decorator-5.1.1
Successfully installed decorator-4.4.2 moviepy-1.0.3 proglog-0.1.10 pyee-9.0.4 python-ffmpeg-1.0.16

from pathlib import Path

REPO_DIR = Path("whisper")
if not REPO_DIR.exists():
    !git clone https://github.com/openai/whisper.git -b v20230124
%cd whisper
!python setup.py develop

Cloning into 'whisper'...
remote: Enumerating objects: 550, done.
remote: Counting objects: 100% (169/169), done.
remote: Compressing objects: 100% (100/100), done.
remote: Total 550 (delta 86), reused 121 (delta 63), pack-reused 381
Receiving objects: 100% (550/550), 8.15 MiB | 4.67 MiB/s, done.
Resolving deltas: 100% (306/306), done.
Note: switching to '55f690af7914c672c69733b7e04ef5a41b2b2774'.

You are in 'detached HEAD' state. You can look around, make experimental
changes and commit them, and you can discard any commits you make in this
state without impacting any branches by switching back to a branch.

If you want to create a new branch to retain commits you create, you may
do so (now or later) by using -c with the switch command. Example:

  git switch -c <new-branch-name>

Or undo this operation with:

  git switch -

Turn off this advice by setting config variable advice.detachedHead to false

/opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/notebooks/227-whisper-subtitles-generation/whisper
/opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages/setuptools/_distutils/dist.py:265: UserWarning: Unknown distribution option: 'readme'
  warnings.warn(msg)
running develop
/opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/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.
  warnings.warn(
/opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/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.
  warnings.warn(
running egg_info
creating openai_whisper.egg-info
writing openai_whisper.egg-info/PKG-INFO
writing dependency_links to openai_whisper.egg-info/dependency_links.txt
writing entry points to openai_whisper.egg-info/entry_points.txt
writing requirements to openai_whisper.egg-info/requires.txt
writing top-level names to openai_whisper.egg-info/top_level.txt
writing manifest file 'openai_whisper.egg-info/SOURCES.txt'
file whisper.py (for module whisper) not found
reading manifest file 'openai_whisper.egg-info/SOURCES.txt'
reading manifest template 'MANIFEST.in'
adding license file 'LICENSE'
writing manifest file 'openai_whisper.egg-info/SOURCES.txt'
running build_ext
Creating /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages/openai-whisper.egg-link (link to .)
Adding openai-whisper 20230124 to easy-install.pth file
Installing whisper script to /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/bin

Installed /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/notebooks/227-whisper-subtitles-generation/whisper
Processing dependencies for openai-whisper==20230124
Searching for ffmpeg-python==0.2.0
Reading https://pypi.org/simple/ffmpeg-python/
Downloading https://files.pythonhosted.org/packages/d7/0c/56be52741f75bad4dc6555991fabd2e07b432d333da82c11ad701123888a/ffmpeg_python-0.2.0-py3-none-any.whl#sha256=ac441a0404e053f8b6a1113a77c0f452f1cfc62f6344a769475ffdc0f56c23c5
Best match: ffmpeg-python 0.2.0
Processing ffmpeg_python-0.2.0-py3-none-any.whl
Installing ffmpeg_python-0.2.0-py3-none-any.whl to /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages
Adding ffmpeg-python 0.2.0 to easy-install.pth file

Installed /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages/ffmpeg_python-0.2.0-py3.8.egg
Searching for more-itertools
Reading https://pypi.org/simple/more-itertools/
Downloading https://files.pythonhosted.org/packages/85/01/e2678ee4e0d7eed4fd6be9e5b043fff9d22d245d06c8c91def8ced664189/more_itertools-9.1.0-py3-none-any.whl#sha256=d2bc7f02446e86a68911e58ded76d6561eea00cddfb2a91e7019bbb586c799f3
Best match: more-itertools 9.1.0
Processing more_itertools-9.1.0-py3-none-any.whl
Installing more_itertools-9.1.0-py3-none-any.whl to /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages
Adding more-itertools 9.1.0 to easy-install.pth file

Installed /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages/more_itertools-9.1.0-py3.8.egg
Searching for transformers==4.26.1
Best match: transformers 4.26.1
Adding transformers 4.26.1 to easy-install.pth file
Installing transformers-cli script to /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/bin

Using /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages
Searching for tqdm==4.65.0
Best match: tqdm 4.65.0
Adding tqdm 4.65.0 to easy-install.pth file
Installing tqdm script to /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/bin

Using /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages
Searching for torch==1.13.1+cpu
Best match: torch 1.13.1+cpu
Adding torch 1.13.1+cpu to easy-install.pth file
Installing convert-caffe2-to-onnx script to /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/bin
Installing convert-onnx-to-caffe2 script to /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/bin
Installing torchrun script to /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/bin

Using /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages
Searching for numpy==1.23.4
Best match: numpy 1.23.4
Adding numpy 1.23.4 to easy-install.pth file
Installing f2py script to /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/bin
Installing f2py3 script to /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/bin
Installing f2py3.8 script to /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/bin

Using /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages
Searching for future==0.18.3
Best match: future 0.18.3
Adding future 0.18.3 to easy-install.pth file
Installing futurize script to /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/bin
Installing pasteurize script to /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/bin

Using /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages
Searching for tokenizers==0.13.2
Best match: tokenizers 0.13.2
Adding tokenizers 0.13.2 to easy-install.pth file

Using /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages
Searching for requests==2.28.1
Best match: requests 2.28.1
Adding requests 2.28.1 to easy-install.pth file

Using /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages
Searching for regex==2022.10.31
Best match: regex 2022.10.31
Adding regex 2022.10.31 to easy-install.pth file

Using /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages
Searching for PyYAML==6.0
Best match: PyYAML 6.0
Adding PyYAML 6.0 to easy-install.pth file

Using /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages
Searching for packaging==23.0
Best match: packaging 23.0
Adding packaging 23.0 to easy-install.pth file

Using /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages
Searching for huggingface-hub==0.13.1
Best match: huggingface-hub 0.13.1
Adding huggingface-hub 0.13.1 to easy-install.pth file
Installing huggingface-cli script to /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/bin

Using /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages
Searching for filelock==3.9.0
Best match: filelock 3.9.0
Adding filelock 3.9.0 to easy-install.pth file

Using /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages
Searching for typing-extensions==4.5.0
Best match: typing-extensions 4.5.0
Adding typing-extensions 4.5.0 to easy-install.pth file

Using /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages
Searching for certifi==2022.12.7
Best match: certifi 2022.12.7
Adding certifi 2022.12.7 to easy-install.pth file

Using /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages
Searching for urllib3==1.26.14
Best match: urllib3 1.26.14
Adding urllib3 1.26.14 to easy-install.pth file

Using /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages
Searching for idna==3.4
Best match: idna 3.4
Adding idna 3.4 to easy-install.pth file

Using /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages
Searching for charset-normalizer==2.1.1
Best match: charset-normalizer 2.1.1
Adding charset-normalizer 2.1.1 to easy-install.pth file
Installing normalizer script to /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/bin

Using /opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages
Finished processing dependencies for openai-whisper==20230124

Instantiate model

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. The Transformer encoder then 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.

You can see the model architecture in the diagram below:

whisper_architecture.svg

There are several models of different sizes and capabilities trained by the model authors. In this tutorial we will use the base model, but the same actions are also applicable to other models from Whisper family.

import whisper

model = whisper.load_model("base")
model.to("cpu")
model.eval()
pass

2023-03-09 23:37:49.286424: I tensorflow/core/util/util.cc:169] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable TF_ENABLE_ONEDNN_OPTS=0.

Convert model to OpenVINO Intermediate Representation (IR) format.

For best results with OpenVINO we recommend converting 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.

Convert Whisper Encoder to OpenVINO IR

import torch
from openvino.tools import mo
from openvino.runtime import serialize

mel = torch.zeros((1, 80, 3000))
audio_features = model.encoder(mel)
torch.onnx.export(
    model.encoder,
    mel,
    "whisper_encoder.onnx",
    input_names=["mel"],
    output_names=["output_features"]
)
encoder_model = mo.convert_model("whisper_encoder.onnx", compress_to_fp16=True)
serialize(encoder_model, xml_path="whisper_encoder.xml")

/opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-358/.workspace/scm/ov-notebook/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!
  assert x.shape[1:] == self.positional_embedding.shape, "incorrect audio shape"

Convert Whisper decoder to OpenVINO IR

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.

import torch
from typing import Optional, Union, List, Dict
from functools import partial

positional_embeddings_size = model.decoder.positional_embedding.shape[0]


def save_to_cache(cache: Dict[str, torch.Tensor], module: str, output: torch.Tensor):
    """
    Saving cached attention hidden states for previous tokens.
    Parameters:
      cache: dictionary with cache.
      module: current attention module name.
      output: predicted hidden state.
    Returns:
      output: cached attention hidden state for specified attention module.
    """
    if module not in cache or output.shape[1] > positional_embeddings_size:
        # save as-is, for the first token or cross attention
        cache[module] = output
    else:
        cache[module] = torch.cat([cache[module], output], dim=1).detach()
    return cache[module]


def attention_forward(
        attention_module,
        x: torch.Tensor,
        xa: Optional[torch.Tensor] = None,
        mask: Optional[torch.Tensor] = None,
        kv_cache: Optional[dict] = None,
        idx: int = 0
):
    """
    Override for forward method of decoder attention module with storing cache values explicitly.
    Parameters:
      attention_module: current attention module
      x: input token ids.
      xa: input audio features (Optional).
      mask: mask for applying attention (Optional).
      kv_cache: dictionary with cached key values for attention modules.
      idx: idx for search in kv_cache.
    Returns:
      attention module output tensor
      updated kv_cache
    """
    q = attention_module.query(x)

    if kv_cache is None or xa is None:
        # hooks, if installed (i.e. kv_cache is not None), will prepend the cached kv tensors;
        # otherwise, perform key/value projections for self- or cross-attention as usual.
        k = attention_module.key(x if xa is None else xa)
        v = attention_module.value(x if xa is None else xa)
        if kv_cache is not None:
            k = save_to_cache(kv_cache, f'k_{idx}', k)
            v = save_to_cache(kv_cache, f'v_{idx}', v)
    else:
        # for cross-attention, calculate keys and values once and reuse in subsequent calls.
        k = kv_cache.get(f'k_{idx}', save_to_cache(
            kv_cache, f'k_{idx}', attention_module.key(xa)))
        v = kv_cache.get(f'v_{idx}', save_to_cache(
            kv_cache, f'v_{idx}', attention_module.value(xa)))

    wv, qk = attention_module.qkv_attention(q, k, v, mask)
    return attention_module.out(wv), kv_cache


def block_forward(
    residual_block,
    x: torch.Tensor,
    xa: Optional[torch.Tensor] = None,
    mask: Optional[torch.Tensor] = None,
    kv_cache: Optional[dict] = None,
    idx: int = 0
):
    """
    Override for residual block forward method for providing kv_cache to attention module.
      Parameters:
        residual_block: current residual block.
        x: input token_ids.
        xa: input audio features (Optional).
        mask: attention mask (Optional).
        kv_cache: cache for storing attention key values.
        idx: index of current residual block for search in kv_cache.
      Returns:
        x: residual block output
        kv_cache: updated kv_cache

    """
    x0, kv_cache = residual_block.attn(residual_block.attn_ln(
        x), mask=mask, kv_cache=kv_cache, idx=f'{idx}a')
    x = x + x0
    if residual_block.cross_attn:
        x1, kv_cache = residual_block.cross_attn(
            residual_block.cross_attn_ln(x), xa, kv_cache=kv_cache, idx=f'{idx}c')
        x = x + x1
    x = x + residual_block.mlp(residual_block.mlp_ln(x))
    return x, kv_cache


# update forward functions
for idx, block in enumerate(model.decoder.blocks):
    block.forward = partial(block_forward, block, idx=idx)
    block.attn.forward = partial(attention_forward, block.attn)
    if block.cross_attn:
        block.cross_attn.forward = partial(attention_forward, block.cross_attn)


def decoder_forward(decoder, x: torch.Tensor, xa: torch.Tensor, kv_cache: Optional[dict] = None):
    """
    Override for decoder forward method.
    Parameters:
      x: torch.LongTensor, shape = (batch_size, <= n_ctx) the text tokens
      xa: torch.Tensor, shape = (batch_size, n_mels, n_audio_ctx)
           the encoded audio features to be attended on
      kv_cache: Dict[str, torch.Tensor], attention modules hidden states cache from previous steps
    """
    offset = next(iter(kv_cache.values())).shape[1] if kv_cache else 0
    x = decoder.token_embedding(
        x) + decoder.positional_embedding[offset: offset + x.shape[-1]]
    x = x.to(xa.dtype)

    for block in decoder.blocks:
        x, kv_cache = block(x, xa, mask=decoder.mask, kv_cache=kv_cache)

    x = decoder.ln(x)
    logits = (
        x @ torch.transpose(decoder.token_embedding.weight.to(x.dtype), 1, 0)).float()

    return logits, kv_cache


# override decoder forward
model.decoder.forward = partial(decoder_forward, model.decoder)

tokens = torch.ones((5, 3), dtype=torch.int64)

logits, kv_cache = model.decoder(tokens, audio_features, kv_cache={})
kv_cache = {k: v for k, v in kv_cache.items()}
tokens = torch.ones((5, 1), dtype=torch.int64)

outputs = [f"out_{k}" for k in kv_cache.keys()]
inputs = [f"in_{k}" for k in kv_cache.keys()]
dynamic_axes = {
    "tokens": {0: "beam_size", 1: "seq_len"},
    "audio_features": {0: "beam_size"},
    "logits": {0: "beam_size", 1: "seq_len"}}
dynamic_outs = {o: {0: "beam_size", 1: "prev_seq_len"} for o in outputs}
dynamic_inp = {i: {0: "beam_size", 1: "prev_seq_len"} for i in inputs}
dynamic_axes.update(dynamic_outs)
dynamic_axes.update(dynamic_inp)
torch.onnx.export(
    model.decoder, {'x': tokens, 'xa': audio_features, 'kv_cache': kv_cache},
    'whisper_decoder.onnx',
    input_names=["tokens", "audio_features"] + inputs,
    output_names=["logits"] + outputs,
    dynamic_axes=dynamic_axes
)

/tmp/ipykernel_3963489/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!
  if module not in cache or output.shape[1] > positional_embeddings_size:

The decoder model autoregressively predicts the next token guided by encoder headen 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, we define an upper bound for dynamic input shapes.

input_shapes = "tokens[1..5 1..224],audio_features[1..5 1500 512]"
for k, v in kv_cache.items():
    if k.endswith('a'):
        input_shapes += f",in_{k}[1..5 0..224 512]"
decoder_model = mo.convert_model(
    input_model="whisper_decoder.onnx",
    compress_to_fp16=True,
    input=input_shapes)
serialize(decoder_model, "whisper_decoder.xml")

Warning: One or more of the values of the Constant can't fit in the float16 data type. Those values were casted to the nearest limit value, the model can produce incorrect results.

Prepare inference pipeline

The image below illustrates the pipeline of video transcribing using the Whisper model.

whisper_pipeline.png

To run the PyTorch Whisper model, you just need to call the model.transcribe(autio, **parameters) function. We will try to reuse original model pipeline for audio transcribing after replacing the original models with OpenVINO IR versions.

class OpenVINOAudioEncoder(torch.nn.Module):
    """
    Helper for inference Whisper encoder model with OpenVINO
    """

    def __init__(self, core, model_path, device='CPU'):
        super().__init__()
        self.model = core.read_model(model_path)
        self.compiled_model = core.compile_model(self.model, device)
        self.output_blob = self.compiled_model.output(0)

    def forward(self, mel: torch.Tensor):
        """
        Inference OpenVINO whisper encoder model.

        Parameters:
          mel: input audio fragment mel spectrogram.
        Returns:
          audio_features: torch tensor with encoded audio features.
        """
        return torch.from_numpy(self.compiled_model(mel)[self.output_blob])

from openvino.runtime import Core


class OpenVINOTextDecoder(torch.nn.Module):
    """
    Helper for inference OpenVINO decoder model
    """

    def __init__(self, core: Core, model_path: Path, device: str = 'CPU'):
        super().__init__()
        self._core = core
        self.model = core.read_model(model_path)
        self._input_names = [inp.any_name for inp in self.model.inputs]
        self.compiled_model = core.compile_model(self.model, device)
        self.device = device

    def init_past_inputs(self, feed_dict):
        """
        Initialize cache input for first step.

        Parameters:
          feed_dict: Dictonary with inputs for inference
        Returns:
          feed_dict: updated feed_dict
        """
        beam_size = feed_dict['tokens'].shape[0]
        audio_len = feed_dict['audio_features'].shape[-1]
        previous_seq_len = 0
        for name in self._input_names:
            if name in ['tokens', 'audio_features']:
                continue
            feed_dict[name] = np.zeros(
                (beam_size, previous_seq_len, audio_len), dtype=np.float32)
        return feed_dict

    def preprocess_kv_cache_inputs(self, feed_dict, kv_cache):
        """
        Transform kv_cache to inputs

        Parameters:
          feed_dict: dictionary with inputs for inference
          kv_cache: dictionary with cached attention hidden states from previous step
        Returns:
          feed_dict: updated feed dictionary with additional inputs
        """
        if not kv_cache:
            return self.init_past_inputs(feed_dict)
        for k, v in kv_cache.items():
            new_k = f'in_{k}'
            if new_k in self._input_names:
                feed_dict[new_k] = v
        return feed_dict

    def postprocess_outputs(self, outputs):
        """
        Transform model output to format expected by the pipeline

        Parameters:
          outputs: outputs: raw inference results.
        Returns:
          logits: decoder predicted token logits
          kv_cache: cached attention hidden states
        """
        logits = None
        kv_cache = {}
        for output_t, out in outputs.items():
            if 'logits' in output_t.get_names():
                logits = torch.from_numpy(out)
            else:
                tensor_name = output_t.any_name
                kv_cache[tensor_name.replace(
                    'out_', '')] = torch.from_numpy(out)
        return logits, kv_cache

    def forward(self, x: torch.Tensor, xa: torch.Tensor, kv_cache: Optional[dict] = None):
        """
        Inference decoder model.

        Parameters:
          x: torch.LongTensor, shape = (batch_size, <= n_ctx) the text tokens
          xa: torch.Tensor, shape = (batch_size, n_mels, n_audio_ctx)
             the encoded audio features to be attended on
          kv_cache: Dict[str, torch.Tensor], attention modules hidden states cache from previous steps
        Returns:
          logits: decoder predicted logits
          kv_cache: updated kv_cache with current step hidden states
        """
        feed_dict = {'tokens': x, 'audio_features': xa}
        feed_dict = (self.preprocess_kv_cache_inputs(feed_dict, kv_cache))
        res = self.compiled_model(feed_dict)
        return self.postprocess_outputs(res)

from whisper.decoding import DecodingTask, Inference, DecodingOptions, DecodingResult


class OpenVINOInference(Inference):
    """
    Wrapper for inference interface
    """

    def __init__(self, model: "Whisper", initial_token_length: int):
        self.model: "Whisper" = model
        self.initial_token_length = initial_token_length
        self.kv_cache = {}

    def logits(self, tokens: torch.Tensor, audio_features: torch.Tensor) -> torch.Tensor:
        """
        getting logits for given tokens sequence and audio features and save kv_cache

        Parameters:
          tokens: input tokens
          audio_features: input audio features
        Returns:
          logits: predicted by decoder logits
        """
        if tokens.shape[-1] > self.initial_token_length:
            # only need to use the last token except in the first forward pass
            tokens = tokens[:, -1:]
        logits, self.kv_cache = self.model.decoder(
            tokens, audio_features, kv_cache=self.kv_cache)
        return logits

    def cleanup_caching(self):
        """
        Reset kv_cache to initial state
        """
        self.kv_cache = {}

    def rearrange_kv_cache(self, source_indices):
        """
        Update hidden states cache for selected sequences
        Parameters:
          source_indicies: sequences indicies
        Returns:
          None
        """
        for module, tensor in self.kv_cache.items():
            # update the key/value cache to contain the selected sequences
            self.kv_cache[module] = tensor[source_indices]


class OpenVINODecodingTask(DecodingTask):
    """
    Class for decoding using OpenVINO
    """

    def __init__(self, model: "Whisper", options: DecodingOptions):
        super().__init__(model, options)
        self.inference = OpenVINOInference(model, len(self.initial_tokens))


@torch.no_grad()
def decode(model: "Whisper", mel: torch.Tensor, options: DecodingOptions = DecodingOptions()) -> Union[DecodingResult, List[DecodingResult]]:
    """
    Performs decoding of 30-second audio segment(s), provided as Mel spectrogram(s).

    Parameters
    ----------
    model: Whisper
        the Whisper model instance

    mel: torch.Tensor, shape = (80, 3000) or (*, 80, 3000)
        A tensor containing the Mel spectrogram(s)

    options: DecodingOptions
        A dataclass that contains all necessary options for decoding 30-second segments

    Returns
    -------
    result: Union[DecodingResult, List[DecodingResult]]
        The result(s) of decoding contained in `DecodingResult` dataclass instance(s)
    """
    single = mel.ndim == 2
    if single:
        mel = mel.unsqueeze(0)

    result = OpenVINODecodingTask(model, options).run(mel)

    if single:
        result = result[0]

    return result

del model.decoder
del model.encoder

from collections import namedtuple

Parameter = namedtuple('Parameter', ['device'])

core = Core()

model.encoder = OpenVINOAudioEncoder(core, 'whisper_encoder.xml')
model.decoder = OpenVINOTextDecoder(core, 'whisper_decoder.xml')
model.decode = partial(decode, model)


def parameters():
    return iter([Parameter(torch.device('cpu'))])


model.parameters = parameters


def logits(model, tokens: torch.Tensor, audio_features: torch.Tensor):
    """
    Override for logits extraction method
    Parameters:
      toekns: input tokens
      audio_features: input audio features
    Returns:
      logits: decoder predicted logits
    """
    return model.decoder(tokens, audio_features, None)[0]


model.logits = partial(logits, model)

The model expects mono-channel audio with a 16000 Hz sample rate, represented in floating point range. In cases when the audio from our input video does not meet these requirements, we will need to apply preprocessing.

import io
from pathlib import Path
import numpy as np
from scipy.io import wavfile
from pytube import YouTube
from moviepy.editor import VideoFileClip


def resample(audio, src_sample_rate, dst_sample_rate):
    """
    Resample audio to specific sample rate

    Parameters:
      audio: input audio signal
      src_sample_rate: source audio sample rate
      dst_sample_rate: destination audio sample rate
    Returns:
      resampled_audio: input audio signal resampled with dst_sample_rate
    """
    if src_sample_rate == dst_sample_rate:
        return audio
    duration = audio.shape[0] / src_sample_rate
    resampled_data = np.zeros(shape=(int(duration * dst_sample_rate)), dtype=np.float32)
    x_old = np.linspace(0, duration, audio.shape[0], dtype=np.float32)
    x_new = np.linspace(0, duration, resampled_data.shape[0], dtype=np.float32)
    resampled_audio = np.interp(x_new, x_old, audio)
    return resampled_audio.astype(np.float32)


def audio_to_float(audio):
    """
    convert audio signal to floating point format
    """
    return audio.astype(np.float32) / np.iinfo(audio.dtype).max


def get_audio(video_file):
    """
    Extract audio signal from a given video file, then convert it to float,
    then mono-channel format and resample it to the expected sample rate

    Parameters:
        video_file: path to input video file
    Returns:
      resampled_audio: mono-channel float audio signal with 16000 Hz sample rate
                       extracted from video
    """
    input_video = VideoFileClip(str(video_file))
    input_video.audio.write_audiofile(video_file.stem + '.wav', verbose=False, logger=None)
    input_audio_file = video_file.stem + '.wav'
    sample_rate, audio = wavfile.read(
        io.BytesIO(open(input_audio_file, 'rb').read()))
    audio = audio_to_float(audio)
    if audio.ndim == 2:
        audio = audio.mean(axis=1)
    resampled_audio = resample(audio, sample_rate, 16000)
    return resampled_audio

Run video transcription pipeline

Now, we are ready to start transcription. Select a video from YouTube that you want to transcribe. Please be patient as downloading the video may take some time.

import ipywidgets as widgets
VIDEO_LINK = "https://www.youtube.com/watch?v=kgL5LBM-hFI"
link = widgets.Text(
    value=VIDEO_LINK,
    placeholder="Type link for video",
    description="Video:",
    disabled=False
)

link

Text(value='https://www.youtube.com/watch?v=kgL5LBM-hFI', description='Video:', placeholder='Type link for vid…

print(f"Downloading video {link.value} started")

output_file = Path("downloaded_video.mp4")
yt = YouTube(link.value)
yt.streams.get_highest_resolution().download(filename=output_file)
print(f"Video saved to {output_file}")

Downloading video https://www.youtube.com/watch?v=kgL5LBM-hFI started
Video saved to downloaded_video.mp4

audio = get_audio(output_file)

Select the task for model: * transcribe - generate audio transcription in the source language (automatically detected). * translate - generate audio transcription with translation to English language.

task = widgets.Select(
    options=["transcribe", "translate"],
    value="translate",
    description="Select task:",
    disabled=False
)
task

Select(description='Select task:', index=1, options=('transcribe', 'translate'), value='translate')

transcription = model.transcribe(audio, beam_size=5, best_of=5, task=task.value)

def format_timestamp(seconds: float):
    """
    format time in srt-file excpected format
    """
    assert seconds >= 0, "non-negative timestamp expected"
    milliseconds = round(seconds * 1000.0)

    hours = milliseconds // 3_600_000
    milliseconds -= hours * 3_600_000

    minutes = milliseconds // 60_000
    milliseconds -= minutes * 60_000

    seconds = milliseconds // 1_000
    milliseconds -= seconds * 1_000

    return (f"{hours}:" if hours > 0 else "00:") + f"{minutes:02d}:{seconds:02d},{milliseconds:03d}"


def prepare_srt(transcription):
    """
    Format transcription into srt file format
    """
    segment_lines = []
    for segment in transcription["segments"]:
        segment_lines.append(str(segment["id"] + 1) + "\n")
        time_start = format_timestamp(segment["start"])
        time_end = format_timestamp(segment["end"])
        time_str = f"{time_start} --> {time_end}\n"
        segment_lines.append(time_str)
        segment_lines.append(segment["text"] + "\n\n")
    return segment_lines

“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.

srt_lines = prepare_srt(transcription)
# save transcription
with output_file.with_suffix(".srt").open("w") as f:
    f.writelines(srt_lines)

Now let’s see the results.

widgets.Video.from_file(output_file, loop=False, width=800, height=800)

Video(value=b'x00x00x00x18ftypmp42x00x00x00x00isommp42x00x00Aimoovx00x00x00lmvhd...', height='800…

print("".join(srt_lines))

1
00:00:00,000 --> 00:00:05,000
 Oh, what's that?

2
00:00:05,000 --> 00:00:09,000
 Oh, wow.

3
00:00:09,000 --> 00:00:10,000
 Hello, humans.

4
00:00:13,000 --> 00:00:15,000
 Focus on me.

5
00:00:15,000 --> 00:00:18,000
 Focus on the guard.

6
00:00:18,000 --> 00:00:22,000
 Don't tell anyone what you've seen in here.

7
00:00:22,000 --> 00:00:30,000
 Have you seen what's in there?