gpt-j-6B.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c2d31f80",
   "metadata": {},
   "source": [
    "# With the official addition of GPTJ to the Transformers library, things are much easier"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f680dc0",
   "metadata": {},
   "source": [
    "# Installation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3d5f6690",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting package metadata (current_repodata.json): done\n",
      "Solving environment: done\n",
      "\n",
      "\n",
      "==> WARNING: A newer version of conda exists. <==\n",
      "  current version: 4.10.1\n",
      "  latest version: 4.10.3\n",
      "\n",
      "Please update conda by running\n",
      "\n",
      "    $ conda update -n base -c defaults conda\n",
      "\n",
      "\n",
      "\n",
      "# All requested packages already installed.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "!conda install pytorch torchvision torchaudio cudatoolkit=11.1 -c pytorch -c nvidia -y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4b0c79d8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: transformers in /home/blake/anaconda3/envs/gptj_notebook/lib/python3.7/site-packages (4.11.2)\n",
      "Requirement already satisfied: importlib-metadata in /home/blake/anaconda3/envs/gptj_notebook/lib/python3.7/site-packages (from transformers) (4.8.1)\n",
      "Requirement already satisfied: regex!=2019.12.17 in /home/blake/anaconda3/envs/gptj_notebook/lib/python3.7/site-packages (from transformers) (2021.9.30)\n",
      "Requirement already satisfied: huggingface-hub>=0.0.17 in /home/blake/anaconda3/envs/gptj_notebook/lib/python3.7/site-packages (from transformers) (0.0.17)\n",
      "Requirement already satisfied: sacremoses in /home/blake/anaconda3/envs/gptj_notebook/lib/python3.7/site-packages (from transformers) (0.0.46)\n",
      "Requirement already satisfied: filelock in /home/blake/anaconda3/envs/gptj_notebook/lib/python3.7/site-packages (from transformers) (3.2.0)\n",
      "Requirement already satisfied: tokenizers<0.11,>=0.10.1 in /home/blake/anaconda3/envs/gptj_notebook/lib/python3.7/site-packages (from transformers) (0.10.3)\n",
      "Requirement already satisfied: tqdm>=4.27 in /home/blake/anaconda3/envs/gptj_notebook/lib/python3.7/site-packages (from transformers) (4.62.3)\n",
      "Requirement already satisfied: requests in /home/blake/anaconda3/envs/gptj_notebook/lib/python3.7/site-packages (from transformers) (2.26.0)\n",
      "Requirement already satisfied: packaging>=20.0 in /home/blake/anaconda3/envs/gptj_notebook/lib/python3.7/site-packages (from transformers) (21.0)\n",
      "Requirement already satisfied: pyyaml>=5.1 in /home/blake/anaconda3/envs/gptj_notebook/lib/python3.7/site-packages (from transformers) (5.4.1)\n",
      "Requirement already satisfied: numpy>=1.17 in /home/blake/anaconda3/envs/gptj_notebook/lib/python3.7/site-packages (from transformers) (1.20.3)\n",
      "Requirement already satisfied: typing-extensions in /home/blake/anaconda3/envs/gptj_notebook/lib/python3.7/site-packages (from huggingface-hub>=0.0.17->transformers) (3.10.0.2)\n",
      "Requirement already satisfied: pyparsing>=2.0.2 in /home/blake/anaconda3/envs/gptj_notebook/lib/python3.7/site-packages (from packaging>=20.0->transformers) (2.4.7)\n",
      "Requirement already satisfied: zipp>=0.5 in /home/blake/anaconda3/envs/gptj_notebook/lib/python3.7/site-packages (from importlib-metadata->transformers) (3.5.0)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /home/blake/anaconda3/envs/gptj_notebook/lib/python3.7/site-packages (from requests->transformers) (3.2)\n",
      "Requirement already satisfied: charset-normalizer~=2.0.0 in /home/blake/anaconda3/envs/gptj_notebook/lib/python3.7/site-packages (from requests->transformers) (2.0.6)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /home/blake/anaconda3/envs/gptj_notebook/lib/python3.7/site-packages (from requests->transformers) (1.26.7)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /home/blake/anaconda3/envs/gptj_notebook/lib/python3.7/site-packages (from requests->transformers) (2021.5.30)\n",
      "Requirement already satisfied: six in /home/blake/anaconda3/envs/gptj_notebook/lib/python3.7/site-packages (from sacremoses->transformers) (1.16.0)\n",
      "Requirement already satisfied: joblib in /home/blake/anaconda3/envs/gptj_notebook/lib/python3.7/site-packages (from sacremoses->transformers) (1.0.1)\n",
      "Requirement already satisfied: click in /home/blake/anaconda3/envs/gptj_notebook/lib/python3.7/site-packages (from sacremoses->transformers) (8.0.1)\n"
     ]
    }
   ],
   "source": [
    "!pip install transformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "90f8bd0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import torch\n",
    "from transformers import GPTJForCausalLM, GPT2Tokenizer\n",
    "import torch\n",
    "import transformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a6475e09",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Will need at least 13-14GB of Vram for CUDA\n",
    "if torch.cuda.is_available():\n",
    "    model =  GPTJForCausalLM.from_pretrained(\"EleutherAI/gpt-j-6B\", torch_dtype=torch.float16).cuda()\n",
    "else:\n",
    "    model =  GPTJForCausalLM.from_pretrained(\"EleutherAI/gpt-j-6B\", torch_dtype=torch.float16)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f75f0a4",
   "metadata": {},
   "source": [
    "# Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0f731069",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"EleutherAI/gpt-j-6B\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "76da2f0c",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "GPTJForCausalLM(\n",
       "  (transformer): GPTJModel(\n",
       "    (wte): Embedding(50400, 4096)\n",
       "    (drop): Dropout(p=0.0, inplace=False)\n",
       "    (h): ModuleList(\n",
       "      (0): GPTJBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTJAttention(\n",
       "          (attn_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (mlp): GPTJMLP(\n",
       "          (fc_in): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (fc_out): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (1): GPTJBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTJAttention(\n",
       "          (attn_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (mlp): GPTJMLP(\n",
       "          (fc_in): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (fc_out): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (2): GPTJBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTJAttention(\n",
       "          (attn_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (mlp): GPTJMLP(\n",
       "          (fc_in): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (fc_out): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (3): GPTJBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTJAttention(\n",
       "          (attn_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (mlp): GPTJMLP(\n",
       "          (fc_in): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (fc_out): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (4): GPTJBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTJAttention(\n",
       "          (attn_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (mlp): GPTJMLP(\n",
       "          (fc_in): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (fc_out): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (5): GPTJBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTJAttention(\n",
       "          (attn_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (mlp): GPTJMLP(\n",
       "          (fc_in): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (fc_out): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (6): GPTJBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTJAttention(\n",
       "          (attn_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (mlp): GPTJMLP(\n",
       "          (fc_in): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (fc_out): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (7): GPTJBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTJAttention(\n",
       "          (attn_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (mlp): GPTJMLP(\n",
       "          (fc_in): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (fc_out): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (8): GPTJBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTJAttention(\n",
       "          (attn_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (mlp): GPTJMLP(\n",
       "          (fc_in): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (fc_out): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (9): GPTJBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTJAttention(\n",
       "          (attn_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (mlp): GPTJMLP(\n",
       "          (fc_in): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (fc_out): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (10): GPTJBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTJAttention(\n",
       "          (attn_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (mlp): GPTJMLP(\n",
       "          (fc_in): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (fc_out): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (11): GPTJBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTJAttention(\n",
       "          (attn_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (mlp): GPTJMLP(\n",
       "          (fc_in): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (fc_out): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (12): GPTJBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTJAttention(\n",
       "          (attn_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (mlp): GPTJMLP(\n",
       "          (fc_in): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (fc_out): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (13): GPTJBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTJAttention(\n",
       "          (attn_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (mlp): GPTJMLP(\n",
       "          (fc_in): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (fc_out): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (14): GPTJBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTJAttention(\n",
       "          (attn_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (mlp): GPTJMLP(\n",
       "          (fc_in): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (fc_out): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (15): GPTJBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTJAttention(\n",
       "          (attn_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (mlp): GPTJMLP(\n",
       "          (fc_in): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (fc_out): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (16): GPTJBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTJAttention(\n",
       "          (attn_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (mlp): GPTJMLP(\n",
       "          (fc_in): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (fc_out): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (17): GPTJBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTJAttention(\n",
       "          (attn_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (mlp): GPTJMLP(\n",
       "          (fc_in): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (fc_out): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (18): GPTJBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTJAttention(\n",
       "          (attn_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (mlp): GPTJMLP(\n",
       "          (fc_in): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (fc_out): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (19): GPTJBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTJAttention(\n",
       "          (attn_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (mlp): GPTJMLP(\n",
       "          (fc_in): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (fc_out): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (20): GPTJBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTJAttention(\n",
       "          (attn_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (mlp): GPTJMLP(\n",
       "          (fc_in): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (fc_out): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (21): GPTJBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTJAttention(\n",
       "          (attn_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (mlp): GPTJMLP(\n",
       "          (fc_in): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (fc_out): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (22): GPTJBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTJAttention(\n",
       "          (attn_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (mlp): GPTJMLP(\n",
       "          (fc_in): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (fc_out): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (23): GPTJBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTJAttention(\n",
       "          (attn_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (mlp): GPTJMLP(\n",
       "          (fc_in): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (fc_out): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (24): GPTJBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTJAttention(\n",
       "          (attn_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (mlp): GPTJMLP(\n",
       "          (fc_in): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (fc_out): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (25): GPTJBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTJAttention(\n",
       "          (attn_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (mlp): GPTJMLP(\n",
       "          (fc_in): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (fc_out): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (26): GPTJBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTJAttention(\n",
       "          (attn_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (mlp): GPTJMLP(\n",
       "          (fc_in): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (fc_out): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (27): GPTJBlock(\n",
       "        (ln_1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): GPTJAttention(\n",
       "          (attn_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (resid_dropout): Dropout(p=0.0, inplace=False)\n",
       "          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        )\n",
       "        (mlp): GPTJMLP(\n",
       "          (fc_in): Linear(in_features=4096, out_features=16384, bias=True)\n",
       "          (fc_out): Linear(in_features=16384, out_features=4096, bias=True)\n",
       "          (dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (ln_f): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n",
       "  )\n",
       "  (lm_head): Linear(in_features=4096, out_features=50400, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "030d4615",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    }
   ],
   "source": [
    "input_text = \"Hello my name is Blake and\"\n",
    "input_ids = tokenizer.encode(str(input_text), return_tensors='pt').cuda()\n",
    "\n",
    "output = model.generate(\n",
    "    input_ids,\n",
    "    do_sample=True,\n",
    "    max_length=20,\n",
    "    top_p=0.7,\n",
    "    top_k=0,\n",
    "    temperature=1.0,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f783672d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Hello my name is Blake and I\n",
      "have been doing computer science for\n",
      "the last 10 years and\n"
     ]
    }
   ],
   "source": [
    "print(tokenizer.decode(output[0], skip_special_tokens=True))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "gptj",
   "language": "python",
   "name": "gptj"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}