Section 5.2 (SST-2 experiments).ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c71eea9e-3c7b-4711-899c-b23aa88716f9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from transformers import DistilBertTokenizerFast, AutoModel, DistilBertModel\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torch import nn\n",
    "from torch.nn import functional as F\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "import copy\n",
    "\n",
    "# code is based on repo: https://github.com/xszheng2020/memorization\n",
    "# data is from https://github.com/xszheng2020/memorization/tree/master/sst/data\n",
    "\n",
    "train = pd.read_csv(\"train.csv\")\n",
    "val = pd.read_csv(\"dev.csv\")\n",
    "test = pd.read_csv(\"test.csv\")\n",
    "num_classes = train[\"label\"].nunique()\n",
    "tokenizer = DistilBertTokenizerFast.from_pretrained(\"distilbert-base-uncased\")\n",
    "\n",
    "\n",
    "class CustomDataset(Dataset):\n",
    "    def __init__(self, data):\n",
    "        self.data = data\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        row = self.data.iloc[index]\n",
    "        label = row[0]\n",
    "        sentence = row[1]\n",
    "        return label, sentence\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "\n",
    "\n",
    "class CustomDatasetWithMask(Dataset):\n",
    "    def __init__(self, data, mask=None):\n",
    "        self.data = data\n",
    "        self.mask = mask\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        row = self.data.iloc[index]\n",
    "        label = row[0]\n",
    "        sentence = row[1]\n",
    "        sample_index = row[2]\n",
    "        mask = self.mask[sample_index]\n",
    "        return label, sentence, mask\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0d062dcd-f4c4-4d25-b4b9-ccf9c6f6870d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import random\n",
    "import glob\n",
    "import re\n",
    "import pickle\n",
    "\n",
    "BATCH_SIZE = 100\n",
    "EPOCHS = 1000\n",
    "early_stop_steps = 30\n",
    "\n",
    "\n",
    "def collate_fn(batch):\n",
    "    label = [b[0] for b in batch]\n",
    "    text_a = [b[1] for b in batch]\n",
    "    my_dict = tokenizer(\n",
    "        text_a,\n",
    "        None,\n",
    "        add_special_tokens=True,\n",
    "        padding=True,\n",
    "        truncation=True,\n",
    "        max_length=512,\n",
    "        return_attention_mask=True,\n",
    "        return_tensors=\"pt\",\n",
    "    )\n",
    "    label = torch.tensor(label)\n",
    "    return label, my_dict[\"input_ids\"], my_dict[\"attention_mask\"]\n",
    "\n",
    "\n",
    "# linear model\n",
    "class CustomModel_nofreeze(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "    ):\n",
    "        super(CustomModel_nofreeze, self).__init__()\n",
    "        self.bert = DistilBertModel.from_pretrained(\n",
    "            \"distilbert-base-uncased\",\n",
    "            output_hidden_states=False,\n",
    "            output_attentions=False,\n",
    "        )\n",
    "        self.classifier = nn.Linear(768, num_classes)\n",
    "\n",
    "    def forward(self, inputs, is_ids, attention_mask):\n",
    "        if is_ids:\n",
    "            last_hidden = self.bert(input_ids=inputs, attention_mask=attention_mask)[0]\n",
    "        else:\n",
    "            last_hidden = self.bert(\n",
    "                inputs_embeds=inputs, attention_mask=attention_mask\n",
    "            )[0]\n",
    "        ####\n",
    "        cls_embedding = last_hidden[:, 0, :]  # (bs, dim) pooled_output = cls_embedding\n",
    "        ####\n",
    "        logits = self.classifier(cls_embedding)  # (bs, num_labels)\n",
    "        ####\n",
    "        return cls_embedding, logits\n",
    "\n",
    "\n",
    "# model with freezed first 3 layers (DNN (3 layers))\n",
    "class CustomModel(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "    ):\n",
    "        super(CustomModel, self).__init__()\n",
    "\n",
    "        self.bert = DistilBertModel.from_pretrained(\n",
    "            \"distilbert-base-uncased\",\n",
    "            output_hidden_states=False,\n",
    "            output_attentions=False,\n",
    "        )\n",
    "\n",
    "        params_to_freeze = [\n",
    "            \"bert.embeddings.\",\n",
    "            \"bert.transformer.layer.0.\",\n",
    "            \"bert.transformer.layer.1.\",\n",
    "            \"bert.transformer.layer.2.\",\n",
    "            \"bert.transformer.layer.3.\",\n",
    "        ]\n",
    "        for name, param in self.named_parameters():\n",
    "            # if \"classifier\" not in name:  # classifier layer\n",
    "            #     param.requires_grad = False\n",
    "\n",
    "            if any(pfreeze in name for pfreeze in params_to_freeze):\n",
    "                param.requires_grad = False\n",
    "\n",
    "        self.classifier = nn.Linear(768, num_classes)\n",
    "\n",
    "    def forward(self, inputs, is_ids, attention_mask):\n",
    "        if is_ids:\n",
    "            last_hidden = self.bert(input_ids=inputs, attention_mask=attention_mask)[0]\n",
    "        else:\n",
    "            last_hidden = self.bert(\n",
    "                inputs_embeds=inputs, attention_mask=attention_mask\n",
    "            )[0]\n",
    "        ####\n",
    "        cls_embedding = last_hidden[:, 0, :]  # (bs, dim) pooled_output = cls_embedding\n",
    "        ####\n",
    "        logits = self.classifier(cls_embedding)  # (bs, num_labels)\n",
    "        ####\n",
    "\n",
    "        return cls_embedding, logits\n",
    "\n",
    "\n",
    "# model without freeze (DNN)\n",
    "class CustomModel_allfreeze(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "    ):\n",
    "        super(CustomModel_allfreeze, self).__init__()\n",
    "        self.bert = DistilBertModel.from_pretrained(\n",
    "            \"distilbert-base-uncased\",\n",
    "            output_hidden_states=False,\n",
    "            output_attentions=False,\n",
    "        )\n",
    "        for param in self.bert.parameters():\n",
    "            param.requires_grad = False\n",
    "        self.classifier = nn.Linear(768, num_classes)\n",
    "\n",
    "    def forward(self, inputs, is_ids, attention_mask):\n",
    "        if is_ids:\n",
    "            last_hidden = self.bert(input_ids=inputs, attention_mask=attention_mask)[0]\n",
    "        else:\n",
    "            last_hidden = self.bert(\n",
    "                inputs_embeds=inputs, attention_mask=attention_mask\n",
    "            )[0]\n",
    "        ####\n",
    "        cls_embedding = last_hidden[:, 0, :]  # (bs, dim) pooled_output = cls_embedding\n",
    "        ####\n",
    "        logits = self.classifier(cls_embedding)  # (bs, num_labels)\n",
    "        ####\n",
    "        return cls_embedding, logits\n",
    "\n",
    "\n",
    "### Train models\n",
    "for ind, SEED in enumerate([321, 123, 456, 654, 345, 789, 876, 907, 697]):\n",
    "    torch.manual_seed(SEED)\n",
    "    random.seed(SEED)\n",
    "    np.random.seed(SEED)\n",
    "    results = []\n",
    "    ### finetune models for each datasets with top m% removed\n",
    "    for filename in glob.glob(\"data/*.csv\"):\n",
    "        from_last_step = 0\n",
    "        perc = int(re.findall(r\"\\d+\", filename)[0])\n",
    "        print(perc)\n",
    "        train = pd.read_csv(filename)\n",
    "        train_dataset = CustomDataset(data=train)\n",
    "        train_dataloader = DataLoader(\n",
    "            train_dataset,\n",
    "            batch_size=BATCH_SIZE,\n",
    "            num_workers=0,\n",
    "            shuffle=False,\n",
    "            collate_fn=collate_fn,\n",
    "        )\n",
    "\n",
    "        val_dataset = CustomDataset(data=val)\n",
    "        val_dataloader = DataLoader(\n",
    "            val_dataset,\n",
    "            batch_size=BATCH_SIZE,\n",
    "            num_workers=0,\n",
    "            shuffle=False,\n",
    "            collate_fn=collate_fn,\n",
    "        )\n",
    "\n",
    "        test_dataset = CustomDataset(data=test)\n",
    "        test_dataloader = DataLoader(\n",
    "            test_dataset,\n",
    "            batch_size=BATCH_SIZE,\n",
    "            num_workers=0,\n",
    "            shuffle=False,\n",
    "            collate_fn=collate_fn,\n",
    "        )\n",
    "\n",
    "        ###DNN\n",
    "        model = CustomModel_nofreeze()\n",
    "        model.cuda()\n",
    "        optimizer = torch.optim.Adam(model.parameters(), lr=0.000001)\n",
    "\n",
    "        best_val_acc = 0\n",
    "        best_state_dict = None\n",
    "\n",
    "        for i in range(EPOCHS):\n",
    "            model.train()\n",
    "            train_correct = 0\n",
    "            total_train_loss = 0\n",
    "            total_val_loss = 0\n",
    "            val_correct = 0\n",
    "\n",
    "            for batch in train_dataloader:\n",
    "                ####\n",
    "                model.zero_grad()\n",
    "                optimizer.zero_grad()\n",
    "                ####\n",
    "                label = batch[0].cuda()\n",
    "                input_ids = batch[1].cuda()\n",
    "                attention_mask = batch[2].cuda()\n",
    "\n",
    "                _, outputs = model(input_ids, True, attention_mask=attention_mask)\n",
    "\n",
    "                loss = F.cross_entropy(outputs, label)\n",
    "                loss.backward()\n",
    "                optimizer.step()\n",
    "                total_train_loss += loss.item()\n",
    "                train_correct += (torch.argmax(outputs, axis=-1) == label).sum().item()\n",
    "\n",
    "            model.eval()\n",
    "            with torch.no_grad():\n",
    "                for batch in val_dataloader:\n",
    "                    label = batch[0].cuda()\n",
    "                    input_ids = batch[1].cuda()\n",
    "                    attention_mask = batch[2].cuda()\n",
    "\n",
    "                    _, outputs = model(input_ids, True, attention_mask=attention_mask)\n",
    "                    total_val_loss += loss.item()\n",
    "                    val_correct += (\n",
    "                        (torch.argmax(outputs, axis=-1) == label).sum().item()\n",
    "                    )\n",
    "\n",
    "            val_acc = val_correct / len(val)\n",
    "            if val_acc >= best_val_acc:\n",
    "                best_val_acc = val_acc\n",
    "                best_state_dict = copy.deepcopy(model.state_dict())\n",
    "                from_last_step = 0\n",
    "\n",
    "            if from_last_step >= early_stop_steps:\n",
    "                break\n",
    "            from_last_step += 1\n",
    "\n",
    "            if (i % 30 == 0) or (i == EPOCHS - 1):\n",
    "                print(\n",
    "                    i,\n",
    "                    \"train_acc:\",\n",
    "                    round(train_correct / len(train), 6),\n",
    "                    \"train_loss:\",\n",
    "                    round(total_train_loss / len(train), 6),\n",
    "                    \"val_loss:\",\n",
    "                    round(total_val_loss / len(val), 6),\n",
    "                    \"val_acc:\",\n",
    "                    round(val_correct / len(val), 6),\n",
    "                )\n",
    "        # load best chechpoint\n",
    "        model.load_state_dict(best_state_dict)\n",
    "        # test accuracy\n",
    "        model.eval()\n",
    "        test_correct = 0\n",
    "        with torch.no_grad():\n",
    "            for batch in test_dataloader:\n",
    "                label = batch[0].cuda()\n",
    "                input_ids = batch[1].cuda()\n",
    "                attention_mask = batch[2].cuda()\n",
    "\n",
    "                _, outputs = model(input_ids, True, attention_mask=attention_mask)\n",
    "                total_val_loss += loss.item()\n",
    "                test_correct += (torch.argmax(outputs, axis=-1) == label).sum().item()\n",
    "        test_acc = test_correct / len(test)\n",
    "\n",
    "        val_res_acc_nofreeze = best_val_acc\n",
    "        test_res_acc_nofreeze = test_acc\n",
    "        print(best_val_acc, test_acc)\n",
    "\n",
    "        # DNN (3 layers)\n",
    "        from_last_step = 0\n",
    "        model = CustomModel()\n",
    "        model.cuda()\n",
    "        optimizer_grouped_parameters = [\n",
    "            {\n",
    "                \"params\": [\n",
    "                    p for n, p in model.named_parameters() if (p.requires_grad == True)\n",
    "                ]\n",
    "            },\n",
    "        ]\n",
    "        optimizer = torch.optim.Adam(optimizer_grouped_parameters, lr=0.00001)\n",
    "\n",
    "        best_val_acc = 0\n",
    "        best_state_dict = None\n",
    "\n",
    "        for i in range(EPOCHS):\n",
    "            model.train()\n",
    "            train_correct = 0\n",
    "            total_train_loss = 0\n",
    "            total_val_loss = 0\n",
    "            val_correct = 0\n",
    "\n",
    "            for batch in train_dataloader:\n",
    "                ####\n",
    "                model.zero_grad()\n",
    "                optimizer.zero_grad()\n",
    "                ####\n",
    "                label = batch[0].cuda()\n",
    "                input_ids = batch[1].cuda()\n",
    "                attention_mask = batch[2].cuda()\n",
    "\n",
    "                _, outputs = model(input_ids, True, attention_mask=attention_mask)\n",
    "\n",
    "                loss = F.cross_entropy(outputs, label)\n",
    "                loss.backward()\n",
    "                optimizer.step()\n",
    "                total_train_loss += loss.item()\n",
    "                train_correct += (torch.argmax(outputs, axis=-1) == label).sum().item()\n",
    "\n",
    "            model.eval()\n",
    "            with torch.no_grad():\n",
    "                for batch in val_dataloader:\n",
    "                    label = batch[0].cuda()\n",
    "                    input_ids = batch[1].cuda()\n",
    "                    attention_mask = batch[2].cuda()\n",
    "\n",
    "                    _, outputs = model(input_ids, True, attention_mask=attention_mask)\n",
    "                    total_val_loss += loss.item()\n",
    "                    val_correct += (\n",
    "                        (torch.argmax(outputs, axis=-1) == label).sum().item()\n",
    "                    )\n",
    "\n",
    "            val_acc = val_correct / len(val)\n",
    "            if val_acc >= best_val_acc:\n",
    "                best_val_acc = val_acc\n",
    "                best_state_dict = copy.deepcopy(model.state_dict())\n",
    "                from_last_step = 0\n",
    "\n",
    "            if from_last_step >= early_stop_steps:\n",
    "                break\n",
    "            from_last_step += 1\n",
    "\n",
    "            if (i % 30 == 0) or (i == EPOCHS - 1):\n",
    "                print(\n",
    "                    i,\n",
    "                    \"train_acc:\",\n",
    "                    round(train_correct / len(train), 6),\n",
    "                    \"train_loss:\",\n",
    "                    round(total_train_loss / len(train), 6),\n",
    "                    \"val_loss:\",\n",
    "                    round(total_val_loss / len(val), 6),\n",
    "                    \"val_acc:\",\n",
    "                    round(val_correct / len(val), 6),\n",
    "                )\n",
    "        # load best chechpoint\n",
    "        model.load_state_dict(best_state_dict)\n",
    "        # test accuracy\n",
    "        model.eval()\n",
    "        test_correct = 0\n",
    "        with torch.no_grad():\n",
    "            for batch in test_dataloader:\n",
    "                label = batch[0].cuda()\n",
    "                input_ids = batch[1].cuda()\n",
    "                attention_mask = batch[2].cuda()\n",
    "\n",
    "                _, outputs = model(input_ids, True, attention_mask=attention_mask)\n",
    "                total_val_loss += loss.item()\n",
    "                test_correct += (torch.argmax(outputs, axis=-1) == label).sum().item()\n",
    "        test_acc = test_correct / len(test)\n",
    "\n",
    "        val_res_acc = best_val_acc\n",
    "        test_res_acc = test_acc\n",
    "        print(best_val_acc, test_acc)\n",
    "\n",
    "        ### Linear\n",
    "        from_last_step = 0\n",
    "        model = CustomModel_allfreeze()\n",
    "        model.cuda()\n",
    "        optimizer_grouped_parameters = [\n",
    "            {\n",
    "                \"params\": [\n",
    "                    p for n, p in model.named_parameters() if (p.requires_grad == True)\n",
    "                ]\n",
    "            },\n",
    "        ]\n",
    "        optimizer = torch.optim.Adam(optimizer_grouped_parameters, lr=0.0001)\n",
    "\n",
    "        best_val_acc = 0\n",
    "        best_state_dict = None\n",
    "\n",
    "        for i in range(EPOCHS):\n",
    "            model.train()\n",
    "            train_correct = 0\n",
    "            total_train_loss = 0\n",
    "            total_val_loss = 0\n",
    "            val_correct = 0\n",
    "\n",
    "            for batch in train_dataloader:\n",
    "                ####\n",
    "                model.zero_grad()\n",
    "                optimizer.zero_grad()\n",
    "                ####\n",
    "                label = batch[0].cuda()\n",
    "                input_ids = batch[1].cuda()\n",
    "                attention_mask = batch[2].cuda()\n",
    "\n",
    "                _, outputs = model(input_ids, True, attention_mask=attention_mask)\n",
    "\n",
    "                loss = F.cross_entropy(outputs, label)\n",
    "                loss.backward()\n",
    "                optimizer.step()\n",
    "                total_train_loss += loss.item()\n",
    "                train_correct += (torch.argmax(outputs, axis=-1) == label).sum().item()\n",
    "\n",
    "            model.eval()\n",
    "            with torch.no_grad():\n",
    "                for batch in val_dataloader:\n",
    "                    label = batch[0].cuda()\n",
    "                    input_ids = batch[1].cuda()\n",
    "                    attention_mask = batch[2].cuda()\n",
    "\n",
    "                    _, outputs = model(input_ids, True, attention_mask=attention_mask)\n",
    "                    total_val_loss += loss.item()\n",
    "                    val_correct += (\n",
    "                        (torch.argmax(outputs, axis=-1) == label).sum().item()\n",
    "                    )\n",
    "\n",
    "            val_acc = val_correct / len(val)\n",
    "            if val_acc >= best_val_acc:\n",
    "                best_val_acc = val_acc\n",
    "                best_state_dict = copy.deepcopy(model.state_dict())\n",
    "                from_last_step = 0\n",
    "\n",
    "            if from_last_step >= early_stop_steps:\n",
    "                break\n",
    "\n",
    "            from_last_step += 1\n",
    "\n",
    "            if (i % 30 == 0) or (i == EPOCHS - 1):\n",
    "                print(\n",
    "                    i,\n",
    "                    \"train_acc:\",\n",
    "                    round(train_correct / len(train), 6),\n",
    "                    \"train_loss:\",\n",
    "                    round(total_train_loss / len(train), 6),\n",
    "                    \"val_loss:\",\n",
    "                    round(total_val_loss / len(val), 6),\n",
    "                    \"val_acc:\",\n",
    "                    round(val_correct / len(val), 6),\n",
    "                )\n",
    "        # load best chechpoint\n",
    "        model.load_state_dict(best_state_dict)\n",
    "        # test accuracy\n",
    "        model.eval()\n",
    "        test_correct = 0\n",
    "        with torch.no_grad():\n",
    "            for batch in test_dataloader:\n",
    "                label = batch[0].cuda()\n",
    "                input_ids = batch[1].cuda()\n",
    "                attention_mask = batch[2].cuda()\n",
    "\n",
    "                _, outputs = model(input_ids, True, attention_mask=attention_mask)\n",
    "                total_val_loss += loss.item()\n",
    "                test_correct += (torch.argmax(outputs, axis=-1) == label).sum().item()\n",
    "        test_acc = test_correct / len(test)\n",
    "\n",
    "        print(best_val_acc, test_acc)\n",
    "        val_res_allfreeze_acc = best_val_acc\n",
    "        test_res_allfreeze_acc = test_acc\n",
    "\n",
    "        results.append(\n",
    "            [\n",
    "                perc,\n",
    "                val_res_acc_nofreeze,\n",
    "                test_res_acc_nofreeze,\n",
    "                val_res_acc,\n",
    "                test_res_acc,\n",
    "                val_res_allfreeze_acc,\n",
    "                test_res_allfreeze_acc,\n",
    "            ]\n",
    "        )\n",
    "    df = pd.DataFrame(results)\n",
    "    df.columns = [\n",
    "        \"perc\",\n",
    "        \"full\",\n",
    "        \"full_test\",\n",
    "        \"3l\",\n",
    "        \"3l_test\",\n",
    "        \"only_classifier\",\n",
    "        \"only_classifier_test\",\n",
    "    ]\n",
    "    df = df.sort_values(\"perc\")\n",
    "    # save results for each SEED\n",
    "    with open(f\"df_{ind}.pkl\", \"wb\") as f:\n",
    "        pickle.dump(df, f)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "KNNLM",
   "language": "python",
   "name": "knnlm"
  },
  "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}