preprocessing_and_training.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "24c67bfc-eab4-4e10-b0a3-c62460d3c4bb",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-10-03 18:22:00.845555: I tensorflow/core/util/port.cc:110] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n",
      "2023-10-03 18:22:01.147946: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2023-10-03 18:22:02.461368: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n"
     ]
    }
   ],
   "source": [
    "from transformers import BertTokenizer, BertForQuestionAnswering, AdamW, AutoTokenizer, AutoModelForSequenceClassification\n",
    "from torch.utils.data import DataLoader\n",
    "import torch\n",
    "from datasets import load_dataset\n",
    "from tqdm import tqdm\n",
    "from transformers import get_linear_schedule_with_warmup\n",
    "import time\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "5a608b47-0e03-408c-ad80-000d6c46c5ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a49ef747-58a0-400e-9e14-24662502fff9",
   "metadata": {},
   "source": [
    "# Dataset is made up of ID, Title, Context, Question, and Answers\n",
    "Here we load in the dataset, now we are going to apply the autotokenizer to it.\n",
    "\n",
    "If you want to reduct training time, create smaller datasets of the existing one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "067f305f-40ec-4537-beea-1c23524301a1",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Found cached dataset squad (/home/mhrrs/.cache/huggingface/datasets/squad/plain_text/1.0.0/d6ec3ceb99ca480ce37cdd35555d6cb2511d223b9150cce08a837ef62ffea453)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dc9b9ac6264540ebb4739ed5c8507f5c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'id': '5733be284776f41900661182', 'title': 'University_of_Notre_Dame', 'context': 'Architecturally, the school has a Catholic character. Atop the Main Building\\'s gold dome is a golden statue of the Virgin Mary. Immediately in front of the Main Building and facing it, is a copper statue of Christ with arms upraised with the legend \"Venite Ad Me Omnes\". Next to the Main Building is the Basilica of the Sacred Heart. Immediately behind the basilica is the Grotto, a Marian place of prayer and reflection. It is a replica of the grotto at Lourdes, France where the Virgin Mary reputedly appeared to Saint Bernadette Soubirous in 1858. At the end of the main drive (and in a direct line that connects through 3 statues and the Gold Dome), is a simple, modern stone statue of Mary.', 'question': 'To whom did the Virgin Mary allegedly appear in 1858 in Lourdes France?', 'answers': {'text': ['Saint Bernadette Soubirous'], 'answer_start': [515]}}\n"
     ]
    }
   ],
   "source": [
    "#Load SQuAD Data\n",
    "dataset = load_dataset(\"squad\")\n",
    "\n",
    "# example output:\n",
    "print(dataset[\"train\"][0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07b5a352-89ad-4f23-9bad-108218fb3add",
   "metadata": {},
   "source": [
    "train_dataloader currently contains: ['id', 'title', 'context', 'question', 'answers', 'input_ids', 'token_type_ids', 'attention_mask']\n",
    "- It needs to contain start and end positions as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "276719e2-1c5c-485d-9a8d-fee1a77f2832",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at /home/mhrrs/.cache/huggingface/datasets/squad/plain_text/1.0.0/d6ec3ceb99ca480ce37cdd35555d6cb2511d223b9150cce08a837ef62ffea453/cache-94cdc3ae0dfac08b.arrow\n",
      "Loading cached processed dataset at /home/mhrrs/.cache/huggingface/datasets/squad/plain_text/1.0.0/d6ec3ceb99ca480ce37cdd35555d6cb2511d223b9150cce08a837ef62ffea453/cache-64af2f1d98310d39.arrow\n"
     ]
    }
   ],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')\n",
    "\n",
    "def tokenize_function(examples):\n",
    "    return tokenizer(examples[\"context\"], padding=\"max_length\", stride=128, truncation=True)\n",
    "\n",
    "tokenized_datasets = dataset.map(tokenize_function, batched=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "66a5f1ab-020c-4ad0-8585-24fa8650b669",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached shuffled indices for dataset at /home/mhrrs/.cache/huggingface/datasets/squad/plain_text/1.0.0/d6ec3ceb99ca480ce37cdd35555d6cb2511d223b9150cce08a837ef62ffea453/cache-3d9ddc5e3cf60fe4.arrow\n",
      "Loading cached shuffled indices for dataset at /home/mhrrs/.cache/huggingface/datasets/squad/plain_text/1.0.0/d6ec3ceb99ca480ce37cdd35555d6cb2511d223b9150cce08a837ef62ffea453/cache-09c923ccbfed43a7.arrow\n"
     ]
    }
   ],
   "source": [
    "# Remove this and use the full set later if you have time to train it\n",
    "train_dataset = tokenized_datasets[\"train\"].shuffle(seed=42)\n",
    "eval_dataset = tokenized_datasets[\"validation\"].shuffle(seed=42)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cdea3cfb-12bf-4970-b9b1-3108d41fdb8a",
   "metadata": {},
   "source": [
    "The start index of the answer is already in the \"answers\" dictionary. Now we just add the length of the \"text\" in the \"answers\" dict to get the \"answer_end\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "48a12721-d434-45c0-88dd-566668e097f1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "contexts: 87599\n",
      "questions: 87599\n",
      "answers: 87599\n"
     ]
    }
   ],
   "source": [
    "# process training and validation data into seperate groups\n",
    "def package_data(dataset):\n",
    "    contexts = []\n",
    "    questions = []\n",
    "    answers = []\n",
    "    \n",
    "    for item in dataset:\n",
    "        contexts.append(item[\"context\"])\n",
    "        questions.append(item[\"question\"])\n",
    "        answers.append(item[\"answers\"])\n",
    "            \n",
    "    return contexts, questions, answers\n",
    "        \n",
    "        \n",
    "train_contexts, train_questions, train_answers = package_data(train_dataset)\n",
    "eval_contexts, eval_questions, eval_answers = package_data(eval_dataset)\n",
    "\n",
    "print(f\"contexts: {len(train_contexts)}\\nquestions: {len(train_questions)}\\nanswers: {len(train_answers)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d3b539dc-a9e8-401f-ba02-0b8eb700c297",
   "metadata": {},
   "outputs": [],
   "source": [
    "# use the following function to add answer_end index to the train_answers/valid_answers data structure\n",
    "def add_end_index(dataset):\n",
    "    for ans in dataset:\n",
    "        text = ans[\"text\"][0]\n",
    "        start_index = ans[\"answer_start\"][0]\n",
    "        end_index = start_index + len(text) #adds the length of text onto the start index\n",
    "        ans[\"answer_start\"] = start_index\n",
    "        ans[\"answer_end\"] = end_index\n",
    "        \n",
    "add_end_index(train_answers)\n",
    "add_end_index(eval_answers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "474fb21f-1095-4b70-91f5-7ff0c00220a7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train_contexts = The Pew Forum on Religion & Public Life ranks Egypt as the fifth worst country in the world for religious freedom. The United States Commission on International Religious Freedom, a bipartisan independent agency of the US government, has placed Egypt on its watch list of countries that require close monitoring due to the nature and extent of violations of religious freedom engaged in or tolerated by the government. According to a 2010 Pew Global Attitudes survey, 84% of Egyptians polled supported the death penalty for those who leave Islam; 77% supported whippings and cutting off of hands for theft and robbery; and 82% support stoning a person who commits adultery.\n",
      "train_questions = What percentage of Egyptians polled support death penalty for those leaving Islam?\n",
      "train_answers = {'text': ['84%'], 'answer_start': 468, 'answer_end': 471}\n"
     ]
    }
   ],
   "source": [
    "print(f\"train_contexts = {train_contexts[0]}\\ntrain_questions = {train_questions[0]}\\ntrain_answers = {train_answers[0]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f838d307-3564-450a-b67e-512304aabf06",
   "metadata": {},
   "source": [
    "# Tokenize the contexts and questions together"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "1e22fd98-72bf-4b4e-8578-11068fb5a84c",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_encodings = tokenizer(train_contexts, train_questions, truncation=True, padding=True)\n",
    "eval_encodings = tokenizer(eval_contexts, eval_questions, truncation=True, padding=True)\n",
    "del train_contexts, train_questions, eval_contexts, eval_questions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "434c806d-7566-4ffc-b55c-ae1479007488",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'[CLS] the pew forum on religion & public life ranks egypt as the fifth worst country in the world for religious freedom. the united states commission on international religious freedom, a bipartisan independent agency of the us government, has placed egypt on its watch list of countries that require close monitoring due to the nature and extent of violations of religious freedom engaged in or tolerated by the government. according to a 2010 pew global attitudes survey, 84 % of egyptians polled supported the death penalty for those who leave islam ; 77 % supported whippings and cutting off of hands for theft and robbery ; and 82 % support stoning a person who commits adultery. [SEP] what percentage of egyptians polled support death penalty for those leaving islam? [SEP] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD]'"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer.decode(train_encodings['input_ids'][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "6f36ff48-bdd8-4c11-817e-d3a0a146612b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_token_positions(encodings, answers):\n",
    "    start_pos = []\n",
    "    end_pos = []\n",
    "    for i in range(len(answers)):\n",
    "        start_pos.append(encodings.char_to_token(i, answers[i][\"answer_start\"]))\n",
    "        end_pos.append(encodings.char_to_token(i, answers[i][\"answer_end\"]-1))\n",
    "        \n",
    "        if start_pos[-1] is None:\n",
    "            start_pos[-1] = tokenizer.model_max_length\n",
    "        if end_pos[-1] is None:\n",
    "            end_pos[-1] = tokenizer.model_max_length\n",
    "        \n",
    "    encodings.update({\"start_positions\": start_pos, \"end_positions\": end_pos})\n",
    "\n",
    "add_token_positions(train_encodings, train_answers)\n",
    "add_token_positions(eval_encodings, eval_answers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "16514293-b8f9-4d0c-8286-84f6ca54cdeb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dict_keys(['input_ids', 'token_type_ids', 'attention_mask', 'start_positions', 'end_positions'])\n",
      "108\n",
      "109\n"
     ]
    }
   ],
   "source": [
    "print(train_encodings.keys())\n",
    "print(f\"{train_encodings['start_positions'][2]}\\n{train_encodings['end_positions'][2]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f0382c8-a008-4cc0-a21f-996cefefa20b",
   "metadata": {},
   "source": [
    "# create DataLoaders from the preprocessed tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "5c0bbf2f-fbf8-43c1-b4be-2bc8ddaa2c1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SQuAD(torch.utils.data.Dataset):\n",
    "    def __init__(self, encodings):\n",
    "        self.encodings = encodings\n",
    "    def __getitem__(self, idx):\n",
    "        return {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}\n",
    "    def __len__(self):\n",
    "        return len(self.encodings.input_ids)\n",
    "    \n",
    "train_dataset = SQuAD(train_encodings)\n",
    "eval_dataset = SQuAD(eval_encodings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "281c00d8-c10c-4a4b-8148-71f8cad099d4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Encoding(num_tokens=512, attributes=[ids, type_ids, tokens, offsets, attention_mask, special_tokens_mask, overflowing])\n",
      "<torch.utils.data.dataloader.DataLoader object at 0x145de9caf880>\n"
     ]
    }
   ],
   "source": [
    "print(train_dataset.encodings[0])\n",
    "\n",
    "# insert converted tensor dataset into DataLoader function\n",
    "train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)\n",
    "eval_loader = DataLoader(train_dataset, batch_size=16)\n",
    "print(train_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "b7d0aeb6-49f9-48da-a0d6-f66233a430ba",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tokenized 84\n",
      "Real ['84%']\n",
      "Real 468\n",
      "Real 471\n"
     ]
    }
   ],
   "source": [
    "def show_answer(idx):\n",
    "    print(\"Tokenized\", tokenizer.decode(train_encodings['input_ids'][idx][train_encodings['start_positions'][idx]: train_encodings['end_positions'][idx]]))\n",
    "    print(\"Real\", train_answers[idx]['text'])\n",
    "    print(\"Real\", train_answers[idx]['answer_start'])\n",
    "    print(\"Real\", train_answers[idx]['answer_end'])\n",
    "    \n",
    "show_answer(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ecd073f6-7146-4f0f-a7fe-34787045bea9",
   "metadata": {},
   "source": [
    "# Create Training Loop\n",
    "Experimental at the moment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "18d9eec5-ab07-4a97-8303-6d602fcf0f0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model(train_loader):\n",
    "    n_epochs = 5\n",
    "    optim = AdamW(model.parameters(), lr=5e-5)\n",
    "    n_train_steps = len(train_loader) * n_epochs\n",
    "    n_warmup_steps = .1 * n_train_steps\n",
    "    scheduler = get_linear_schedule_with_warmup(optim, n_warmup_steps, n_t_steps)\n",
    "\n",
    "    model.to(device)\n",
    "    model.train()\n",
    "    for epoch in range(n_epochs):\n",
    "        loop = tqdm(train_loader, leave=True)\n",
    "        for batch in loop:\n",
    "            optim.zero_grad()\n",
    "            input_ids = batch['input_ids'].to(device)\n",
    "            attention_mask = batch['attention_mask'].to(device)\n",
    "            start_pos = batch['start_positions'].to(device)\n",
    "            end_pos = batch['end_positions'].to(device)\n",
    "            outputs = model(input_ids, attention_mask=attention_mask, start_positions=start_pos, end_positions=end_pos)\n",
    "            loss = outputs[0]\n",
    "            loss.backward()\n",
    "            optim.step()\n",
    "            scheduler.step()\n",
    "\n",
    "            loop.set_description(f\"Epoch {epoch+1}\")\n",
    "            loop.set_postfix(loss=loss.item())\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d38e5a6-6d07-44fb-ae8a-0af26f89710a",
   "metadata": {},
   "source": [
    "# Updated Training Loop\n",
    "The following training loop uses:\n",
    "- Automatic Mixed Precision\n",
    "- Linear Learning Rate Decay\n",
    "- Gradient Accumulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "bf39a5d2-e264-4dd8-8728-3d2396b1c596",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_amp_model(train_loader):\n",
    "    n_epochs = 5\n",
    "    optim = AdamW(model.parameters(), lr=5e-5)\n",
    "    \n",
    "    # linear learning rate decay setup\n",
    "    n_train_steps = len(train_loader) * n_epochs\n",
    "    n_warmup_steps = .1 * n_train_steps\n",
    "    scheduler = get_linear_schedule_with_warmup(optim, n_warmup_steps, n_t_steps)\n",
    "    \n",
    "    # automatic mixed precision setup\n",
    "    scaler = torch.cuda.amp.GradScaler()\n",
    "    \n",
    "    # other\n",
    "    total_loss = 0\n",
    "    total_time = 0.0\n",
    "\n",
    "    #batch accumulation parameter\n",
    "    accum_iter = 4\n",
    "    \n",
    "    model.to(device)\n",
    "    model.train()\n",
    "    for epoch in range(n_epochs):\n",
    "        start_time = time.perf_counter()\n",
    "        \n",
    "        for step, batch in enumerate(train_loader):\n",
    "            input_ids = batch['input_ids'].to(device)\n",
    "            attention_mask = batch['attention_mask'].to(device)\n",
    "            start_pos = batch['start_positions'].to(device)\n",
    "            end_pos = batch['end_positions'].to(device)\n",
    "            \n",
    "            # automatic mixed precision\n",
    "            with torch.cuda.amp.autocast():\n",
    "                outputs = model(input_ids, attention_mask=attention_mask, start_positions=start_pos, end_positions=end_pos)\n",
    "                loss = outputs[0]\n",
    "                \n",
    "            scaler.scale(loss).backward()\n",
    "            total_loss += loss.item()\n",
    "            \n",
    "            # Gradient accumulation implementation\n",
    "            if ((step+1)%accum_iter == 0) or (step+1 == len(train_loader)):\n",
    "                scaler.step(optim)\n",
    "                scaler.update()\n",
    "                optim.zero_grad()\n",
    "                scheduler.step()\n",
    "            \n",
    "            if step % 100 == 0:\n",
    "                epoch_time = time.perf_counter()\n",
    "                print(f\"Epoch: {epoch+1} | step: {step}/{len(train_loader)} | loss: {total_loss/(step+1):.4f} | time: {(epoch_time-start_time)/60:.1f} (minutes)\")\n",
    "                \n",
    "        total_time += (epoch_time-start_time)\n",
    "            \n",
    "    print(f\"Total time: {total_time/360} (hours)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "0d0cb2c6-3e76-4954-8cf8-43b4afa61afe",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at bert-base-uncased were not used when initializing BertForQuestionAnswering: ['cls.seq_relationship.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.decoder.weight', 'cls.predictions.bias', 'cls.predictions.transform.dense.weight', 'cls.predictions.transform.LayerNorm.bias', 'cls.predictions.transform.dense.bias', 'cls.seq_relationship.weight']\n",
      "- This IS expected if you are initializing BertForQuestionAnswering from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertForQuestionAnswering from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some weights of BertForQuestionAnswering were not initialized from the model checkpoint at bert-base-uncased and are newly initialized: ['qa_outputs.weight', 'qa_outputs.bias']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n",
      "/home/mhrrs/.local/lib/python3.9/site-packages/transformers/optimization.py:391: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 1 | step: 0/5475 | loss: 6.1989 | time: 0.1 (minutes)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/mhrrs/.local/lib/python3.9/site-packages/torch/optim/lr_scheduler.py:138: UserWarning: Detected call of `lr_scheduler.step()` before `optimizer.step()`. In PyTorch 1.1.0 and later, you should call them in the opposite order: `optimizer.step()` before `lr_scheduler.step()`.  Failure to do this will result in PyTorch skipping the first value of the learning rate schedule. See more details at https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate\n",
      "  warnings.warn(\"Detected call of `lr_scheduler.step()` before `optimizer.step()`. \"\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 1 | step: 100/5475 | loss: 6.2183 | time: 0.4 (minutes)\n",
      "Epoch: 1 | step: 200/5475 | loss: 6.1947 | time: 0.7 (minutes)\n",
      "Epoch: 1 | step: 300/5475 | loss: 6.1534 | time: 1.0 (minutes)\n",
      "Epoch: 1 | step: 400/5475 | loss: 6.0792 | time: 1.3 (minutes)\n",
      "Epoch: 1 | step: 500/5475 | loss: 5.9394 | time: 1.6 (minutes)\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Input \u001b[0;32mIn [18]\u001b[0m, in \u001b[0;36m<cell line: 16>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      9\u001b[0m model \u001b[38;5;241m=\u001b[39m BertForQuestionAnswering\u001b[38;5;241m.\u001b[39mfrom_pretrained(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mbert-base-uncased\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[1;32m     11\u001b[0m \u001b[38;5;66;03m# TOOK 45 MIN for 1 EPOCH\u001b[39;00m\n\u001b[1;32m     12\u001b[0m \u001b[38;5;66;03m# train_model(train_loader)\u001b[39;00m\n\u001b[1;32m     13\u001b[0m \n\u001b[1;32m     14\u001b[0m \n\u001b[1;32m     15\u001b[0m \u001b[38;5;66;03m# TOOK LESS THAN 25 MIN for 1 EPOCH\u001b[39;00m\n\u001b[0;32m---> 16\u001b[0m \u001b[43mtrain_amp_model\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtrain_loader\u001b[49m\u001b[43m)\u001b[49m\n",
      "Input \u001b[0;32mIn [17]\u001b[0m, in \u001b[0;36mtrain_amp_model\u001b[0;34m(train_loader)\u001b[0m\n\u001b[1;32m     33\u001b[0m     outputs \u001b[38;5;241m=\u001b[39m model(input_ids, attention_mask\u001b[38;5;241m=\u001b[39mattention_mask, start_positions\u001b[38;5;241m=\u001b[39mstart_pos, end_positions\u001b[38;5;241m=\u001b[39mend_pos)\n\u001b[1;32m     34\u001b[0m     loss \u001b[38;5;241m=\u001b[39m outputs[\u001b[38;5;241m0\u001b[39m]\n\u001b[0;32m---> 36\u001b[0m \u001b[43mscaler\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mscale\u001b[49m\u001b[43m(\u001b[49m\u001b[43mloss\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbackward\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     37\u001b[0m total_loss \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m loss\u001b[38;5;241m.\u001b[39mitem()\n\u001b[1;32m     39\u001b[0m \u001b[38;5;66;03m# Gradient accumulation implementation\u001b[39;00m\n",
      "File \u001b[0;32m~/.local/lib/python3.9/site-packages/torch/_tensor.py:488\u001b[0m, in \u001b[0;36mTensor.backward\u001b[0;34m(self, gradient, retain_graph, create_graph, inputs)\u001b[0m\n\u001b[1;32m    478\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m has_torch_function_unary(\u001b[38;5;28mself\u001b[39m):\n\u001b[1;32m    479\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m handle_torch_function(\n\u001b[1;32m    480\u001b[0m         Tensor\u001b[38;5;241m.\u001b[39mbackward,\n\u001b[1;32m    481\u001b[0m         (\u001b[38;5;28mself\u001b[39m,),\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    486\u001b[0m         inputs\u001b[38;5;241m=\u001b[39minputs,\n\u001b[1;32m    487\u001b[0m     )\n\u001b[0;32m--> 488\u001b[0m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mautograd\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbackward\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    489\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mgradient\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mretain_graph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcreate_graph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minputs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minputs\u001b[49m\n\u001b[1;32m    490\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/.local/lib/python3.9/site-packages/torch/autograd/__init__.py:197\u001b[0m, in \u001b[0;36mbackward\u001b[0;34m(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)\u001b[0m\n\u001b[1;32m    192\u001b[0m     retain_graph \u001b[38;5;241m=\u001b[39m create_graph\n\u001b[1;32m    194\u001b[0m \u001b[38;5;66;03m# The reason we repeat same the comment below is that\u001b[39;00m\n\u001b[1;32m    195\u001b[0m \u001b[38;5;66;03m# some Python versions print out the first line of a multi-line function\u001b[39;00m\n\u001b[1;32m    196\u001b[0m \u001b[38;5;66;03m# calls in the traceback and some print out the last line\u001b[39;00m\n\u001b[0;32m--> 197\u001b[0m \u001b[43mVariable\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_execution_engine\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun_backward\u001b[49m\u001b[43m(\u001b[49m\u001b[43m  \u001b[49m\u001b[38;5;66;43;03m# Calls into the C++ engine to run the backward pass\u001b[39;49;00m\n\u001b[1;32m    198\u001b[0m \u001b[43m    \u001b[49m\u001b[43mtensors\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mgrad_tensors_\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mretain_graph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcreate_graph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minputs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    199\u001b[0m \u001b[43m    \u001b[49m\u001b[43mallow_unreachable\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43maccumulate_grad\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m)\u001b[49m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "import torch.utils.data as data_utils\n",
    "import gc\n",
    "\n",
    "# clear cuda\n",
    "torch.cuda.empty_cache()\n",
    "gc.collect()\n",
    "# del model\n",
    "\n",
    "model = BertForQuestionAnswering.from_pretrained('bert-base-uncased')\n",
    "\n",
    "# TOOK 45 MIN for 1 EPOCH\n",
    "# train_model(train_loader)\n",
    "\n",
    "# TOOK LESS THAN 25 MIN for 1 EPOCH\n",
    "train_amp_model(train_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "12ae6fff-1f5c-40d7-96d5-f77515fb7c87",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('./bert-squadv6/tokenizer_config.json',\n",
       " './bert-squadv6/special_tokens_map.json',\n",
       " './bert-squadv6/vocab.txt',\n",
       " './bert-squadv6/added_tokens.json',\n",
       " './bert-squadv6/tokenizer.json')"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.save_pretrained(\"./bert-squadv7\")\n",
    "tokenizer.save_pretrained(\"./bert-squadv7\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ccd10e83-34d9-412d-b2bb-e1fd08831d05",
   "metadata": {},
   "source": [
    "# Evaluate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "f9140e85-6366-459f-a582-bc66f5d6d12f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def eval_model(eval_loader):\n",
    "    model.eval()\n",
    "\n",
    "    # accuracy list\n",
    "    acc = []\n",
    "\n",
    "    for batch in tqdm(eval_loader):\n",
    "        with torch.no_grad():\n",
    "            input_ids = batch['input_ids'].to(device)\n",
    "            attention_mask = batch['attention_mask'].to(device)\n",
    "            start_true = batch['start_positions'].to(device)\n",
    "            end_true = batch['end_positions'].to(device)\n",
    "\n",
    "            outputs = model(input_ids, attention_mask=attention_mask)\n",
    "\n",
    "            start_pred = torch.argmax(outputs['start_logits'], dim=1)\n",
    "            end_pred = torch.argmax(outputs['end_logits'], dim=1)\n",
    "\n",
    "            acc.append(((start_pred == start_true).sum()/len(start_pred)).item())\n",
    "            acc.append(((end_pred == end_true).sum()/len(end_pred)).item())\n",
    "\n",
    "    acc = sum(acc)/len(acc)\n",
    "\n",
    "    print(\"\\n\\nT/P\\tanswer_start\\tanswer_end\\n\")\n",
    "    for i in range(len(start_true)):\n",
    "        print(f\"true\\t{start_true[i]}\\t{end_true[i]}\\n\"\n",
    "            f\"pred\\t{start_pred[i]}\\t{end_pred[i]}\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "746c181d-3c71-4c61-a877-7ba4ac36fe87",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_path = './bert-squadv6'\n",
    "model = BertForQuestionAnswering.from_pretrained(model_path)\n",
    "model = model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "352d4dcd-f051-4284-9624-80f6c679fe96",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 5475/5475 [10:43<00:00,  8.51it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "T/P\tanswer_start\tanswer_end\n",
      "\n",
      "true\t97\t98\n",
      "pred\t97\t98\n",
      "\n",
      "true\t45\t45\n",
      "pred\t45\t45\n",
      "\n",
      "true\t1\t2\n",
      "pred\t1\t2\n",
      "\n",
      "true\t4\t4\n",
      "pred\t4\t4\n",
      "\n",
      "true\t4\t8\n",
      "pred\t4\t8\n",
      "\n",
      "true\t21\t22\n",
      "pred\t21\t22\n",
      "\n",
      "true\t55\t90\n",
      "pred\t55\t90\n",
      "\n",
      "true\t34\t49\n",
      "pred\t34\t49\n",
      "\n",
      "true\t111\t111\n",
      "pred\t111\t111\n",
      "\n",
      "true\t52\t52\n",
      "pred\t52\t52\n",
      "\n",
      "true\t36\t36\n",
      "pred\t36\t36\n",
      "\n",
      "true\t16\t20\n",
      "pred\t16\t19\n",
      "\n",
      "true\t5\t5\n",
      "pred\t5\t5\n",
      "\n",
      "true\t101\t103\n",
      "pred\t101\t103\n",
      "\n",
      "true\t78\t79\n",
      "pred\t78\t79\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "eval_model(eval_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "677b2629-d0b2-475c-8a09-a2308a001683",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}