sciq.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9685aa79-2af1-424a-87e0-99ddf6a9eb05",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install datasets -q\n",
    "!pip install transformers -q\n",
    "!pip install accelerate -U -q\n",
    "!pip install torch torchvision torchaudio -q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a5fe6e8e-01fa-47c4-868e-d12df4010f4c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['question', 'distractor3', 'distractor1', 'distractor2', 'correct_answer', 'support'],\n",
       "        num_rows: 11679\n",
       "    })\n",
       "    validation: Dataset({\n",
       "        features: ['question', 'distractor3', 'distractor1', 'distractor2', 'correct_answer', 'support'],\n",
       "        num_rows: 1000\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['question', 'distractor3', 'distractor1', 'distractor2', 'correct_answer', 'support'],\n",
       "        num_rows: 1000\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "import torch\n",
    "from transformers import AutoTokenizer, AutoModelForMultipleChoice, TrainingArguments, Trainer\n",
    "from datetime import datetime\n",
    "\n",
    "project = \"sciq\"\n",
    "dataset = load_dataset(project)\n",
    "dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5976518c-01bc-400e-9ff7-68eccd34164e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of examples in the 'sciq' dataset: 11679\n",
      "Sample Data:\n",
      "{'question': 'What type of organism is commonly used in preparation of foods such as cheese and yogurt?', 'distractor3': 'viruses', 'distractor1': 'protozoa', 'distractor2': 'gymnosperms', 'correct_answer': 'mesophilic organisms', 'support': 'Mesophiles grow best in moderate temperature, typically between 25°C and 40°C (77°F and 104°F). Mesophiles are often found living in or on the bodies of humans or other animals. The optimal growth temperature of many pathogenic mesophiles is 37°C (98°F), the normal human body temperature. Mesophilic organisms have important uses in food preparation, including cheese, yogurt, beer and wine.'}\n"
     ]
    }
   ],
   "source": [
    "# Number of examples in the dataset\n",
    "num_examples = len(dataset[\"train\"])\n",
    "\n",
    "# Print the dataset size\n",
    "print(f\"Number of examples in the 'sciq' dataset: {num_examples}\")\n",
    "\n",
    "# Display some sample data\n",
    "sample_data = dataset[\"train\"][0]  # You can choose any index\n",
    "print(\"Sample Data:\")\n",
    "print(sample_data)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b760985f-40f6-4856-acff-acee33307419",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of BertForMultipleChoice were not initialized from the model checkpoint at bert-large-uncased and are newly initialized: ['classifier.weight', 'classifier.bias']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "base_model_name = \"bert-large-uncased\"\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(base_model_name)\n",
    "model = AutoModelForMultipleChoice.from_pretrained(base_model_name)\n",
    "\n",
    "def preprocess(example):\n",
    "    # Combine question with each distractor and correct answer\n",
    "    choices = [example['distractor1'], example['distractor2'], example['distractor3'], example['correct_answer']]\n",
    "    combined_texts = [example['question'] + \" \" + choice for choice in choices]\n",
    "    \n",
    "    # Encode combined texts\n",
    "    encoded_choices = tokenizer(combined_texts, padding='max_length', truncation=True, max_length=128, return_tensors='pt')\n",
    "    \n",
    "    return {\n",
    "        'input_ids': encoded_choices['input_ids'],\n",
    "        'attention_mask': encoded_choices['attention_mask'],\n",
    "        'label': torch.tensor(3)  # 3 because the correct answer is always the last option in our format\n",
    "    }\n",
    "\n",
    "import random\n",
    "\n",
    "def preprocess(example):\n",
    "    choices = [example['distractor1'], example['distractor2'], example['distractor3'], example['correct_answer']]\n",
    "    \n",
    "    # Shuffle choices with their indices\n",
    "    zipped_choices = list(zip(choices, range(4)))  # zip with original indices\n",
    "    random.shuffle(zipped_choices)\n",
    "    shuffled_choices, original_indices = zip(*zipped_choices)\n",
    "    \n",
    "    # Determine label after shuffling\n",
    "    label = original_indices.index(3)  # Find the new position of the correct answer\n",
    "    \n",
    "    combined_texts = [example['question'] + \" \" + choice for choice in shuffled_choices]\n",
    "    \n",
    "    # Encode combined texts\n",
    "    encoded_choices = tokenizer(combined_texts, padding='max_length', truncation=True, max_length=128, return_tensors='pt')\n",
    "    \n",
    "    return {\n",
    "        'input_ids': encoded_choices['input_ids'],\n",
    "        'attention_mask': encoded_choices['attention_mask'],\n",
    "        'label': torch.tensor(label)\n",
    "    }\n",
    "\n",
    "\n",
    "train_data = dataset[\"train\"]\n",
    "validation_data = dataset[\"validation\"]\n",
    "\n",
    "encoded_train_data = train_data.map(preprocess, remove_columns=train_data.column_names)\n",
    "encoded_validation_data = validation_data.map(preprocess, remove_columns=validation_data.column_names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "31f5d369-558d-48c1-bb7d-0445a2113fea",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'input_ids': tensor([[  101,  2054,  2828,  1997, 15923,  2003,  4141,  2109,  1999,  7547,\n",
       "           1997,  9440,  2107,  2004,  8808,  1998, 10930, 27390,  2102,  1029,\n",
       "          15053,  6844,  2050,   102,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0],\n",
       "         [  101,  2054,  2828,  1997, 15923,  2003,  4141,  2109,  1999,  7547,\n",
       "           1997,  9440,  2107,  2004,  8808,  1998, 10930, 27390,  2102,  1029,\n",
       "           9726, 15460,  4842,  5244,   102,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0],\n",
       "         [  101,  2054,  2828,  1997, 15923,  2003,  4141,  2109,  1999,  7547,\n",
       "           1997,  9440,  2107,  2004,  8808,  1998, 10930, 27390,  2102,  1029,\n",
       "          18191,   102,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0],\n",
       "         [  101,  2054,  2828,  1997, 15923,  2003,  4141,  2109,  1999,  7547,\n",
       "           1997,  9440,  2107,  2004,  8808,  1998, 10930, 27390,  2102,  1029,\n",
       "           2033, 28793, 19466,  2594, 11767,   102,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0]]),\n",
       " 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0, 0, 0, 0],\n",
       "         [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "          1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0, 0, 0, 0],\n",
       "         [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0,\n",
       "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0, 0, 0, 0],\n",
       "         [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "          1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0, 0, 0, 0]]),\n",
       " 'label': tensor(3)}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "preprocess(sample_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "9db7b56a-f036-4fc1-82be-b527bffd9658",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Dataset({\n",
       "     features: ['question', 'distractor3', 'distractor1', 'distractor2', 'correct_answer', 'support'],\n",
       "     num_rows: 11679\n",
       " }),\n",
       " Dataset({\n",
       "     features: ['input_ids', 'attention_mask', 'label'],\n",
       "     num_rows: 11679\n",
       " }))"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_data, encoded_train_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b953eebb-a36d-43ba-89fb-ffda3460ab2c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mabhijoy-sar\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
     ]
    }
   ],
   "source": [
    "!pip install -q wandb -U\n",
    "\n",
    "import wandb, os\n",
    "wandb.login()\n",
    "\n",
    "wandb_project = \"sciq-finetune\"\n",
    "if len(wandb_project) > 0:\n",
    "    os.environ[\"WANDB_PROJECT\"] = wandb_project"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "6d671f9c-573c-4e7e-827f-8b1c9fe7799f",
   "metadata": {},
   "outputs": [],
   "source": [
    "if torch.cuda.device_count() > 1: # If more than 1 GPU\n",
    "    model.is_parallelizable = True\n",
    "    model.model_parallel = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "3909c3db-7683-4957-aa90-d90ed899f222",
   "metadata": {},
   "outputs": [],
   "source": [
    "run_name = base_model_name + \"-\" + wandb_project"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "b59576ae-8c18-420a-9910-2b213d8b9c9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "training_args = TrainingArguments(\n",
    "    output_dir='./results',\n",
    "    evaluation_strategy='steps',\n",
    "    eval_steps=100,\n",
    "    logging_steps=50,\n",
    "    save_strategy='epoch',\n",
    "    learning_rate=3e-5,  # Slightly increased learning rate\n",
    "    per_device_train_batch_size=16,  # Increased train batch size\n",
    "    per_device_eval_batch_size=16,   # Increased eval batch size\n",
    "    gradient_accumulation_steps=2,  # Use gradient accumulation\n",
    "    num_train_epochs=3,\n",
    "    weight_decay=0.01,\n",
    "    push_to_hub=False,\n",
    "    logging_dir='./logs',\n",
    "    do_eval=True,\n",
    "    report_to=\"wandb\",\n",
    "    run_name=f\"{run_name}-{datetime.now().strftime('%Y-%m-%d-%H-%M')}\",\n",
    "    # If you choose to include a learning rate scheduler:\n",
    "    lr_scheduler_type=\"linear\",\n",
    "    warmup_ratio=0.1  # Warm up over 10% of total steps\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "9b37166f-6500-43a0-b63b-ec7efbe94b31",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=encoded_train_data,\n",
    "    eval_dataset=encoded_validation_data\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "88f92ecd-a8fd-4e48-a45b-27276fa5e260",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "Changes to your `wandb` environment variables will be ignored because your `wandb` session has already started. For more information on how to modify your settings with `wandb.init()` arguments, please refer to <a href='https://wandb.me/wandb-init' target=\"_blank\">the W&B docs</a>."
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.15.12"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>/home/abhijoysarkar/exp/wandb/run-20231012_113057-aazuwhng</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href='https://wandb.ai/abhijoy-sar/sciq-finetune/runs/aazuwhng' target=\"_blank\">bert-large-uncased-sciq-finetune-2023-10-12-11-30</a></strong> to <a href='https://wandb.ai/abhijoy-sar/sciq-finetune' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View project at <a href='https://wandb.ai/abhijoy-sar/sciq-finetune' target=\"_blank\">https://wandb.ai/abhijoy-sar/sciq-finetune</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run at <a href='https://wandb.ai/abhijoy-sar/sciq-finetune/runs/aazuwhng' target=\"_blank\">https://wandb.ai/abhijoy-sar/sciq-finetune/runs/aazuwhng</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='1095' max='1095' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [1095/1095 38:48, Epoch 3/3]\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",
       "      <th>Validation Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>100</td>\n",
       "      <td>1.332500</td>\n",
       "      <td>1.278223</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>200</td>\n",
       "      <td>1.153500</td>\n",
       "      <td>1.033543</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>300</td>\n",
       "      <td>0.998800</td>\n",
       "      <td>0.890164</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>400</td>\n",
       "      <td>0.842200</td>\n",
       "      <td>0.826168</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>500</td>\n",
       "      <td>0.707700</td>\n",
       "      <td>0.805956</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>600</td>\n",
       "      <td>0.692200</td>\n",
       "      <td>0.837888</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>700</td>\n",
       "      <td>0.671500</td>\n",
       "      <td>0.762519</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>800</td>\n",
       "      <td>0.369000</td>\n",
       "      <td>0.870426</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>900</td>\n",
       "      <td>0.345700</td>\n",
       "      <td>0.843378</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1000</td>\n",
       "      <td>0.398700</td>\n",
       "      <td>0.822918</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=1095, training_loss=0.7403869498265933, metrics={'train_runtime': 2332.4925, 'train_samples_per_second': 15.021, 'train_steps_per_second': 0.469, 'total_flos': 3.2651968852417536e+16, 'train_loss': 0.7403869498265933, 'epoch': 3.0})"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "a5f5b863-3451-447f-9cd9-41ae3df1f4ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_data = dataset[\"test\"]\n",
    "encoded_test_data = test_data.map(preprocess, remove_columns=test_data.column_names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "05ccc20c-b538-4852-893f-1b31ff31b03a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "trainer.eval_dataset = encoded_test_data\n",
    "results = trainer.evaluate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "71e0a8ff-6f9e-4af2-a1f1-4882a0a095a5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.8196806311607361, 'eval_runtime': 20.1356, 'eval_samples_per_second': 49.663, 'eval_steps_per_second': 3.129, 'epoch': 3.0}\n"
     ]
    }
   ],
   "source": [
    "print(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "9de120de-6f13-44ab-9f98-e0251d942c31",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PredictionOutput(predictions=array([[-1.4230529 , -6.4143987 , -6.7155604 ,  0.8424742 ],\n",
       "       [-5.562239  , -7.2469797 , -3.7093894 ,  4.5314064 ],\n",
       "       [-0.51601464, -1.2431774 , -6.3135324 ,  4.4964776 ],\n",
       "       ...,\n",
       "       [-6.3059225 , -4.8675466 ,  2.5583303 ,  5.1088743 ],\n",
       "       [-7.419259  ,  3.008205  , -7.409169  ,  3.1992912 ],\n",
       "       [-3.996625  ,  4.8065557 ,  3.4317608 ,  3.787773  ]],\n",
       "      dtype=float32), label_ids=array([3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,\n",
       "       3, 3, 3, 3, 3, 3, 3, 3, 3, 3]), metrics={'test_loss': 0.8196806311607361, 'test_runtime': 20.1956, 'test_samples_per_second': 49.516, 'test_steps_per_second': 3.119})"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "predictions = trainer.predict(encoded_test_data)\n",
    "predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "680c0a0e-0c89-4119-b1e3-491d3fdab5f0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Counter({3: 705, 1: 105, 2: 98, 0: 92})\n"
     ]
    }
   ],
   "source": [
    "from collections import Counter\n",
    "\n",
    "prediction_counts = Counter(predicted_labels)\n",
    "print(prediction_counts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "4adae7c9-0e3b-410b-a296-50b788e18a36",
   "metadata": {},
   "outputs": [],
   "source": [
    "predicted_labels = predictions.predictions.argmax(axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "c4156baa-4525-4092-a8aa-87825cb6713e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "Changes to your `wandb` environment variables will be ignored because your `wandb` session has already started. For more information on how to modify your settings with `wandb.init()` arguments, please refer to <a href='https://wandb.me/wandb-init' target=\"_blank\">the W&B docs</a>."
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Finishing last run (ID:aazuwhng) before initializing another..."
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Waiting for W&B process to finish... <strong style=\"color:green\">(success).</strong>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "656be7974c424fd58bb4202f67b4fed8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(Label(value='0.002 MB of 0.002 MB uploaded (0.000 MB deduped)\\r'), FloatProgress(value=1.0, max…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<style>\n",
       "    table.wandb td:nth-child(1) { padding: 0 10px; text-align: left ; width: auto;} td:nth-child(2) {text-align: left ; width: 100%}\n",
       "    .wandb-row { display: flex; flex-direction: row; flex-wrap: wrap; justify-content: flex-start; width: 100% }\n",
       "    .wandb-col { display: flex; flex-direction: column; flex-basis: 100%; flex: 1; padding: 10px; }\n",
       "    </style>\n",
       "<div class=\"wandb-row\"><div class=\"wandb-col\"><h3>Run history:</h3><br/><table class=\"wandb\"><tr><td>eval/loss</td><td>█▅▃▂▂▂▁▂▂▂▂</td></tr><tr><td>eval/runtime</td><td>▂▄▁▂▁▁▂▁▂▃█</td></tr><tr><td>eval/samples_per_second</td><td>▇▅█▇██▇█▇▆▁</td></tr><tr><td>eval/steps_per_second</td><td>▇▆█▇██▇█▇▆▁</td></tr><tr><td>train/epoch</td><td>▁▁▁▂▂▂▂▃▃▃▃▃▄▄▄▄▅▅▅▅▅▆▆▆▆▇▇▇▇▇███</td></tr><tr><td>train/global_step</td><td>▁▁▁▂▂▂▂▃▃▃▃▃▄▄▄▄▅▅▅▅▅▆▆▆▆▇▇▇▇▇███</td></tr><tr><td>train/learning_rate</td><td>▄███▇▇▆▆▆▅▅▅▄▄▃▃▃▂▂▁▁</td></tr><tr><td>train/loss</td><td>██▇▆▆▅▅▄▃▃▄▃▃▃▂▁▁▁▁▁▁</td></tr><tr><td>train/total_flos</td><td>▁</td></tr><tr><td>train/train_loss</td><td>▁</td></tr><tr><td>train/train_runtime</td><td>▁</td></tr><tr><td>train/train_samples_per_second</td><td>▁</td></tr><tr><td>train/train_steps_per_second</td><td>▁</td></tr></table><br/></div><div class=\"wandb-col\"><h3>Run summary:</h3><br/><table class=\"wandb\"><tr><td>eval/loss</td><td>0.81968</td></tr><tr><td>eval/runtime</td><td>20.1356</td></tr><tr><td>eval/samples_per_second</td><td>49.663</td></tr><tr><td>eval/steps_per_second</td><td>3.129</td></tr><tr><td>train/epoch</td><td>3.0</td></tr><tr><td>train/global_step</td><td>1095</td></tr><tr><td>train/learning_rate</td><td>0.0</td></tr><tr><td>train/loss</td><td>0.3518</td></tr><tr><td>train/total_flos</td><td>3.2651968852417536e+16</td></tr><tr><td>train/train_loss</td><td>0.74039</td></tr><tr><td>train/train_runtime</td><td>2332.4925</td></tr><tr><td>train/train_samples_per_second</td><td>15.021</td></tr><tr><td>train/train_steps_per_second</td><td>0.469</td></tr></table><br/></div></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run <strong style=\"color:#cdcd00\">bert-large-uncased-sciq-finetune-2023-10-12-11-30</strong> at: <a href='https://wandb.ai/abhijoy-sar/sciq-finetune/runs/aazuwhng' target=\"_blank\">https://wandb.ai/abhijoy-sar/sciq-finetune/runs/aazuwhng</a><br/>Synced 5 W&B file(s), 0 media file(s), 0 artifact file(s) and 0 other file(s)"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Find logs at: <code>./wandb/run-20231012_113057-aazuwhng/logs</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Successfully finished last run (ID:aazuwhng). Initializing new run:<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "155c11020dc444329a3aa07993e69f66",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(Label(value='Waiting for wandb.init()...\\r'), FloatProgress(value=0.011112459988887874, max=1.0…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.15.12"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>/home/abhijoysarkar/exp/wandb/run-20231012_121030-xt982b54</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href='https://wandb.ai/abhijoy-sar/uncategorized/runs/xt982b54' target=\"_blank\">feasible-donkey-2</a></strong> to <a href='https://wandb.ai/abhijoy-sar/uncategorized' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View project at <a href='https://wandb.ai/abhijoy-sar/uncategorized' target=\"_blank\">https://wandb.ai/abhijoy-sar/uncategorized</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run at <a href='https://wandb.ai/abhijoy-sar/uncategorized/runs/xt982b54' target=\"_blank\">https://wandb.ai/abhijoy-sar/uncategorized/runs/xt982b54</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Waiting for W&B process to finish... <strong style=\"color:green\">(success).</strong>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b7fc3e82670c495dbdc5d2ec25a31df6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(Label(value='0.327 MB of 0.327 MB uploaded (0.000 MB deduped)\\r'), FloatProgress(value=1.0, max…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run <strong style=\"color:#cdcd00\">feasible-donkey-2</strong> at: <a href='https://wandb.ai/abhijoy-sar/uncategorized/runs/xt982b54' target=\"_blank\">https://wandb.ai/abhijoy-sar/uncategorized/runs/xt982b54</a><br/>Synced 4 W&B file(s), 1 media file(s), 1 artifact file(s) and 0 other file(s)"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Find logs at: <code>./wandb/run-20231012_121030-xt982b54/logs</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import wandb\n",
    "\n",
    "wandb.init()\n",
    "\n",
    "# Create a table\n",
    "table = wandb.Table(columns=[\"Question\", \"Possible Answers\", \"Correct Answer\", \"Predicted Answer\"])\n",
    "\n",
    "for idx, example in enumerate(test_data):\n",
    "    question = example[\"question\"]\n",
    "    possible_answers = [example[\"distractor1\"], example[\"distractor2\"], example[\"distractor3\"], example[\"correct_answer\"]]\n",
    "    correct_answer = example[\"correct_answer\"]\n",
    "    predicted_answer = possible_answers[predicted_labels[idx]]\n",
    "    \n",
    "    table.add_data(question, possible_answers, correct_answer, predicted_answer)\n",
    "\n",
    "# Log the table\n",
    "wandb.log({\"questions_predictions\": table})\n",
    "\n",
    "wandb.finish()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "025eb629-ca8e-43f1-a3f0-c3d320c720a1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2f80687a1f3643ceb5f43d8cfb16dbd4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/1000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import random\n",
    "\n",
    "def shuffle_choices(example):\n",
    "    choices = [example['distractor1'], example['distractor2'], example['distractor3'], example['correct_answer']]\n",
    "    correct_idx = 3\n",
    "    zipped_choices = list(zip(choices, range(4)))  # zip with original indices\n",
    "    random.shuffle(zipped_choices)\n",
    "    shuffled_choices, original_indices = zip(*zipped_choices)\n",
    "    \n",
    "    combined_texts = [example['question'] + \" \" + choice for choice in shuffled_choices]\n",
    "    encoded_choices = tokenizer(combined_texts, padding='max_length', truncation=True, max_length=128, return_tensors='pt')\n",
    "    \n",
    "    return {\n",
    "        'input_ids': encoded_choices['input_ids'],\n",
    "        'attention_mask': encoded_choices['attention_mask'],\n",
    "        'label': torch.tensor(original_indices.index(correct_idx))\n",
    "    }\n",
    "\n",
    "encoded_test_data_shuffled = test_data.map(shuffle_choices, remove_columns=test_data.column_names)\n",
    "predictions_shuffled = trainer.predict(encoded_test_data_shuffled)\n",
    "predicted_labels_shuffled = predictions_shuffled.predictions.argmax(axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "6a5ff9ba-c34f-4e40-aa96-f4049bc8f6d1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Counter({3: 269, 1: 260, 2: 237, 0: 234})\n"
     ]
    }
   ],
   "source": [
    "prediction_counts_shuffled = Counter(predicted_labels_shuffled)\n",
    "print(prediction_counts_shuffled)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "b521d5b2-b0cf-4b13-b7b7-8c120bde98f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "single_example = test_data[1]  # This picks the first example from the test dataset\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "4c656270-7c2b-4ee7-9b97-5da596fe5a88",
   "metadata": {},
   "outputs": [],
   "source": [
    "# encoded_single_example = preprocess(single_example)\n",
    "encoded_single_example = shuffle_choices(single_example)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "68fa0d2c-1904-4d2c-a487-7cda9893d35b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "\n",
    "dataset_single = [encoded_single_example]\n",
    "loader = DataLoader(dataset_single, batch_size=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "2f82d334-1b95-4f6d-a450-13d63d099626",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.eval()  # Set the model to evaluation mode\n",
    "\n",
    "with torch.no_grad():\n",
    "    for batch in loader:\n",
    "        input_ids = batch[\"input_ids\"].to(model.device)\n",
    "        attention_mask = batch[\"attention_mask\"].to(model.device)\n",
    "        \n",
    "        outputs = model(input_ids=input_ids, attention_mask=attention_mask)\n",
    "        logits = outputs.logits\n",
    "        predicted_label = logits.argmax(dim=1).item()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "12320bcb-9630-4603-b91b-9bc5ae37636a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Question: What term in biotechnology means a genetically exact copy of an organism?\n",
      "\n",
      "Choices:\n",
      "0. adult\n",
      "1. male\n",
      "2. phenotype\n",
      "3. clone\n",
      "\n",
      "Model's Prediction: phenotype\n",
      "\n",
      "Was the model correct? False\n"
     ]
    }
   ],
   "source": [
    "choices = [\n",
    "    single_example[\"distractor1\"],\n",
    "    single_example[\"distractor2\"],\n",
    "    single_example[\"distractor3\"],\n",
    "    single_example[\"correct_answer\"]\n",
    "]\n",
    "\n",
    "print(\"Question:\", single_example[\"question\"])\n",
    "print(\"\\nChoices:\")\n",
    "for idx, choice in enumerate(choices):\n",
    "    print(f\"{idx}. {choice}\")\n",
    "\n",
    "print(\"\\nModel's Prediction:\", choices[predicted_label])\n",
    "\n",
    "correct = predicted_label == 3  # Since the correct answer is always at index 3 in your setup\n",
    "print(\"\\nWas the model correct?\", correct)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f45c66b6-d581-459d-9fc1-15b6b53ec587",
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}