notebooks/gpt2/linear_probe.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "54b4a6ba",
   "metadata": {
    "papermill": {
     "duration": 0.005449,
     "end_time": "2024-03-23T20:47:06.033910",
     "exception": false,
     "start_time": "2024-03-23T20:47:06.028461",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Training a simple linear probe on a transformer model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8574ff07",
   "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": "51b3e12a",
   "metadata": {
    "papermill": {
     "duration": 0.003441,
     "end_time": "2024-03-23T20:47:06.041407",
     "exception": false,
     "start_time": "2024-03-23T20:47:06.037966",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "In this introductory notebook, we will train a simple linear probe for a transformer model to check if causal language modelling understanding for the wikitext dataset is present in a transformer model even at some intermediate layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "80043d10",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-23T20:47:06.049329Z",
     "iopub.status.busy": "2024-03-23T20:47:06.049083Z",
     "iopub.status.idle": "2024-03-23T20:47:42.632268Z",
     "shell.execute_reply": "2024-03-23T20:47:42.631606Z"
    },
    "papermill": {
     "duration": 36.589463,
     "end_time": "2024-03-23T20:47:42.634311",
     "exception": false,
     "start_time": "2024-03-23T20:47:06.044848",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Standard imports\n",
    "from transformers import (\n",
    "    AutoTokenizer,\n",
    "    MistralForCausalLM,\n",
    "    Trainer,\n",
    "    BitsAndBytesConfig,\n",
    "    TrainingArguments,\n",
    "    GPT2Model,\n",
    "    GPT2LMHeadModel,\n",
    ")\n",
    "from datasets import load_dataset\n",
    "from peft import LoraConfig\n",
    "import torch\n",
    "\n",
    "# Imports from the transformer_heads library\n",
    "from transformer_heads import load_headed\n",
    "from transformer_heads.util.helpers import DataCollatorWithPadding, get_model_params\n",
    "from transformer_heads.config import HeadConfig\n",
    "from transformer_heads.util.model import print_trainable_parameters\n",
    "from transformer_heads.util.evaluate import evaluate_head_wise, get_top_n_preds"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f42e762b",
   "metadata": {
    "papermill": {
     "duration": 0.003609,
     "end_time": "2024-03-23T20:47:42.643557",
     "exception": false,
     "start_time": "2024-03-23T20:47:42.639948",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Set the model and it's model parameters. Default is GPT2, but you can also use Mistral 7b if you have enough GPU RAM (and are willing to wait longer for training to complete)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "7aaf7d13",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-23T20:47:42.651847Z",
     "iopub.status.busy": "2024-03-23T20:47:42.651401Z",
     "iopub.status.idle": "2024-03-23T20:47:42.654781Z",
     "shell.execute_reply": "2024-03-23T20:47:42.654313Z"
    },
    "papermill": {
     "duration": 0.008796,
     "end_time": "2024-03-23T20:47:42.655833",
     "exception": false,
     "start_time": "2024-03-23T20:47:42.647037",
     "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9c5f4b4f",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-23T20:47:42.681451Z",
     "iopub.status.busy": "2024-03-23T20:47:42.681268Z",
     "iopub.status.idle": "2024-03-23T20:47:43.195160Z",
     "shell.execute_reply": "2024-03-23T20:47:43.194683Z"
    },
    "papermill": {
     "duration": 0.519624,
     "end_time": "2024-03-23T20:47:43.196768",
     "exception": false,
     "start_time": "2024-03-23T20:47:42.677144",
     "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": "8e414c4c",
   "metadata": {
    "papermill": {
     "duration": 0.003605,
     "end_time": "2024-03-23T20:47:43.204351",
     "exception": false,
     "start_time": "2024-03-23T20:47:43.200746",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Start out by configuring the linear probing head. In this example we hook at layer -4. This is using the python indexing format. E.g. gpt-2 has 12 transformer blocks. Hooking at layer -4 means that the linear probe processes the hidden state after the 9th layer while hooking at layer -1 would mean processing the hidden state after the last (12th) transformer block.\n",
    "\n",
    "Otherwise, we are setting *num_layers* to 1, to make sure that we are actually training a *linear* probe and not an mlp. With *is_causal_lm* we specify the type of task that the model is supposed to learn."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "e6444d2d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-23T20:47:43.213445Z",
     "iopub.status.busy": "2024-03-23T20:47:43.213253Z",
     "iopub.status.idle": "2024-03-23T20:47:43.216598Z",
     "shell.execute_reply": "2024-03-23T20:47:43.216117Z"
    },
    "papermill": {
     "duration": 0.008649,
     "end_time": "2024-03-23T20:47:43.217627",
     "exception": false,
     "start_time": "2024-03-23T20:47:43.208978",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "heads_configs = [\n",
    "    HeadConfig(\n",
    "        name=\"wikitext_head\",\n",
    "        layer_hook=-4,  # Hook to layer [-4] (Drop 3 layers from the end)\n",
    "        in_size=hidden_size,\n",
    "        num_layers=1,\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",
    "    )\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "264f547a",
   "metadata": {
    "papermill": {
     "duration": 0.003675,
     "end_time": "2024-03-23T20:47:43.224940",
     "exception": false,
     "start_time": "2024-03-23T20:47:43.221265",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Now we load and format our dataset. We need to make sure that the dataset has labels stored for each head that we want to train. In case of the causal language modelling task, these labels are just copys of the input_ids."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "5faaf4d1",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-23T20:47:43.233162Z",
     "iopub.status.busy": "2024-03-23T20:47:43.232763Z",
     "iopub.status.idle": "2024-03-23T20:47:53.573073Z",
     "shell.execute_reply": "2024-03-23T20:47:53.572502Z"
    },
    "papermill": {
     "duration": 10.346162,
     "end_time": "2024-03-23T20:47:53.574742",
     "exception": false,
     "start_time": "2024-03-23T20:47:43.228580",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "dd = load_dataset(\"wikitext\", \"wikitext-2-v1\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "4e07be6c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-23T20:47:53.583697Z",
     "iopub.status.busy": "2024-03-23T20:47:53.583429Z",
     "iopub.status.idle": "2024-03-23T20:47:56.453080Z",
     "shell.execute_reply": "2024-03-23T20:47:56.452561Z"
    },
    "papermill": {
     "duration": 2.875777,
     "end_time": "2024-03-23T20:47:56.454668",
     "exception": false,
     "start_time": "2024-03-23T20:47:53.578891",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e1870d6c6b1947cbbf6c1767d7d41ba4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/23627 [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 tokenize_function(examples):\n",
    "    out = tokenizer(examples[\"text\"], padding=False, truncation=True)\n",
    "    out[heads_configs[0].name] = out[\"input_ids\"].copy()\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].map(tokenize_function, batched=True)\n",
    "dd.set_format(\n",
    "    type=\"torch\", columns=[\"input_ids\", \"attention_mask\", heads_configs[0].name]\n",
    ")\n",
    "for split in dd.keys():\n",
    "    dd[split] = dd[split].remove_columns(\"text\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "78131ee5",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-23T20:47:56.463973Z",
     "iopub.status.busy": "2024-03-23T20:47:56.463769Z",
     "iopub.status.idle": "2024-03-23T20:47:56.467731Z",
     "shell.execute_reply": "2024-03-23T20:47:56.467265Z"
    },
    "papermill": {
     "duration": 0.009843,
     "end_time": "2024-03-23T20:47:56.468878",
     "exception": false,
     "start_time": "2024-03-23T20:47:56.459035",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['input_ids', 'attention_mask', 'wikitext_head'],\n",
       "    num_rows: 23627\n",
       "})"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dd[\"train\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f75133e",
   "metadata": {
    "papermill": {
     "duration": 0.003909,
     "end_time": "2024-03-23T20:47:56.476928",
     "exception": false,
     "start_time": "2024-03-23T20:47:56.473019",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Now it is time to load our model. The load_headed function of transformer_heads is great for loading frozen models with a linear probe. To save GPU memory, we will load the model in a quantized state. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "416838c1",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-23T20:47:56.485616Z",
     "iopub.status.busy": "2024-03-23T20:47:56.485432Z",
     "iopub.status.idle": "2024-03-23T20:48:05.710241Z",
     "shell.execute_reply": "2024-03-23T20:48:05.709607Z"
    },
    "papermill": {
     "duration": 9.230758,
     "end_time": "2024-03-23T20:48:05.711621",
     "exception": false,
     "start_time": "2024-03-23T20:47:56.480863",
     "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.wikitext_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",
    "\n",
    "model = load_headed(\n",
    "    model_class,\n",
    "    model_path,\n",
    "    head_configs=heads_configs,\n",
    "    quantization_config=quantization_config,\n",
    "    device_map={\"\": torch.cuda.current_device()},\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5ba0b48",
   "metadata": {
    "papermill": {
     "duration": 0.004219,
     "end_time": "2024-03-23T20:48:05.720535",
     "exception": false,
     "start_time": "2024-03-23T20:48:05.716316",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "That warning about weights not initialized from the model checkpoint is exactly what we want to see. We want a newly initialized linear probe that is not in the pretrained gpt2 checkpoint."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30323ad5",
   "metadata": {
    "papermill": {
     "duration": 0.004282,
     "end_time": "2024-03-23T20:48:05.728887",
     "exception": false,
     "start_time": "2024-03-23T20:48:05.724605",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Let's check some data about our model using the convenience method *print_trainable_parameters*."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "fe7ca798",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-23T20:48:05.738352Z",
     "iopub.status.busy": "2024-03-23T20:48:05.737995Z",
     "iopub.status.idle": "2024-03-23T20:48:05.742241Z",
     "shell.execute_reply": "2024-03-23T20:48:05.741654Z"
    },
    "papermill": {
     "duration": 0.010246,
     "end_time": "2024-03-23T20:48:05.743401",
     "exception": false,
     "start_time": "2024-03-23T20:48:05.733155",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "all params: 120569856 || trainable params: 38597376 || trainable%: 32.01245923359152\n",
      "params by dtype: defaultdict(<class 'int'>, {torch.float32: 78102528, torch.uint8: 42467328})\n",
      "trainable params by dtype: defaultdict(<class 'int'>, {torch.float32: 38597376})\n"
     ]
    }
   ],
   "source": [
    "print_trainable_parameters(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f83bae86",
   "metadata": {
    "papermill": {
     "duration": 0.004089,
     "end_time": "2024-03-23T20:48:05.751733",
     "exception": false,
     "start_time": "2024-03-23T20:48:05.747644",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Given that gpt2 is a fairly small model with a large vocab size, our single linear probe already has quite a lot of parameters compared to the rest of the model. Every parameter in the model that is not part of the linear probe is frozen (has requires_grad set to false)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d63b570",
   "metadata": {
    "papermill": {
     "duration": 0.004147,
     "end_time": "2024-03-23T20:48:05.760142",
     "exception": false,
     "start_time": "2024-03-23T20:48:05.755995",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Let's see how our linear probe performs before it is trained"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "6032324a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-23T20:48:05.769183Z",
     "iopub.status.busy": "2024-03-23T20:48:05.768986Z",
     "iopub.status.idle": "2024-03-23T20:48:07.742818Z",
     "shell.execute_reply": "2024-03-23T20:48:07.742135Z"
    },
    "papermill": {
     "duration": 1.979807,
     "end_time": "2024-03-23T20:48:07.744051",
     "exception": false,
     "start_time": "2024-03-23T20:48:05.764244",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'wikitext_head': ['aunder', ' mesh', 'inki', 'Thread', ' Prediction']}\n"
     ]
    }
   ],
   "source": [
    "print(\n",
    "    get_top_n_preds(\n",
    "        n=5, model=model, text=\"The historical significance of\", tokenizer=tokenizer\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71e528e3",
   "metadata": {
    "papermill": {
     "duration": 0.004313,
     "end_time": "2024-03-23T20:48:07.753141",
     "exception": false,
     "start_time": "2024-03-23T20:48:07.748828",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "As expected, this is pretty random.\n",
    "\n",
    "Let's train the linear probe now using huggingfaces simple to use Trainer class. Note that we are using a custom collator here, to handle the labels under the heads_configs names correctly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "fce0ef54",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-23T20:48:07.762938Z",
     "iopub.status.busy": "2024-03-23T20:48:07.762561Z",
     "iopub.status.idle": "2024-03-23T20:55:07.276825Z",
     "shell.execute_reply": "2024-03-23T20:55:07.276197Z"
    },
    "papermill": {
     "duration": 419.520737,
     "end_time": "2024-03-23T20:55:07.278207",
     "exception": false,
     "start_time": "2024-03-23T20:48:07.757470",
     "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_214812-zpv7qa6m\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[33mmajor-sunset-208\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/zpv7qa6m\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='2954' max='2954' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [2954/2954 06:40, 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>10.900800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>200</td>\n",
       "      <td>7.522400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>300</td>\n",
       "      <td>6.429900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>400</td>\n",
       "      <td>5.989900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>500</td>\n",
       "      <td>5.740400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>600</td>\n",
       "      <td>5.528000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>700</td>\n",
       "      <td>5.313400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>800</td>\n",
       "      <td>5.220800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>900</td>\n",
       "      <td>5.102600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1000</td>\n",
       "      <td>5.031800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1100</td>\n",
       "      <td>5.025600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1200</td>\n",
       "      <td>4.886800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1300</td>\n",
       "      <td>4.827500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1400</td>\n",
       "      <td>4.839700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1500</td>\n",
       "      <td>4.762400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1600</td>\n",
       "      <td>4.690500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1700</td>\n",
       "      <td>4.663900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1800</td>\n",
       "      <td>4.635100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1900</td>\n",
       "      <td>4.664100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2000</td>\n",
       "      <td>4.580400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2100</td>\n",
       "      <td>4.481000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2200</td>\n",
       "      <td>4.525600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2300</td>\n",
       "      <td>4.505800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2400</td>\n",
       "      <td>4.496100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2500</td>\n",
       "      <td>4.519800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2600</td>\n",
       "      <td>4.482800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2700</td>\n",
       "      <td>4.484200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2800</td>\n",
       "      <td>4.465100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2900</td>\n",
       "      <td>4.469000</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 linear_probe_test/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 linear_probe_test/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 linear_probe_test/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 linear_probe_test/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 linear_probe_test/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"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=2954, training_loss=5.186163890014731, metrics={'train_runtime': 415.5038, 'train_samples_per_second': 56.864, 'train_steps_per_second': 7.109, 'total_flos': 4454006226937344.0, 'train_loss': 5.186163890014731, 'epoch': 1.0})"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "args = TrainingArguments(\n",
    "    output_dir=\"linear_probe_test\",\n",
    "    learning_rate=0.0002,\n",
    "    num_train_epochs=train_epochs,\n",
    "    logging_steps=logging_steps,\n",
    "    do_eval=False,\n",
    "    remove_unused_columns=False,  # Important to set to False, otherwise things will fail\n",
    ")\n",
    "collator = DataCollatorWithPadding(\n",
    "    feature_name_to_padding_value={\n",
    "        \"input_ids\": tokenizer.pad_token_id,\n",
    "        heads_configs[0].name: -100,\n",
    "        \"attention_mask\": 0,\n",
    "    }\n",
    ")\n",
    "trainer = Trainer(\n",
    "    model,\n",
    "    args=args,\n",
    "    train_dataset=dd[\"train\"],\n",
    "    data_collator=collator,\n",
    ")\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d83b7f4d",
   "metadata": {
    "papermill": {
     "duration": 0.005657,
     "end_time": "2024-03-23T20:55:07.290960",
     "exception": false,
     "start_time": "2024-03-23T20:55:07.285303",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "So this is nice to see, the probe is learning something and the training loss decreases. But how about evaluation on the validation set?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "79a8227e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-23T20:55:07.303299Z",
     "iopub.status.busy": "2024-03-23T20:55:07.303077Z",
     "iopub.status.idle": "2024-03-23T20:55:53.067372Z",
     "shell.execute_reply": "2024-03-23T20:55:53.066686Z"
    },
    "papermill": {
     "duration": 45.772106,
     "end_time": "2024-03-23T20:55:53.068574",
     "exception": false,
     "start_time": "2024-03-23T20:55:07.296468",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r\n",
      "Evaluating: 100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 308/308 [00:45<00:00,  6.74it/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",
      "\r\n",
      "Evaluating:   0%|          | 1/308 [00:00<00:42,  7.20it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(4.339219330967246, {'wikitext_head': 4.339219330967246})\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "print(evaluate_head_wise(model, dd[\"validation\"], collator, epochs=eval_epochs))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be87d2d2",
   "metadata": {
    "papermill": {
     "duration": 0.017376,
     "end_time": "2024-03-23T20:55:53.104048",
     "exception": false,
     "start_time": "2024-03-23T20:55:53.086672",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Yep, the evaluation loss is similar to the training loss, indicating no overfitting. How do things look in a practical example?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "276c547e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-23T20:55:53.142209Z",
     "iopub.status.busy": "2024-03-23T20:55:53.141859Z",
     "iopub.status.idle": "2024-03-23T20:55:53.170856Z",
     "shell.execute_reply": "2024-03-23T20:55:53.170377Z"
    },
    "papermill": {
     "duration": 0.050597,
     "end_time": "2024-03-23T20:55:53.171996",
     "exception": false,
     "start_time": "2024-03-23T20:55:53.121399",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'wikitext_head': [' the', ' this', ' his', ' its', ' these']}\n"
     ]
    }
   ],
   "source": [
    "print(get_top_n_preds(5, model, \"The historical significance of\", tokenizer))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f2d06f9",
   "metadata": {
    "papermill": {
     "duration": 0.017585,
     "end_time": "2024-03-23T20:55:53.207396",
     "exception": false,
     "start_time": "2024-03-23T20:55:53.189811",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "We see that the linear probe has learned to predict tokens that are pretty likely to follow that sentence."
   ]
  }
 ],
 "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": 536.439992,
   "end_time": "2024-03-23T20:55:55.964259",
   "environment_variables": {},
   "exception": null,
   "input_path": "notebooks/base/linear_probe.ipynb",
   "output_path": "notebooks/gpt2/linear_probe.ipynb",
   "parameters": {
    "eval_epochs": 1,
    "logging_steps": 100,
    "model_path": "gpt2",
    "train_epochs": 1
   },
   "start_time": "2024-03-23T20:46:59.524267",
   "version": "2.5.0"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {
     "03d42341242c49fb9053f18e7eeb0b85": {
      "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_23ffaad2326e4dd380556ff8a5f2540c",
       "max": 23627,
       "min": 0,
       "orientation": "horizontal",
       "style": "IPY_MODEL_db1643a32d7846d380f13ba510f528ec",
       "tabbable": null,
       "tooltip": null,
       "value": 23627
      }
     },
     "23ffaad2326e4dd380556ff8a5f2540c": {
      "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
      }
     },
     "2f7c7d3674ce49d9aced54bd7107cfb4": {
      "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
      }
     },
     "2fadd86620504fb3bef7b4fab0a4a065": {
      "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
      }
     },
     "498fa0330b0846fd8bebe10ca1fc9a55": {
      "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_2fadd86620504fb3bef7b4fab0a4a065",
       "placeholder": "\u200b",
       "style": "IPY_MODEL_80b505178aef4bb3835f85f8b63703ca",
       "tabbable": null,
       "tooltip": null,
       "value": "\u200723627/23627\u2007[00:02&lt;00:00,\u200710248.94\u2007examples/s]"
      }
     },
     "58396c007b99483cbe0e1d87a45f5308": {
      "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
      }
     },
     "7cbe61a9403f4939a36dd1a6a1d9a5aa": {
      "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
      }
     },
     "80b505178aef4bb3835f85f8b63703ca": {
      "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
      }
     },
     "db1643a32d7846d380f13ba510f528ec": {
      "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": ""
      }
     },
     "e1870d6c6b1947cbbf6c1767d7d41ba4": {
      "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_f31652bba9b042d7a78f3c85d7eedfea",
        "IPY_MODEL_03d42341242c49fb9053f18e7eeb0b85",
        "IPY_MODEL_498fa0330b0846fd8bebe10ca1fc9a55"
       ],
       "layout": "IPY_MODEL_2f7c7d3674ce49d9aced54bd7107cfb4",
       "tabbable": null,
       "tooltip": null
      }
     },
     "f31652bba9b042d7a78f3c85d7eedfea": {
      "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_58396c007b99483cbe0e1d87a45f5308",
       "placeholder": "\u200b",
       "style": "IPY_MODEL_7cbe61a9403f4939a36dd1a6a1d9a5aa",
       "tabbable": null,
       "tooltip": null,
       "value": "Map:\u2007100%"
      }
     }
    },
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}