docs/tutorials/pytorch/question-answering/distillation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This tutorial demonstrates how to use the distillation approach based on [Intel® Neural Compressor](https://github.com/intel/neural-compressor) for question-answering."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prerequisite"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install packages"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* Follow [installation](https://github.com/intel/intel-extension-for-transformers#installation) to install **intel-extension-for-transformers**. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# install model dependency\n",
    "! pip install accelerate datasets >= 1.1.3 sentencepiece != 0.1.92 protobuf torch >= 1.10 transformers >= 4.12.0 wandb\n",
    "! pip install -r requirements.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import os\n",
    "import sys\n",
    "from dataclasses import dataclass, field\n",
    "from typing import Optional\n",
    "\n",
    "import datasets\n",
    "from datasets import load_dataset, load_metric\n",
    "\n",
    "import functools\n",
    "import numpy as np\n",
    "import time\n",
    "import torch\n",
    "import transformers\n",
    "from intel_extension_for_transformers.transformers import metrics, OptimizedModel, DistillationConfig\n",
    "from torch.utils.data import DataLoader\n",
    "from tqdm import tqdm\n",
    "from trainer_qa import QuestionAnsweringTrainer\n",
    "from transformers import (\n",
    "    AutoConfig,\n",
    "    AutoModelForQuestionAnswering,\n",
    "    AutoTokenizer,\n",
    "    DataCollatorWithPadding,\n",
    "    EvalPrediction,\n",
    "    HfArgumentParser,\n",
    "    PreTrainedTokenizerFast,\n",
    "    TrainingArguments,\n",
    "    default_data_collator,\n",
    "    set_seed,\n",
    ")\n",
    "from transformers.trainer_utils import get_last_checkpoint\n",
    "from transformers.utils import check_min_version\n",
    "from transformers.utils.versions import require_version\n",
    "from typing import Optional\n",
    "from utils_qa import postprocess_qa_predictions\n",
    "\n",
    "\n",
    "# Will error if the minimal version of Transformers is not installed. Remove at your own risks.\n",
    "check_min_version(\"4.12.0\")\n",
    "\n",
    "require_version(\"datasets>=1.8.0\", \"To fix: pip install -r examples/huggingface/pytorch/question-answering/distillation/requirements.txt\")\n",
    "\n",
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "os.environ[\"WANDB_DISABLED\"] = \"true\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define arguments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ========== Define arguments =========\n",
    "@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",
    "@dataclass\n",
    "class OptimizationArguments:\n",
    "    \"\"\"\n",
    "    Arguments pertaining to what type of optimization we are going to apply on the model.\n",
    "    \"\"\"\n",
    "\n",
    "    distillation: bool = field(\n",
    "        default=False,\n",
    "        metadata={\"help\": \"Whether or not to apply distillation.\"},\n",
    "    )\n",
    "    teacher_model_name_or_path: str = field(\n",
    "        default=False,\n",
    "        metadata={\"help\": \"Path to pretrained model or model identifier from huggingface.co/models\"},\n",
    "    )\n",
    "    run_teacher_logits: bool = field(\n",
    "        default=False,\n",
    "        metadata={\"help\": \"Whether or not to obtain teacher model's logits on train dataset before training.\"},\n",
    "        )\n",
    "    metric_name: Optional[str] = field(\n",
    "        default=\"eval_f1\",\n",
    "        metadata={\"help\": \"Metric used for the tuning strategy.\"},\n",
    "    )\n",
    "    tolerance_mode: Optional[str] = field(\n",
    "        default=\"absolute\",\n",
    "        metadata={\"help\": \"Metric tolerance model, expected to be relative or absolute.\"},\n",
    "    )\n",
    "    perf_tol: Optional[float] = field(\n",
    "        default=0.02,\n",
    "        metadata={\"help\": \"Performance tolerance when optimizing the model.\"},\n",
    "    )\n",
    "    benchmark: bool = field(\n",
    "        default=False,\n",
    "        metadata={\"help\": \"run benchmark.\"}\n",
    "    )\n",
    "    accuracy_only: bool = field(\n",
    "        default=False,\n",
    "        metadata={\"help\":\"Whether to only test accuracy for model tuned by Neural Compressor.\"}\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We now keep distinct sets of args, for a cleaner separation of concerns.\n",
    "parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments, OptimizationArguments))\n",
    "if len(sys.argv) == 2 and sys.argv[1].endswith(\".json\"):\n",
    "    # If we pass only one argument to the script and it's the path to a json file,\n",
    "    # let's parse it to get our arguments.\n",
    "    model_args, data_args, training_args, optim_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))\n",
    "else:\n",
    "    model_args, data_args, training_args, optim_args = parser.parse_args_into_dataclasses()\n",
    "\n",
    "# Setup logging\n",
    "logging.basicConfig(\n",
    "    format=\"%(asctime)s - %(levelname)s - %(name)s - %(message)s\",\n",
    "    datefmt=\"%m/%d/%Y %H:%M:%S\",\n",
    "    handlers=[logging.StreamHandler(sys.stdout)],\n",
    ")\n",
    "\n",
    "log_level = training_args.get_process_log_level()\n",
    "logger.setLevel(log_level)\n",
    "datasets.utils.logging.set_verbosity(log_level)\n",
    "transformers.utils.logging.set_verbosity(log_level)\n",
    "transformers.utils.logging.enable_default_handler()\n",
    "transformers.utils.logging.enable_explicit_format()\n",
    "\n",
    "# Log on each process the small summary:\n",
    "logger.warning(\n",
    "    f\"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}\"\n",
    "    + f\"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}\"\n",
    ")\n",
    "logger.info(f\"Training/evaluation parameters {training_args}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download dataset from the hub"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "raw_datasets = load_dataset(\n",
    "    data_args.dataset_name, data_args.dataset_config_name, cache_dir=model_args.cache_dir\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download fp32 model from the hub"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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=True if model_args.use_auth_token else None,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preprocessing the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Preprocessing the datasets.\n",
    "# Preprocessing is slightly 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, tokenizer=tokenizer):\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",
    "    if data_args.max_eval_samples is not None:\n",
    "        # We will select sample from whole data\n",
    "        eval_examples = eval_examples.select(range(data_args.max_eval_samples))\n",
    "    # Validation Feature Creation\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",
    "            num_proc=data_args.preprocessing_num_workers,\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",
    "        eval_dataset = eval_dataset.select(range(data_args.max_eval_samples))\n",
    "\n",
    "# Data collator\n",
    "data_collator = (\n",
    "    DataCollatorWithPadding(tokenizer, pad_to_multiple_of=8 if training_args.fp16 else None)\n",
    ")\n",
    "\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)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Distillation & Benchmark"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model Distillation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class QAModel_output_reshaped(torch.nn.Module):\n",
    "    def __init__(self, model):\n",
    "        super(QAModel_output_reshaped, self).__init__()\n",
    "        self.model = model\n",
    "\n",
    "    def forward(self, *args, **kwargs):\n",
    "        outputs = self.model(*args, **kwargs)\n",
    "        outputs_reshaped = torch.vstack([torch.vstack([sx, ex]) \\\n",
    "                for sx, ex in zip(outputs['start_logits'], outputs['end_logits'])])\n",
    "        return outputs_reshaped\n",
    "\n",
    "teacher_config = AutoConfig.from_pretrained(\n",
    "    optim_args.teacher_model_name_or_path,\n",
    "    use_auth_token=True if model_args.use_auth_token else None,\n",
    ")\n",
    "teacher_tokenizer = AutoTokenizer.from_pretrained(\n",
    "    optim_args.teacher_model_name_or_path,\n",
    "    use_fast=True,\n",
    "    use_auth_token=True if model_args.use_auth_token else None,\n",
    ")\n",
    "teacher_model = AutoModelForQuestionAnswering.from_pretrained(\n",
    "    optim_args.teacher_model_name_or_path,\n",
    "    from_tf=bool(\".ckpt\" in model_args.model_name_or_path),\n",
    "    config=teacher_config,\n",
    "    use_auth_token=True if model_args.use_auth_token else None,\n",
    ")\n",
    "teacher_model.to(training_args.device)\n",
    "\n",
    "# Prepare datasets for teacher model\n",
    "# Create train feature from dataset\n",
    "with training_args.main_process_first(desc=\"train dataset map pre-processing\"):\n",
    "    teacher_train_dataset = train_examples.map(\n",
    "        functools.partial(prepare_train_features, tokenizer=teacher_tokenizer),\n",
    "        batched=True,\n",
    "        num_proc=data_args.preprocessing_num_workers,\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",
    "    teacher_train_dataset = teacher_train_dataset.select(range(data_args.max_train_samples))\n",
    "\n",
    "# Validation Feature Creation\n",
    "with training_args.main_process_first(desc=\"validation dataset map pre-processing\"):\n",
    "    teacher_eval_dataset = eval_examples.map(\n",
    "        functools.partial(prepare_validation_features, tokenizer=teacher_tokenizer),\n",
    "        batched=True,\n",
    "        num_proc=data_args.preprocessing_num_workers,\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",
    "    teacher_eval_dataset = teacher_eval_dataset.select(range(data_args.max_eval_samples))\n",
    "    \n",
    "# get logits of teacher model\n",
    "if optim_args.run_teacher_logits:\n",
    "    def dict_tensor_to_model_device(batch, model):\n",
    "        device = next(model.parameters()).device\n",
    "        for k in batch:\n",
    "            batch[k] = batch[k].to(device)\n",
    "\n",
    "    def get_logits(teacher_model, train_dataset, teacher_train_dataset):\n",
    "        logger.info(\"***** Getting logits of teacher model *****\")\n",
    "        logger.info(f\"  Num examples = {len(train_dataset) }\")\n",
    "        teacher_model.eval()\n",
    "        npy_file = os.path.join(os.path.dirname(os.path.abspath(__file__)),\n",
    "            '{}.{}.npy'.format(data_args.dataset_name, \n",
    "                            optim_args.teacher_model_name_or_path.replace('/', '.')))\n",
    "        if os.path.exists(npy_file):\n",
    "            teacher_logits = [list(x) for x in np.load(npy_file, allow_pickle=True)]\n",
    "        else:\n",
    "            sampler = None\n",
    "            if training_args.world_size > 1:\n",
    "                from transformers.trainer_pt_utils import ShardSampler\n",
    "                sampler = ShardSampler(\n",
    "                    teacher_train_dataset,\n",
    "                    batch_size=training_args.per_device_eval_batch_size,\n",
    "                    num_processes=training_args.world_size,\n",
    "                    process_index=training_args.process_index,\n",
    "                )\n",
    "                teacher_model = torch.nn.parallel.DistributedDataParallel(\n",
    "                    teacher_model,\n",
    "                    device_ids=[training_args.local_rank] \\\n",
    "                        if training_args._n_gpu != 0 else None,\n",
    "                    output_device=training_args.local_rank \\\n",
    "                        if training_args._n_gpu != 0 else None,\n",
    "                )\n",
    "            train_dataloader = DataLoader(teacher_train_dataset, \n",
    "                                        collate_fn=data_collator, \n",
    "                                        sampler=sampler,\n",
    "                                        batch_size=training_args.per_device_eval_batch_size)\n",
    "            train_dataloader = tqdm(train_dataloader, desc=\"Evaluating\")\n",
    "            teacher_logits = []\n",
    "            for step, batch in enumerate(train_dataloader):\n",
    "                dict_tensor_to_model_device(batch, teacher_model)\n",
    "                outputs = teacher_model(**batch).cpu().numpy()\n",
    "                if training_args.world_size > 1:\n",
    "                    outputs_list = [None for i in range(training_args.world_size)]\n",
    "                    torch.distributed.all_gather_object(outputs_list, outputs)\n",
    "                    outputs = np.concatenate(outputs_list, axis=0)\n",
    "                teacher_logits += [[s,e] for s,e in zip(outputs[0::2], outputs[1::2])]\n",
    "            if training_args.world_size > 1:\n",
    "                teacher_logits = teacher_logits[:len(teacher_train_dataset)]\n",
    "            if training_args.local_rank in [-1, 0]:\n",
    "                np.save(npy_file, teacher_logits, allow_pickle=True)\n",
    "        return train_dataset.add_column('teacher_logits', teacher_logits[:data_args.max_train_samples])\n",
    "    with torch.no_grad():\n",
    "        train_dataset = get_logits(QAModel_output_reshaped(teacher_model), train_dataset, teacher_train_dataset)\n",
    "        \n",
    "para_counter = lambda model:sum(p.numel() for p in model.parameters())\n",
    "logger.info(\"***** Number of teacher model parameters: {:.2f}M *****\".format(\\\n",
    "            para_counter(teacher_model)/10**6))\n",
    "logger.info(\"***** Number of student model parameters: {:.2f}M *****\".format(\\\n",
    "            para_counter(model)/10**6))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "set_seed(training_args.seed)\n",
    "# Initialize our Trainer\n",
    "trainer = 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",
    "tune_metric = metrics.Metric(name=optim_args.metric_name)\n",
    "distillation_conf = DistillationConfig(metrics=tune_metric)\n",
    "model = trainer.distill(\n",
    "    distillation_config=distillation_conf, teacher_model=teacher_model\n",
    ")\n",
    "trainer.save_model(training_args.output_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run Benchmark after Distillation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the model obtained after Intel Neural Compressor (INC) quantization\n",
    "model = OptimizedModel.from_pretrained(\n",
    "    training_args.output_dir,\n",
    ")\n",
    "model.eval()\n",
    "trainer.model = model\n",
    "start_time = timeit.default_timer()\n",
    "results = trainer.evaluate()\n",
    "evalTime = timeit.default_timer() - start_time\n",
    "max_eval_samples = data_args.max_eval_samples \\\n",
    "    if data_args.max_eval_samples is not None else len(eval_dataset)\n",
    "eval_samples = min(max_eval_samples, len(eval_dataset))\n",
    "samples = eval_samples - (eval_samples % batch_size) \\\n",
    "    if training_args.dataloader_drop_last else eval_samples\n",
    "logger.info(\"metrics keys: {}\".format(results.keys()))\n",
    "bert_task_acc_keys = ['eval_f1', 'eval_accuracy', 'eval_matthews_correlation',\n",
    "                        'eval_pearson', 'eval_mcc', 'eval_spearmanr']\n",
    "ret = False\n",
    "for key in bert_task_acc_keys:\n",
    "    if key in results.keys():\n",
    "        ret = True\n",
    "        print('Batch size = ', training_args.per_device_eval_batch_size)\n",
    "        print(\"Finally Eval {} Accuracy: {}\".format(key, results[key]))\n",
    "        print(\"Latency: {:.5f} ms\".format(evalTime / samples * 1000))\n",
    "        print(\"Throughput: {:.5f} samples/sec\".format(samples/evalTime))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.6 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.8.6"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "f54fd8d6160ddfbc370985ee3ad2925997e28943a671b1747496a6859c59cd26"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}