docs/tutorials/pytorch/summarization/pegasus-samsum.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 Pegasus model with 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": "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/intel-extension-for-transformers#installation) to install **intel-extension-for-transformers**. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef103e44",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# install model dependency\n",
    "!pip install accelerate datasets >= 1.8.0 sentencepiece != 0.1.92 protobuf rouge-score nltk py7zr torch >= 1.10 transformers>=4.19.0.dev0"
   ]
  },
  {
   "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 sys\n",
    "from dataclasses import dataclass, field\n",
    "\n",
    "import datasets\n",
    "import nltk  # Here to have a nice missing dependency error message early on\n",
    "import numpy as np\n",
    "import transformers\n",
    "from datasets import load_dataset, load_metric\n",
    "\n",
    "from filelock import FileLock\n",
    "from intel_extension_for_transformers.transformers import OptimizedModel\n",
    "from neural_compressor.config import PostTrainingQuantConfig\n",
    "from intel_extension_for_transformers.transformers import metrics as nlp_metrics\n",
    "from intel_extension_for_transformers.transformers.trainer import NLPSeq2SeqTrainer\n",
    "from transformers import (\n",
    "    AutoConfig,\n",
    "    AutoModelForSeq2SeqLM,\n",
    "    AutoTokenizer,\n",
    "    DataCollatorForSeq2Seq,\n",
    "    HfArgumentParser,\n",
    "    MBart50Tokenizer,\n",
    "    MBart50TokenizerFast,\n",
    "    MBartTokenizer,\n",
    "    MBartTokenizerFast,\n",
    "    Seq2SeqTrainingArguments,\n",
    "    set_seed,\n",
    ")\n",
    "from transformers.trainer_utils import get_last_checkpoint\n",
    "from transformers.utils import check_min_version, is_offline_mode\n",
    "from transformers.utils.versions import require_version\n",
    "from typing import Optional\n",
    "\n",
    "# Will error if the minimal version of Transformers is not installed. Remove at your own risks.\n",
    "check_min_version(\"4.19.0.dev0\")\n",
    "\n",
    "require_version(\"datasets>=1.17.0\", \"To fix: pip install -r examples/huggingface/pytorch/summarization/quantization/requirements.txt\")\n",
    "\n",
    "\n",
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "try:\n",
    "    nltk.data.find(\"tokenizers/punkt\")\n",
    "except (LookupError, OSError):\n",
    "    if is_offline_mode():\n",
    "        raise LookupError(\n",
    "            \"Offline mode: run this script without TRANSFORMERS_OFFLINE first to download nltk data files\"\n",
    "        )\n",
    "    with FileLock(\".lock\") as lock:\n",
    "        nltk.download(\"punkt\", quiet=True)\n",
    "\n",
    "# A list of all multilicngual tokenizer which require lang attribute.\n",
    "MULTILINGUAL_TOKENIZERS = [MBartTokenizer, MBartTokenizerFast, MBart50Tokenizer, MBart50TokenizerFast]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "048665d4",
   "metadata": {
    "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",
    "@dataclass\n",
    "class DataTrainingArguments:\n",
    "    \"\"\"\n",
    "    Arguments pertaining to what data we are going to input our model for training and eval.\n",
    "    \"\"\"\n",
    "    lang: str = field(default=None, metadata={\"help\": \"Language id for summarization.\"})\n",
    "    dataset_name: Optional[str] = field(\n",
    "        default=None, metadata={\"help\": \"The name of the dataset to use (via the datasets library).\"}\n",
    "    )\n",
    "    dataset_config_name: Optional[str] = field(\n",
    "        default=None, metadata={\"help\": \"The configuration name of the dataset to use (via the datasets library).\"}\n",
    "    )\n",
    "    overwrite_cache: bool = field(\n",
    "        default=False, metadata={\"help\": \"Overwrite the cached training and evaluation sets\"}\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",
    "    train_file: Optional[str] = field(\n",
    "        default=None, metadata={\"help\": \"The input training data file (a jsonlines or csv file).\"}\n",
    "    )\n",
    "    validation_file: Optional[str] = field(\n",
    "        default=None,\n",
    "        metadata={\n",
    "            \"help\": \"An optional input evaluation data file to evaluate the metrics (rouge) on \"\n",
    "                    \"(a jsonlines or csv file).\"\n",
    "        },\n",
    "    )\n",
    "    max_target_length: Optional[int] = field(\n",
    "        default=128,\n",
    "        metadata={\n",
    "            \"help\": \"The maximum total sequence length for target text after tokenization. Sequences longer \"\n",
    "                    \"than this will be truncated, sequences shorter will be padded.\"\n",
    "        },\n",
    "    )\n",
    "    val_max_target_length: Optional[int] = field(\n",
    "        default=None,\n",
    "        metadata={\n",
    "            \"help\": \"The maximum total sequence length for validation target text after tokenization. Sequences longer \"\n",
    "                    \"than this will be truncated, sequences shorter will be padded. Will default to `max_target_length`.\"\n",
    "                    \"This argument is also used to override the ``max_length`` param of ``model.generate``, which is used \"\n",
    "                    \"during ``evaluate`` and ``predict``.\"\n",
    "        },\n",
    "    )\n",
    "    ignore_pad_token_for_loss: bool = field(\n",
    "        default=True,\n",
    "        metadata={\n",
    "            \"help\": \"Whether to ignore the tokens corresponding to padded labels in the loss computation or not.\"\n",
    "        },\n",
    "    )\n",
    "    num_beams: Optional[int] = field(\n",
    "        default=None,\n",
    "        metadata={\n",
    "            \"help\": \"Number of beams to use for evaluation. This argument will be passed to ``model.generate``, \"\n",
    "                    \"which is used during ``evaluate`` and ``predict``.\"\n",
    "        },\n",
    "    )\n",
    "\n",
    "    def __post_init__(self):\n",
    "        if self.dataset_name is None and self.train_file is None and self.validation_file is None:\n",
    "            raise ValueError(\"Need either a dataset name or a training/validation file.\")\n",
    "        else:\n",
    "            if self.train_file is not None:\n",
    "                extension = self.train_file.split(\".\")[-1]\n",
    "                assert extension in [\"csv\", \"json\"], \"`train_file` should be a csv or a json file.\"\n",
    "            if self.validation_file is not None:\n",
    "                extension = self.validation_file.split(\".\")[-1]\n",
    "                assert extension in [\"csv\", \"json\"], \"`validation_file` should be a csv or a json file.\"\n",
    "        if self.val_max_target_length is None:\n",
    "            self.val_max_target_length = self.max_target_length\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",
    "    tune: bool = field(\n",
    "        default=False,\n",
    "        metadata={\"help\": \"Whether or not to apply quantization.\"},\n",
    "    )\n",
    "    quantization_approach: Optional[str] = field(\n",
    "        default=\"static\",\n",
    "        metadata={\"help\": \"Quantization approach. Supported approach are static, \"\n",
    "                          \"dynamic and qat.\"},\n",
    "    )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11b035ee",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "model_args = ModelArguments(\n",
    "    model_name_or_path=\"lvwerra/pegasus-samsum\",\n",
    ")\n",
    "data_args = DataTrainingArguments(\n",
    "    dataset_name=\"samsum\",\n",
    "    overwrite_cache=True,\n",
    "    max_train_samples=10000,\n",
    "    max_eval_samples=500\n",
    ")\n",
    "training_args = Seq2SeqTrainingArguments(\n",
    "    output_dir=\"./saved_results_dynamic\",\n",
    "    do_eval=True,\n",
    "    do_train=True,\n",
    "    no_cuda=True,\n",
    "    predict_with_generate=True,\n",
    "    overwrite_output_dir=True,\n",
    "    per_device_train_batch_size=8,\n",
    "    per_device_eval_batch_size=8,\n",
    ")\n",
    "optim_args = OptimizationArguments(\n",
    "    tune=True\n",
    ")\n",
    "summarization_name_mapping = {\n",
    "    \"amazon_reviews_multi\": (\"review_body\", \"review_title\"),\n",
    "    \"big_patent\": (\"description\", \"abstract\"),\n",
    "    \"cnn_dailymail\": (\"article\", \"highlights\"),\n",
    "    \"orange_sum\": (\"text\", \"summary\"),\n",
    "    \"pn_summary\": (\"article\", \"summary\"),\n",
    "    \"psc\": (\"extract_text\", \"summary_text\"),\n",
    "    \"samsum\": (\"dialogue\", \"summary\"),\n",
    "    \"thaisum\": (\"body\", \"summary\"),\n",
    "    \"xglue\": (\"news_body\", \"news_title\"),\n",
    "    \"xsum\": (\"document\", \"summary\"),\n",
    "    \"wiki_summary\": (\"article\", \"highlights\"),\n",
    "}"
   ]
  },
  {
   "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(data_args.dataset_name)"
   ]
  },
  {
   "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": [
    "# download model & vocab.\n",
    "config = AutoConfig.from_pretrained(model_args.model_name_or_path, revision=\"main\")\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path, use_fast=True, revision=\"main\")\n",
    "model = AutoModelForSeq2SeqLM.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",
    "    revision=\"main\"\n",
    ")\n",
    "model.resize_token_embeddings(len(tokenizer))\n",
    "\n",
    "if model.config.decoder_start_token_id is None and isinstance(tokenizer, (MBartTokenizer, MBartTokenizerFast)):\n",
    "    if isinstance(tokenizer, MBartTokenizer):\n",
    "        model.config.decoder_start_token_id = tokenizer.lang_code_to_id[data_args.lang]\n",
    "    else:\n",
    "        model.config.decoder_start_token_id = tokenizer.convert_tokens_to_ids(data_args.lang)\n",
    "\n",
    "if model.config.decoder_start_token_id is None:\n",
    "    raise ValueError(\"Make sure that `config.decoder_start_token_id` is correctly defined\")\n",
    "\n",
    "if (\n",
    "    hasattr(model.config, \"max_position_embeddings\")\n",
    "    and model.config.max_position_embeddings < 1024\n",
    "):\n",
    "    model.resize_position_embeddings(1024)\n",
    "\n",
    "prefix = \"\""
   ]
  },
  {
   "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",
    "# We need to tokenize inputs and targets.\n",
    "if training_args.do_train:\n",
    "    column_names = raw_datasets[\"train\"].column_names\n",
    "elif training_args.do_eval:\n",
    "    column_names = raw_datasets[\"validation\"].column_names\n",
    "elif training_args.do_predict:\n",
    "    column_names = raw_datasets[\"test\"].column_names\n",
    "else:\n",
    "    logger.info(\"There is nothing to do. Please pass `do_train`, `do_eval` and/or `do_predict`.\")\n",
    "\n",
    "if isinstance(tokenizer, tuple(MULTILINGUAL_TOKENIZERS)):\n",
    "    assert (\n",
    "        data_args.lang is not None\n",
    "    ), f\"{tokenizer.__class__.__name__} is a multilingual tokenizer which requires --lang argument\"\n",
    "\n",
    "    tokenizer.src_lang = data_args.lang\n",
    "    tokenizer.tgt_lang = data_args.lang\n",
    "    model.config.forced_bos_token_id = None\n",
    "\n",
    "\n",
    "# Get the column names for input/target.\n",
    "dataset_columns = summarization_name_mapping.get(data_args.dataset_name, None)\n",
    "text_column = dataset_columns[0] if dataset_columns is not None else column_names[0]\n",
    "summary_column = dataset_columns[1] if dataset_columns is not None else column_names[1]\n",
    "\n",
    "# Temporarily set max_target_length for training.\n",
    "max_target_length = 128\n",
    "padding = False\n",
    "\n",
    "def preprocess_function(examples):\n",
    "    # remove pairs where at least one record is None\n",
    "\n",
    "    inputs, targets = [], []\n",
    "    for i in range(len(examples[text_column])):\n",
    "        if examples[text_column][i] is not None and examples[summary_column][i] is not None:\n",
    "            inputs.append(examples[text_column][i])\n",
    "            targets.append(examples[summary_column][i])\n",
    "\n",
    "    inputs = [prefix + inp for inp in inputs]\n",
    "    model_inputs = tokenizer(inputs, max_length=1024, padding=padding, truncation=True)\n",
    "\n",
    "    # Setup the tokenizer for targets\n",
    "    with tokenizer.as_target_tokenizer():\n",
    "        labels = tokenizer(targets, max_length=max_target_length, padding=padding, truncation=True)\n",
    "\n",
    "    # If we are padding here, replace all tokenizer.pad_token_id in the labels by -100 when we want to ignore\n",
    "    # padding in the loss.\n",
    "    if padding == \"max_length\":\n",
    "        labels[\"input_ids\"] = [\n",
    "            [(l if l != tokenizer.pad_token_id else -100) for l in label] for label in labels[\"input_ids\"]\n",
    "        ]\n",
    "\n",
    "    model_inputs[\"labels\"] = labels[\"input_ids\"]\n",
    "    return model_inputs\n",
    "\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",
    "    if data_args.max_train_samples is not None:\n",
    "        max_train_samples = min(len(train_dataset), data_args.max_train_samples)\n",
    "        train_dataset = train_dataset.select(range(max_train_samples))\n",
    "    with training_args.main_process_first(desc=\"train dataset map pre-processing\"):\n",
    "        train_dataset = train_dataset.map(\n",
    "            preprocess_function,\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",
    "\n",
    "if training_args.do_eval:\n",
    "    max_target_length = data_args.val_max_target_length\n",
    "    if \"validation\" not in raw_datasets:\n",
    "        raise ValueError(\"--do_eval requires a validation dataset\")\n",
    "    eval_dataset = raw_datasets[\"validation\"]\n",
    "    if data_args.max_eval_samples is not None:\n",
    "        max_eval_samples = min(len(eval_dataset), data_args.max_eval_samples)\n",
    "        eval_dataset = eval_dataset.select(range(max_eval_samples))\n",
    "    with training_args.main_process_first(desc=\"validation dataset map pre-processing\"):\n",
    "        eval_dataset = eval_dataset.map(\n",
    "            preprocess_function,\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",
    "\n",
    "# Data collator\n",
    "label_pad_token_id = -100\n",
    "data_collator = DataCollatorForSeq2Seq(\n",
    "    tokenizer,\n",
    "    model=model,\n",
    "    label_pad_token_id=label_pad_token_id,\n",
    "    pad_to_multiple_of=8 if training_args.fp16 else None,\n",
    ")\n",
    "\n",
    "# Metric\n",
    "metric = load_metric(\"rouge\")\n",
    "\n",
    "def postprocess_text(preds, labels):\n",
    "    preds = [pred.strip() for pred in preds]\n",
    "    labels = [label.strip() for label in labels]\n",
    "\n",
    "    # rougeLSum expects newline after each sentence\n",
    "    preds = [\"\\n\".join(nltk.sent_tokenize(pred)) for pred in preds]\n",
    "    labels = [\"\\n\".join(nltk.sent_tokenize(label)) for label in labels]\n",
    "\n",
    "    return preds, labels\n",
    "\n",
    "def compute_metrics(eval_preds):\n",
    "    preds, labels = eval_preds\n",
    "    if isinstance(preds, tuple):\n",
    "        preds = preds[0]\n",
    "    decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)\n",
    "    if data_args.ignore_pad_token_for_loss:\n",
    "        # Replace -100 in the labels as we can't decode them.\n",
    "        labels = np.where(labels != -100, labels, tokenizer.pad_token_id)\n",
    "    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)\n",
    "\n",
    "    # Some simple post-processing\n",
    "    decoded_preds, decoded_labels = postprocess_text(decoded_preds, decoded_labels)\n",
    "\n",
    "    result = metric.compute(predictions=decoded_preds, references=decoded_labels, use_stemmer=True)\n",
    "    # Extract a few results from ROUGE\n",
    "    result = {key: value.mid.fmeasure * 100 for key, value in result.items()}\n",
    "\n",
    "    prediction_lens = [np.count_nonzero(pred != tokenizer.pad_token_id) for pred in preds]\n",
    "    result[\"gen_len\"] = np.mean(prediction_lens)\n",
    "    result = {k: round(v, 4) for k, v in result.items()}\n",
    "    return result\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": [
    "## Dynamic 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 = NLPSeq2SeqTrainer(\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",
    "    tokenizer=tokenizer,\n",
    "    data_collator=data_collator,\n",
    "    compute_metrics=compute_metrics if training_args.predict_with_generate else None,\n",
    ")\n",
    "\n",
    "results = {}\n",
    "max_length = (\n",
    "    training_args.generation_max_length\n",
    "    if training_args.generation_max_length is not None\n",
    "    else data_args.val_max_target_length\n",
    ")\n",
    "num_beams = data_args.num_beams if data_args.num_beams is not None else training_args.generation_num_beams\n",
    "\n",
    "metric_name = \"eval_rougeLsum\"\n",
    "\n",
    "# tuning\n",
    "model.config.save_pretrained(\"./saved_results_dynamic\")\n",
    "trainer.save_model(\"./saved_results_dynamic\")\n",
    "\n",
    "tune_metric = nlp_metrics.Metric(\n",
    "    name=metric_name, is_relative=True, criterion=0.25\n",
    ")\n",
    "trainer.metrics = tune_metric\n",
    "quantization_config = PostTrainingQuantConfig(\n",
    "    approach=\"dynamic\",\n",
    ")\n",
    "trainer.max_length = max_length\n",
    "trainer.num_beams = num_beams\n",
    "trainer.quantize(quant_config=quantization_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a7e93de",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Run Benchmark after Dynamic Post Training Quantization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf54c034",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "results = trainer.evaluate(max_length=max_length, num_beams=num_beams)\n",
    "logger.info(\"metrics keys: {}\".format(results.keys()))\n",
    "throughput = results.get(\"eval_samples_per_second\")\n",
    "print('Batch size = {:d}'.format(training_args.per_device_eval_batch_size))\n",
    "print(\"Finally Eval {} Accuracy: {:.5f}\".format(metric_name, results[metric_name]))\n",
    "print(\"Latency: {:.5f} ms\".format(1000 / throughput))\n",
    "print(\"Throughput: {:.5f} samples/sec\".format(throughput))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f35818b8",
   "metadata": {},
   "source": [
    "## Run Benchmark after Dynamic Post Training Quantization with Multi-Instance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bac69198",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "results = os.system('bash ../multi_instance.sh --model=saved_results_dynamic --core_per_instance=4 --data_type=int8')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44cca2a1",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Run Benchmark for FP32 model"
   ]
  },
  {
   "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 = NLPSeq2SeqTrainer(\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",
    "    tokenizer=tokenizer,\n",
    "    data_collator=data_collator,\n",
    "    compute_metrics=compute_metrics if training_args.predict_with_generate else None,\n",
    ")\n",
    "results_fp32 = trainer_fp32.evaluate(max_length=max_length, num_beams=num_beams)\n",
    "throughput_fp32 = results_fp32.get(\"eval_samples_per_second\")\n",
    "print('Batch size = {:d}'.format(training_args.per_device_eval_batch_size))\n",
    "print(\"Finally Eval {} Accuracy: {:.5f}\".format(metric_name, results[metric_name]))\n",
    "print(\"Latency: {:.5f} ms\".format(1000 / throughput_fp32))\n",
    "print(\"Throughput: {:.5f} samples/sec\".format(throughput_fp32))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05df025f",
   "metadata": {},
   "source": [
    "## Run Benchmark for FP32 Model with Multi-Instance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71e4e7da",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "results = os.system('bash ../multi_instance.sh --model=textattack/bert-base-uncased-MRPC --core_per_instance=4 --data_type=fp32')"
   ]
  }
 ],
 "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
}