three-conventional-methods/1_distilbert-feature-extractor.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3c5d72f4",
   "metadata": {},
   "source": [
    "# LLM as Feature Extractor"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb9d0299-8fc0-48f0-9b02-4c19214d479a",
   "metadata": {},
   "source": [
    "In this feature-based approach, we are using the embeddings from a pretrained transormer to train a random forest and logistic regression model in scikit-learn:\n",
    "\n",
    "<img src=\"figures/1_feature-based.png\" width=500>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6fd9cda8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# pip install transformers datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "df18e3de-577a-43c5-8b9d-868397a6d7da",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# conda install sklearn --yes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "033b75c5",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch       : 2.0.0\n",
      "transformers: 4.27.4\n",
      "datasets    : 2.11.0\n",
      "sklearn     : 1.2.2\n",
      "\n",
      "conda environment: finetuning-blog\n",
      "\n"
     ]
    }
   ],
   "source": [
    "%load_ext watermark\n",
    "%watermark --conda -p torch,transformers,datasets,sklearn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "602ba8a0",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda:0\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4cfd724d",
   "metadata": {
    "tags": []
   },
   "source": [
    "# 1 Loading the Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "e39e2228-5f0b-4fb9-b762-df26c2052b45",
   "metadata": {},
   "outputs": [],
   "source": [
    "# pip install datasets\n",
    "\n",
    "import os.path as op\n",
    "\n",
    "from datasets import load_dataset\n",
    "\n",
    "import lightning as L\n",
    "from lightning.pytorch.loggers import CSVLogger\n",
    "from lightning.pytorch.callbacks import ModelCheckpoint\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "\n",
    "from sklearn.feature_extraction.text import CountVectorizer\n",
    "\n",
    "from local_dataset_utilities import download_dataset, load_dataset_into_to_dataframe, partition_dataset\n",
    "from local_dataset_utilities import IMDBDataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "fb31ac90-9e3a-41d0-baf1-8e613043924b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████| 50000/50000 [00:25<00:00, 1973.05it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Class distribution:\n"
     ]
    }
   ],
   "source": [
    "download_dataset()\n",
    "\n",
    "df = load_dataset_into_to_dataframe()\n",
    "partition_dataset(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "221f30a1-b433-4304-a18d-8d03abd42b58",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_train = pd.read_csv(\"train.csv\")\n",
    "df_val = pd.read_csv(\"val.csv\")\n",
    "df_test = pd.read_csv(\"test.csv\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "846d83b1",
   "metadata": {},
   "source": [
    "# 2 Tokenization and Numericalization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "21114d27-2697-4132-9714-b259bd63f5a1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading and preparing dataset csv/default to /home/sebastian/.cache/huggingface/datasets/csv/default-2417067d5b75d213/0.0.0/6954658bab30a358235fa864b05cf819af0e179325c740e4bc853bcc7ec513e1...\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0f3dbdca454a4e7d8ebfe80e8e946e7d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data files:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1d913db1678e4636849970ff87653992",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Extracting data files:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating train split: 0 examples [00:00, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating validation split: 0 examples [00:00, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating test split: 0 examples [00:00, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset csv downloaded and prepared to /home/sebastian/.cache/huggingface/datasets/csv/default-2417067d5b75d213/0.0.0/6954658bab30a358235fa864b05cf819af0e179325c740e4bc853bcc7ec513e1. Subsequent calls will reuse this data.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0afb64f5f1f945248ac5ed71b2bda5d7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DatasetDict({\n",
      "    train: Dataset({\n",
      "        features: ['index', 'text', 'label'],\n",
      "        num_rows: 35000\n",
      "    })\n",
      "    validation: Dataset({\n",
      "        features: ['index', 'text', 'label'],\n",
      "        num_rows: 5000\n",
      "    })\n",
      "    test: Dataset({\n",
      "        features: ['index', 'text', 'label'],\n",
      "        num_rows: 10000\n",
      "    })\n",
      "})\n"
     ]
    }
   ],
   "source": [
    "imdb_dataset = load_dataset(\n",
    "    \"csv\",\n",
    "    data_files={\n",
    "        \"train\": \"train.csv\",\n",
    "        \"validation\": \"val.csv\",\n",
    "        \"test\": \"test.csv\",\n",
    "    },\n",
    ")\n",
    "\n",
    "print(imdb_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "5ea762ba",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tokenizer input max length: 512\n",
      "Tokenizer vocabulary size: 30522\n"
     ]
    }
   ],
   "source": [
    "from transformers import AutoTokenizer\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"distilbert-base-uncased\")\n",
    "print(\"Tokenizer input max length:\", tokenizer.model_max_length)\n",
    "print(\"Tokenizer vocabulary size:\", tokenizer.vocab_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "8432c15c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def tokenize_text(batch):\n",
    "    return tokenizer(batch[\"text\"], truncation=True, padding=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "0bb392cf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/35000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/5000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/10000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "imdb_tokenized = imdb_dataset.map(tokenize_text, batched=True, batch_size=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "6d4103c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "del imdb_dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bfeb1553",
   "metadata": {},
   "source": [
    "# 3 Using DistilBERT as a Feature Extractor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "9f2c474d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at distilbert-base-uncased were not used when initializing DistilBertModel: ['vocab_layer_norm.bias', 'vocab_transform.bias', 'vocab_transform.weight', 'vocab_projector.weight', 'vocab_projector.bias', 'vocab_layer_norm.weight']\n",
      "- This IS expected if you are initializing DistilBertModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing DistilBertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    }
   ],
   "source": [
    "from transformers import AutoModel\n",
    "\n",
    "model = AutoModel.from_pretrained(\"distilbert-base-uncased\")\n",
    "model.to(device);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "c6686adc",
   "metadata": {},
   "outputs": [],
   "source": [
    "imdb_tokenized.set_format(\"torch\", columns=[\"input_ids\", \"attention_mask\", \"label\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "07122e49",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([3, 512, 768])"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_batch = {\"attention_mask\": imdb_tokenized[\"train\"][:3][\"attention_mask\"].to(device),\n",
    "              \"input_ids\": imdb_tokenized[\"train\"][:3][\"input_ids\"].to(device)}\n",
    "\n",
    "with torch.inference_mode():\n",
    "    test_output = model(**test_batch)\n",
    "    \n",
    "test_output.last_hidden_state.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "083e61f1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([3, 768])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cls_token_output = test_output.last_hidden_state[:, 0]\n",
    "cls_token_output.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "316d0450",
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.inference_mode()\n",
    "def get_output_embeddings(batch):\n",
    "    output = model(\n",
    "        batch[\"input_ids\"].to(device),\n",
    "        attention_mask=batch[\"attention_mask\"].to(device)).last_hidden_state[:, 0]\n",
    "    return {\"features\": output.cpu().numpy()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "2629aaa3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/35000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/5000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/10000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import time\n",
    "start = time.time()\n",
    "\n",
    "imdb_features = imdb_tokenized.map(get_output_embeddings, batched=True, batch_size=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "0fe91178",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train = np.array(imdb_features[\"train\"][\"features\"])\n",
    "y_train = np.array(imdb_features[\"train\"][\"label\"])\n",
    "\n",
    "X_val = np.array(imdb_features[\"validation\"][\"features\"])\n",
    "y_val = np.array(imdb_features[\"validation\"][\"label\"])\n",
    "\n",
    "X_test = np.array(imdb_features[\"test\"][\"features\"])\n",
    "y_test = np.array(imdb_features[\"test\"][\"label\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e76e2e95-e9b3-4a54-b778-0bdcef59f098",
   "metadata": {},
   "source": [
    "# 4 Train Model on Embeddings (Extracted Features)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "81c31cf9-ec66-41a9-aa54-3e5b6ca33cf6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training accuracy 0.8866285714285714\n",
      "Validation accuracy 0.883\n",
      "test accuracy 0.8795\n",
      "Time elapsed 3.28 min\n"
     ]
    }
   ],
   "source": [
    "from sklearn.linear_model import LogisticRegression\n",
    "\n",
    "clf = LogisticRegression(max_iter=1000)\n",
    "clf.fit(X_train, y_train)\n",
    "\n",
    "print(\"Training accuracy\", clf.score(X_train, y_train))\n",
    "print(\"Validation accuracy\", clf.score(X_val, y_val))\n",
    "print(\"test accuracy\", clf.score(X_test, y_test))\n",
    "\n",
    "end = time.time()\n",
    "elapsed = end - start\n",
    "print(f\"Time elapsed {elapsed/60:.2f} min\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "201a4329-7a91-4501-9c75-4d18f4646fa5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training accuracy 1.0\n",
      "Validation accuracy 0.8408\n",
      "test accuracy 0.8324\n"
     ]
    }
   ],
   "source": [
    "from sklearn.ensemble import RandomForestClassifier\n",
    "\n",
    "clf = RandomForestClassifier()\n",
    "clf.fit(X_train, y_train)\n",
    "\n",
    "print(\"Training accuracy\", clf.score(X_train, y_train))\n",
    "print(\"Validation accuracy\", clf.score(X_val, y_val))\n",
    "print(\"test accuracy\", clf.score(X_test, y_test))"
   ]
  }
 ],
 "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}