Finetuning Optuna.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "53d86bf4",
   "metadata": {},
   "source": [
    "### Hyperparameter Optimization using Optuna"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "045130a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install ipywidgets\n",
    "!jupyter nbextension enable --py widgetsnbextension\n",
    "!pip install transformers==4.17 \n",
    "!pip install datasets\n",
    "!pip install optuna\n",
    "!pip install ray"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a3a7aa7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "from sklearn.metrics import accuracy_score, precision_recall_fscore_support\n",
    "\n",
    "import torch.nn as nn \n",
    "import torch.optim as optim \n",
    "import matplotlib.pyplot as plt\n",
    "from tqdm.notebook import tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b328087",
   "metadata": {},
   "source": [
    "Load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "207b4f26",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "dataset = load_dataset('tweet_eval','emotion')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ce7da4b",
   "metadata": {},
   "source": [
    "Parameters for model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cdd58a8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "pretrained_model_name = \"roberta-base\" \n",
    "max_length = 50\n",
    "batch_size = 32\n",
    "d_in = 768 \n",
    "d_h = 512 \n",
    "d_out = num_labels \n",
    "freeze_pretrained = False #True\n",
    "\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "331f9981",
   "metadata": {},
   "source": [
    "Tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b71d212e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer,Trainer, TrainingArguments\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name, cache_dir = \".\")\n",
    "# Tokenize text, add padding or truncate the text to the max length    \n",
    "def preprocess(examples):\n",
    "    return tokenizer(examples['text'],truncation=True,   padding='max_length',max_length=max_length)   \n",
    "\n",
    "dataset_op = dataset.map(preprocess, batched=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e49487d",
   "metadata": {},
   "source": [
    "Helper functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7b8210b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import f1_score\n",
    "\n",
    "def f1_score_func(preds, labels):\n",
    "    preds_flat = np.argmax(preds, axis=1).flatten()\n",
    "    labels_flat = labels.flatten()\n",
    "    return f1_score(labels_flat, preds_flat, average='weighted')\n",
    "\n",
    "def accuracy_per_class(preds, labels):\n",
    "    #print per class accuracy and return averaged accuracy over all classes\n",
    "    preds_flat = np.argmax(preds, axis=1).flatten()\n",
    "    labels_flat = labels.flatten()\n",
    "    num, den =0,0\n",
    "    for label in np.unique(labels_flat):\n",
    "        y_preds = preds_flat[labels_flat==label]\n",
    "        y_true = labels_flat[labels_flat==label]\n",
    "        print(f'Class: {label}')\n",
    "        print(f'Accuracy: {len(y_preds[y_preds==label])}/{len(y_true)}\\n')\n",
    "        num+=len(y_preds[y_preds==label])\n",
    "        den+=len(y_true)\n",
    "        \n",
    "    return num/den"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91c1cc79",
   "metadata": {},
   "source": [
    "Optuna for Hyperparameter Optimization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a5e4569",
   "metadata": {},
   "outputs": [],
   "source": [
    "import optuna \n",
    "from transformers import RobertaForSequenceClassification\n",
    "\n",
    "def objective(trial: optuna.Trial):     \n",
    "    model = RobertaForSequenceClassification.from_pretrained(pretrained_model_name, num_labels=d_out, output_attentions=False, output_hidden_states=False)     \n",
    "    \n",
    "    training_args = TrainingArguments(         \n",
    "        output_dir='optuna-test',                 \n",
    "        learning_rate=trial.suggest_loguniform('learning_rate', low=4e-5, high=0.01),         \n",
    "        weight_decay=trial.suggest_loguniform('weight_decay', 4e-5, 0.01),         \n",
    "        num_train_epochs=trial.suggest_int('num_train_epochs', low = 2,high= 12),         \n",
    "        per_device_train_batch_size=8,         \n",
    "        per_device_eval_batch_size=8,         \n",
    "        disable_tqdm=True\n",
    "        )     \n",
    "    trainer = Trainer(\n",
    "        model=model,\n",
    "        args=training_args,\n",
    "        train_dataset=dataset_op['train'],         \n",
    "        eval_dataset=dataset_op['validation']     \n",
    "        )      \n",
    "    \n",
    "    result = trainer.train()     \n",
    "    return result.training_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a76fd81",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We want to minimize the loss\n",
    "study = optuna.create_study(study_name='hyper-parameter-search', direction='minimize',\n",
    "                           pruner=optuna.pruners.SuccessiveHalvingPruner()) \n",
    "study.optimize(func=objective, n_trials=10, gc_after_trial=True)   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb24d966",
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Minimum objective value: ' + str(study.best_value))\n",
    "print(study.best_trial)\n",
    "print('Best parameter: ' + str(study.best_params))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7cb9ec7",
   "metadata": {},
   "source": [
    "### Custom BPE Tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12eb9812",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from tokenizers import ByteLevelBPETokenizer\n",
    "\n",
    "ofile = open(\"data/train.txt\",'w')\n",
    "for i in dataset['train']['text']:\n",
    "  ofile.write(i)\n",
    "ofile.close()\n",
    "ofile = open(\"data/validation.txt\",'w')\n",
    "for i in dataset['validation']['text']:\n",
    "  ofile.write(i)\n",
    "ofile.close()\n",
    "ofile = open(\"data/test.txt\",'w')\n",
    "for i in dataset['test']['text']:\n",
    "  ofile.write(i)\n",
    "ofile.close()\n",
    "\n",
    "paths = [str(x) for x in Path(\".\").glob(\"data/*.txt\")]\n",
    "\n",
    "\n",
    "tokenizer = ByteLevelBPETokenizer()\n",
    "# Customize training\n",
    "tokenizer.train(files=paths, vocab_size=50265, min_frequency=2, special_tokens=[\n",
    "    \"<s>\",\n",
    "    \"<pad>\",\n",
    "    \"</s>\",\n",
    "    \"<unk>\",\n",
    "    \"<mask>\",\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8117c4d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir wikiTokens\n",
    "tokenizer.save_model(\"wikiTokens\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ddae8608",
   "metadata": {},
   "source": [
    "Tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1283d821",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import RobertaTokenizerFast\n",
    "tokenizer = RobertaTokenizerFast.from_pretrained(\"./wikiTokens\", max_len=510)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36fa73db",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import LineByLineTextDataset\n",
    "\n",
    "train_dataset = LineByLineTextDataset(\n",
    "    tokenizer=tokenizer,\n",
    "    file_path=\"./data/train.txt\",\n",
    "    block_size=510,\n",
    ")\n",
    "\n",
    "val_dataset = LineByLineTextDataset(\n",
    "    tokenizer=tokenizer,\n",
    "    file_path=\"./data/validation.txt\",\n",
    "    block_size=510,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "695c2626",
   "metadata": {},
   "outputs": [],
   "source": [
    "max_length=510\n",
    "encoded_data_train = tokenizer.batch_encode_plus(\n",
    "    dataset[\"train\"][\"text\"], \n",
    "    add_special_tokens=True, \n",
    "    return_attention_mask=True, \n",
    "    truncation=True,\n",
    "    padding=\"max_length\", \n",
    "    max_length=max_length, \n",
    "    return_tensors='pt'\n",
    ")\n",
    "\n",
    "encoded_data_test = tokenizer.batch_encode_plus(\n",
    "    dataset[\"test\"][\"text\"], \n",
    "    add_special_tokens=True, \n",
    "    return_attention_mask=True, \n",
    "    truncation=True,\n",
    "    padding=\"max_length\",  \n",
    "    max_length=max_length, \n",
    "    return_tensors='pt'\n",
    ")\n",
    "\n",
    "encoded_data_val = tokenizer.batch_encode_plus(\n",
    "    dataset[\"validation\"][\"text\"], \n",
    "    add_special_tokens=True, \n",
    "    return_attention_mask=True, \n",
    "    truncation=True,\n",
    "    padding=\"max_length\", \n",
    "    max_length=max_length, \n",
    "    return_tensors='pt'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91155026",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import TensorDataset\n",
    "\n",
    "input_ids_train = encoded_data_train['input_ids']\n",
    "attention_masks_train = encoded_data_train['attention_mask']\n",
    "labels_train = torch.tensor(dataset[\"train\"][\"label\"])\n",
    "\n",
    "input_ids_test = encoded_data_test['input_ids']\n",
    "attention_masks_test = encoded_data_test['attention_mask']\n",
    "labels_test = torch.tensor(dataset[\"test\"][\"label\"])\n",
    "\n",
    "input_ids_val = encoded_data_val['input_ids']\n",
    "attention_masks_val = encoded_data_val['attention_mask']\n",
    "labels_val = torch.tensor(dataset[\"validation\"][\"label\"])\n",
    "\n",
    "dataset_train = TensorDataset(input_ids_train, attention_masks_train, labels_train)\n",
    "dataset_test = TensorDataset(input_ids_test, attention_masks_test, labels_test)\n",
    "dataset_val = TensorDataset(input_ids_val, attention_masks_val, labels_val)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "969132e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader, RandomSampler, SequentialSampler\n",
    "\n",
    "batch_size = 5 # greater value causes memory error\n",
    "\n",
    "dataloader_train = DataLoader(dataset_train, \n",
    "                              sampler=RandomSampler(dataset_train), \n",
    "                              batch_size=batch_size)\n",
    "\n",
    "dataloader_validation = DataLoader(dataset_val, \n",
    "                                   sampler=SequentialSampler(dataset_val), \n",
    "                                   batch_size=batch_size)\n",
    "\n",
    "dataloader_test = DataLoader(dataset_test, \n",
    "                                   sampler=SequentialSampler(dataset_test), \n",
    "                                   batch_size=batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76d99cbc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import numpy as np\n",
    "\n",
    "seed_val = 0\n",
    "random.seed(seed_val)\n",
    "np.random.seed(seed_val)\n",
    "torch.manual_seed(seed_val)\n",
    "torch.cuda.manual_seed_all(seed_val)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e600f036",
   "metadata": {},
   "source": [
    "Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28adb3be",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import RobertaForSequenceClassification\n",
    "\n",
    "model = RobertaForSequenceClassification.from_pretrained(\"roberta-base\", num_labels=4, output_attentions=False, output_hidden_states=False)\n",
    "for param in model.base_model.parameters():\n",
    "    param.requires_grad = False\n",
    "n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "print(\"\\nThe pre-trained model has {} trainable parameters\".format(n_params))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0aa30a84",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(dataloader_test):\n",
    "\n",
    "    model.eval()\n",
    "    \n",
    "    loss_val_total = 0\n",
    "    predictions, true_vals = [], []\n",
    "    \n",
    "    for batch in dataloader_test:\n",
    "        \n",
    "        batch = tuple(b.to(device) for b in batch)\n",
    "        \n",
    "        inputs = {'input_ids':      batch[0],\n",
    "                  'attention_mask': batch[1],\n",
    "                  'labels':         batch[2],\n",
    "                 }\n",
    "\n",
    "        with torch.no_grad():        \n",
    "            outputs = model(**inputs)\n",
    "            \n",
    "        # print(outputs)\n",
    "        loss = outputs[0]\n",
    "        # loss = loss_fct(outs, labels)\n",
    "        logits = outputs[1]\n",
    "        loss_val_total += loss.item()\n",
    "\n",
    "        logits = logits.detach().cpu().numpy()\n",
    "        label_ids = inputs['labels'].cpu().numpy()\n",
    "        predictions.append(logits)\n",
    "        true_vals.append(label_ids)\n",
    "    \n",
    "    loss_val_avg = loss_val_total/len(dataloader_test) \n",
    "    \n",
    "    predictions = np.concatenate(predictions, axis=0)\n",
    "    true_vals = np.concatenate(true_vals, axis=0)\n",
    "            \n",
    "    return loss_val_avg, predictions, true_vals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f1ee7d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm.notebook import tqdm\n",
    "for epoch in tqdm(range(1, n_epochs+1)):\n",
    "    model.train()\n",
    "    \n",
    "    loss_train_total = 0\n",
    "\n",
    "    progress_bar = tqdm(dataloader_train, desc='Epoch {:1d}'.format(epoch), leave=False, disable=False)\n",
    "    for batch in progress_bar:\n",
    "\n",
    "        model.zero_grad()\n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        batch = tuple(b.to(device) for b in batch)\n",
    "        \n",
    "        inputs = {'input_ids':      batch[0],\n",
    "                  'attention_mask': batch[1],\n",
    "                  'labels':         batch[2],\n",
    "                 }       \n",
    "        # print(inputs)\n",
    "        outputs = model(**inputs)\n",
    "        # print(outputs)\n",
    "        loss = outputs[0]\n",
    "        # print(loss)\n",
    "        loss_train_total += loss.item()\n",
    "        loss.backward()\n",
    "\n",
    "        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",
    "\n",
    "        optimizer.step()\n",
    "        scheduler.step()\n",
    "        optimizer.zero_grad()\n",
    "        progress_bar.set_postfix({'training_loss': '{:.3f}'.format(loss.item()/len(batch))})\n",
    "        # break\n",
    "        \n",
    "    # torch.save(model.state_dict(), f'./finetuned_RoBERTa_epoch_{epoch}.model')\n",
    "        \n",
    "    tqdm.write(f'\\nEpoch {epoch}')\n",
    "    \n",
    "    loss_train_avg = loss_train_total/len(dataloader_train)            \n",
    "    tqdm.write(f'Training loss: {loss_train_avg}')\n",
    "    #On validation set\n",
    "    val_loss, predictions, true_vals = evaluate(dataloader_validation)\n",
    "    val_f1 = f1_score_func(predictions, true_vals)\n",
    "    tqdm.write(f'Validation loss: {val_loss}')\n",
    "    tqdm.write(f'F1 Score (Weighted): {val_f1}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5eef18f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "#On test set\n",
    "val_loss, predictions, true_vals = evaluate(dataloader_test)\n",
    "val_f1 = f1_score_func(predictions, true_vals)\n",
    "print('Test loss: {}'.format(val_loss))\n",
    "print('F1 Score (Weighted): {}'.format(val_f1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92f30186",
   "metadata": {},
   "outputs": [],
   "source": [
    "training_args = TrainingArguments(\n",
    "    output_dir=\"./wikiTokens\",\n",
    "    overwrite_output_dir=True,\n",
    "    num_train_epochs=10,\n",
    "    save_steps=10_000,\n",
    "    save_total_limit=2,\n",
    "    prediction_loss_only=True,\n",
    "    do_eval=True,\n",
    "    evaluation_strategy = \"epoch\",\n",
    "    load_best_model_at_end = True,\n",
    "    save_strategy = \"epoch\"\n",
    ")\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=train_dataset,\n",
    "    eval_dataset= val_dataset\n",
    ")\n",
    "\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53a9a3a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "best_model = RobertaForSequenceClassification.from_pretrained(pretrained_model_name, num_labels=d_out, output_attentions=False, output_hidden_states=False)\n",
    "\n",
    "best_model.to(device)\n",
    "\n",
    "best_model.load_state_dict(torch.load('./finetuned_RoBERTa_epoch_7.model', map_location=torch.device('cpu')))\n",
    "\n",
    "best_model.eval()\n",
    "\n",
    "# _, predictions, true_vals = evaluate(dataloader_test)\n",
    "# accuracy_per_class(predictions, true_vals)"
   ]
  }
 ],
 "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}