docs/tutorials/pytorch/translation/t5-small.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 T5 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 sentencepiece != 0.1.92 protobuf sacrebleu >= 1.4.12 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",
    "from typing import Optional\n",
    "\n",
    "import datasets\n",
    "import numpy as np\n",
    "from datasets import load_dataset, load_metric\n",
    "\n",
    "from intel_extension_for_transformers.transformers import OptimizedModel, QuantizationConfig\n",
    "from intel_extension_for_transformers.transformers import metrics as nlp_metrics\n",
    "from intel_extension_for_transformers.transformers.trainer import NLPSeq2SeqTrainer\n",
    "import transformers\n",
    "from transformers import (\n",
    "    AutoConfig,\n",
    "    AutoModelForSeq2SeqLM,\n",
    "    AutoTokenizer,\n",
    "    DataCollatorForSeq2Seq,\n",
    "    HfArgumentParser,\n",
    "    M2M100Tokenizer,\n",
    "    MBart50Tokenizer,\n",
    "    MBart50TokenizerFast,\n",
    "    MBartTokenizer,\n",
    "    MBartTokenizerFast,\n",
    "    Seq2SeqTrainer,\n",
    "    Seq2SeqTrainingArguments,\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",
    "\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.8.0\", \"To fix: pip install -r examples/huggingface/pytorch/translation/quantization/requirements.txt\")\n",
    "\n",
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "# A list of all multilingual tokenizer which require src_lang and tgt_lang attributes.\n",
    "MULTILINGUAL_TOKENIZERS = [MBartTokenizer, MBartTokenizerFast, MBart50Tokenizer, MBart50TokenizerFast, M2M100Tokenizer]\n"
   ]
  },
  {
   "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",
    "\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",
    "\n",
    "    source_lang: str = field(default=None, metadata={\"help\": \"Source language id for translation.\"})\n",
    "    target_lang: str = field(default=None, metadata={\"help\": \"Target language id for translation.\"})\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",
    "    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",
    "    source_prefix: Optional[str] = field(\n",
    "        default=None, metadata={\"help\": \"A prefix to add before every source text (useful for T5 models).\"}\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",
    "    tune: bool = field(\n",
    "        default=False,\n",
    "        metadata={\"help\": \"Whether or not to apply quantization.\"},\n",
    "    )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11b035ee",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "model_args = ModelArguments(\n",
    "    model_name_or_path=\"t5-small\",\n",
    ")\n",
    "data_args = DataTrainingArguments(\n",
    "    source_lang=\"en\",\n",
    "    target_lang=\"ro\",\n",
    "    dataset_name=\"wmt16\",\n",
    "    dataset_config_name=\"ro-en\",\n",
    "    overwrite_cache=True,\n",
    "    max_eval_samples=400,\n",
    "    source_prefix=\"translate English to Romanian: \"\n",
    ")\n",
    "training_args = Seq2SeqTrainingArguments(\n",
    "    output_dir=\"./saved_results_dynamic\",\n",
    "    do_eval=True,\n",
    "    do_train=True,\n",
    "    no_cuda=True,\n",
    "    overwrite_output_dir=True,\n",
    "    per_device_eval_batch_size=8,\n",
    "    predict_with_generate=True\n",
    ")\n",
    "optim_args = OptimizationArguments(\n",
    "    tune=True,\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(\"wmt16\", \"ro-en\")"
   ]
  },
  {
   "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",
    "# download model & vocab.\n",
    "config = AutoConfig.from_pretrained(\n",
    "    \"t5-small\",\n",
    "    revision=\"main\"\n",
    ")\n",
    "tokenizer = AutoTokenizer.from_pretrained(\n",
    "    \"t5-small\",\n",
    "    revision=\"main\",\n",
    "    use_fast=True\n",
    ")\n",
    "model = AutoModelForSeq2SeqLM.from_pretrained(\n",
    "    \"t5-small\",\n",
    "    from_tf=bool(\".ckpt\" in model_args.model_name_or_path),\n",
    "    config=config,\n",
    "    revision=\"main\"\n",
    ")\n",
    "\n",
    "model.resize_token_embeddings(len(tokenizer))\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": [
    "# We need to tokenize inputs and targets.\n",
    "column_names = raw_datasets[\"train\"].column_names\n",
    "\n",
    "# Get the language codes for input/target.\n",
    "source_lang = data_args.source_lang.split(\"_\")[0]\n",
    "target_lang = data_args.target_lang.split(\"_\")[0]\n",
    "\n",
    "# Temporarily set max_target_length for training.\n",
    "max_target_length = 128\n",
    "padding = False\n",
    "\n",
    "def preprocess_function(examples):\n",
    "    inputs = [ex[source_lang] for ex in examples[\"translation\"]]\n",
    "    targets = [ex[target_lang] for ex in examples[\"translation\"]]\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",
    "    model_inputs[\"labels\"] = labels[\"input_ids\"]\n",
    "    return model_inputs\n",
    "\n",
    "# define 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",
    "# define eval 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(\"sacrebleu\")\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",
    "    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",
    "    # 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)\n",
    "    result = {\"bleu\": result[\"score\"]}\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",
    "\n",
    "metric_name = \"eval_bleu\"\n",
    "max_length = 128\n",
    "num_beams = None"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "942d2967",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Quantization & Benchmark"
   ]
  },
  {
   "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 before initializing model.\n",
    "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",
    "\n",
    "# # tuning\n",
    "model.config.save_pretrained(training_args.output_dir)\n",
    "trainer.save_model(training_args.output_dir)\n",
    "\n",
    "tune_metric = nlp_metrics.Metric(\n",
    "    name=metric_name, is_relative=True, criterion=0.25\n",
    ")\n",
    "quantization_config = QuantizationConfig(\n",
    "    approach=\"PostTrainingDynamic\",\n",
    "    max_trials=200,\n",
    "    metrics=[tune_metric],\n",
    ")\n",
    "trainer.max_length = max_length\n",
    "trainer.num_beams = num_beams\n",
    "trainer.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": [
    "results = trainer.evaluate(max_length=max_length, num_beams=num_beams)\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": "4f477616",
   "metadata": {},
   "source": [
    "## Run Benchmark after Dynamic Post Training Quantization with Multi-Instance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d40db91b",
   "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 before initializing model.\n",
    "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\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_fp32[metric_name]))\n",
    "print(\"Latency: {:.5f} ms\".format(1000 / throughput_fp32))\n",
    "print(\"Throughput: {:.5f} samples/sec\".format(throughput_fp32))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96e88ce1",
   "metadata": {},
   "source": [
    "## Run Benchmark for FP32 Model with Multi-Instance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ade4c930",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "results = os.system('bash ../multi_instance.sh --model=t5-small --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
}