Final Project Finetune.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0a9de720",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using cpu\n"
     ]
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "import numpy as np\n",
    "import torch\n",
    "import pprint\n",
    "device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n",
    "print(\"Using {}\".format(device))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b7b3be9b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "dataset = load_dataset(\"ai2_arc\", \"ARC-Challenge\")\n",
    "dataset = dataset.filter(lambda item: len(item[\"choices\"][\"label\"]) == 4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "963ea0fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"bert-base-cased\")\n",
    "\n",
    "def tokenize_function(examples):\n",
    "    choice_prompt = [[examples[\"question\"][i] + \" Choose one from the following: {}. {}; {}. {}; {}. {}; {}. {}\".format(examples[\"choices\"][i][\"label\"][0],examples[\"choices\"][i][\"text\"][0], examples[\"choices\"][i][\"label\"][1], examples[\"choices\"][i][\"text\"][1], examples[\"choices\"][i][\"label\"][2], examples[\"choices\"][i][\"text\"][2], examples[\"choices\"][i][\"label\"][3], examples[\"choices\"][i][\"text\"][3])] for i in range(len(examples[\"id\"]))]\n",
    "    gt_sentences = [\"The correct answer is: {}. {}\".format(examples[\"answerKey\"][i], examples[\"choices\"][i][\"text\"][examples[\"choices\"][i][\"label\"].index(examples[\"answerKey\"][i])]) for i in range(len(examples[\"id\"]))] \n",
    "    # Flatten everything\n",
    "    choice_prompt = sum(choice_prompt, [])\n",
    "    \n",
    "    # Tokenize\n",
    "    tokenized_examples = tokenizer(choice_prompt)\n",
    "    tokenized_examples[\"labels\"] = tokenizer(gt_sentences)[\"input_ids\"]\n",
    "    # Un-flatten\n",
    "    return tokenized_examples\n",
    "\n",
    "def show_one(example):\n",
    "    '''show one example in a nicely formatted way'''\n",
    "    print(f\"Question: {example['question']}\")\n",
    "    print(f\"  A - {example['choices']['text'][0]}\")\n",
    "    print(f\"  B - {example['choices']['text'][1]}\")\n",
    "    print(f\"  C - {example['choices']['text'][2]}\")\n",
    "    print(f\"  D - {example['choices']['text'][3]}\")\n",
    "    print(f\"Ground truth: {example['answerKey']}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "022ede5d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dict_keys(['input_ids', 'token_type_ids', 'attention_mask', 'labels'])\n",
      "= = = Intuitive Demo = = =\n",
      "Question: George wants to warm his hands quickly by rubbing them. Which skin surface will produce the most heat?\n",
      "  A - dry palms\n",
      "  B - wet palms\n",
      "  C - palms covered with oil\n",
      "  D - palms covered with lotion\n",
      "Ground truth: A\n",
      "\n",
      "= = = Processed Format = = =\n",
      "-->Prompt:\n",
      "[CLS] George wants to warm his hands quickly by rubbing them. Which skin surface will produce the most heat? Choose one from the following : A. dry palms ; B. wet palms ; C. palms covered with oil ; D. palms covered with lotion [SEP]\n",
      "-->Ground truth answer:\n",
      "[CLS] The correct answer is : A. dry palms [SEP]\n",
      "-->Prompt:\n",
      "[CLS] Which of the following statements best explains why magnets usually stick to a refrigerator door? Choose one from the following : A. The refrigerator door is smooth. ; B. The refrigerator door contains iron. ; C. The refrigerator door is a good conductor. ; D. The refrigerator door has electric wires in it. [SEP]\n",
      "-->Ground truth answer:\n",
      "[CLS] The correct answer is : B. The refrigerator door contains iron. [SEP]\n",
      "-->Prompt:\n",
      "[CLS] A fold observed in layers of sedimentary rock most likely resulted from the Choose one from the following : A. cooling of flowing magma. ; B. converging of crustal plates. ; C. deposition of river sediments. ; D. solution of carbonate minerals. [SEP]\n",
      "-->Ground truth answer:\n",
      "[CLS] The correct answer is : B. converging of crustal plates. [SEP]\n"
     ]
    }
   ],
   "source": [
    "examples = dataset[\"train\"][:3]\n",
    "features = tokenize_function(examples)\n",
    "print(features.keys())\n",
    "print('= = = Intuitive Demo = = =')\n",
    "show_one(dataset[\"train\"][0])\n",
    "\n",
    "print('\\n= = = Processed Format = = =')\n",
    "for i in range(3):\n",
    "    print(\"-->Prompt:\")\n",
    "    print(tokenizer.decode(features[\"input_ids\"][i]))\n",
    "    print(\"-->Ground truth answer:\")\n",
    "    print(tokenizer.decode(features[\"labels\"][i]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ed7eb54e",
   "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",
      "Map: 100%|██████████| 1117/1117 [00:00<00:00, 5022.91 examples/s]\n",
      "Map: 100%|██████████| 1165/1165 [00:00<00:00, 3619.34 examples/s]\n",
      "Map: 100%|██████████| 295/295 [00:00<00:00, 3512.03 examples/s]\n"
     ]
    }
   ],
   "source": [
    "tokenized_datasets = dataset.map(tokenize_function, batched=True)\n",
    "accepted_keys = [\"input_ids\", \"attention_mask\", \"labels\"]\n",
    "for key in tokenized_datasets['train'].features.keys():\n",
    "    if key not in accepted_keys:\n",
    "        tokenized_datasets = tokenized_datasets.remove_columns(key)\n",
    "tokenized_datasets.set_format(\"torch\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d662e605",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import DataCollatorForSeq2Seq\n",
    "from torch.utils.data import DataLoader\n",
    "collator = DataCollatorForSeq2Seq(tokenizer, model, padding=\"max_length\")\n",
    "\n",
    "# TODO: load model here\n",
    "\n",
    "batch_size = 2\n",
    "# data_collator = DataCollatorForMultipleChoice(tokenizer)\n",
    "\n",
    "train_dataset = tokenized_datasets[\"train\"].shuffle(seed=123).select(range(1000))\n",
    "train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=batch_size, collate_fn=collator)\n",
    "eval_dataset = tokenized_datasets[\"validation\"]\n",
    "eval_dataloader = DataLoader(eval_dataset, batch_size=batch_size, collate_fn=collator)\n",
    "test_dataset = tokenized_datasets[\"test\"]\n",
    "test_dataloader = DataLoader(test_dataset, batch_size=batch_size, collate_fn=collator)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "428a82e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Training Loop\n",
    "from torch.optim import AdamW\n",
    "from tqdm import tqdm\n",
    "\n",
    "optimizer = AdamW(model.parameters(), lr=5e-5)\n",
    "\n",
    "num_epochs = 1\n",
    "model.train()\n",
    "for epoch in tqdm(range(num_epochs)):  # Define num_epochs\n",
    "    epoch_loss = 0\n",
    "    for batch in train_dataloader:\n",
    "        optimizer.zero_grad()\n",
    "        input_ids = batch['input_ids'].to(device)\n",
    "        attention_mask = batch['attention_mask'].to(device)\n",
    "        \n",
    "        # Debugging: Print shapes\n",
    "        print(f\"Input IDs shape: {input_ids.shape}\")\n",
    "        print(f\"Attention mask shape: {attention_mask.shape}\")\n",
    "        labels = input_ids.clone()  # Assuming a denoising objective\n",
    "        outputs = model(input_ids, attention_mask=attention_mask, labels=labels)\n",
    "        loss = outputs.loss\n",
    "        epoch_loss += loss\n",
    "        loss.backward()\n",
    "        optimizer.step()    \n",
    "    print(f\"epoch {epoch+1} loss\", epoch_loss)"
   ]
  }
 ],
 "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.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}