src/models/reranker/reranker-train.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6f114ed9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# from reranker import RerankerForInference\n",
    "# rk = RerankerForInference.from_pretrained(\"results/checkpoint-140000/\")  # load checkpoint\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "7d737664",
   "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": 7,
   "id": "f40dc37e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from Reranker.src.reranker.arguments import ModelArguments, DataArguments, \\\n",
    "    RerankerTrainingArguments as TrainingArguments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "fa286904",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "19b97b80",
   "metadata": {},
   "outputs": [],
   "source": [
    "from reranker import Reranker \n",
    "from transformers import AutoTokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "74ab46c1",
   "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 Reranker(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",
    "        \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": null,
   "id": "62b981fa",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "3c31f0fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_args ={\n",
    "    \"train_dir\"\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "071f9d9d",
   "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": [
    "model = Reranker.from_pretrained(\"results/checkpoint-140000/\") "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "7781c5d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# rk.hf_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9597f2cb",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a27ffc0",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75bb4d0c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 150,
   "id": "dcb00a2f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'best_match': 'coe : list΄ pfun := λ n, (list',\n",
       " 'score': 0.3333333333333333,\n",
       " 'real': 'haveI := has_reflect',\n",
       " 'all': ['exact (has_lift_',\n",
       "  'exact (reflect l).symm',\n",
       "  'assumption',\n",
       "  'induction n with n IH',\n",
       "  'induction l with',\n",
       "  'convert to_sublist_of_reflect_nat_abs',\n",
       "  'apply_instance',\n",
       "  'induction l with n l ih',\n",
       "  'coe : list΄ pfun := λ n, (list',\n",
       "  'convert reflect_reflect l with t ht',\n",
       "  \"rw [mirror_eq_write' l]\",\n",
       "  'exact rintro ⟨_, rfl⟩',\n",
       "  'rw ←@reflect.reflect tactic.rcases_patt',\n",
       "  'exact (reflect_eq_',\n",
       "  'convert reflection_reflect tactic.rcases_patt',\n",
       "  'induction n with n l IH',\n",
       "  'exact trans (mfld_trans_gen.mpr (reflect_reflect_reflect_modeq_right t ht)) rfl',\n",
       "  'rw [← reflect_dvd_erson_aux, ← reflect_dvd_iff]',\n",
       "  'exact rintro ⟨a, b, ⟨rfl⟩, rfl⟩',\n",
       "  'exact (reflect_eq_self',\n",
       "  'xt',\n",
       "  'induction ih',\n",
       "  'exact (reverse_rec_',\n",
       "  'rcases this.exists_inv with ⟨R, hR⟩',\n",
       "  'induction l with n l IH',\n",
       "  'induction n with n l IH generalizing l',\n",
       "  'exact (has_coe_to',\n",
       "  'exact ⟨λ pos h1, forall_path h1, pos.reflect, diff h1.symm, diff pos.fst⟩',\n",
       "  'rw [← forall_and_distrib, ←reflecting_C]',\n",
       "  'exact (reflect_iff_mod'],\n",
       " 'prompt': 'GOAL has_reflect : _root_.has_reflect tactic.rcases_patt,\\tl : listΠ tactic.rcases_patt\\t⊢ reflected l\\n PROOFSTEP '}"
      ]
     },
     "execution_count": 150,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "each = json_c[999]\n",
    "each"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 151,
   "id": "b4d585e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "test = each[\"prompt\"].replace(\"GOAL\",\"<GOAL>\").replace(\"PROOFSTEP\",\"<PROOFSTEP>\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 152,
   "id": "47fa1511",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = [i for i in each[\"all\"]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 156,
   "id": "9d01f286",
   "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][1],i))\n",
    "    l.sort(key=lambda x:x[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 157,
   "id": "24f28a1c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(-2.0707211, 'rw [← reflect_dvd_erson_aux, ← reflect_dvd_iff]'),\n",
       " (-1.9039781, 'assumption'),\n",
       " (-1.7388834, 'rw [← forall_and_distrib, ←reflecting_C]'),\n",
       " (-1.6454703,\n",
       "  'exact trans (mfld_trans_gen.mpr (reflect_reflect_reflect_modeq_right t ht)) rfl'),\n",
       " (-1.3055459, 'exact rintro ⟨_, rfl⟩'),\n",
       " (-1.235331, 'exact rintro ⟨a, b, ⟨rfl⟩, rfl⟩'),\n",
       " (-1.2266147, 'convert to_sublist_of_reflect_nat_abs'),\n",
       " (-1.1743462,\n",
       "  'exact ⟨λ pos h1, forall_path h1, pos.reflect, diff h1.symm, diff pos.fst⟩'),\n",
       " (-1.0940642, \"rw [mirror_eq_write' l]\"),\n",
       " (-0.9648125, 'rcases this.exists_inv with ⟨R, hR⟩'),\n",
       " (-0.9581935, 'induction n with n l IH generalizing l'),\n",
       " (-0.92618155, 'rw ←@reflect.reflect tactic.rcases_patt'),\n",
       " (-0.82797396, 'exact (reflect_eq_self'),\n",
       " (-0.8064505, 'convert reflection_reflect tactic.rcases_patt'),\n",
       " (-0.43701363, 'induction n with n l IH'),\n",
       " (-0.38078067, 'induction l with n l IH'),\n",
       " (-0.29419905, 'exact (reflect l).symm'),\n",
       " (-0.24308807, 'induction n with n IH'),\n",
       " (-0.14588475, 'haveI := has_reflect'),\n",
       " (-0.11893135, 'convert reflect_reflect l with t ht'),\n",
       " (-0.04986143, 'exact (has_lift_'),\n",
       " (-0.024886109, 'induction l with n l ih'),\n",
       " (0.02763924, 'induction ih'),\n",
       " (0.031445414, 'coe : list΄ pfun := λ n, (list'),\n",
       " (0.11387861, 'exact (reflect_iff_mod'),\n",
       " (0.14731786, 'exact (has_coe_to'),\n",
       " (0.1782412, 'xt'),\n",
       " (0.21090153, 'apply_instance'),\n",
       " (0.4576639, 'exact (reflect_eq_'),\n",
       " (0.57259375, 'induction l with'),\n",
       " (1.2164512, 'exact (reverse_rec_')]"
      ]
     },
     "execution_count": 157,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "l"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f39d0e6c",
   "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": "36659dea",
   "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": "025d4138",
   "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": "8757814d",
   "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": "b7c4c9c5",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d2943cf",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b82f0d9",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ab3ad558",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "9c253b6f",
   "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": "058c5d59",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import GPTNeoModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "abda3208",
   "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": "d8494197",
   "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": 11,
   "id": "a52bbbce",
   "metadata": {},
   "outputs": [],
   "source": [
    "# train_dataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "38809f4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from reranker import RerankerTrainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "5d234ae6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from reranker.data import GroupedTrainDataset, PredictionDataset, GroupCollator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "77657eb0",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'model' 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 [19]\u001b[0m, in \u001b[0;36m<cell line: 1>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m trainer \u001b[38;5;241m=\u001b[39m RerankerTrainer(\n\u001b[0;32m----> 2\u001b[0m         model\u001b[38;5;241m=\u001b[39m\u001b[43mmodel\u001b[49m,\n\u001b[1;32m      3\u001b[0m         train_dataset\u001b[38;5;241m=\u001b[39mtrain_dataset,\n\u001b[1;32m      4\u001b[0m         data_collator\u001b[38;5;241m=\u001b[39mGroupCollator(hf_tokenizer)\n\u001b[1;32m      5\u001b[0m     )\n",
      "\u001b[0;31mNameError\u001b[0m: name 'model' is not defined"
     ]
    }
   ],
   "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": "c99cedba",
   "metadata": {},
   "outputs": [],
   "source": [
    "# training_args = TrainingArguments(\"test-trainer\",fp16=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "5c902c76",
   "metadata": {},
   "outputs": [],
   "source": [
    "# training_args"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "37dcdae8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# from transformers import TrainingArguments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "bae11f81",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'trainer' 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 [13]\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;241m.\u001b[39mtrain()\n",
      "\u001b[0;31mNameError\u001b[0m: name 'trainer' is not defined"
     ]
    }
   ],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa6231f1",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "021cb691",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a138f88",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68db64c8",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "13779610",
   "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": "08221fdd",
   "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": 81,
   "id": "436712aa",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'prompt': \"GOAL α : Type u,\\tb : buffer α,\\ti : ℕ,\\th : i < b.size,\\tv : α\\t⊢ b.write ⟨i, h⟩ v = b.write' i v\\n PROOFSTEP \",\n",
       " 'completion': \" cases b; unfold write write'; simp [array.write_eq_write']\\n\"}"
      ]
     },
     "execution_count": 81,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "json_c_all[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "id": "94581431",
   "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": "c37563dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "id": "a65b11c1",
   "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": "e3d8b8d3",
   "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": "257a145d",
   "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": "0d074ea0",
   "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
}