docs/tutorials/pytorch/text-classification/SWEET.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Separating Weights for Early-Exit Transformers (SWEET) for Text Classification with BERT or DeBERTa"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This code implements SWEET Early-Exit fine-tuning method described in the paper [**Finding the SWEET spot:** Improvement and Analysis of Adaptive Inference in Low Resource Settings](https://arxiv.org/abs/2306.02307) and is based on [SWEET repo](https://github.com/schwartz-lab-NLP/SWEET).\n",
    "\n",
    "\n",
    "Below is an illustration of Early-Exit and Multi-Model adaptive inference approaches. In both methods, multiple classifiers of increasing sizes are run serially, until a confident prediction is made. In Early-Exit (left), a single model with multiple classifiers is used, such that early computations are reused by later classifiers. In Multi-Model (right), a sequence of independent models is used, allowing each classifier to decouple its parameters from other classifiers.\n",
    "<p align=\"center\">\n",
    "    <img src='../../../../examples/huggingface/pytorch/text-classification/early-exit/figures/EE_MM_illustration.png'  width=30% height=30% align=\"center\" alt=\"Illustration of Early-Exit and Multi-Model approaches.\">\n",
    "</p>\n",
    "\n",
    "Below is an illustration of the standard Early-Exit fine-tuning method (left), where lower layers get gradient updates from multiple classifiers, with the SWEET method (right), in which each layer parameters are updated only by the next classifier.\n",
    "\n",
    "<p align=\"center\">\n",
    "    <img src='../../../../examples/huggingface/pytorch/text-classification/early-exit/figures/SWEET_illustration.png'  width=30% height=30% align=\"center\" alt=\"Illustration of the SWEET method\">\n",
    "</p>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prerequisite"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install torch transformers datasets allennlp scipy sklearn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running Text Classification on MNLI with SWEET"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import os\n",
    "import random\n",
    "import sys\n",
    "from dataclasses import dataclass, field\n",
    "from os import path\n",
    "from typing import Optional\n",
    "\n",
    "import datasets\n",
    "import numpy as np\n",
    "import torch\n",
    "import transformers\n",
    "from datasets import load_dataset, load_metric\n",
    "from transformers import (\n",
    "    AutoTokenizer,\n",
    "    DataCollatorWithPadding,\n",
    "    EvalPrediction,\n",
    "    PretrainedConfig,\n",
    "    Trainer,\n",
    "    TrainingArguments,\n",
    "    default_data_collator,\n",
    "    set_seed,\n",
    ")\n",
    "from transformers.trainer_utils import get_last_checkpoint\n",
    "from transformers.utils import send_example_telemetry\n",
    "from transformers.utils.versions import require_version\n",
    "\n",
    "from models.bert.configuration_bert import BertConfig\n",
    "from models.bert.modeling_bert import BertForSequenceClassification\n",
    "from models.deberta.configuration_deberta import DebertaConfig\n",
    "from models.deberta.modeling_deberta import DebertaForSequenceClassification\n",
    "from utils.temperature_scaling import set_temperature\n",
    "\n",
    "require_version(\"datasets>=1.8.0\", \"To fix: pip install -r examples/pytorch/text-classification/requirements.txt\")\n",
    "\n",
    "os.environ[\"WANDB_DISABLED\"] = \"true\"\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__)\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",
    "    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\": (\n",
    "                \"The maximum total input sequence length after tokenization. Sequences longer \"\n",
    "                \"than this will be truncated, sequences shorter will be padded.\"\n",
    "            )\n",
    "        },\n",
    "    )\n",
    "    overwrite_cache: bool = field(\n",
    "        default=False, metadata={\"help\": \"Overwrite the cached preprocessed datasets or not.\"}\n",
    "    )\n",
    "\n",
    "    pad_to_max_length: bool = field(\n",
    "        default=True,\n",
    "        metadata={\n",
    "            \"help\": (\n",
    "                \"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",
    "    )\n",
    "    max_train_samples: Optional[int] = field(\n",
    "        default=None,\n",
    "        metadata={\n",
    "            \"help\": (\n",
    "                \"For debugging purposes or quicker training, truncate the number of training examples to this \"\n",
    "                \"value if set.\"\n",
    "            )\n",
    "        },\n",
    "    )\n",
    "    max_eval_samples: Optional[int] = field(\n",
    "        default=None,\n",
    "        metadata={\n",
    "            \"help\": (\n",
    "                \"For debugging purposes or quicker training, truncate the number of evaluation examples to this \"\n",
    "                \"value if set.\"\n",
    "            )\n",
    "        },\n",
    "    )\n",
    "    max_predict_samples: Optional[int] = field(\n",
    "        default=None,\n",
    "        metadata={\n",
    "            \"help\": (\n",
    "                \"For debugging purposes or quicker training, truncate the number of prediction examples to this \"\n",
    "                \"value if set.\"\n",
    "            )\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",
    "    do_calibration: Optional[bool] = field(\n",
    "        default=False,\n",
    "        metadata={\"help\": \"A boolean flag indicating whether the model should be calibrated prior to evaluation\"}\n",
    "    )\n",
    "    eval_seperate_layers: Optional[bool] = field(\n",
    "        default=False,\n",
    "        metadata={\n",
    "            \"help\": \"A boolean flag indicating whether the model should not be saved after evaluation to save space\"}\n",
    "    )\n",
    "    test_file: Optional[str] = field(default=None, metadata={\"help\": \"A csv or a json file containing the test data.\"})\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",
    "    exit_layers: str = field(\n",
    "        default=None,\n",
    "        metadata={\"help\": \"layers of the model for which to train classifiers.\"},\n",
    "    )\n",
    "    SWEET: bool = field(\n",
    "        default=True,\n",
    "        metadata={\"help\": \"A flag indicating whether to use the SWEET fine tuning method.\"},\n",
    "    )\n",
    "    gold_exit_layer: int = field(\n",
    "        default=None,\n",
    "        metadata={\"help\": \"An exit layer to use during evaluation.\"},\n",
    "    )\n",
    "    exit_thresholds: str = field(\n",
    "        default=1,\n",
    "        metadata={\"help\": \"Number of exit threshold for inference early exit (ranged  1/num_labels to 1\"},\n",
    "    )\n",
    "    exit_strategy: str = field(\n",
    "        default='confidence_threshold',\n",
    "        metadata={\"help\": \"A function implementing a decision rule for \\\"early exit\\\"\"}\n",
    "    )\n",
    "    exit_kwargs: str = field(\n",
    "        default=None,\n",
    "        metadata={\"help\": \"arguments for the exit strategy if needed\"}\n",
    "    )\n",
    "    deberta: bool = field(\n",
    "        default=False,\n",
    "        metadata={\n",
    "            \"help\": \"A flag indicating whether the pretrained model is DeBERTa (by default, the pretrained model is BERT).\"},\n",
    "    )\n",
    "    use_auth_token: bool = field(\n",
    "        default=False,\n",
    "        metadata={\n",
    "            \"help\": (\n",
    "                \"Will use the token generated when running `transformers-cli login` (necessary to use this script \"\n",
    "                \"with private models).\"\n",
    "            )\n",
    "        },\n",
    "    )\n",
    "    ignore_mismatched_sizes: bool = field(\n",
    "        default=False,\n",
    "        metadata={\"help\": \"Will enable to load a pretrained model whose head dimensions are different.\"},\n",
    "    )\n",
    "    do_not_save_model: bool = field(\n",
    "        default=False, metadata={\"help\": \"erase the saved model weights.\"}\n",
    "    )\n",
    "\n",
    "\n",
    "def parse_exit_thresholds(num_labels, num_thresholds, is_lte=False):\n",
    "    if num_thresholds is not None:\n",
    "        if is_lte:\n",
    "            lower = 0\n",
    "        else:\n",
    "            lower = 1 / num_labels\n",
    "        return list(np.linspace(lower, 1, num_thresholds))\n",
    "    else:\n",
    "        return [1.0]\n",
    "\n",
    "\n",
    "def parse_exit_layers(exit_layers):\n",
    "    if exit_layers is not None:\n",
    "        return [int(i) for i in exit_layers.split('_')]\n",
    "    return [-1]\n",
    "\n",
    "\n",
    "def parse_exit_kwargs(exit_kwargs):\n",
    "    if exit_kwargs is not None:\n",
    "        ks_and_vs = exit_kwargs.split(',')\n",
    "        exit_kwargs = {}\n",
    "        for k_v in ks_and_vs:\n",
    "            k, v = k_v.split(':')\n",
    "            exit_kwargs[k] = v\n",
    "        return exit_kwargs\n",
    "    return {}\n",
    "\n",
    "\n",
    "def move_to_device(batch, device='cuda'):\n",
    "    for key, tensor in batch.items():\n",
    "        batch[key] = tensor.to(device=device)\n",
    "    return batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# See all possible arguments in src/transformers/training_args.py\n",
    "# or by passing the --help flag to this script.`\n",
    "# We now keep distinct sets of args, for a cleaner separation of concerns.\n",
    "\n",
    "model_args = ModelArguments(\n",
    "    model_name_or_path=\"bert-base-uncased\",\n",
    "    exit_layers=\"0_3_5_11\",\n",
    "    exit_thresholds=11,\n",
    "    SWEET=True\n",
    ")\n",
    "data_args = DataTrainingArguments(\n",
    "    task_name=\"mnli\",\n",
    "    do_calibration=True,\n",
    "    max_train_samples=6000,\n",
    "    max_seq_length=256,\n",
    ")\n",
    "training_args = TrainingArguments(\n",
    "    per_device_train_batch_size=16,\n",
    "    per_device_eval_batch_size=1,\n",
    "    do_train=True,\n",
    "    do_eval=True,\n",
    "    output_dir=\"sweet\",\n",
    "    learning_rate=5e-5,\n",
    "    num_train_epochs=2,\n",
    "    overwrite_output_dir=True\n",
    ")\n",
    "\n",
    "# Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The\n",
    "# information sent is the one passed as arguments along with your Python/PyTorch versions.\n",
    "send_example_telemetry(\"run_glue\", model_args, data_args)\n",
    "\n",
    "if training_args.do_eval and training_args.per_device_eval_batch_size != 1:\n",
    "    training_args.per_device_eval_batch_size = 1\n",
    "    print(\"Setting per_device_eval_batch_size to 1 since early exit works only with batch size 1.\")\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}\")\n",
    "\n",
    "# Detecting last checkpoint.\n",
    "last_checkpoint = None\n",
    "if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:\n",
    "    last_checkpoint = get_last_checkpoint(training_args.output_dir)\n",
    "    if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:\n",
    "        raise ValueError(\n",
    "            f\"Output directory ({training_args.output_dir}) already exists and is not empty. \"\n",
    "            \"Use --overwrite_output_dir to overcome.\"\n",
    "        )\n",
    "    elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:\n",
    "        logger.info(\n",
    "            f\"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change \"\n",
    "            \"the `--output_dir` or add `--overwrite_output_dir` to train from scratch.\"\n",
    "        )\n",
    "\n",
    "# Set seed before initializing model.\n",
    "set_seed(training_args.seed)\n",
    "\n",
    "# Get the datasets: you can either provide your own CSV/JSON training and evaluation files (see below)\n",
    "# or specify a GLUE benchmark task (the dataset will be downloaded automatically from the datasets Hub).\n",
    "#\n",
    "# For CSV/JSON files, this script will use as labels the column called 'label' and as pair of sentences the\n",
    "# sentences in columns called 'sentence1' and 'sentence2' if such column exists or the first two columns not named\n",
    "# label if at least two columns are provided.\n",
    "#\n",
    "# If the CSVs/JSONs contain only one non-label column, the script does single sentence classification on this\n",
    "# single column. You can easily tweak this behavior (see below)\n",
    "#\n",
    "# In distributed training, the load_dataset function guarantee that only one local process can concurrently\n",
    "# download the dataset.\n",
    "if data_args.task_name is not None:\n",
    "    # Downloading and loading a dataset from the hub.\n",
    "    raw_datasets = load_dataset(\n",
    "        \"glue\",\n",
    "        data_args.task_name,\n",
    "        cache_dir=model_args.cache_dir,\n",
    "        use_auth_token=True if model_args.use_auth_token else None,\n",
    "    )\n",
    "elif data_args.dataset_name is not None:\n",
    "    # Downloading and loading a dataset from the hub.\n",
    "    raw_datasets = load_dataset(\n",
    "        data_args.dataset_name,\n",
    "        data_args.dataset_config_name,\n",
    "        cache_dir=model_args.cache_dir,\n",
    "        use_auth_token=True if model_args.use_auth_token else None)\n",
    "else:\n",
    "    # Loading a dataset from your local files.\n",
    "    # CSV/JSON training and evaluation files are needed.\n",
    "    data_files = {\"train\": data_args.train_file, \"validation\": data_args.validation_file}\n",
    "\n",
    "    # Get the test dataset: you can provide your own CSV/JSON test file (see below)\n",
    "    # when you use `do_predict` without specifying a GLUE benchmark task.\n",
    "    if training_args.do_predict:\n",
    "        if data_args.test_file is not None:\n",
    "            train_extension = data_args.train_file.split(\".\")[-1]\n",
    "            test_extension = data_args.test_file.split(\".\")[-1]\n",
    "            assert (\n",
    "                    test_extension == train_extension\n",
    "            ), \"`test_file` should have the same extension (csv or json) as `train_file`.\"\n",
    "            data_files[\"test\"] = data_args.test_file\n",
    "        else:\n",
    "            raise ValueError(\"Need either a GLUE task or a test file for `do_predict`.\")\n",
    "\n",
    "    for key in data_files.keys():\n",
    "        logger.info(f\"load a local file for {key}: {data_files[key]}\")\n",
    "\n",
    "    if data_args.train_file.endswith(\".csv\"):\n",
    "        # Loading a dataset from local csv files\n",
    "        raw_datasets = load_dataset(\n",
    "            \"csv\",\n",
    "            data_files=data_files,\n",
    "            cache_dir=model_args.cache_dir,\n",
    "            use_auth_token=True if model_args.use_auth_token else None,\n",
    "        )\n",
    "\n",
    "    else:\n",
    "        # Loading a dataset from local json files\n",
    "        raw_datasets = load_dataset(\n",
    "            \"json\",\n",
    "            data_files=data_files,\n",
    "            cache_dir=model_args.cache_dir,\n",
    "            use_auth_token=True if model_args.use_auth_token else None,\n",
    "        )\n",
    "# See more about loading any type of standard or custom dataset at\n",
    "# https://huggingface.co/docs/datasets/loading_datasets.html.\n",
    "\n",
    "# Labels\n",
    "if data_args.task_name is not None:\n",
    "    is_regression = data_args.task_name == \"stsb\"\n",
    "    if not is_regression:\n",
    "        label_list = raw_datasets[\"train\"].features[\"label\"].names\n",
    "        num_labels = len(label_list)\n",
    "    else:\n",
    "        num_labels = 1\n",
    "\n",
    "\n",
    "else:\n",
    "    # Trying to have good defaults here, don't hesitate to tweak to your needs.\n",
    "    is_regression = raw_datasets[\"train\"].features[\"label\"].dtype in [\"float32\", \"float64\"]\n",
    "    if is_regression:\n",
    "\n",
    "        num_labels = 1\n",
    "    else:\n",
    "        # A useful fast method:\n",
    "        # https://huggingface.co/docs/datasets/package_reference/main_classes.html#datasets.Dataset.unique\n",
    "        label_list = raw_datasets[\"train\"].unique(\"label\")\n",
    "        label_list.sort()  # Let's sort it for determinism\n",
    "        num_labels = len(label_list)\n",
    "\n",
    "# Load pretrained model and tokenizer\n",
    "#\n",
    "# In distributed training, the .from_pretrained methods guarantee that only one local process can concurrently\n",
    "# download model & vocab.\n",
    "exit_layers = parse_exit_layers(model_args.exit_layers)\n",
    "exit_kwargs = parse_exit_kwargs(model_args.exit_kwargs)\n",
    "exit_thresholds = parse_exit_thresholds(int(num_labels), int(model_args.exit_thresholds),\n",
    "                                        model_args.exit_strategy == 'lte')\n",
    "\n",
    "if model_args.deberta:\n",
    "    my_config = DebertaConfig\n",
    "    my_model = DebertaForSequenceClassification\n",
    "else:\n",
    "    my_config = BertConfig\n",
    "    my_model = BertForSequenceClassification\n",
    "\n",
    "config = my_config.from_pretrained(\n",
    "    model_args.config_name if model_args.config_name else model_args.model_name_or_path,\n",
    "    num_labels=num_labels,\n",
    "    finetuning_task=data_args.task_name,\n",
    "    cache_dir=model_args.cache_dir,\n",
    "    revision=model_args.model_revision,\n",
    "    use_auth_token=True if model_args.use_auth_token else None,\n",
    "    exit_layers=exit_layers,\n",
    "    SWEET=model_args.SWEET,\n",
    "    gold_exit_layer=model_args.gold_exit_layer,\n",
    "    exit_thresholds=exit_thresholds,\n",
    "    exit_strategy=model_args.exit_strategy,\n",
    "    exit_kwargs=exit_kwargs)\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(\n",
    "    model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,\n",
    "    cache_dir=model_args.cache_dir,\n",
    "    use_fast=model_args.use_fast_tokenizer,\n",
    "    revision=model_args.model_revision,\n",
    "    use_auth_token=True if model_args.use_auth_token else None,\n",
    ")\n",
    "model = my_model.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",
    "    cache_dir=model_args.cache_dir,\n",
    "    revision=model_args.model_revision,\n",
    "    use_auth_token=True if model_args.use_auth_token else None,\n",
    "    ignore_mismatched_sizes=model_args.ignore_mismatched_sizes,\n",
    ")\n",
    "# Preprocessing the raw_datasets\n",
    "if data_args.task_name is not None:\n",
    "    sentence1_key, sentence2_key = task_to_keys[data_args.task_name]\n",
    "\n",
    "else:\n",
    "    # Again, we try to have some nice defaults but don't hesitate to tweak to your use case.\n",
    "    non_label_column_names = [name for name in raw_datasets[\"train\"].column_names if name != \"label\"]\n",
    "    if \"sentence1\" in non_label_column_names and \"sentence2\" in non_label_column_names:\n",
    "        sentence1_key, sentence2_key = \"sentence1\", \"sentence2\"\n",
    "    else:\n",
    "        if len(non_label_column_names) >= 2:\n",
    "            sentence1_key, sentence2_key = non_label_column_names[:2]\n",
    "        else:\n",
    "            sentence1_key, sentence2_key = non_label_column_names[0], None\n",
    "\n",
    "# Padding strategy\n",
    "if data_args.pad_to_max_length:\n",
    "    padding = \"max_length\"\n",
    "else:\n",
    "    # We will pad later, dynamically at batch creation, to the max sequence length in each batch\n",
    "    padding = False\n",
    "\n",
    "# Some models have set the order of the labels to use, so let's make sure we do use it.\n",
    "# You can define your custom compute_metrics function. It takes an `EvalPrediction` object (a namedtuple with a\n",
    "# predictions and label_ids field) and has to return a dictionary string to float.\n",
    "\n",
    "def calculate_speedup_ratio(counts):\n",
    "    logger.info(f'counts = {counts}, exit layers = {exit_layers}')\n",
    "    counts_times_layers = [count * (exit_layer + 1) for count, exit_layer in zip(counts, exit_layers)]\n",
    "    return sum(counts_times_layers) / (sum(counts) * (exit_layers[-1] + 1))\n",
    "\n",
    "def compute_multiexit_metrics(p: EvalPrediction, gold_layer=None):\n",
    "    preds = p.predictions[0] if isinstance(p.predictions, tuple) else p.predictions\n",
    "\n",
    "    # This entire segments it a workaround aimed at not changing the trainer file. We rely on all logits in a\n",
    "    # prediction to be equal as an indicator of the previous layer being the exit point for the sample\n",
    "    final_predictions = np.zeros(\n",
    "        (preds.shape[0], preds.shape[-1]))  # will hold the predictions of the final exit layer to give a prediction\n",
    "    layer_prediction_indices = {}\n",
    "    layer_counts = [0 for _ in range(preds.shape[1])]\n",
    "\n",
    "    for i, prediction in enumerate(preds):\n",
    "        layer_ind = 0\n",
    "        for layer_logits in prediction[1:]:\n",
    "            if np.all(layer_logits == layer_logits[0]) and gold_layer is None:\n",
    "                break\n",
    "            elif gold_layer == exit_layers[layer_ind]:\n",
    "                break\n",
    "            else:\n",
    "                layer_ind += 1\n",
    "        layer_counts[layer_ind] += 1  # add a sample prediction count to the previous layer\n",
    "        final_predictions[i] = prediction[layer_ind]  # Update the actual prediction on the i-th sample\n",
    "        if layer_ind not in layer_prediction_indices:  # Update that the i-th sample was classified by the\n",
    "            # (j-1)st exit layer\n",
    "            layer_prediction_indices[layer_ind] = []\n",
    "        layer_prediction_indices[layer_ind].append(i)\n",
    "    preds = np.squeeze(final_predictions) if is_regression else np.argmax(final_predictions, axis=1)\n",
    "\n",
    "    results = metric.compute(predictions=preds, references=p.label_ids)\n",
    "    results['speedup_ratio'] = calculate_speedup_ratio(layer_counts)\n",
    "\n",
    "    metric_name = list(results.keys())[0]\n",
    "    if data_args.task_name is not None or data_args.dataset_name is not None:\n",
    "        for layer_ind, exit_layer in enumerate(exit_layers):\n",
    "            # for layer_ind ,sample_indices in layer_prediction_indices.items():\n",
    "            #     exit_layer = exit_layers[layer_ind]\n",
    "            if layer_ind not in layer_prediction_indices:\n",
    "                results[f'layer_{exit_layer}'] = {metric_name: 0, 'counts': layer_counts[layer_ind]}\n",
    "            else:\n",
    "                sample_indices = layer_prediction_indices[layer_ind]\n",
    "                results[f'layer_{exit_layer}'] = metric.compute(predictions=preds[sample_indices],\n",
    "                                                                references=(p.label_ids)[sample_indices])\n",
    "                results[f'layer_{exit_layer}']['counts'] = layer_counts[layer_ind]\n",
    "\n",
    "        return results\n",
    "    elif is_regression:\n",
    "        return {\"mse\": ((preds - p.label_ids) ** 2).mean().item()}\n",
    "    else:\n",
    "        return {\"accuracy\": (preds == p.label_ids).astype(np.float32).mean().item()}\n",
    "\n",
    "label_to_id = None\n",
    "if (model.config.label2id != PretrainedConfig(num_labels=num_labels).label2id\n",
    "        and data_args.task_name is not None\n",
    "        and not is_regression):\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",
    "            \"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",
    "            \"\\nIgnoring the model labels as a result.\",\n",
    "        )\n",
    "elif data_args.task_name is None and not is_regression:\n",
    "    label_to_id = {v: i for i, v in enumerate(label_list)}\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",
    "elif data_args.task_name is not None and not is_regression:\n",
    "    model.config.label2id = {l: i for i, l in enumerate(label_list)}\n",
    "    model.config.id2label = {id: label for label, id in config.label2id.items()}\n",
    "if data_args.max_seq_length > tokenizer.model_max_length:\n",
    "    logger.warning(\n",
    "        f\"The max_seq_length passed ({data_args.max_seq_length}) is larger than the maximum length for the\"\n",
    "        f\"model ({tokenizer.model_max_length}). Using max_seq_length={tokenizer.model_max_length}.\"\n",
    "    )\n",
    "max_seq_length = min(data_args.max_seq_length, tokenizer.model_max_length)\n",
    "\n",
    "def preprocess_function(examples):\n",
    "    # Tokenize the texts\n",
    "    args = (\n",
    "        (examples[sentence1_key],) if sentence2_key is None else (\n",
    "            examples[sentence1_key], examples[sentence2_key])\n",
    "    )\n",
    "    result = tokenizer(*args, padding=padding, max_length=max_seq_length, truncation=True)\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,\n",
    "        batched=True,\n",
    "        load_from_cache_file=not data_args.overwrite_cache,\n",
    "        desc=\"Running tokenizer on dataset\",\n",
    "    )\n",
    "\n",
    "# Shuffle and make sure samples are identical for same seeds across experiments\n",
    "for dataset in raw_datasets.values():\n",
    "    dataset.shuffle(seed=training_args.seed)\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",
    "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",
    "    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",
    "if training_args.do_predict or data_args.task_name is not None or data_args.test_file is not None:\n",
    "    if \"test\" not in raw_datasets and \"test_matched\" not in raw_datasets:\n",
    "        raise ValueError(\"--do_predict requires a test dataset\")\n",
    "    predict_dataset = raw_datasets[\"test_matched\" if data_args.task_name == \"mnli\" else \"test\"]\n",
    "    if data_args.max_predict_samples is not None:\n",
    "        max_predict_samples = min(len(predict_dataset), data_args.max_predict_samples)\n",
    "        predict_dataset = predict_dataset.select(range(max_predict_samples))\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",
    "if data_args.task_name is not None:\n",
    "    metric = load_metric(\"glue\", data_args.task_name)\n",
    "else:\n",
    "    metric = load_metric(\"accuracy\")\n",
    "\n",
    "# Data collator will default to DataCollatorWithPadding when the tokenizer is passed to Trainer, so we change it if\n",
    "# we already did the padding.\n",
    "if data_args.pad_to_max_length:\n",
    "    data_collator = default_data_collator\n",
    "elif training_args.fp16:\n",
    "    data_collator = DataCollatorWithPadding(tokenizer, pad_to_multiple_of=8)\n",
    "else:\n",
    "    data_collator = None\n",
    "\n",
    "# Initialize our Trainer\n",
    "trainer = Trainer(\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_multiexit_metrics,\n",
    "    tokenizer=tokenizer,\n",
    "    data_collator=data_collator,\n",
    ")\n",
    "# Training\n",
    "\n",
    "if training_args.do_train:\n",
    "    checkpoint = None\n",
    "    if training_args.resume_from_checkpoint is not None:\n",
    "        checkpoint = training_args.resume_from_checkpoint\n",
    "    elif last_checkpoint is not None:\n",
    "        checkpoint = last_checkpoint\n",
    "    train_result = trainer.train(resume_from_checkpoint=checkpoint)\n",
    "    metrics = train_result.metrics\n",
    "    max_train_samples = (\n",
    "        data_args.max_train_samples if data_args.max_train_samples is not None else len(train_dataset)\n",
    "    )\n",
    "    metrics[\"train_samples\"] = min(max_train_samples, len(train_dataset))\n",
    "    if not model_args.do_not_save_model:\n",
    "        trainer.save_model()  # Saves the tokenizer too for easy upload\n",
    "    trainer.log_metrics(\"train\", metrics)\n",
    "    trainer.save_metrics(\"train\", metrics)\n",
    "    trainer.save_state()\n",
    "\n",
    "# Calibration\n",
    "if data_args.do_calibration:\n",
    "    calibration_dataset = eval_dataset\n",
    "    if model_args.exit_strategy == 'lte':\n",
    "        min_lte, max_lte = None, None\n",
    "        eval_dataloader = trainer.get_eval_dataloader(calibration_dataset)\n",
    "        for step, batch in enumerate(eval_dataloader):\n",
    "            batch = move_to_device(batch, training_args.device)\n",
    "            with torch.no_grad():\n",
    "                outputs = model(**batch, output_hidden_states=True)\n",
    "            if type(outputs.hidden_states) is tuple:\n",
    "                logger.info(f'hidden states is a tuple: {outputs.hidden_states}')\n",
    "            lte_logits = outputs.hidden_states.cpu()\n",
    "            if min_lte is None:\n",
    "                min_lte = lte_logits\n",
    "            if max_lte is None:\n",
    "                max_lte = lte_logits\n",
    "            min_lte = torch.minimum(min_lte, lte_logits)\n",
    "            max_lte = torch.maximum(max_lte, lte_logits)\n",
    "        model.set_min_max_lte(min_lte, max_lte)\n",
    "    else:\n",
    "        temperature_file_path = path.join(training_args.output_dir, 'scaling_temperatures.txt')\n",
    "        if path.exists(temperature_file_path):\n",
    "            with open(temperature_file_path, 'r') as temp_file:\n",
    "                temps = temp_file.readlines()\n",
    "        else:\n",
    "            cuda_device = -1 if training_args.device.type == 'cpu' else 0\n",
    "            temps = set_temperature(model, len(model.exit_layers), calibration_dataset, cuda_device)\n",
    "            with open(temperature_file_path, 'a') as temp_file:\n",
    "                for t in temps:\n",
    "                    temp_file.write(str(t) + '\\n')\n",
    "        logger.info(f'Calibrated temperatures: {temps}')\n",
    "        model.set_scaling_temperatures(temps)\n",
    "model = model.to('cpu')\n",
    "# Evaluation\n",
    "\n",
    "if training_args.do_eval:\n",
    "    # training_args.output_dir = base_output_dir\n",
    "    trainer = Trainer(\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_multiexit_metrics,\n",
    "        tokenizer=tokenizer,\n",
    "        data_collator=data_collator,\n",
    "    )\n",
    "\n",
    "    logger.info(\"*** Evaluate ***\")\n",
    "\n",
    "    # Loop to handle MNLI double evaluation (matched, mis-matched)\n",
    "    tasks = [data_args.task_name]\n",
    "    eval_datasets = [eval_dataset]\n",
    "\n",
    "    # If you want to automatically handle mismatched when handling mnli, uncomment this segment (and debug):\n",
    "    # if data_args.task_name == \"mnli\":\n",
    "    #     tasks.append(\"mnli-mm\")\n",
    "    #     mm_dataset = raw_datasets[\"validation_mismatched\"]\n",
    "    #     if data_args.max_eval_samples is not None:\n",
    "    #         mm_dataset = mm_dataset.select(range(data_args.max_eval_samples))\n",
    "    #     eval_datasets.append(mm_dataset)\n",
    "    total_metrics = {}\n",
    "    for eval_dataset, task in zip(eval_datasets, tasks):\n",
    "        for thr in exit_thresholds:\n",
    "            model.set_exit_strategy(model_args.exit_strategy, threshold=thr)\n",
    "            metrics = trainer.evaluate(eval_dataset=eval_dataset)\n",
    "            max_eval_samples = (\n",
    "                data_args.max_eval_samples if data_args.max_eval_samples is not None else len(eval_dataset))\n",
    "            metrics[\"eval_samples\"] = min(max_eval_samples, len(eval_dataset))\n",
    "            if task == \"mnli-mm\":\n",
    "                metrics = {k + \"_mm\": v for k, v in metrics.items()}\n",
    "\n",
    "            # If you want to automatically handle mismatched when handling mnli, uncomment this segment (and debug):\n",
    "            # if task is not None and \"mnli\" in task:\n",
    "            #     combined.update(metrics)\n",
    "            # total_metrics[thr] = combined if task is not None and \"mnli\" in task else metrics\n",
    "            total_metrics[thr] = metrics\n",
    "            metrics_for_logging = {k: v for k, v in metrics.items() if 'layer' not in k}\n",
    "            trainer.log_metrics(f\"eval_for_thr {thr}\", metrics_for_logging)\n",
    "        trainer.save_metrics(f\"threshold\", total_metrics, combined=False)\n",
    "    model.to(training_args.device)\n",
    "\n",
    "if training_args.do_predict:\n",
    "    logger.info(\"*** Predict ***\")\n",
    "\n",
    "    # Loop to handle MNLI double evaluation (matched, mis-matched)\n",
    "    tasks = [data_args.task_name]\n",
    "    predict_datasets = [predict_dataset]\n",
    "    if data_args.task_name == \"mnli\":\n",
    "        tasks.append(\"mnli-mm\")\n",
    "        predict_datasets.append(raw_datasets[\"test_mismatched\"])\n",
    "\n",
    "    for predict_dataset, task in zip(predict_datasets, tasks):\n",
    "        # Removing the `label` columns because it contains -1 and Trainer won't like that.\n",
    "        predict_dataset = predict_dataset.remove_columns(\"label\")\n",
    "        predictions = trainer.predict(predict_dataset, metric_key_prefix=\"predict\").predictions\n",
    "        predictions = np.squeeze(predictions) if is_regression else np.argmax(predictions, axis=1)\n",
    "\n",
    "        output_predict_file = os.path.join(training_args.output_dir, f\"predict_results_{task}.txt\")\n",
    "        if trainer.is_world_process_zero():\n",
    "            with open(output_predict_file, \"w\") as writer:\n",
    "                logger.info(f\"***** Predict results {task} *****\")\n",
    "                writer.write(\"index\\tprediction\\n\")\n",
    "                for index, item in enumerate(predictions):\n",
    "                    if is_regression:\n",
    "                        writer.write(f\"{index}\\t{item:3.3f}\\n\")\n",
    "                    else:\n",
    "                        item = label_list[item]\n",
    "                        writer.write(f\"{index}\\t{item}\\n\")\n",
    "\n",
    "kwargs = {\"finetuned_from\": model_args.model_name_or_path, \"tasks\": \"text-classification\"}\n",
    "if data_args.task_name is not None:\n",
    "    kwargs[\"language\"] = \"en\"\n",
    "    kwargs[\"dataset_tags\"] = \"glue\"\n",
    "    kwargs[\"dataset_args\"] = data_args.task_name\n",
    "    kwargs[\"dataset\"] = f\"GLUE {data_args.task_name.upper()}\"\n",
    "\n",
    "if training_args.push_to_hub:\n",
    "    trainer.push_to_hub(**kwargs)\n",
    "else:\n",
    "    trainer.create_model_card(**kwargs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Citation\n",
    "\n",
    "```\n",
    "@misc{rotem2023finding,\n",
    "      title={Finding the SWEET Spot: Analysis and Improvement of Adaptive Inference in Low Resource Settings}, \n",
    "      author={Daniel Rotem and Michael Hassid and Jonathan Mamou and Roy Schwartz},\n",
    "      year={2023},\n",
    "      eprint={2306.02307},\n",
    "      archivePrefix={arXiv},\n",
    "      primaryClass={cs.CL}\n",
    "}\n",
    "```"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py39",
   "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.9.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}