notebooks/107-speech-recognition-quantization/107-speech-recognition-quantization-data2vec.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9d91c190-956f-44ef-9e3a-1428aa020493",
   "metadata": {},
   "source": [
    "# Quantize Speech Recognition Models with OpenVINO™ Post-Training Optimization Tool ​\n",
    "This tutorial demonstrates how to apply `INT8` quantization to the speech recognition model, known as [Data2Vec](https://arxiv.org/abs/2202.03555), using the [Post-Training Optimization Tool API (POT API)](https://docs.openvino.ai/latest/pot_compression_api_README.html) (part of the [OpenVINO Toolkit](https://docs.openvino.ai/)). This notebook uses a fine-tuned [data2vec-audio-base-960h](https://huggingface.co/facebook/data2vec-audio-base-960h) [PyTorch](https://pytorch.org/) model trained on the [LibriSpeech ASR corpus](https://www.openslr.org/12). The tutorial is designed to be extendable to custom models and datasets. It consists of the following steps:\n",
    "\n",
    "- Download and prepare model.\n",
    "- Define data loading and accuracy validation functionality.\n",
    "- Prepare the model for quantization.\n",
    "- Run optimization pipeline.\n",
    "- Compare performance of the original and quantized models."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "03f79f99-6baf-4f4b-8bc4-a2b0a4e40030",
   "metadata": {},
   "source": [
    "## Download and prepare model\n",
    "\n",
    "data2vec is a framework for self-supervised representation learning for images, speech, and text as described in [data2vec: A General Framework for Self-supervised Learning in Speech, Vision and Language (Baevski et al., 2022)](https://ai.facebook.com/research/data2vec-a-general-framework-for-self-supervised-learning-in-speech-vision-and-language). The algorithm uses the same learning mechanism for different modalities.\n",
    "\n",
    "![pre-trained pipeline](https://raw.githubusercontent.com/patrickvonplaten/scientific_images/master/data2vec.png)\n",
    "\n",
    "In our case, we will use `data2vec-audio-base-960h` model, which was finetuned on 960 hours of audio from LibriSpeech Automatic Speech Recognition corpus and distributed as part of HuggingFace transformers.\n",
    "\n",
    "### Obtain Pytorch model representation\n",
    "\n",
    "For instantiating PyTorch model class, we should use `Data2VecAudioForCTC.from_pretrained` method with providing model ID for downloading from HuggingFace hub. Model weights and configuration files will be downloaded automatically in first time usage.\n",
    "Keep in mind that downloading the files can take several minutes and depends on your internet connection.\n",
    "\n",
    "Additionally, we can create processor class which is responsible for model specific pre- and post-processing steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b18e8a87-1064-4100-af21-8d392f236d89",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import Wav2Vec2Processor, Data2VecAudioForCTC\n",
    "\n",
    "processor = Wav2Vec2Processor.from_pretrained(\"facebook/data2vec-audio-base-960h\")\n",
    "model = Data2VecAudioForCTC.from_pretrained(\"facebook/data2vec-audio-base-960h\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58f16b2f-fae8-4a4b-a8dc-5b49aec358a8",
   "metadata": {},
   "source": [
    "### Convert model to OpenVINO Intermediate Representation\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06b90468-a7a1-4bab-9b52-d8ed3420e24d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "# Set model directory\n",
    "MODEL_DIR = Path(\"model\")\n",
    "MODEL_DIR.mkdir(exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52880bfb-0f26-44b5-b60d-dc707ea708c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from openvino.tools import mo\n",
    "from openvino.runtime import serialize, Core\n",
    "import torch\n",
    "\n",
    "core = Core()\n",
    "\n",
    "BATCH_SIZE = 1\n",
    "MAX_SEQ_LENGTH = 30480\n",
    "\n",
    "\n",
    "def export_model_to_onnx(model, path):\n",
    "    # switch model to evaluation mode \n",
    "    model.eval()\n",
    "    # disallow gradient propagation for reducing memory during export\n",
    "    with torch.no_grad():\n",
    "        # define dummy input with specific shape\n",
    "        default_input = torch.zeros([1, MAX_SEQ_LENGTH], dtype=torch.float)\n",
    "        inputs = {\n",
    "            \"inputs\": default_input\n",
    "        }\n",
    "\n",
    "        # define names for dynamic dimentions\n",
    "        symbolic_names = {0: \"batch_size\", 1: \"sequence_len\"}\n",
    "        # export model\n",
    "        torch.onnx.export(\n",
    "            model,\n",
    "            (inputs[\"inputs\"]),\n",
    "            path,\n",
    "            opset_version=11,\n",
    "            input_names=[\"inputs\"],\n",
    "            output_names=[\"logits\"],\n",
    "            dynamic_axes={\n",
    "                \"inputs\": symbolic_names,\n",
    "                \"logits\": symbolic_names,\n",
    "            },\n",
    "        )\n",
    "        print(\"ONNX model saved to {}\".format(path))\n",
    "\n",
    "\n",
    "onnx_model_path = MODEL_DIR / \"data2vec-audo-base.onnx\"\n",
    "ir_model_path = onnx_model_path.with_suffix('.xml')\n",
    "\n",
    "if not ir_model_path.exists():\n",
    "    if not onnx_model_path.exists():\n",
    "        export_model_to_onnx(model, onnx_model_path)\n",
    "    ov_model = mo.convert_model(onnx_model_path, compress_to_fp16=True)\n",
    "    serialize(ov_model, str(ir_model_path))\n",
    "    print(\"IR model saved to {}\".format(ir_model_path))\n",
    "else:\n",
    "    print(\"Read IR model from {}\".format(ir_model_path))\n",
    "    ov_model = core.read_model(ir_model_path)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "ea7106bc-7cfb-4158-a9c7-bbbb76787523",
   "metadata": {},
   "source": [
    "### Prepare inference data\n",
    "\n",
    "For demonstration purposes, we will use short dummy version of librispeach dataset - `patrickvonplaten/librispeech_asr_dummy` to speed up model evaluation. Model accuracy can be different from reported in the paper. For reproducing original accuracy, use `librispeech_asr` dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c312b02e-2c99-4211-b7f1-4e491550b568",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0873e12-ae73-4ac3-847f-d431818b7b5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "ds = load_dataset(\"patrickvonplaten/librispeech_asr_dummy\", \"clean\", split=\"validation\")\n",
    "\n",
    "\n",
    "# define preprocessing function for converting audio to input values for model\n",
    "def map_to_input(batch):\n",
    "    preprocessed_signal = processor(batch[\"audio\"][\"array\"], return_tensors=\"pt\", padding=\"longest\", sampling_rate=batch['audio']['sampling_rate'])\n",
    "    input_values = preprocessed_signal.input_values\n",
    "    batch['input_values'] = input_values\n",
    "    return batch\n",
    "\n",
    "\n",
    "# apply preprocessing function to dataset and remove audio column, to save memory as we do not need it anymore\n",
    "dataset = ds.map(map_to_input, batched=False, remove_columns=[\"audio\"])\n",
    "\n",
    "test_sample = ds[0][\"audio\"]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "dc35d830-4b64-4083-bd0c-aaef68ba3d86",
   "metadata": {},
   "source": [
    "## Check model inference result\n",
    "\n",
    "The code below is used for running model inference on a single sample from the dataset. It contains the following steps:\n",
    "\n",
    "* Get the input_values tensor as model input.\n",
    "* Run model inference and obtain logits.\n",
    "* Find logits ids with highest probability, using argmax.\n",
    "* Decode predicted token ids, using processor.\n",
    "\n",
    "For reference, see the same function provided for OpenVINO model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6a11f44-2fe1-4aa9-a360-ba4b23bba87d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "\n",
    "# inference function for pytorch\n",
    "def torch_infer(model, sample):\n",
    "    logits = model(torch.Tensor(sample['input_values'])).logits\n",
    "    # take argmax and decode\n",
    "    predicted_ids = torch.argmax(logits, dim=-1)\n",
    "    transcription = processor.batch_decode(predicted_ids)\n",
    "    return transcription\n",
    "\n",
    "\n",
    "# inference function for openvino\n",
    "def ov_infer(model, sample):\n",
    "    output = model.output(0)\n",
    "    logits = model(np.array(sample['input_values']))[output]\n",
    "    predicted_ids = np.argmax(logits, axis=-1)\n",
    "    transcription = processor.batch_decode(torch.from_numpy(predicted_ids))\n",
    "    return transcription"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb36df3c-b461-4400-9ebc-19d9f307032c",
   "metadata": {},
   "outputs": [],
   "source": [
    "pt_transcription = torch_infer(model, dataset[0])\n",
    "compiled_model = core.compile_model(ov_model)\n",
    "ov_transcription = ov_infer(compiled_model, dataset[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d02ca6e-9939-4d10-a65f-f03b19c4153a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import IPython.display as ipd\n",
    "\n",
    "print(f\"[Reference]:     {dataset[0]['text']}\")\n",
    "print(f\"[PyTorch]:       {pt_transcription[0]}\")\n",
    "print(f\"[OpenVINO FP16]: {ov_transcription[0]}\")\n",
    "ipd.Audio(test_sample[\"array\"], rate=16000)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "d7281085-47ef-4f6e-b778-8d5921a97644",
   "metadata": {},
   "source": [
    "## Validate model accuracy on dataset\n",
    "\n",
    "For model accuracy evaluation, [Word Error Rate](https://en.wikipedia.org/wiki/Word_error_rate) metric can be used. Word Error Rate or WER is the ratio of errors in a transcript to the total words spoken. A lower WER in speech-to-text means better accuracy in recognizing speech.\n",
    "\n",
    "For WER calculation, we will use [torchmetrics](https://torchmetrics.readthedocs.io/en/stable/text/word_error_rate.html) library."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0cdf626f-a1c3-4bd7-b31d-4542fed04236",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchmetrics import WordErrorRate\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "\n",
    "def compute_wer(dataset, model, infer_fn):\n",
    "    wer = WordErrorRate()\n",
    "    for sample in tqdm(dataset):\n",
    "        # run infer function on sample\n",
    "        transcription = infer_fn(model, sample)\n",
    "        # update metric on sample result\n",
    "        wer.update(transcription, [sample['text']])\n",
    "    # finalize metric calculation\n",
    "    result = wer.compute()\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da4ceda6-1535-4816-ab1b-bf3dc80a636f",
   "metadata": {},
   "outputs": [],
   "source": [
    "pt_result = compute_wer(dataset, model, torch_infer)\n",
    "ov_result = compute_wer(dataset, compiled_model, ov_infer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3879e414-42c2-4b12-a06f-bcc0538c8834",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f'[PyTorch]   Word Error Rate: {pt_result:.4f}')\n",
    "print(f'[OpenVino]  Word Error Rate: {ov_result:.4f}')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "88f6eff7-c68e-476b-92c3-771ce9786337",
   "metadata": {},
   "source": [
    "## Prepare quantization pipeline\n",
    "\n",
    "Post-Training Optimization Tool designed to accelerate the inference of DL models by converting them into a more hardware-friendly representation by applying specific methods that do not require re-training, for example, post-training quantization. For more details about the low-precision flow in OpenVINO™, refer to the [Low Precision Optimization Guide](https://docs.openvino.ai/2020.4/pot_docs_LowPrecisionOptimizationGuide.html).\n",
    "\n",
    "[The Python POT API](https://docs.openvino.ai/2020.4/pot_compression_api_README.html) provides simple interfaces for implementing custom model inference with data loading and pre-processing on an arbitrary dataset and implementing custom accuracy metrics to make it possible to use optimization algorithms from the POT.\n",
    "The Python POT API represented by `Pipeline` class for creating and configuring the optimization pipeline and applying it to the model. The `Pipeline` class depends on the implementation of the following model specific interfaces which should be implemented according to the custom DL model:\n",
    "\n",
    "* `Engine` is responsible for model inference and provides statistical data and accuracy metrics for the model.\n",
    "* `DataLoader` is responsible for the dataset loading, including the data pre-processing.\n",
    "* `Metric` is responsible for calculating the accuracy metric for the model.\n",
    "\n",
    "The diagram below shows relationships between the classes:\n",
    "\n",
    "![pot pipeline](https://docs.openvino.ai/2020.4/custom_optimization_pipeline.png)\n",
    "\n",
    "\n",
    "### Define DataLoader class\n",
    "\n",
    "Define `DataLoader` based on POT API, as it will be used to collect statistics for quantization and run model evaluation.\n",
    "Data22Vec model accepts a raw waveform of the speech signal as input and produces vocabulary class estimations as output. We already have prepared dataset above for accuracy measurement. It will serve as data source for quantization. DataLoader class encapsulates logic for iteration over dataset samples and gets input data and label by index using `__getitem__` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab3e90f7-7223-4a1d-b4af-278869adea70",
   "metadata": {},
   "outputs": [],
   "source": [
    "from openvino.tools.pot import Metric, DataLoader, IEEngine, load_model, save_model, compress_model_weights, create_pipeline\n",
    "\n",
    "\n",
    "class LibriSpeechDataLoader(DataLoader):\n",
    "\n",
    "    # Required methods\n",
    "    def __init__(self, dataset, sample_limit=None):\n",
    "        \"\"\"Constructor\n",
    "        :param config: data loader specific config\n",
    "        \"\"\"\n",
    "        super().__init__({})\n",
    "        self._ds = dataset\n",
    "        self.sample_limit = None \n",
    "        \n",
    "    def __len__(self):\n",
    "        \"\"\"Returns size of the dataset\"\"\"\n",
    "        return self.sample_limit or len(self._ds)\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        \"\"\"\n",
    "        Returns annotation, data and metadata at the specified index.\n",
    "        Possible formats:\n",
    "        (index, annotation), data\n",
    "        (index, annotation), data, metadata\n",
    "        \"\"\"\n",
    "        if self.sample_limit is not None and index >= self.sample_limit:\n",
    "            raise StopIteration\n",
    "        sample = self._ds[index]\n",
    "        inputs = {'inputs': np.array(sample['input_values'])}\n",
    "        label = [sample['text']]\n",
    "        return inputs, label"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "2d992fe2-586f-461f-9848-14f52db0990a",
   "metadata": {},
   "source": [
    "### Define Evaluation Metric class\n",
    "\n",
    "In this step, the `Metric` interface for WER metric is implemented. To make the metric compatible with running inside POT Pipeline, we should inherit it from `openvino.tools.pot.Metric` class and override following properties and methods:\n",
    "* `value` - returns the accuracy metric value for the last model output.\n",
    "* `avg_value` - returns the average accuracy metric value for all model outputs.\n",
    "* `attributes` - returns a dictionary of metric attributes: `direction` - metric growing direction (`higher-better` or `higher-worse`), `type` - type of metric.\n",
    "* `update(output, annotation)` - calculates and updates the accuracy metric value using last model output and annotation.\n",
    "* `reset()` - resets collected accuracy metric."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "125c16bf-3bad-45ab-abee-530fe7135070",
   "metadata": {},
   "outputs": [],
   "source": [
    "class WERMetric(Metric):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self._name = \"WER\"\n",
    "\n",
    "    def reset(self):\n",
    "        \"\"\"\n",
    "        Resets collected matches\n",
    "        \"\"\"\n",
    "        self._wer = WordErrorRate()\n",
    "        self._last_result = None\n",
    "\n",
    "    def get_attributes(self):\n",
    "        \"\"\"\n",
    "        Returns a dictionary of metric attributes {metric_name: {attribute_name: value}}.\n",
    "        Required attributes: 'direction': 'higher-better' or 'higher-worse'\n",
    "                             'type': metric type\n",
    "        \"\"\"\n",
    "        return {self._name: {\"direction\": \"higher-worse\", \"type\": \"WER\"}}\n",
    "\n",
    "    @property\n",
    "    def value(self):\n",
    "        \"\"\"Returns accuracy metric value for the last model output.\"\"\"\n",
    "        return {self._name: self._last_result if self._last_result is not None else self._wer.compute().item()}\n",
    "\n",
    "    @property\n",
    "    def avg_value(self):\n",
    "        \"\"\"Returns accuracy metric value for all model outputs.\"\"\"\n",
    "        return {self._name: self._wer.compute().item()}\n",
    "\n",
    "    def update(self, output, target):\n",
    "        \"\"\"\n",
    "        Updates prediction matches.\n",
    "\n",
    "        :param output: model output\n",
    "        :param target: annotations\n",
    "        \"\"\"\n",
    "        res = output[0]\n",
    "        predicted_ids = np.argmax(res, axis=-1)\n",
    "        predicted_transcription = processor.batch_decode(torch.from_numpy(predicted_ids))\n",
    "        res = []\n",
    "        for pred, gt in zip(predicted_transcription, target):\n",
    "            res.append(self._wer.forward([pred], gt).item())\n",
    "        self._last_result = res\n",
    "        return res"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "fa9af259-42c7-4a2d-8c46-b3a94002a4f8",
   "metadata": {},
   "source": [
    "### Define quantization configuration and optimization pipeline\n",
    "\n",
    "The code below defines a configuration for the quantization pipeline and runs it. To keep example minimalistic, built-in `IEEngine` implementation of `Engine` interface from the POT API for model inference is used here.\n",
    "We will use DefaultQuantization algorithm with `performance` preset and additional specification of quantization algorithm for activations. For information about configuration parameters, refer to [POT documentation](https://docs.openvino.ai/latest/pot_compression_algorithms_quantization_default_README.html).\n",
    "Our model architecture is transformer-based, so `model_type: transformer` should be selected. For better accuracy, part of layers should be kept in floating point representation using `ignored` parameter. The ignored layers can be selected using [AccuracyAwareQuantization](https://docs.openvino.ai/latest/pot_accuracyaware_usage.html) algorithm, which aim to find layers that have the most significant impact on accuracy drop and revert them back to floating point precision. This process can be time consuming, that is why we keep this experiment out of this tutorial and reuse its result, using DefaultQuantization algorithm.\n",
    "> **NOTE**: Consider increasing `stat_subset_size` to get more precise results. A suggested value is `300` or more, as it will take longer time to process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7967a79-7ea9-4267-ace7-b45a58b123d9",
   "metadata": {
    "test_replace": {
     "\"stat_subset_size\": 300,": "\"stat_subset_size\": 3,"
    }
   },
   "outputs": [],
   "source": [
    "model_config = {\"model_name\": \"data2vec_base\", \"model\": ir_model_path, \"weights\": ir_model_path.with_suffix(\".bin\")}\n",
    "\n",
    "engine_config = {\"device\": \"CPU\"}\n",
    "\n",
    "algorithms = [\n",
    "    {\n",
    "        \"name\": \"DefaultQuantization\",\n",
    "        \"params\": {\n",
    "            \"target_device\": \"ANY\",\n",
    "            \"model_type\": \"transformer\",\n",
    "            \"preset\": \"performance\",\n",
    "            \"stat_subset_size\": 300,\n",
    "            \"activations\": {\n",
    "                \"range_estimator\": {\n",
    "                    \"min\": {\n",
    "                        \"aggregator\": \"min\",\n",
    "                        \"type\": \"min\"\n",
    "                    },\n",
    "                    \"max\": {\n",
    "                        \"aggregator\": \"mean\",\n",
    "                        \"type\": \"quantile\",\n",
    "                        \"outlier_prob\": 0.0001\n",
    "                    }\n",
    "                }\n",
    "            },\n",
    "            \"ignored\": {\n",
    "                \"scope\": [\n",
    "                    \"/data2vec_audio/encoder/layers.3/feed_forward/intermediate_dense/MatMul\", \n",
    "                    \"/data2vec_audio/feature_extractor/conv_layers.2/conv/Conv\", \n",
    "                    \"/data2vec_audio/encoder/layers.3/Add_1\", \n",
    "                    \"/data2vec_audio/encoder/layers.2/feed_forward/intermediate_dense/MatMul\", \n",
    "                    \"/data2vec_audio/feature_extractor/conv_layers.0/conv/Conv\", \n",
    "                    \"/data2vec_audio/encoder/layers.4/Add_1\", \n",
    "                    \"/data2vec_audio/encoder/layers.4/feed_forward/intermediate_dense/MatMul\", \n",
    "                    \"/data2vec_audio/encoder/layers.4/final_layer_norm/Div\", \n",
    "                    \"/data2vec_audio/encoder/layers.4/feed_forward/output_dense/MatMul\", \n",
    "                    \"/data2vec_audio/encoder/layers.8/attention/MatMul_1\", \n",
    "                    \"/data2vec_audio/feature_extractor/conv_layers.1/conv/Conv\", \n",
    "                    \"/data2vec_audio/encoder/layers.2/Add_1\", \n",
    "                    \"/data2vec_audio/feature_extractor/conv_layers.0/layer_norm/Div\", \n",
    "                    \"/data2vec_audio/encoder/layers.1/feed_forward/intermediate_dense/MatMul\", \n",
    "                    \"/data2vec_audio/encoder/layers.1/Add_1\", \n",
    "                    \"/data2vec_audio/feature_extractor/conv_layers.3/layer_norm/Div\"\n",
    "                ]\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "]\n",
    "\n",
    "# Step 1: Load the model.\n",
    "model = load_model(model_config=model_config)\n",
    "\n",
    "# Step 2: Initialize the data loader.\n",
    "data_loader = LibriSpeechDataLoader(dataset)\n",
    "\n",
    "# Step 3 (Optional. Required for AccuracyAwareQuantization): Initialize the metric.\n",
    "metric = WERMetric()\n",
    "\n",
    "# Step 4: Initialize the engine for metric calculation and statistics collection.\n",
    "engine = IEEngine(config=engine_config, data_loader=data_loader, metric=metric)\n",
    "\n",
    "# Step 5: Create a pipeline of compression algorithms.\n",
    "pipeline = create_pipeline(algo_config=algorithms, engine=engine)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "cb1057b4-0ea9-4b35-8f5e-a232d52818a0",
   "metadata": {},
   "source": [
    "## Run model quantization\n",
    "\n",
    "Now, when all parts of compression pipeline are collected, we can start quantization.\n",
    ">**NOTE**: quantization process is time and memory consuming. It may takes several minutes depending on your hardware configuration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbe5c67f-efd0-4095-84a8-64a1bf2aadfe",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "# Step 6: Run compression pipeline\n",
    "print(f\"Quantizing model with {algorithms[0]['params']['preset']} preset and {algorithms[0]['name']}\")\n",
    "start_time = time.perf_counter()\n",
    "compressed_model = pipeline.run(model=model)\n",
    "end_time = time.perf_counter()\n",
    "print(f\"Quantization finished in {end_time - start_time:.2f} seconds\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2fc9e942-58d4-49d3-b115-758583661093",
   "metadata": {},
   "source": [
    "After quantization is finished, compressed model representation can be saved using `save_model` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "898ecbdd-1f89-47c7-8877-b7acf57ec8f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 7 (Optional): Compress model weights to quantized precision\n",
    "#                    in order to reduce the size of the final .bin file.\n",
    "compress_model_weights(model=compressed_model)\n",
    "\n",
    "# Step 8: Save the compressed model to the desired path.\n",
    "compressed_model_paths = save_model(model=compressed_model, save_path=MODEL_DIR, model_name=\"quantized_data2vec_base\")\n",
    "compressed_model_path = compressed_model_paths[0][\"model\"]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "9fda5391-9a7c-4606-8d5e-fed5fd6cf0fc",
   "metadata": {},
   "source": [
    "## Check INT8 model inference result\n",
    "\n",
    "`INT8` model is the same in usage like the original one. We need to read it, using the `core.read_model` method and load on the device, using `core.compile_model`. After that, we can reuse the same `ov_infer` function for getting model inference result on test sample."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0ea3977-c5bd-42ec-b700-4bcc43a8cd29",
   "metadata": {},
   "outputs": [],
   "source": [
    "ov_int8_model = core.read_model(compressed_model_path)\n",
    "int8_compiled_model = core.compile_model(ov_int8_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "518b56cb-5617-4ad6-a7b5-747e9698829a",
   "metadata": {},
   "outputs": [],
   "source": [
    "transcription = ov_infer(int8_compiled_model, dataset[0])\n",
    "print(f\"[Reference]:     {dataset[0]['text']}\")\n",
    "print(f\"[OpenVINO INT8]: {transcription[0]}\")\n",
    "ipd.Audio(test_sample[\"array\"], rate=16000)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "98777233-6cc5-4d97-bbcd-f8338c75b55e",
   "metadata": {},
   "source": [
    "## Compare Performance of the Original and Quantized Models\n",
    "[Benchmark Tool](https://docs.openvino.ai/latest/openvino_inference_engine_tools_benchmark_tool_README.html) is used to measure the inference performance of the `FP16` and `INT8` models.\n",
    "\n",
    "> **NOTE**: For more accurate performance, it is recommended to run `benchmark_app` in a terminal/command prompt after closing other applications. Run `benchmark_app -m model.xml -d CPU` to benchmark async inference on CPU for one minute. Change `CPU` to `GPU` to benchmark on GPU. Run `benchmark_app --help` to see an overview of all command-line options."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c65ef565-5b65-4efb-b663-526c8e9f0eed",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Inference FP16 model (OpenVINO IR)\n",
    "! benchmark_app -m $ir_model_path -shape [1,30480] -d CPU -api async -t 15"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32960bce-373a-4421-ba7c-d9a7e3f9b678",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Inference INT8 model (OpenVINO IR)\n",
    "! benchmark_app -m $compressed_model_path -shape [1,30480] -d CPU -api async -t 15"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "34347a28-73bb-4bfc-8048-25e03d70f80e",
   "metadata": {},
   "source": [
    "## Compare Accuracy of the Original and Quantized Models\n",
    "\n",
    "Finally, calculate WER metric for the `INT8` model representation and compare it with the `FP16` result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb4937fb-ba00-490b-b116-d6c731d35450",
   "metadata": {},
   "outputs": [],
   "source": [
    "int8_ov_result = compute_wer(dataset, int8_compiled_model, ov_infer)\n",
    "print(f'[OpenVino FP16] Word Error Rate: {ov_result:.4}')\n",
    "print(f'[OpenVino INT8] Word Error Rate: {int8_ov_result:.4f}')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  },
  "vscode": {
   "interpreter": {
    "hash": "916dbcbb3f70747c44a77c7bcd40155683ae19c65e1c03b4aa3499c5328201f1"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}