docs/tutorials/pytorch/question-answering/bert-large-uncased-whole-word-masking-finetuned-squad.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "17ec3098",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Introduction "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69e00856",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "This tutorial demonstrates how to quantize a BERT model with both static and dynamic post training quantization based on [Intel® Neural Compressor](https://github.com/intel/neural-compressor) and benchmark the quantized models. For better int8 performance benefit, multi-instance benchmarking with 4 cores/instance is recommended."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a95b92e7",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "0b9ed365",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "0cbd8bd4",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Prerequisite"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86b20e2b",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Install packages"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1816be1",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "* Follow [installation](https://github.com/intel-innersource/frameworks.ai.nlp-toolkit.intel-nlp-toolkit#installation) to install **nlp-toolkit**. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef103e44",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# install model dependency\n",
    "!pip install datasets>=1.8.0 torch>=1.10.0 transformers>=4.12.0 wandb"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e35d9563",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Import packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b424903",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import logging\n",
    "import os\n",
    "import timeit\n",
    "import transformers\n",
    "from dataclasses import dataclass, field\n",
    "from datasets import load_dataset, load_metric\n",
    "from intel_extension_for_transformers.transformers import metrics , QuantizationConfig\n",
    "from transformers import (\n",
    "    AutoConfig,\n",
    "    AutoModelForQuestionAnswering,\n",
    "    AutoTokenizer,\n",
    "    DataCollatorWithPadding,\n",
    "    EvalPrediction,\n",
    "    TrainingArguments,\n",
    "    set_seed,\n",
    ")\n",
    "from transformers.utils.versions import require_version\n",
    "from typing import Optional\n",
    "\n",
    "\n",
    "require_version(\"datasets>=1.8.0\", \"To fix: pip install -r examples/pytorch/huggingface/question-answering/quantization/requirements.txt\")\n",
    "\n",
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "os.environ[\"WANDB_DISABLED\"] = \"true\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## import non-local Class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from transformers import  is_torch_tpu_available\n",
    "from intel_extension_for_transformers.transformers.trainer import NLPTrainer\n",
    "from transformers.trainer_utils import PredictionOutput\n",
    "\n",
    "if is_torch_tpu_available():\n",
    "    import torch_xla.core.xla_model as xm\n",
    "    import torch_xla.debug.metrics as met\n",
    "\n",
    "class QuestionAnsweringTrainer(NLPTrainer):\n",
    "    def __init__(self, *args, eval_examples=None, post_process_function=None, **kwargs):\n",
    "        super().__init__(*args, **kwargs)\n",
    "        self.eval_examples = eval_examples\n",
    "        self.post_process_function = post_process_function\n",
    "\n",
    "    def evaluate(self, eval_dataset=None, eval_examples=None, ignore_keys=None, metric_key_prefix: str = \"eval\"):\n",
    "        eval_dataset = self.eval_dataset if eval_dataset is None else eval_dataset\n",
    "        eval_dataloader = self.get_eval_dataloader(eval_dataset)\n",
    "        eval_examples = self.eval_examples if eval_examples is None else eval_examples\n",
    "\n",
    "        # Temporarily disable metric computation, we will do it in the loop here.\n",
    "        compute_metrics = self.compute_metrics\n",
    "        self.compute_metrics = None\n",
    "        eval_loop = self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop\n",
    "        try:\n",
    "            output = eval_loop(\n",
    "                eval_dataloader,\n",
    "                description=\"Evaluation\",\n",
    "                # No point gathering the predictions if there are no metrics, otherwise we defer to\n",
    "                # self.args.prediction_loss_only\n",
    "                prediction_loss_only=True if compute_metrics is None else None,\n",
    "                ignore_keys=ignore_keys,\n",
    "            )\n",
    "        finally:\n",
    "            self.compute_metrics = compute_metrics\n",
    "\n",
    "        if self.post_process_function is not None and self.compute_metrics is not None:\n",
    "            eval_preds = self.post_process_function(eval_examples, eval_dataset, output.predictions)\n",
    "            metrics = self.compute_metrics(eval_preds)\n",
    "\n",
    "            # Prefix all keys with metric_key_prefix + '_'\n",
    "            for key in list(metrics.keys()):\n",
    "                if not key.startswith(f\"{metric_key_prefix}_\"):\n",
    "                    metrics[f\"{metric_key_prefix}_{key}\"] = metrics.pop(key)\n",
    "\n",
    "            self.log(metrics)\n",
    "        else:\n",
    "            metrics = {}\n",
    "\n",
    "        if self.args.tpu_metrics_debug or self.args.debug:\n",
    "            # tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)\n",
    "            xm.master_print(met.metrics_report())\n",
    "\n",
    "        self.control = self.callback_handler.on_evaluate(self.args, self.state, self.control, metrics)\n",
    "        return metrics\n",
    "\n",
    "    def predict(self, predict_dataset, predict_examples, ignore_keys=None, metric_key_prefix: str = \"test\"):\n",
    "        predict_dataloader = self.get_test_dataloader(predict_dataset)\n",
    "\n",
    "        # Temporarily disable metric computation, we will do it in the loop here.\n",
    "        compute_metrics = self.compute_metrics\n",
    "        self.compute_metrics = None\n",
    "        eval_loop = self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop\n",
    "        try:\n",
    "            output = eval_loop(\n",
    "                predict_dataloader,\n",
    "                description=\"Prediction\",\n",
    "                # No point gathering the predictions if there are no metrics, otherwise we defer to\n",
    "                # self.args.prediction_loss_only\n",
    "                prediction_loss_only=True if compute_metrics is None else None,\n",
    "                ignore_keys=ignore_keys,\n",
    "            )\n",
    "        finally:\n",
    "            self.compute_metrics = compute_metrics\n",
    "\n",
    "        if self.post_process_function is None or self.compute_metrics is None:\n",
    "            return output\n",
    "\n",
    "        predictions = self.post_process_function(predict_examples, predict_dataset, output.predictions, \"predict\")\n",
    "        metrics = self.compute_metrics(predictions)\n",
    "\n",
    "        # Prefix all keys with metric_key_prefix + '_'\n",
    "        for key in list(metrics.keys()):\n",
    "            if not key.startswith(f\"{metric_key_prefix}_\"):\n",
    "                metrics[f\"{metric_key_prefix}_{key}\"] = metrics.pop(key)\n",
    "\n",
    "        return PredictionOutput(predictions=predictions.predictions, label_ids=predictions.label_ids, metrics=metrics)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## import non-local Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import collections\n",
    "import json\n",
    "import logging\n",
    "import os\n",
    "from typing import Optional, Tuple\n",
    "\n",
    "import numpy as np\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "def postprocess_qa_predictions(\n",
    "    examples,\n",
    "    features,\n",
    "    predictions: Tuple[np.ndarray, np.ndarray],\n",
    "    version_2_with_negative: bool = False,\n",
    "    n_best_size: int = 20,\n",
    "    max_answer_length: int = 30,\n",
    "    null_score_diff_threshold: float = 0.0,\n",
    "    output_dir: Optional[str] = None,\n",
    "    prefix: Optional[str] = None,\n",
    "    log_level: Optional[int] = logging.WARNING,\n",
    "):\n",
    "    \"\"\"\n",
    "    Post-processes the predictions of a question-answering model to convert them to answers that are substrings of the\n",
    "    original contexts. This is the base postprocessing functions for models that only return start and end logits.\n",
    "\n",
    "    Args:\n",
    "        examples: The non-preprocessed dataset (see the main script for more information).\n",
    "        features: The processed dataset (see the main script for more information).\n",
    "        predictions (:obj:`Tuple[np.ndarray, np.ndarray]`):\n",
    "            The predictions of the model: two arrays containing the start logits and the end logits respectively. Its\n",
    "            first dimension must match the number of elements of :obj:`features`.\n",
    "        version_2_with_negative (:obj:`bool`, `optional`, defaults to :obj:`False`):\n",
    "            Whether or not the underlying dataset contains examples with no answers.\n",
    "        n_best_size (:obj:`int`, `optional`, defaults to 20):\n",
    "            The total number of n-best predictions to generate when looking for an answer.\n",
    "        max_answer_length (:obj:`int`, `optional`, defaults to 30):\n",
    "            The maximum length of an answer that can be generated. This is needed because the start and end predictions\n",
    "            are not conditioned on one another.\n",
    "        null_score_diff_threshold (:obj:`float`, `optional`, defaults to 0):\n",
    "            The threshold used to select the null answer: if the best answer has a score that is less than the score of\n",
    "            the null answer minus this threshold, the null answer is selected for this example (note that the score of\n",
    "            the null answer for an example giving several features is the minimum of the scores for the null answer on\n",
    "            each feature: all features must be aligned on the fact they `want` to predict a null answer).\n",
    "\n",
    "            Only useful when :obj:`version_2_with_negative` is :obj:`True`.\n",
    "        output_dir (:obj:`str`, `optional`):\n",
    "            If provided, the dictionaries of predictions, n_best predictions (with their scores and logits) and, if\n",
    "            :obj:`version_2_with_negative=True`, the dictionary of the scores differences between best and null\n",
    "            answers, are saved in `output_dir`.\n",
    "        prefix (:obj:`str`, `optional`):\n",
    "            If provided, the dictionaries mentioned above are saved with `prefix` added to their names.\n",
    "        log_level (:obj:`int`, `optional`, defaults to ``logging.WARNING``):\n",
    "            ``logging`` log level (e.g., ``logging.WARNING``)\n",
    "    \"\"\"\n",
    "\n",
    "    if len(predictions) != 2:\n",
    "        raise ValueError(\"`predictions` should be a tuple with two elements (start_logits, end_logits).\")\n",
    "    all_start_logits, all_end_logits = predictions\n",
    "\n",
    "    if len(predictions[0]) != len(features):\n",
    "        raise ValueError(f\"Got {len(predictions[0])} predictions and {len(features)} features.\")\n",
    "\n",
    "    # Build a map example to its corresponding features.\n",
    "    example_id_to_index = {k: i for i, k in enumerate(examples[\"id\"])}\n",
    "    features_per_example = collections.defaultdict(list)\n",
    "    for i, feature in enumerate(features):\n",
    "        features_per_example[example_id_to_index[feature[\"example_id\"]]].append(i)\n",
    "\n",
    "    # The dictionaries we have to fill.\n",
    "    all_predictions = collections.OrderedDict()\n",
    "    all_nbest_json = collections.OrderedDict()\n",
    "    if version_2_with_negative:\n",
    "        scores_diff_json = collections.OrderedDict()\n",
    "\n",
    "    # Logging.\n",
    "    logger.setLevel(log_level)\n",
    "    logger.info(f\"Post-processing {len(examples)} example predictions split into {len(features)} features.\")\n",
    "\n",
    "    # Let's loop over all the examples!\n",
    "    for example_index, example in enumerate(tqdm(examples)):\n",
    "        # Those are the indices of the features associated to the current example.\n",
    "        feature_indices = features_per_example[example_index]\n",
    "\n",
    "        min_null_prediction = None\n",
    "        prelim_predictions = []\n",
    "\n",
    "        # Looping through all the features associated to the current example.\n",
    "        for feature_index in feature_indices:\n",
    "            # We grab the predictions of the model for this feature.\n",
    "            start_logits = all_start_logits[feature_index]\n",
    "            end_logits = all_end_logits[feature_index]\n",
    "            # This is what will allow us to map some the positions in our logits to span of texts in the original\n",
    "            # context.\n",
    "            offset_mapping = features[feature_index][\"offset_mapping\"]\n",
    "            # Optional `token_is_max_context`, if provided we will remove answers that do not have the maximum context\n",
    "            # available in the current feature.\n",
    "            token_is_max_context = features[feature_index].get(\"token_is_max_context\", None)\n",
    "\n",
    "            # Update minimum null prediction.\n",
    "            feature_null_score = start_logits[0] + end_logits[0]\n",
    "            if min_null_prediction is None or min_null_prediction[\"score\"] > feature_null_score:\n",
    "                min_null_prediction = {\n",
    "                    \"offsets\": (0, 0),\n",
    "                    \"score\": feature_null_score,\n",
    "                    \"start_logit\": start_logits[0],\n",
    "                    \"end_logit\": end_logits[0],\n",
    "                }\n",
    "\n",
    "            # Go through all possibilities for the `n_best_size` greater start and end logits.\n",
    "            start_indexes = np.argsort(start_logits)[-1 : -n_best_size - 1 : -1].tolist()\n",
    "            end_indexes = np.argsort(end_logits)[-1 : -n_best_size - 1 : -1].tolist()\n",
    "            for start_index in start_indexes:\n",
    "                for end_index in end_indexes:\n",
    "                    # Don't consider out-of-scope answers, either because the indices are out of bounds or correspond\n",
    "                    # to part of the input_ids that are not in the context.\n",
    "                    if (\n",
    "                        start_index >= len(offset_mapping)\n",
    "                        or end_index >= len(offset_mapping)\n",
    "                        or offset_mapping[start_index] is None\n",
    "                        or len(offset_mapping[start_index]) < 2\n",
    "                        or offset_mapping[end_index] is None\n",
    "                        or len(offset_mapping[end_index]) < 2\n",
    "                    ):\n",
    "                        continue\n",
    "                    # Don't consider answers with a length that is either < 0 or > max_answer_length.\n",
    "                    if end_index < start_index or end_index - start_index + 1 > max_answer_length:\n",
    "                        continue\n",
    "                    # Don't consider answer that don't have the maximum context available (if such information is\n",
    "                    # provided).\n",
    "                    if token_is_max_context is not None and not token_is_max_context.get(str(start_index), False):\n",
    "                        continue\n",
    "\n",
    "                    prelim_predictions.append(\n",
    "                        {\n",
    "                            \"offsets\": (offset_mapping[start_index][0], offset_mapping[end_index][1]),\n",
    "                            \"score\": start_logits[start_index] + end_logits[end_index],\n",
    "                            \"start_logit\": start_logits[start_index],\n",
    "                            \"end_logit\": end_logits[end_index],\n",
    "                        }\n",
    "                    )\n",
    "        if version_2_with_negative:\n",
    "            # Add the minimum null prediction\n",
    "            prelim_predictions.append(min_null_prediction)\n",
    "            null_score = min_null_prediction[\"score\"]\n",
    "\n",
    "        # Only keep the best `n_best_size` predictions.\n",
    "        predictions = sorted(prelim_predictions, key=lambda x: x[\"score\"], reverse=True)[:n_best_size]\n",
    "\n",
    "        # Add back the minimum null prediction if it was removed because of its low score.\n",
    "        if version_2_with_negative and not any(p[\"offsets\"] == (0, 0) for p in predictions):\n",
    "            predictions.append(min_null_prediction)\n",
    "\n",
    "        # Use the offsets to gather the answer text in the original context.\n",
    "        context = example[\"context\"]\n",
    "        for pred in predictions:\n",
    "            offsets = pred.pop(\"offsets\")\n",
    "            pred[\"text\"] = context[offsets[0] : offsets[1]]\n",
    "\n",
    "        # In the very rare edge case we have not a single non-null prediction, we create a fake prediction to avoid\n",
    "        # failure.\n",
    "        if len(predictions) == 0 or (len(predictions) == 1 and predictions[0][\"text\"] == \"\"):\n",
    "            predictions.insert(0, {\"text\": \"empty\", \"start_logit\": 0.0, \"end_logit\": 0.0, \"score\": 0.0})\n",
    "\n",
    "        # Compute the softmax of all scores (we do it with numpy to stay independent from torch/tf in this file, using\n",
    "        # the LogSumExp trick).\n",
    "        scores = np.array([pred.pop(\"score\") for pred in predictions])\n",
    "        exp_scores = np.exp(scores - np.max(scores))\n",
    "        probs = exp_scores / exp_scores.sum()\n",
    "\n",
    "        # Include the probabilities in our predictions.\n",
    "        for prob, pred in zip(probs, predictions):\n",
    "            pred[\"probability\"] = prob\n",
    "\n",
    "        # Pick the best prediction. If the null answer is not possible, this is easy.\n",
    "        if not version_2_with_negative:\n",
    "            all_predictions[example[\"id\"]] = predictions[0][\"text\"]\n",
    "        else:\n",
    "            # Otherwise we first need to find the best non-empty prediction.\n",
    "            i = 0\n",
    "            while predictions[i][\"text\"] == \"\":\n",
    "                i += 1\n",
    "            best_non_null_pred = predictions[i]\n",
    "\n",
    "            # Then we compare to the null prediction using the threshold.\n",
    "            score_diff = null_score - best_non_null_pred[\"start_logit\"] - best_non_null_pred[\"end_logit\"]\n",
    "            scores_diff_json[example[\"id\"]] = float(score_diff)  # To be JSON-serializable.\n",
    "            if score_diff > null_score_diff_threshold:\n",
    "                all_predictions[example[\"id\"]] = \"\"\n",
    "            else:\n",
    "                all_predictions[example[\"id\"]] = best_non_null_pred[\"text\"]\n",
    "\n",
    "        # Make `predictions` JSON-serializable by casting np.float back to float.\n",
    "        all_nbest_json[example[\"id\"]] = [\n",
    "            {k: (float(v) if isinstance(v, (np.float16, np.float32, np.float64)) else v) for k, v in pred.items()}\n",
    "            for pred in predictions\n",
    "        ]\n",
    "\n",
    "    # If we have an output_dir, let's save all those dicts.\n",
    "    if output_dir is not None:\n",
    "        if not os.path.isdir(output_dir):\n",
    "            raise EnvironmentError(f\"{output_dir} is not a directory.\")\n",
    "\n",
    "        prediction_file = os.path.join(\n",
    "            output_dir, \"predictions.json\" if prefix is None else f\"{prefix}_predictions.json\"\n",
    "        )\n",
    "        nbest_file = os.path.join(\n",
    "            output_dir, \"nbest_predictions.json\" if prefix is None else f\"{prefix}_nbest_predictions.json\"\n",
    "        )\n",
    "        if version_2_with_negative:\n",
    "            null_odds_file = os.path.join(\n",
    "                output_dir, \"null_odds.json\" if prefix is None else f\"{prefix}_null_odds.json\"\n",
    "            )\n",
    "\n",
    "        logger.info(f\"Saving predictions to {prediction_file}.\")\n",
    "        with open(prediction_file, \"w\") as writer:\n",
    "            writer.write(json.dumps(all_predictions, indent=4) + \"\\n\")\n",
    "        logger.info(f\"Saving nbest_preds to {nbest_file}.\")\n",
    "        with open(nbest_file, \"w\") as writer:\n",
    "            writer.write(json.dumps(all_nbest_json, indent=4) + \"\\n\")\n",
    "        if version_2_with_negative:\n",
    "            logger.info(f\"Saving null_odds to {null_odds_file}.\")\n",
    "            with open(null_odds_file, \"w\") as writer:\n",
    "                writer.write(json.dumps(scores_diff_json, indent=4) + \"\\n\")\n",
    "\n",
    "    return all_predictions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Define arguments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aead749f",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class ModelArguments:\n",
    "    \"\"\"\n",
    "    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.\n",
    "    \"\"\"\n",
    "    model_name_or_path: str = field(\n",
    "        metadata={\"help\": \"Path to pretrained model or model identifier from huggingface.co/models\"}\n",
    "    )\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class DataTrainingArguments:\n",
    "    \"\"\"\n",
    "    Arguments pertaining to what data we are going to input our model for training and eval.\n",
    "    \"\"\"\n",
    "    dataset_name: Optional[str] = field(\n",
    "        default=None, metadata={\"help\": \"The name of the dataset to use (via the datasets library).\"}\n",
    "    )\n",
    "    max_seq_length: int = field(\n",
    "        default=384,\n",
    "        metadata={\n",
    "            \"help\": \"The maximum total input sequence length after tokenization. Sequences longer \"\n",
    "            \"than this will be truncated, sequences shorter will be padded.\"\n",
    "        },\n",
    "    )\n",
    "    max_train_samples: Optional[int] = field(\n",
    "        default=None,\n",
    "        metadata={\n",
    "            \"help\": \"For debugging purposes or quicker training, truncate the number of training examples to this \"\n",
    "            \"value if set.\"\n",
    "        },\n",
    "    )\n",
    "    max_eval_samples: Optional[int] = field(\n",
    "        default=None,\n",
    "        metadata={\n",
    "            \"help\": \"For debugging purposes or quicker training, truncate the number of evaluation examples to this \"\n",
    "            \"value if set.\"\n",
    "        },\n",
    "    )\n",
    "    overwrite_cache: bool = field(\n",
    "        default=False, metadata={\"help\": \"Overwrite the cached training and evaluation sets\"}\n",
    "    )\n",
    "    doc_stride: int = field(\n",
    "        default=128,\n",
    "        metadata={\"help\": \"When splitting up a long document into chunks, how much stride to take between chunks.\"},\n",
    "    )\n",
    "    pad_to_max_length: bool = field(\n",
    "        default=True,\n",
    "        metadata={\n",
    "            \"help\": \"Whether to pad all samples to `max_seq_length`. \"\n",
    "            \"If False, will pad the samples dynamically when batching to the maximum length in the batch (which can \"\n",
    "            \"be faster on GPU but will be slower on TPU).\"\n",
    "        },\n",
    "    )\n",
    "    version_2_with_negative: bool = field(\n",
    "        default=False, metadata={\"help\": \"If true, some of the examples do not have an answer.\"}\n",
    "    )\n",
    "    null_score_diff_threshold: float = field(\n",
    "        default=0.0,\n",
    "        metadata={\n",
    "            \"help\": \"The threshold used to select the null answer: if the best answer has a score that is less than \"\n",
    "            \"the score of the null answer minus this threshold, the null answer is selected for this example. \"\n",
    "            \"Only useful when `version_2_with_negative=True`.\"\n",
    "        },\n",
    "    )\n",
    "    n_best_size: int = field(\n",
    "        default=20,\n",
    "        metadata={\"help\": \"The total number of n-best predictions to generate when looking for an answer.\"},\n",
    "    )\n",
    "    max_answer_length: int = field(\n",
    "        default=30,\n",
    "        metadata={\n",
    "            \"help\": \"The maximum length of an answer that can be generated. This is needed because the start \"\n",
    "            \"and end predictions are not conditioned on one another.\"\n",
    "        },\n",
    "    )\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class OptimizationArguments:\n",
    "    \"\"\"\n",
    "    Arguments pertaining to what type of optimization we are going to apply on the model.\n",
    "    \"\"\"\n",
    "\n",
    "    tune: bool = field(\n",
    "        default=False,\n",
    "        metadata={\"help\": \"Whether or not to apply quantization.\"},\n",
    "    )\n",
    "    quantization_approach: Optional[str] = field(\n",
    "        default=\"POSTTRAININGSTATIC\",\n",
    "        metadata={\"help\": \"Quantization approach. Supported approach are POSTTRAININGSTATIC, \"\n",
    "                  \"POSTTRAININGDYNAMIC and QUANTIZATIONAWARETRAINING.\"},\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11b035ee",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "model_args = ModelArguments(\n",
    "    model_name_or_path=\"distilbert-base-uncased-distilled-squad\",\n",
    ")\n",
    "data_args = DataTrainingArguments(\n",
    "    dataset_name=\"squad\",\n",
    "    max_seq_length=384,\n",
    "    max_eval_samples=5000\n",
    ")\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=\"./tmp/squad_output\",\n",
    "    do_eval=True,\n",
    "    do_train=True,\n",
    "    no_cuda=True,\n",
    "    overwrite_output_dir=True,\n",
    "    per_device_train_batch_size=8,\n",
    ")\n",
    "optim_args = OptimizationArguments(\n",
    "    tune=True,\n",
    "    quantization_approach=\"PostTrainingStatic\"\n",
    ")\n",
    "log_level = training_args.get_process_log_level()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "522f4153",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Download dataset from the hub"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22bcb0fb",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "raw_datasets = load_dataset(\"squad\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d162019c",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Download fp32 model from the hub"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2620a5f",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Set seed before initializing model.\n",
    "set_seed(training_args.seed)\n",
    "\n",
    "# get fp32 model\n",
    "config = AutoConfig.from_pretrained(model_args.model_name_or_path)\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path, use_fast=True)\n",
    "model = AutoModelForQuestionAnswering.from_pretrained(\n",
    "    model_args.model_name_or_path,\n",
    "    from_tf=bool(\".ckpt\" in model_args.model_name_or_path),\n",
    "    config=config,\n",
    "    use_auth_token=None\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2634a186",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Preprocessing the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a873a324",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Preprocessing the datasets.\n",
    "# Preprocessing is slighlty different for training and evaluation.\n",
    "column_names = raw_datasets[\"train\"].column_names\n",
    "question_column_name = \"question\" if \"question\" in column_names else column_names[0]\n",
    "context_column_name = \"context\" if \"context\" in column_names else column_names[1]\n",
    "answer_column_name = \"answers\" if \"answers\" in column_names else column_names[2]\n",
    "\n",
    "# Padding side determines if we do (question|context) or (context|question).\n",
    "pad_on_right = tokenizer.padding_side == \"right\"\n",
    "\n",
    "max_seq_length = min(data_args.max_seq_length, tokenizer.model_max_length)\n",
    "\n",
    "# Training preprocessing\n",
    "def prepare_train_features(examples):\n",
    "    # Some of the questions have lots of whitespace on the left, which is not useful and will make the\n",
    "    # truncation of the context fail (the tokenized question will take a lots of space). So we remove that\n",
    "    # left whitespace\n",
    "    examples[question_column_name] = [q.lstrip() for q in examples[question_column_name]]\n",
    "\n",
    "    # Tokenize our examples with truncation and maybe padding, but keep the overflows using a stride. This results\n",
    "    # in one example possible giving several features when a context is long, each of those features having a\n",
    "    # context that overlaps a bit the context of the previous feature.\n",
    "    tokenized_examples = tokenizer(\n",
    "        examples[question_column_name if pad_on_right else context_column_name],\n",
    "        examples[context_column_name if pad_on_right else question_column_name],\n",
    "        truncation=\"only_second\" if pad_on_right else \"only_first\",\n",
    "        max_length=max_seq_length,\n",
    "        stride=data_args.doc_stride,\n",
    "        return_overflowing_tokens=True,\n",
    "        return_offsets_mapping=True,\n",
    "        padding=\"max_length\" if data_args.pad_to_max_length else False,\n",
    "    )\n",
    "\n",
    "    # Since one example might give us several features if it has a long context, we need a map from a feature to\n",
    "    # its corresponding example. This key gives us just that.\n",
    "    sample_mapping = tokenized_examples.pop(\"overflow_to_sample_mapping\")\n",
    "    # The offset mappings will give us a map from token to character position in the original context. This will\n",
    "    # help us compute the start_positions and end_positions.\n",
    "    offset_mapping = tokenized_examples.pop(\"offset_mapping\")\n",
    "\n",
    "    # Let's label those examples!\n",
    "    tokenized_examples[\"start_positions\"] = []\n",
    "    tokenized_examples[\"end_positions\"] = []\n",
    "\n",
    "    for i, offsets in enumerate(offset_mapping):\n",
    "        # We will label impossible answers with the index of the CLS token.\n",
    "        input_ids = tokenized_examples[\"input_ids\"][i]\n",
    "        cls_index = input_ids.index(tokenizer.cls_token_id)\n",
    "\n",
    "        # Grab the sequence corresponding to that example (to know what is the context and what is the question).\n",
    "        sequence_ids = tokenized_examples.sequence_ids(i)\n",
    "\n",
    "        # One example can give several spans, this is the index of the example containing this span of text.\n",
    "        sample_index = sample_mapping[i]\n",
    "        answers = examples[answer_column_name][sample_index]\n",
    "        # If no answers are given, set the cls_index as answer.\n",
    "        if len(answers[\"answer_start\"]) == 0:\n",
    "            tokenized_examples[\"start_positions\"].append(cls_index)\n",
    "            tokenized_examples[\"end_positions\"].append(cls_index)\n",
    "        else:\n",
    "            # Start/end character index of the answer in the text.\n",
    "            start_char = answers[\"answer_start\"][0]\n",
    "            end_char = start_char + len(answers[\"text\"][0])\n",
    "\n",
    "            # Start token index of the current span in the text.\n",
    "            token_start_index = 0\n",
    "            while sequence_ids[token_start_index] != (1 if pad_on_right else 0):\n",
    "                token_start_index += 1\n",
    "\n",
    "            # End token index of the current span in the text.\n",
    "            token_end_index = len(input_ids) - 1\n",
    "            while sequence_ids[token_end_index] != (1 if pad_on_right else 0):\n",
    "                token_end_index -= 1\n",
    "\n",
    "            # Detect if the answer is out of the span (in which case this feature is labeled with the CLS index).\n",
    "            if not (offsets[token_start_index][0] <= start_char and offsets[token_end_index][1] >= end_char):\n",
    "                tokenized_examples[\"start_positions\"].append(cls_index)\n",
    "                tokenized_examples[\"end_positions\"].append(cls_index)\n",
    "            else:\n",
    "                # Otherwise move the token_start_index and token_end_index to the two ends of the answer.\n",
    "                # Note: we could go after the last offset if the answer is the last word (edge case).\n",
    "                while token_start_index < len(offsets) and offsets[token_start_index][0] <= start_char:\n",
    "                    token_start_index += 1\n",
    "                tokenized_examples[\"start_positions\"].append(token_start_index - 1)\n",
    "                while offsets[token_end_index][1] >= end_char:\n",
    "                    token_end_index -= 1\n",
    "                tokenized_examples[\"end_positions\"].append(token_end_index + 1)\n",
    "\n",
    "    return tokenized_examples\n",
    "\n",
    "if training_args.do_train:\n",
    "    if \"train\" not in raw_datasets:\n",
    "        raise ValueError(\"--do_train requires a train dataset\")\n",
    "    train_dataset = raw_datasets[\"train\"]\n",
    "    with training_args.main_process_first(desc=\"train dataset map pre-processing\"):\n",
    "        train_dataset = train_dataset.map(\n",
    "            prepare_train_features,\n",
    "            batched=True,\n",
    "            remove_columns=column_names,\n",
    "            load_from_cache_file=not data_args.overwrite_cache,\n",
    "            desc=\"Running tokenizer on train dataset\"\n",
    "        )\n",
    "    if data_args.max_train_samples is not None:\n",
    "        # Number of samples might increase during Feature Creation, We select only specified max samples\n",
    "        max_train_samples = min(len(train_dataset), data_args.max_train_samples)\n",
    "        train_dataset = train_dataset.select(range(max_train_samples))\n",
    "\n",
    "# Validation preprocessing\n",
    "def prepare_validation_features(examples):\n",
    "    # Some of the questions have lots of whitespace on the left, which is not useful and will make the\n",
    "    # truncation of the context fail (the tokenized question will take a lots of space). So we remove that\n",
    "    # left whitespace\n",
    "    examples[question_column_name] = [q.lstrip() for q in examples[question_column_name]]\n",
    "\n",
    "    # Tokenize our examples with truncation and maybe padding, but keep the overflows using a stride. This results\n",
    "    # in one example possible giving several features when a context is long, each of those features having a\n",
    "    # context that overlaps a bit the context of the previous feature.\n",
    "    tokenized_examples = tokenizer(\n",
    "        examples[question_column_name if pad_on_right else context_column_name],\n",
    "        examples[context_column_name if pad_on_right else question_column_name],\n",
    "        truncation=\"only_second\" if pad_on_right else \"only_first\",\n",
    "        max_length=max_seq_length,\n",
    "        stride=data_args.doc_stride,\n",
    "        return_overflowing_tokens=True,\n",
    "        return_offsets_mapping=True,\n",
    "        padding=\"max_length\" if data_args.pad_to_max_length else False,\n",
    "    )\n",
    "\n",
    "    # Since one example might give us several features if it has a long context, we need a map from a feature to\n",
    "    # its corresponding example. This key gives us just that.\n",
    "    sample_mapping = tokenized_examples.pop(\"overflow_to_sample_mapping\")\n",
    "\n",
    "    # For evaluation, we will need to convert our predictions to substrings of the context, so we keep the\n",
    "    # corresponding example_id and we will store the offset mappings.\n",
    "    tokenized_examples[\"example_id\"] = []\n",
    "\n",
    "    for i in range(len(tokenized_examples[\"input_ids\"])):\n",
    "        # Grab the sequence corresponding to that example (to know what is the context and what is the question).\n",
    "        sequence_ids = tokenized_examples.sequence_ids(i)\n",
    "        context_index = 1 if pad_on_right else 0\n",
    "\n",
    "        # One example can give several spans, this is the index of the example containing this span of text.\n",
    "        sample_index = sample_mapping[i]\n",
    "        tokenized_examples[\"example_id\"].append(examples[\"id\"][sample_index])\n",
    "\n",
    "        # Set to None the offset_mapping that are not part of the context so it's easy to determine if a token\n",
    "        # position is part of the context or not.\n",
    "        tokenized_examples[\"offset_mapping\"][i] = [\n",
    "            (o if sequence_ids[k] == context_index else None)\n",
    "            for k, o in enumerate(tokenized_examples[\"offset_mapping\"][i])\n",
    "        ]\n",
    "\n",
    "    return tokenized_examples\n",
    "\n",
    "if training_args.do_eval:\n",
    "    if \"validation\" not in raw_datasets:\n",
    "        raise ValueError(\"--do_eval requires a validation dataset\")\n",
    "    eval_examples = raw_datasets[\"validation\"]\n",
    "    with training_args.main_process_first(desc=\"validation dataset map pre-processing\"):\n",
    "        eval_dataset = eval_examples.map(\n",
    "            prepare_validation_features,\n",
    "            batched=True,\n",
    "            remove_columns=column_names,\n",
    "            load_from_cache_file=not data_args.overwrite_cache,\n",
    "            desc=\"Running tokenizer on validation dataset\",\n",
    "        )\n",
    "    if data_args.max_eval_samples is not None:\n",
    "        # During Feature creation dataset samples might increase, we will select required samples again\n",
    "        max_eval_samples = min(len(eval_dataset), data_args.max_eval_samples)\n",
    "        eval_dataset = eval_dataset.select(range(max_eval_samples))\n",
    "\n",
    "# Data collator\n",
    "data_collator = DataCollatorWithPadding(tokenizer, pad_to_multiple_of=8 if training_args.fp16 else None)\n",
    "\n",
    "# Post-processing:\n",
    "def post_processing_function(examples, features, predictions, stage=\"eval\"):\n",
    "    # Post-processing: we match the start logits and end logits to answers in the original context.\n",
    "    predictions = postprocess_qa_predictions(\n",
    "        examples=examples,\n",
    "        features=features,\n",
    "        predictions=predictions,\n",
    "        version_2_with_negative=data_args.version_2_with_negative,\n",
    "        n_best_size=data_args.n_best_size,\n",
    "        max_answer_length=data_args.max_answer_length,\n",
    "        null_score_diff_threshold=data_args.null_score_diff_threshold,\n",
    "        output_dir=training_args.output_dir,\n",
    "        log_level=log_level,\n",
    "        prefix=stage,\n",
    "    )\n",
    "    # Format the result to the format the metric expects.\n",
    "    if data_args.version_2_with_negative:\n",
    "        formatted_predictions = [\n",
    "            {\"id\": k, \"prediction_text\": v, \"no_answer_probability\": 0.0} for k, v in predictions.items()\n",
    "        ]\n",
    "    else:\n",
    "        formatted_predictions = [{\"id\": k, \"prediction_text\": v} for k, v in predictions.items()]\n",
    "\n",
    "    references = [{\"id\": ex[\"id\"], \"answers\": ex[answer_column_name]} for ex in examples]\n",
    "    return EvalPrediction(predictions=formatted_predictions, label_ids=references)\n",
    "\n",
    "metric = load_metric(\"squad_v2\" if data_args.version_2_with_negative else \"squad\")\n",
    "\n",
    "def compute_metrics(p: EvalPrediction):\n",
    "    return metric.compute(predictions=p.predictions, references=p.label_ids)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "942d2967",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Quantization & Benchmark"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a07a8735",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Static Post Training Quantization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df04531d",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "set_seed(training_args.seed)\n",
    "# Initialize our Trainer\n",
    "trainer_static = QuestionAnsweringTrainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=train_dataset if training_args.do_train else None,\n",
    "    eval_dataset=eval_dataset if training_args.do_eval else None,\n",
    "    eval_examples=eval_examples if training_args.do_eval else None,\n",
    "    tokenizer=tokenizer,\n",
    "    data_collator=data_collator,\n",
    "    post_process_function=post_processing_function,\n",
    "    compute_metrics=compute_metrics,\n",
    ")\n",
    "\n",
    "trainer_static.save_model(training_args.output_dir+'/saved_results_static')\n",
    "model.config.save_pretrained(training_args.output_dir+'/saved_results_static')\n",
    "\n",
    "tune_metric = metrics.Metric(\n",
    "    name=\"eval_f1\", # Metric used for the tuning strategy.\n",
    "    is_relative=True, # Metric tolerance mode, True is for relative, otherwise for absolute.\n",
    "    criterion=0.25, # Performance tolerance when optimizing the model.\n",
    ")\n",
    "quantization_config = QuantizationConfig(\n",
    "    approach=\"PostTrainingStatic\",\n",
    "    max_trials=200,\n",
    "    metrics=[tune_metric],\n",
    ")\n",
    "\n",
    "# run quantization\n",
    "trainer_static.quantize(quant_config=quantization_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a7e93de",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Run Benchmark after Static Post Training Quantization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf54c034",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "set_seed(training_args.seed)\n",
    "start_time = timeit.default_timer()\n",
    "results = trainer_static.evaluate()\n",
    "evalTime = timeit.default_timer() - start_time\n",
    "max_eval_samples = data_args.max_eval_samples\n",
    "samples = min(max_eval_samples, len(eval_dataset))\n",
    "\n",
    "eval_f1_static = results.get(\"eval_f1\")\n",
    "print('Batch size = {}'.format(training_args.per_device_eval_batch_size))\n",
    "print(\"Finally Eval eval_f1 Accuracy: {}\".format(eval_f1_static))\n",
    "print(\"Latency: {:.3f} ms\".format(evalTime / samples * 1000))\n",
    "print(\"Throughput: {} samples/sec\".format(samples/evalTime))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a51f6ca",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Dynamic Post Training Quantization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a57aa0d6",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "set_seed(training_args.seed)\n",
    "# Initialize our Trainer\n",
    "trainer_dynamic = QuestionAnsweringTrainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=train_dataset if training_args.do_train else None,\n",
    "    eval_dataset=eval_dataset if training_args.do_eval else None,\n",
    "    eval_examples=eval_examples if training_args.do_eval else None,\n",
    "    tokenizer=tokenizer,\n",
    "    data_collator=data_collator,\n",
    "    post_process_function=post_processing_function,\n",
    "    compute_metrics=compute_metrics,\n",
    ")\n",
    "\n",
    "trainer_dynamic.save_model(training_args.output_dir+'/saved_results_dynamic')\n",
    "model.config.save_pretrained(training_args.output_dir+'/saved_results_dynamic')\n",
    "\n",
    "tune_metric = metrics.Metric(\n",
    "    name=\"eval_f1\", # Metric used for the tuning strategy.\n",
    "    is_relative=True, # Metric tolerance mode, True is for relative, otherwise for absolute.\n",
    "    criterion=0.25, # Performance tolerance when optimizing the model.\n",
    ")\n",
    "quantization_config = QuantizationConfig(\n",
    "    approach=\"PostTrainingDynamic\",\n",
    "    max_trials=200,\n",
    "    metrics=[tune_metric],\n",
    ")\n",
    "\n",
    "# run quantization\n",
    "trainer_dynamic.quantize(quant_config=quantization_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9dd5b041",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Run Benchmark after Dynamic Post Training Quantization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "689e288d",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "set_seed(training_args.seed)\n",
    "start_time = timeit.default_timer()\n",
    "results = trainer_dynamic.evaluate()\n",
    "evalTime = timeit.default_timer() - start_time\n",
    "max_eval_samples = data_args.max_eval_samples\n",
    "samples = min(max_eval_samples, len(eval_dataset))\n",
    "\n",
    "eval_f1_dynamic = results.get(\"eval_f1\")\n",
    "print('Batch size = {}'.format(training_args.per_device_eval_batch_size))\n",
    "print(\"Finally Eval eval_f1 Accuracy: {}\".format(eval_f1_dynamic))\n",
    "print(\"Latency: {:.3f} ms\".format(evalTime / samples * 1000))\n",
    "print(\"Throughput: {} samples/sec\".format(samples/evalTime))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Run Benchmark for FP32 model\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fda9f985",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "set_seed(training_args.seed)\n",
    "# Initialize our Trainer\n",
    "trainer_fp32 = QuestionAnsweringTrainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=train_dataset if training_args.do_train else None,\n",
    "    eval_dataset=eval_dataset if training_args.do_eval else None,\n",
    "    eval_examples=eval_examples if training_args.do_eval else None,\n",
    "    tokenizer=tokenizer,\n",
    "    data_collator=data_collator,\n",
    "    post_process_function=post_processing_function,\n",
    "    compute_metrics=compute_metrics,\n",
    ")\n",
    "\n",
    "start_time = timeit.default_timer()\n",
    "results = trainer_fp32.evaluate()\n",
    "evalTime = timeit.default_timer() - start_time\n",
    "max_eval_samples = data_args.max_eval_samples\n",
    "samples = eval_samples = min(max_eval_samples, len(eval_dataset))\n",
    "\n",
    "eval_f1_fp32 = results.get(\"eval_f1\")\n",
    "print('Batch size = {}'.format(training_args.per_device_eval_batch_size))\n",
    "print(\"Finally Eval eval_f1 Accuracy: {}\".format(eval_f1_fp32))\n",
    "print(\"Latency: {:.3f} ms\".format(evalTime / samples * 1000))\n",
    "print(\"Throughput: {} samples/sec\".format(samples/evalTime))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.6 64-bit",
   "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.6"
  },
  "vscode": {
   "interpreter": {
    "hash": "f54fd8d6160ddfbc370985ee3ad2925997e28943a671b1747496a6859c59cd26"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}