layer_prunning.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cuda computing enabled\n"
     ]
    }
   ],
   "source": [
    "import re\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "from copy import deepcopy\n",
    "from sklearn.linear_model import LogisticRegression\n",
    "from tqdm import tqdm\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModel, RobertaModel, RobertaTokenizer\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    device = 'cuda'\n",
    "    print(\"Cuda computing enabled\")\n",
    "else:\n",
    "    device = 'cpu'    \n",
    "    print(\"No cuda found\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading the model and data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of RobertaModel were not initialized from the model checkpoint at roberta-base and are newly initialized: ['roberta.pooler.dense.bias', 'roberta.pooler.dense.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n",
      "/home/eduard/.local/lib/python3.8/site-packages/torch/cuda/__init__.py:628: UserWarning: Can't initialize NVML\n",
      "  warnings.warn(\"Can't initialize NVML\")\n"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "Loading the model and tokenizer\n",
    "\"\"\" \n",
    "\n",
    "CACHE_PATH = \"/mnt/hdd_drive/huggingface/hub/\"\n",
    "MODEL_NAME = \"roberta-base\"\n",
    "\n",
    "model = RobertaModel.from_pretrained(MODEL_NAME, cache_dir=CACHE_PATH)\n",
    "tokenizer = RobertaTokenizer.from_pretrained(MODEL_NAME, cache_dir=CACHE_PATH)\n",
    "\n",
    "model = model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "RobertaModel(\n",
       "  (embeddings): RobertaEmbeddings(\n",
       "    (word_embeddings): Embedding(50265, 768, padding_idx=1)\n",
       "    (position_embeddings): Embedding(514, 768, padding_idx=1)\n",
       "    (token_type_embeddings): Embedding(1, 768)\n",
       "    (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "    (dropout): Dropout(p=0.1, inplace=False)\n",
       "  )\n",
       "  (encoder): RobertaEncoder(\n",
       "    (layer): ModuleList(\n",
       "      (0-11): 12 x RobertaLayer(\n",
       "        (attention): RobertaAttention(\n",
       "          (self): RobertaSelfAttention(\n",
       "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (output): RobertaSelfOutput(\n",
       "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (intermediate): RobertaIntermediate(\n",
       "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          (intermediate_act_fn): GELUActivation()\n",
       "        )\n",
       "        (output): RobertaOutput(\n",
       "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (pooler): RobertaPooler(\n",
       "    (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "    (activation): Tanh()\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"\n",
    "Check the model architecture.\n",
    "\"\"\" \n",
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Loading the datasets\n",
    "\"\"\"\n",
    "\n",
    "DATA_PATH = \"...\"\n",
    "\n",
    "df_gpt4o_wiki = pd.read_csv(DATA_PATH + \"gpt-4-o-wiki-correct-1500.csv\")\n",
    "df_gpt4o_reddit = pd.read_csv(DATA_PATH + \"gpt-4-o-reddit-1500.csv\")\n",
    "df_gpt4o_stackex = pd.read_csv(DATA_PATH + \"gpt-4-o-stackexchange-1500.csv\")\n",
    "\n",
    "df_gpt3_wiki = pd.read_json(DATA_PATH + \"gpt3_davinci_003_wikip.jsonl_pp\", lines=True)[:1500]\n",
    "df_gpt3_reddit = pd.read_json(DATA_PATH + \"gpt3_davinci_003_reddit.jsonl_pp\", lines=True)[:1500]\n",
    "df_gpt3_stackex = pd.read_json(DATA_PATH + \"gpt3_davinci_003_300_len.jsonl\", lines=True)[:1500]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Main functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_folder = \"cache/\"\n",
    "def prune_layers(model, layers=[0,], return_weights=True):\n",
    "    \"\"\"\n",
    "    Function to remove attention heads by entire layers. Standart method model.prune_heads({...}) \n",
    "    of BERT-like models might not work with GPU-acceleration if there an entire layer of heads is removed (prunned)  \n",
    "    \n",
    "    Parameters:\n",
    "        model          --- transformer model to prune\n",
    "        layers         --- list of layer numbers (zero-based) to be removed\n",
    "        return_weights --- boolean flag, whenever to return the pruned weights (needed to repair the model)\n",
    "    \"\"\"\n",
    "    layers_to_remove = layers\n",
    "    preserved_copies = { }\n",
    "    print(layers_to_remove)\n",
    "    with torch.no_grad():\n",
    "        for j in layers_to_remove:\n",
    "            # If you plan to use other model than BERT/RoBERTa,\n",
    "            # change the following 2 lines according to the architecture of your model\n",
    "            preserved_copies[j] = deepcopy(model.encoder.layer[j].attention.output.dense.weight)\n",
    "            model.encoder.layer[j].attention.output.dense.weight *= 0.0\n",
    "    \n",
    "    if return_weights:  \n",
    "        return preserved_copies\n",
    "    \n",
    "    \n",
    "def restore_layers(model, preserved_copies):\n",
    "    \"\"\"\n",
    "    Function to repair model pruned by function prune_layers( )\n",
    "    Parameters:\n",
    "        model            --- transformer model to prune\n",
    "        preserved_copies --- dictionary of layer weights returned by prune_layers\n",
    "    \"\"\"\n",
    "   \n",
    "    with torch.no_grad():\n",
    "        for j in preserved_copies.keys():\n",
    "            # If you plan to use other model than BERT/RoBERTa,\n",
    "            # change the following line according to the architecture of your model\n",
    "            model.encoder.layer[j].attention.output.dense.weight += preserved_copies[j]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_avg_pool(series):\n",
    "    \"Returns mean-pooled embeddings for each text in the container <series> \"\n",
    "    \n",
    "    cls_set = []\n",
    "    for text in tqdm(series):\n",
    "        text = re.sub(r'\\s+', ' ', text).strip()\n",
    "        inpt = tokenizer(text, truncation=True, max_length=510, return_tensors=\"pt\").to(device)\n",
    "        with torch.no_grad():\n",
    "            outp = model(**inpt)[0][0]\n",
    "\n",
    "        cls_set.append(torch.mean(outp, 0).cpu().numpy())\n",
    "\n",
    "    return np.vstack(cls_set)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### baseline --- no pruning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1500/1500 [00:38<00:00, 39.28it/s]\n",
      "100%|██████████| 1500/1500 [00:33<00:00, 44.54it/s]\n",
      "100%|██████████| 1500/1500 [00:18<00:00, 81.42it/s]\n",
      "100%|██████████| 1500/1500 [00:31<00:00, 47.50it/s]\n",
      "100%|██████████| 1500/1500 [00:34<00:00, 43.30it/s]\n",
      "100%|██████████| 1500/1500 [00:40<00:00, 37.34it/s]\n"
     ]
    }
   ],
   "source": [
    "cls_lst_w1 = get_avg_pool(df_gpt4o_wiki[\"gold_completion\"])\n",
    "cls_lst_w2 = get_avg_pool(df_gpt4o_wiki[\"gen_completion\"])    \n",
    "np.save(output_folder + 'avg_roberta_human4o_wiki_none.npy', cls_lst_w1)\n",
    "np.save(output_folder + 'avg_roberta_gpt4o_wiki_none.npy', cls_lst_w2)\n",
    "\n",
    "cls_lst_r1 = get_avg_pool(df_gpt4o_reddit[\"gold_completion\"])\n",
    "cls_lst_r2 = get_avg_pool(df_gpt4o_reddit[\"gen_completion\"])\n",
    "np.save(output_folder + 'avg_roberta_human4o_reddit_none.npy', cls_lst_r1)\n",
    "np.save(output_folder + 'avg_roberta_gpt4o_reddit_none.npy', cls_lst_r2)\n",
    "    \n",
    "cls_lst_s1 = get_avg_pool(df_gpt4o_stackex[\"gold_completion\"])\n",
    "cls_lst_s2 = get_avg_pool(df_gpt4o_stackex[\"gen_completion\"])\n",
    "np.save(output_folder + 'avg_roberta_human4o_stackexchange_none.npy', cls_lst_s1)\n",
    "np.save(output_folder + 'avg_roberta_gpt4o_stackexchange_none.npy', cls_lst_s2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1500/1500 [00:35<00:00, 42.35it/s]\n",
      "100%|██████████| 1500/1500 [00:20<00:00, 74.04it/s]\n",
      "100%|██████████| 1500/1500 [00:20<00:00, 73.79it/s]\n",
      "100%|██████████| 1500/1500 [00:22<00:00, 66.43it/s]\n",
      "100%|██████████| 1500/1500 [00:36<00:00, 41.19it/s]\n",
      "100%|██████████| 1500/1500 [00:28<00:00, 53.26it/s]\n"
     ]
    }
   ],
   "source": [
    "cls_lst_w1 = get_avg_pool(df_gpt3_wiki[\"gold_completion\"])\n",
    "cls_lst_w2 = get_avg_pool(df_gpt3_wiki[\"gen_completion\"])    \n",
    "np.save(output_folder + 'avg_roberta_human3_wiki_none.npy', cls_lst_w1)\n",
    "np.save(output_folder + 'avg_roberta_gpt3_wiki_none.npy', cls_lst_w2)\n",
    "\n",
    "cls_lst_r1 = get_avg_pool(df_gpt3_reddit[\"gold_completion\"])\n",
    "cls_lst_r2 = get_avg_pool(df_gpt3_reddit[\"gen_completion\"])\n",
    "np.save(output_folder + 'avg_roberta_human3_reddit_none.npy', cls_lst_r1)\n",
    "np.save(output_folder + 'avg_roberta_gpt3_reddit_none.npy', cls_lst_r2)\n",
    "    \n",
    "cls_lst_s1 = get_avg_pool(df_gpt3_stackex[\"gold_completion\"])\n",
    "cls_lst_s2 = get_avg_pool(df_gpt3_stackex[\"gen_completion\"])\n",
    "np.save(output_folder + 'avg_roberta_human3_stackexchange_none.npy', cls_lst_s1)\n",
    "np.save(output_folder + 'avg_roberta_gpt3_stackexchange_none.npy', cls_lst_s2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### prunning applied"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Proceesing layer #1:\n",
      "[1]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1500/1500 [00:36<00:00, 41.59it/s]\n",
      "100%|██████████| 1500/1500 [00:17<00:00, 84.04it/s]\n",
      "100%|██████████| 1500/1500 [00:20<00:00, 74.80it/s]\n",
      "100%|██████████| 1500/1500 [00:26<00:00, 57.05it/s]\n",
      "100%|██████████| 1500/1500 [00:35<00:00, 41.98it/s]\n",
      "100%|██████████| 1500/1500 [00:29<00:00, 51.61it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\n",
      "Proceesing layer #2:\n",
      "[2]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1500/1500 [00:34<00:00, 43.68it/s]\n",
      "100%|██████████| 1500/1500 [00:19<00:00, 75.29it/s]\n",
      "100%|██████████| 1500/1500 [00:20<00:00, 73.70it/s]\n",
      "100%|██████████| 1500/1500 [00:22<00:00, 65.25it/s]\n",
      "100%|██████████| 1500/1500 [00:37<00:00, 40.07it/s]\n",
      "100%|██████████| 1500/1500 [00:25<00:00, 59.85it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\n",
      "Proceesing layer #3:\n",
      "[3]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1500/1500 [00:36<00:00, 40.70it/s]\n",
      "100%|██████████| 1500/1500 [00:19<00:00, 76.08it/s]\n",
      "100%|██████████| 1500/1500 [00:18<00:00, 82.65it/s] \n",
      "100%|██████████| 1500/1500 [00:25<00:00, 58.11it/s]\n",
      "100%|██████████| 1500/1500 [00:35<00:00, 42.53it/s]\n",
      "100%|██████████| 1500/1500 [00:24<00:00, 60.42it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\n",
      "Proceesing layer #4:\n",
      "[4]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1500/1500 [00:36<00:00, 40.83it/s]\n",
      "100%|██████████| 1500/1500 [00:18<00:00, 80.63it/s] \n",
      "100%|██████████| 1500/1500 [00:19<00:00, 76.18it/s]\n",
      "100%|██████████| 1500/1500 [00:26<00:00, 56.70it/s]\n",
      "100%|██████████| 1500/1500 [00:33<00:00, 44.69it/s]\n",
      "100%|██████████| 1500/1500 [00:27<00:00, 54.50it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\n",
      "Proceesing layer #5:\n",
      "[5]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1500/1500 [00:35<00:00, 41.93it/s]\n",
      "100%|██████████| 1500/1500 [00:18<00:00, 79.52it/s]\n",
      "100%|██████████| 1500/1500 [00:20<00:00, 73.62it/s]\n",
      "100%|██████████| 1500/1500 [00:24<00:00, 62.18it/s]\n",
      "100%|██████████| 1500/1500 [00:36<00:00, 41.39it/s]\n",
      "100%|██████████| 1500/1500 [00:27<00:00, 54.37it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\n",
      "Proceesing layer #6:\n",
      "[6]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1500/1500 [00:35<00:00, 42.48it/s]\n",
      "100%|██████████| 1500/1500 [00:20<00:00, 73.96it/s]\n",
      "100%|██████████| 1500/1500 [00:19<00:00, 78.29it/s] \n",
      "100%|██████████| 1500/1500 [00:21<00:00, 68.43it/s]\n",
      "100%|██████████| 1500/1500 [00:35<00:00, 41.68it/s]\n",
      "100%|██████████| 1500/1500 [00:26<00:00, 56.14it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\n",
      "Proceesing layer #7:\n",
      "[7]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1500/1500 [00:36<00:00, 40.63it/s]\n",
      "100%|██████████| 1500/1500 [00:19<00:00, 75.64it/s]\n",
      "100%|██████████| 1500/1500 [00:17<00:00, 87.83it/s] \n",
      "100%|██████████| 1500/1500 [00:25<00:00, 59.45it/s]\n",
      "100%|██████████| 1500/1500 [00:36<00:00, 41.39it/s]\n",
      "100%|██████████| 1500/1500 [00:24<00:00, 60.05it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\n",
      "Proceesing layer #8:\n",
      "[8]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1500/1500 [00:37<00:00, 40.24it/s]\n",
      "100%|██████████| 1500/1500 [00:16<00:00, 89.61it/s]\n",
      "100%|██████████| 1500/1500 [00:20<00:00, 74.93it/s]\n",
      "100%|██████████| 1500/1500 [00:25<00:00, 59.83it/s]\n",
      "100%|██████████| 1500/1500 [00:34<00:00, 42.95it/s]\n",
      "100%|██████████| 1500/1500 [00:28<00:00, 52.54it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\n",
      "Proceesing layer #9:\n",
      "[9]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1500/1500 [00:34<00:00, 43.09it/s]\n",
      "100%|██████████| 1500/1500 [00:19<00:00, 75.81it/s]\n",
      "100%|██████████| 1500/1500 [00:20<00:00, 73.25it/s] \n",
      "100%|██████████| 1500/1500 [00:25<00:00, 59.61it/s]\n",
      "100%|██████████| 1500/1500 [00:37<00:00, 40.34it/s]\n",
      "100%|██████████| 1500/1500 [00:27<00:00, 54.75it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\n",
      "Proceesing layer #10:\n",
      "[10]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1500/1500 [00:38<00:00, 38.50it/s]\n",
      "100%|██████████| 1500/1500 [00:28<00:00, 53.21it/s]\n",
      "100%|██████████| 1500/1500 [00:27<00:00, 54.09it/s]\n",
      "100%|██████████| 1500/1500 [00:30<00:00, 49.70it/s]\n",
      "100%|██████████| 1500/1500 [00:42<00:00, 35.55it/s]\n",
      "100%|██████████| 1500/1500 [00:32<00:00, 46.63it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\n",
      "Proceesing layer #11:\n",
      "[11]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1500/1500 [00:41<00:00, 36.49it/s]\n",
      "100%|██████████| 1500/1500 [00:29<00:00, 51.38it/s]\n",
      "100%|██████████| 1500/1500 [00:25<00:00, 59.39it/s]\n",
      "100%|██████████| 1500/1500 [00:32<00:00, 46.30it/s]\n",
      "100%|██████████| 1500/1500 [00:40<00:00, 36.69it/s]\n",
      "100%|██████████| 1500/1500 [00:30<00:00, 48.65it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "for i in range(12):\n",
    "    print(\"Proceesing layer #\" + str(i) + \":\")\n",
    "    preserved_copies = prune_layers(model, [i])\n",
    "    \n",
    "    cls_lst_w1 = get_avg_pool(df_gpt3_wiki[\"gold_completion\"])\n",
    "    cls_lst_w2 = get_avg_pool(df_gpt3_wiki[\"gen_completion\"])\n",
    "    np.save(output_folder + 'avg_roberta_human3_wiki_l' + str(i) + '.npy', cls_lst_w1)\n",
    "    np.save(output_folder + 'avg_roberta_gpt3_wiki_l' + str(i) + '.npy', cls_lst_w2)\n",
    "    \n",
    "    cls_lst_r1 = get_avg_pool(df_gpt3_reddit[\"gold_completion\"])\n",
    "    cls_lst_r2 = get_avg_pool(df_gpt3_reddit[\"gen_completion\"])\n",
    "    np.save(output_folder + 'avg_roberta_human3_reddit_l' + str(i) + '.npy', cls_lst_r1)\n",
    "    np.save(output_folder + 'avg_roberta_gpt3_reddit_l' + str(i) + '.npy', cls_lst_r2)\n",
    "    \n",
    "    cls_lst_s1 = get_avg_pool(df_gpt3_stackex[\"gold_completion\"])\n",
    "    cls_lst_s2 = get_avg_pool(df_gpt3_stackex[\"gen_completion\"])\n",
    "    np.save(output_folder + 'avg_roberta_human3_stackexchange_l' + str(i) + '.npy', cls_lst_s1)\n",
    "    np.save(output_folder + 'avg_roberta_gpt3_stackexchange_l' + str(i) + '.npy', cls_lst_s2)\n",
    "    \n",
    "    cls_lst_w1 = get_avg_pool(df_gpt4o_wiki[\"gold_completion\"])\n",
    "    cls_lst_w2 = get_avg_pool(df_gpt4o_wiki[\"gen_completion\"])\n",
    "    np.save(output_folder + 'avg_roberta_human4o_wiki_l' + str(i) + '.npy', cls_lst_w1)\n",
    "    np.save(output_folder + 'avg_roberta_gpt4o_wiki_l' + str(i) + '.npy', cls_lst_w2)\n",
    "    \n",
    "    cls_lst_r1 = get_avg_pool(df_gpt4o_reddit[\"gold_completion\"])\n",
    "    cls_lst_r2 = get_avg_pool(df_gpt4o_reddit[\"gen_completion\"])\n",
    "    np.save(output_folder + 'avg_roberta_human4o_reddit_l' + str(i) + '.npy', cls_lst_r1)\n",
    "    np.save(output_folder + 'avg_roberta_gpt4o_reddit_l' + str(i) + '.npy', cls_lst_r2)\n",
    "    \n",
    "    cls_lst_s1 = get_avg_pool(df_gpt4o_stackex[\"gold_completion\"])\n",
    "    cls_lst_s2 = get_avg_pool(df_gpt4o_stackex[\"gen_completion\"])\n",
    "    np.save(output_folder + 'avg_roberta_human4o_stackexchange_l' + str(i) + '.npy', cls_lst_s1)\n",
    "    np.save(output_folder + 'avg_roberta_gpt4o_stackexchange_l' + str(i) + '.npy', cls_lst_s2)\n",
    "    \n",
    "    restore_layers(model, preserved_copies)\n",
    "    print('\\n\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### print the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_crossdomen_results(model_name=\"roberta\", config_name=\"none\", C=1):\n",
    "    avg_in = 0\n",
    "    avg_out = 0\n",
    "    avg_3 = 0\n",
    "    \n",
    "    line_prefixes = [\" GPT-3  \", \" GPT-4o \"]\n",
    "    suffixes = [\"3\", \"4o\"]\n",
    "    print(' ' * 21 + \"        GPT-3     \" + \"        GPT-4o    \" )\n",
    "    print(' ' * 21 + \"  Wiki Redd. Stac.\" * 2)\n",
    "    for iteration in range(2):\n",
    "        suffix_t = suffixes[iteration]\n",
    "        suffix_ht = suffix_t\n",
    " \n",
    "        for train in [\"wiki\", \"reddit\", \"stackexchange\"]:\n",
    "            print(line_prefixes[iteration] + train + ' ' * (13 - len(train)), end=\" \")\n",
    "\n",
    "            X_train = np.vstack([\n",
    "                np.load(\"{}avg_{}_human{}_{}_{}.npy\".format(output_folder, model_name, suffix_ht, train, config_name))[:1300],\n",
    "                np.load(\"{}avg_{}_gpt{}_{}_{}.npy\".format(output_folder, model_name, suffix_t, train, config_name))[:1300]\n",
    "            ])\n",
    "            y_train = np.zeros(len(X_train))\n",
    "            y_train[len(X_train) // 2:] = 1\n",
    "            cls = LogisticRegression(max_iter=1000, C=C).fit(X_train, y_train)\n",
    "\n",
    "            for suffix_v in suffixes:\n",
    "                suffix_hv = suffix_v\n",
    "\n",
    "                for valid in [\"wiki\", \"reddit\", \"stackexchange\"]:\n",
    "                    X_valid = np.vstack([\n",
    "                        np.load(\"{}avg_{}_human{}_{}_{}.npy\".format(output_folder, model_name, suffix_hv, valid, config_name))[1300:],\n",
    "                        np.load(\"{}avg_{}_gpt{}_{}_{}.npy\".format(output_folder, model_name, suffix_v, valid, config_name))[1300:]\n",
    "                    ])\n",
    "                    y_val = np.zeros(len(X_valid))\n",
    "                    y_val[len(X_valid) // 2:] = 1\n",
    "\n",
    "                    print(format(cls.score(X_valid, y_val), '.3f'), end=\" \")\n",
    "                    if suffix_v != suffix_t:\n",
    "                        if train == valid:\n",
    "                            avg_out += cls.score(X_valid, y_val) / 6.0\n",
    "                        else:\n",
    "                            avg_3 += cls.score(X_valid, y_val) / 12.0\n",
    "                    else:\n",
    "                        if train != valid:\n",
    "                            avg_in += cls.score(X_valid, y_val) / 12.0      \n",
    "            print(\"\")\n",
    "    print(\"Cross-domain:\", avg_in, \"; Cross-model: \", avg_out, \"; Cross-domain&model: \", avg_3, \"\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "All layers intact. Baseline RoBERTa Cross-domain accuracy: \n",
      "                             GPT-3             GPT-4o    \n",
      "                       Wiki Redd. Stac.  Wiki Redd. Stac.\n",
      " GPT-3  wiki          0.988 0.593 0.965 0.500 0.323 0.495 \n",
      " GPT-3  reddit        0.507 1.000 0.973 0.775 1.000 0.828 \n",
      " GPT-3  stackexchange 0.938 0.820 0.995 0.915 0.780 0.578 \n",
      " GPT-4o wiki          0.527 0.848 0.802 1.000 0.980 0.907 \n",
      " GPT-4o reddit        0.297 0.993 0.920 0.650 0.998 0.968 \n",
      " GPT-4o stackexchange 0.733 0.770 0.667 0.890 0.953 1.000 \n",
      "Cross-domain: 0.8452083333333333 ; Cross-model:  0.7108333333333333 ; Cross-domain&model:  0.7070833333333334 \n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(\"All layers intact. Baseline RoBERTa Cross-domain accuracy: \")    \n",
    "print_crossdomen_results(model_name=\"roberta\", config_name=\"none\", C=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "RoBERTa, heads from layer#1 is prunned; Cross-domain accuracy: \n",
      "                             GPT-3             GPT-4o    \n",
      "                       Wiki Redd. Stac.  Wiki Redd. Stac.\n",
      " GPT-3  wiki          0.988 0.593 0.958 0.530 0.335 0.485 \n",
      " GPT-3  reddit        0.598 1.000 0.980 0.875 1.000 0.772 \n",
      " GPT-3  stackexchange 0.950 0.762 0.990 0.927 0.715 0.608 \n",
      " GPT-4o wiki          0.545 0.935 0.850 1.000 0.985 0.902 \n",
      " GPT-4o reddit        0.390 0.988 0.915 0.750 0.993 0.963 \n",
      " GPT-4o stackexchange 0.795 0.750 0.693 0.877 0.858 1.000 \n",
      "Cross-domain: 0.8479166666666667 ; Cross-model:  0.7270833333333333 ; Cross-domain&model:  0.72875 \n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(\"RoBERTa, heads from layer#1 is prunned; Cross-domain accuracy: \")    \n",
    "print_crossdomen_results(model_name=\"roberta\", config_name=\"l0\", C=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "RoBERTa, heads from layer#4 is prunned; Cross-domain accuracy: \n",
      "                             GPT-3             GPT-4o    \n",
      "                       Wiki Redd. Stac.  Wiki Redd. Stac.\n",
      " GPT-3  wiki          0.985 0.580 0.960 0.512 0.355 0.492 \n",
      " GPT-3  reddit        0.522 1.000 0.973 0.787 1.000 0.838 \n",
      " GPT-3  stackexchange 0.932 0.850 0.993 0.917 0.812 0.632 \n",
      " GPT-4o wiki          0.532 0.875 0.772 1.000 0.980 0.907 \n",
      " GPT-4o reddit        0.325 0.985 0.905 0.672 0.995 0.963 \n",
      " GPT-4o stackexchange 0.745 0.818 0.675 0.887 0.955 0.998 \n",
      "Cross-domain: 0.8485416666666667 ; Cross-model:  0.7229166666666667 ; Cross-domain&model:  0.7202083333333333 \n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(\"RoBERTa, heads from layer#4 is prunned; Cross-domain accuracy: \")    \n",
    "print_crossdomen_results(model_name=\"roberta\", config_name=\"l3\", C=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}