docs/tutorials/pytorch/text-classification/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 text-classification."
   ]
  },
  {
   "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 datasets\n",
    "import functools\n",
    "import logging\n",
    "import os\n",
    "import numpy as np\n",
    "import random\n",
    "import sys\n",
    "import torch\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 (\n",
    "    metrics,\n",
    "    DistillationConfig,\n",
    "    OptimizedModel,\n",
    ")\n",
    "from intel_extension_for_transformers.transformers.trainer import NLPTrainer\n",
    "from torch.utils.data import DataLoader\n",
    "from tqdm.auto import tqdm\n",
    "from transformers import (\n",
    "    AutoConfig,\n",
    "    AutoModelForSequenceClassification,\n",
    "    AutoTokenizer,\n",
    "    DataCollatorWithPadding,\n",
    "    EvalPrediction,\n",
    "    HfArgumentParser,\n",
    "    PretrainedConfig,\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 typing import Optional\n",
    "\n",
    "\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"\"\n",
    "os.environ[\"WANDB_DISABLED\"] = \"true\"\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",
    "\n",
    "task_to_keys = {\n",
    "    \"cola\": (\"sentence\", None),\n",
    "    \"mnli\": (\"premise\", \"hypothesis\"),\n",
    "    \"mrpc\": (\"sentence1\", \"sentence2\"),\n",
    "    \"qnli\": (\"question\", \"sentence\"),\n",
    "    \"qqp\": (\"question1\", \"question2\"),\n",
    "    \"rte\": (\"sentence1\", \"sentence2\"),\n",
    "    \"sst2\": (\"sentence\", None),\n",
    "    \"stsb\": (\"sentence1\", \"sentence2\"),\n",
    "    \"wnli\": (\"sentence1\", \"sentence2\"),\n",
    "}\n",
    "\n",
    "logger = logging.getLogger(__name__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define arguments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ========== Define arguments =========\n",
    "@dataclass\n",
    "class DataTrainingArguments:\n",
    "    \"\"\"\n",
    "    Arguments pertaining to what data we are going to input our model for training and eval.\n",
    "    Using `HfArgumentParser` we can turn this class\n",
    "    into argparse arguments to be able to specify them on\n",
    "    the command line.\n",
    "    \"\"\"\n",
    "\n",
    "    task_name: Optional[str] = field(\n",
    "        default=None,\n",
    "        metadata={\"help\": \"The name of the task to train on: \" + \", \".join(task_to_keys.keys())},\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",
    "    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",
    "    max_seq_length: int = field(\n",
    "        default=128,\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",
    "    overwrite_cache: bool = field(\n",
    "        default=False, metadata={\"help\": \"Overwrite the cached preprocessed datasets or not.\"}\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.\"\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",
    "    max_predict_samples: Optional[int] = field(\n",
    "        default=None,\n",
    "        metadata={\n",
    "            \"help\": \"For debugging purposes or quicker training, truncate the number of prediction examples to this \"\n",
    "            \"value if set.\"\n",
    "        },\n",
    "    )\n",
    "    train_file: Optional[str] = field(\n",
    "        default=None, metadata={\"help\": \"A csv or a json file containing the training data.\"}\n",
    "    )\n",
    "    validation_file: Optional[str] = field(\n",
    "        default=None, metadata={\"help\": \"A csv or a json file containing the validation data.\"}\n",
    "    )\n",
    "\n",
    "    def __post_init__(self):\n",
    "        if self.task_name is not None:\n",
    "            self.task_name = self.task_name.lower()\n",
    "            if self.task_name not in task_to_keys.keys():\n",
    "                raise ValueError(\"Unknown task, you should pick one in \" + \",\".join(task_to_keys.keys()))\n",
    "        elif self.dataset_name is not None:\n",
    "            pass\n",
    "        elif self.train_file is None or self.validation_file is None:\n",
    "            raise ValueError(\"Need either a GLUE task, a training/validation file or a dataset name.\")\n",
    "        else:\n",
    "            train_extension = self.train_file.split(\".\")[-1]\n",
    "            assert train_extension in [\"csv\", \"json\"], \"`train_file` should be a csv or a json file.\"\n",
    "            validation_extension = self.validation_file.split(\".\")[-1]\n",
    "            assert (\n",
    "                validation_extension == train_extension\n",
    "            ), \"`validation_file` should have the same extension (csv or json) as `train_file`.\"\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class ModelArguments:\n",
    "    \"\"\"\n",
    "    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.\n",
    "    \"\"\"\n",
    "\n",
    "    model_name_or_path: str = field(\n",
    "        metadata={\"help\": \"Path to pretrained model or model identifier from huggingface.co/models\"}\n",
    "    )\n",
    "    config_name: Optional[str] = field(\n",
    "        default=None, metadata={\"help\": \"Pretrained config name or path if not the same as model_name\"}\n",
    "    )\n",
    "    tokenizer_name: Optional[str] = field(\n",
    "        default=None, metadata={\"help\": \"Pretrained tokenizer name or path if not the same as model_name\"}\n",
    "    )\n",
    "    cache_dir: Optional[str] = field(\n",
    "        default=None,\n",
    "        metadata={\"help\": \"Where do you want to store the pretrained models downloaded from huggingface.co\"},\n",
    "    )\n",
    "    use_fast_tokenizer: bool = field(\n",
    "        default=True,\n",
    "        metadata={\"help\": \"Whether to use one of the fast tokenizer (backed by the tokenizers library) or not.\"},\n",
    "    )\n",
    "    model_revision: str = field(\n",
    "        default=\"main\",\n",
    "        metadata={\"help\": \"The specific model version to use (can be a branch name, tag name or commit id).\"},\n",
    "    )\n",
    "    use_auth_token: bool = field(\n",
    "        default=False,\n",
    "        metadata={\n",
    "            \"help\": \"Will use the token generated when running `transformers-cli login` (necessary to use this script \"\n",
    "            \"with private models).\"\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",
    "    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",
    "    metric_name: Optional[str] = field(\n",
    "        default=None,\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",
    "    accuracy_only: bool = field(\n",
    "        default=False,\n",
    "        metadata={\"help\":\"Whether to only test accuracy for model tuned by Neural Compressor.\"})"
   ]
  },
  {
   "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(\n",
    "        json_file=os.path.abspath(sys.argv[1])\n",
    "    )\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": [
    "# download the dataset.\n",
    "raw_datasets = load_dataset(\"glue\", data_args.task_name)\n",
    "# Labels\n",
    "label_list = raw_datasets[\"train\"].features[\"label\"].names\n",
    "num_labels = len(label_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download fp32 model from the hub"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load pretrained model and tokenizer\n",
    "config = AutoConfig.from_pretrained(\n",
    "    model_args.model_name_or_path,\n",
    "    num_labels=num_labels,\n",
    "    finetuning_task=data_args.task_name,\n",
    "    revision=\"main\"\n",
    ")\n",
    "tokenizer = AutoTokenizer.from_pretrained(\n",
    "    model_args.model_name_or_path,\n",
    "    use_fast=True,\n",
    "    revision=\"main\"\n",
    ")\n",
    "model = AutoModelForSequenceClassification.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",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preprocessing the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Preprocessing the raw_datasets\n",
    "sentence1_key, sentence2_key = task_to_keys[data_args.task_name]\n",
    "# Padding strategy\n",
    "padding = False\n",
    "# Some models have set the order of the labels to use, so let's make sure we do use it.\n",
    "label_to_id = None\n",
    "if model.config.label2id != PretrainedConfig(num_labels=num_labels).label2id:\n",
    "    # Some have all caps in their config, some don't.\n",
    "    label_name_to_id = {k.lower(): v for k, v in model.config.label2id.items()}\n",
    "    if list(sorted(label_name_to_id.keys())) == list(sorted(label_list)):\n",
    "        label_to_id = {i: int(label_name_to_id[label_list[i]]) for i in range(num_labels)}\n",
    "    else:\n",
    "        logger.warning(\n",
    "            f\"Your model seems to have been trained with labels, but they don't match the dataset: \"\n",
    "            f\"model labels: {list(sorted(label_name_to_id.keys()))}, dataset labels: {list(sorted(label_list))}.\\n\"\n",
    "            f\"Ignoring the model labels as a result.\"\n",
    "        )\n",
    "if label_to_id is not None:\n",
    "    model.config.label2id = label_to_id\n",
    "    model.config.id2label = {id: label for label, id in config.label2id.items()}\n",
    "max_seq_length = min(data_args.max_seq_length, tokenizer.model_max_length)\n",
    "\n",
    "def preprocess_function(examples, tokenizer=tokenizer):\n",
    "    # Tokenize the texts\n",
    "    args = (\n",
    "        (examples[sentence1_key],) if sentence2_key is None else (examples[sentence1_key], examples[sentence2_key])\n",
    "    )\n",
    "    result = tokenizer(*args, padding=padding, max_length=max_seq_length, truncation=True)\n",
    "\n",
    "    # Map labels to IDs (not necessary for GLUE tasks)\n",
    "    if label_to_id is not None and \"label\" in examples:\n",
    "        result[\"label\"] = [(label_to_id[l] if l != -1 else -1) for l in examples[\"label\"]]\n",
    "    return result\n",
    "\n",
    "with training_args.main_process_first(desc=\"dataset map pre-processing\"):\n",
    "    raw_datasets = raw_datasets.map(\n",
    "        preprocess_function, batched=True, load_from_cache_file=not data_args.overwrite_cache\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",
    "\n",
    "if training_args.do_eval:\n",
    "    if \"validation\" not in raw_datasets and \"validation_matched\" not in raw_datasets:\n",
    "        raise ValueError(\"--do_eval requires a validation dataset\")\n",
    "    eval_dataset = raw_datasets[\"validation_matched\" if data_args.task_name == \"mnli\" else \"validation\"]\n",
    "\n",
    "# Log a few random samples from the training set:\n",
    "if training_args.do_train:\n",
    "    for index in random.sample(range(len(train_dataset)), 3):\n",
    "        logger.info(f\"Sample {index} of the training set: {train_dataset[index]}.\")\n",
    "\n",
    "# Get the metric function\n",
    "metric = load_metric(\"glue\", data_args.task_name)\n",
    "\n",
    "metric_name = \"eval_accuracy\"\n",
    "\n",
    "# You can define your custom compute_metrics function. It takes an `EvalPrediction` object (a namedtuple with\n",
    "# predictions and label_ids field) and has to return a dictionary string to float.\n",
    "def compute_metrics(p: EvalPrediction):\n",
    "    preds = p.predictions[0] if isinstance(p.predictions, tuple) else p.predictions\n",
    "    preds =  np.argmax(preds, axis=1)\n",
    "    if data_args.task_name is not None:\n",
    "        result = metric.compute(predictions=preds, references=p.label_ids)\n",
    "        if len(result) > 1:\n",
    "            result[\"combined_score\"] = np.mean(list(result.values())).item()\n",
    "        return result\n",
    "    else:\n",
    "        return {\"accuracy\": (preds == p.label_ids).astype(np.float32).mean().item()}\n",
    "\n",
    "# Data collator will default to DataCollatorWithPadding, so we change it if we already did the padding.\n",
    "data_collator = None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Distillation & Benchmark"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model Distillation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BertModelforLogitsOutputOnly(torch.nn.Module):\n",
    "    def __init__(self, model):\n",
    "        super(BertModelforLogitsOutputOnly, self).__init__()\n",
    "        self.model = model\n",
    "    def forward(self, *args, **kwargs):\n",
    "        output = self.model(*args, **kwargs)\n",
    "        return output['logits']\n",
    "\n",
    "teacher_config = AutoConfig.from_pretrained(optim_args.teacher_model_name_or_path, \\\n",
    "                    num_labels=num_labels, finetuning_task=data_args.task_name)\n",
    "teacher_tokenizer = AutoTokenizer.from_pretrained(optim_args.teacher_model_name_or_path, \\\n",
    "                    use_fast=model_args.use_fast_tokenizer)\n",
    "teacher_model = AutoModelForSequenceClassification.from_pretrained(\n",
    "    optim_args.teacher_model_name_or_path,\n",
    "    from_tf=bool(\".ckpt\" in optim_args.teacher_model_name_or_path),\n",
    "    config=teacher_config,\n",
    ")\n",
    "teacher_model.to(training_args.device)\n",
    "\n",
    "# prepare datasets for teacher model\n",
    "teacher_processed_datasets = raw_datasets.map(\n",
    "    functools.partial(preprocess_function, tokenizer=teacher_tokenizer), \n",
    "    batched=True, remove_columns=raw_datasets[\"train\"].column_names\n",
    ")\n",
    "teacher_train_dataset = teacher_processed_datasets[\"train\"]\n",
    "teacher_train_dataset = teacher_train_dataset.select(range(data_args.max_train_samples))\n",
    "teacher_eval_dataset = teacher_processed_datasets[\"validation_matched\" \\\n",
    "                            if data_args.task_name == \"mnli\" else \"validation\"]\n",
    "teacher_eval_dataset = teacher_eval_dataset.select(range(data_args.max_eval_samples))\n",
    "    \n",
    "# get logits of teacher model\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.task_name, \n",
    "                            optim_args.teacher_model_name_or_path.replace('/', '.')))\n",
    "    if os.path.exists(npy_file):\n",
    "        teacher_logits = [x for x in np.load(npy_file)]\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)\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 = torch.concat(outputs_list, dim=0)\n",
    "            teacher_logits += [x for x in outputs.cpu().numpy()]\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, np.array(teacher_logits))\n",
    "    return train_dataset.add_column('teacher_logits', teacher_logits)\n",
    "\n",
    "with torch.no_grad():\n",
    "    train_dataset = get_logits(BertModelforLogitsOutputOnly(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 before initializing model.\n",
    "set_seed(training_args.seed)\n",
    "# Initialize our Trainer\n",
    "trainer = NLPTrainer(\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",
    "    compute_metrics=compute_metrics,\n",
    "    tokenizer=tokenizer,\n",
    "    data_collator=data_collator,\n",
    ")\n",
    "# distillation\n",
    "if not training_args.do_eval:\n",
    "    raise ValueError(\"do_eval must be set to True for distillation.\")\n",
    "\n",
    "tune_metric = metrics.Metric(name=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",
    "results = trainer.evaluate()\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",
    "\n",
    "for key in bert_task_acc_keys:\n",
    "    if key in results.keys():\n",
    "        ret = True\n",
    "        throughput = results.get(\"eval_samples_per_second\")\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(1000 / throughput))\n",
    "        print(\"Throughput: {:.5f} samples/sec\".format(throughput))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.6 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10.6"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "a3ed54c68abdb79eabea0140062ffa976ea4d8132b937aa83ca919a8d862edf2"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}