notebooks/language-quantize-bert/language-quantize-bert.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "cacbe6b4",
   "metadata": {
    "id": "rQc-wXjqrEuR"
   },
   "source": [
    "# Quantize NLP models with Post-Training Quantization ​in NNCF\n",
    "This tutorial demonstrates how to apply `INT8` quantization to the Natural Language Processing model known as [BERT](https://en.wikipedia.org/wiki/BERT_(language_model)), using the [Post-Training Quantization API](https://docs.openvino.ai/2024/openvino-workflow/model-optimization-guide/quantizing-models-post-training/basic-quantization-flow.html) (NNCF library). A fine-tuned [HuggingFace BERT](https://huggingface.co/transformers/model_doc/bert.html) [PyTorch](https://pytorch.org/) model, trained on the [Microsoft Research Paraphrase Corpus (MRPC)](https://www.microsoft.com/en-us/download/details.aspx?id=52398), will be used. The tutorial is designed to be extendable to custom models and datasets. It consists of the following steps:\n",
    "\n",
    "- Download and prepare the BERT model and MRPC dataset.\n",
    "- Define data loading and accuracy validation functionality.\n",
    "- Prepare the model for quantization.\n",
    "- Run optimization pipeline.\n",
    "- Load and test quantized model.\n",
    "- Compare the performance of the original, converted and quantized models.\n",
    "\n",
    "\n",
    "<img referrerpolicy=\"no-referrer-when-downgrade\" src=\"https://static.scarf.sh/a.png?x-pxid=5b5a4db0-7875-4bfb-bdbd-01698b5b1a77&file=notebooks/language-quantize-bert/language-quantize-bert.ipynb\" />\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "d89f8a6b",
   "metadata": {},
   "source": [
    "\n",
    "#### Table of contents:\n",
    "\n",
    "- [Imports](#Imports)\n",
    "- [Settings](#Settings)\n",
    "- [Prepare the Model](#Prepare-the-Model)\n",
    "- [Prepare the Dataset](#Prepare-the-Dataset)\n",
    "- [Optimize model using NNCF Post-training Quantization API](#Optimize-model-using-NNCF-Post-training-Quantization-API)\n",
    "- [Load and Test OpenVINO Model](#Load-and-Test-OpenVINO-Model)\n",
    "    - [Select inference device](#Select-inference-device)\n",
    "- [Compare F1-score of FP32 and INT8 models](#Compare-F1-score-of-FP32-and-INT8-models)\n",
    "- [Compare Performance of the Original, Converted and Quantized Models](#Compare-Performance-of-the-Original,-Converted-and-Quantized-Models)\n",
    "\n",
    "\n",
    "### Installation Instructions\n",
    "\n",
    "This is a self-contained example that relies solely on its own code.\n",
    "\n",
    "We recommend  running the notebook in a virtual environment. You only need a Jupyter server to start.\n",
    "For details, please refer to [Installation Guide](https://github.com/openvinotoolkit/openvino_notebooks/blob/latest/README.md#-installation-guide)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "694d9fc1-501c-4b86-a747-637e2aad64ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -q \"nncf>=2.5.0\"\n",
    "%pip install -q torch transformers \"torch>=2.1\" datasets evaluate tqdm  --extra-index-url https://download.pytorch.org/whl/cpu\n",
    "%pip install -q \"openvino>=2023.1.0\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "4d6b41e6-132b-40da-b3b9-91bacba29e31",
   "metadata": {},
   "source": [
    "## Imports\n",
    "[back to top ⬆️](#Table-of-contents:)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "771388d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "from pathlib import Path\n",
    "from zipfile import ZipFile\n",
    "from typing import Iterable\n",
    "from typing import Any\n",
    "\n",
    "import datasets\n",
    "import evaluate\n",
    "import numpy as np\n",
    "import nncf\n",
    "from nncf.parameters import ModelType\n",
    "import openvino as ov\n",
    "import torch\n",
    "from transformers import BertForSequenceClassification, BertTokenizer\n",
    "\n",
    "# Fetch `notebook_utils` module\n",
    "import requests\n",
    "\n",
    "if not Path(\"notebook_utils.py\").exists():\n",
    "    r = requests.get(\n",
    "        url=\"https://raw.githubusercontent.com/openvinotoolkit/openvino_notebooks/latest/utils/notebook_utils.py\",\n",
    "    )\n",
    "\n",
    "    open(\"notebook_utils.py\", \"w\").write(r.text)\n",
    "from notebook_utils import download_file, device_widget\n",
    "\n",
    "# Read more about telemetry collection at https://github.com/openvinotoolkit/openvino_notebooks?tab=readme-ov-file#-telemetry\n",
    "from notebook_utils import collect_telemetry\n",
    "\n",
    "collect_telemetry(\"language-quantize-bert.ipynb\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e9e66896-d439-4065-868a-65b44d31525a",
   "metadata": {},
   "source": [
    "## Settings\n",
    "[back to top ⬆️](#Table-of-contents:)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "284e9a4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the data and model directories, source URL and the filename of the model.\n",
    "DATA_DIR = Path(\"data\")\n",
    "MODEL_DIR = Path(\"model\")\n",
    "MODEL_LINK = \"https://download.pytorch.org/tutorial/MRPC.zip\"\n",
    "FILE_NAME = MODEL_LINK.split(\"/\")[-1]\n",
    "PRETRAINED_MODEL_DIR = MODEL_DIR / \"MRPC\"\n",
    "\n",
    "DATA_DIR.mkdir(exist_ok=True)\n",
    "MODEL_DIR.mkdir(exist_ok=True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "44dc335d",
   "metadata": {
    "id": "YytHDzLE0uOJ",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Prepare the Model\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "Perform the following:\n",
    "\n",
    "- Download and unpack pre-trained BERT model for MRPC by PyTorch.\n",
    "- Convert the model to the OpenVINO Intermediate Representation (OpenVINO IR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be9fc64c",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not PRETRAINED_MODEL_DIR.exists():\n",
    "    download_file(MODEL_LINK, directory=MODEL_DIR, show_progress=True)\n",
    "    with ZipFile(f\"{MODEL_DIR}/{FILE_NAME}\", \"r\") as zip_ref:\n",
    "        zip_ref.extractall(MODEL_DIR)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "5dca2fa0",
   "metadata": {
    "id": "ehX7F6KB0uPu"
   },
   "source": [
    "Convert the original PyTorch model to the OpenVINO Intermediate Representation.\n",
    "\n",
    "From OpenVINO 2023.0, we can directly convert a model from the PyTorch format to the OpenVINO IR format using model conversion API. Following PyTorch model formats are supported:\n",
    "\n",
    "- `torch.nn.Module`\n",
    "- `torch.jit.ScriptModule`\n",
    "- `torch.jit.ScriptFunction`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb2f6d66",
   "metadata": {},
   "outputs": [],
   "source": [
    "MAX_SEQ_LENGTH = 128\n",
    "input_shape = ov.PartialShape([1, -1])\n",
    "ir_model_xml = Path(MODEL_DIR) / \"bert_mrpc.xml\"\n",
    "core = ov.Core()\n",
    "\n",
    "torch_model = BertForSequenceClassification.from_pretrained(PRETRAINED_MODEL_DIR)\n",
    "torch_model.eval\n",
    "\n",
    "input_info = [\n",
    "    (\"input_ids\", input_shape, np.int64),\n",
    "    (\"attention_mask\", input_shape, np.int64),\n",
    "    (\"token_type_ids\", input_shape, np.int64),\n",
    "]\n",
    "default_input = torch.ones(1, MAX_SEQ_LENGTH, dtype=torch.int64)\n",
    "inputs = {\n",
    "    \"input_ids\": default_input,\n",
    "    \"attention_mask\": default_input,\n",
    "    \"token_type_ids\": default_input,\n",
    "}\n",
    "\n",
    "# Convert the PyTorch model to OpenVINO IR FP32.\n",
    "if not ir_model_xml.exists():\n",
    "    model = ov.convert_model(torch_model, example_input=inputs, input=input_info)\n",
    "    ov.save_model(model, str(ir_model_xml))\n",
    "else:\n",
    "    model = core.read_model(ir_model_xml)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "17f79b5f",
   "metadata": {},
   "source": [
    "## Prepare the Dataset\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "We download the [General Language Understanding Evaluation (GLUE)](https://gluebenchmark.com/) dataset for the MRPC task from HuggingFace datasets.\n",
    "Then, we tokenize the data with a pre-trained BERT tokenizer from HuggingFace."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "632fb1fc",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "def create_data_source():\n",
    "    raw_dataset = datasets.load_dataset(\"glue\", \"mrpc\", split=\"validation\")\n",
    "    tokenizer = BertTokenizer.from_pretrained(PRETRAINED_MODEL_DIR)\n",
    "\n",
    "    def _preprocess_fn(examples):\n",
    "        texts = (examples[\"sentence1\"], examples[\"sentence2\"])\n",
    "        result = tokenizer(*texts, padding=\"max_length\", max_length=MAX_SEQ_LENGTH, truncation=True)\n",
    "        result[\"labels\"] = examples[\"label\"]\n",
    "        return result\n",
    "\n",
    "    processed_dataset = raw_dataset.map(_preprocess_fn, batched=True, batch_size=1)\n",
    "\n",
    "    return processed_dataset\n",
    "\n",
    "\n",
    "data_source = create_data_source()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e082b01d",
   "metadata": {},
   "source": [
    "## Optimize model using NNCF Post-training Quantization API\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "[NNCF](https://github.com/openvinotoolkit/nncf) provides a suite of advanced algorithms for Neural Networks inference optimization in OpenVINO with minimal accuracy drop.\n",
    "We will use 8-bit quantization in post-training mode (without the fine-tuning pipeline) to optimize BERT.\n",
    "\n",
    "The optimization process contains the following steps:\n",
    "\n",
    "1. Create a Dataset for quantization\n",
    "2. Run `nncf.quantize` for getting an optimized model\n",
    "3. Serialize OpenVINO IR model using `openvino.save_model` function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e089ea99",
   "metadata": {
    "test_replace": {
     "quantized_model = nncf.quantize(model, calibration_dataset, model_type=ModelType.TRANSFORMER)": "quantized_model = nncf.quantize(model, calibration_dataset, model_type=ModelType.TRANSFORMER, subset_size=10)"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:nncf:202 ignored nodes was found by types in the NNCFGraph\n",
      "INFO:nncf:24 ignored nodes was found by name in the NNCFGraph\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 22 aten::rsub_16\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 25 aten::rsub_17\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 30 aten::mul_18\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 11 aten::add_40\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 14 aten::add__46\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 17 aten::layer_norm_48\n",
      "20 aten::layer_norm_49\n",
      "23 aten::layer_norm_50\n",
      "\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 36 aten::add_108\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 55 aten::softmax_109\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 74 aten::matmul_110\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 26 aten::add_126\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 31 aten::layer_norm_128\n",
      "47 aten::layer_norm_129\n",
      "66 aten::layer_norm_130\n",
      "\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 85 aten::add_140\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 103 aten::layer_norm_142\n",
      "133 aten::layer_norm_143\n",
      "171 aten::layer_norm_144\n",
      "\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 38 aten::add_202\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 57 aten::softmax_203\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 76 aten::matmul_204\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 209 aten::add_220\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 236 aten::layer_norm_222\n",
      "250 aten::layer_norm_223\n",
      "267 aten::layer_norm_224\n",
      "\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 287 aten::add_234\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 316 aten::layer_norm_236\n",
      "342 aten::layer_norm_237\n",
      "364 aten::layer_norm_238\n",
      "\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 39 aten::add_296\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 58 aten::softmax_297\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 77 aten::matmul_298\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 221 aten::add_314\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 242 aten::layer_norm_316\n",
      "259 aten::layer_norm_317\n",
      "279 aten::layer_norm_318\n",
      "\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 300 aten::add_328\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 326 aten::layer_norm_330\n",
      "348 aten::layer_norm_331\n",
      "370 aten::layer_norm_332\n",
      "\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 40 aten::add_390\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 59 aten::softmax_391\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 78 aten::matmul_392\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 223 aten::add_408\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 243 aten::layer_norm_410\n",
      "260 aten::layer_norm_411\n",
      "280 aten::layer_norm_412\n",
      "\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 302 aten::add_422\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 328 aten::layer_norm_424\n",
      "350 aten::layer_norm_425\n",
      "372 aten::layer_norm_426\n",
      "\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 41 aten::add_484\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 60 aten::softmax_485\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 79 aten::matmul_486\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 225 aten::add_502\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 244 aten::layer_norm_504\n",
      "261 aten::layer_norm_505\n",
      "281 aten::layer_norm_506\n",
      "\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 304 aten::add_516\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 330 aten::layer_norm_518\n",
      "352 aten::layer_norm_519\n",
      "374 aten::layer_norm_520\n",
      "\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 42 aten::add_578\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 61 aten::softmax_579\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 80 aten::matmul_580\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 227 aten::add_596\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 245 aten::layer_norm_598\n",
      "262 aten::layer_norm_599\n",
      "282 aten::layer_norm_600\n",
      "\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 306 aten::add_610\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 332 aten::layer_norm_612\n",
      "354 aten::layer_norm_613\n",
      "376 aten::layer_norm_614\n",
      "\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 43 aten::add_672\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 62 aten::softmax_673\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 81 aten::matmul_674\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 229 aten::add_690\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 246 aten::layer_norm_692\n",
      "263 aten::layer_norm_693\n",
      "283 aten::layer_norm_694\n",
      "\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 308 aten::add_704\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 334 aten::layer_norm_706\n",
      "356 aten::layer_norm_707\n",
      "378 aten::layer_norm_708\n",
      "\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 44 aten::add_766\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 63 aten::softmax_767\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 82 aten::matmul_768\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 231 aten::add_784\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 247 aten::layer_norm_786\n",
      "264 aten::layer_norm_787\n",
      "284 aten::layer_norm_788\n",
      "\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 310 aten::add_798\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 336 aten::layer_norm_800\n",
      "358 aten::layer_norm_801\n",
      "380 aten::layer_norm_802\n",
      "\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 45 aten::add_860\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 64 aten::softmax_861\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 83 aten::matmul_862\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 233 aten::add_878\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 248 aten::layer_norm_880\n",
      "265 aten::layer_norm_881\n",
      "285 aten::layer_norm_882\n",
      "\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 312 aten::add_892\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 338 aten::layer_norm_894\n",
      "360 aten::layer_norm_895\n",
      "382 aten::layer_norm_896\n",
      "\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 46 aten::add_954\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 65 aten::softmax_955\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 84 aten::matmul_956\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 235 aten::add_972\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 249 aten::layer_norm_974\n",
      "266 aten::layer_norm_975\n",
      "286 aten::layer_norm_976\n",
      "\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 314 aten::add_986\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 340 aten::layer_norm_988\n",
      "362 aten::layer_norm_989\n",
      "384 aten::layer_norm_990\n",
      "\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 35 aten::add_1048\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 54 aten::softmax_1049\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 73 aten::matmul_1050\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 215 aten::add_1066\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 240 aten::layer_norm_1068\n",
      "257 aten::layer_norm_1069\n",
      "277 aten::layer_norm_1070\n",
      "\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 296 aten::add_1080\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 322 aten::layer_norm_1082\n",
      "344 aten::layer_norm_1083\n",
      "366 aten::layer_norm_1084\n",
      "\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 37 aten::add_1142\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 56 aten::softmax_1143\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 75 aten::matmul_1144\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 218 aten::add_1160\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 241 aten::layer_norm_1162\n",
      "258 aten::layer_norm_1163\n",
      "278 aten::layer_norm_1164\n",
      "\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 298 aten::add_1174\n",
      "INFO:nncf:Not adding activation input quantizer for operation: 324 aten::layer_norm_1176\n",
      "346 aten::layer_norm_1177\n",
      "368 aten::layer_norm_1178\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Statistics collection: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 300/300 [00:26<00:00, 11.28it/s]\n",
      "Biases correction: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 74/74 [00:25<00:00,  2.89it/s]\n"
     ]
    }
   ],
   "source": [
    "INPUT_NAMES = [key for key in inputs.keys()]\n",
    "\n",
    "\n",
    "def transform_fn(data_item):\n",
    "    \"\"\"\n",
    "    Extract the model's input from the data item.\n",
    "    The data item here is the data item that is returned from the data source per iteration.\n",
    "    This function should be passed when the data item cannot be used as model's input.\n",
    "    \"\"\"\n",
    "    inputs = {name: np.asarray([data_item[name]], dtype=np.int64) for name in INPUT_NAMES}\n",
    "    return inputs\n",
    "\n",
    "\n",
    "calibration_dataset = nncf.Dataset(data_source, transform_fn)\n",
    "# Quantize the model. By specifying model_type, we specify additional transformer patterns in the model.\n",
    "quantized_model = nncf.quantize(model, calibration_dataset, model_type=ModelType.TRANSFORMER)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "da83574c-7abc-40a8-ae30-431c1b2bd823",
   "metadata": {},
   "outputs": [],
   "source": [
    "compressed_model_xml = Path(MODEL_DIR) / \"quantized_bert_mrpc.xml\"\n",
    "ov.save_model(quantized_model, compressed_model_xml)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "7c30ab44",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Load and Test OpenVINO Model\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "To load and test converted model, perform the following:\n",
    "\n",
    "* Load the model and compile it for selected device.\n",
    "* Prepare the input.\n",
    "* Run the inference.\n",
    "* Get the answer from the model output."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "93749c47-073f-4ffe-a507-4d38447159f5",
   "metadata": {},
   "source": [
    "### Select inference device\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "select device from dropdown list for running inference using OpenVINO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "6436bcf3-e446-4fee-a6ed-58235119a18d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3cdd9d273ae64886993d433f91a7289a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dropdown(description='Device:', index=2, options=('CPU', 'GPU', 'AUTO'), value='AUTO')"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "device = device_widget()\n",
    "\n",
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "4d79b1a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compile the model for a specific device.\n",
    "compiled_quantized_model = core.compile_model(model=quantized_model, device_name=device.value)\n",
    "output_layer = compiled_quantized_model.outputs[0]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "ef1d846e",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "The Data Source returns a pair of sentences (indicated by `sample_idx`) and the inference compares these sentences and outputs whether their meaning is the same. You can test other sentences by changing `sample_idx` to another value (from 0 to 407)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "e72504b9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Text 1: Wal-Mart said it would check all of its million-plus domestic workers to ensure they were legally employed .\n",
      "Text 2: It has also said it would review all of its domestic employees more than 1 million to ensure they have legal status .\n",
      "The same meaning: yes\n"
     ]
    }
   ],
   "source": [
    "sample_idx = 5\n",
    "sample = data_source[sample_idx]\n",
    "inputs = {k: torch.unsqueeze(torch.tensor(sample[k]), 0) for k in [\"input_ids\", \"token_type_ids\", \"attention_mask\"]}\n",
    "\n",
    "result = compiled_quantized_model(inputs)[output_layer]\n",
    "result = np.argmax(result)\n",
    "\n",
    "print(f\"Text 1: {sample['sentence1']}\")\n",
    "print(f\"Text 2: {sample['sentence2']}\")\n",
    "print(f\"The same meaning: {'yes' if result == 1 else 'no'}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "89920c37-dc2f-4177-b25f-bd8b1d0e34d3",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Compare F1-score of FP32 and INT8 models\n",
    "[back to top ⬆️](#Table-of-contents:)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "aeea7cc8-3eed-4474-8f59-ae63197368d1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Checking the accuracy of the original model:\n",
      "F1 score: 0.9019\n",
      "Checking the accuracy of the quantized model:\n",
      "F1 score: 0.8995\n"
     ]
    }
   ],
   "source": [
    "def validate(model: ov.Model, dataset: Iterable[Any]) -> float:\n",
    "    \"\"\"\n",
    "    Evaluate the model on GLUE dataset.\n",
    "    Returns F1 score metric.\n",
    "    \"\"\"\n",
    "    compiled_model = core.compile_model(model, device_name=device.value)\n",
    "    output_layer = compiled_model.output(0)\n",
    "\n",
    "    metric = evaluate.load(\"glue\", \"mrpc\")\n",
    "    for batch in dataset:\n",
    "        inputs = [np.expand_dims(np.asarray(batch[key], dtype=np.int64), 0) for key in INPUT_NAMES]\n",
    "        outputs = compiled_model(inputs)[output_layer]\n",
    "        predictions = outputs[0].argmax(axis=-1)\n",
    "        metric.add_batch(predictions=[predictions], references=[batch[\"labels\"]])\n",
    "    metrics = metric.compute()\n",
    "    f1_score = metrics[\"f1\"]\n",
    "\n",
    "    return f1_score\n",
    "\n",
    "\n",
    "print(\"Checking the accuracy of the original model:\")\n",
    "metric = validate(model, data_source)\n",
    "print(f\"F1 score: {metric:.4f}\")\n",
    "\n",
    "print(\"Checking the accuracy of the quantized model:\")\n",
    "metric = validate(quantized_model, data_source)\n",
    "print(f\"F1 score: {metric:.4f}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "4f67f6a2",
   "metadata": {
    "id": "vQACMfAUo52V",
    "tags": []
   },
   "source": [
    "## Compare Performance of the Original, Converted and Quantized Models\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "Compare the original PyTorch model with OpenVINO converted and quantized models (`FP32`, `INT8`) to see the difference in performance. It is expressed in Sentences Per Second (SPS) measure, which is the same as Frames Per Second (FPS) for images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "734ae69a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compile the model for a specific device.\n",
    "compiled_model = core.compile_model(model=model, device_name=device.value)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "f484fff2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PyTorch model on CPU: 0.080 seconds per sentence, SPS: 12.47\n",
      "IR FP32 model in OpenVINO Runtime/AUTO: 0.024 seconds per sentence, SPS: 41.92\n",
      "OpenVINO IR INT8 model in OpenVINO Runtime/AUTO: 0.012 seconds per sentence, SPS: 84.38\n"
     ]
    }
   ],
   "source": [
    "num_samples = 50\n",
    "sample = data_source[0]\n",
    "inputs = {k: torch.unsqueeze(torch.tensor(sample[k]), 0) for k in [\"input_ids\", \"token_type_ids\", \"attention_mask\"]}\n",
    "\n",
    "with torch.no_grad():\n",
    "    start = time.perf_counter()\n",
    "    for _ in range(num_samples):\n",
    "        torch_model(torch.vstack(list(inputs.values())))\n",
    "    end = time.perf_counter()\n",
    "    time_torch = end - start\n",
    "print(f\"PyTorch model on CPU: {time_torch / num_samples:.3f} seconds per sentence, \" f\"SPS: {num_samples / time_torch:.2f}\")\n",
    "\n",
    "start = time.perf_counter()\n",
    "for _ in range(num_samples):\n",
    "    compiled_model(inputs)\n",
    "end = time.perf_counter()\n",
    "time_ir = end - start\n",
    "print(f\"IR FP32 model in OpenVINO Runtime/{device.value}: {time_ir / num_samples:.3f} \" f\"seconds per sentence, SPS: {num_samples / time_ir:.2f}\")\n",
    "\n",
    "start = time.perf_counter()\n",
    "for _ in range(num_samples):\n",
    "    compiled_quantized_model(inputs)\n",
    "end = time.perf_counter()\n",
    "time_ir = end - start\n",
    "print(f\"OpenVINO IR INT8 model in OpenVINO Runtime/{device.value}: {time_ir / num_samples:.3f} \" f\"seconds per sentence, SPS: {num_samples / time_ir:.2f}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "add78af0",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Finally, measure the inference performance of OpenVINO `FP32` and `INT8` models. For this purpose, use [Benchmark Tool](https://docs.openvino.ai/2024/learn-openvino/openvino-samples/benchmark-tool.html) in OpenVINO.\n",
    "\n",
    "> **Note**: The `benchmark_app` tool is able to measure the performance of the OpenVINO Intermediate Representation (OpenVINO IR) models only. For more accurate performance, 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": "f71b38a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Inference FP32 model (OpenVINO IR)\n",
    "!benchmark_app -m $ir_model_xml -shape [1,128],[1,128],[1,128] -d {device.value} -api sync"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fdf41525",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Inference INT8 model (OpenVINO IR)\n",
    "! benchmark_app -m $compressed_model_xml -shape [1,128],[1,128],[1,128] -d {device.value} -api sync"
   ]
  }
 ],
 "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"
  },
  "openvino_notebooks": {
   "imageUrl": "",
   "tags": {
    "categories": [
     "API Overview",
     "Optimize"
    ],
    "libraries": [],
    "other": [],
    "tasks": [
     "Text Classification"
    ]
   }
  },
  "vscode": {
   "interpreter": {
    "hash": "cec18e25feb9469b5ff1085a8097bdcd86db6a4ac301d6aeff87d0f3e7ce4ca5"
   }
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}