src/models/reranker/reranker-show.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ded10377",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "__init__.py   arguments.py  \u001b[0m\u001b[01;34mdist\u001b[0m/        trainer.py\r\n",
      "\u001b[01;34m__pycache__\u001b[0m/  data.py       modeling.py\r\n"
     ]
    }
   ],
   "source": [
    "ls Reranker/src/reranker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8f2eca39",
   "metadata": {},
   "outputs": [],
   "source": [
    "from Reranker.src.reranker.arguments import ModelArguments, DataArguments, \\\n",
    "    RerankerTrainingArguments as TrainingArguments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "efff1461",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e2fa027b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from reranker import Reranker \n",
    "from transformers import AutoTokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "015dfe05",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Optional\n",
    "\n",
    "import torch\n",
    "import torch.functional as F\n",
    "import copy\n",
    "from transformers import AutoModelForSequenceClassification, AutoTokenizer,\\\n",
    "    PreTrainedModel, PreTrainedTokenizer, GPTNeoForSequenceClassification\n",
    "\n",
    "from transformers.modeling_outputs import SequenceClassifierOutput, BaseModelOutputWithPooling\n",
    "from torch import nn\n",
    "import torch.distributed as dist\n",
    "\n",
    "class RerankerForInference(nn.Module):\n",
    "    def __init__(\n",
    "            self,\n",
    "            hf_model: Optional[PreTrainedModel] = None,\n",
    "            tokenizer: Optional[PreTrainedTokenizer] = None\n",
    "    ):\n",
    "        super().__init__()\n",
    "        self.hf_model = hf_model\n",
    "        self.tokenizer = tokenizer\n",
    "\n",
    "    def tokenize(self, *args, **kwargs):\n",
    "        return self.tokenizer(*args, **kwargs)\n",
    "\n",
    "    def forward(self, batch):\n",
    "        return self.hf_model(**batch)\n",
    "\n",
    "    @classmethod\n",
    "    def from_pretrained(cls, pretrained_model_name_or_path: str):\n",
    "        hf_model = GPTNeoForSequenceClassification.from_pretrained(\n",
    "            pretrained_model_name_or_path,num_labels=1)\n",
    "        hf_tokenizer = AutoTokenizer.from_pretrained(\"EleutherAI/gpt-neo-125M\",    \n",
    "                            bos_token=\"<|startoftext|>\",\n",
    "                            eos_token=\"<|endoftext|>\",\n",
    "                            pad_token=\"<|pad|>\")\n",
    "        hf_model.cuda()\n",
    "        hf_model.eval()\n",
    "        return cls(hf_model, hf_tokenizer)\n",
    "\n",
    "    def load_pretrained_model(self, pretrained_model_name_or_path, *model_args, **kwargs):\n",
    "        self.hf_model = AutoModelForSequenceClassification.from_pretrained(\n",
    "            pretrained_model_name_or_path, *model_args, **kwargs\n",
    "        )\n",
    "\n",
    "    def load_pretrained_tokenizer(self, pretrained_model_name_or_path, *inputs, **kwargs):\n",
    "        self.tokenizer = AutoTokenizer.from_pretrained(\n",
    "            pretrained_model_name_or_path, *inputs, **kwargs\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "01041eba",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_args ={\n",
    "    \"train_dir\"\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "477fae25",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of GPTNeoForSequenceClassification were not initialized from the model checkpoint at results/checkpoint-140000/ and are newly initialized: ['score.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n",
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n",
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    }
   ],
   "source": [
    "rk = RerankerForInference.from_pretrained(\"results/checkpoint-140000/\") "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "3aa5a124",
   "metadata": {},
   "outputs": [],
   "source": [
    "# rk.hf_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ecac7e65",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "d88c1af0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "with open('all_gpt_combined.json', 'r') as json_file:\n",
    "    json_c = [json.loads(i) for i in list(json_file)][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "7cfc0b5f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "7128"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(json_c)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "b558698c",
   "metadata": {},
   "outputs": [],
   "source": [
    "exa = json_c[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "4ec633bf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.5"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "exa[\"score\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "3a330205",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'cases i'"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "exa[\"best_match\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "e44d9562",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'cases b'"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "exa[\"real\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "e4a27c2f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'1'"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "str(hash(1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "3d44b938",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_embed(exa):\n",
    "    if exa[\"score\"]!=0:\n",
    "        positive =  [exa[\"best_match\"],exa[\"real\"]]\n",
    "        negative = [i for i in exa[\"all\"] if i not in positive]\n",
    "    else:\n",
    "        negative = exa[\"all\"]\n",
    "        positive =  [exa[\"best_match\"]]\n",
    "\n",
    "    positive_token = [rk.tokenize(i, return_tensors='pt') for i in positive]\n",
    "    negative_token = [rk.tokenize(i, return_tensors='pt') for i in negative]\n",
    "    prompt = exa[\"prompt\"].replace(\"GOAL\",\"<GOAL>\").replace(\"PROOFSTEP\",\"<PROOFSTEP>\")\n",
    "    prompt_token = rk.tokenize(prompt, return_tensors='pt').input_ids\n",
    "    idi = str(hash(prompt))\n",
    "    return { \"qry\": {  \"qid\": idi,  \"query\": prompt_token, }, \"pos\": [ {\"pid\": str(hash(i)),  \"passage\": j  } for i,j in zip(positive,positive_token)]        ,\"neg\": [ {\"pid\": str(hash(i)),  \"passage\": j  } for i,j in zip(negative,negative_token)]}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "522b13be",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 102,
   "id": "3cb57ff0",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100% 7128/7128 [00:13<00:00, 519.24it/s]\n"
     ]
    }
   ],
   "source": [
    "json_c_train = []\n",
    "\n",
    "for i in tqdm(json_c):\n",
    "    json_c_train.append(get_embed(i))\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ef34bf3",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "5fe0d464",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'best_match': 'cases i',\n",
       " 'score': 0.5,\n",
       " 'real': 'cases b',\n",
       " 'all': ['cases i',\n",
       "  \"apply b.reverse_corec'\",\n",
       "  'apply b.append_right_',\n",
       "  'apply b',\n",
       "  'apply fin',\n",
       "  'dsimp [read] at h',\n",
       "  'apply_instance',\n",
       "  'simp! * at *',\n",
       "  'dsimp only [read, dsimp] at h',\n",
       "  \"set i := b.reverse_core i with h'\",\n",
       "  'apply b.append_left_inj',\n",
       "  'rw [size, h]',\n",
       "  '⊢ b.read',\n",
       "  'apply mem',\n",
       "  'simp [append_right]',\n",
       "  'delta read',\n",
       "  'simp [size, h]',\n",
       "  'subst i',\n",
       "  \"rw [array.read_write', array.write'_write, buffer.write_mem, array.write_mem, h]\",\n",
       "  'rw [←append_right_inj]',\n",
       "  'simp',\n",
       "  'dsimp [read]',\n",
       "  \"apply b.append_right'\",\n",
       "  'set_of_eq_eq_succ_of_eq',\n",
       "  \"convert b.read'\",\n",
       "  'contradiction',\n",
       "  'rw [←append_right_inverse_iff, ←',\n",
       "  \"apply b.read'_eq_of_mem\",\n",
       "  '_match : ∀ (_a _a_1 :append (list α) ⟨i,',\n",
       "  'apply b.append_right',\n",
       "  'apply b.append_left_injective',\n",
       "  '_inst : inhabited α :=',\n",
       "  'apply b.reverse_core_eq',\n",
       "  'convert read_inj b',\n",
       "  '⊢ b.reverse_'],\n",
       " 'prompt': \"GOAL α : Type u,\\t_inst_1 : inhabited α,\\tb : buffer α,\\ti : ℕ,\\th : i < b.size\\t⊢ b.read ⟨i, h⟩ = b.read' i\\n PROOFSTEP \"}"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "json_c[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad4b293b",
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = rk.tokenize(json_c[\"promp\"], 'it is cold today in new york', return_tensors='pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "8c40b74e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'qry': {'qid': '-6574462426148653399',\n",
       "  'query': tensor([[   27, 11230,  1847,    29, 26367,  1058,  5994,   334,    11,   197,\n",
       "              62,  8625,    62,    16,  1058, 30671, 26367,    11,   197,    65,\n",
       "            1058, 11876, 26367,    11,   197,    72,  1058,  2343,   226,   243,\n",
       "              11,   197,    71,  1058,  1312,  1279,   275,    13,  7857,   197,\n",
       "             158,   232,    95,   275,    13,   961,  2343,   253,   101,    72,\n",
       "              11,   289,   158,   253,   102,   796,   275,    13,   961,     6,\n",
       "            1312,   198,  1279,  4805,  6684,    37, 42135,    29,   220]])},\n",
       " 'pos': [{'pid': '8592327939438459529',\n",
       "   'passage': {'input_ids': tensor([[33964,  1312]]), 'attention_mask': tensor([[1, 1]])}},\n",
       "  {'pid': '-6997269615945621261',\n",
       "   'passage': {'input_ids': tensor([[33964,   275]]), 'attention_mask': tensor([[1, 1]])}}],\n",
       " 'neg': [{'pid': '-1059230821898650523',\n",
       "   'passage': {'input_ids': tensor([[39014,   275,    13, 50188,    62,  7295,    66,     6]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1]])}},\n",
       "  {'pid': '5979920950827081735',\n",
       "   'passage': {'input_ids': tensor([[39014,   275,    13, 33295,    62,  3506,    62]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1]])}},\n",
       "  {'pid': '5006218284396414035',\n",
       "   'passage': {'input_ids': tensor([[39014,   275]]), 'attention_mask': tensor([[1, 1]])}},\n",
       "  {'pid': '8519025108267685379',\n",
       "   'passage': {'input_ids': tensor([[39014,   957]]), 'attention_mask': tensor([[1, 1]])}},\n",
       "  {'pid': '-2029295515473291484',\n",
       "   'passage': {'input_ids': tensor([[ 9310, 11011,   685,   961,    60,   379,   289]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1]])}},\n",
       "  {'pid': '-7086168812316445213',\n",
       "   'passage': {'input_ids': tensor([[39014,    62, 39098]]), 'attention_mask': tensor([[1, 1, 1]])}},\n",
       "  {'pid': '2194103192624150216',\n",
       "   'passage': {'input_ids': tensor([[   82, 11011,     0,  1635,   379,  1635]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1]])}},\n",
       "  {'pid': '675688915442226575',\n",
       "   'passage': {'input_ids': tensor([[ 9310, 11011,   691,   685,   961,    11,   288,    82, 11011,    60,\n",
       "              379,   289]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])}},\n",
       "  {'pid': '-532060175982832960',\n",
       "   'passage': {'input_ids': tensor([[ 2617,  1312, 19039,   275,    13, 50188,    62,  7295,  1312,   351,\n",
       "              289,     6]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])}},\n",
       "  {'pid': '-425050867452080408',\n",
       "   'passage': {'input_ids': tensor([[39014,   275,    13, 33295,    62,  9464,    62,   259,    73]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1]])}},\n",
       "  {'pid': '-566168252729545530',\n",
       "   'passage': {'input_ids': tensor([[31653,   685,  7857,    11,   289,    60]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1]])}},\n",
       "  {'pid': '-7671123778674527926',\n",
       "   'passage': {'input_ids': tensor([[158, 232,  95, 275,  13, 961]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1]])}},\n",
       "  {'pid': '1079449569097250766',\n",
       "   'passage': {'input_ids': tensor([[39014,  1066]]), 'attention_mask': tensor([[1, 1]])}},\n",
       "  {'pid': '-6162255274703520873',\n",
       "   'passage': {'input_ids': tensor([[   82, 11011,   685, 33295,    62,  3506,    60]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1]])}},\n",
       "  {'pid': '1504670812205882816',\n",
       "   'passage': {'input_ids': tensor([[   67, 12514,  1100]]), 'attention_mask': tensor([[1, 1, 1]])}},\n",
       "  {'pid': '-7837678799220824417',\n",
       "   'passage': {'input_ids': tensor([[   82, 11011,   685,  7857,    11,   289,    60]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1]])}},\n",
       "  {'pid': '-5594607461795628077',\n",
       "   'passage': {'input_ids': tensor([[7266,  301, 1312]]), 'attention_mask': tensor([[1, 1, 1]])}},\n",
       "  {'pid': '846899338491892079',\n",
       "   'passage': {'input_ids': tensor([[31653,   685, 18747,    13,   961,    62, 13564,  3256,  7177,    13,\n",
       "            13564,     6,    62, 13564,    11, 11876,    13, 13564,    62, 11883,\n",
       "               11,  7177,    13, 13564,    62, 11883,    11,   289,    60]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "            1, 1, 1, 1, 1]])}},\n",
       "  {'pid': '34868178422759404',\n",
       "   'passage': {'input_ids': tensor([[31653,   685, 29705,   238, 33295,    62,  3506,    62,   259,    73,\n",
       "               60]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])}},\n",
       "  {'pid': '-4291723548435929039',\n",
       "   'passage': {'input_ids': tensor([[   82, 11011]]), 'attention_mask': tensor([[1, 1]])}},\n",
       "  {'pid': '5905105816470661935',\n",
       "   'passage': {'input_ids': tensor([[ 9310, 11011,   685,   961,    60]]), 'attention_mask': tensor([[1, 1, 1, 1, 1]])}},\n",
       "  {'pid': '-240506235358478009',\n",
       "   'passage': {'input_ids': tensor([[39014,   275,    13, 33295,    62,  3506,     6]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1]])}},\n",
       "  {'pid': '-7014387072194972932',\n",
       "   'passage': {'input_ids': tensor([[ 2617,    62,  1659,    62, 27363,    62, 27363,    62,  2385,   535,\n",
       "               62,  1659,    62, 27363]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])}},\n",
       "  {'pid': '-9163172994500598878',\n",
       "   'passage': {'input_ids': tensor([[1102, 1851,  275,   13,  961,    6]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1]])}},\n",
       "  {'pid': '8180726172081474272',\n",
       "   'passage': {'input_ids': tensor([[3642, 6335, 2867]]), 'attention_mask': tensor([[1, 1, 1]])}},\n",
       "  {'pid': '-5168536336062798448',\n",
       "   'passage': {'input_ids': tensor([[31653,   685, 29705,   238, 33295,    62,  3506,    62,   259,  4399,\n",
       "               62,   733,    11, 17804,   238]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])}},\n",
       "  {'pid': '3276683918839931364',\n",
       "   'passage': {'input_ids': tensor([[39014,   275,    13,   961,     6,    62, 27363,    62,  1659,    62,\n",
       "            11883]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])}},\n",
       "  {'pid': '-3015987546996079227',\n",
       "   'passage': {'input_ids': tensor([[   62, 15699,  1058, 18872,   222, 44104,    64,  4808,    64,    62,\n",
       "               16,  1058, 33295,   357,  4868, 26367,     8,  2343,   253,   101,\n",
       "               72,    11]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])}},\n",
       "  {'pid': '4991927682312661887',\n",
       "   'passage': {'input_ids': tensor([[39014,   275,    13, 33295,    62,  3506]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1]])}},\n",
       "  {'pid': '3280197074985918113',\n",
       "   'passage': {'input_ids': tensor([[39014,   275,    13, 33295,    62,  9464,    62,   259,   752,   425]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])}},\n",
       "  {'pid': '-553055782530395591',\n",
       "   'passage': {'input_ids': tensor([[   62,  8625,  1058, 30671, 26367, 19039]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1]])}},\n",
       "  {'pid': '8547743317388738779',\n",
       "   'passage': {'input_ids': tensor([[39014,   275,    13, 50188,    62,  7295,    62, 27363]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1]])}},\n",
       "  {'pid': '-2125485750154732633',\n",
       "   'passage': {'input_ids': tensor([[1102, 1851, 1100,   62,  259,   73,  275]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1]])}},\n",
       "  {'pid': '-7360264406821898777',\n",
       "   'passage': {'input_ids': tensor([[  158,   232,    95,   275,    13, 50188,    62]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1]])}}]}"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "json_c_train[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de050b25",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1490d724",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e85189a4",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6a76609",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 237,
   "id": "7f88b760",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.model_selection import train_test_split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 238,
   "id": "9198827f",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train,x_dev = train_test_split(json_c_train,test_size= 0.15)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 239,
   "id": "6f0d0adf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "6058"
      ]
     },
     "execution_count": 239,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(X_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 240,
   "id": "e3e6353b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'qry': {'qid': '918643313112262782',\n",
       "  'query': [27,\n",
       "   11230,\n",
       "   1847,\n",
       "   29,\n",
       "   299,\n",
       "   1058,\n",
       "   2343,\n",
       "   226,\n",
       "   243,\n",
       "   11,\n",
       "   197,\n",
       "   37,\n",
       "   1058,\n",
       "   2099,\n",
       "   35138,\n",
       "   357,\n",
       "   77,\n",
       "   1343,\n",
       "   352,\n",
       "   8,\n",
       "   15168,\n",
       "   5994,\n",
       "   334,\n",
       "   11,\n",
       "   197,\n",
       "   62,\n",
       "   8625,\n",
       "   62,\n",
       "   16,\n",
       "   1058,\n",
       "   285,\n",
       "   85,\n",
       "   12543,\n",
       "   2715,\n",
       "   376,\n",
       "   11,\n",
       "   197,\n",
       "   80,\n",
       "   1058,\n",
       "   285,\n",
       "   85,\n",
       "   80,\n",
       "   79,\n",
       "   69,\n",
       "   376,\n",
       "   11,\n",
       "   197,\n",
       "   17394,\n",
       "   1058,\n",
       "   2099,\n",
       "   35138,\n",
       "   299,\n",
       "   11,\n",
       "   197,\n",
       "   81,\n",
       "   1058,\n",
       "   285,\n",
       "   85,\n",
       "   80,\n",
       "   79,\n",
       "   69,\n",
       "   13,\n",
       "   1073,\n",
       "   13049,\n",
       "   376,\n",
       "   26367,\n",
       "   15168,\n",
       "   285,\n",
       "   85,\n",
       "   80,\n",
       "   79,\n",
       "   69,\n",
       "   13,\n",
       "   1073,\n",
       "   13049,\n",
       "   376,\n",
       "   26367,\n",
       "   15168,\n",
       "   8772,\n",
       "   11,\n",
       "   197,\n",
       "   71,\n",
       "   1058,\n",
       "   18872,\n",
       "   222,\n",
       "   357,\n",
       "   87,\n",
       "   331,\n",
       "   1058,\n",
       "   285,\n",
       "   85,\n",
       "   80,\n",
       "   79,\n",
       "   69,\n",
       "   13,\n",
       "   1073,\n",
       "   13049,\n",
       "   376,\n",
       "   26367,\n",
       "   828,\n",
       "   374,\n",
       "   2124,\n",
       "   331,\n",
       "   15168,\n",
       "   285,\n",
       "   85,\n",
       "   12543,\n",
       "   2715,\n",
       "   13,\n",
       "   26282,\n",
       "   81,\n",
       "   357,\n",
       "   17394,\n",
       "   13,\n",
       "   2411,\n",
       "   62,\n",
       "   12957,\n",
       "   374,\n",
       "   8,\n",
       "   2124,\n",
       "   13,\n",
       "   16520,\n",
       "   331,\n",
       "   13,\n",
       "   16520,\n",
       "   11,\n",
       "   197,\n",
       "   87,\n",
       "   331,\n",
       "   1058,\n",
       "   285,\n",
       "   85,\n",
       "   80,\n",
       "   79,\n",
       "   69,\n",
       "   13,\n",
       "   1073,\n",
       "   13049,\n",
       "   376,\n",
       "   26367,\n",
       "   11,\n",
       "   197,\n",
       "   81,\n",
       "   5431,\n",
       "   1058,\n",
       "   374,\n",
       "   2124,\n",
       "   331,\n",
       "   11,\n",
       "   197,\n",
       "   64,\n",
       "   1058,\n",
       "   357,\n",
       "   76,\n",
       "   85,\n",
       "   80,\n",
       "   79,\n",
       "   69,\n",
       "   13,\n",
       "   47,\n",
       "   376,\n",
       "   737,\n",
       "   32,\n",
       "   11,\n",
       "   197,\n",
       "   69,\n",
       "   158,\n",
       "   224,\n",
       "   222,\n",
       "   277,\n",
       "   158,\n",
       "   224,\n",
       "   223,\n",
       "   1058,\n",
       "   357,\n",
       "   76,\n",
       "   85,\n",
       "   80,\n",
       "   79,\n",
       "   69,\n",
       "   13,\n",
       "   47,\n",
       "   376,\n",
       "   737,\n",
       "   33,\n",
       "   257,\n",
       "   2343,\n",
       "   253,\n",
       "   117,\n",
       "   7377,\n",
       "   119,\n",
       "   357,\n",
       "   72,\n",
       "   1058,\n",
       "   957,\n",
       "   17,\n",
       "   357,\n",
       "   77,\n",
       "   1343,\n",
       "   352,\n",
       "   36911,\n",
       "   357,\n",
       "   17394,\n",
       "   1058,\n",
       "   3712,\n",
       "   285,\n",
       "   85,\n",
       "   80,\n",
       "   79,\n",
       "   69,\n",
       "   13,\n",
       "   1073,\n",
       "   13049,\n",
       "   376,\n",
       "   26367,\n",
       "   8,\n",
       "   1312,\n",
       "   11,\n",
       "   197,\n",
       "   34350,\n",
       "   27363,\n",
       "   1058,\n",
       "   2124,\n",
       "   13,\n",
       "   16520,\n",
       "   796,\n",
       "   285,\n",
       "   85,\n",
       "   80,\n",
       "   79,\n",
       "   69,\n",
       "   13,\n",
       "   8937,\n",
       "   2343,\n",
       "   253,\n",
       "   101,\n",
       "   64,\n",
       "   11,\n",
       "   277,\n",
       "   158,\n",
       "   224,\n",
       "   222,\n",
       "   158,\n",
       "   253,\n",
       "   102,\n",
       "   11,\n",
       "   197,\n",
       "   67,\n",
       "   5948,\n",
       "   80,\n",
       "   1058,\n",
       "   331,\n",
       "   13,\n",
       "   16520,\n",
       "   796,\n",
       "   285,\n",
       "   85,\n",
       "   80,\n",
       "   79,\n",
       "   69,\n",
       "   13,\n",
       "   8937,\n",
       "   2343,\n",
       "   253,\n",
       "   101,\n",
       "   64,\n",
       "   11,\n",
       "   277,\n",
       "   158,\n",
       "   224,\n",
       "   223,\n",
       "   158,\n",
       "   253,\n",
       "   102,\n",
       "   11,\n",
       "   197,\n",
       "   71,\n",
       "   6,\n",
       "   1058,\n",
       "   18872,\n",
       "   222,\n",
       "   357,\n",
       "   72,\n",
       "   1058,\n",
       "   957,\n",
       "   17,\n",
       "   357,\n",
       "   77,\n",
       "   1343,\n",
       "   352,\n",
       "   4008,\n",
       "   357,\n",
       "   73,\n",
       "   1058,\n",
       "   357,\n",
       "   76,\n",
       "   85,\n",
       "   80,\n",
       "   79,\n",
       "   69,\n",
       "   13,\n",
       "   47,\n",
       "   376,\n",
       "   737,\n",
       "   33,\n",
       "   257,\n",
       "   1312,\n",
       "   828,\n",
       "   26367,\n",
       "   13,\n",
       "   2411,\n",
       "   62,\n",
       "   12957,\n",
       "   374,\n",
       "   357,\n",
       "   69,\n",
       "   158,\n",
       "   224,\n",
       "   222,\n",
       "   1312,\n",
       "   474,\n",
       "   8,\n",
       "   357,\n",
       "   69,\n",
       "   158,\n",
       "   224,\n",
       "   223,\n",
       "   1312,\n",
       "   474,\n",
       "   8,\n",
       "   197,\n",
       "   158,\n",
       "   232,\n",
       "   95,\n",
       "   357,\n",
       "   4906,\n",
       "   35138,\n",
       "   13,\n",
       "   312,\n",
       "   1058,\n",
       "   3712,\n",
       "   23611,\n",
       "   13,\n",
       "   28015,\n",
       "   357,\n",
       "   39377,\n",
       "   357,\n",
       "   87,\n",
       "   331,\n",
       "   1058,\n",
       "   285,\n",
       "   85,\n",
       "   80,\n",
       "   79,\n",
       "   69,\n",
       "   13,\n",
       "   1073,\n",
       "   13049,\n",
       "   376,\n",
       "   26367,\n",
       "   828,\n",
       "   374,\n",
       "   2124,\n",
       "   331,\n",
       "   4008,\n",
       "   1279,\n",
       "   13702,\n",
       "   29,\n",
       "   2124,\n",
       "   13,\n",
       "   16520,\n",
       "   796,\n",
       "   357,\n",
       "   4906,\n",
       "   35138,\n",
       "   13,\n",
       "   312,\n",
       "   1058,\n",
       "   3712,\n",
       "   23611,\n",
       "   13,\n",
       "   28015,\n",
       "   357,\n",
       "   39377,\n",
       "   357,\n",
       "   87,\n",
       "   331,\n",
       "   1058,\n",
       "   285,\n",
       "   85,\n",
       "   80,\n",
       "   79,\n",
       "   69,\n",
       "   13,\n",
       "   1073,\n",
       "   13049,\n",
       "   376,\n",
       "   26367,\n",
       "   828,\n",
       "   374,\n",
       "   2124,\n",
       "   331,\n",
       "   4008,\n",
       "   1279,\n",
       "   13702,\n",
       "   29,\n",
       "   331,\n",
       "   13,\n",
       "   16520,\n",
       "   198,\n",
       "   1279,\n",
       "   4805,\n",
       "   6684,\n",
       "   37,\n",
       "   42135,\n",
       "   29,\n",
       "   220]},\n",
       " 'pos': [{'pid': '6579643507496699425',\n",
       "   'passage': [31653, 685, 34350, 27363, 11, 288]},\n",
       "  {'pid': '-2978845321307905187',\n",
       "   'passage': [31653,\n",
       "    685,\n",
       "    34350,\n",
       "    27363,\n",
       "    11,\n",
       "    30121,\n",
       "    80,\n",
       "    11,\n",
       "    17804,\n",
       "    238,\n",
       "    8937,\n",
       "    62,\n",
       "    8899,\n",
       "    11,\n",
       "    17804,\n",
       "    238,\n",
       "    8937,\n",
       "    62,\n",
       "    8899,\n",
       "    11,\n",
       "    285,\n",
       "    36133,\n",
       "    12543,\n",
       "    2715,\n",
       "    13,\n",
       "    8899,\n",
       "    62,\n",
       "    27363,\n",
       "    11,\n",
       "    285,\n",
       "    36133,\n",
       "    12543,\n",
       "    2715,\n",
       "    13,\n",
       "    8899,\n",
       "    62,\n",
       "    27363,\n",
       "    60]}],\n",
       " 'neg': [{'pid': '6575927600422562800',\n",
       "   'passage': [31653,\n",
       "    685,\n",
       "    29705,\n",
       "    238,\n",
       "    23611,\n",
       "    13,\n",
       "    23661,\n",
       "    357,\n",
       "    421,\n",
       "    313,\n",
       "    13,\n",
       "    28015,\n",
       "    62,\n",
       "    27363,\n",
       "    62,\n",
       "    28015,\n",
       "    62,\n",
       "    1659,\n",
       "    62,\n",
       "    27363,\n",
       "    4808,\n",
       "    828,\n",
       "    17804,\n",
       "    238,\n",
       "    23611,\n",
       "    13,\n",
       "    12543,\n",
       "    62,\n",
       "    15699,\n",
       "    62,\n",
       "    16]},\n",
       "  {'pid': '928431572601992639',\n",
       "   'passage': [31653,\n",
       "    685,\n",
       "    29705,\n",
       "    238,\n",
       "    23611,\n",
       "    13,\n",
       "    23661,\n",
       "    357,\n",
       "    421,\n",
       "    313,\n",
       "    13,\n",
       "    28015,\n",
       "    62,\n",
       "    27363,\n",
       "    62,\n",
       "    28015,\n",
       "    62,\n",
       "    1659,\n",
       "    62,\n",
       "    27363,\n",
       "    4808,\n",
       "    828,\n",
       "    17804,\n",
       "    238,\n",
       "    23611,\n",
       "    13,\n",
       "    26282,\n",
       "    62,\n",
       "    261,\n",
       "    6,\n",
       "    62]},\n",
       "  {'pid': '7552767570238416096',\n",
       "   'passage': [31653,\n",
       "    685,\n",
       "    29705,\n",
       "    238,\n",
       "    23611,\n",
       "    13,\n",
       "    23661,\n",
       "    357,\n",
       "    421,\n",
       "    313,\n",
       "    13,\n",
       "    28015,\n",
       "    62,\n",
       "    27363,\n",
       "    62,\n",
       "    28015,\n",
       "    4808,\n",
       "    4808,\n",
       "    828,\n",
       "    17804,\n",
       "    238,\n",
       "    23611,\n",
       "    13,\n",
       "    28015,\n",
       "    62,\n",
       "    27363,\n",
       "    62,\n",
       "    28015,\n",
       "    11,\n",
       "    17804,\n",
       "    238]},\n",
       "  {'pid': '-2234788984240462035', 'passage': [1069, 529]},\n",
       "  {'pid': '-4411941448410955259',\n",
       "   'passage': [31653,\n",
       "    685,\n",
       "    29705,\n",
       "    238,\n",
       "    23611,\n",
       "    13,\n",
       "    23661,\n",
       "    357,\n",
       "    421,\n",
       "    313,\n",
       "    13,\n",
       "    28015,\n",
       "    62,\n",
       "    27363,\n",
       "    62,\n",
       "    28015,\n",
       "    62,\n",
       "    1659,\n",
       "    62,\n",
       "    27363,\n",
       "    4808,\n",
       "    828,\n",
       "    17804,\n",
       "    238,\n",
       "    23611,\n",
       "    13,\n",
       "    23661,\n",
       "    357,\n",
       "    421,\n",
       "    313,\n",
       "    13]},\n",
       "  {'pid': '3122899436888549272', 'passage': [742]},\n",
       "  {'pid': '-8839975914427080847', 'passage': [90]},\n",
       "  {'pid': '-2432619728502266055', 'passage': [67]},\n",
       "  {'pid': '-7361431018201337745', 'passage': [7266, 301]},\n",
       "  {'pid': '4194638930665964425', 'passage': [81]}]}"
      ]
     },
     "execution_count": 240,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X_train[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 273,
   "id": "210138fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# with open('train.json', 'w') as fout:\n",
    "#     json.dump(X_train , fout)\n",
    "    \n",
    "with open('data/train/train.json', 'w') as outfile:\n",
    "    for entry in X_train:\n",
    "        json.dump(entry, outfile)\n",
    "        outfile.write('\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 274,
   "id": "22622cae",
   "metadata": {},
   "outputs": [],
   "source": [
    "# with open('dev.json', 'w') as fout:\n",
    "#     json.dump(x_dev , fout)\n",
    "    \n",
    "with open('data/dev/dev.json', 'w') as outfile:\n",
    "    for entry in x_dev:\n",
    "        json.dump(entry, outfile)\n",
    "        outfile.write('\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "70f8d8b1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "432K\tReranker\n",
      "6.2M\tall_gpt-ada.json\n",
      "11M\tall_gpt_combined.json\n",
      "538M\tclass\n",
      "26M\tdata\n",
      "562M\tfiles_upload\n",
      "7.9M\tgdrive\n",
      "20K\tgptNEO_eval.ipynb\n",
      "1012K\tjson_data_0.json\n",
      "7.9M\tjson_data_1.json\n",
      "4.3M\tjson_data__hard_ones.json\n",
      "740K\tneo-Copy1.ipynb\n",
      "908K\tneo.ipynb\n",
      "0\tonstart.log\n",
      "4.0K\tonstart.sh\n",
      "116K\treranker.ipynb\n",
      "11G\tresults\n",
      "0\ttest.json\n",
      "0\ttmp_trainer\n"
     ]
    }
   ],
   "source": [
    "!du -sh *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8e19348",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "id": "a3d0e5ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_best_ranking(each):\n",
    "    test = each[\"prompt\"].replace(\"GOAL\",\"<GOAL>\").replace(\"PROOFSTEP\",\"<PROOFSTEP>\")\n",
    "    \n",
    "    results = [i for i in each[\"all\"]]\n",
    "    l = []\n",
    "    for i in results+[each[\"real\"]]:\n",
    "        inputs = rk.tokenize(test, i, return_tensors='pt')\n",
    "        inputs.to(torch.device(\"cuda:0\"))\n",
    "        score = rk(inputs).logits\n",
    "        l.append((score.cpu().detach().numpy()[0][0],i))\n",
    "    l.sort(key=lambda x:x[0])\n",
    "    l.reverse()\n",
    "    l_e = [i[1] for i in l]\n",
    "    return l_e.index(each[\"real\"]),l_e.index(each[\"best_match\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "id": "8eb03e1e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100% 1000/1000 [03:01<00:00,  5.52it/s]\n"
     ]
    }
   ],
   "source": [
    "resi_search = []\n",
    "    \n",
    "for i in tqdm(json_c[:1000]):\n",
    "    resi_search.append(get_best_ranking(i))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "id": "aeb1d717",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(15, 23),\n",
       " (17, 25),\n",
       " (16, 15),\n",
       " (15, 5),\n",
       " (13, 20),\n",
       " (9, 4),\n",
       " (20, 20),\n",
       " (2, 12),\n",
       " (11, 16),\n",
       " (0, 6),\n",
       " (15, 15),\n",
       " (2, 12),\n",
       " (8, 8),\n",
       " (2, 8),\n",
       " (2, 3),\n",
       " (8, 8),\n",
       " (14, 14),\n",
       " (0, 6),\n",
       " (1, 1),\n",
       " (11, 11),\n",
       " (13, 13),\n",
       " (4, 10),\n",
       " (3, 3),\n",
       " (16, 16),\n",
       " (5, 5),\n",
       " (10, 1),\n",
       " (5, 5),\n",
       " (5, 5),\n",
       " (6, 19),\n",
       " (6, 6),\n",
       " (5, 5),\n",
       " (8, 21),\n",
       " (3, 10),\n",
       " (4, 11),\n",
       " (10, 16),\n",
       " (6, 6),\n",
       " (8, 12),\n",
       " (14, 17),\n",
       " (21, 9),\n",
       " (2, 15),\n",
       " (21, 12),\n",
       " (5, 6),\n",
       " (7, 6),\n",
       " (17, 14),\n",
       " (2, 15),\n",
       " (21, 12),\n",
       " (12, 4),\n",
       " (3, 3),\n",
       " (22, 2),\n",
       " (7, 0),\n",
       " (15, 15),\n",
       " (15, 10),\n",
       " (14, 10),\n",
       " (11, 13),\n",
       " (1, 11),\n",
       " (20, 16),\n",
       " (1, 12),\n",
       " (11, 12),\n",
       " (15, 16),\n",
       " (7, 7),\n",
       " (12, 12),\n",
       " (5, 5),\n",
       " (14, 15),\n",
       " (8, 8),\n",
       " (6, 10),\n",
       " (15, 4),\n",
       " (9, 14),\n",
       " (15, 2),\n",
       " (12, 17),\n",
       " (14, 14),\n",
       " (15, 15),\n",
       " (11, 6),\n",
       " (5, 6),\n",
       " (14, 14),\n",
       " (12, 13),\n",
       " (9, 18),\n",
       " (8, 8),\n",
       " (0, 20),\n",
       " (13, 21),\n",
       " (1, 14),\n",
       " (10, 11),\n",
       " (9, 13),\n",
       " (31, 29),\n",
       " (16, 11),\n",
       " (25, 24),\n",
       " (4, 4),\n",
       " (3, 20),\n",
       " (5, 20),\n",
       " (4, 17),\n",
       " (11, 11),\n",
       " (2, 10),\n",
       " (4, 16),\n",
       " (3, 7),\n",
       " (15, 6),\n",
       " (11, 11),\n",
       " (8, 7),\n",
       " (10, 9),\n",
       " (2, 2),\n",
       " (0, 12),\n",
       " (10, 8),\n",
       " (17, 16),\n",
       " (14, 11),\n",
       " (0, 7),\n",
       " (5, 4),\n",
       " (6, 4),\n",
       " (13, 20),\n",
       " (3, 20),\n",
       " (12, 6),\n",
       " (17, 15),\n",
       " (10, 10),\n",
       " (19, 10),\n",
       " (0, 0),\n",
       " (0, 0),\n",
       " (11, 9),\n",
       " (1, 1),\n",
       " (14, 16),\n",
       " (16, 16),\n",
       " (0, 0),\n",
       " (2, 2),\n",
       " (0, 0),\n",
       " (11, 13),\n",
       " (0, 0),\n",
       " (2, 10),\n",
       " (14, 18),\n",
       " (9, 18),\n",
       " (5, 14),\n",
       " (4, 2),\n",
       " (0, 5),\n",
       " (8, 24),\n",
       " (14, 12),\n",
       " (4, 13),\n",
       " (4, 4),\n",
       " (6, 8),\n",
       " (4, 13),\n",
       " (0, 12),\n",
       " (16, 6),\n",
       " (3, 3),\n",
       " (4, 10),\n",
       " (14, 14),\n",
       " (3, 13),\n",
       " (3, 1),\n",
       " (11, 6),\n",
       " (2, 2),\n",
       " (11, 0),\n",
       " (2, 11),\n",
       " (2, 25),\n",
       " (1, 2),\n",
       " (12, 7),\n",
       " (14, 20),\n",
       " (22, 8),\n",
       " (1, 3),\n",
       " (15, 1),\n",
       " (8, 15),\n",
       " (1, 22),\n",
       " (2, 7),\n",
       " (5, 14),\n",
       " (4, 3),\n",
       " (12, 15),\n",
       " (19, 18),\n",
       " (12, 19),\n",
       " (5, 1),\n",
       " (8, 13),\n",
       " (5, 9),\n",
       " (15, 22),\n",
       " (7, 16),\n",
       " (10, 13),\n",
       " (11, 15),\n",
       " (0, 5),\n",
       " (11, 14),\n",
       " (18, 15),\n",
       " (22, 12),\n",
       " (13, 10),\n",
       " (5, 6),\n",
       " (6, 4),\n",
       " (0, 11),\n",
       " (11, 5),\n",
       " (21, 9),\n",
       " (14, 19),\n",
       " (14, 15),\n",
       " (12, 7),\n",
       " (8, 8),\n",
       " (15, 15),\n",
       " (19, 6),\n",
       " (12, 9),\n",
       " (2, 2),\n",
       " (0, 0),\n",
       " (0, 0),\n",
       " (22, 22),\n",
       " (13, 13),\n",
       " (5, 5),\n",
       " (7, 7),\n",
       " (15, 15),\n",
       " (8, 8),\n",
       " (14, 14),\n",
       " (4, 3),\n",
       " (4, 4),\n",
       " (0, 8),\n",
       " (11, 11),\n",
       " (10, 12),\n",
       " (11, 5),\n",
       " (4, 4),\n",
       " (0, 0),\n",
       " (1, 12),\n",
       " (9, 11),\n",
       " (11, 11),\n",
       " (22, 8),\n",
       " (3, 5),\n",
       " (4, 11),\n",
       " (2, 13),\n",
       " (5, 16),\n",
       " (13, 5),\n",
       " (13, 13),\n",
       " (4, 3),\n",
       " (24, 1),\n",
       " (20, 13),\n",
       " (20, 17),\n",
       " (0, 11),\n",
       " (0, 8),\n",
       " (5, 4),\n",
       " (3, 1),\n",
       " (10, 5),\n",
       " (7, 8),\n",
       " (10, 9),\n",
       " (8, 12),\n",
       " (6, 5),\n",
       " (1, 2),\n",
       " (6, 10),\n",
       " (9, 1),\n",
       " (10, 0),\n",
       " (1, 6),\n",
       " (8, 4),\n",
       " (8, 0),\n",
       " (5, 0),\n",
       " (6, 9),\n",
       " (12, 9),\n",
       " (8, 8),\n",
       " (7, 7),\n",
       " (7, 5),\n",
       " (8, 7),\n",
       " (9, 2),\n",
       " (12, 13),\n",
       " (9, 8),\n",
       " (6, 11),\n",
       " (8, 10),\n",
       " (16, 13),\n",
       " (3, 15),\n",
       " (15, 20),\n",
       " (12, 11),\n",
       " (4, 1),\n",
       " (11, 20),\n",
       " (11, 17),\n",
       " (13, 11),\n",
       " (13, 15),\n",
       " (11, 7),\n",
       " (12, 11),\n",
       " (19, 16),\n",
       " (16, 7),\n",
       " (17, 11),\n",
       " (3, 3),\n",
       " (9, 15),\n",
       " (21, 11),\n",
       " (7, 0),\n",
       " (14, 14),\n",
       " (4, 2),\n",
       " (7, 1),\n",
       " (8, 6),\n",
       " (15, 19),\n",
       " (18, 15),\n",
       " (6, 17),\n",
       " (14, 20),\n",
       " (2, 8),\n",
       " (12, 19),\n",
       " (8, 8),\n",
       " (19, 12),\n",
       " (9, 10),\n",
       " (8, 16),\n",
       " (4, 10),\n",
       " (11, 9),\n",
       " (6, 9),\n",
       " (13, 5),\n",
       " (16, 1),\n",
       " (19, 19),\n",
       " (20, 21),\n",
       " (5, 14),\n",
       " (3, 11),\n",
       " (3, 18),\n",
       " (3, 19),\n",
       " (0, 18),\n",
       " (20, 20),\n",
       " (1, 3),\n",
       " (17, 9),\n",
       " (22, 24),\n",
       " (14, 13),\n",
       " (18, 3),\n",
       " (11, 15),\n",
       " (18, 12),\n",
       " (9, 5),\n",
       " (7, 16),\n",
       " (13, 6),\n",
       " (5, 0),\n",
       " (9, 16),\n",
       " (11, 0),\n",
       " (4, 4),\n",
       " (12, 12),\n",
       " (7, 7),\n",
       " (0, 2),\n",
       " (15, 3),\n",
       " (7, 15),\n",
       " (8, 11),\n",
       " (17, 16),\n",
       " (9, 25),\n",
       " (14, 3),\n",
       " (5, 17),\n",
       " (4, 8),\n",
       " (4, 5),\n",
       " (6, 10),\n",
       " (12, 13),\n",
       " (5, 4),\n",
       " (8, 6),\n",
       " (5, 5),\n",
       " (5, 8),\n",
       " (15, 14),\n",
       " (5, 9),\n",
       " (2, 1),\n",
       " (6, 2),\n",
       " (2, 6),\n",
       " (6, 8),\n",
       " (5, 0),\n",
       " (13, 14),\n",
       " (0, 0),\n",
       " (13, 1),\n",
       " (9, 11),\n",
       " (6, 6),\n",
       " (3, 3),\n",
       " (6, 6),\n",
       " (8, 8),\n",
       " (5, 5),\n",
       " (8, 8),\n",
       " (4, 4),\n",
       " (5, 4),\n",
       " (6, 2),\n",
       " (3, 3),\n",
       " (13, 20),\n",
       " (1, 6),\n",
       " (21, 21),\n",
       " (9, 11),\n",
       " (23, 20),\n",
       " (6, 6),\n",
       " (6, 19),\n",
       " (5, 0),\n",
       " (2, 15),\n",
       " (10, 12),\n",
       " (3, 11),\n",
       " (12, 9),\n",
       " (9, 5),\n",
       " (14, 10),\n",
       " (6, 18),\n",
       " (4, 1),\n",
       " (6, 2),\n",
       " (17, 15),\n",
       " (8, 15),\n",
       " (3, 6),\n",
       " (3, 4),\n",
       " (8, 18),\n",
       " (8, 7),\n",
       " (13, 7),\n",
       " (21, 24),\n",
       " (15, 13),\n",
       " (8, 11),\n",
       " (15, 16),\n",
       " (3, 11),\n",
       " (20, 5),\n",
       " (10, 13),\n",
       " (12, 11),\n",
       " (8, 8),\n",
       " (6, 5),\n",
       " (12, 7),\n",
       " (13, 13),\n",
       " (14, 11),\n",
       " (19, 13),\n",
       " (15, 15),\n",
       " (19, 11),\n",
       " (0, 3),\n",
       " (4, 4),\n",
       " (7, 12),\n",
       " (8, 16),\n",
       " (11, 11),\n",
       " (9, 6),\n",
       " (12, 12),\n",
       " (12, 12),\n",
       " (13, 13),\n",
       " (7, 7),\n",
       " (6, 6),\n",
       " (10, 10),\n",
       " (2, 2),\n",
       " (6, 1),\n",
       " (14, 0),\n",
       " (1, 16),\n",
       " (13, 11),\n",
       " (12, 12),\n",
       " (2, 12),\n",
       " (20, 10),\n",
       " (8, 10),\n",
       " (0, 11),\n",
       " (0, 0),\n",
       " (3, 3),\n",
       " (3, 3),\n",
       " (11, 3),\n",
       " (8, 8),\n",
       " (11, 10),\n",
       " (16, 10),\n",
       " (7, 1),\n",
       " (19, 19),\n",
       " (12, 10),\n",
       " (11, 12),\n",
       " (10, 10),\n",
       " (6, 6),\n",
       " (7, 13),\n",
       " (5, 9),\n",
       " (13, 5),\n",
       " (8, 16),\n",
       " (7, 7),\n",
       " (13, 1),\n",
       " (2, 17),\n",
       " (15, 2),\n",
       " (14, 17),\n",
       " (5, 13),\n",
       " (2, 2),\n",
       " (3, 15),\n",
       " (10, 5),\n",
       " (11, 11),\n",
       " (2, 9),\n",
       " (11, 4),\n",
       " (16, 18),\n",
       " (9, 12),\n",
       " (4, 11),\n",
       " (15, 12),\n",
       " (0, 0),\n",
       " (4, 3),\n",
       " (5, 4),\n",
       " (12, 10),\n",
       " (19, 19),\n",
       " (18, 18),\n",
       " (0, 17),\n",
       " (0, 0),\n",
       " (0, 7),\n",
       " (5, 6),\n",
       " (3, 4),\n",
       " (13, 8),\n",
       " (15, 15),\n",
       " (17, 17),\n",
       " (0, 0),\n",
       " (0, 0),\n",
       " (3, 1),\n",
       " (3, 7),\n",
       " (9, 19),\n",
       " (12, 9),\n",
       " (18, 18),\n",
       " (9, 9),\n",
       " (18, 0),\n",
       " (3, 19),\n",
       " (12, 3),\n",
       " (11, 11),\n",
       " (3, 3),\n",
       " (11, 11),\n",
       " (2, 17),\n",
       " (11, 11),\n",
       " (9, 9),\n",
       " (10, 10),\n",
       " (16, 16),\n",
       " (0, 0),\n",
       " (13, 2),\n",
       " (6, 6),\n",
       " (2, 4),\n",
       " (11, 12),\n",
       " (20, 15),\n",
       " (15, 0),\n",
       " (0, 0),\n",
       " (20, 18),\n",
       " (15, 12),\n",
       " (8, 7),\n",
       " (13, 4),\n",
       " (15, 10),\n",
       " (3, 3),\n",
       " (7, 15),\n",
       " (17, 11),\n",
       " (17, 11),\n",
       " (22, 14),\n",
       " (22, 23),\n",
       " (4, 1),\n",
       " (8, 18),\n",
       " (2, 0),\n",
       " (3, 1),\n",
       " (10, 3),\n",
       " (5, 4),\n",
       " (0, 0),\n",
       " (2, 3),\n",
       " (0, 0),\n",
       " (1, 1),\n",
       " (9, 8),\n",
       " (7, 10),\n",
       " (20, 10),\n",
       " (13, 11),\n",
       " (12, 15),\n",
       " (8, 12),\n",
       " (19, 10),\n",
       " (15, 19),\n",
       " (27, 27),\n",
       " (26, 26),\n",
       " (19, 12),\n",
       " (9, 9),\n",
       " (10, 5),\n",
       " (18, 1),\n",
       " (8, 2),\n",
       " (19, 18),\n",
       " (4, 14),\n",
       " (6, 7),\n",
       " (15, 16),\n",
       " (8, 6),\n",
       " (14, 22),\n",
       " (6, 13),\n",
       " (19, 22),\n",
       " (5, 16),\n",
       " (10, 1),\n",
       " (5, 11),\n",
       " (5, 3),\n",
       " (6, 16),\n",
       " (22, 9),\n",
       " (0, 7),\n",
       " (15, 15),\n",
       " (21, 6),\n",
       " (15, 4),\n",
       " (17, 22),\n",
       " (12, 6),\n",
       " (4, 11),\n",
       " (7, 13),\n",
       " (5, 10),\n",
       " (22, 19),\n",
       " (23, 16),\n",
       " (6, 8),\n",
       " (7, 0),\n",
       " (6, 8),\n",
       " (4, 4),\n",
       " (17, 17),\n",
       " (15, 8),\n",
       " (0, 0),\n",
       " (17, 13),\n",
       " (0, 0),\n",
       " (10, 14),\n",
       " (7, 6),\n",
       " (2, 11),\n",
       " (9, 12),\n",
       " (21, 8),\n",
       " (15, 12),\n",
       " (16, 14),\n",
       " (16, 16),\n",
       " (0, 0),\n",
       " (17, 21),\n",
       " (2, 2),\n",
       " (7, 7),\n",
       " (17, 13),\n",
       " (10, 11),\n",
       " (13, 13),\n",
       " (10, 10),\n",
       " (10, 10),\n",
       " (22, 4),\n",
       " (19, 20),\n",
       " (16, 17),\n",
       " (6, 2),\n",
       " (5, 15),\n",
       " (20, 21),\n",
       " (16, 15),\n",
       " (16, 1),\n",
       " (16, 8),\n",
       " (15, 4),\n",
       " (4, 12),\n",
       " (5, 9),\n",
       " (7, 8),\n",
       " (15, 9),\n",
       " (16, 16),\n",
       " (3, 1),\n",
       " (3, 13),\n",
       " (4, 4),\n",
       " (15, 15),\n",
       " (18, 5),\n",
       " (11, 9),\n",
       " (8, 12),\n",
       " (16, 7),\n",
       " (1, 6),\n",
       " (13, 14),\n",
       " (13, 13),\n",
       " (3, 4),\n",
       " (13, 13),\n",
       " (11, 11),\n",
       " (15, 10),\n",
       " (8, 8),\n",
       " (4, 5),\n",
       " (7, 18),\n",
       " (9, 1),\n",
       " (18, 10),\n",
       " (0, 0),\n",
       " (14, 7),\n",
       " (11, 12),\n",
       " (13, 10),\n",
       " (13, 17),\n",
       " (13, 6),\n",
       " (5, 12),\n",
       " (3, 8),\n",
       " (21, 9),\n",
       " (6, 2),\n",
       " (14, 15),\n",
       " (14, 2),\n",
       " (9, 11),\n",
       " (14, 17),\n",
       " (11, 13),\n",
       " (3, 7),\n",
       " (2, 0),\n",
       " (4, 3),\n",
       " (3, 0),\n",
       " (5, 5),\n",
       " (7, 7),\n",
       " (5, 15),\n",
       " (7, 7),\n",
       " (4, 4),\n",
       " (5, 5),\n",
       " (6, 6),\n",
       " (0, 0),\n",
       " (2, 2),\n",
       " (0, 8),\n",
       " (3, 5),\n",
       " (4, 5),\n",
       " (4, 5),\n",
       " (3, 7),\n",
       " (3, 16),\n",
       " (5, 13),\n",
       " (3, 2),\n",
       " (2, 8),\n",
       " (5, 9),\n",
       " (4, 6),\n",
       " (5, 5),\n",
       " (3, 2),\n",
       " (0, 0),\n",
       " (0, 0),\n",
       " (13, 12),\n",
       " (17, 2),\n",
       " (19, 0),\n",
       " (18, 2),\n",
       " (1, 1),\n",
       " (1, 1),\n",
       " (0, 0),\n",
       " (7, 6),\n",
       " (5, 5),\n",
       " (3, 18),\n",
       " (8, 14),\n",
       " (16, 13),\n",
       " (12, 3),\n",
       " (12, 5),\n",
       " (11, 6),\n",
       " (1, 16),\n",
       " (14, 4),\n",
       " (10, 10),\n",
       " (7, 2),\n",
       " (20, 13),\n",
       " (2, 4),\n",
       " (1, 4),\n",
       " (5, 8),\n",
       " (11, 12),\n",
       " (7, 16),\n",
       " (14, 13),\n",
       " (10, 0),\n",
       " (8, 10),\n",
       " (9, 23),\n",
       " (18, 11),\n",
       " (4, 2),\n",
       " (8, 4),\n",
       " (7, 7),\n",
       " (7, 7),\n",
       " (5, 5),\n",
       " (0, 0),\n",
       " (5, 5),\n",
       " (7, 7),\n",
       " (7, 7),\n",
       " (6, 6),\n",
       " (4, 4),\n",
       " (9, 9),\n",
       " (5, 5),\n",
       " (2, 2),\n",
       " (7, 10),\n",
       " (5, 4),\n",
       " (4, 4),\n",
       " (2, 2),\n",
       " (3, 3),\n",
       " (5, 5),\n",
       " (2, 2),\n",
       " (1, 1),\n",
       " (3, 3),\n",
       " (6, 6),\n",
       " (2, 2),\n",
       " (9, 9),\n",
       " (9, 9),\n",
       " (5, 5),\n",
       " (6, 6),\n",
       " (6, 6),\n",
       " (2, 2),\n",
       " (3, 3),\n",
       " (4, 4),\n",
       " (2, 2),\n",
       " (3, 3),\n",
       " (4, 4),\n",
       " (5, 5),\n",
       " (7, 13),\n",
       " (17, 22),\n",
       " (0, 0),\n",
       " (10, 6),\n",
       " (18, 18),\n",
       " (12, 8),\n",
       " (8, 9),\n",
       " (21, 10),\n",
       " (24, 2),\n",
       " (3, 12),\n",
       " (3, 10),\n",
       " (5, 13),\n",
       " (1, 14),\n",
       " (12, 3),\n",
       " (13, 16),\n",
       " (6, 3),\n",
       " (12, 11),\n",
       " (5, 5),\n",
       " (19, 23),\n",
       " (0, 12),\n",
       " (13, 21),\n",
       " (2, 2),\n",
       " (9, 12),\n",
       " (6, 13),\n",
       " (10, 19),\n",
       " (3, 1),\n",
       " (7, 22),\n",
       " (4, 14),\n",
       " (7, 19),\n",
       " (12, 12),\n",
       " (2, 0),\n",
       " (9, 1),\n",
       " (7, 7),\n",
       " (14, 3),\n",
       " (4, 17),\n",
       " (6, 16),\n",
       " (3, 8),\n",
       " (16, 9),\n",
       " (4, 4),\n",
       " (6, 10),\n",
       " (7, 7),\n",
       " (1, 1),\n",
       " (2, 4),\n",
       " (5, 15),\n",
       " (17, 14),\n",
       " (11, 11),\n",
       " (8, 0),\n",
       " (5, 5),\n",
       " (2, 2),\n",
       " (4, 5),\n",
       " (5, 4),\n",
       " (12, 0),\n",
       " (6, 7),\n",
       " (4, 7),\n",
       " (19, 18),\n",
       " (1, 1),\n",
       " (9, 0),\n",
       " (4, 10),\n",
       " (1, 3),\n",
       " (7, 1),\n",
       " (6, 12),\n",
       " (5, 13),\n",
       " (11, 17),\n",
       " (2, 2),\n",
       " (8, 13),\n",
       " (10, 1),\n",
       " (8, 10),\n",
       " (26, 12),\n",
       " (11, 20),\n",
       " (16, 17),\n",
       " (8, 16),\n",
       " (12, 20),\n",
       " (6, 15),\n",
       " (1, 1),\n",
       " (6, 17),\n",
       " (15, 16),\n",
       " (21, 22),\n",
       " (2, 13),\n",
       " (3, 17),\n",
       " (10, 16),\n",
       " (18, 15),\n",
       " (21, 14),\n",
       " (8, 14),\n",
       " (15, 4),\n",
       " (4, 11),\n",
       " (17, 16),\n",
       " (5, 5),\n",
       " (5, 5),\n",
       " (3, 3),\n",
       " (3, 7),\n",
       " (6, 4),\n",
       " (7, 12),\n",
       " (15, 18),\n",
       " (2, 20),\n",
       " (1, 11),\n",
       " (14, 15),\n",
       " (3, 15),\n",
       " (17, 24),\n",
       " (2, 0),\n",
       " (1, 0),\n",
       " (9, 9),\n",
       " (16, 8),\n",
       " (7, 0),\n",
       " (13, 6),\n",
       " (5, 1),\n",
       " (14, 6),\n",
       " (9, 17),\n",
       " (10, 20),\n",
       " (5, 4),\n",
       " (19, 1),\n",
       " (3, 6),\n",
       " (4, 3),\n",
       " (19, 12),\n",
       " (5, 7),\n",
       " (8, 11),\n",
       " (8, 6),\n",
       " (2, 1),\n",
       " (3, 0),\n",
       " (1, 7),\n",
       " (1, 0),\n",
       " (3, 4),\n",
       " (7, 1),\n",
       " (5, 5),\n",
       " (10, 8),\n",
       " (4, 8),\n",
       " (2, 0),\n",
       " (7, 11),\n",
       " (11, 7),\n",
       " (12, 17),\n",
       " (7, 14),\n",
       " (9, 17),\n",
       " (4, 5),\n",
       " (5, 16),\n",
       " (9, 21),\n",
       " (8, 10),\n",
       " (4, 6),\n",
       " (8, 9),\n",
       " (4, 5),\n",
       " (19, 18),\n",
       " (5, 0),\n",
       " (8, 3),\n",
       " (5, 6),\n",
       " (7, 1),\n",
       " (3, 2),\n",
       " (6, 7),\n",
       " (5, 7),\n",
       " (12, 9),\n",
       " (11, 8),\n",
       " (0, 3),\n",
       " (10, 2),\n",
       " (18, 18),\n",
       " (3, 3),\n",
       " (0, 4),\n",
       " (20, 20),\n",
       " (22, 1),\n",
       " (1, 9),\n",
       " (1, 2),\n",
       " (1, 12),\n",
       " (7, 21),\n",
       " (6, 2),\n",
       " (1, 1),\n",
       " (3, 5),\n",
       " (4, 10),\n",
       " (2, 1),\n",
       " (15, 9),\n",
       " (11, 2),\n",
       " (15, 13),\n",
       " (6, 4),\n",
       " (11, 13),\n",
       " (5, 5),\n",
       " (7, 4),\n",
       " (9, 4),\n",
       " (10, 5),\n",
       " (6, 7),\n",
       " (6, 6),\n",
       " (5, 6),\n",
       " (5, 11),\n",
       " (0, 0),\n",
       " (7, 1),\n",
       " (8, 0),\n",
       " (2, 6),\n",
       " (7, 9),\n",
       " (3, 5),\n",
       " (3, 2),\n",
       " (4, 5),\n",
       " (3, 15),\n",
       " (0, 5),\n",
       " (14, 14),\n",
       " (1, 10),\n",
       " (25, 25),\n",
       " (23, 23),\n",
       " (20, 20),\n",
       " (23, 23),\n",
       " (15, 12),\n",
       " (13, 9),\n",
       " (16, 9),\n",
       " (12, 18),\n",
       " (9, 9),\n",
       " (10, 10),\n",
       " (18, 15),\n",
       " (0, 3),\n",
       " (2, 8),\n",
       " (0, 0),\n",
       " (21, 14),\n",
       " (1, 4),\n",
       " (17, 17),\n",
       " (11, 4),\n",
       " (6, 6),\n",
       " (14, 14),\n",
       " (17, 6),\n",
       " (0, 0),\n",
       " (4, 3),\n",
       " (12, 8),\n",
       " (11, 10),\n",
       " (1, 1),\n",
       " (3, 16),\n",
       " (10, 10),\n",
       " (14, 14),\n",
       " (13, 10),\n",
       " (11, 11),\n",
       " (1, 16),\n",
       " (10, 10),\n",
       " (4, 4),\n",
       " (14, 10),\n",
       " (3, 4),\n",
       " (5, 20),\n",
       " (12, 12),\n",
       " (8, 8),\n",
       " (17, 0),\n",
       " (12, 12),\n",
       " (8, 8),\n",
       " (24, 14),\n",
       " (14, 16),\n",
       " (11, 9),\n",
       " (11, 7),\n",
       " (0, 6),\n",
       " (11, 7),\n",
       " (15, 6),\n",
       " (3, 3),\n",
       " (5, 5),\n",
       " (25, 25),\n",
       " (24, 24),\n",
       " (12, 24),\n",
       " (9, 16),\n",
       " (10, 7),\n",
       " (1, 1),\n",
       " (6, 6),\n",
       " (6, 6),\n",
       " (7, 3),\n",
       " (11, 4),\n",
       " (15, 20),\n",
       " (14, 18),\n",
       " (18, 19),\n",
       " (18, 12),\n",
       " (26, 20),\n",
       " (14, 16),\n",
       " (13, 14),\n",
       " (11, 2),\n",
       " (13, 8),\n",
       " (1, 16),\n",
       " (14, 14),\n",
       " (6, 8),\n",
       " (15, 17),\n",
       " (19, 21),\n",
       " (0, 6),\n",
       " (6, 17),\n",
       " (3, 11),\n",
       " (21, 3),\n",
       " (3, 11),\n",
       " (10, 14),\n",
       " (3, 4),\n",
       " (8, 13),\n",
       " (10, 4),\n",
       " (5, 11),\n",
       " (8, 8),\n",
       " (21, 9),\n",
       " (9, 4),\n",
       " (7, 11),\n",
       " (18, 6),\n",
       " (10, 13),\n",
       " (4, 7),\n",
       " (0, 3),\n",
       " (20, 21),\n",
       " (19, 20),\n",
       " (18, 19),\n",
       " (2, 15),\n",
       " (14, 17),\n",
       " (11, 11),\n",
       " (21, 10),\n",
       " (21, 12)]"
      ]
     },
     "execution_count": 104,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "resi_search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 106,
   "id": "190124c6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>0</th>\n",
       "      <th>1</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>count</th>\n",
       "      <td>1000.000000</td>\n",
       "      <td>1000.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>mean</th>\n",
       "      <td>8.878000</td>\n",
       "      <td>9.300000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>6.124592</td>\n",
       "      <td>6.173177</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>min</th>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>25%</th>\n",
       "      <td>4.000000</td>\n",
       "      <td>4.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>50%</th>\n",
       "      <td>8.000000</td>\n",
       "      <td>9.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>75%</th>\n",
       "      <td>13.000000</td>\n",
       "      <td>14.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>max</th>\n",
       "      <td>31.000000</td>\n",
       "      <td>29.000000</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                 0            1\n",
       "count  1000.000000  1000.000000\n",
       "mean      8.878000     9.300000\n",
       "std       6.124592     6.173177\n",
       "min       0.000000     0.000000\n",
       "25%       4.000000     4.000000\n",
       "50%       8.000000     9.000000\n",
       "75%      13.000000    14.000000\n",
       "max      31.000000    29.000000"
      ]
     },
     "execution_count": 106,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "pd.DataFrame(resi_search).describe()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d436c06",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "02152bac",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'best_match': 'let k := (has_reflect l).arrow',\n",
       " 'score': 0.3333333333333333,\n",
       " 'real': 'haveI := has_reflect',\n",
       " 'all': ['exact (has_lift_',\n",
       "  'rw [mirror_eq_rel.2 _]',\n",
       "  \"rcases exists_mul_eq_mul_aux' (reflect tactic.rcases_patt).exists with ⟨b⟩|rfl\",\n",
       "  'exact is_colimit.reflect_rel (TM1.reflect l)',\n",
       "  \"exact pnat.reflect.reflect' (λ (e : l = 0), by rw [this, (mt (reflect e) l, eq_comm)] )\",\n",
       "  'induction n with n IH',\n",
       "  'exact (has',\n",
       "  'exact rfl',\n",
       "  'rw ← reflection_patt',\n",
       "  'apply_instance',\n",
       "  'induction l with n l ih',\n",
       "  'exact reflect_reflect tactic.rcases_patt',\n",
       "  'induction n with n l IH',\n",
       "  'exact refl_reflect_aux l l',\n",
       "  'exact (reverse_rec _ _',\n",
       "  'exact [rcases_patt, diff_empty.elim]',\n",
       "  'exact (reflect_eq_self',\n",
       "  'conv { to_lhs | (_root_.has_reflect tactic.rcases_patt) <immigrant} { rw ← l, exact (reflect (mt_trans (eq_iff_iff.mp aux.trans_le this))) }',\n",
       "  'let k := (has_reflect l).arrow',\n",
       "  'exact reflection_patt_reflect tactic.rcases_patt',\n",
       "  'induction ih',\n",
       "  'exact (l.',\n",
       "  'induction l with n l IH',\n",
       "  'induction n with n l IH generalizing l',\n",
       "  'rcases l with (_ | ⟨x, hx⟩)',\n",
       "  'coe : listΣ (append_nil : ℕ →',\n",
       "  'induction n with n IH generalizing l',\n",
       "  'substs l l',\n",
       "  'exact (has_lift_t',\n",
       "  'rw [←reflect_out, ←reflect]',\n",
       "  'rw ←_root_.reflect_comp'],\n",
       " 'prompt': 'GOAL has_reflect : _root_.has_reflect tactic.rcases_patt,\\tl : listΣ tactic.rcases_patt\\t⊢ reflected l\\n PROOFSTEP '}"
      ]
     },
     "execution_count": 57,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "each = json_c[1000]\n",
    "each"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "947ef08a",
   "metadata": {},
   "outputs": [],
   "source": [
    "test = each[\"prompt\"].replace(\"GOAL\",\"<GOAL>\").replace(\"PROOFSTEP\",\"<PROOFSTEP>\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "ae99ed61",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = [i for i in each[\"all\"]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "id": "b607bdf4",
   "metadata": {},
   "outputs": [],
   "source": [
    "l = []\n",
    "for i in results+[each[\"real\"]]:\n",
    "    inputs = rk.tokenize(test, i, return_tensors='pt')\n",
    "    inputs.to(torch.device(\"cuda:0\"))\n",
    "    score = rk(inputs).logits\n",
    "    l.append((score.cpu().detach().numpy()[0][0],i))\n",
    "    l.sort(key=lambda x:x[0])\n",
    "    l.reverse()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "id": "d5c2f434",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(1.0750328, 'exact (has_lift_t'),\n",
       " (0.51294786, 'exact (has_lift_'),\n",
       " (-0.49501634, 'rw ← reflection_patt'),\n",
       " (-0.5962713, 'induction l with n l IH'),\n",
       " (-0.88799334, 'let k := (has_reflect l).arrow'),\n",
       " (-0.8978929, 'exact reflect_reflect tactic.rcases_patt'),\n",
       " (-0.98143125, 'induction n with n IH'),\n",
       " (-1.0099363, 'apply_instance'),\n",
       " (-1.0381949, 'rcases l with (_ | ⟨x, hx⟩)'),\n",
       " (-1.0494938, 'induction n with n l IH'),\n",
       " (-1.0732956, 'rw ←_root_.reflect_comp'),\n",
       " (-1.0913246, 'substs l l'),\n",
       " (-1.0934861, 'rw [←reflect_out, ←reflect]'),\n",
       " (-1.1548848, 'exact (reflect_eq_self'),\n",
       " (-1.1796035, 'induction l with n l ih'),\n",
       " (-1.1827998, 'haveI := has_reflect'),\n",
       " (-1.2198576, 'exact (has'),\n",
       " (-1.2420163, 'induction ih'),\n",
       " (-1.2687602, 'exact reflection_patt_reflect tactic.rcases_patt'),\n",
       " (-1.3217957, 'rw [mirror_eq_rel.2 _]'),\n",
       " (-1.4752369, 'exact refl_reflect_aux l l'),\n",
       " (-1.5429989, 'exact (l.'),\n",
       " (-1.6189396,\n",
       "  \"rcases exists_mul_eq_mul_aux' (reflect tactic.rcases_patt).exists with ⟨b⟩|rfl\"),\n",
       " (-1.9771941, 'exact [rcases_patt, diff_empty.elim]'),\n",
       " (-2.3252175, 'exact is_colimit.reflect_rel (TM1.reflect l)'),\n",
       " (-2.509611,\n",
       "  \"exact pnat.reflect.reflect' (λ (e : l = 0), by rw [this, (mt (reflect e) l, eq_comm)] )\"),\n",
       " (-2.5570722, 'induction n with n IH generalizing l'),\n",
       " (-2.742371, 'coe : listΣ (append_nil : ℕ →'),\n",
       " (-2.7648082, 'exact (reverse_rec _ _'),\n",
       " (-3.0234487, 'induction n with n l IH generalizing l'),\n",
       " (-3.05518,\n",
       "  'conv { to_lhs | (_root_.has_reflect tactic.rcases_patt) <immigrant} { rw ← l, exact (reflect (mt_trans (eq_iff_iff.mp aux.trans_le this))) }'),\n",
       " (-3.2146297, 'exact rfl')]"
      ]
     },
     "execution_count": 89,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "l"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "id": "3cb4fe04",
   "metadata": {},
   "outputs": [],
   "source": [
    "l_e = [i[1] for i in l]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "id": "1021dac4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['exact (has_lift_t',\n",
       " 'exact (has_lift_',\n",
       " 'rw ← reflection_patt',\n",
       " 'induction l with n l IH',\n",
       " 'let k := (has_reflect l).arrow',\n",
       " 'exact reflect_reflect tactic.rcases_patt',\n",
       " 'induction n with n IH',\n",
       " 'apply_instance',\n",
       " 'rcases l with (_ | ⟨x, hx⟩)',\n",
       " 'induction n with n l IH',\n",
       " 'rw ←_root_.reflect_comp',\n",
       " 'substs l l',\n",
       " 'rw [←reflect_out, ←reflect]',\n",
       " 'exact (reflect_eq_self',\n",
       " 'induction l with n l ih',\n",
       " 'haveI := has_reflect',\n",
       " 'exact (has',\n",
       " 'induction ih',\n",
       " 'exact reflection_patt_reflect tactic.rcases_patt',\n",
       " 'rw [mirror_eq_rel.2 _]',\n",
       " 'exact refl_reflect_aux l l',\n",
       " 'exact (l.',\n",
       " \"rcases exists_mul_eq_mul_aux' (reflect tactic.rcases_patt).exists with ⟨b⟩|rfl\",\n",
       " 'exact [rcases_patt, diff_empty.elim]',\n",
       " 'exact is_colimit.reflect_rel (TM1.reflect l)',\n",
       " \"exact pnat.reflect.reflect' (λ (e : l = 0), by rw [this, (mt (reflect e) l, eq_comm)] )\",\n",
       " 'induction n with n IH generalizing l',\n",
       " 'coe : listΣ (append_nil : ℕ →',\n",
       " 'exact (reverse_rec _ _',\n",
       " 'induction n with n l IH generalizing l',\n",
       " 'conv { to_lhs | (_root_.has_reflect tactic.rcases_patt) <immigrant} { rw ← l, exact (reflect (mt_trans (eq_iff_iff.mp aux.trans_le this))) }',\n",
       " 'exact rfl']"
      ]
     },
     "execution_count": 87,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "l_e"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "id": "ba988df8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'haveI := has_reflect'"
      ]
     },
     "execution_count": 76,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "each[\"real\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "id": "b030ba0c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "15"
      ]
     },
     "execution_count": 79,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "l_e.index(each[\"real\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "id": "ae009872",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4"
      ]
     },
     "execution_count": 80,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "l_e.index(each[\"best_match\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "id": "8d25571b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'let k := (has_reflect l).arrow'"
      ]
     },
     "execution_count": 88,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "each[\"best_match\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d7c74de",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2a8ef80",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a85aca1",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f2b53fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"GOAL α : Type u,\\t_inst_1 : inhabited α,\\tb : buffer α,\\ti : ℕ,\\th : i < b.size\\t⊢ b.read ⟨i, h⟩ = b.read i\\n <PROOFSTEP> cases \""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "7a05ff29",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n",
      "data_test.jsonl           data_train_prepared.jsonl  \u001b[0m\u001b[01;34mwandb\u001b[0m/\r\n",
      "data_test_prepared.jsonl  data_valid.jsonl\r\n",
      "data_train.jsonl          data_valid_prepared.jsonl\r\n"
     ]
    }
   ],
   "source": [
    "ls files_upload"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4909b495",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "with open('files_upload/data_train_prepared.jsonl', 'r') as json_file:\n",
    "    json_c_all = [json.loads(i) for i in list(json_file)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "id": "64230281",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "168590"
      ]
     },
     "execution_count": 65,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(json_c_all)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "969ed9ea",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0ac6ff6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5de484a2",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "e304e0a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f6198d05",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n",
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    }
   ],
   "source": [
    "\n",
    "hf_tokenizer = AutoTokenizer.from_pretrained(\"EleutherAI/gpt-neo-125M\",    \n",
    "                            bos_token=\"<|startoftext|>\",\n",
    "                            eos_token=\"<|endoftext|>\",\n",
    "                            pad_token=\"<|pad|>\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b9f0b9ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import GPTNeoModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "064cf265",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_args = {\n",
    "    \"train_dir\":\"data/train\",\n",
    "    \"train_path\":\"data/train/train.json\",\n",
    "    \"pred_dir\":\"data/dev\",\n",
    "    \"pred_path\":\"data/train/dev.json\",\n",
    "    \"train_group_size\":8\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "dd04a9f5",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using custom data configuration default-21c912046bcb6d4c\n",
      "Reusing dataset json (/root/.cache/huggingface/datasets/json/default-21c912046bcb6d4c/0.0.0/ac0ca5f5289a6cf108e706efcf040422dbbfa8e658dee6a819f20d76bb84d26b)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4ebdb9d53f1441018c6bdf00699582b3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from reranker.data import GroupedTrainDataset\n",
    "train_dataset = GroupedTrainDataset(\n",
    "    data_args, data_args[\"train_path\"], \n",
    "    tokenizer=hf_tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "f513388b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{'pid': '6575927600422562800', 'passage': [31653, 685, 29705, 238, 23611, 13, 23661, 357, 421, 313, 13, 28015, 62, 27363, 62, 28015, 62, 1659, 62, 27363, 4808, 828, 17804, 238, 23611, 13, 12543, 62, 15699, 62, 16]}, {'pid': '928431572601992639', 'passage': [31653, 685, 29705, 238, 23611, 13, 23661, 357, 421, 313, 13, 28015, 62, 27363, 62, 28015, 62, 1659, 62, 27363, 4808, 828, 17804, 238, 23611, 13, 26282, 62, 261, 6, 62]}, {'pid': '7552767570238416096', 'passage': [31653, 685, 29705, 238, 23611, 13, 23661, 357, 421, 313, 13, 28015, 62, 27363, 62, 28015, 4808, 4808, 828, 17804, 238, 23611, 13, 28015, 62, 27363, 62, 28015, 11, 17804, 238]}, {'pid': '-2234788984240462035', 'passage': [1069, 529]}, {'pid': '-4411941448410955259', 'passage': [31653, 685, 29705, 238, 23611, 13, 23661, 357, 421, 313, 13, 28015, 62, 27363, 62, 28015, 62, 1659, 62, 27363, 4808, 828, 17804, 238, 23611, 13, 23661, 357, 421, 313, 13]}, {'pid': '3122899436888549272', 'passage': [742]}, {'pid': '-8839975914427080847', 'passage': [90]}, {'pid': '-2432619728502266055', 'passage': [67]}, {'pid': '-7361431018201337745', 'passage': [7266, 301]}, {'pid': '4194638930665964425', 'passage': [81]}]\n"
     ]
    },
    {
     "ename": "TypeError",
     "evalue": "TextInputSequence must be str",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "Input \u001b[0;32mIn [10]\u001b[0m, in \u001b[0;36m<cell line: 1>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mtrain_dataset\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m]\u001b[49m\n",
      "File \u001b[0;32m/opt/conda/lib/python3.8/site-packages/reranker/data.py:103\u001b[0m, in \u001b[0;36mGroupedTrainDataset.__getitem__\u001b[0;34m(self, item)\u001b[0m\n\u001b[1;32m    100\u001b[0m     examples \u001b[38;5;241m=\u001b[39m examples[\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mchunk_start: \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mchunk_end]\n\u001b[1;32m    102\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m e \u001b[38;5;129;01min\u001b[39;00m examples:\n\u001b[0;32m--> 103\u001b[0m     group_batch\u001b[38;5;241m.\u001b[39mappend(\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcreate_one_example\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43me\u001b[49m\u001b[43m)\u001b[49m)\n\u001b[1;32m    104\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m group_batch\n",
      "File \u001b[0;32m/opt/conda/lib/python3.8/site-packages/reranker/data.py:68\u001b[0m, in \u001b[0;36mGroupedTrainDataset.create_one_example\u001b[0;34m(self, qry_encoding, doc_encoding)\u001b[0m\n\u001b[1;32m     67\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mcreate_one_example\u001b[39m(\u001b[38;5;28mself\u001b[39m, qry_encoding: List[\u001b[38;5;28mint\u001b[39m], doc_encoding: List[\u001b[38;5;28mint\u001b[39m]):\n\u001b[0;32m---> 68\u001b[0m     item \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtok\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mencode_plus\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m     69\u001b[0m \u001b[43m        \u001b[49m\u001b[43mqry_encoding\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     70\u001b[0m \u001b[43m        \u001b[49m\u001b[43mdoc_encoding\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     71\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtruncation\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43monly_second\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     72\u001b[0m \u001b[43m        \u001b[49m\u001b[43mmax_length\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m128\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     73\u001b[0m \u001b[43m        \u001b[49m\u001b[43mpadding\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mFalse\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m     74\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     75\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m item\n",
      "File \u001b[0;32m/opt/conda/lib/python3.8/site-packages/transformers/tokenization_utils_base.py:2548\u001b[0m, in \u001b[0;36mPreTrainedTokenizerBase.encode_plus\u001b[0;34m(self, text, text_pair, add_special_tokens, padding, truncation, max_length, stride, is_split_into_words, pad_to_multiple_of, return_tensors, return_token_type_ids, return_attention_mask, return_overflowing_tokens, return_special_tokens_mask, return_offsets_mapping, return_length, verbose, **kwargs)\u001b[0m\n\u001b[1;32m   2538\u001b[0m \u001b[38;5;66;03m# Backward compatibility for 'truncation_strategy', 'pad_to_max_length'\u001b[39;00m\n\u001b[1;32m   2539\u001b[0m padding_strategy, truncation_strategy, max_length, kwargs \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_get_padding_truncation_strategies(\n\u001b[1;32m   2540\u001b[0m     padding\u001b[38;5;241m=\u001b[39mpadding,\n\u001b[1;32m   2541\u001b[0m     truncation\u001b[38;5;241m=\u001b[39mtruncation,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m   2545\u001b[0m     \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs,\n\u001b[1;32m   2546\u001b[0m )\n\u001b[0;32m-> 2548\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_encode_plus\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   2549\u001b[0m \u001b[43m    \u001b[49m\u001b[43mtext\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtext\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   2550\u001b[0m \u001b[43m    \u001b[49m\u001b[43mtext_pair\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtext_pair\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   2551\u001b[0m \u001b[43m    \u001b[49m\u001b[43madd_special_tokens\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43madd_special_tokens\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   2552\u001b[0m \u001b[43m    \u001b[49m\u001b[43mpadding_strategy\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mpadding_strategy\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   2553\u001b[0m \u001b[43m    \u001b[49m\u001b[43mtruncation_strategy\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtruncation_strategy\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   2554\u001b[0m \u001b[43m    \u001b[49m\u001b[43mmax_length\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mmax_length\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   2555\u001b[0m \u001b[43m    \u001b[49m\u001b[43mstride\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mstride\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   2556\u001b[0m \u001b[43m    \u001b[49m\u001b[43mis_split_into_words\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mis_split_into_words\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   2557\u001b[0m \u001b[43m    \u001b[49m\u001b[43mpad_to_multiple_of\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mpad_to_multiple_of\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   2558\u001b[0m \u001b[43m    \u001b[49m\u001b[43mreturn_tensors\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mreturn_tensors\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   2559\u001b[0m \u001b[43m    \u001b[49m\u001b[43mreturn_token_type_ids\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mreturn_token_type_ids\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   2560\u001b[0m \u001b[43m    \u001b[49m\u001b[43mreturn_attention_mask\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mreturn_attention_mask\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   2561\u001b[0m \u001b[43m    \u001b[49m\u001b[43mreturn_overflowing_tokens\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mreturn_overflowing_tokens\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   2562\u001b[0m \u001b[43m    \u001b[49m\u001b[43mreturn_special_tokens_mask\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mreturn_special_tokens_mask\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   2563\u001b[0m \u001b[43m    \u001b[49m\u001b[43mreturn_offsets_mapping\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mreturn_offsets_mapping\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   2564\u001b[0m \u001b[43m    \u001b[49m\u001b[43mreturn_length\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mreturn_length\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   2565\u001b[0m \u001b[43m    \u001b[49m\u001b[43mverbose\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mverbose\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   2566\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   2567\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/opt/conda/lib/python3.8/site-packages/transformers/models/gpt2/tokenization_gpt2_fast.py:174\u001b[0m, in \u001b[0;36mGPT2TokenizerFast._encode_plus\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    167\u001b[0m is_split_into_words \u001b[38;5;241m=\u001b[39m kwargs\u001b[38;5;241m.\u001b[39mget(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mis_split_into_words\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;28;01mFalse\u001b[39;00m)\n\u001b[1;32m    169\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39madd_prefix_space \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m is_split_into_words, (\n\u001b[1;32m    170\u001b[0m     \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mYou need to instantiate \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__class__\u001b[39m\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m with add_prefix_space=True \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    171\u001b[0m     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mto use it with pretokenized inputs.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    172\u001b[0m )\n\u001b[0;32m--> 174\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_encode_plus\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/opt/conda/lib/python3.8/site-packages/transformers/tokenization_utils_fast.py:498\u001b[0m, in \u001b[0;36mPreTrainedTokenizerFast._encode_plus\u001b[0;34m(self, text, text_pair, add_special_tokens, padding_strategy, truncation_strategy, max_length, stride, is_split_into_words, pad_to_multiple_of, return_tensors, return_token_type_ids, return_attention_mask, return_overflowing_tokens, return_special_tokens_mask, return_offsets_mapping, return_length, verbose, **kwargs)\u001b[0m\n\u001b[1;32m    475\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_encode_plus\u001b[39m(\n\u001b[1;32m    476\u001b[0m     \u001b[38;5;28mself\u001b[39m,\n\u001b[1;32m    477\u001b[0m     text: Union[TextInput, PreTokenizedInput],\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    494\u001b[0m     \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs\n\u001b[1;32m    495\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m BatchEncoding:\n\u001b[1;32m    497\u001b[0m     batched_input \u001b[38;5;241m=\u001b[39m [(text, text_pair)] \u001b[38;5;28;01mif\u001b[39;00m text_pair \u001b[38;5;28;01melse\u001b[39;00m [text]\n\u001b[0;32m--> 498\u001b[0m     batched_output \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_batch_encode_plus\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    499\u001b[0m \u001b[43m        \u001b[49m\u001b[43mbatched_input\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    500\u001b[0m \u001b[43m        \u001b[49m\u001b[43mis_split_into_words\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mis_split_into_words\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    501\u001b[0m \u001b[43m        \u001b[49m\u001b[43madd_special_tokens\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43madd_special_tokens\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    502\u001b[0m \u001b[43m        \u001b[49m\u001b[43mpadding_strategy\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mpadding_strategy\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    503\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtruncation_strategy\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtruncation_strategy\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    504\u001b[0m \u001b[43m        \u001b[49m\u001b[43mmax_length\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mmax_length\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    505\u001b[0m \u001b[43m        \u001b[49m\u001b[43mstride\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mstride\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    506\u001b[0m \u001b[43m        \u001b[49m\u001b[43mpad_to_multiple_of\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mpad_to_multiple_of\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    507\u001b[0m \u001b[43m        \u001b[49m\u001b[43mreturn_tensors\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mreturn_tensors\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    508\u001b[0m \u001b[43m        \u001b[49m\u001b[43mreturn_token_type_ids\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mreturn_token_type_ids\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    509\u001b[0m \u001b[43m        \u001b[49m\u001b[43mreturn_attention_mask\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mreturn_attention_mask\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    510\u001b[0m \u001b[43m        \u001b[49m\u001b[43mreturn_overflowing_tokens\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mreturn_overflowing_tokens\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    511\u001b[0m \u001b[43m        \u001b[49m\u001b[43mreturn_special_tokens_mask\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mreturn_special_tokens_mask\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    512\u001b[0m \u001b[43m        \u001b[49m\u001b[43mreturn_offsets_mapping\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mreturn_offsets_mapping\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    513\u001b[0m \u001b[43m        \u001b[49m\u001b[43mreturn_length\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mreturn_length\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    514\u001b[0m \u001b[43m        \u001b[49m\u001b[43mverbose\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mverbose\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    515\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    516\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    518\u001b[0m     \u001b[38;5;66;03m# Return tensor is None, then we can remove the leading batch axis\u001b[39;00m\n\u001b[1;32m    519\u001b[0m     \u001b[38;5;66;03m# Overflowing tokens are returned as a batch of output so we keep them in this case\u001b[39;00m\n\u001b[1;32m    520\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m return_tensors \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m return_overflowing_tokens:\n",
      "File \u001b[0;32m/opt/conda/lib/python3.8/site-packages/transformers/models/gpt2/tokenization_gpt2_fast.py:164\u001b[0m, in \u001b[0;36mGPT2TokenizerFast._batch_encode_plus\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    158\u001b[0m is_split_into_words \u001b[38;5;241m=\u001b[39m kwargs\u001b[38;5;241m.\u001b[39mget(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mis_split_into_words\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;28;01mFalse\u001b[39;00m)\n\u001b[1;32m    159\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39madd_prefix_space \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m is_split_into_words, (\n\u001b[1;32m    160\u001b[0m     \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mYou need to instantiate \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__class__\u001b[39m\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m with add_prefix_space=True \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    161\u001b[0m     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mto use it with pretokenized inputs.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    162\u001b[0m )\n\u001b[0;32m--> 164\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_batch_encode_plus\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/opt/conda/lib/python3.8/site-packages/transformers/tokenization_utils_fast.py:425\u001b[0m, in \u001b[0;36mPreTrainedTokenizerFast._batch_encode_plus\u001b[0;34m(self, batch_text_or_text_pairs, add_special_tokens, padding_strategy, truncation_strategy, max_length, stride, is_split_into_words, pad_to_multiple_of, return_tensors, return_token_type_ids, return_attention_mask, return_overflowing_tokens, return_special_tokens_mask, return_offsets_mapping, return_length, verbose)\u001b[0m\n\u001b[1;32m    416\u001b[0m \u001b[38;5;66;03m# Set the truncation and padding strategy and restore the initial configuration\u001b[39;00m\n\u001b[1;32m    417\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mset_truncation_and_padding(\n\u001b[1;32m    418\u001b[0m     padding_strategy\u001b[38;5;241m=\u001b[39mpadding_strategy,\n\u001b[1;32m    419\u001b[0m     truncation_strategy\u001b[38;5;241m=\u001b[39mtruncation_strategy,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    422\u001b[0m     pad_to_multiple_of\u001b[38;5;241m=\u001b[39mpad_to_multiple_of,\n\u001b[1;32m    423\u001b[0m )\n\u001b[0;32m--> 425\u001b[0m encodings \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_tokenizer\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mencode_batch\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    426\u001b[0m \u001b[43m    \u001b[49m\u001b[43mbatch_text_or_text_pairs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    427\u001b[0m \u001b[43m    \u001b[49m\u001b[43madd_special_tokens\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43madd_special_tokens\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    428\u001b[0m \u001b[43m    \u001b[49m\u001b[43mis_pretokenized\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mis_split_into_words\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    429\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    431\u001b[0m \u001b[38;5;66;03m# Convert encoding to dict\u001b[39;00m\n\u001b[1;32m    432\u001b[0m \u001b[38;5;66;03m# `Tokens` has type: Tuple[\u001b[39;00m\n\u001b[1;32m    433\u001b[0m \u001b[38;5;66;03m#                       List[Dict[str, List[List[int]]]] or List[Dict[str, 2D-Tensor]],\u001b[39;00m\n\u001b[1;32m    434\u001b[0m \u001b[38;5;66;03m#                       List[EncodingFast]\u001b[39;00m\n\u001b[1;32m    435\u001b[0m \u001b[38;5;66;03m#                    ]\u001b[39;00m\n\u001b[1;32m    436\u001b[0m \u001b[38;5;66;03m# with nested dimensions corresponding to batch, overflows, sequence length\u001b[39;00m\n\u001b[1;32m    437\u001b[0m tokens_and_encodings \u001b[38;5;241m=\u001b[39m [\n\u001b[1;32m    438\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_convert_encoding(\n\u001b[1;32m    439\u001b[0m         encoding\u001b[38;5;241m=\u001b[39mencoding,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    448\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m encoding \u001b[38;5;129;01min\u001b[39;00m encodings\n\u001b[1;32m    449\u001b[0m ]\n",
      "\u001b[0;31mTypeError\u001b[0m: TextInputSequence must be str"
     ]
    }
   ],
   "source": [
    "train_dataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "8008a3b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from reranker import RerankerTrainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "04a5caf7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from reranker.data import GroupedTrainDataset, PredictionDataset, GroupCollator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "adf3c64d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "No `TrainingArguments` passed, using `output_dir=tmp_trainer`.\n",
      "PyTorch: setting up devices\n",
      "The default value for the training argument `--report_to` will change in v5 (from all installed integrations to none). In v5, you will need to use `--report_to all` to get the same behavior as now. You should start updating your code and make this info disappear :-).\n"
     ]
    }
   ],
   "source": [
    "trainer = RerankerTrainer(\n",
    "        model=model,\n",
    "        train_dataset=train_dataset,\n",
    "        data_collator=GroupCollator(hf_tokenizer)\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "483f19cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# training_args = TrainingArguments(\"test-trainer\",fp16=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "7b95edc5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# training_args"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "175e9777",
   "metadata": {},
   "outputs": [],
   "source": [
    "# from transformers import TrainingArguments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "e617b023",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "***** Running training *****\n",
      "  Num examples = 6058\n",
      "  Num Epochs = 3\n",
      "  Instantaneous batch size per device = 8\n",
      "  Total train batch size (w. parallel, distributed & accumulation) = 8\n",
      "  Gradient Accumulation steps = 1\n",
      "  Total optimization steps = 2271\n"
     ]
    },
    {
     "ename": "AttributeError",
     "evalue": "'dict' object has no attribute 'max_len'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "Input \u001b[0;32mIn [27]\u001b[0m, in \u001b[0;36m<cell line: 1>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mtrainer\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtrain\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/opt/conda/lib/python3.8/site-packages/transformers/trainer.py:1396\u001b[0m, in \u001b[0;36mTrainer.train\u001b[0;34m(self, resume_from_checkpoint, trial, ignore_keys_for_eval, **kwargs)\u001b[0m\n\u001b[1;32m   1393\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcontrol \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcallback_handler\u001b[38;5;241m.\u001b[39mon_epoch_begin(args, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstate, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcontrol)\n\u001b[1;32m   1395\u001b[0m step \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m\n\u001b[0;32m-> 1396\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m step, inputs \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28menumerate\u001b[39m(epoch_iterator):\n\u001b[1;32m   1397\u001b[0m \n\u001b[1;32m   1398\u001b[0m     \u001b[38;5;66;03m# Skip past any already trained steps if resuming training\u001b[39;00m\n\u001b[1;32m   1399\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m steps_trained_in_current_epoch \u001b[38;5;241m>\u001b[39m \u001b[38;5;241m0\u001b[39m:\n\u001b[1;32m   1400\u001b[0m         steps_trained_in_current_epoch \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m1\u001b[39m\n",
      "File \u001b[0;32m/opt/conda/lib/python3.8/site-packages/torch/utils/data/dataloader.py:530\u001b[0m, in \u001b[0;36m_BaseDataLoaderIter.__next__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    528\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_sampler_iter \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    529\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_reset()\n\u001b[0;32m--> 530\u001b[0m data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_next_data\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    531\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_num_yielded \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m1\u001b[39m\n\u001b[1;32m    532\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_dataset_kind \u001b[38;5;241m==\u001b[39m _DatasetKind\u001b[38;5;241m.\u001b[39mIterable \u001b[38;5;129;01mand\u001b[39;00m \\\n\u001b[1;32m    533\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_IterableDataset_len_called \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mand\u001b[39;00m \\\n\u001b[1;32m    534\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_num_yielded \u001b[38;5;241m>\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_IterableDataset_len_called:\n",
      "File \u001b[0;32m/opt/conda/lib/python3.8/site-packages/torch/utils/data/dataloader.py:570\u001b[0m, in \u001b[0;36m_SingleProcessDataLoaderIter._next_data\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    568\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_next_data\u001b[39m(\u001b[38;5;28mself\u001b[39m):\n\u001b[1;32m    569\u001b[0m     index \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_next_index()  \u001b[38;5;66;03m# may raise StopIteration\u001b[39;00m\n\u001b[0;32m--> 570\u001b[0m     data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_dataset_fetcher\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfetch\u001b[49m\u001b[43m(\u001b[49m\u001b[43mindex\u001b[49m\u001b[43m)\u001b[49m  \u001b[38;5;66;03m# may raise StopIteration\u001b[39;00m\n\u001b[1;32m    571\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_pin_memory:\n\u001b[1;32m    572\u001b[0m         data \u001b[38;5;241m=\u001b[39m _utils\u001b[38;5;241m.\u001b[39mpin_memory\u001b[38;5;241m.\u001b[39mpin_memory(data)\n",
      "File \u001b[0;32m/opt/conda/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py:49\u001b[0m, in \u001b[0;36m_MapDatasetFetcher.fetch\u001b[0;34m(self, possibly_batched_index)\u001b[0m\n\u001b[1;32m     47\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mfetch\u001b[39m(\u001b[38;5;28mself\u001b[39m, possibly_batched_index):\n\u001b[1;32m     48\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mauto_collation:\n\u001b[0;32m---> 49\u001b[0m         data \u001b[38;5;241m=\u001b[39m [\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdataset[idx] \u001b[38;5;28;01mfor\u001b[39;00m idx \u001b[38;5;129;01min\u001b[39;00m possibly_batched_index]\n\u001b[1;32m     50\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m     51\u001b[0m         data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdataset[possibly_batched_index]\n",
      "File \u001b[0;32m/opt/conda/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py:49\u001b[0m, in \u001b[0;36m<listcomp>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m     47\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mfetch\u001b[39m(\u001b[38;5;28mself\u001b[39m, possibly_batched_index):\n\u001b[1;32m     48\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mauto_collation:\n\u001b[0;32m---> 49\u001b[0m         data \u001b[38;5;241m=\u001b[39m [\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdataset\u001b[49m\u001b[43m[\u001b[49m\u001b[43midx\u001b[49m\u001b[43m]\u001b[49m \u001b[38;5;28;01mfor\u001b[39;00m idx \u001b[38;5;129;01min\u001b[39;00m possibly_batched_index]\n\u001b[1;32m     50\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m     51\u001b[0m         data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdataset[possibly_batched_index]\n",
      "File \u001b[0;32m/opt/conda/lib/python3.8/site-packages/reranker/data.py:103\u001b[0m, in \u001b[0;36m__getitem__\u001b[0;34m(self, item)\u001b[0m\n\u001b[1;32m    101\u001b[0m         padding=False,\n\u001b[1;32m    102\u001b[0m     ) @dataclass class GroupCollator(DataCollatorWithPadding):\n\u001b[0;32m--> 103\u001b[0m \"\"\"\n\u001b[1;32m    104\u001b[0m Wrapper that does conversion from List[Tuple[encode_qry, encode_psg]] to List[qry], List[psg]\n\u001b[1;32m    105\u001b[0m and pass batch separately to the actual collator.\n",
      "File \u001b[0;32m/opt/conda/lib/python3.8/site-packages/reranker/data.py:72\u001b[0m, in \u001b[0;36mcreate_one_example\u001b[0;34m(self, qry_encoding, doc_encoding)\u001b[0m\n\u001b[1;32m     70\u001b[0m     negs = random.choices(group['neg'], k=8 - 1)\n\u001b[1;32m     71\u001b[0m else:\n\u001b[0;32m---> 72\u001b[0m     negs = random.sample(group['neg'], k=8- 1)\n\u001b[1;32m     73\u001b[0m for neg_entry in negs:\n\u001b[1;32m     74\u001b[0m     _, neg_psg = [neg_entry[k] for k in self.document_columns]\n",
      "\u001b[0;31mAttributeError\u001b[0m: 'dict' object has no attribute 'max_len'"
     ]
    }
   ],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4cb3c52",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d542f90",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eee24662",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e9d4e42",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "3b1ca9b7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['last_hidden_state', 'past_key_values']"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[i for i in hf_model(**inputs).last_hidden_state]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "id": "3f165b5a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "GPTNeoModel(\n",
       "  (wte): Embedding(50259, 768)\n",
       "  (wpe): Embedding(2048, 768)\n",
       "  (drop): Dropout(p=0, inplace=False)\n",
       "  (h): ModuleList(\n",
       "    (0): GPTNeoBlock(\n",
       "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (attn): GPTNeoAttention(\n",
       "        (attention): GPTNeoSelfAttention(\n",
       "          (attn_dropout): Dropout(p=0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0, inplace=False)\n",
       "          (k_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (v_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (q_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (out_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (mlp): GPTNeoMLP(\n",
       "        (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "        (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "        (act): NewGELUActivation()\n",
       "        (dropout): Dropout(p=0, inplace=False)\n",
       "      )\n",
       "    )\n",
       "    (1): GPTNeoBlock(\n",
       "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (attn): GPTNeoAttention(\n",
       "        (attention): GPTNeoSelfAttention(\n",
       "          (attn_dropout): Dropout(p=0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0, inplace=False)\n",
       "          (k_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (v_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (q_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (out_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (mlp): GPTNeoMLP(\n",
       "        (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "        (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "        (act): NewGELUActivation()\n",
       "        (dropout): Dropout(p=0, inplace=False)\n",
       "      )\n",
       "    )\n",
       "    (2): GPTNeoBlock(\n",
       "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (attn): GPTNeoAttention(\n",
       "        (attention): GPTNeoSelfAttention(\n",
       "          (attn_dropout): Dropout(p=0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0, inplace=False)\n",
       "          (k_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (v_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (q_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (out_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (mlp): GPTNeoMLP(\n",
       "        (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "        (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "        (act): NewGELUActivation()\n",
       "        (dropout): Dropout(p=0, inplace=False)\n",
       "      )\n",
       "    )\n",
       "    (3): GPTNeoBlock(\n",
       "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (attn): GPTNeoAttention(\n",
       "        (attention): GPTNeoSelfAttention(\n",
       "          (attn_dropout): Dropout(p=0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0, inplace=False)\n",
       "          (k_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (v_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (q_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (out_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (mlp): GPTNeoMLP(\n",
       "        (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "        (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "        (act): NewGELUActivation()\n",
       "        (dropout): Dropout(p=0, inplace=False)\n",
       "      )\n",
       "    )\n",
       "    (4): GPTNeoBlock(\n",
       "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (attn): GPTNeoAttention(\n",
       "        (attention): GPTNeoSelfAttention(\n",
       "          (attn_dropout): Dropout(p=0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0, inplace=False)\n",
       "          (k_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (v_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (q_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (out_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (mlp): GPTNeoMLP(\n",
       "        (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "        (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "        (act): NewGELUActivation()\n",
       "        (dropout): Dropout(p=0, inplace=False)\n",
       "      )\n",
       "    )\n",
       "    (5): GPTNeoBlock(\n",
       "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (attn): GPTNeoAttention(\n",
       "        (attention): GPTNeoSelfAttention(\n",
       "          (attn_dropout): Dropout(p=0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0, inplace=False)\n",
       "          (k_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (v_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (q_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (out_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (mlp): GPTNeoMLP(\n",
       "        (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "        (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "        (act): NewGELUActivation()\n",
       "        (dropout): Dropout(p=0, inplace=False)\n",
       "      )\n",
       "    )\n",
       "    (6): GPTNeoBlock(\n",
       "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (attn): GPTNeoAttention(\n",
       "        (attention): GPTNeoSelfAttention(\n",
       "          (attn_dropout): Dropout(p=0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0, inplace=False)\n",
       "          (k_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (v_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (q_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (out_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (mlp): GPTNeoMLP(\n",
       "        (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "        (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "        (act): NewGELUActivation()\n",
       "        (dropout): Dropout(p=0, inplace=False)\n",
       "      )\n",
       "    )\n",
       "    (7): GPTNeoBlock(\n",
       "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (attn): GPTNeoAttention(\n",
       "        (attention): GPTNeoSelfAttention(\n",
       "          (attn_dropout): Dropout(p=0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0, inplace=False)\n",
       "          (k_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (v_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (q_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (out_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (mlp): GPTNeoMLP(\n",
       "        (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "        (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "        (act): NewGELUActivation()\n",
       "        (dropout): Dropout(p=0, inplace=False)\n",
       "      )\n",
       "    )\n",
       "    (8): GPTNeoBlock(\n",
       "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (attn): GPTNeoAttention(\n",
       "        (attention): GPTNeoSelfAttention(\n",
       "          (attn_dropout): Dropout(p=0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0, inplace=False)\n",
       "          (k_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (v_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (q_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (out_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (mlp): GPTNeoMLP(\n",
       "        (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "        (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "        (act): NewGELUActivation()\n",
       "        (dropout): Dropout(p=0, inplace=False)\n",
       "      )\n",
       "    )\n",
       "    (9): GPTNeoBlock(\n",
       "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (attn): GPTNeoAttention(\n",
       "        (attention): GPTNeoSelfAttention(\n",
       "          (attn_dropout): Dropout(p=0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0, inplace=False)\n",
       "          (k_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (v_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (q_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (out_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (mlp): GPTNeoMLP(\n",
       "        (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "        (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "        (act): NewGELUActivation()\n",
       "        (dropout): Dropout(p=0, inplace=False)\n",
       "      )\n",
       "    )\n",
       "    (10): GPTNeoBlock(\n",
       "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (attn): GPTNeoAttention(\n",
       "        (attention): GPTNeoSelfAttention(\n",
       "          (attn_dropout): Dropout(p=0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0, inplace=False)\n",
       "          (k_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (v_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (q_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (out_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (mlp): GPTNeoMLP(\n",
       "        (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "        (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "        (act): NewGELUActivation()\n",
       "        (dropout): Dropout(p=0, inplace=False)\n",
       "      )\n",
       "    )\n",
       "    (11): GPTNeoBlock(\n",
       "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (attn): GPTNeoAttention(\n",
       "        (attention): GPTNeoSelfAttention(\n",
       "          (attn_dropout): Dropout(p=0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0, inplace=False)\n",
       "          (k_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (v_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (q_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (out_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (mlp): GPTNeoMLP(\n",
       "        (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "        (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "        (act): NewGELUActivation()\n",
       "        (dropout): Dropout(p=0, inplace=False)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (ln_f): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       ")"
      ]
     },
     "execution_count": 79,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hf_model.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "28337770",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'json_c_all' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Input \u001b[0;32mIn [30]\u001b[0m, in \u001b[0;36m<cell line: 1>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mjson_c_all\u001b[49m[\u001b[38;5;241m0\u001b[39m]\n",
      "\u001b[0;31mNameError\u001b[0m: name 'json_c_all' is not defined"
     ]
    }
   ],
   "source": [
    "json_c_all[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "id": "b780dc8c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'input_ids': tensor([[11230,  1847,   371,  1058,  5994,   334,    11,   197,    43,  1058,\n",
       "          5994,   410,    11,   197,    44,  1058,  5994,   266,    11,   197,\n",
       "            62,  8625,    62,    16,  1058,   725,    62,  1806,   371,    11,\n",
       "           197,    62,  8625,    62,    17,  1058,  6486,    62,  1806,   406,\n",
       "            11,   197,    62,  8625,    62,    18,  1058,  6486,    62,   282,\n",
       "         29230,   371,   406,    11,   197,    62,  8625,    62,    19,  1058,\n",
       "           751,    62,  9503,    62,  8094,   337,    11,   197,    62,  8625,\n",
       "            62,    20,  1058,  8265,   371,   337,    11,   197,    62,  8625,\n",
       "            62,    21,  1058,  6486,    62,  1806,    62, 21412,   406,   337,\n",
       "            11,   197,    62,  8625,    62,    22,  1058,  6486,    62, 21412,\n",
       "           371,   406,   337,    11,   197,    45,   399,     6,  1058,  6486,\n",
       "            62,  7266, 21412,   371,   406,   337,    11,   197,    71,  1058,\n",
       "         24935,    45,   796, 24935,    45,  3256,   197,    76,  1058,   337,\n",
       "           197,   158,   232,    95,   285, 18872,   230,   399, 17804,   242,\n",
       "           285, 18872,   230,   399,     6,   198,  1279,  4805,  6684,    37,\n",
       "         42135,    29,   220, 31653,   685, 29705,   238,  1066,    62, 49270,\n",
       "            62,  7266, 21412,    11,   289,    60,   198]], device='cuda:0'), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]],\n",
       "       device='cuda:0')}"
      ]
     },
     "execution_count": 70,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "inputs.to(torch.device(\"cuda:0\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "9ceb2b6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "id": "32ef40f1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 34.9 ms, sys: 20.5 ms, total: 55.3 ms\n",
      "Wall time: 51.9 ms\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[ 2.5791e-01,  2.2195e-01,  6.1754e-01, -4.1989e-01, -1.1693e+00,\n",
       "         -4.2416e-01, -1.2921e-02, -6.3473e-02,  1.8195e-01, -2.2530e-02,\n",
       "          6.7583e-01, -5.1123e-01,  8.9316e-02,  7.8981e-01, -2.5015e-01,\n",
       "          1.9408e-01,  6.3822e-01,  9.7251e-01, -7.8701e-02, -6.0445e-01,\n",
       "         -9.8325e-01,  2.0833e-01, -1.9024e+00,  9.8087e-01,  2.8801e-02,\n",
       "         -5.0874e-01,  1.5393e+00, -3.0677e-01,  1.4771e-01, -1.6751e+00,\n",
       "         -1.6980e+00,  8.8647e-01,  8.6847e-01,  8.0590e-01, -1.2548e+00,\n",
       "         -1.8401e-02, -5.4263e-01, -6.5954e-01, -3.4339e-01, -4.9939e-01,\n",
       "          5.0794e-01,  4.8216e-01,  6.5291e-02,  6.9481e-02,  9.3943e-01,\n",
       "          1.1056e-01, -9.0834e-01,  5.9255e-01,  6.8692e-01,  2.9180e-01,\n",
       "          3.6947e-01,  7.4422e-01,  1.4877e-01,  4.4374e-01,  1.0055e+00,\n",
       "         -7.2418e-01,  1.5478e-01,  2.2642e-01,  2.0349e-01,  6.3065e-01,\n",
       "         -4.8900e-01, -1.1986e+00, -2.5607e-01,  9.0008e-01,  1.2902e-01,\n",
       "          5.3663e+00, -1.3429e-01,  7.9810e-01, -9.2527e-01,  4.1761e-01,\n",
       "         -2.6186e-01, -2.0151e-01, -3.0563e-01, -2.9163e-01,  4.9363e-01,\n",
       "         -1.0187e+00,  8.9439e-02, -1.5086e-01,  2.4725e-01, -1.8182e-01,\n",
       "         -5.4884e-01,  1.3904e+00, -7.3445e-01,  1.7192e-01, -8.2985e-01,\n",
       "          3.2465e-02, -5.8382e+00, -1.2886e-01,  1.2982e+00, -2.2217e-01,\n",
       "          1.2597e+00, -2.0946e-01, -1.3153e+00,  1.8629e-01,  5.6036e-01,\n",
       "         -8.0274e-03,  3.7031e-02, -7.2689e-02, -4.2259e-01, -7.3990e-01,\n",
       "         -7.2865e-01, -3.5358e-01, -1.0636e+00, -7.8831e-02, -1.8759e-01,\n",
       "          8.3637e-01,  3.1117e-01,  6.1989e-01, -1.2414e-01,  5.0989e-01,\n",
       "          1.1214e+00, -1.0783e+00, -5.9599e-01,  7.6019e-01, -7.4424e-01,\n",
       "          7.1560e-02,  2.7303e-01, -9.0587e-01,  4.4073e-01,  1.9910e-01,\n",
       "         -6.6675e-02,  3.5972e-01,  2.7237e-01,  4.4705e-01,  1.5214e+00,\n",
       "          2.8673e-01, -9.4406e-01, -1.5607e-01, -8.6961e-01,  3.1393e-01,\n",
       "          1.2868e-01,  7.9200e-02, -2.5765e-01, -5.2737e-01,  2.6680e-02,\n",
       "          1.1275e+00,  1.4169e-01, -4.6904e-01,  4.5877e-01, -1.5315e-01,\n",
       "          1.8020e-01,  2.1740e-01, -5.3287e-01, -2.6836e-01,  2.9287e-01,\n",
       "         -4.4783e-01,  1.4212e-01,  4.7490e-01,  6.1778e-02,  4.4987e-01,\n",
       "          1.7372e-01,  5.6562e-02, -1.5217e+00,  1.1232e+00,  5.5890e-01,\n",
       "          4.6813e-01, -7.9550e-01, -1.7125e+00,  5.0035e-01, -4.3188e-01,\n",
       "         -1.6289e+00,  1.5501e-02, -1.0270e+00,  2.3431e-01, -3.7337e-01,\n",
       "         -3.9971e-01,  5.9703e-02,  5.1988e-01,  3.0298e-01, -9.9593e-01,\n",
       "          9.9599e-02, -1.8770e-01,  1.0976e-03,  3.1389e-01, -6.7635e-02,\n",
       "         -2.5654e-01, -1.0079e+00, -3.3847e-01,  1.0830e+00, -8.5420e-02,\n",
       "         -6.8329e-01, -2.2209e-01,  9.8029e-01, -3.5572e-01, -9.7739e-02,\n",
       "          3.1339e-01, -7.5790e-01, -1.1114e-01,  1.3195e+00, -2.6651e-01,\n",
       "         -8.6284e-02, -6.1514e-02, -9.6365e-01,  5.8111e-01, -7.9026e-01,\n",
       "          3.4359e-02,  1.4805e-01, -3.9151e-01, -3.3965e-01,  1.3782e+00,\n",
       "          1.5073e-01,  5.0220e-01,  9.1980e-02, -1.2238e-01, -5.1866e-02,\n",
       "         -1.1727e-01, -8.5541e-01, -3.1244e-01,  1.1770e-01,  1.7958e+00,\n",
       "         -2.2563e-01, -5.6103e-01, -1.1069e+00, -9.7789e-01, -1.2189e+00,\n",
       "         -1.0113e+00,  3.3071e-01,  1.4034e+00, -5.0933e-01, -6.8410e-02,\n",
       "          3.9110e-01, -2.6448e-01, -1.6042e-01, -6.9321e-01,  3.4150e-01,\n",
       "         -5.1736e-01,  1.5081e-01, -7.2690e-02, -4.4258e-01,  9.6587e-01,\n",
       "          1.1243e+00,  3.6307e-01, -8.3291e-02,  9.4453e-01,  5.9471e-01,\n",
       "         -1.9810e-01, -1.2600e+00,  1.1084e-01,  6.1746e-01, -4.3833e-01,\n",
       "          4.9749e-01, -6.7068e-01,  7.3220e-01, -3.6291e-01,  3.4117e+00,\n",
       "          4.2510e-01, -1.3726e+00,  9.1664e-01,  3.2722e-01,  4.9647e-02,\n",
       "          2.7723e-01,  1.3886e+00,  2.0157e+00,  3.6911e-01,  2.2851e-01,\n",
       "         -1.3195e+00, -2.0280e-01,  2.3085e-01, -1.2566e+00, -6.1733e-02,\n",
       "         -4.9790e-01, -4.2601e-01, -8.4731e-01, -1.2078e-01, -8.5055e-01,\n",
       "          8.4002e-02, -8.0585e-01, -8.6329e-01,  5.9033e-02,  5.8254e-01,\n",
       "         -7.6936e-01,  1.0379e+00,  1.8664e-01,  2.5744e-01, -3.4262e-01,\n",
       "          1.7607e+00, -2.5542e-01, -5.9538e-01,  4.8452e+00,  5.2938e-01,\n",
       "          4.5253e-01,  5.4308e-01,  7.8852e-01, -8.5713e-01,  5.7179e-01,\n",
       "         -2.5038e-02,  4.0092e-01,  1.6203e+00,  8.9025e-01, -3.7051e-01,\n",
       "          2.8878e-01, -4.0672e-01, -2.8992e-01,  1.2436e+00,  8.4281e-01,\n",
       "          3.3404e-01,  3.5212e+00, -6.2464e-01,  1.6889e-01,  5.0950e-01,\n",
       "          6.0882e-01, -1.0021e+00, -1.4234e+00, -1.8274e-01, -1.1205e+00,\n",
       "          5.8818e-01,  2.0395e-01, -1.5636e-01, -2.7134e-01,  1.7388e-01,\n",
       "          3.6655e-01,  3.5106e-02,  1.1893e+00, -9.5504e-01, -1.7335e+00,\n",
       "          2.6121e-01,  1.5959e-01,  2.1152e-01,  4.6638e-01,  3.0343e-02,\n",
       "          1.1501e+00,  7.6918e-01, -1.0978e-01,  2.6012e-01,  7.4091e-01,\n",
       "         -4.8964e-02, -2.2023e-01,  3.8122e-01,  5.8171e-01, -3.0145e-01,\n",
       "         -2.5431e-01, -5.6066e-02, -1.4672e-02, -7.7137e-01, -7.6175e-02,\n",
       "          7.0992e-03,  3.3074e-01,  2.0417e-01, -2.2275e-02, -1.8799e+00,\n",
       "          3.2200e-01, -3.4532e-01, -1.2397e+00,  7.1686e-01, -9.5655e-01,\n",
       "          6.4862e-01, -1.4755e-01, -1.1373e+00, -4.0796e-01, -9.8583e-01,\n",
       "         -7.0172e-01,  3.0339e-02,  9.0085e-01, -5.3020e-01,  1.2186e-01,\n",
       "          1.3400e-01,  3.0865e-01,  4.4937e-01, -9.9567e-01,  1.1413e+00,\n",
       "         -9.1110e-01,  4.7118e-01,  4.6427e-01, -1.9850e+00, -4.4892e-02,\n",
       "         -8.3049e-01, -1.2660e-01, -5.3937e-02,  6.8228e-01, -7.8966e-01,\n",
       "          4.0273e-01, -7.6949e-02, -4.4711e-01, -3.7368e-02, -1.4817e-01,\n",
       "         -2.0367e-01,  6.8044e-01, -2.3592e-01,  2.1292e-02, -1.2443e-02,\n",
       "         -2.8481e-01, -1.0849e+00,  2.0934e-02, -9.8099e-01, -2.8614e-01,\n",
       "         -1.1888e-01,  2.3289e+00, -8.4848e-01,  1.0650e+00, -4.8800e-01,\n",
       "         -7.3174e-01,  1.7405e+00, -1.1867e-01,  6.0654e-01, -3.2369e-01,\n",
       "          4.2891e-03,  1.3777e-01, -8.7465e-01,  3.9065e-01,  5.8552e-01,\n",
       "         -1.1822e-01, -2.8735e-01,  4.5932e-01, -1.0558e+00, -2.0773e-01,\n",
       "         -4.2257e-02, -7.4723e-01, -1.0418e-01,  1.1046e-01,  2.5200e-01,\n",
       "         -7.7008e-01, -6.9124e-01,  1.2328e-02,  9.2941e-01, -9.5823e-01,\n",
       "          7.8868e-01,  4.8766e-01,  1.2000e+00,  2.0337e+00,  1.3819e+00,\n",
       "         -7.0696e-01, -3.9163e+00, -4.0065e-01,  2.4839e-01,  6.8717e-01,\n",
       "          2.4711e-01, -2.3006e-01,  8.9260e-01, -1.1988e-01, -1.0582e+00,\n",
       "         -6.1881e-01,  1.9576e-01, -3.5435e-01,  3.7132e-01, -1.0532e+00,\n",
       "          5.1836e-01,  1.4250e-01,  3.1409e-01,  6.4463e-01, -1.0462e+00,\n",
       "          6.6189e-02, -1.4131e-02, -5.6984e-01,  7.3621e-01,  4.5947e-01,\n",
       "         -3.4809e-02, -5.7257e-01, -2.4561e-01, -3.6593e-01, -2.4191e-02,\n",
       "         -8.4636e-01, -2.9382e-01,  8.2950e-01, -2.8704e-01, -6.1543e-01,\n",
       "         -2.8003e-02,  4.3479e-01,  1.0016e+00,  4.4122e-01, -6.2453e-02,\n",
       "         -1.7907e+00,  1.3726e+00, -1.7483e-01, -6.3347e-01, -3.6389e-01,\n",
       "          5.5712e-01,  1.7147e-01, -1.7688e-02,  6.7811e-01,  4.3906e-02,\n",
       "         -5.6973e-01, -4.2983e-01,  7.6825e-01,  1.2913e+00, -4.2006e-01,\n",
       "          5.0041e-01, -5.7768e-01,  1.0845e-01, -5.6376e-01, -3.3116e-01,\n",
       "          1.9473e-01,  1.6148e+00, -7.9727e-01,  1.8664e-01,  6.9911e-01,\n",
       "         -7.8633e-01,  3.7915e+00, -1.5934e-01,  1.3869e-01, -4.3771e-01,\n",
       "         -6.2996e-01,  2.0706e+00,  2.8263e-01, -4.5241e-01, -5.7089e-02,\n",
       "          1.1213e+00, -4.9894e-02,  2.0719e-01, -9.7199e-01, -1.4038e-01,\n",
       "          1.8437e+00, -5.0838e-01, -4.5824e-01,  1.8325e-02,  1.6349e-01,\n",
       "          2.3760e-01,  9.2985e-01,  4.2633e-01,  6.1916e-01, -4.9182e-01,\n",
       "          1.0441e-01,  1.4710e+00, -1.0186e+00,  6.7099e-02,  1.1068e+00,\n",
       "          3.1317e-01, -7.6338e-02,  8.3072e-01,  2.1650e-02, -1.7168e-01,\n",
       "          6.7508e-01,  3.4145e-01, -1.5327e+00,  5.8343e-01, -1.1839e+00,\n",
       "         -6.2361e-02,  5.3014e-02,  1.2807e+00,  1.0915e+00, -8.6758e-01,\n",
       "          1.0344e+00,  3.5148e-01, -9.8511e-01, -7.7864e-01,  1.2014e-02,\n",
       "         -8.4390e-01, -1.0858e+00,  5.2799e-01,  2.7797e-01,  8.7393e-01,\n",
       "          1.6785e+00,  4.9440e-01,  2.5744e-02, -1.8004e+00, -5.2505e-01,\n",
       "         -1.7587e-01, -6.3750e-01,  3.1973e-01, -7.9784e-02, -2.1100e-01,\n",
       "          8.2759e-01, -1.2025e+00, -1.4338e+00,  1.4730e+00, -1.0535e+00,\n",
       "          4.8772e-01,  8.9435e-01, -1.0030e+00,  3.5396e-01, -3.7726e-01,\n",
       "         -4.0844e-01, -6.5378e-01, -2.3089e-01, -1.0119e-01,  4.6434e-01,\n",
       "         -9.2134e-01,  4.6914e-01,  3.7056e-01, -7.9002e-01,  4.0602e-03,\n",
       "         -7.7099e-01,  6.6663e-01,  3.4323e-01, -1.0957e+00,  9.0008e-01,\n",
       "         -4.0340e-01, -8.0557e-01,  2.0100e+00, -1.1193e+00, -3.3462e-01,\n",
       "         -4.5713e-01,  1.7430e+00, -1.3358e+00, -7.0499e-02, -5.5699e-01,\n",
       "          4.2485e-01,  5.6796e-01,  3.1242e-01,  1.5269e+00, -1.2979e+00,\n",
       "         -5.4608e-02,  1.1424e+00,  1.0122e+00,  6.8035e-01,  1.3340e+00,\n",
       "          3.0875e-01,  6.4539e-01, -3.2952e-01, -1.0114e+00,  7.5653e-01,\n",
       "         -5.5787e-01, -1.6231e+00, -4.5545e-01,  6.2937e-02,  7.6943e-01,\n",
       "         -1.0360e-01, -7.2392e-01,  1.4437e+00,  3.0526e-02,  1.0234e-01,\n",
       "         -4.8389e-02,  1.0221e+00,  2.0669e-01, -2.4083e+00,  8.9863e-01,\n",
       "          4.7944e-01, -8.4005e-01, -6.3425e-01, -7.0860e-01, -3.2115e-01,\n",
       "          6.3741e-01, -1.9570e-01, -4.6066e-01,  5.6871e-01,  3.1966e-01,\n",
       "          6.6062e-01, -2.6750e-01,  8.9014e-01, -5.9421e-01,  1.2518e+00,\n",
       "         -6.4939e-01,  9.3352e-01,  1.8120e-01, -2.7597e-01,  1.2792e-01,\n",
       "         -3.4255e-01,  7.5101e-01, -8.0686e-01,  8.6254e-01, -9.3071e-01,\n",
       "          5.9072e-01,  5.2779e-01,  7.6271e-01, -3.3893e-01,  3.8799e-01,\n",
       "         -4.3208e-01, -1.6288e-01,  5.6881e-02,  1.9117e-01,  2.8365e-02,\n",
       "          9.5397e-01, -1.8657e-01,  3.1514e-01, -5.9065e-01,  1.5293e+00,\n",
       "          3.0394e-01, -5.9831e-01, -4.2810e-01,  7.7838e-01, -7.6792e-01,\n",
       "         -7.1487e-01, -1.4372e-01, -1.2557e+00,  7.7492e-02, -5.5653e-01,\n",
       "          1.7257e-01,  6.0696e-01,  4.0132e-01, -4.6850e-01, -3.3263e-01,\n",
       "         -1.7245e+00,  1.2748e+00,  3.4657e-01, -2.1751e-01,  4.7821e-01,\n",
       "          8.0060e-01,  2.2041e-01,  4.5081e-01,  1.1787e+00,  4.2144e-02,\n",
       "         -6.8987e-01,  8.5284e-01,  8.6362e-01,  5.5671e-01, -3.9498e-01,\n",
       "         -8.1284e-01,  3.5733e-01,  2.3513e-01,  8.8358e-01,  7.1591e-01,\n",
       "          2.1957e-01,  2.4697e-01, -7.2402e-01, -3.5483e-01, -2.1268e+00,\n",
       "          1.0866e+00, -7.8985e-01, -1.0879e-01,  7.2387e-02,  1.6729e+00,\n",
       "          1.5013e-01,  5.7198e-01,  4.3690e-01, -7.7675e-01, -1.0156e+00,\n",
       "          7.6743e-01, -9.8260e-02, -9.4869e-01,  9.8678e-01,  5.2385e-01,\n",
       "          3.7922e-01,  2.1336e-01, -1.0158e-01,  1.2670e+00,  5.4585e-01,\n",
       "          7.0972e-01, -1.7041e+00,  1.3264e-01, -2.2960e-01,  8.2514e-01,\n",
       "         -1.0594e+00, -9.2191e-01, -2.6381e+00,  4.5535e-01, -1.5036e+00,\n",
       "         -1.3049e+00,  3.2866e-01, -6.8583e-01, -4.3822e-02, -1.0406e+00,\n",
       "          1.4157e+00,  1.4056e+00,  2.3406e-01, -1.4198e+00, -4.1428e-01,\n",
       "          7.1460e-01, -9.9952e-01, -8.9575e-02, -9.3440e-01,  2.4660e-01,\n",
       "         -1.5474e-02,  1.0905e+00, -2.7923e-01,  7.2442e-02, -1.7743e-01,\n",
       "          7.7267e-03, -2.7213e-01,  2.9660e-02, -5.0046e-01, -1.1490e+00,\n",
       "          9.7945e-01, -1.2734e+00, -6.2709e-01, -3.5033e-01,  2.5188e-01,\n",
       "          3.4297e-01,  1.1887e+00, -2.2985e-01,  8.3060e-01,  5.9193e-02,\n",
       "         -1.3219e+00, -7.2062e-01,  7.3293e-02, -1.1011e+00, -2.7101e-01,\n",
       "          8.7101e-02,  2.4765e-01, -2.5952e-01]], device='cuda:0',\n",
       "       grad_fn=<MeanBackward1>)"
      ]
     },
     "execution_count": 80,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "inputs = rk.tokenize(\"GOAL R : Type u,\\tL : Type v,\\tM : Type w,\\t_inst_1 : comm_ring R,\\t_inst_2 : lie_ring L,\\t_inst_3 : lie_algebra R L,\\t_inst_4 : add_comm_group M,\\t_inst_5 : module R M,\\t_inst_6 : lie_ring_module L M,\\t_inst_7 : lie_module R L M,\\tN N' : lie_submodule R L M,\\th : ↑N = ↑N',\\tm : M\\t⊢ m ∈ N ↔ m ∈ N'\\n <PROOFSTEP> \", \"rw [\\u2190 mem_coe_submodule, h]\\n\", return_tensors='pt')\n",
    "inputs.to(torch.device(\"cuda:0\"))\n",
    "torch.mean(hf_model(**inputs).last_hidden_state,axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "id": "984526f3",
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "expected Tensor as element 0 in argument 0, but got BatchEncoding",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "Input \u001b[0;32mIn [75]\u001b[0m, in \u001b[0;36m<cell line: 1>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcat\u001b[49m\u001b[43m(\u001b[49m\u001b[43m(\u001b[49m\u001b[43minputs\u001b[49m\u001b[43m,\u001b[49m\u001b[43minputs\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[0;31mTypeError\u001b[0m: expected Tensor as element 0 in argument 0, but got BatchEncoding"
     ]
    }
   ],
   "source": [
    "torch.cat((inputs,inputs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "id": "32d6c8c6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "GPTNeoModel(\n",
       "  (wte): Embedding(50259, 768)\n",
       "  (wpe): Embedding(2048, 768)\n",
       "  (drop): Dropout(p=0, inplace=False)\n",
       "  (h): ModuleList(\n",
       "    (0): GPTNeoBlock(\n",
       "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (attn): GPTNeoAttention(\n",
       "        (attention): GPTNeoSelfAttention(\n",
       "          (attn_dropout): Dropout(p=0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0, inplace=False)\n",
       "          (k_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (v_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (q_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (out_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (mlp): GPTNeoMLP(\n",
       "        (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "        (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "        (act): NewGELUActivation()\n",
       "        (dropout): Dropout(p=0, inplace=False)\n",
       "      )\n",
       "    )\n",
       "    (1): GPTNeoBlock(\n",
       "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (attn): GPTNeoAttention(\n",
       "        (attention): GPTNeoSelfAttention(\n",
       "          (attn_dropout): Dropout(p=0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0, inplace=False)\n",
       "          (k_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (v_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (q_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (out_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (mlp): GPTNeoMLP(\n",
       "        (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "        (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "        (act): NewGELUActivation()\n",
       "        (dropout): Dropout(p=0, inplace=False)\n",
       "      )\n",
       "    )\n",
       "    (2): GPTNeoBlock(\n",
       "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (attn): GPTNeoAttention(\n",
       "        (attention): GPTNeoSelfAttention(\n",
       "          (attn_dropout): Dropout(p=0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0, inplace=False)\n",
       "          (k_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (v_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (q_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (out_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (mlp): GPTNeoMLP(\n",
       "        (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "        (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "        (act): NewGELUActivation()\n",
       "        (dropout): Dropout(p=0, inplace=False)\n",
       "      )\n",
       "    )\n",
       "    (3): GPTNeoBlock(\n",
       "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (attn): GPTNeoAttention(\n",
       "        (attention): GPTNeoSelfAttention(\n",
       "          (attn_dropout): Dropout(p=0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0, inplace=False)\n",
       "          (k_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (v_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (q_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (out_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (mlp): GPTNeoMLP(\n",
       "        (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "        (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "        (act): NewGELUActivation()\n",
       "        (dropout): Dropout(p=0, inplace=False)\n",
       "      )\n",
       "    )\n",
       "    (4): GPTNeoBlock(\n",
       "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (attn): GPTNeoAttention(\n",
       "        (attention): GPTNeoSelfAttention(\n",
       "          (attn_dropout): Dropout(p=0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0, inplace=False)\n",
       "          (k_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (v_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (q_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (out_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (mlp): GPTNeoMLP(\n",
       "        (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "        (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "        (act): NewGELUActivation()\n",
       "        (dropout): Dropout(p=0, inplace=False)\n",
       "      )\n",
       "    )\n",
       "    (5): GPTNeoBlock(\n",
       "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (attn): GPTNeoAttention(\n",
       "        (attention): GPTNeoSelfAttention(\n",
       "          (attn_dropout): Dropout(p=0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0, inplace=False)\n",
       "          (k_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (v_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (q_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (out_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (mlp): GPTNeoMLP(\n",
       "        (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "        (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "        (act): NewGELUActivation()\n",
       "        (dropout): Dropout(p=0, inplace=False)\n",
       "      )\n",
       "    )\n",
       "    (6): GPTNeoBlock(\n",
       "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (attn): GPTNeoAttention(\n",
       "        (attention): GPTNeoSelfAttention(\n",
       "          (attn_dropout): Dropout(p=0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0, inplace=False)\n",
       "          (k_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (v_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (q_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (out_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (mlp): GPTNeoMLP(\n",
       "        (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "        (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "        (act): NewGELUActivation()\n",
       "        (dropout): Dropout(p=0, inplace=False)\n",
       "      )\n",
       "    )\n",
       "    (7): GPTNeoBlock(\n",
       "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (attn): GPTNeoAttention(\n",
       "        (attention): GPTNeoSelfAttention(\n",
       "          (attn_dropout): Dropout(p=0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0, inplace=False)\n",
       "          (k_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (v_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (q_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (out_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (mlp): GPTNeoMLP(\n",
       "        (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "        (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "        (act): NewGELUActivation()\n",
       "        (dropout): Dropout(p=0, inplace=False)\n",
       "      )\n",
       "    )\n",
       "    (8): GPTNeoBlock(\n",
       "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (attn): GPTNeoAttention(\n",
       "        (attention): GPTNeoSelfAttention(\n",
       "          (attn_dropout): Dropout(p=0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0, inplace=False)\n",
       "          (k_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (v_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (q_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (out_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (mlp): GPTNeoMLP(\n",
       "        (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "        (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "        (act): NewGELUActivation()\n",
       "        (dropout): Dropout(p=0, inplace=False)\n",
       "      )\n",
       "    )\n",
       "    (9): GPTNeoBlock(\n",
       "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (attn): GPTNeoAttention(\n",
       "        (attention): GPTNeoSelfAttention(\n",
       "          (attn_dropout): Dropout(p=0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0, inplace=False)\n",
       "          (k_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (v_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (q_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (out_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (mlp): GPTNeoMLP(\n",
       "        (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "        (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "        (act): NewGELUActivation()\n",
       "        (dropout): Dropout(p=0, inplace=False)\n",
       "      )\n",
       "    )\n",
       "    (10): GPTNeoBlock(\n",
       "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (attn): GPTNeoAttention(\n",
       "        (attention): GPTNeoSelfAttention(\n",
       "          (attn_dropout): Dropout(p=0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0, inplace=False)\n",
       "          (k_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (v_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (q_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (out_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (mlp): GPTNeoMLP(\n",
       "        (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "        (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "        (act): NewGELUActivation()\n",
       "        (dropout): Dropout(p=0, inplace=False)\n",
       "      )\n",
       "    )\n",
       "    (11): GPTNeoBlock(\n",
       "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (attn): GPTNeoAttention(\n",
       "        (attention): GPTNeoSelfAttention(\n",
       "          (attn_dropout): Dropout(p=0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0, inplace=False)\n",
       "          (k_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (v_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (q_proj): Linear(in_features=768, out_features=768, bias=False)\n",
       "          (out_proj): Linear(in_features=768, out_features=768, bias=True)\n",
       "        )\n",
       "      )\n",
       "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "      (mlp): GPTNeoMLP(\n",
       "        (c_fc): Linear(in_features=768, out_features=3072, bias=True)\n",
       "        (c_proj): Linear(in_features=3072, out_features=768, bias=True)\n",
       "        (act): NewGELUActivation()\n",
       "        (dropout): Dropout(p=0, inplace=False)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (ln_f): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       ")"
      ]
     },
     "execution_count": 76,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hf_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "182200bf",
   "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.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}