notebooks/gpt2/joint_multitask_learning.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4bfe1985-c81c-4593-bfd8-999742b9938d",
   "metadata": {
    "papermill": {
     "duration": 0.004448,
     "end_time": "2024-03-23T20:55:59.915880",
     "exception": false,
     "start_time": "2024-03-23T20:55:59.911432",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Joint multitask"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f10fe697",
   "metadata": {},
   "source": [
    "**GPU Requirements:** For running with GPT-2 you may be fine with just 8GB of GPU RAM. With about 24GB you should be able to run any 7B or 13B model. With 80GB (A100) GPU you may be able to run a 70B model."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c49a3f5-8c7f-4128-8c40-0b134035fad7",
   "metadata": {
    "papermill": {
     "duration": 0.003695,
     "end_time": "2024-03-23T20:55:59.923451",
     "exception": false,
     "start_time": "2024-03-23T20:55:59.919756",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "The point of this notebook is not to do anything useful, but to show what is possible with relatively low effort using the *transformer_heads* library. In this example, we will train four heads on a transformer model while using qlora to finetune the transformer block weights. The first head will be hooked at layer 9 (-4) and predict the sentiment of imdb reviews (text classification). The second head will be hooked at the last layer (-1 or 12) and does causal language modelling on imdb reviews. The third head will be hooked at layer 6 (-7) and will learn to count the number of occurences of each letter of the alphabet occuring in imdb reviews (Text-level regression). The final head will be hooked at layer 4 (-9) and will predict how many tokens will follow before the review ends for each token in imdb reviews (Token-level regression). The final head will also be a small mlp instead of a linear head.\n",
    "\n",
    "All heads and the qlora parameters will be trained jointly (multi-task learning)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "13030870-2149-4d09-83d3-2bfcb591ba23",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-23T20:55:59.931807Z",
     "iopub.status.busy": "2024-03-23T20:55:59.931611Z",
     "iopub.status.idle": "2024-03-23T20:56:06.141560Z",
     "shell.execute_reply": "2024-03-23T20:56:06.140956Z"
    },
    "papermill": {
     "duration": 6.216004,
     "end_time": "2024-03-23T20:56:06.143203",
     "exception": false,
     "start_time": "2024-03-23T20:55:59.927199",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from transformer_heads import create_headed_qlora, load_lora_with_heads\n",
    "from datasets import load_dataset\n",
    "from transformers import (\n",
    "    AutoTokenizer,\n",
    "    MistralForCausalLM,\n",
    "    Trainer,\n",
    "    BitsAndBytesConfig,\n",
    "    TrainingArguments,\n",
    "    GPT2Model,\n",
    "    GPT2LMHeadModel,\n",
    ")\n",
    "from transformer_heads.util.helpers import DataCollatorWithPadding, get_model_params\n",
    "from peft import LoraConfig\n",
    "from transformer_heads.config import HeadConfig\n",
    "from transformer_heads.util.model import print_trainable_parameters\n",
    "from transformer_heads.util.evaluate import (\n",
    "    evaluate_head_wise,\n",
    ")\n",
    "import torch\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6fc42b9c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-23T20:56:06.151615Z",
     "iopub.status.busy": "2024-03-23T20:56:06.151225Z",
     "iopub.status.idle": "2024-03-23T20:56:06.154531Z",
     "shell.execute_reply": "2024-03-23T20:56:06.154016Z"
    },
    "papermill": {
     "duration": 0.008537,
     "end_time": "2024-03-23T20:56:06.155645",
     "exception": false,
     "start_time": "2024-03-23T20:56:06.147108",
     "status": "completed"
    },
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "# GPT2 is the fastest and requires fewest memory. However, this works just the same with any Llama or Mistral model. Just change model_path to its huggingface path.\n",
    "model_path = \"gpt2\"\n",
    "train_epochs = 1\n",
    "eval_epochs = 1\n",
    "logging_steps = 100\n",
    "train_batch_size = 4\n",
    "eval_batch_size = 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "0b50d4c4-7a22-4b7a-a092-e28165dca6ed",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-23T20:56:06.174268Z",
     "iopub.status.busy": "2024-03-23T20:56:06.173905Z",
     "iopub.status.idle": "2024-03-23T20:56:06.639276Z",
     "shell.execute_reply": "2024-03-23T20:56:06.638775Z"
    },
    "papermill": {
     "duration": 0.470451,
     "end_time": "2024-03-23T20:56:06.640448",
     "exception": false,
     "start_time": "2024-03-23T20:56:06.169997",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'vocab_size': 50257, 'n_positions': 1024, 'n_embd': 768, 'n_layer': 12, 'n_head': 12, 'n_inner': None, 'activation_function': 'gelu_new', 'resid_pdrop': 0.1, 'embd_pdrop': 0.1, 'attn_pdrop': 0.1, 'layer_norm_epsilon': 1e-05, 'initializer_range': 0.02, 'summary_type': 'cls_index', 'summary_use_proj': True, 'summary_activation': None, 'summary_first_dropout': 0.1, 'summary_proj_to_labels': True, 'scale_attn_weights': True, 'use_cache': True, 'scale_attn_by_inverse_layer_idx': False, 'reorder_and_upcast_attn': False, 'bos_token_id': 50256, 'eos_token_id': 50256, 'return_dict': True, 'output_hidden_states': False, 'output_attentions': False, 'torchscript': False, 'torch_dtype': None, 'use_bfloat16': False, 'tf_legacy_loss': False, 'pruned_heads': {}, 'tie_word_embeddings': True, 'chunk_size_feed_forward': 0, 'is_encoder_decoder': False, 'is_decoder': False, 'cross_attention_hidden_size': None, 'add_cross_attention': False, 'tie_encoder_decoder': False, 'max_length': 20, 'min_length': 0, 'do_sample': False, 'early_stopping': False, 'num_beams': 1, 'num_beam_groups': 1, 'diversity_penalty': 0.0, 'temperature': 1.0, 'top_k': 50, 'top_p': 1.0, 'typical_p': 1.0, 'repetition_penalty': 1.0, 'length_penalty': 1.0, 'no_repeat_ngram_size': 0, 'encoder_no_repeat_ngram_size': 0, 'bad_words_ids': None, 'num_return_sequences': 1, 'output_scores': False, 'return_dict_in_generate': False, 'forced_bos_token_id': None, 'forced_eos_token_id': None, 'remove_invalid_values': False, 'exponential_decay_length_penalty': None, 'suppress_tokens': None, 'begin_suppress_tokens': None, 'architectures': ['GPT2LMHeadModel'], 'finetuning_task': None, 'id2label': {0: 'LABEL_0', 1: 'LABEL_1'}, 'label2id': {'LABEL_0': 0, 'LABEL_1': 1}, 'tokenizer_class': None, 'prefix': None, 'pad_token_id': None, 'sep_token_id': None, 'decoder_start_token_id': None, 'task_specific_params': {'text-generation': {'do_sample': True, 'max_length': 50}}, 'problem_type': None, '_name_or_path': 'gpt2', 'transformers_version': '4.37.2', 'model_type': 'gpt2', 'n_ctx': 1024, 'model_class': <class 'transformers.models.gpt2.modeling_gpt2.GPT2LMHeadModel'>, 'hidden_size': 768}\n"
     ]
    }
   ],
   "source": [
    "model_params = get_model_params(model_path)\n",
    "model_class = model_params[\"model_class\"]\n",
    "hidden_size = model_params[\"hidden_size\"]\n",
    "vocab_size = model_params[\"vocab_size\"]\n",
    "print(model_params)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f544f84a-c846-4701-9ae9-3ea63f43cc24",
   "metadata": {
    "papermill": {
     "duration": 0.003426,
     "end_time": "2024-03-23T20:56:06.647637",
     "exception": false,
     "start_time": "2024-03-23T20:56:06.644211",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Define the various different heads. Given the differences in loss functions and magnitudes of label data in the dataset, it is important to weigh the losses of each head so that training is given similar importance for all of them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "41be78c0-e80c-436d-b763-bbee21fb5d35",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-23T20:56:06.655660Z",
     "iopub.status.busy": "2024-03-23T20:56:06.655281Z",
     "iopub.status.idle": "2024-03-23T20:56:06.660789Z",
     "shell.execute_reply": "2024-03-23T20:56:06.660340Z"
    },
    "papermill": {
     "duration": 0.01071,
     "end_time": "2024-03-23T20:56:06.661901",
     "exception": false,
     "start_time": "2024-03-23T20:56:06.651191",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "head_configs = [\n",
    "    HeadConfig(\n",
    "        name=f\"sentiment_head\",\n",
    "        layer_hook=-4,\n",
    "        in_size=hidden_size,\n",
    "        output_activation=\"linear\",\n",
    "        pred_for_sequence=True,\n",
    "        loss_fct=\"cross_entropy\",\n",
    "        num_outputs=2,\n",
    "        loss_weight=2.0,\n",
    "    ),\n",
    "    HeadConfig(\n",
    "        name=f\"causal_lm\",\n",
    "        layer_hook=-1,\n",
    "        in_size=hidden_size,\n",
    "        output_activation=\"linear\",\n",
    "        is_causal_lm=True,\n",
    "        loss_fct=\"cross_entropy\",\n",
    "        num_outputs=vocab_size,\n",
    "        is_regression=False,\n",
    "        output_bias=False,\n",
    "        loss_weight=1.0,\n",
    "    ),\n",
    "    HeadConfig(\n",
    "        name=f\"alphabet_regression\",\n",
    "        layer_hook=-7,\n",
    "        in_size=hidden_size,\n",
    "        output_activation=\"linear\",\n",
    "        is_causal_lm=False,\n",
    "        pred_for_sequence=True,\n",
    "        loss_fct=\"mse\",\n",
    "        num_outputs=26,  # 26 letters in the alphabet\n",
    "        is_regression=True,\n",
    "        loss_weight=0.002,\n",
    "    ),\n",
    "    HeadConfig(\n",
    "        name=f\"num_tokens_regression\",\n",
    "        layer_hook=-7,\n",
    "        hidden_size=128,  # MLP hidden size\n",
    "        num_layers=3,  # 2 hidden layers in MLP\n",
    "        in_size=hidden_size,\n",
    "        output_activation=\"linear\",\n",
    "        is_causal_lm=False,\n",
    "        pred_for_sequence=False,\n",
    "        loss_fct=\"mse\",\n",
    "        num_outputs=1,\n",
    "        is_regression=True,\n",
    "        loss_weight=0.0002,\n",
    "    ),\n",
    "    HeadConfig(\n",
    "        name=f\"lm_head\",  # Let's also keep the original lm head for comparison\n",
    "        layer_hook=-1,\n",
    "        in_size=hidden_size,\n",
    "        output_activation=\"linear\",\n",
    "        is_causal_lm=True,\n",
    "        pred_for_sequence=False,\n",
    "        loss_fct=\"cross_entropy\",\n",
    "        num_outputs=vocab_size,\n",
    "        is_regression=False,\n",
    "        trainable=False,  # Keep it in it's pretrained state\n",
    "    ),\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "46ffdd47-e302-4ee5-a4fc-bbd956e6a4a6",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-23T20:56:06.669833Z",
     "iopub.status.busy": "2024-03-23T20:56:06.669458Z",
     "iopub.status.idle": "2024-03-23T20:56:14.888180Z",
     "shell.execute_reply": "2024-03-23T20:56:14.887653Z"
    },
    "papermill": {
     "duration": 8.22439,
     "end_time": "2024-03-23T20:56:14.889792",
     "exception": false,
     "start_time": "2024-03-23T20:56:06.665402",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "dd = load_dataset(\"imdb\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "4df5fed8-17df-4745-bbf3-f3880b269d7d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-23T20:56:14.898501Z",
     "iopub.status.busy": "2024-03-23T20:56:14.898155Z",
     "iopub.status.idle": "2024-03-23T20:56:54.299011Z",
     "shell.execute_reply": "2024-03-23T20:56:54.298453Z"
    },
    "papermill": {
     "duration": 39.406891,
     "end_time": "2024-03-23T20:56:54.300764",
     "exception": false,
     "start_time": "2024-03-23T20:56:14.893873",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c84751918dc04134b5fa195f0ca2d4b4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/25000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "da8cde3491ce4af38cd5264274ddf8cf",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/25000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d6e5a1e307914a1d9a62e28b9d870e43",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/50000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(model_path)\n",
    "if tokenizer.pad_token_id is None:\n",
    "    tokenizer.pad_token = tokenizer.eos_token\n",
    "\n",
    "\n",
    "def processing_function(examples):\n",
    "    out = tokenizer(examples[\"text\"], padding=False, truncation=True)\n",
    "    out[\"sentiment_head\"] = examples[\"label\"]\n",
    "    out[\"causal_lm\"] = out[\"lm_head\"] = out[\"input_ids\"].copy()\n",
    "    out[\"num_tokens_regression\"] = [\n",
    "        list(map(float, range(len(ids) - 1, -1, -1))) for ids in out[\"input_ids\"]\n",
    "    ]\n",
    "    out[\"alphabet_regression\"] = [\n",
    "        [\n",
    "            float(text.count(x) + text.count(x.upper()))\n",
    "            for x in \"abcdefghijklmnopqrstuvwxyz\"\n",
    "        ]\n",
    "        for text in examples[\"text\"]\n",
    "    ]\n",
    "    return out\n",
    "\n",
    "\n",
    "for split in dd.keys():\n",
    "    dd[split] = dd[split].filter(function=lambda example: len(example[\"text\"]) > 10)\n",
    "    dd[split] = dd[split].shuffle()\n",
    "    dd[split] = dd[split].map(processing_function, batched=True)\n",
    "\n",
    "dd.set_format(\n",
    "    type=\"torch\",\n",
    "    columns=[\"input_ids\", \"attention_mask\"] + [x.name for x in head_configs],\n",
    ")\n",
    "for split in dd.keys():\n",
    "    dd[split] = dd[split].remove_columns([\"text\", \"label\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "c1ddb342-5be7-4ff4-afc6-2fb0cd104eeb",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-23T20:56:54.313445Z",
     "iopub.status.busy": "2024-03-23T20:56:54.313157Z",
     "iopub.status.idle": "2024-03-23T20:56:54.317295Z",
     "shell.execute_reply": "2024-03-23T20:56:54.316803Z"
    },
    "papermill": {
     "duration": 0.011592,
     "end_time": "2024-03-23T20:56:54.319114",
     "exception": false,
     "start_time": "2024-03-23T20:56:54.307522",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['input_ids', 'attention_mask', 'sentiment_head', 'causal_lm', 'lm_head', 'num_tokens_regression', 'alphabet_regression'],\n",
       "    num_rows: 25000\n",
       "})"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dd[\"test\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a278656d",
   "metadata": {
    "papermill": {
     "duration": 0.003963,
     "end_time": "2024-03-23T20:56:54.327272",
     "exception": false,
     "start_time": "2024-03-23T20:56:54.323309",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Setting *target_modules=None* in the qlora config will make *create_headed_qlora* create LoRA modules for all linear layers in the transformer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "6b35c555-df0e-47b6-a772-1d25c7280fed",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-23T20:56:54.336340Z",
     "iopub.status.busy": "2024-03-23T20:56:54.335969Z",
     "iopub.status.idle": "2024-03-23T20:56:58.169023Z",
     "shell.execute_reply": "2024-03-23T20:56:58.168490Z"
    },
    "papermill": {
     "duration": 3.839035,
     "end_time": "2024-03-23T20:56:58.170331",
     "exception": false,
     "start_time": "2024-03-23T20:56:54.331296",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of TransformerWithHeads were not initialized from the model checkpoint at gpt2 and are newly initialized: ['heads.alphabet_regression.lins.0.weight', 'heads.causal_lm.lins.0.weight', 'heads.num_tokens_regression.lins.0.bias', 'heads.num_tokens_regression.lins.0.weight', 'heads.num_tokens_regression.lins.1.bias', 'heads.num_tokens_regression.lins.1.weight', 'heads.num_tokens_regression.lins.2.weight', 'heads.sentiment_head.lins.0.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "quantization_config = BitsAndBytesConfig(\n",
    "    load_in_4bit=True,\n",
    "    load_in_8bit=False,\n",
    "    llm_int8_threshold=6.0,\n",
    "    llm_int8_has_fp16_weight=False,\n",
    "    bnb_4bit_compute_dtype=torch.float32,\n",
    "    bnb_4bit_use_double_quant=True,\n",
    "    bnb_4bit_quant_type=\"nf4\",\n",
    ")\n",
    "lora_config = LoraConfig(\n",
    "    r=32,\n",
    "    lora_alpha=16,\n",
    "    target_modules=None,\n",
    "    lora_dropout=0.0,\n",
    "    bias=\"none\",\n",
    "    task_type=\"CAUSAL_LM\",\n",
    ")\n",
    "model = create_headed_qlora(\n",
    "    base_model_class=model_class,\n",
    "    model_name=model_path,\n",
    "    quantization_config=quantization_config,\n",
    "    lora_config=lora_config,\n",
    "    head_configs=head_configs,\n",
    "    fully_trained_heads=True,\n",
    "    device_map={\"\": torch.cuda.current_device()},\n",
    "    gradient_checkpointing=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "1e69fa9b-4ad4-4d1f-ad0d-ee98f66925e3",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-23T20:56:58.180289Z",
     "iopub.status.busy": "2024-03-23T20:56:58.179912Z",
     "iopub.status.idle": "2024-03-23T20:56:58.184905Z",
     "shell.execute_reply": "2024-03-23T20:56:58.184418Z"
    },
    "papermill": {
     "duration": 0.01113,
     "end_time": "2024-03-23T20:56:58.186016",
     "exception": false,
     "start_time": "2024-03-23T20:56:58.174886",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "all params: 125425024 || trainable params: 43452544 || trainable%: 34.64423813863492\n",
      "params by dtype: defaultdict(<class 'int'>, {torch.float32: 82957696, torch.uint8: 42467328})\n",
      "trainable params by dtype: defaultdict(<class 'int'>, {torch.float32: 43452544})\n"
     ]
    }
   ],
   "source": [
    "print_trainable_parameters(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "aa750893-dcae-4227-b075-87e78c4f112b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-23T20:56:58.195435Z",
     "iopub.status.busy": "2024-03-23T20:56:58.195094Z",
     "iopub.status.idle": "2024-03-23T20:56:58.198326Z",
     "shell.execute_reply": "2024-03-23T20:56:58.197796Z"
    },
    "papermill": {
     "duration": 0.009233,
     "end_time": "2024-03-23T20:56:58.199441",
     "exception": false,
     "start_time": "2024-03-23T20:56:58.190208",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "collator = DataCollatorWithPadding(\n",
    "    feature_name_to_padding_value={\n",
    "        \"input_ids\": tokenizer.pad_token_id,\n",
    "        \"attention_mask\": 0,\n",
    "        \"causal_lm\": -100,\n",
    "        \"lm_head\": -100,\n",
    "        \"num_tokens_regression\": 0,\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "f195ef53-8bbf-4dcc-a375-34aaa3c8a407",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-23T20:56:58.208994Z",
     "iopub.status.busy": "2024-03-23T20:56:58.208669Z",
     "iopub.status.idle": "2024-03-23T21:28:49.661874Z",
     "shell.execute_reply": "2024-03-23T21:28:49.661153Z"
    },
    "papermill": {
     "duration": 1911.459433,
     "end_time": "2024-03-23T21:28:49.663119",
     "exception": false,
     "start_time": "2024-03-23T20:56:58.203686",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r\n",
      "Evaluating: 100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 6250/6250 [31:51<00:00,  3.27it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(51.503634765014645, {'sentiment_head': 5.004700186492335, 'causal_lm': 22.38208692779541, 'alphabet_regression': 4078.0613999365232, 'num_tokens_regression': 36263.337341796876, 'lm_head': 3.703357028274536})\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "print(\n",
    "    evaluate_head_wise(\n",
    "        model, dd[\"test\"], collator, epochs=eval_epochs, batch_size=eval_batch_size\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "f65dd030-50e8-44e3-869f-a2250f432494",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-23T21:28:50.275065Z",
     "iopub.status.busy": "2024-03-23T21:28:50.274737Z",
     "iopub.status.idle": "2024-03-23T21:54:17.323889Z",
     "shell.execute_reply": "2024-03-23T21:54:17.323355Z"
    },
    "papermill": {
     "duration": 1527.339569,
     "end_time": "2024-03-23T21:54:17.325017",
     "exception": false,
     "start_time": "2024-03-23T21:28:49.985448",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mykeller\u001b[0m (\u001b[33mchm-hci\u001b[0m). Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: wandb version 0.16.4 is available!  To upgrade, please run:\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m:  $ pip install wandb --upgrade\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: Tracking run with wandb version 0.16.3\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: Run data is saved locally in \u001b[35m\u001b[1m/raven/u/ykeller/transformer_heads/wandb/run-20240323_222852-2unafm5b\u001b[0m\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Run \u001b[1m`wandb offline`\u001b[0m to turn off syncing.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: Syncing run \u001b[33mfrosty-salad-210\u001b[0m\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: \u2b50\ufe0f View project at \u001b[34m\u001b[4mhttps://wandb.ai/chm-hci/huggingface\u001b[0m\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: \ud83d\ude80 View run at \u001b[34m\u001b[4mhttps://wandb.ai/chm-hci/huggingface/runs/2unafm5b\u001b[0m\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/u/ykeller/conda-envs/sh_finetuning/lib/python3.10/site-packages/torch/utils/checkpoint.py:460: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='6250' max='6250' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [6250/6250 25:18, Epoch 1/1]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>100</td>\n",
       "      <td>32.494500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>200</td>\n",
       "      <td>22.742800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>300</td>\n",
       "      <td>19.527100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>400</td>\n",
       "      <td>16.875700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>500</td>\n",
       "      <td>15.876800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>600</td>\n",
       "      <td>13.926000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>700</td>\n",
       "      <td>14.789600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>800</td>\n",
       "      <td>14.329600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>900</td>\n",
       "      <td>14.772700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1000</td>\n",
       "      <td>14.079500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1100</td>\n",
       "      <td>13.717200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1200</td>\n",
       "      <td>13.369500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1300</td>\n",
       "      <td>13.409500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1400</td>\n",
       "      <td>12.704300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1500</td>\n",
       "      <td>13.387900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1600</td>\n",
       "      <td>13.235000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1700</td>\n",
       "      <td>12.415100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1800</td>\n",
       "      <td>12.419800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1900</td>\n",
       "      <td>13.192000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2000</td>\n",
       "      <td>12.811300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2100</td>\n",
       "      <td>12.604200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2200</td>\n",
       "      <td>12.917700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2300</td>\n",
       "      <td>12.113800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2400</td>\n",
       "      <td>12.045700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2500</td>\n",
       "      <td>11.922300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2600</td>\n",
       "      <td>12.398900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2700</td>\n",
       "      <td>12.346600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2800</td>\n",
       "      <td>12.027800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2900</td>\n",
       "      <td>12.072100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3000</td>\n",
       "      <td>12.221600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3100</td>\n",
       "      <td>12.009200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3200</td>\n",
       "      <td>11.656600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3300</td>\n",
       "      <td>11.998200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3400</td>\n",
       "      <td>11.843900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3500</td>\n",
       "      <td>12.657400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3600</td>\n",
       "      <td>11.940500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3700</td>\n",
       "      <td>12.569700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3800</td>\n",
       "      <td>11.859200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3900</td>\n",
       "      <td>11.724500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4000</td>\n",
       "      <td>12.437400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4100</td>\n",
       "      <td>11.899900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4200</td>\n",
       "      <td>12.014900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4300</td>\n",
       "      <td>11.854900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4400</td>\n",
       "      <td>11.695200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4500</td>\n",
       "      <td>12.032300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4600</td>\n",
       "      <td>11.510000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4700</td>\n",
       "      <td>11.603700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4800</td>\n",
       "      <td>11.456200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4900</td>\n",
       "      <td>12.111700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5000</td>\n",
       "      <td>11.779700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5100</td>\n",
       "      <td>11.387700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5200</td>\n",
       "      <td>11.760800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5300</td>\n",
       "      <td>12.003100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5400</td>\n",
       "      <td>12.064100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5500</td>\n",
       "      <td>12.490600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5600</td>\n",
       "      <td>11.467500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5700</td>\n",
       "      <td>12.329100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5800</td>\n",
       "      <td>11.223600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5900</td>\n",
       "      <td>11.940400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>6000</td>\n",
       "      <td>11.661200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>6100</td>\n",
       "      <td>11.468700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>6200</td>\n",
       "      <td>11.421300</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Checkpoint destination directory imdb_linear_probe/checkpoint-500 already exists and is non-empty.Saving will proceed but saved results may be invalid.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/u/ykeller/conda-envs/sh_finetuning/lib/python3.10/site-packages/torch/utils/checkpoint.py:460: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Checkpoint destination directory imdb_linear_probe/checkpoint-1000 already exists and is non-empty.Saving will proceed but saved results may be invalid.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/u/ykeller/conda-envs/sh_finetuning/lib/python3.10/site-packages/torch/utils/checkpoint.py:460: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Checkpoint destination directory imdb_linear_probe/checkpoint-1500 already exists and is non-empty.Saving will proceed but saved results may be invalid.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/u/ykeller/conda-envs/sh_finetuning/lib/python3.10/site-packages/torch/utils/checkpoint.py:460: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Checkpoint destination directory imdb_linear_probe/checkpoint-2000 already exists and is non-empty.Saving will proceed but saved results may be invalid.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/u/ykeller/conda-envs/sh_finetuning/lib/python3.10/site-packages/torch/utils/checkpoint.py:460: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Checkpoint destination directory imdb_linear_probe/checkpoint-2500 already exists and is non-empty.Saving will proceed but saved results may be invalid.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/u/ykeller/conda-envs/sh_finetuning/lib/python3.10/site-packages/torch/utils/checkpoint.py:460: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Checkpoint destination directory imdb_linear_probe/checkpoint-3000 already exists and is non-empty.Saving will proceed but saved results may be invalid.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/u/ykeller/conda-envs/sh_finetuning/lib/python3.10/site-packages/torch/utils/checkpoint.py:460: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Checkpoint destination directory imdb_linear_probe/checkpoint-3500 already exists and is non-empty.Saving will proceed but saved results may be invalid.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/u/ykeller/conda-envs/sh_finetuning/lib/python3.10/site-packages/torch/utils/checkpoint.py:460: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Checkpoint destination directory imdb_linear_probe/checkpoint-4000 already exists and is non-empty.Saving will proceed but saved results may be invalid.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/u/ykeller/conda-envs/sh_finetuning/lib/python3.10/site-packages/torch/utils/checkpoint.py:460: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Checkpoint destination directory imdb_linear_probe/checkpoint-4500 already exists and is non-empty.Saving will proceed but saved results may be invalid.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/u/ykeller/conda-envs/sh_finetuning/lib/python3.10/site-packages/torch/utils/checkpoint.py:460: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Checkpoint destination directory imdb_linear_probe/checkpoint-5000 already exists and is non-empty.Saving will proceed but saved results may be invalid.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/u/ykeller/conda-envs/sh_finetuning/lib/python3.10/site-packages/torch/utils/checkpoint.py:460: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Checkpoint destination directory imdb_linear_probe/checkpoint-5500 already exists and is non-empty.Saving will proceed but saved results may be invalid.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/u/ykeller/conda-envs/sh_finetuning/lib/python3.10/site-packages/torch/utils/checkpoint.py:460: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Checkpoint destination directory imdb_linear_probe/checkpoint-6000 already exists and is non-empty.Saving will proceed but saved results may be invalid.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/u/ykeller/conda-envs/sh_finetuning/lib/python3.10/site-packages/torch/utils/checkpoint.py:460: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=6250, training_loss=13.10343646484375, metrics={'train_runtime': 1526.4757, 'train_samples_per_second': 16.378, 'train_steps_per_second': 4.094, 'total_flos': 1.0162662301258752e+16, 'train_loss': 13.10343646484375, 'epoch': 1.0})"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "args = TrainingArguments(\n",
    "    output_dir=\"imdb_linear_probe\",\n",
    "    learning_rate=0.0002,\n",
    "    num_train_epochs=train_epochs,  # To speed things up set to 0.1, set to 1 for better performance\n",
    "    logging_steps=logging_steps,\n",
    "    do_eval=False,\n",
    "    remove_unused_columns=False,\n",
    "    optim=\"paged_adamw_32bit\",\n",
    "    gradient_checkpointing=True,\n",
    "    lr_scheduler_type=\"constant\",\n",
    "    ddp_find_unused_parameters=False,\n",
    "    per_device_train_batch_size=train_batch_size,\n",
    "    per_device_eval_batch_size=eval_batch_size,\n",
    ")\n",
    "trainer = Trainer(\n",
    "    model,\n",
    "    args=args,\n",
    "    train_dataset=dd[\"train\"],\n",
    "    data_collator=collator,\n",
    ")\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "b83e5316-33b9-4882-a266-3515aacfc745",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-23T21:54:17.892327Z",
     "iopub.status.busy": "2024-03-23T21:54:17.892028Z",
     "iopub.status.idle": "2024-03-23T22:26:21.126220Z",
     "shell.execute_reply": "2024-03-23T22:26:21.125597Z"
    },
    "papermill": {
     "duration": 1923.519378,
     "end_time": "2024-03-23T22:26:21.127324",
     "exception": false,
     "start_time": "2024-03-23T21:54:17.607946",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r\n",
      "Evaluating: 100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 6250/6250 [32:03<00:00,  3.25it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/u/ykeller/conda-envs/sh_finetuning/lib/python3.10/site-packages/torch/utils/checkpoint.py:90: UserWarning: None of the inputs have requires_grad=True. Gradients will be None\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(11.534467044143677, {'sentiment_head': 0.32214655859653885, 'causal_lm': 4.127306442260743, 'alphabet_regression': 100.85618892456054, 'num_tokens_regression': 13946.417831679688, 'lm_head': 3.771871594276428})\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "evals = evaluate_head_wise(\n",
    "    model, dd[\"test\"], collator, epochs=eval_epochs, batch_size=eval_batch_size\n",
    ")\n",
    "print(evals)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dbf3b8d6",
   "metadata": {
    "papermill": {
     "duration": 0.594538,
     "end_time": "2024-03-23T22:26:22.289930",
     "exception": false,
     "start_time": "2024-03-23T22:26:21.695392",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Saving and loading\n",
    "Now let's how to save a complicated mulit-headed model to then load it again for inference. Saving is super easy. Just call save_pretrained and all trained parameters will be saved correctly. The saving will also work correctly for checkpoints created during training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "46978204-0c69-41b4-b36a-9da84bb4b4da",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-23T22:26:23.452949Z",
     "iopub.status.busy": "2024-03-23T22:26:23.452704Z",
     "iopub.status.idle": "2024-03-23T22:26:24.472662Z",
     "shell.execute_reply": "2024-03-23T22:26:24.472153Z"
    },
    "papermill": {
     "duration": 1.616777,
     "end_time": "2024-03-23T22:26:24.474299",
     "exception": false,
     "start_time": "2024-03-23T22:26:22.857522",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "model.save_pretrained(\"qlora_multitask_imdb\")\n",
    "del model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb41b2f8",
   "metadata": {
    "papermill": {
     "duration": 0.595746,
     "end_time": "2024-03-23T22:26:25.639375",
     "exception": false,
     "start_time": "2024-03-23T22:26:25.043629",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "While loading the model, we need to make sure to correctly attach and initialize all the heads, so that won't easily work with the huggingface api. Instead, *transformer_heads* provides the *load_lora_with_heads* function. Note that giving a quantization config is optional here. We could also give a different quantization config or none at all."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "c4d15108",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-23T22:26:26.797433Z",
     "iopub.status.busy": "2024-03-23T22:26:26.797154Z",
     "iopub.status.idle": "2024-03-23T22:26:28.950617Z",
     "shell.execute_reply": "2024-03-23T22:26:28.950047Z"
    },
    "papermill": {
     "duration": 2.747734,
     "end_time": "2024-03-23T22:26:28.952091",
     "exception": false,
     "start_time": "2024-03-23T22:26:26.204357",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "model = load_lora_with_heads(\n",
    "    model_class,\n",
    "    \"qlora_multitask_imdb\",\n",
    "    quantization_config,\n",
    "    device_map={\"\": torch.cuda.current_device()},\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "150f5bc9",
   "metadata": {
    "papermill": {
     "duration": 0.595812,
     "end_time": "2024-03-23T22:26:30.114610",
     "exception": false,
     "start_time": "2024-03-23T22:26:29.518798",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Let's now find out if the loaded model behaves the same as the saved model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "6eb62bce",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-23T22:26:31.276135Z",
     "iopub.status.busy": "2024-03-23T22:26:31.275857Z",
     "iopub.status.idle": "2024-03-23T22:58:14.979284Z",
     "shell.execute_reply": "2024-03-23T22:58:14.978662Z"
    },
    "papermill": {
     "duration": 1904.299677,
     "end_time": "2024-03-23T22:58:14.980431",
     "exception": false,
     "start_time": "2024-03-23T22:26:30.680754",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r\n",
      "Evaluating: 100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 6250/6250 [31:43<00:00,  3.28it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(11.163953929519653, {'sentiment_head': 0.3076160474227845, 'causal_lm': 4.014423089256287, 'alphabet_regression': 105.97276580749512, 'num_tokens_regression': 13420.149343261719, 'lm_head': 3.6383234076690676})\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "new_evals = evaluate_head_wise(\n",
    "    model, dd[\"test\"], collator, epochs=eval_epochs, batch_size=eval_batch_size\n",
    ")\n",
    "print(new_evals)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "sh_finetuning",
   "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.10.13"
  },
  "papermill": {
   "default_parameters": {},
   "duration": 7340.088484,
   "end_time": "2024-03-23T22:58:18.604223",
   "environment_variables": {},
   "exception": null,
   "input_path": "notebooks/base/joint_multitask_learning.ipynb",
   "output_path": "notebooks/gpt2/joint_multitask_learning.ipynb",
   "parameters": {
    "eval_epochs": 1,
    "logging_steps": 100,
    "model_path": "gpt2",
    "train_epochs": 1
   },
   "start_time": "2024-03-23T20:55:58.515739",
   "version": "2.5.0"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {
     "00409361c3c64065813a1824514c1ccd": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "FloatProgressModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "FloatProgressModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "ProgressView",
       "bar_style": "success",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_594d4a7eafb74776b595c31feb69c423",
       "max": 25000,
       "min": 0,
       "orientation": "horizontal",
       "style": "IPY_MODEL_195d19b4a06947daba9ee2accd93f3c7",
       "tabbable": null,
       "tooltip": null,
       "value": 25000
      }
     },
     "077618c3fae541f18c5b23f7eaabf9bf": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_43cf00abbef144b590616417220eaa6e",
       "placeholder": "\u200b",
       "style": "IPY_MODEL_902c8d85782c4a0ebb216862c3230c61",
       "tabbable": null,
       "tooltip": null,
       "value": "\u200725000/25000\u2007[00:09&lt;00:00,\u20072726.95\u2007examples/s]"
      }
     },
     "13aa26cc8c104d54822d7162ae69ed05": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "195d19b4a06947daba9ee2accd93f3c7": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "ProgressStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "ProgressStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "bar_color": null,
       "description_width": ""
      }
     },
     "21beb5bbee8e408ea6f26a40c0e32248": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "22a3b2861289413aba66a57f02ab64c2": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "32d16e5d5edd4a9cac536aea883606ab": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "354535af61de4254bd06dae31aeb0212": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "3bf1090602204a10a1b7353d7da3c8f9": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "FloatProgressModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "FloatProgressModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "ProgressView",
       "bar_style": "success",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_354535af61de4254bd06dae31aeb0212",
       "max": 25000,
       "min": 0,
       "orientation": "horizontal",
       "style": "IPY_MODEL_7743fc7b281e4625b0e54cab261d4acf",
       "tabbable": null,
       "tooltip": null,
       "value": 25000
      }
     },
     "3e20b9f55c534cc1bb6652a74b255eb1": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "43cf00abbef144b590616417220eaa6e": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "4d466e93321346f985db8a59a31b8a85": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "5433809f471941a1a262b0bb4c2ed3f9": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_c2679fcf286d4775b80ae9d54fcee44c",
       "placeholder": "\u200b",
       "style": "IPY_MODEL_d19aeb05221a4cceb03e17dc70948f1d",
       "tabbable": null,
       "tooltip": null,
       "value": "\u200725000/25000\u2007[00:09&lt;00:00,\u20072634.66\u2007examples/s]"
      }
     },
     "5838563bb5d64c919996f9475680666a": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_4d466e93321346f985db8a59a31b8a85",
       "placeholder": "\u200b",
       "style": "IPY_MODEL_6895196ca01143a0b04f017e57b547be",
       "tabbable": null,
       "tooltip": null,
       "value": "Map:\u2007100%"
      }
     },
     "594d4a7eafb74776b595c31feb69c423": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "60404243247347d4b138e93dc997e3c8": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_32d16e5d5edd4a9cac536aea883606ab",
       "placeholder": "\u200b",
       "style": "IPY_MODEL_ca7f6f1fae6140a480d0dc2610d22b75",
       "tabbable": null,
       "tooltip": null,
       "value": "\u200750000/50000\u2007[00:19&lt;00:00,\u20072618.89\u2007examples/s]"
      }
     },
     "64f6746f26534ffe82e6fae47f2185f9": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_3e20b9f55c534cc1bb6652a74b255eb1",
       "placeholder": "\u200b",
       "style": "IPY_MODEL_a930a9faafac4c72a3bdea294ff465c8",
       "tabbable": null,
       "tooltip": null,
       "value": "Map:\u2007100%"
      }
     },
     "6895196ca01143a0b04f017e57b547be": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "background": null,
       "description_width": "",
       "font_size": null,
       "text_color": null
      }
     },
     "7743fc7b281e4625b0e54cab261d4acf": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "ProgressStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "ProgressStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "bar_color": null,
       "description_width": ""
      }
     },
     "902c8d85782c4a0ebb216862c3230c61": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "background": null,
       "description_width": "",
       "font_size": null,
       "text_color": null
      }
     },
     "9f39b5db6cf543458bbb6131e7a6f063": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_e53e9a95ae334148a40cef385648c5dd",
       "placeholder": "\u200b",
       "style": "IPY_MODEL_e2c8e66a626c460fab010e8e91a4540e",
       "tabbable": null,
       "tooltip": null,
       "value": "Map:\u2007100%"
      }
     },
     "a930a9faafac4c72a3bdea294ff465c8": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "background": null,
       "description_width": "",
       "font_size": null,
       "text_color": null
      }
     },
     "bc86d427763a401082858ce152d29519": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "c2679fcf286d4775b80ae9d54fcee44c": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "c84751918dc04134b5fa195f0ca2d4b4": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HBoxModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HBoxModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HBoxView",
       "box_style": "",
       "children": [
        "IPY_MODEL_64f6746f26534ffe82e6fae47f2185f9",
        "IPY_MODEL_3bf1090602204a10a1b7353d7da3c8f9",
        "IPY_MODEL_5433809f471941a1a262b0bb4c2ed3f9"
       ],
       "layout": "IPY_MODEL_13aa26cc8c104d54822d7162ae69ed05",
       "tabbable": null,
       "tooltip": null
      }
     },
     "ca7f6f1fae6140a480d0dc2610d22b75": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "background": null,
       "description_width": "",
       "font_size": null,
       "text_color": null
      }
     },
     "d19aeb05221a4cceb03e17dc70948f1d": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "background": null,
       "description_width": "",
       "font_size": null,
       "text_color": null
      }
     },
     "d6e5a1e307914a1d9a62e28b9d870e43": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HBoxModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HBoxModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HBoxView",
       "box_style": "",
       "children": [
        "IPY_MODEL_5838563bb5d64c919996f9475680666a",
        "IPY_MODEL_f3afcd0ab47b4d48b2780415578f21f7",
        "IPY_MODEL_60404243247347d4b138e93dc997e3c8"
       ],
       "layout": "IPY_MODEL_22a3b2861289413aba66a57f02ab64c2",
       "tabbable": null,
       "tooltip": null
      }
     },
     "da8cde3491ce4af38cd5264274ddf8cf": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HBoxModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HBoxModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HBoxView",
       "box_style": "",
       "children": [
        "IPY_MODEL_9f39b5db6cf543458bbb6131e7a6f063",
        "IPY_MODEL_00409361c3c64065813a1824514c1ccd",
        "IPY_MODEL_077618c3fae541f18c5b23f7eaabf9bf"
       ],
       "layout": "IPY_MODEL_21beb5bbee8e408ea6f26a40c0e32248",
       "tabbable": null,
       "tooltip": null
      }
     },
     "dff71698d75441bf94802eddd133a080": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "ProgressStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "ProgressStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "bar_color": null,
       "description_width": ""
      }
     },
     "e2c8e66a626c460fab010e8e91a4540e": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "background": null,
       "description_width": "",
       "font_size": null,
       "text_color": null
      }
     },
     "e53e9a95ae334148a40cef385648c5dd": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "f3afcd0ab47b4d48b2780415578f21f7": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "FloatProgressModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "FloatProgressModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "ProgressView",
       "bar_style": "success",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_bc86d427763a401082858ce152d29519",
       "max": 50000,
       "min": 0,
       "orientation": "horizontal",
       "style": "IPY_MODEL_dff71698d75441bf94802eddd133a080",
       "tabbable": null,
       "tooltip": null,
       "value": 50000
      }
     }
    },
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}