extract_lstm_features.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a40ac286",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "pd.options.display.max_colwidth = 100\n",
    "pd.set_option('display.max_colwidth', None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "712161ea",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using cuda device\n"
     ]
    }
   ],
   "source": [
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "print(f\"Using {device} device\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9720c10",
   "metadata": {},
   "outputs": [],
   "source": [
    "#!pip install -Uq fastai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "8116c59a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0f9f2d1e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'2.6.3'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import fastai\n",
    "fastai.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c91eefd1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from fastai.basics import *\n",
    "from fastai.callback.all import *\n",
    "from fastai.text.all import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "31d6efa6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import Dataset\n",
    "from torch.utils.data import DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "9f20ed78",
   "metadata": {},
   "outputs": [],
   "source": [
    "path = untar_data(URLs.WIKITEXT)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a13b6fb5",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_train = pd.read_csv(path/'train.csv', header=None)\n",
    "df_valid = pd.read_csv(path/'test.csv', header=None)\n",
    "df_all = pd.concat([df_train, df_valid])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "c5447483",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "<style>\n",
       "    /* Turns off some styling */\n",
       "    progress {\n",
       "        /* gets rid of default border in Firefox and Opera. */\n",
       "        border: none;\n",
       "        /* Needs to be in here for Safari polyfill so background images work as expected. */\n",
       "        background-size: auto;\n",
       "    }\n",
       "    .progress-bar-interrupted, .progress-bar-interrupted::-webkit-progress-bar {\n",
       "        background: #F44336;\n",
       "    }\n",
       "</style>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "splits = [list(range_of(df_train)), list(range(len(df_train), len(df_all)))]\n",
    "tfms = [attrgetter(\"text\"), Tokenizer.from_df(0), Numericalize()]\n",
    "dsets = Datasets(df_all, [tfms], splits=splits, dl_type=LMDataLoader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "8d1963c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "bs,sl = 5,512\n",
    "dls = dsets.dataloaders(bs=bs, seq_len=sl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "b915f141",
   "metadata": {},
   "outputs": [],
   "source": [
    "lm = language_model_learner(dls, AWD_LSTM, drop_mult=0.5, metrics=Perplexity(), pretrained=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "c6c8c3f1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "<style>\n",
       "    /* Turns off some styling */\n",
       "    progress {\n",
       "        /* gets rid of default border in Firefox and Opera. */\n",
       "        border: none;\n",
       "        /* Needs to be in here for Safari polyfill so background images work as expected. */\n",
       "        background-size: auto;\n",
       "    }\n",
       "    .progress-bar-interrupted, .progress-bar-interrupted::-webkit-progress-bar {\n",
       "        background: #F44336;\n",
       "    }\n",
       "</style>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "(#2) [3.2393643856048584,25.517498016357422]"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lm.validate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "3507b0e9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SequentialRNN(\n",
       "  (0): AWD_LSTM(\n",
       "    (encoder): Embedding(60008, 400, padding_idx=1)\n",
       "    (encoder_dp): EmbeddingDropout(\n",
       "      (emb): Embedding(60008, 400, padding_idx=1)\n",
       "    )\n",
       "    (rnns): ModuleList(\n",
       "      (0): WeightDropout(\n",
       "        (module): LSTM(400, 1152, batch_first=True)\n",
       "      )\n",
       "      (1): WeightDropout(\n",
       "        (module): LSTM(1152, 1152, batch_first=True)\n",
       "      )\n",
       "      (2): WeightDropout(\n",
       "        (module): LSTM(1152, 400, batch_first=True)\n",
       "      )\n",
       "    )\n",
       "    (input_dp): RNNDropout()\n",
       "    (hidden_dps): ModuleList(\n",
       "      (0): RNNDropout()\n",
       "      (1): RNNDropout()\n",
       "      (2): RNNDropout()\n",
       "    )\n",
       "  )\n",
       "  (1): LinearDecoder(\n",
       "    (decoder): Linear(in_features=400, out_features=60008, bias=True)\n",
       "    (output_dp): RNNDropout()\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lm.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3c15ae13",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_train_final = pd.read_json(\"train_final.json\")\n",
    "df_val_final = pd.read_json(\"val_final.json\")\n",
    "df_test_final = pd.read_json(\"test_final.json\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "44ab51c0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The size of the lm train dataset is 70749\n",
      "The size of the lm val dataset is 18223\n",
      "The size of the lm test dataset is 38513\n"
     ]
    }
   ],
   "source": [
    "print(\"The size of the lm train dataset is {}\".format(len(df_train_final)))\n",
    "print(\"The size of the lm val dataset is {}\".format(len(df_val_final)))\n",
    "print(\"The size of the lm test dataset is {}\".format(len(df_test_final)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9c310d90",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_all_final= pd.concat([df_train_final, df_val_final, df_test_final])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d72be394",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "positive training examples are 57348 and negative training examples are 70137\n"
     ]
    }
   ],
   "source": [
    "print(\"positive training examples are {} and negative training examples are {}\"\n",
    "      .format(len(df_all_final[df_all_final['label']==1]), len(df_all_final[df_all_final['label']==0])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "d8df2798",
   "metadata": {},
   "outputs": [],
   "source": [
    "max_seq_len = 512\n",
    "batch_size = 128"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "5ba030df",
   "metadata": {},
   "outputs": [],
   "source": [
    "def tokenize_fastai(x, len_df):\n",
    "    tokenized = np.zeros((len_df, max_seq_len))\n",
    "    for idx, c in tqdm(enumerate(x)):\n",
    "        tokens = np.array(dsets.numericalize(dsets.tokenizer(c)).tolist())[:512]\n",
    "        tokenized[idx, max_seq_len-len(tokens):] = tokens\n",
    "    \n",
    "    return tokenized  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "4f67e9d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def rep_fastai(tokenized, len_df):\n",
    "    count = 0\n",
    "    rep_list = []\n",
    "    while count < len_df:\n",
    "        rep_list.append(lm.model[0](\n",
    "            torch.Tensor(tokenized[count:count+batch_size]).long().to(device)).detach().cpu().numpy()[:,-1,:])\n",
    "        count += batch_size\n",
    "        #print(count)\n",
    "        \n",
    "    return torch.tensor(np.vstack(rep_list))   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "c85f9d02",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "70749it [04:18, 273.38it/s]\n",
      "70749it [00:25, 2788.64it/s]\n",
      "18223it [01:06, 273.24it/s]\n",
      "18223it [00:06, 2921.08it/s]\n",
      "38513it [02:20, 273.61it/s]\n",
      "38513it [00:13, 2841.38it/s]\n"
     ]
    }
   ],
   "source": [
    "tokenized_context_train = tokenize_fastai(df_train_final['context'], len(df_train_final))\n",
    "tokenized_query_train= tokenize_fastai(df_train_final['query'], len(df_train_final))\n",
    "\n",
    "tokenized_context_val = tokenize_fastai(df_val_final['context'], len(df_val_final))\n",
    "tokenized_query_val= tokenize_fastai(df_val_final['query'], len(df_val_final))\n",
    "\n",
    "tokenized_context_test = tokenize_fastai(df_test_final['context'], len(df_test_final))\n",
    "tokenized_query_test = tokenize_fastai(df_test_final['query'], len(df_test_final))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "1000f7a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_context_train = rep_fastai(tokenized_context_train, len(df_train_final))\n",
    "X_query_train = rep_fastai(tokenized_query_train, len(df_train_final))\n",
    "\n",
    "X_context_val = rep_fastai(tokenized_context_val, len(df_val_final))\n",
    "X_query_val = rep_fastai(tokenized_query_val, len(df_val_final))\n",
    "\n",
    "X_context_test = rep_fastai(tokenized_context_test, len(df_test_final))\n",
    "X_query_test = rep_fastai(tokenized_query_test, len(df_test_final))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "fcb58b05",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([38513, 400])"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X_query_test.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "a8e0ce73",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(X_context_train, 'train_context_lstm_reps.pt')\n",
    "torch.save(X_context_val, 'val_context_lstm_reps.pt')\n",
    "torch.save(X_context_test, 'test_context_lstm_reps.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "3b4971f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(X_query_train, 'train_query_lstm_reps.pt')\n",
    "torch.save(X_query_val, 'val_query_lstm_reps.pt')\n",
    "torch.save(X_query_test, 'test_query_lstm_reps.pt')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0195c928",
   "metadata": {},
   "source": [
    "# suprisal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "16ea9d5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "id": "b830688e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "8da44dac",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "38513it [02:07, 301.40it/s]\n",
      "38513it [00:12, 2983.38it/s]\n"
     ]
    }
   ],
   "source": [
    "tokenized_context_test = tokenize_fastai(df_test_final['context'], len(df_test_final))\n",
    "tokenized_query_test = tokenize_fastai(df_test_final['query'], len(df_test_final))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "id": "672cce9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def surp_fastai(tokenized, query_id=None):\n",
    "    count = 0\n",
    "    surp = np.zeros(len(tokenized))\n",
    "    \n",
    "    with tqdm(total=len(tokenized)) as pbar:\n",
    "        while count < len(tokenized):\n",
    "            log_prob = -F.log_softmax(\n",
    "                lm.model(\n",
    "                torch.Tensor(tokenized[count:count+batch_size]).long().to(device))[0][:,-1].detach().cpu(),\n",
    "                dim=1).gather(dim=1, index=torch.Tensor(query_id[count:count+batch_size]).long().view(-1,1))\n",
    "            #print(log_prob)\n",
    "            surp[count:count+batch_size] = log_prob.view(-1).numpy()\n",
    "            count += batch_size\n",
    "            pbar.update(batch_size)\n",
    "        \n",
    "    return surp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "id": "081a8928",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "38528it [00:48, 797.65it/s]                                                                                                                                                                               \n"
     ]
    }
   ],
   "source": [
    "surp_lstm_test = surp_fastai(tokenized_context_test, tokenized_query_test[:,-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "id": "dc1da9e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(surp_lstm_test, 'surp_lstm_test.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7625e032",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}