heads_prunning.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "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:0'\n",
    "    print(\"Cuda computing enabled\")\n",
    "else:\n",
    "    device = 'cpu'    \n",
    "    print(\"No cuda found\")\n",
    "    \n",
    "output_folder = \"cache/\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading the model and data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "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": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_gpt3_wiki = pd.read_json(\"gpt3_davinci_003_wikip.jsonl_pp\", lines=True)[:1500]\n",
    "df_gpt3_reddit = pd.read_json(\"gpt3_davinci_003_reddit.jsonl_pp\", lines=True)[:1500]\n",
    "df_gpt3_stackex = pd.read_json(\"gpt3_davinci_003_300_len.jsonl\", lines=True)[:1500]\n",
    "df_gpt4o_wiki = pd.read_csv(\"gpt-4o/gpt-4-o-wiki-correct-1500.csv\")\n",
    "df_gpt4o_reddit = pd.read_csv(\"gpt-4o/gpt-4-o-reddit-1500.csv\")\n",
    "df_gpt4o_stackex = pd.read_csv(\"gpt-4o/gpt-4-o-stackexchange-1500.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_avg_pool(series, flag=True):\n",
    "    \"Returns mean-pooled embeddings for each text in the container <series> \"\n",
    "    cls_set = []\n",
    "    for text in series:\n",
    "        if flag and len(text) < 25:\n",
    "            continue\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",
    "        cls_set.append(torch.mean(outp, 0).cpu().numpy())\n",
    "\n",
    "        if flag and len(cls_set) >= 600:\n",
    "            break\n",
    "\n",
    "    return np.vstack(cls_set)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### auxillary functions for the head search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calc_crossdomen_results(model_name=\"roberta\", config_name=\"none\", C=1):\n",
    "    avg = 0\n",
    "    suffixes = [\"3\", \"4o\"]\n",
    "    for iteration in range(len(suffixes)):\n",
    "        suffix_t = suffixes[iteration]\n",
    "        for train in [\"wiki\", \"reddit\", \"stackexchange\"]:\n",
    "            X_train = np.vstack([\n",
    "                np.load(\"{}avg_{}_human{}_{}_{}.npy\".format(output_folder, model_name, suffix_t, train, config_name))[:400],\n",
    "                np.load(\"{}avg_{}_gpt{}_{}_{}.npy\".format(output_folder, model_name, suffix_t, train, config_name))[:400]\n",
    "            ])\n",
    "            y_train = np.zeros(len(X_train))\n",
    "            y_train[len(X_train) // 2:] = 1\n",
    "\n",
    "            cls = LogisticRegression(max_iter=1000, C=C).fit(X_train, y_train)\n",
    "\n",
    "            for suffix_v in suffixes:\n",
    "                for valid in [\"wiki\", \"reddit\", \"stackexchange\"]:\n",
    "                    X_valid = np.vstack([\n",
    "                        np.load(\"{}avg_{}_human{}_{}_{}.npy\".format(output_folder, model_name, suffix_v, valid, config_name))[:-200],\n",
    "                        np.load(\"{}avg_{}_gpt{}_{}_{}.npy\".format(output_folder, model_name, suffix_v, valid, config_name))[:-200]\n",
    "                    ])\n",
    "                    y_val = np.zeros(len(X_valid))\n",
    "                    y_val[len(X_valid) // 2:] = 1\n",
    "\n",
    "                    if train != valid or suffix_v != suffix_t:\n",
    "                        avg += cls.score(X_valid, y_val) / 30.0      \n",
    "    return avg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calc_quality():\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_tmp.npy', cls_lst_w1)\n",
    "    np.save(output_folder + 'avg_roberta_gpt4o_wiki_tmp.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_tmp.npy', cls_lst_r1)\n",
    "    np.save(output_folder + 'avg_roberta_gpt4o_reddit_tmp.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_tmp.npy', cls_lst_s1)\n",
    "    np.save(output_folder + 'avg_roberta_gpt4o_stackexchange_tmp.npy', cls_lst_s2)\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_tmp.npy', cls_lst_w1)\n",
    "    np.save(output_folder + 'avg_roberta_gpt3_wiki_tmp.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_tmp.npy', cls_lst_r1)\n",
    "    np.save(output_folder + 'avg_roberta_gpt3_reddit_tmp.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_tmp.npy', cls_lst_s1)\n",
    "    np.save(output_folder + 'avg_roberta_gpt3_stackexchange_tmp.npy', cls_lst_s2)\n",
    "    \n",
    "    results= calc_crossdomen_results(model_name=\"roberta\", config_name=\"tmp\")\n",
    "\n",
    "    return results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### greedy search for the best set of heads"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "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"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.5425 {0: [7, 0]}\n"
     ]
    },
    {
     "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",
      "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",
      "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",
      "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"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.54375 {0: [7, 0, 4]}\n"
     ]
    },
    {
     "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"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-8-d0e18d904bff>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     20\u001b[0m             \u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mprune_heads\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m{\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mj\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m}\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     21\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 22\u001b[0;31m             \u001b[0mqual\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcalc_quality\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     23\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mqual\u001b[0m \u001b[0;34m>\u001b[0m \u001b[0mbest_result\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     24\u001b[0m                 \u001b[0mbest_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mqual\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-7-d36c914ea43b>\u001b[0m in \u001b[0;36mcalc_quality\u001b[0;34m()\u001b[0m\n\u001b[1;32m     23\u001b[0m     \u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msave\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moutput_folder\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;34m'avg_roberta_gpt4o_stackexchange_tmp.npy'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcls_lst_s2\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     24\u001b[0m     \"\"\"\n\u001b[0;32m---> 25\u001b[0;31m     \u001b[0mcls_lst_w1\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mget_avg_pool\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdf_gpt3_wiki\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m\"gold_completion\"\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     26\u001b[0m     \u001b[0mcls_lst_w2\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mget_avg_pool\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdf_gpt3_wiki\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m\"gen_completion\"\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     27\u001b[0m     \u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msave\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moutput_folder\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;34m'avg_rob_human3_wiki_tmp.npy'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcls_lst_w1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-5-ec9333871421>\u001b[0m in \u001b[0;36mget_avg_pool\u001b[0;34m(series, flag)\u001b[0m\n\u001b[1;32m      8\u001b[0m         \u001b[0minpt\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtokenizer\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtext\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtruncation\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmax_length\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m510\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mreturn_tensors\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m\"pt\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      9\u001b[0m         \u001b[0;32mwith\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mno_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 10\u001b[0;31m             \u001b[0moutp\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m**\u001b[0m\u001b[0minpt\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     11\u001b[0m         \u001b[0mcls_set\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmean\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moutp\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcpu\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnumpy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     12\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1509\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_compiled_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[misc]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1510\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1511\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1512\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1513\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1518\u001b[0m                 \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_pre_hooks\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_hooks\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1519\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1520\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1521\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1522\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/transformers/models/roberta/modeling_roberta.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input_ids, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, encoder_hidden_states, encoder_attention_mask, past_key_values, use_cache, output_attentions, output_hidden_states, return_dict)\u001b[0m\n\u001b[1;32m    833\u001b[0m             \u001b[0mpast_key_values_length\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mpast_key_values_length\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    834\u001b[0m         )\n\u001b[0;32m--> 835\u001b[0;31m         encoder_outputs = self.encoder(\n\u001b[0m\u001b[1;32m    836\u001b[0m             \u001b[0membedding_output\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    837\u001b[0m             \u001b[0mattention_mask\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mextended_attention_mask\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1509\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_compiled_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[misc]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1510\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1511\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1512\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1513\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1518\u001b[0m                 \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_pre_hooks\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_hooks\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1519\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1520\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1521\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1522\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/transformers/models/roberta/modeling_roberta.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, hidden_states, attention_mask, head_mask, encoder_hidden_states, encoder_attention_mask, past_key_values, use_cache, output_attentions, output_hidden_states, return_dict)\u001b[0m\n\u001b[1;32m    522\u001b[0m                 )\n\u001b[1;32m    523\u001b[0m             \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 524\u001b[0;31m                 layer_outputs = layer_module(\n\u001b[0m\u001b[1;32m    525\u001b[0m                     \u001b[0mhidden_states\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    526\u001b[0m                     \u001b[0mattention_mask\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1509\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_compiled_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[misc]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1510\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1511\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1512\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1513\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1518\u001b[0m                 \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_pre_hooks\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_hooks\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1519\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1520\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1521\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1522\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/transformers/models/roberta/modeling_roberta.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, hidden_states, attention_mask, head_mask, encoder_hidden_states, encoder_attention_mask, past_key_value, output_attentions)\u001b[0m\n\u001b[1;32m    453\u001b[0m             \u001b[0mpresent_key_value\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpresent_key_value\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mcross_attn_present_key_value\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    454\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 455\u001b[0;31m         layer_output = apply_chunking_to_forward(\n\u001b[0m\u001b[1;32m    456\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfeed_forward_chunk\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mchunk_size_feed_forward\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mseq_len_dim\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mattention_output\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    457\u001b[0m         )\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/transformers/pytorch_utils.py\u001b[0m in \u001b[0;36mapply_chunking_to_forward\u001b[0;34m(forward_fn, chunk_size, chunk_dim, *input_tensors)\u001b[0m\n\u001b[1;32m    204\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    205\u001b[0m     \u001b[0;31m# inspect.signature exist since python 3.5 and is a python method -> no problem with backward compatibility\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 206\u001b[0;31m     \u001b[0mnum_args_in_forward_chunk_fn\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minspect\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msignature\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mforward_fn\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mparameters\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    207\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mnum_args_in_forward_chunk_fn\u001b[0m \u001b[0;34m!=\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput_tensors\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    208\u001b[0m         raise ValueError(\n",
      "\u001b[0;32m/usr/lib/python3.8/inspect.py\u001b[0m in \u001b[0;36msignature\u001b[0;34m(obj, follow_wrapped)\u001b[0m\n\u001b[1;32m   3103\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0msignature\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mobj\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfollow_wrapped\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3104\u001b[0m     \u001b[0;34m\"\"\"Get a signature object for the passed callable.\"\"\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 3105\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0mSignature\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfrom_callable\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mobj\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfollow_wrapped\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mfollow_wrapped\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   3106\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3107\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/lib/python3.8/inspect.py\u001b[0m in \u001b[0;36mfrom_callable\u001b[0;34m(cls, obj, follow_wrapped)\u001b[0m\n\u001b[1;32m   2852\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mfrom_callable\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcls\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mobj\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfollow_wrapped\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2853\u001b[0m         \u001b[0;34m\"\"\"Constructs Signature for the given callable object.\"\"\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2854\u001b[0;31m         return _signature_from_callable(obj, sigcls=cls,\n\u001b[0m\u001b[1;32m   2855\u001b[0m                                         follow_wrapper_chains=follow_wrapped)\n\u001b[1;32m   2856\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/lib/python3.8/inspect.py\u001b[0m in \u001b[0;36m_signature_from_callable\u001b[0;34m(obj, follow_wrapper_chains, skip_bound_arg, sigcls)\u001b[0m\n\u001b[1;32m   2231\u001b[0m         \u001b[0;31m# In this case we skip the first parameter of the underlying\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2232\u001b[0m         \u001b[0;31m# function (usually `self` or `cls`).\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2233\u001b[0;31m         sig = _signature_from_callable(\n\u001b[0m\u001b[1;32m   2234\u001b[0m             \u001b[0mobj\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__func__\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2235\u001b[0m             \u001b[0mfollow_wrapper_chains\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mfollow_wrapper_chains\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/lib/python3.8/inspect.py\u001b[0m in \u001b[0;36m_signature_from_callable\u001b[0;34m(obj, follow_wrapper_chains, skip_bound_arg, sigcls)\u001b[0m\n\u001b[1;32m   2244\u001b[0m     \u001b[0;31m# Was this function wrapped by a decorator?\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2245\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mfollow_wrapper_chains\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2246\u001b[0;31m         \u001b[0mobj\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0munwrap\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mobj\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstop\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;32mlambda\u001b[0m \u001b[0mf\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mhasattr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"__signature__\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   2247\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mobj\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtypes\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mMethodType\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2248\u001b[0m             \u001b[0;31m# If the unwrapped object is a *method*, we might want to\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/lib/python3.8/inspect.py\u001b[0m in \u001b[0;36munwrap\u001b[0;34m(func, stop)\u001b[0m\n\u001b[1;32m    516\u001b[0m     \u001b[0;31m# Memoise by id to tolerate non-hashable objects, but store objects to\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    517\u001b[0m     \u001b[0;31m# ensure they aren't destroyed, which would allow their IDs to be reused.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 518\u001b[0;31m     \u001b[0mmemo\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m{\u001b[0m\u001b[0mid\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mf\u001b[0m\u001b[0;34m}\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    519\u001b[0m     \u001b[0mrecursion_limit\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0msys\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgetrecursionlimit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    520\u001b[0m     \u001b[0;32mwhile\u001b[0m \u001b[0m_is_wrapper\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfunc\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "We use some sort of a greedy Depth-first search on a space of all subsets of heads (strictly speaking, ordered subsets)\n",
    "\"\"\"\n",
    "\n",
    "heads = {}\n",
    "current_heads_stack = []\n",
    "best_result = 0.0\n",
    "best_heads = {}\n",
    "start_i, start_j = 0, 0\n",
    "\n",
    "while True:\n",
    "    path_has_been_continued = False\n",
    "    for i in range(start_i, 12):\n",
    "        for j in range(start_j, 12):\n",
    "            if i in heads.keys() and j in heads[i]:\n",
    "                continue\n",
    "\n",
    "            model = RobertaModel.from_pretrained(MODEL_NAME, cache_dir=CACHE_PATH).to(device)\n",
    "            model.prune_heads(heads)\n",
    "            model.prune_heads({i:[j]})\n",
    "\n",
    "            qual = calc_quality()\n",
    "            if qual > best_result:\n",
    "                best_result = qual\n",
    "                if i in heads.keys():\n",
    "                    heads[i].append(j)\n",
    "                else:\n",
    "                    heads[i] = [j]\n",
    "                    \n",
    "                path_has_been_continued = True\n",
    "                current_heads_stack.append((i, j))\n",
    "                best_heads = deepcopy(heads)\n",
    "                print(best_result, heads)\n",
    "                \n",
    "    if path_has_been_continued:\n",
    "        print(\"Iteration finished succesfully\")\n",
    "        start_i, start_j = 0, 0\n",
    "        continue\n",
    "    \n",
    "    print(\"No extension has been found, backtracking\")\n",
    "    if len(current_heads_stack) == 0:\n",
    "        break\n",
    "    last_vert = current_heads_stack[-1]\n",
    "    current_heads_stack = current_heads_stack[:-1]\n",
    "    \n",
    "    start_i, start_j = last_vert\n",
    "    head[start_i] = [XXX for XXX in head[start_i] if XXX != start_j]\n",
    "    \n",
    "print(\"Search completely exhausted\")\n",
    "print(best_result, best_heads)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### print the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "best_heads_roberta = {\n",
    "    0: [0, 1, 4, 6, 8, 9], \n",
    "    1: [3, 4, 6, 8, 9, 10], \n",
    "    2: [0, 1, 2, 7], \n",
    "    3: [6, 7, 10, 11], \n",
    "    4: [2, 4, 5, 6, 11], \n",
    "    5: [0, 7, 9, 10], \n",
    "    6: [0, 10], \n",
    "    7: [4, 6, 8], \n",
    "    8: [4, 7, 10, 11], \n",
    "    9: [0, 3, 4, 6, 8, 10], \n",
    "    10: [1, 4, 5, 7, 10], \n",
    "    11: [3, 6, 7]}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = RobertaModel.from_pretrained(MODEL_NAME, cache_dir=CACHE_PATH).to(device).to(device)\n",
    "model.prune_heads(heads)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cls_lst_w1 = get_avg_pool(df_gpt4o_wiki[\"gold_completion\"], False)\n",
    "cls_lst_w2 = get_avg_pool(df_gpt4o_wiki[\"gen_completion\"], False)    \n",
    "np.save(output_folder + 'avg_rob_human4o_wiki_tmp2.npy', cls_lst_w1)\n",
    "np.save(output_folder + 'avg_rob_gpt4o_wiki_tmp2.npy', cls_lst_w2)\n",
    "\n",
    "cls_lst_r1 = get_avg_pool(df_gpt4o_reddit[\"gold_completion\"], False)\n",
    "cls_lst_r2 = get_avg_pool(df_gpt4o_reddit[\"gen_completion\"], False)\n",
    "np.save(output_folder + 'avg_rob_human4o_reddit_tmp2.npy', cls_lst_r1)\n",
    "np.save(output_folder + 'avg_rob_gpt4o_reddit_tmp2.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_rob_human4o_stackexchange_tmp2.npy', cls_lst_s1)\n",
    "np.save(output_folder + 'avg_rob_gpt4o_stackexchange_tmp2.npy', cls_lst_s2)\n",
    "\n",
    "cls_lst_w1 = get_avg_pool(df_gpt3_wiki[\"gold_completion\"], False)\n",
    "cls_lst_w2 = get_avg_pool(df_gpt3_wiki[\"gen_completion\"], False)    \n",
    "np.save(output_folder + 'avg_rob_human3_wiki_tmp2.npy', cls_lst_w1)\n",
    "np.save(output_folder + 'avg_rob_gpt3_wiki_tmp2.npy', cls_lst_w2)\n",
    "\n",
    "cls_lst_r1 = get_avg_pool(df_gpt3_reddit[\"gold_completion\"], False)\n",
    "cls_lst_r2 = get_avg_pool(df_gpt3_reddit[\"gen_completion\"], False)\n",
    "np.save(output_folder + 'avg_rob_human3_reddit_tmp2.npy', cls_lst_r1)\n",
    "np.save(output_folder + 'avg_rob_gpt3_reddit_tmp2.npy', cls_lst_r2)\n",
    "    \n",
    "cls_lst_s1 = get_avg_pool(df_gpt3_stackex[\"gold_completion\"], False)\n",
    "cls_lst_s2 = get_avg_pool(df_gpt3_stackex[\"gen_completion\"], False)\n",
    "np.save(output_folder + 'avg_rob_human3_stackexchange_tmp2.npy', cls_lst_s1)\n",
    "np.save(output_folder + 'avg_rob_gpt3_stackexchange_tmp2.npy', cls_lst_s2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print_crossdomen_results(model_name=\"rob\", config_name=\"tmp2\", 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
}