fine_tuning_local.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kLB3I4FKZ5Lr"
   },
   "source": [
    "# Fine-tuning BERT for GitHub issue label prediction on a local installation\n",
    "\n",
    "Based on: [Fine-tuning BERT (and friends) for multi-label text classification](https://github.com/NielsRogge/Transformers-Tutorials/blob/master/BERT/Fine_tuning_BERT_(and_friends)_for_multi_label_text_classification.ipynb)\n",
    "Adapted with: [Multi-Label Classification Model From Scratch: Step-by-Step Tutorial](https://huggingface.co/blog/Valerii-Knowledgator/multi-label-classification)\n",
    "\n",
    "This is just a local exploration of the whole ML training space. The script will be migrated to be runnable inside the kubeflow installation. It's mostly copy paste from above with a few adaptions for handling the custom dataset. \n",
    "\n",
    "*\"In this notebook, we are going to fine-tune BERT to predict one or more labels for a given piece of text. Note that this notebook illustrates how to fine-tune a bert-base-uncased model, but you can also fine-tune a RoBERTa, DeBERTa, DistilBERT, CANINE, ... checkpoint in the same way.* \n",
    "\n",
    "*All of those work in the same way: they add a linear layer on top of the base model, which is used to produce a tensor of shape (batch_size, num_labels), indicating the unnormalized scores for a number of labels for every example in the batch.\"*\n",
    "\n",
    "## Set-up environment\n",
    "\n",
    "First, we install the libraries which we'll use: HuggingFace Transformers and Datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "4wxY3x-ZZz8h"
   },
   "outputs": [],
   "source": [
    "!pip install -q transformers datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bIH9NP0MZ6-O"
   },
   "source": [
    "## Load dataset\n",
    "\n",
    "Here we load the previously prepared dataset. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 86,
     "referenced_widgets": [
      "9a1aa9f2cc29473f9f8e5459d2641e76",
      "308fa6a7348140ec981a8d6c7d31f346",
      "8bc92587e35443488445e7521fbd0a13",
      "091f8220f33241f288faa0612853585f",
      "1045bb16e3694410898a73cf1b848917",
      "cb95e545fbdd4e99903bf634df694c9f",
      "5080d322a8034924b652b379c04667ed",
      "8b1899a0c4b144d7a5e6599f8afb8b65",
      "6111a73e684a47769bda7183a836ee91",
      "cd3570ddf67541d7818d97e236c54e54",
      "bbba60f793c14100934a268063f63d26"
     ]
    },
    "id": "sd1LiXGjZ420",
    "outputId": "1b5783cd-0e4e-4c92-c67c-57b9288f2381"
   },
   "outputs": [],
   "source": [
    "from datasets import load_dataset, DatasetDict\n",
    "\n",
    "dataset = load_dataset(\"json\", data_files=\"prepared-issues-reduced.json\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "pRd1kXQZjYIY",
    "outputId": "8021590c-5607-4a9c-a474-bc57da503c93"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['title', 'body', 'enhancement', 'bug'],\n",
       "        num_rows: 359\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['title', 'body', 'enhancement', 'bug'],\n",
       "        num_rows: 287\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['title', 'body', 'enhancement', 'bug'],\n",
       "        num_rows: 72\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds = dataset[\"train\"].train_test_split(test_size=0.2)\n",
    "ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "unjuTtKUjZI3",
    "outputId": "6f1e5051-8272-40f9-ced8-ba17a105e904"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'title': 'CLI: Create namespace if necessary when installing a namespace-scoped Package',\n",
       " 'body': \"**Is your feature request related to a problem? Please describe.**\\r\\nInstalling a namespace-scoped package can be more complicated than necessary if the target namespace does not exist yet. In that case, a user has to resort to `kubectl` or other means to create the namespace before being able to install the package. This should not be necessary.\\r\\n\\r\\n**Describe the solution you'd like**\\r\\nWhen installing a namespace-scoped package, the CLI should check if the target namespace exists and, if it does not, create it after the user has confirmed the install operation (if no `--yes` flag is present) and before the `Package` is created. The installation summary should inform the user that the namespace is going to be created. Example output:\\r\\n\\r\\n```\\r\\n$ glasskube install quickwit -n test\\r\\nSummary:\\r\\n * The following packages will be installed in your cluster (minikube):\\r\\n    1. quickwit of type quickwit in namespace test (version v0.8.1+5)\\r\\n * Namespace test does not exist and will be created\\r\\n * Automatic updates will be not enabled\\r\\n```\\r\\n\\r\\n**Additional context**\\r\\nFor context, see #1084\",\n",
       " 'enhancement': True,\n",
       " 'bug': False}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "example = ds['train'][0]\n",
    "example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6DV0Rtetxgd4"
   },
   "source": [
    "Let's create a list that contains the labels, as well as 2 dictionaries that map labels to integers and back."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "e5vZhQpvkE8s",
    "outputId": "5d513b30-f209-492f-c6ab-245d64a67d40"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['enhancement', 'bug']"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "labels = [label for label in dataset['train'].features.keys() if label not in ['body', 'title']]\n",
    "id2label = {idx:label for idx, label in enumerate(labels)}\n",
    "label2id = {label:idx for idx, label in enumerate(labels)}\n",
    "labels"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nJ3Teyjmank2"
   },
   "source": [
    "## Preprocess data\n",
    "\n",
    "*As models like BERT don't expect text as direct input, but rather `input_ids`, etc., we tokenize the text using the tokenizer. Here I'm using the `AutoTokenizer` API, which will automatically load the appropriate tokenizer based on the checkpoint on the hub.*\n",
    "\n",
    "*What's a bit tricky is that we also need to provide labels to the model. For multi-label text classification, this is a matrix of shape (batch_size, num_labels). Also important: this should be a tensor of floats rather than integers, otherwise PyTorch' `BCEWithLogitsLoss` (which the model will use) will complain, as explained [here](https://discuss.pytorch.org/t/multi-label-binary-classification-result-type-float-cant-be-cast-to-the-desired-output-type-long/117915/3).*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "AFWlSsbZaRLc"
   },
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer\n",
    "import numpy as np\n",
    "import math\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"bert-base-uncased\")\n",
    "\n",
    "def preprocess_data(example):\n",
    "  text = f'{example[\"title\"]}\\n{example[\"body\"]}'\n",
    "  # encode them\n",
    "  encoding = tokenizer(text, padding=True, truncation=True)\n",
    "\n",
    "  lbls = [0. for i in range(len(labels))]\n",
    "  for label in labels:\n",
    "    if label in example and example[label] == True:\n",
    "      label_id = label2id[label]\n",
    "      lbls[label_id] = 1.\n",
    "\n",
    "  encoding[\"labels\"] = lbls  \n",
    "  return encoding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "i4ENBTdulBEI",
    "outputId": "02554a1f-4961-461a-bf29-555b8debeabf"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Map: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████| 287/287 [00:00<00:00, 3302.66 examples/s]\n",
      "Map: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████| 72/72 [00:00<00:00, 2996.59 examples/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['input_ids', 'token_type_ids', 'attention_mask', 'labels'],\n",
       "        num_rows: 287\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['input_ids', 'token_type_ids', 'attention_mask', 'labels'],\n",
       "        num_rows: 72\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "encoded_dataset = ds.map(preprocess_data, remove_columns=ds['train'].column_names)\n",
    "encoded_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "0enAb0W9o25W",
    "outputId": "55bc5ba6-d169-49c6-f562-bb7ea4143866"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dict_keys(['input_ids', 'token_type_ids', 'attention_mask', 'labels'])\n"
     ]
    }
   ],
   "source": [
    "example = encoded_dataset['train'][0]\n",
    "print(example.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 137
    },
    "id": "D0McCtJ8HRJY",
    "outputId": "82fb0336-51a3-40ad-ebc0-65eeb7cf4b6c"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"[CLS] cli : create namespace if necessary when installing a namespace - scoped package * * is your feature request related to a problem? please describe. * * installing a namespace - scoped package can be more complicated than necessary if the target namespace does not exist yet. in that case, a user has to resort to ` kubectl ` or other means to create the namespace before being able to install the package. this should not be necessary. * * describe the solution you ' d like * * when installing a namespace - scoped package, the cli should check if the target namespace exists and, if it does not, create it after the user has confirmed the install operation ( if no ` - - yes ` flag is present ) and before the ` package ` is created. the installation summary should inform the user that the namespace is going to be created. example output : ` ` ` $ glasskube install quickwit - n test summary : * the following packages will be installed in your cluster ( minikube ) : 1. quickwit of type quickwit in namespace test ( version v0. 8. 1 + 5 ) * namespace test does not exist and will be created * automatic updates will be not enabled ` ` ` * * additional context * * for context, see # 1084 [SEP]\""
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer.decode(example['input_ids'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "VdIvj6WjHeZQ",
    "outputId": "418c14d2-cca3-44e9-d0a2-6ad4ca7a007c",
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[1.0, 0.0]"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "example['labels']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "q4Dx95t2o6N9",
    "outputId": "3ce6c923-0b45-4743-bc4b-7771d088b03e"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['enhancement']"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[id2label[idx] for idx, label in enumerate(example['labels']) if label == 1.0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HgpKXDfvKBxn"
   },
   "source": [
    "Finally, we set the format of our data to PyTorch tensors. This will turn the training, validation and test sets into standard PyTorch [datasets](https://pytorch.org/docs/stable/data.html). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "id": "Lk6Cq9duKBkA"
   },
   "outputs": [],
   "source": [
    "encoded_dataset.set_format(\"torch\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "w5qSmCgWefWs"
   },
   "source": [
    "## Define model\n",
    "\n",
    "*Here we define a model that includes a pre-trained base (i.e. the weights from bert-base-uncased) are loaded, with a random initialized classification head (linear layer) on top. One should fine-tune this head, together with the pre-trained base on a labeled dataset.*\n",
    "\n",
    "*This is also printed by the warning.*\n",
    "\n",
    "*We set the `problem_type` to be \"multi_label_classification\", as this will make sure the appropriate loss function is used (namely [`BCEWithLogitsLoss`](https://pytorch.org/docs/stable/generated/torch.nn.BCEWithLogitsLoss.html)). We also make sure the output layer has `len(labels)` output neurons, and we set the id2label and label2id mappings.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "6XPL1Z_RegBF",
    "outputId": "22994300-8c93-421e-faa4-678d6cc14aab"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of BertForSequenceClassification were not initialized from the model checkpoint at bert-base-uncased and are newly initialized: ['classifier.bias', 'classifier.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": [
    "from transformers import AutoModelForSequenceClassification\n",
    "\n",
    "model = AutoModelForSequenceClassification.from_pretrained(\"bert-base-uncased\", \n",
    "                                                           problem_type=\"multi_label_classification\", \n",
    "                                                           num_labels=len(labels),\n",
    "                                                           id2label=id2label,\n",
    "                                                           label2id=label2id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mjJGEXShp7te"
   },
   "source": [
    "## Train the model!\n",
    "\n",
    "We are going to train the model using HuggingFace's Trainer API. This requires us to define 2 things: \n",
    "\n",
    "* `TrainingArguments`, which specify training hyperparameters. All options can be found in the [docs](https://huggingface.co/transformers/main_classes/trainer.html#trainingarguments). Below, we for example specify that we want to evaluate after every epoch of training, we would like to save the model every epoch, we set the learning rate, the batch size to use for training/evaluation, how many epochs to train for, and so on.\n",
    "* a `Trainer` object (docs can be found [here](https://huggingface.co/transformers/main_classes/trainer.html#id1))."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "id": "K5a8_vIKqr7P"
   },
   "outputs": [],
   "source": [
    "batch_size = 3\n",
    "metric_name = \"f1\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "dR2GmpvDqbuZ"
   },
   "outputs": [],
   "source": [
    "from transformers import TrainingArguments, Trainer\n",
    "\n",
    "args = TrainingArguments(\n",
    "    \"bert-finetuned-trainer-output\",\n",
    "    evaluation_strategy = \"epoch\",\n",
    "    save_strategy = \"epoch\",\n",
    "    learning_rate=2e-5,\n",
    "    per_device_train_batch_size=batch_size,\n",
    "    per_device_eval_batch_size=batch_size,\n",
    "    num_train_epochs=2,\n",
    "    weight_decay=0.01,\n",
    "    load_best_model_at_end=True,\n",
    "    metric_for_best_model=metric_name,\n",
    "    #push_to_hub=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1_v2fPFFJ3-v"
   },
   "source": [
    "We are also going to compute metrics while training. For this, we need to define a `compute_metrics` function, that returns a dictionary with the desired metric values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "id": "797b2WHJqUgZ"
   },
   "outputs": [],
   "source": [
    "import evaluate\n",
    "import numpy as np\n",
    "\n",
    "clf_metrics = evaluate.combine([\"accuracy\", \"f1\", \"precision\", \"recall\"])\n",
    "\n",
    "def sigmoid(x):\n",
    "   return 1/(1 + np.exp(-x))\n",
    "\n",
    "def compute_metrics(eval_pred):\n",
    "\n",
    "   predictions, labels = eval_pred\n",
    "   predictions = sigmoid(predictions)\n",
    "   predictions = (predictions > 0.5).astype(int).reshape(-1)\n",
    "   return clf_metrics.compute(predictions=predictions, references=labels.astype(int).reshape(-1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fxNo4_TsvzDm"
   },
   "source": [
    "Let's verify a batch as well as a forward pass:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "id": "IlOgGiojuWwG",
    "outputId": "cd0b2c99-b520-468d-8ffc-c36211b7820a"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'torch.FloatTensor'"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "encoded_dataset['train'][0]['labels'].type()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Y41Kre_jvD7x",
    "outputId": "b6ca888b-6371-40fb-ab83-3dc24d28320a"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([  101, 18856,  2072,  1024,  3443,  3415, 15327,  2065,  4072,  2043,\n",
       "        23658,  1037,  3415, 15327,  1011,  9531,  2094,  7427,  1008,  1008,\n",
       "         2003,  2115,  3444,  5227,  3141,  2000,  1037,  3291,  1029,  3531,\n",
       "         6235,  1012,  1008,  1008, 23658,  1037,  3415, 15327,  1011,  9531,\n",
       "         2094,  7427,  2064,  2022,  2062,  8552,  2084,  4072,  2065,  1996,\n",
       "         4539,  3415, 15327,  2515,  2025,  4839,  2664,  1012,  1999,  2008,\n",
       "         2553,  1010,  1037,  5310,  2038,  2000,  7001,  2000,  1036, 13970,\n",
       "         4783,  6593,  2140,  1036,  2030,  2060,  2965,  2000,  3443,  1996,\n",
       "         3415, 15327,  2077,  2108,  2583,  2000, 16500,  1996,  7427,  1012,\n",
       "         2023,  2323,  2025,  2022,  4072,  1012,  1008,  1008,  6235,  1996,\n",
       "         5576,  2017,  1005,  1040,  2066,  1008,  1008,  2043, 23658,  1037,\n",
       "         3415, 15327,  1011,  9531,  2094,  7427,  1010,  1996, 18856,  2072,\n",
       "         2323,  4638,  2065,  1996,  4539,  3415, 15327,  6526,  1998,  1010,\n",
       "         2065,  2009,  2515,  2025,  1010,  3443,  2009,  2044,  1996,  5310,\n",
       "         2038,  4484,  1996, 16500,  3169,  1006,  2065,  2053,  1036,  1011,\n",
       "         1011,  2748,  1036,  5210,  2003,  2556,  1007,  1998,  2077,  1996,\n",
       "         1036,  7427,  1036,  2003,  2580,  1012,  1996,  8272, 12654,  2323,\n",
       "        12367,  1996,  5310,  2008,  1996,  3415, 15327,  2003,  2183,  2000,\n",
       "         2022,  2580,  1012,  2742,  6434,  1024,  1036,  1036,  1036,  1002,\n",
       "         3221,  5283,  4783, 16500,  4248,  9148,  2102,  1011,  1050,  3231,\n",
       "        12654,  1024,  1008,  1996,  2206, 14555,  2097,  2022,  5361,  1999,\n",
       "         2115,  9324,  1006,  7163,  5283,  4783,  1007,  1024,  1015,  1012,\n",
       "         4248,  9148,  2102,  1997,  2828,  4248,  9148,  2102,  1999,  3415,\n",
       "        15327,  3231,  1006,  2544,  1058,  2692,  1012,  1022,  1012,  1015,\n",
       "         1009,  1019,  1007,  1008,  3415, 15327,  3231,  2515,  2025,  4839,\n",
       "         1998,  2097,  2022,  2580,  1008,  6882, 14409,  2097,  2022,  2025,\n",
       "         9124,  1036,  1036,  1036,  1008,  1008,  3176,  6123,  1008,  1008,\n",
       "         2005,  6123,  1010,  2156,  1001, 10715,  2549,   102])"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "encoded_dataset['train']['input_ids'][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "sxWcnZ8ku12V",
    "outputId": "26522911-c3cd-466a-ae2d-d81d23003c23"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SequenceClassifierOutput(loss=tensor(0.7555, grad_fn=<BinaryCrossEntropyWithLogitsBackward0>), logits=tensor([[-0.3859, -0.1816]], grad_fn=<AddmmBackward0>), hidden_states=None, attentions=None)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#forward pass\n",
    "outputs = model(input_ids=encoded_dataset['train']['input_ids'][0].unsqueeze(0), labels=encoded_dataset['train'][0]['labels'].unsqueeze(0))\n",
    "outputs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "f-X2brZcv0X6"
   },
   "source": [
    "Let's start training!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "id": "chq_3nUz73ib"
   },
   "outputs": [],
   "source": [
    "trainer = Trainer(\n",
    "    model,\n",
    "    args,\n",
    "    train_dataset=encoded_dataset[\"train\"],\n",
    "    eval_dataset=encoded_dataset[\"test\"],\n",
    "    tokenizer=tokenizer,\n",
    "    compute_metrics=compute_metrics\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "KXmFds8js6P8",
    "outputId": "66ebb2ab-f93f-48aa-a4dc-36f55f2f8559"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='192' max='192' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [192/192 05:02, Epoch 2/2]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Epoch</th>\n",
       "      <th>Training Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "      <th>Accuracy</th>\n",
       "      <th>F1</th>\n",
       "      <th>Precision</th>\n",
       "      <th>Recall</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>No log</td>\n",
       "      <td>0.362327</td>\n",
       "      <td>0.875000</td>\n",
       "      <td>0.823529</td>\n",
       "      <td>0.724138</td>\n",
       "      <td>0.954545</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>No log</td>\n",
       "      <td>0.281655</td>\n",
       "      <td>0.937500</td>\n",
       "      <td>0.905263</td>\n",
       "      <td>0.843137</td>\n",
       "      <td>0.977273</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=192, training_loss=0.4105148712793986, metrics={'train_runtime': 302.937, 'train_samples_per_second': 1.895, 'train_steps_per_second': 0.634, 'total_flos': 80538191267940.0, 'train_loss': 0.4105148712793986, 'epoch': 2.0})"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hiloh9eMK91o"
   },
   "source": [
    "## Evaluate\n",
    "\n",
    "After training, we evaluate our model on the validation set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 224
    },
    "id": "cMlebJ83LRYG",
    "outputId": "b18102e7-2198-4beb-c874-d39636f740ed"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='24' max='24' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [24/24 00:07]\n",
       "    </div>\n",
       "    "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "{'eval_loss': 0.28165486454963684,\n",
       " 'eval_accuracy': 0.9375,\n",
       " 'eval_f1': 0.9052631578947369,\n",
       " 'eval_precision': 0.8431372549019608,\n",
       " 'eval_recall': 0.9772727272727273,\n",
       " 'eval_runtime': 7.7696,\n",
       " 'eval_samples_per_second': 9.267,\n",
       " 'eval_steps_per_second': 3.089,\n",
       " 'epoch': 2.0}"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.evaluate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "# trained model can be saved locally\n",
    "trainer.save_model(\"local-bert-reduced\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3nmvJp0pLq-3"
   },
   "source": [
    "## Inference\n",
    "\n",
    "Let's test the model on new issues:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "id": "3fxjfr8PLD42"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[{'label': 'bug', 'score': 0.7489847540855408}, {'label': 'enhancement', 'score': 0.3993261158466339}]]\n",
      "[[{'label': 'enhancement', 'score': 0.8171513080596924}, {'label': 'bug', 'score': 0.06645093113183975}]]\n"
     ]
    }
   ],
   "source": [
    "from transformers import BertModel, pipeline\n",
    "\n",
    "# load existing fine-tuned model\n",
    "nlp = pipeline(\"sentiment-analysis\", model=\"./local-bert-reduced\", top_k=5)\n",
    "\n",
    "bug = \"\"\"\n",
    "**Describe the bug**\n",
    "\n",
    "**To reproduce**\n",
    "\n",
    "**Cluster Info (please complete the following information):**\n",
    "\"\"\"\n",
    "\n",
    "enhancement = \"\"\"\n",
    "**Is your feature request related to a problem? Please describe.**\n",
    "\n",
    "**Describe the solution you'd like**\n",
    "\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "print(nlp(bug))\n",
    "print(nlp(enhancement))"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "authorship_tag": "ABX9TyMiblo1Ci0GTlAbbA5wB3mn",
   "collapsed_sections": [],
   "include_colab_link": true,
   "name": "Fine-tuning BERT (and friends) for multi-label text classification.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12.7"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "091f8220f33241f288faa0612853585f": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_6111a73e684a47769bda7183a836ee91",
      "max": 3,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_8b1899a0c4b144d7a5e6599f8afb8b65",
      "value": 3
     }
    },
    "1045bb16e3694410898a73cf1b848917": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_bbba60f793c14100934a268063f63d26",
      "placeholder": "​",
      "style": "IPY_MODEL_cd3570ddf67541d7818d97e236c54e54",
      "value": " 3/3 [00:00&lt;00:00, 75.93it/s]"
     }
    },
    "308fa6a7348140ec981a8d6c7d31f346": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": 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,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "5080d322a8034924b652b379c04667ed": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": 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,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "6111a73e684a47769bda7183a836ee91": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": 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,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "8b1899a0c4b144d7a5e6599f8afb8b65": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "8bc92587e35443488445e7521fbd0a13": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_5080d322a8034924b652b379c04667ed",
      "placeholder": "​",
      "style": "IPY_MODEL_cb95e545fbdd4e99903bf634df694c9f",
      "value": "100%"
     }
    },
    "9a1aa9f2cc29473f9f8e5459d2641e76": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_8bc92587e35443488445e7521fbd0a13",
       "IPY_MODEL_091f8220f33241f288faa0612853585f",
       "IPY_MODEL_1045bb16e3694410898a73cf1b848917"
      ],
      "layout": "IPY_MODEL_308fa6a7348140ec981a8d6c7d31f346"
     }
    },
    "bbba60f793c14100934a268063f63d26": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": 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,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "cb95e545fbdd4e99903bf634df694c9f": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "cd3570ddf67541d7818d97e236c54e54": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}