Fuat_Arslan_NLP_Assignment2_FinetuningBERTforTextClassification.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyP+cdP2WlRt+bQsU/jtDlaF",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/fuat-arslan/NLP_Course/blob/main/Fuat_Arslan_NLP_Assignment2_FinetuningBERTforTextClassification.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OCYZ8_cF9tbT"
      },
      "outputs": [],
      "source": [
        "#if colab used\n",
        "!pip -q install optuna transformers datasets"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#Imports"
      ],
      "metadata": {
        "id": "tBz35JO9967d"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import optuna\n",
        "\n",
        "from transformers import logging\n",
        "import time\n",
        "import datetime\n",
        "\n",
        "import random\n",
        "from random import seed\n",
        "\n",
        "import numpy as np\n",
        "from sklearn.metrics import matthews_corrcoef\n",
        "import torch.utils.data as Data\n",
        "\n",
        "from datasets import load_dataset\n",
        "from transformers import AutoTokenizer\n",
        "import optuna.visualization as vis\n",
        "\n",
        "\n",
        "logging.set_verbosity_error()"
      ],
      "metadata": {
        "id": "Cv6aWpwC94m-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#Utils.py from Tutorial"
      ],
      "metadata": {
        "id": "x7fV_vqd-hmD"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def get_device():\n",
        "    if torch.backends.cuda.is_built():\n",
        "        print(\"CUDA\")\n",
        "        device = torch.device(\"cuda\")\n",
        "    elif torch.backends.mps.is_built():\n",
        "        print(\"mps\")\n",
        "        device = torch.device(\"mps\")\n",
        "    else:\n",
        "        device = torch.device(\"cpu\")\n",
        "        raise Exception(\"GPU is not avalaible!\")\n",
        "    return device\n",
        "\n",
        "\n",
        "# Function to calculate the accuracy of our predictions vs labels\n",
        "def flat_accuracy(preds, labels):\n",
        "    pred_flat = np.argmax(preds, axis=1).flatten()\n",
        "    labels_flat = labels.flatten()\n",
        "    return np.sum(pred_flat == labels_flat) / len(labels_flat)\n",
        "\n",
        "\n",
        "def format_time(elapsed):\n",
        "    elapsed_rounded = int(round((elapsed)))\n",
        "    return str(datetime.timedelta(seconds=elapsed_rounded))\n",
        "\n",
        "\n",
        "def train_eval_loop(\n",
        "    model, loader, optimizer, scheduler, device, n_epochs=2, seed_val=42\n",
        "):\n",
        "    # Set the seed value all over the place to make this reproducible.\n",
        "\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)\n",
        "\n",
        "    loss_values = []\n",
        "    val_loss = []\n",
        "    val_MCC_list = []\n",
        "    t00 = time.time()\n",
        "    for epoch_i in range(0, n_epochs):\n",
        "        print(\"\")\n",
        "        print(\"======== Epoch {:} / {:} ========\".format(epoch_i + 1, n_epochs))\n",
        "        print(\"Training...\")\n",
        "\n",
        "        # Measure how long the training epoch takes.\n",
        "        t0 = time.time()\n",
        "\n",
        "        # Reset the total loss for this epoch.\n",
        "        total_loss = 0\n",
        "        model.train()\n",
        "\n",
        "        # For each batch of training data...\n",
        "        for step, batch in enumerate(loader[\"train\"]):\n",
        "            # print('Memory Usage:')\n",
        "            # print('Allocated:', round(torch.mps.driver_allocated_memory()/1024**3,1), 'GB')\n",
        "\n",
        "            b_input_ids, b_input_mask, b_labels = tuple(t.to(device) for t in batch)\n",
        "            model.zero_grad()\n",
        "\n",
        "            loss = model(b_input_ids, attention_mask=b_input_mask, labels=b_labels).loss\n",
        "            total_loss += loss.item()\n",
        "            loss.backward()\n",
        "\n",
        "            # Clip the norm of the gradients to 1.0, this is to help prevent the \"exploding gradients\" problem.\n",
        "            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",
        "            optimizer.step()\n",
        "            scheduler.step()\n",
        "\n",
        "        avg_train_loss = total_loss / len(loader[\"train\"])\n",
        "        loss_values.append(avg_train_loss)\n",
        "\n",
        "        print(\"\\nAverage training loss: {0:.2f}\".format(avg_train_loss))\n",
        "        print(\"  Training epoch took: {:}\".format(format_time(time.time() - t00)))\n",
        "\n",
        "        print(\"\\nRunning Validation...\")\n",
        "        t0 = time.time()\n",
        "        model.eval()\n",
        "        val_mcc, nb_eval_steps = 0, 0\n",
        "\n",
        "        for batch in loader[\"validation\"]:\n",
        "            b_input_ids, b_input_mask, b_labels = tuple(t.to(device) for t in batch)\n",
        "            # print('Memory Usage:')\n",
        "            # print('Allocated:', round(torch.mps.driver_allocated_memory()/1024**3,1), 'GB')\n",
        "\n",
        "            with torch.no_grad():\n",
        "                logits = model(b_input_ids, attention_mask=b_input_mask).logits\n",
        "\n",
        "            logits = logits.detach().cpu().numpy()\n",
        "            logits = np.argmax(logits, axis=1).flatten()\n",
        "            label_ids = b_labels.to(\"cpu\").numpy()\n",
        "\n",
        "            val_mcc += matthews_corrcoef(logits, label_ids)\n",
        "            nb_eval_steps += 1\n",
        "\n",
        "        val_mcc = 100 * (val_mcc / nb_eval_steps)\n",
        "        val_MCC_list.append(val_mcc)\n",
        "        print(\"  Validation MCC: {0:.2f}\".format(val_mcc))\n",
        "        print(\"  Validation took: {:}\".format(format_time(time.time() - t0)))\n",
        "\n",
        "    return val_mcc, (loss_values,val_MCC_list)\n",
        "\n",
        "\n",
        "def init_loader(max_length=16, batch_size=32, test_size=0.2, random_state=2023):\n",
        "    model_checkpoint = \"bert-base-uncased\"\n",
        "\n",
        "    dataset = load_dataset(\"glue\", \"cola\")\n",
        "\n",
        "    df_s = {}\n",
        "    x = {}\n",
        "    y = {}\n",
        "    input_ids, attention_mask = {}, {}\n",
        "    datasets, loader = {}, {}\n",
        "\n",
        "    tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, use_fast=True)\n",
        "\n",
        "    for split in [\"train\", \"validation\", \"test\"]:\n",
        "        df_s[split] = dataset[split].to_pandas()\n",
        "        x[split] = dataset[split].to_pandas().sentence.values\n",
        "        y[split] = dataset[split].to_pandas().label.values\n",
        "\n",
        "        input = tokenizer(\n",
        "            list(x[split]),\n",
        "            max_length=max_length,\n",
        "            truncation=True,\n",
        "            padding=True,\n",
        "            return_tensors=\"pt\",\n",
        "        )\n",
        "\n",
        "        input_ids[split], attention_mask[split] = input.input_ids, input.attention_mask\n",
        "\n",
        "        datasets[split] = Data.TensorDataset(\n",
        "            input_ids[split], attention_mask[split], torch.LongTensor(y[split])\n",
        "        )\n",
        "\n",
        "        loader[split] = Data.DataLoader(\n",
        "            datasets[split], batch_size=batch_size, shuffle=False\n",
        "        )\n",
        "    return loader, y\n",
        "\n",
        "\n",
        "from transformers import get_linear_schedule_with_warmup\n",
        "from transformers import BertForSequenceClassification\n",
        "\n",
        "\n",
        "def init_objects(\n",
        "    lr, n_epochs, dropout_p=0.1, max_length=16, batch_size=32, test_size=0.2, random_state=2023\n",
        "):\n",
        "    loader, _ = init_loader(max_length=max_length, batch_size=batch_size)\n",
        "\n",
        "    model = BertForSequenceClassification.from_pretrained(\n",
        "        \"bert-base-uncased\", num_labels=2\n",
        "    )\n",
        "    model.dropout.p = dropout_p\n",
        "\n",
        "    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, eps=1e-8)\n",
        "\n",
        "    total_steps = len(loader[\"train\"]) * n_epochs\n",
        "\n",
        "    scheduler = get_linear_schedule_with_warmup(\n",
        "        optimizer, num_warmup_steps=0, num_training_steps=total_steps\n",
        "    )\n",
        "    return model, loader, optimizer, scheduler\n"
      ],
      "metadata": {
        "id": "if8sFRE_-f10"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "###Sample train"
      ],
      "metadata": {
        "id": "QrbbiUTn-oMM"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# %%\n",
        "lr = 2e-5\n",
        "n_epochs = 1\n",
        "max_length = 16\n",
        "batch_size = 32\n",
        "test_size = 0.2\n",
        "random_state = 2023\n",
        "device = torch.device(\"cuda\")\n",
        "\n",
        "model, loader, optimizer, scheduler = init_objects(\n",
        "    lr, n_epochs, max_length, batch_size, test_size, random_state\n",
        ")\n",
        "model.to(device)\n",
        "_, _ = train_eval_loop(\n",
        "    model, loader, optimizer, scheduler, device, n_epochs=n_epochs, seed_val=42\n",
        ")"
      ],
      "metadata": {
        "id": "lPOljGZ8-k97"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#Tuner"
      ],
      "metadata": {
        "id": "hVVTl1tc-tjT"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Hyperparam Tune\n",
        "param_dict = {\n",
        "    \"lr\": [1e-5, 2e-5],\n",
        "    \"n_epochs\": [1, 2, 3],\n",
        "    \"max_length\": [16, 32, 64],\n",
        "    \"batch_size\": [32],\n",
        "    \"dropout_p\": [0.0, 0.5]\n",
        "}\n",
        "\n",
        "\n",
        "class BertObjective:\n",
        "    def __init__(self, d, device):\n",
        "        self.d = d\n",
        "        self.device = device\n",
        "\n",
        "    def __call__(self, trial: optuna.trial.Trial):\n",
        "        self.lr = trial.suggest_float(\"lr\", self.d[\"lr\"][0], self.d[\"lr\"][1], log=True)\n",
        "        self.n_epochs = trial.suggest_categorical(\"n_epochs\", self.d[\"n_epochs\"])\n",
        "        self.max_length = trial.suggest_categorical(\"max_length\", self.d[\"max_length\"])\n",
        "        self.batch_size = trial.suggest_categorical(\"batch_size\", self.d[\"batch_size\"])\n",
        "        self.dp = trial.suggest_float(\"dropout_p\", self.d[\"dropout_p\"][0], self.d[\"dropout_p\"][1])\n",
        "\n",
        "        model, loader, optimizer, scheduler = init_objects(\n",
        "            self.lr, self.n_epochs, self.dp, self.max_length, self.batch_size\n",
        "        )\n",
        "        model.to(self.device)\n",
        "        val_mcc, _ = train_eval_loop(\n",
        "            model, loader, optimizer, scheduler, self.device, self.n_epochs\n",
        "        )\n",
        "\n",
        "        return val_mcc\n",
        "\n",
        "\n",
        "device = torch.device(\"cuda\")\n",
        "study = optuna.create_study(study_name=\"Stduy 0\", direction=\"maximize\")\n",
        "study.optimize(BertObjective(param_dict, device), n_trials=20)\n",
        "\n",
        "# Train again with best parameters\n",
        "lr = study.best_params[\"lr\"]\n",
        "n_epochs = study.best_params[\"n_epochs\"]\n",
        "max_length = study.best_params[\"max_length\"]\n",
        "batch_size = study.best_params[\"batch_size\"]\n",
        "dropout_p = study.best_params[\"dropout_p\"]\n",
        "\n",
        "model, loader, optimizer, scheduler = init_objects(lr, n_epochs, dropout_p, max_length, batch_size)\n",
        "model.to(device)\n",
        "val_mcc, _ = train_eval_loop(model, loader, optimizer, scheduler, device, n_epochs)\n",
        "# Obtain Test Results"
      ],
      "metadata": {
        "id": "g9lCE8LL-sgL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "study.best_params"
      ],
      "metadata": {
        "id": "CjHJgcSZ-7AO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Visualize reuslts"
      ],
      "metadata": {
        "id": "nXsEkzWC_Azs"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "vis.plot_parallel_coordinate(study)"
      ],
      "metadata": {
        "id": "q0i8tsSh_APX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "vis.plot_param_importances(study)"
      ],
      "metadata": {
        "id": "VTUTptqh_6pd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#Train Best model"
      ],
      "metadata": {
        "id": "3dHfBpR0APEH"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "b_params = study.best_params"
      ],
      "metadata": {
        "id": "x-icmBm-ADIE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# %%\n",
        "lr = b_params['lr']\n",
        "n_epochs = b_params['n_epochs']\n",
        "max_length = b_params['max_length']\n",
        "batch_size = b_params['batch_size']\n",
        "dropout_p = b_params['dropout_p']\n",
        "test_size = 0.2\n",
        "random_state = 2023\n",
        "device = torch.device(\"cuda\")\n",
        "\n",
        "model, loader, optimizer, scheduler = init_objects(\n",
        "    lr, n_epochs,dropout_p, max_length, batch_size, test_size, random_state\n",
        ")\n",
        "model.to(device)\n",
        "val_mcc, losses = train_eval_loop(\n",
        "    model, loader, optimizer, scheduler, device, n_epochs=n_epochs, seed_val=42\n",
        ")"
      ],
      "metadata": {
        "id": "wVkoatpbAKCT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "# Sample data\n",
        "val_MCC_list = losses[1]\n",
        "loss_val = losses[0]\n",
        "# Create a figure with two subplots\n",
        "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 5))\n",
        "\n",
        "# Plot val_MCC_list\n",
        "ax1.plot(val_MCC_list, marker='o')\n",
        "ax1.set_xlabel('Index')\n",
        "ax1.set_ylabel('MCC Value')\n",
        "ax1.set_title('val_MCC_list')\n",
        "\n",
        "# Plot loss_val\n",
        "ax2.plot(loss_val, marker='o')\n",
        "ax2.set_xlabel('Index')\n",
        "ax2.set_ylabel('Loss Value')\n",
        "ax2.set_title('loss_val')\n",
        "\n",
        "# Adjust the layout to prevent overlapping labels\n",
        "plt.tight_layout()\n",
        "\n",
        "# Display the plot\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "OFGKwZbJAMpQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Test Run"
      ],
      "metadata": {
        "id": "-tqcoXf4AXHG"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def test_run(model,loader):\n",
        "    val_MCC_list = []\n",
        "    print(\"\\nRunning Test...\")\n",
        "    t0 = time.time()\n",
        "    model.eval()\n",
        "    val_mcc, nb_eval_steps = 0, 0\n",
        "\n",
        "    for batch in loader[\"test\"]:\n",
        "        b_input_ids, b_input_mask, b_labels = tuple(t.to(device) for t in batch)\n",
        "        # print('Memory Usage:')\n",
        "        # print('Allocated:', round(torch.mps.driver_allocated_memory()/1024**3,1), 'GB')\n",
        "\n",
        "        with torch.no_grad():\n",
        "            logits = model(b_input_ids, attention_mask=b_input_mask).logits\n",
        "\n",
        "\n",
        "        logits = logits.detach().cpu().numpy()\n",
        "        logits = np.argmax(logits, axis=1).flatten()\n",
        "        print(logits)\n",
        "        label_ids = b_labels.to(\"cpu\").numpy()\n",
        "        print(label_ids)\n",
        "        val_mcc += matthews_corrcoef(logits, label_ids)\n",
        "        nb_eval_steps += 1\n",
        "        val_MCC_list.append(val_mcc)\n",
        "\n",
        "    val_mcc = 100 * (val_mcc / nb_eval_steps)\n",
        "\n",
        "    print(\"  Test MCC: {0:.2f}\".format(val_mcc))\n",
        "    print(\"  Test took: {:}\".format(format_time(time.time() - t0)))\n",
        "\n",
        "    return val_MCC_list"
      ],
      "metadata": {
        "id": "3b1P8WdlAVfS"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#Push Model"
      ],
      "metadata": {
        "id": "-d1IDadKAR1d"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import huggingface_hub\n",
        "from huggingface_hub import notebook_login\n",
        "your_token = None #Please generate a token from huggingface\n",
        "huggingface_hub.login(token = your_token)"
      ],
      "metadata": {
        "id": "6lOozwNJAdEF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model.push_to_hub(repo_id = 'bert_fine_tune')"
      ],
      "metadata": {
        "id": "zNPjROqeAezB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#Custom Model"
      ],
      "metadata": {
        "id": "h9EpolrxAmcZ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def custom_train_eval_loop(\n",
        "    model, loader, optimizer, scheduler, device, n_epochs=2, seed_val=42\n",
        "):\n",
        "    # Set the seed value all over the place to make this reproducible.\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)\n",
        "\n",
        "    loss_values = []\n",
        "    val_loss = []\n",
        "    val_MCC_list = []\n",
        "    t00 = time.time()\n",
        "\n",
        "    for epoch_i in range(0, n_epochs):\n",
        "        print(\"\")\n",
        "        print(\"======== Epoch {:} / {:} ========\".format(epoch_i + 1, n_epochs))\n",
        "        print(\"Training...\")\n",
        "\n",
        "        # Measure how long the training epoch takes.\n",
        "        t0 = time.time()\n",
        "\n",
        "        # Reset the total loss for this epoch.\n",
        "        total_loss = 0\n",
        "        model.train()\n",
        "\n",
        "        # For each batch of training data...\n",
        "        for step, batch in enumerate(loader[\"train\"]):\n",
        "            b_input_ids, b_input_mask, b_labels = tuple(t.to(device) for t in batch)\n",
        "            model.zero_grad()\n",
        "\n",
        "            # Forward pass\n",
        "            loss, logits = model(b_input_ids, attention_mask=b_input_mask, labels=b_labels)\n",
        "\n",
        "            total_loss += loss.item()\n",
        "            loss.backward()\n",
        "\n",
        "            # Clip the norm of the gradients to 1.0\n",
        "            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",
        "            optimizer.step()\n",
        "            scheduler.step()\n",
        "\n",
        "        avg_train_loss = total_loss / len(loader[\"train\"])\n",
        "        loss_values.append(avg_train_loss)\n",
        "\n",
        "        print(\"\\nAverage training loss: {0:.2f}\".format(avg_train_loss))\n",
        "        print(\"Training epoch took: {:}\".format(format_time(time.time() - t00)))\n",
        "\n",
        "        print(\"\\nRunning Validation...\")\n",
        "        t0 = time.time()\n",
        "        model.eval()\n",
        "        val_mcc, nb_eval_steps = 0, 0\n",
        "\n",
        "        for batch in loader[\"validation\"]:\n",
        "            b_input_ids, b_input_mask, b_labels = tuple(t.to(device) for t in batch)\n",
        "\n",
        "            with torch.no_grad():\n",
        "                _, logits = model(b_input_ids, attention_mask=b_input_mask)\n",
        "\n",
        "            logits = logits.detach().cpu().numpy()\n",
        "            logits = np.argmax(logits, axis=1).flatten()\n",
        "            label_ids = b_labels.to(\"cpu\").numpy()\n",
        "\n",
        "            val_mcc += matthews_corrcoef(logits, label_ids)\n",
        "            nb_eval_steps += 1\n",
        "\n",
        "        val_mcc = 100 * (val_mcc / nb_eval_steps)\n",
        "        val_MCC_list.append(val_mcc)\n",
        "        print(\"Validation MCC: {0:.2f}\".format(val_mcc))\n",
        "        print(\"Validation took: {:}\".format(format_time(time.time() - t0)))\n",
        "\n",
        "    return val_mcc, (loss_values, val_MCC_list)\n",
        "\n",
        "def custom_init_objects(\n",
        "    lr, n_epochs, dropout_p=0.1, max_length=16, batch_size=32, test_size=0.2, random_state=2023\n",
        "):\n",
        "    loader, _ = init_loader(max_length=max_length, batch_size=batch_size)\n",
        "\n",
        "    model = BertClassifier(pretrained_model_name='bert-base-uncased', num_classes=2, pooling_fn = max_pooling)\n",
        "\n",
        "    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, eps=1e-8)\n",
        "\n",
        "    total_steps = len(loader[\"train\"]) * n_epochs\n",
        "\n",
        "    scheduler = get_linear_schedule_with_warmup(\n",
        "        optimizer, num_warmup_steps=0, num_training_steps=total_steps\n",
        "    )\n",
        "    return model, loader, optimizer, scheduler"
      ],
      "metadata": {
        "id": "8AZilPhvAn0S"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "from transformers import BertModel"
      ],
      "metadata": {
        "id": "Ujuk7oyOA12a"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Max pooling function\n",
        "def max_pooling(tensor, dim):\n",
        "    return torch.max(tensor, dim)[0]"
      ],
      "metadata": {
        "id": "KUL9CWJHA3Bq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class BertClassifier(nn.Module):\n",
        "    def __init__(self, pretrained_model_name, num_classes, pooling_fn=torch.mean):\n",
        "        super(BertClassifier, self).__init__()\n",
        "        self.bert = BertModel.from_pretrained(pretrained_model_name)\n",
        "        self.pooling_fn = pooling_fn\n",
        "        self.classifier = nn.Linear(self.bert.config.hidden_size, num_classes)\n",
        "        self.loss_fn = nn.CrossEntropyLoss()\n",
        "\n",
        "    def forward(self, input_ids, attention_mask, labels=None):\n",
        "        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)\n",
        "        pooled_output = self.pooling_fn(outputs.last_hidden_state, dim=1)\n",
        "        logits = self.classifier(pooled_output)\n",
        "\n",
        "        if labels is not None:\n",
        "            loss = self.loss_fn(logits, labels)\n",
        "            return loss, logits\n",
        "        else:\n",
        "            return 0, logits"
      ],
      "metadata": {
        "id": "JJ2kDcxzA3xW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Hyperparam Tune\n",
        "param_dict = {\n",
        "    \"lr\": [1e-5, 2e-5],\n",
        "    \"n_epochs\": [1,2,3],\n",
        "    \"max_length\": [16, 32, 64],\n",
        "    \"batch_size\": [32]\n",
        "}\n",
        "\n",
        "\n",
        "class BertObjective:\n",
        "    def __init__(self, d, device):\n",
        "        self.d = d\n",
        "        self.device = device\n",
        "\n",
        "    def __call__(self, trial: optuna.trial.Trial):\n",
        "        self.lr = trial.suggest_float(\"lr\", self.d[\"lr\"][0], self.d[\"lr\"][1], log=True)\n",
        "        self.n_epochs = trial.suggest_categorical(\"n_epochs\", self.d[\"n_epochs\"])\n",
        "        self.max_length = trial.suggest_categorical(\"max_length\", self.d[\"max_length\"])\n",
        "        self.batch_size = trial.suggest_categorical(\"batch_size\", self.d[\"batch_size\"])\n",
        "\n",
        "        model, loader, optimizer, scheduler = custom_init_objects(\n",
        "            self.lr, self.n_epochs,  0.1, self.max_length, self.batch_size\n",
        "        )\n",
        "        model.to(self.device)\n",
        "        val_mcc, _ = custom_train_eval_loop(\n",
        "            model, loader, optimizer, scheduler, self.device, self.n_epochs\n",
        "        )\n",
        "\n",
        "        return val_mcc\n",
        "\n",
        "\n",
        "device = torch.device(\"cuda\")\n",
        "study = optuna.create_study(study_name=\"Stduy 0\", direction=\"maximize\")\n",
        "study.optimize(BertObjective(param_dict, device), n_trials=20)\n",
        "\n",
        "# Train again with best parameters\n",
        "lr = study.best_params[\"lr\"]\n",
        "n_epochs = study.best_params[\"n_epochs\"]\n",
        "max_length = study.best_params[\"max_length\"]\n",
        "batch_size = study.best_params[\"batch_size\"]\n",
        "\n",
        "\n",
        "model, loader, optimizer, scheduler = custom_init_objects(lr, n_epochs, 0.1, max_length, batch_size)\n",
        "model.to(device)\n",
        "val_mcc, losses_l = custom_train_eval_loop(model, loader, optimizer, scheduler, device, n_epochs)\n",
        "# Obtain Test Results"
      ],
      "metadata": {
        "id": "HCwl0C_WA6a9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Sample data\n",
        "val_MCC_list = losses_l[1]\n",
        "loss_val = losses_l[0]\n",
        "# Create a figure with two subplots\n",
        "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 5))\n",
        "\n",
        "# Plot val_MCC_list\n",
        "ax1.plot(val_MCC_list, marker='o')\n",
        "ax1.set_xlabel('Index')\n",
        "ax1.set_ylabel('MCC Value')\n",
        "ax1.set_title('val_MCC_list')\n",
        "\n",
        "# Plot loss_val\n",
        "ax2.plot(loss_val, marker='o')\n",
        "ax2.set_xlabel('Index')\n",
        "ax2.set_ylabel('Loss Value')\n",
        "ax2.set_title('loss_val')\n",
        "\n",
        "# Adjust the layout to prevent overlapping labels\n",
        "plt.tight_layout()\n",
        "\n",
        "# Display the plot\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "6XgDza5XA9Ya"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "vis.plot_param_importances(study)"
      ],
      "metadata": {
        "id": "EINn7Z8pBBs8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "vis.plot_parallel_coordinate(study)"
      ],
      "metadata": {
        "id": "IwjqppYyBERo"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}