Part3_RoBERTA_Finetuning.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Utilities/Misc\n",
    "import os\n",
    "from tqdm import tqdm\n",
    "from pathlib import Path\n",
    "import pickle as pkl\n",
    "from typing import Tuple\n",
    "\n",
    "# Data handling\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import torch\n",
    "from sklearn.model_selection import train_test_split\n",
    "from torch.utils.data import DataLoader, TensorDataset, WeightedRandomSampler\n",
    "from sklearn.preprocessing import LabelEncoder\n",
    "\n",
    "# Visualization\n",
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# ML/Huggingface tools\n",
    "from transformers import RobertaTokenizer, RobertaForSequenceClassification\n",
    "from torch.optim import AdamW\n",
    "\n",
    "# Scoring\n",
    "import torch.nn.functional as F\n",
    "from sklearn.metrics import (\n",
    "    f1_score,\n",
    "    balanced_accuracy_score,\n",
    "    roc_auc_score,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Paths and directories\n",
    "CWD = Path(os.getcwd())\n",
    "DATA_DIR = CWD / \"data\"\n",
    "CORPUS_DIR = DATA_DIR / \"corpus_files\"\n",
    "OBJ_DIR = DATA_DIR / \"objects\"\n",
    "MODEL_SEARCH = DATA_DIR / \"model_search_results\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load/Prepare Data\n",
    "doc_corpus_path = CORPUS_DIR / \"document_corpus.pkl\"\n",
    "doc_df: pd.DataFrame = pd.read_pickle(doc_corpus_path)\n",
    "doc_df.info()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get Training (65%)/Testing(25%)/Validation(10%) Sets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if \"split\" not in doc_df.columns:\n",
    "    model_data = doc_df[~doc_df[\"submission_flair\"].isin([\"Background\", \"Compendium\"])]\n",
    "    train_df, test_df = train_test_split(\n",
    "        model_data, train_size=0.75, random_state=29359\n",
    "    )\n",
    "    train_df, valid_df = train_test_split(\n",
    "        train_df, train_size=65/76, random_state=29359\n",
    "    )\n",
    "    doc_df[\"split\"] = None\n",
    "    doc_df.loc[doc_df[\"UID\"].isin(train_df[\"UID\"]), \"split\"] = \"train\"\n",
    "    doc_df.loc[doc_df[\"UID\"].isin(test_df[\"UID\"]), \"split\"] = \"test\"\n",
    "    doc_df.loc[doc_df[\"UID\"].isin(valid_df[\"UID\"]), \"split\"] = \"validation\"    \n",
    "    doc_df.to_pickle(doc_corpus_path)\n",
    "elif \"validation\" not in doc_df[\"split\"].unique():\n",
    "    train_df = doc_df[doc_df[\"split\"] == \"train\"].copy()\n",
    "    test_df = doc_df[doc_df[\"split\"] == \"test\"].copy()\n",
    "    train_df, valid_df = train_test_split(\n",
    "        train_df, train_size=65/76, random_state=29359\n",
    "    )\n",
    "    doc_df.to_pickle(doc_corpus_path)\n",
    "else:\n",
    "    train_df = doc_df[doc_df[\"split\"] == \"train\"].copy()\n",
    "    valid_df = doc_df[doc_df[\"split\"] == \"validation\"].copy()\n",
    "    test_df = doc_df[doc_df[\"split\"] == \"test\"].copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lab_enc_path = OBJ_DIR/\"label_encoder.pkl\"\n",
    "if lab_enc_path.is_file():\n",
    "    with open(lab_enc_path, \"rb\") as le:\n",
    "        lab_enc = pkl.load(le)\n",
    "    train_df[\"label\"] = lab_enc.transform(train_df[\"submission_flair\"])\n",
    "else:\n",
    "    lab_enc = LabelEncoder()\n",
    "    train_df[\"label\"] = lab_enc.fit_transform(train_df[\"submission_flair\"])\n",
    "    with open(lab_enc_path, \"wb\") as le:\n",
    "        pkl.dump(lab_enc, le)\n",
    "\n",
    "test_df[\"label\"] = lab_enc.transform(test_df[\"submission_flair\"])\n",
    "valid_df[\"label\"] = lab_enc.transform(valid_df[\"submission_flair\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to tokenize and format the input data\n",
    "def tokenize_data(data: pd.DataFrame, tokenizer: RobertaTokenizer, data_str: str):\n",
    "    input_ids = []\n",
    "    attention_masks = []\n",
    "\n",
    "    for text in tqdm(data['clean_text'], desc=f\"Tokenizing ({data_str} data)\"):\n",
    "        tokens = tokenizer.encode_plus(\n",
    "            text,\n",
    "            max_length=512,\n",
    "            padding='max_length',\n",
    "            truncation=True,\n",
    "            add_special_tokens=True,\n",
    "            return_tensors='pt'\n",
    "        )\n",
    "        input_ids.append(tokens['input_ids'])\n",
    "        attention_masks.append(tokens['attention_mask'])\n",
    "\n",
    "    input_ids = torch.cat(input_ids, dim=0)\n",
    "\n",
    "    attention_masks = torch.cat(attention_masks, dim=0)\n",
    "    labels = torch.tensor(data[\"label\"].values)\n",
    "\n",
    "    return TensorDataset(input_ids, attention_masks, labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the RoBERTa tokenizer and model\n",
    "roberta_tokenizer = RobertaTokenizer.from_pretrained(\"roberta-base\")\n",
    "# Tokenize and format the data\n",
    "train_dataset = tokenize_data(train_df, roberta_tokenizer, \"training\")\n",
    "test_dataset = tokenize_data(test_df, roberta_tokenizer, \"testing\")\n",
    "val_dataset = tokenize_data(valid_df, roberta_tokenizer, \"validation\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import/get pre-trained RoBERTA model\n",
    "model = RobertaForSequenceClassification.from_pretrained(\"roberta-base\", num_labels=len(train_df['label'].unique()))\n",
    "# Define the optimizer and loss function\n",
    "OPTIMIZER = AdamW(model.parameters(), lr=2e-5)\n",
    "CRITERION = torch.nn.CrossEntropyLoss()\n",
    "# Send model to correct device\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define some training params\n",
    "BATCH_SIZE = 2 # Batch size for training\n",
    "MAX_EPOCHS = 100 # Maximum of training epochs\n",
    "MIN_EPOCHS = 1 # Minimum number of training epochs\n",
    "\n",
    "# Create a WeightedRandomSampler for class-balanced sampling\n",
    "class_weights = [1.0 / count for count in train_df['label'].value_counts().sort_index().values]\n",
    "class_weights = [i/sum(class_weights) for i in class_weights]\n",
    "train_sample_weights = torch.tensor(train_df['label'].map(lambda i: class_weights[i]).values)\n",
    "train_weighted_rand_sampler = WeightedRandomSampler(weights=train_sample_weights, num_samples=len(train_dataset), replacement=True)\n",
    "val_sample_weights = torch.tensor(valid_df['label'].map(lambda i: class_weights[i]).values)\n",
    "val_weighted_rand_sampler = WeightedRandomSampler(weights=val_sample_weights, num_samples=len(val_dataset), replacement=True)\n",
    "\n",
    "# Create data loaders\n",
    "train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, sampler=train_weighted_rand_sampler)\n",
    "val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, sampler=val_weighted_rand_sampler)\n",
    "test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def numeric_labels_to_one_hot(labels, num_classes: int = 8):\n",
    "    \"\"\"\n",
    "    Convert numeric label encodings to one-hot encoded numpy array.\n",
    "\n",
    "    Args:\n",
    "    labels (list or numpy array): Numeric label encodings.\n",
    "    num_classes (int): Total number of classes.\n",
    "\n",
    "    Returns:\n",
    "    numpy array: One-hot encoded labels.\n",
    "    \"\"\"\n",
    "    one_hot_labels = np.zeros((len(labels), num_classes))\n",
    "    one_hot_labels[np.arange(len(labels)), labels] = 1\n",
    "    return one_hot_labels\n",
    "\n",
    "def eval_model(model: RobertaForSequenceClassification, loader: DataLoader, data_str: str, n_iter: int = 5) -> Tuple[float, float, float, float]:\n",
    "    \"\"\"\n",
    "    Evaluate the model for mean loss, balanaced accuracy, macro F1 score, macro OVR ROC-AUC.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    model : RobertaForSequenceClassification\n",
    "        The RoBERTA model\n",
    "    loader : DataLoader\n",
    "        PyTorch data load which returns data in batches\n",
    "    data_str : str\n",
    "        A string describing the data loader\n",
    "    n_iter : int\n",
    "        A string describing the data loader\n",
    "\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    Tuple[float, float, float, float]\n",
    "        mean_loss, balanced_accuracy, macro_f1, macro_ovr_auc\n",
    "    \"\"\"\n",
    "    # Calculate validation loss and metrics\n",
    "    model.eval()\n",
    "    loss = 0.0\n",
    "    predictions = []\n",
    "    labels = []\n",
    "    probs = []\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for _ in range(n_iter):\n",
    "            for batch in loader:\n",
    "                input_ids, attention_mask, _labels = batch\n",
    "                input_ids, attention_mask, _labels = input_ids.to(device), attention_mask.to(device), _labels.to(device)\n",
    "\n",
    "                outputs = model(input_ids, attention_mask=attention_mask)\n",
    "                logits = outputs.logits\n",
    "                _loss = CRITERION(logits, _labels)\n",
    "                loss += _loss.item()\n",
    "                _predictions = torch.argmax(logits, dim=1)\n",
    "                _probs = F.softmax(logits, dim=1)\n",
    "                predictions.extend(_predictions.cpu().numpy())\n",
    "                labels.extend(_labels.cpu().numpy())\n",
    "                probs.extend(_probs.cpu().numpy())\n",
    "\n",
    "    one_hot = numeric_labels_to_one_hot(labels)\n",
    "    accuracy = balanced_accuracy_score(labels, predictions)\n",
    "    f1 = f1_score(labels, predictions, average='macro')\n",
    "    roc_auc = roc_auc_score(one_hot, probs, multi_class=\"ovr\", average=\"macro\")\n",
    "    mean_loss = loss / len(predictions)\n",
    "\n",
    "    print(f\"Mean {data_str} Loss: {mean_loss}\")\n",
    "    print(f\"{data_str} Balanced Accuracy: {accuracy}\")\n",
    "    print(f\"{data_str} Macro F1 Score: {f1}\")\n",
    "    print(f\"{data_str} Macro OVR AUC-ROC: {roc_auc}\")\n",
    "\n",
    "    return mean_loss, accuracy, f1, roc_auc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_output_path = OBJ_DIR / \"roberta_final.pth\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not model_output_path.is_file():\n",
    "    cont_training = True\n",
    "    saved = False\n",
    "    prev_val_loss = np.inf\n",
    "    results = {\n",
    "        \"epoch\": [],\n",
    "        \"train_loss\": [],\n",
    "        \"train_accuracy\": [],\n",
    "        \"train_f1\": [],\n",
    "        \"train_auc\": [],\n",
    "        \"val_loss\": [],\n",
    "        \"val_accuracy\": [],\n",
    "        \"val_f1\": [],\n",
    "        \"val_auc\": [],\n",
    "        }\n",
    "    for epoch in range(MAX_EPOCHS+1):\n",
    "        torch.save(model.state_dict(), model_output_path)\n",
    "        if cont_training and (epoch <= MAX_EPOCHS):\n",
    "            model.train()\n",
    "            train_loss = 0.0\n",
    "\n",
    "            # Train in batches\n",
    "            for batch in tqdm(train_loader, desc=\"Epoch {}\".format(epoch + 1)):\n",
    "                input_ids, attention_mask, labels = batch\n",
    "                input_ids, attention_mask, labels = input_ids.to(device), attention_mask.to(device), labels.to(device)\n",
    "                \n",
    "                OPTIMIZER.zero_grad()\n",
    "                outputs = model(input_ids, attention_mask=attention_mask, labels=labels.long())\n",
    "                loss = outputs.loss\n",
    "                loss.backward()\n",
    "                OPTIMIZER.step()\n",
    "            \n",
    "            # Report and record training and validation performance\n",
    "            print(f\"Epoch {epoch + 1}/{MAX_EPOCHS}\")\n",
    "            train_loss, train_accuracy, train_f1, train_roc_auc = eval_model(model, train_loader, \"Training\", n_iter=3)\n",
    "            val_loss, val_accuracy, val_f1, val_roc_auc = eval_model(model, val_loader, \"Validation\", n_iter=3)\n",
    "\n",
    "            results[\"epoch\"].append(epoch+1)\n",
    "            results[\"train_loss\"].append(train_loss)\n",
    "            results[\"train_accuracy\"].append(train_accuracy)\n",
    "            results[\"train_f1\"].append(train_f1)\n",
    "            results[\"train_auc\"].append(train_roc_auc)\n",
    "            results[\"val_loss\"].append(val_loss)\n",
    "            results[\"val_accuracy\"].append(val_accuracy)\n",
    "            results[\"val_f1\"].append(val_f1)\n",
    "            results[\"val_auc\"].append(val_roc_auc)\n",
    "\n",
    "            # Save the model if it validation loss has stopped decreasing (give a tolerance of 1% previous loss)\n",
    "            if (prev_val_loss < val_loss*.99) & (epoch >= MIN_EPOCHS):\n",
    "                best_val_loss = val_loss\n",
    "                cont_training = False\n",
    "            prev_val_loss = val_loss\n",
    "    \n",
    "        # If we are at max epochs OR we are not continuing (validation loss dropped)\n",
    "        else:\n",
    "            break\n",
    "\n",
    "    results_df = pd.DataFrame(results)\n",
    "    results_df.to_pickle(MODEL_SEARCH/\"roberta_training_progress.pkl\")\n",
    "else:\n",
    "    # Load the best model\n",
    "    model.load_state_dict(torch.load(model_output_path))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_loss, test_accuracy, test_f1, test_roc_auc = eval_model(model, test_loader, \"Evaluating Test Performance\", n_iter=1)\n",
    "print(f\"Mean Testing Loss: {test_loss / len(test_loader)}\")\n",
    "print(f\"Testing Balanced Accuracy: {test_accuracy}\")\n",
    "print(f\"Testing Macro F1 Score: {test_f1}\")\n",
    "print(f\"TestingMacro OVR AUC-ROC: {test_roc_auc}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for c in [\"loss\", \"accuracy\", \"f1\",\"auc\"]:\n",
    "    dat = results_df.set_index(\"epoch\", drop=True)\n",
    "    dat = dat[[f\"val_{c}\", f\"train_{c}\"]]\n",
    "    sns.lineplot(data=dat)\n",
    "    plt.title(f\"{c.capitalize()} vs Epoch\")\n",
    "    plt.show()\n",
    "    plt.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py310",
   "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.9"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}