OpenAI Transformer testing.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/share/anaconda3/lib/python3.6/site-packages/h5py/__init__.py:36: FutureWarning: Conversion of the second argument of issubdtype from `float` to `np.floating` is deprecated. In future, it will be treated as `np.float64 == np.dtype(float).type`.\n",
      "  from ._conv import register_converters as _register_converters\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensorflow 1.8.0\n",
      "pytorch 0.4.0\n"
     ]
    }
   ],
   "source": [
    "import sys, os, re, json, random\n",
    "import time\n",
    "from importlib import reload\n",
    "\n",
    "from typing import List, Dict\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "import tensorflow as tf\n",
    "print(\"tensorflow\", tf.__version__)\n",
    "\n",
    "import torch\n",
    "print(\"pytorch\", torch.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Additional dependencies (install with `conda`):\n",
    "- `ftfy`\n",
    "- `spacy`\n",
    "\n",
    "Also download spaCy `en` model with:\n",
    "```\n",
    "python -m spacy download en\n",
    "```\n",
    "(see https://spacy.io/usage/models)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ftfy 5.4.1\n",
      "spacy 2.0.11\n"
     ]
    }
   ],
   "source": [
    "import ftfy\n",
    "print(\"ftfy\", ftfy.__version__)\n",
    "import spacy\n",
    "print(\"spacy\", spacy.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tokenize text and apply BPE\n",
    "\n",
    "The `TextEncoder` object handles tokenization and applying BPE to raw text, giving a list of IDs. See https://github.com/openai/finetune-transformer-lm/blob/master/text_utils.py and https://github.com/openai/finetune-transformer-lm/blob/master/utils.py#L14\n",
    "\n",
    "In order to align with original text tokenization, we probably want a two-stage process:\n",
    "1. Recover the processed spaCy tokens by a reverse lookup, and project annotations to spaCy tokenization.\n",
    "2. Project annotations from the spaCy tokenization to the BPE pieces."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<module 'src.openai_transformer_lm.utils' from '/nfs/jsalt/home/iftenney/jiant/src/openai_transformer_lm/utils.py'>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from src.openai_transformer_lm import utils as openai_utils\n",
    "reload(openai_utils)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[[718, 889, 2510, 636, 246, 8210, 7961, 7961]]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "text = \"how much wood would a woodchuck chuck\"\n",
    "e = openai_utils.encode([text])\n",
    "e"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    },
    {
     "data": {
      "text/plain": [
       "['how</w>',\n",
       " 'much</w>',\n",
       " 'wood</w>',\n",
       " 'would</w>',\n",
       " 'a</w>',\n",
       " 'wood',\n",
       " 'chuck</w>',\n",
       " 'chuck</w>']"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "openai_utils.tokenize(text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "openai_dir /nfs/jsalt/home/iftenney/jiant/src/openai_transformer_lm\n",
      "openai_data_dir /nfs/jsalt/home/iftenney/jiant/src/openai_transformer_lm/tf_original/model\n"
     ]
    }
   ],
   "source": [
    "openai_dir = os.path.dirname(openai_utils.__file__)\n",
    "print('openai_dir', openai_dir)\n",
    "openai_data_dir = openai_utils.OPENAI_DATA_DIR\n",
    "print('openai_data_dir', openai_data_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[['how</w>',\n",
       "  'much</w>',\n",
       "  'wood</w>',\n",
       "  'would</w>',\n",
       "  'a</w>',\n",
       "  'wood',\n",
       "  'chuck</w>',\n",
       "  'chuck</w>']]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "list(openai_utils.decode_partial(e))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['how much wood would a woodchuck chuck']"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "list(openai_utils.decode_full(e))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For their model, they process the ID list by adding start and end IDs, `encoder['_start_']` and `encoder['_delimiter_']`, then pad with `clf_token = encoder['_classify_']`. All of these are set equal to `n_vocab = len(encoder)`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[['a</w>', 'recent</w>', 'report</w>']]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "openai_utils.N_VOCAB\n",
    "list(openai_utils.decode_partial([[246, 6264, 4144]]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Project Annotations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "from src import utils\n",
    "from src import retokenize\n",
    "\n",
    "fname = \"/nfs/jsalt/share/glue_data/edges/spr2/train.edges.json\"\n",
    "records = list(utils.load_json_data(fname))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def space_tokenize_with_eow(sentence):\n",
    "    \"\"\"Add </w> markers to ensure word-boundary alignment.\"\"\"\n",
    "    return [t + \"</w>\" for t in sentence.split()]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TokenAligner(7, 8):\n",
      "  0 -> [0]\n",
      "  1 -> [1]\n",
      "  2 -> [2]\n",
      "  3 -> [3]\n",
      "  4 -> [4]\n",
      "  5 -> [5, 6]\n",
      "  6 -> [7]\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r"
     ]
    }
   ],
   "source": [
    "text = \"how much wood would a woodchuck chuck\"\n",
    "ta = retokenize.TokenAligner(space_tokenize_with_eow(text), openai_utils.tokenize(text))\n",
    "print(ta)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TokenAligner(17, 19):\n",
      "  0 -> [0]\n",
      "  1 -> [1]\n",
      "  2 -> [2]\n",
      "  3 -> [3, 4]\n",
      "  4 -> [5]\n",
      "  5 -> [6]\n",
      "  6 -> [7]\n",
      "  7 -> [8]\n",
      "  8 -> [9]\n",
      "  9 -> [10]\n",
      "  10 -> [11]\n",
      "  11 -> [12, 13]\n",
      "  12 -> [14]\n",
      "  13 -> [15]\n",
      "  14 -> [16]\n",
      "  15 -> [17]\n",
      "  16 -> [18]\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r"
     ]
    }
   ],
   "source": [
    "record = records[10]\n",
    "text = record['text']\n",
    "ta = retokenize.TokenAligner(space_tokenize_with_eow(text), openai_utils.tokenize(text))\n",
    "print(ta)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    },
    {
     "data": {
      "text/plain": [
       "['i</w>',\n",
       " 'have</w>',\n",
       " 'a</w>',\n",
       " 'pre',\n",
       " 'order</w>',\n",
       " 'and</w>',\n",
       " 'am</w>',\n",
       " 'even</w>',\n",
       " 'considering</w>',\n",
       " 'getting</w>',\n",
       " 'a</w>',\n",
       " 'second</w>',\n",
       " 'pre',\n",
       " 'order</w>',\n",
       " 'to</w>',\n",
       " 'have</w>',\n",
       " 'multiple</w>',\n",
       " 'accounts</w>',\n",
       " '.</w>']"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "openai_utils.tokenize(text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# OpenAI TensorFlow Model\n",
    "\n",
    "Adapted from https://github.com/openai/finetune-transformer-lm/blob/master/train.py#L163 to just export weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /nfs/jsalt/home/iftenney/jiant/src/openai_transformer_lm/transformer_tf_simplified.py:51: calling reduce_mean (from tensorflow.python.ops.math_ops) with keep_dims is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "keep_dims is deprecated, use keepdims instead\n",
      "Running initializer...\n",
      "Loading pre-trained params...\n",
      "Assigning pre-trained params...\n",
      "Done!\n"
     ]
    }
   ],
   "source": [
    "from src.openai_transformer_lm import transformer_tf_simplified\n",
    "reload(transformer_tf_simplified)\n",
    "from src.openai_transformer_lm.tf_original import utils as openai_tf_utils\n",
    "assert openai_utils.N_VOCAB == transformer_tf_simplified.n_vocab\n",
    "\n",
    "SHAPES_FILE = os.path.join(openai_data_dir, \"params_shapes.json\")\n",
    "PARAMS_FILE_TMPL = os.path.join(openai_data_dir, \"params_{}.npy\")\n",
    "\n",
    "seed = 42\n",
    "random.seed(seed)\n",
    "np.random.seed(seed)\n",
    "tf.set_random_seed(seed)\n",
    "\n",
    "with tf.Graph().as_default():\n",
    "    X_in = tf.placeholder(tf.int32, [None, transformer_tf_simplified.n_ctx, 2])\n",
    "    h = transformer_tf_simplified.model_abbreviated(tf.expand_dims(X_in, 1))\n",
    "    \n",
    "    params = openai_tf_utils.find_trainable_variables(\"model\")\n",
    "    print(\"Running initializer...\")\n",
    "    sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))\n",
    "    sess.run(tf.global_variables_initializer())\n",
    "    transformer_tf_simplified.load_params(sess, params, SHAPES_FILE, PARAMS_FILE_TMPL)\n",
    "    \n",
    "    h_val = sess.run(h, {X_in:openai_utils.prep_ids(e)})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1, 512, 768)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "h_val.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[ 0.16202518, -0.00763676,  0.14871871, ...,  0.17555687,\n",
       "         -0.0820915 ,  0.15243863],\n",
       "        [-0.22960284, -0.7131108 ,  0.25970727, ...,  0.29409307,\n",
       "         -0.28634194, -0.3005669 ],\n",
       "        [-0.33067906, -0.11411758,  0.37495106, ..., -0.23809306,\n",
       "         -1.2012712 , -0.27857336],\n",
       "        ...,\n",
       "        [ 0.50163245,  0.84288955,  0.8027246 , ..., -0.03076262,\n",
       "          1.1621311 ,  0.7543984 ],\n",
       "        [ 0.51369536,  0.7608172 ,  0.8307392 , ..., -0.05948838,\n",
       "          1.0978718 ,  0.6896318 ],\n",
       "        [ 0.47562498,  0.79977787,  0.8108253 , ..., -0.0525886 ,\n",
       "          1.1286504 ,  0.7206507 ]]], dtype=float32)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "h_val"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Huggingface PyTorch port\n",
    "\n",
    "Code from https://github.com/huggingface/pytorch-openai-transformer-lm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading weights...\n"
     ]
    }
   ],
   "source": [
    "from src.openai_transformer_lm.pytorch_huggingface import model_pytorch\n",
    "reload(model_pytorch)\n",
    "\n",
    "seed = 42\n",
    "random.seed(seed)\n",
    "np.random.seed(seed)\n",
    "torch.manual_seed(seed)\n",
    "\n",
    "args = model_pytorch.DEFAULT_CONFIG\n",
    "n_special = transformer_tf_simplified.n_special\n",
    "model = model_pytorch.TransformerModel(args, vocab=40990+n_special)\n",
    "loader_args = dict(n_special=n_special)\n",
    "loader_args['path'] = openai_data_dir + \"/\"\n",
    "loader_args['path_names'] = os.path.dirname(model_pytorch.__file__) + \"/\"\n",
    "model_pytorch.load_openai_pretrained_model(model, **loader_args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'n_embd': 768,\n",
       " 'n_head': 12,\n",
       " 'n_layer': 12,\n",
       " 'embd_pdrop': 0.1,\n",
       " 'attn_pdrop': 0.1,\n",
       " 'resid_pdrop': 0.1,\n",
       " 'afn': 'gelu',\n",
       " 'clf_pdrop': 0.1}"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "args"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([4, 16, 2])"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "t = torch.arange(16).repeat(4,1)\n",
    "t.size()[0]\n",
    "torch.stack([t,t], dim=2).size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.eval()\n",
    "ids = torch.LongTensor(openai_utils.prep_ids(e))\n",
    "h_val_pytorch = model(ids).detach().numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1, 512, 768)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "h_val_pytorch.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[ 0.1620266 , -0.00763825,  0.14871885, ...,  0.17555933,\n",
       "         -0.08209433,  0.15244102],\n",
       "        [-0.22960448, -0.71311027,  0.25970703, ...,  0.2940926 ,\n",
       "         -0.28634158, -0.3005667 ],\n",
       "        [-0.33067977, -0.11411723,  0.37495226, ..., -0.23809452,\n",
       "         -1.2012721 , -0.2785728 ],\n",
       "        ...,\n",
       "        [ 0.50163686,  0.8428891 ,  0.80272746, ..., -0.03076063,\n",
       "          1.1621205 ,  0.7543992 ],\n",
       "        [ 0.51369643,  0.76081246,  0.8307426 , ..., -0.05948911,\n",
       "          1.0978615 ,  0.6896316 ],\n",
       "        [ 0.47562388,  0.7997696 ,  0.81082684, ..., -0.05258931,\n",
       "          1.1286354 ,  0.7206483 ]]], dtype=float32)"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "h_val_pytorch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3.0493553e-05"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.sqrt(np.mean((h_val - h_val_pytorch)**2.0))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Woohoo! Looks like the PyTorch implementation loads the weights correctly and matches the original implementation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TensorFlow Checkpoints\n",
    "\n",
    "Inspect these and see if we can load the weights into the PyTorch model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('model/h0/attn/c_attn/b', [2304]),\n",
       " ('model/h0/attn/c_attn/w', [1, 768, 2304]),\n",
       " ('model/h0/attn/c_proj/b', [768]),\n",
       " ('model/h0/attn/c_proj/w', [1, 768, 768]),\n",
       " ('model/h0/ln_1/b', [768]),\n",
       " ('model/h0/ln_1/g', [768]),\n",
       " ('model/h0/ln_2/b', [768]),\n",
       " ('model/h0/ln_2/g', [768]),\n",
       " ('model/h0/mlp/c_fc/b', [3072]),\n",
       " ('model/h0/mlp/c_fc/w', [1, 768, 3072]),\n",
       " ('model/h0/mlp/c_proj/b', [768]),\n",
       " ('model/h0/mlp/c_proj/w', [1, 3072, 768]),\n",
       " ('model/h1/attn/c_attn/b', [2304]),\n",
       " ('model/h1/attn/c_attn/w', [1, 768, 2304]),\n",
       " ('model/h1/attn/c_proj/b', [768]),\n",
       " ('model/h1/attn/c_proj/w', [1, 768, 768]),\n",
       " ('model/h1/ln_1/b', [768]),\n",
       " ('model/h1/ln_1/g', [768]),\n",
       " ('model/h1/ln_2/b', [768]),\n",
       " ('model/h1/ln_2/g', [768]),\n",
       " ('model/h1/mlp/c_fc/b', [3072]),\n",
       " ('model/h1/mlp/c_fc/w', [1, 768, 3072]),\n",
       " ('model/h1/mlp/c_proj/b', [768]),\n",
       " ('model/h1/mlp/c_proj/w', [1, 3072, 768]),\n",
       " ('model/h10/attn/c_attn/b', [2304]),\n",
       " ('model/h10/attn/c_attn/w', [1, 768, 2304]),\n",
       " ('model/h10/attn/c_proj/b', [768]),\n",
       " ('model/h10/attn/c_proj/w', [1, 768, 768]),\n",
       " ('model/h10/ln_1/b', [768]),\n",
       " ('model/h10/ln_1/g', [768]),\n",
       " ('model/h10/ln_2/b', [768]),\n",
       " ('model/h10/ln_2/g', [768]),\n",
       " ('model/h10/mlp/c_fc/b', [3072]),\n",
       " ('model/h10/mlp/c_fc/w', [1, 768, 3072]),\n",
       " ('model/h10/mlp/c_proj/b', [768]),\n",
       " ('model/h10/mlp/c_proj/w', [1, 3072, 768]),\n",
       " ('model/h11/attn/c_attn/b', [2304]),\n",
       " ('model/h11/attn/c_attn/w', [1, 768, 2304]),\n",
       " ('model/h11/attn/c_proj/b', [768]),\n",
       " ('model/h11/attn/c_proj/w', [1, 768, 768]),\n",
       " ('model/h11/ln_1/b', [768]),\n",
       " ('model/h11/ln_1/g', [768]),\n",
       " ('model/h11/ln_2/b', [768]),\n",
       " ('model/h11/ln_2/g', [768]),\n",
       " ('model/h11/mlp/c_fc/b', [3072]),\n",
       " ('model/h11/mlp/c_fc/w', [1, 768, 3072]),\n",
       " ('model/h11/mlp/c_proj/b', [768]),\n",
       " ('model/h11/mlp/c_proj/w', [1, 3072, 768]),\n",
       " ('model/h2/attn/c_attn/b', [2304]),\n",
       " ('model/h2/attn/c_attn/w', [1, 768, 2304]),\n",
       " ('model/h2/attn/c_proj/b', [768]),\n",
       " ('model/h2/attn/c_proj/w', [1, 768, 768]),\n",
       " ('model/h2/ln_1/b', [768]),\n",
       " ('model/h2/ln_1/g', [768]),\n",
       " ('model/h2/ln_2/b', [768]),\n",
       " ('model/h2/ln_2/g', [768]),\n",
       " ('model/h2/mlp/c_fc/b', [3072]),\n",
       " ('model/h2/mlp/c_fc/w', [1, 768, 3072]),\n",
       " ('model/h2/mlp/c_proj/b', [768]),\n",
       " ('model/h2/mlp/c_proj/w', [1, 3072, 768]),\n",
       " ('model/h3/attn/c_attn/b', [2304]),\n",
       " ('model/h3/attn/c_attn/w', [1, 768, 2304]),\n",
       " ('model/h3/attn/c_proj/b', [768]),\n",
       " ('model/h3/attn/c_proj/w', [1, 768, 768]),\n",
       " ('model/h3/ln_1/b', [768]),\n",
       " ('model/h3/ln_1/g', [768]),\n",
       " ('model/h3/ln_2/b', [768]),\n",
       " ('model/h3/ln_2/g', [768]),\n",
       " ('model/h3/mlp/c_fc/b', [3072]),\n",
       " ('model/h3/mlp/c_fc/w', [1, 768, 3072]),\n",
       " ('model/h3/mlp/c_proj/b', [768]),\n",
       " ('model/h3/mlp/c_proj/w', [1, 3072, 768]),\n",
       " ('model/h4/attn/c_attn/b', [2304]),\n",
       " ('model/h4/attn/c_attn/w', [1, 768, 2304]),\n",
       " ('model/h4/attn/c_proj/b', [768]),\n",
       " ('model/h4/attn/c_proj/w', [1, 768, 768]),\n",
       " ('model/h4/ln_1/b', [768]),\n",
       " ('model/h4/ln_1/g', [768]),\n",
       " ('model/h4/ln_2/b', [768]),\n",
       " ('model/h4/ln_2/g', [768]),\n",
       " ('model/h4/mlp/c_fc/b', [3072]),\n",
       " ('model/h4/mlp/c_fc/w', [1, 768, 3072]),\n",
       " ('model/h4/mlp/c_proj/b', [768]),\n",
       " ('model/h4/mlp/c_proj/w', [1, 3072, 768]),\n",
       " ('model/h5/attn/c_attn/b', [2304]),\n",
       " ('model/h5/attn/c_attn/w', [1, 768, 2304]),\n",
       " ('model/h5/attn/c_proj/b', [768]),\n",
       " ('model/h5/attn/c_proj/w', [1, 768, 768]),\n",
       " ('model/h5/ln_1/b', [768]),\n",
       " ('model/h5/ln_1/g', [768]),\n",
       " ('model/h5/ln_2/b', [768]),\n",
       " ('model/h5/ln_2/g', [768]),\n",
       " ('model/h5/mlp/c_fc/b', [3072]),\n",
       " ('model/h5/mlp/c_fc/w', [1, 768, 3072]),\n",
       " ('model/h5/mlp/c_proj/b', [768]),\n",
       " ('model/h5/mlp/c_proj/w', [1, 3072, 768]),\n",
       " ('model/h6/attn/c_attn/b', [2304]),\n",
       " ('model/h6/attn/c_attn/w', [1, 768, 2304]),\n",
       " ('model/h6/attn/c_proj/b', [768]),\n",
       " ('model/h6/attn/c_proj/w', [1, 768, 768]),\n",
       " ('model/h6/ln_1/b', [768]),\n",
       " ('model/h6/ln_1/g', [768]),\n",
       " ('model/h6/ln_2/b', [768]),\n",
       " ('model/h6/ln_2/g', [768]),\n",
       " ('model/h6/mlp/c_fc/b', [3072]),\n",
       " ('model/h6/mlp/c_fc/w', [1, 768, 3072]),\n",
       " ('model/h6/mlp/c_proj/b', [768]),\n",
       " ('model/h6/mlp/c_proj/w', [1, 3072, 768]),\n",
       " ('model/h7/attn/c_attn/b', [2304]),\n",
       " ('model/h7/attn/c_attn/w', [1, 768, 2304]),\n",
       " ('model/h7/attn/c_proj/b', [768]),\n",
       " ('model/h7/attn/c_proj/w', [1, 768, 768]),\n",
       " ('model/h7/ln_1/b', [768]),\n",
       " ('model/h7/ln_1/g', [768]),\n",
       " ('model/h7/ln_2/b', [768]),\n",
       " ('model/h7/ln_2/g', [768]),\n",
       " ('model/h7/mlp/c_fc/b', [3072]),\n",
       " ('model/h7/mlp/c_fc/w', [1, 768, 3072]),\n",
       " ('model/h7/mlp/c_proj/b', [768]),\n",
       " ('model/h7/mlp/c_proj/w', [1, 3072, 768]),\n",
       " ('model/h8/attn/c_attn/b', [2304]),\n",
       " ('model/h8/attn/c_attn/w', [1, 768, 2304]),\n",
       " ('model/h8/attn/c_proj/b', [768]),\n",
       " ('model/h8/attn/c_proj/w', [1, 768, 768]),\n",
       " ('model/h8/ln_1/b', [768]),\n",
       " ('model/h8/ln_1/g', [768]),\n",
       " ('model/h8/ln_2/b', [768]),\n",
       " ('model/h8/ln_2/g', [768]),\n",
       " ('model/h8/mlp/c_fc/b', [3072]),\n",
       " ('model/h8/mlp/c_fc/w', [1, 768, 3072]),\n",
       " ('model/h8/mlp/c_proj/b', [768]),\n",
       " ('model/h8/mlp/c_proj/w', [1, 3072, 768]),\n",
       " ('model/h9/attn/c_attn/b', [2304]),\n",
       " ('model/h9/attn/c_attn/w', [1, 768, 2304]),\n",
       " ('model/h9/attn/c_proj/b', [768]),\n",
       " ('model/h9/attn/c_proj/w', [1, 768, 768]),\n",
       " ('model/h9/ln_1/b', [768]),\n",
       " ('model/h9/ln_1/g', [768]),\n",
       " ('model/h9/ln_2/b', [768]),\n",
       " ('model/h9/ln_2/g', [768]),\n",
       " ('model/h9/mlp/c_fc/b', [3072]),\n",
       " ('model/h9/mlp/c_fc/w', [1, 768, 3072]),\n",
       " ('model/h9/mlp/c_proj/b', [768]),\n",
       " ('model/h9/mlp/c_proj/w', [1, 3072, 768]),\n",
       " ('model/pe', [512, 768]),\n",
       " ('model/we', [40481, 768])]"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# ckpt_path = \"/nfs/jsalt/home/iftenney/checkpoints/bwb_shuffled/model.ckpt-1000000\"\n",
    "ckpt_path = \"/nfs/jsalt/home/iftenney/checkpoints/orig_openai_checkpoint/init_model.ckpt\"\n",
    "tf.train.list_variables(ckpt_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "40993"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "40481 + 512"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "embed.weight [40993, 768]\n",
      "h.0.attn.c_attn.w [768, 2304]\n",
      "h.0.attn.c_attn.b [2304]\n",
      "h.0.attn.c_proj.w [768, 768]\n",
      "h.0.attn.c_proj.b [768]\n",
      "h.0.ln_1.g [768]\n",
      "h.0.ln_1.b [768]\n",
      "h.0.mlp.c_fc.w [768, 3072]\n",
      "h.0.mlp.c_fc.b [3072]\n",
      "h.0.mlp.c_proj.w [3072, 768]\n",
      "h.0.mlp.c_proj.b [768]\n",
      "h.0.ln_2.g [768]\n",
      "h.0.ln_2.b [768]\n",
      "h.1.attn.c_attn.w [768, 2304]\n",
      "h.1.attn.c_attn.b [2304]\n",
      "h.1.attn.c_proj.w [768, 768]\n",
      "h.1.attn.c_proj.b [768]\n",
      "h.1.ln_1.g [768]\n",
      "h.1.ln_1.b [768]\n",
      "h.1.mlp.c_fc.w [768, 3072]\n",
      "h.1.mlp.c_fc.b [3072]\n",
      "h.1.mlp.c_proj.w [3072, 768]\n",
      "h.1.mlp.c_proj.b [768]\n",
      "h.1.ln_2.g [768]\n",
      "h.1.ln_2.b [768]\n",
      "h.2.attn.c_attn.w [768, 2304]\n",
      "h.2.attn.c_attn.b [2304]\n",
      "h.2.attn.c_proj.w [768, 768]\n",
      "h.2.attn.c_proj.b [768]\n",
      "h.2.ln_1.g [768]\n",
      "h.2.ln_1.b [768]\n",
      "h.2.mlp.c_fc.w [768, 3072]\n",
      "h.2.mlp.c_fc.b [3072]\n",
      "h.2.mlp.c_proj.w [3072, 768]\n",
      "h.2.mlp.c_proj.b [768]\n",
      "h.2.ln_2.g [768]\n",
      "h.2.ln_2.b [768]\n",
      "h.3.attn.c_attn.w [768, 2304]\n",
      "h.3.attn.c_attn.b [2304]\n",
      "h.3.attn.c_proj.w [768, 768]\n",
      "h.3.attn.c_proj.b [768]\n",
      "h.3.ln_1.g [768]\n",
      "h.3.ln_1.b [768]\n",
      "h.3.mlp.c_fc.w [768, 3072]\n",
      "h.3.mlp.c_fc.b [3072]\n",
      "h.3.mlp.c_proj.w [3072, 768]\n",
      "h.3.mlp.c_proj.b [768]\n",
      "h.3.ln_2.g [768]\n",
      "h.3.ln_2.b [768]\n",
      "h.4.attn.c_attn.w [768, 2304]\n",
      "h.4.attn.c_attn.b [2304]\n",
      "h.4.attn.c_proj.w [768, 768]\n",
      "h.4.attn.c_proj.b [768]\n",
      "h.4.ln_1.g [768]\n",
      "h.4.ln_1.b [768]\n",
      "h.4.mlp.c_fc.w [768, 3072]\n",
      "h.4.mlp.c_fc.b [3072]\n",
      "h.4.mlp.c_proj.w [3072, 768]\n",
      "h.4.mlp.c_proj.b [768]\n",
      "h.4.ln_2.g [768]\n",
      "h.4.ln_2.b [768]\n",
      "h.5.attn.c_attn.w [768, 2304]\n",
      "h.5.attn.c_attn.b [2304]\n",
      "h.5.attn.c_proj.w [768, 768]\n",
      "h.5.attn.c_proj.b [768]\n",
      "h.5.ln_1.g [768]\n",
      "h.5.ln_1.b [768]\n",
      "h.5.mlp.c_fc.w [768, 3072]\n",
      "h.5.mlp.c_fc.b [3072]\n",
      "h.5.mlp.c_proj.w [3072, 768]\n",
      "h.5.mlp.c_proj.b [768]\n",
      "h.5.ln_2.g [768]\n",
      "h.5.ln_2.b [768]\n",
      "h.6.attn.c_attn.w [768, 2304]\n",
      "h.6.attn.c_attn.b [2304]\n",
      "h.6.attn.c_proj.w [768, 768]\n",
      "h.6.attn.c_proj.b [768]\n",
      "h.6.ln_1.g [768]\n",
      "h.6.ln_1.b [768]\n",
      "h.6.mlp.c_fc.w [768, 3072]\n",
      "h.6.mlp.c_fc.b [3072]\n",
      "h.6.mlp.c_proj.w [3072, 768]\n",
      "h.6.mlp.c_proj.b [768]\n",
      "h.6.ln_2.g [768]\n",
      "h.6.ln_2.b [768]\n",
      "h.7.attn.c_attn.w [768, 2304]\n",
      "h.7.attn.c_attn.b [2304]\n",
      "h.7.attn.c_proj.w [768, 768]\n",
      "h.7.attn.c_proj.b [768]\n",
      "h.7.ln_1.g [768]\n",
      "h.7.ln_1.b [768]\n",
      "h.7.mlp.c_fc.w [768, 3072]\n",
      "h.7.mlp.c_fc.b [3072]\n",
      "h.7.mlp.c_proj.w [3072, 768]\n",
      "h.7.mlp.c_proj.b [768]\n",
      "h.7.ln_2.g [768]\n",
      "h.7.ln_2.b [768]\n",
      "h.8.attn.c_attn.w [768, 2304]\n",
      "h.8.attn.c_attn.b [2304]\n",
      "h.8.attn.c_proj.w [768, 768]\n",
      "h.8.attn.c_proj.b [768]\n",
      "h.8.ln_1.g [768]\n",
      "h.8.ln_1.b [768]\n",
      "h.8.mlp.c_fc.w [768, 3072]\n",
      "h.8.mlp.c_fc.b [3072]\n",
      "h.8.mlp.c_proj.w [3072, 768]\n",
      "h.8.mlp.c_proj.b [768]\n",
      "h.8.ln_2.g [768]\n",
      "h.8.ln_2.b [768]\n",
      "h.9.attn.c_attn.w [768, 2304]\n",
      "h.9.attn.c_attn.b [2304]\n",
      "h.9.attn.c_proj.w [768, 768]\n",
      "h.9.attn.c_proj.b [768]\n",
      "h.9.ln_1.g [768]\n",
      "h.9.ln_1.b [768]\n",
      "h.9.mlp.c_fc.w [768, 3072]\n",
      "h.9.mlp.c_fc.b [3072]\n",
      "h.9.mlp.c_proj.w [3072, 768]\n",
      "h.9.mlp.c_proj.b [768]\n",
      "h.9.ln_2.g [768]\n",
      "h.9.ln_2.b [768]\n",
      "h.10.attn.c_attn.w [768, 2304]\n",
      "h.10.attn.c_attn.b [2304]\n",
      "h.10.attn.c_proj.w [768, 768]\n",
      "h.10.attn.c_proj.b [768]\n",
      "h.10.ln_1.g [768]\n",
      "h.10.ln_1.b [768]\n",
      "h.10.mlp.c_fc.w [768, 3072]\n",
      "h.10.mlp.c_fc.b [3072]\n",
      "h.10.mlp.c_proj.w [3072, 768]\n",
      "h.10.mlp.c_proj.b [768]\n",
      "h.10.ln_2.g [768]\n",
      "h.10.ln_2.b [768]\n",
      "h.11.attn.c_attn.w [768, 2304]\n",
      "h.11.attn.c_attn.b [2304]\n",
      "h.11.attn.c_proj.w [768, 768]\n",
      "h.11.attn.c_proj.b [768]\n",
      "h.11.ln_1.g [768]\n",
      "h.11.ln_1.b [768]\n",
      "h.11.mlp.c_fc.w [768, 3072]\n",
      "h.11.mlp.c_fc.b [3072]\n",
      "h.11.mlp.c_proj.w [3072, 768]\n",
      "h.11.mlp.c_proj.b [768]\n",
      "h.11.ln_2.g [768]\n",
      "h.11.ln_2.b [768]\n"
     ]
    }
   ],
   "source": [
    "for name, p in model.named_parameters():\n",
    "    print(name, list(p.shape))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Verify that converted checkpoint weights match the directly-loaded weights in the PyTorch model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.all(p.shape == (768,))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "embed.weight -> ['model/we', 'model/pe']\n",
      "torch.Size([40993, 768]) (40993, 768)\n",
      "h.0.attn.c_attn.w -> model/h0/attn/c_attn/w\n",
      "torch.Size([768, 2304]) (1, 768, 2304)\n",
      "h.0.attn.c_attn.b -> model/h0/attn/c_attn/b\n",
      "torch.Size([2304]) (2304,)\n",
      "h.0.attn.c_proj.w -> model/h0/attn/c_proj/w\n",
      "torch.Size([768, 768]) (1, 768, 768)\n",
      "h.0.attn.c_proj.b -> model/h0/attn/c_proj/b\n",
      "torch.Size([768]) (768,)\n",
      "h.0.ln_1.g -> model/h0/ln_1/g\n",
      "torch.Size([768]) (768,)\n",
      "h.0.ln_1.b -> model/h0/ln_1/b\n",
      "torch.Size([768]) (768,)\n",
      "h.0.mlp.c_fc.w -> model/h0/mlp/c_fc/w\n",
      "torch.Size([768, 3072]) (1, 768, 3072)\n",
      "h.0.mlp.c_fc.b -> model/h0/mlp/c_fc/b\n",
      "torch.Size([3072]) (3072,)\n",
      "h.0.mlp.c_proj.w -> model/h0/mlp/c_proj/w\n",
      "torch.Size([3072, 768]) (1, 3072, 768)\n",
      "h.0.mlp.c_proj.b -> model/h0/mlp/c_proj/b\n",
      "torch.Size([768]) (768,)\n",
      "h.0.ln_2.g -> model/h0/ln_2/g\n",
      "torch.Size([768]) (768,)\n",
      "h.0.ln_2.b -> model/h0/ln_2/b\n",
      "torch.Size([768]) (768,)\n",
      "h.1.attn.c_attn.w -> model/h1/attn/c_attn/w\n",
      "torch.Size([768, 2304]) (1, 768, 2304)\n",
      "h.1.attn.c_attn.b -> model/h1/attn/c_attn/b\n",
      "torch.Size([2304]) (2304,)\n",
      "h.1.attn.c_proj.w -> model/h1/attn/c_proj/w\n",
      "torch.Size([768, 768]) (1, 768, 768)\n",
      "h.1.attn.c_proj.b -> model/h1/attn/c_proj/b\n",
      "torch.Size([768]) (768,)\n",
      "h.1.ln_1.g -> model/h1/ln_1/g\n",
      "torch.Size([768]) (768,)\n",
      "h.1.ln_1.b -> model/h1/ln_1/b\n",
      "torch.Size([768]) (768,)\n",
      "h.1.mlp.c_fc.w -> model/h1/mlp/c_fc/w\n",
      "torch.Size([768, 3072]) (1, 768, 3072)\n",
      "h.1.mlp.c_fc.b -> model/h1/mlp/c_fc/b\n",
      "torch.Size([3072]) (3072,)\n",
      "h.1.mlp.c_proj.w -> model/h1/mlp/c_proj/w\n",
      "torch.Size([3072, 768]) (1, 3072, 768)\n",
      "h.1.mlp.c_proj.b -> model/h1/mlp/c_proj/b\n",
      "torch.Size([768]) (768,)\n",
      "h.1.ln_2.g -> model/h1/ln_2/g\n",
      "torch.Size([768]) (768,)\n",
      "h.1.ln_2.b -> model/h1/ln_2/b\n",
      "torch.Size([768]) (768,)\n",
      "h.2.attn.c_attn.w -> model/h2/attn/c_attn/w\n",
      "torch.Size([768, 2304]) (1, 768, 2304)\n",
      "h.2.attn.c_attn.b -> model/h2/attn/c_attn/b\n",
      "torch.Size([2304]) (2304,)\n",
      "h.2.attn.c_proj.w -> model/h2/attn/c_proj/w\n",
      "torch.Size([768, 768]) (1, 768, 768)\n",
      "h.2.attn.c_proj.b -> model/h2/attn/c_proj/b\n",
      "torch.Size([768]) (768,)\n",
      "h.2.ln_1.g -> model/h2/ln_1/g\n",
      "torch.Size([768]) (768,)\n",
      "h.2.ln_1.b -> model/h2/ln_1/b\n",
      "torch.Size([768]) (768,)\n",
      "h.2.mlp.c_fc.w -> model/h2/mlp/c_fc/w\n",
      "torch.Size([768, 3072]) (1, 768, 3072)\n",
      "h.2.mlp.c_fc.b -> model/h2/mlp/c_fc/b\n",
      "torch.Size([3072]) (3072,)\n",
      "h.2.mlp.c_proj.w -> model/h2/mlp/c_proj/w\n",
      "torch.Size([3072, 768]) (1, 3072, 768)\n",
      "h.2.mlp.c_proj.b -> model/h2/mlp/c_proj/b\n",
      "torch.Size([768]) (768,)\n",
      "h.2.ln_2.g -> model/h2/ln_2/g\n",
      "torch.Size([768]) (768,)\n",
      "h.2.ln_2.b -> model/h2/ln_2/b\n",
      "torch.Size([768]) (768,)\n",
      "h.3.attn.c_attn.w -> model/h3/attn/c_attn/w\n",
      "torch.Size([768, 2304]) (1, 768, 2304)\n",
      "h.3.attn.c_attn.b -> model/h3/attn/c_attn/b\n",
      "torch.Size([2304]) (2304,)\n",
      "h.3.attn.c_proj.w -> model/h3/attn/c_proj/w\n",
      "torch.Size([768, 768]) (1, 768, 768)\n",
      "h.3.attn.c_proj.b -> model/h3/attn/c_proj/b\n",
      "torch.Size([768]) (768,)\n",
      "h.3.ln_1.g -> model/h3/ln_1/g\n",
      "torch.Size([768]) (768,)\n",
      "h.3.ln_1.b -> model/h3/ln_1/b\n",
      "torch.Size([768]) (768,)\n",
      "h.3.mlp.c_fc.w -> model/h3/mlp/c_fc/w\n",
      "torch.Size([768, 3072]) (1, 768, 3072)\n",
      "h.3.mlp.c_fc.b -> model/h3/mlp/c_fc/b\n",
      "torch.Size([3072]) (3072,)\n",
      "h.3.mlp.c_proj.w -> model/h3/mlp/c_proj/w\n",
      "torch.Size([3072, 768]) (1, 3072, 768)\n",
      "h.3.mlp.c_proj.b -> model/h3/mlp/c_proj/b\n",
      "torch.Size([768]) (768,)\n",
      "h.3.ln_2.g -> model/h3/ln_2/g\n",
      "torch.Size([768]) (768,)\n",
      "h.3.ln_2.b -> model/h3/ln_2/b\n",
      "torch.Size([768]) (768,)\n",
      "h.4.attn.c_attn.w -> model/h4/attn/c_attn/w\n",
      "torch.Size([768, 2304]) (1, 768, 2304)\n",
      "h.4.attn.c_attn.b -> model/h4/attn/c_attn/b\n",
      "torch.Size([2304]) (2304,)\n",
      "h.4.attn.c_proj.w -> model/h4/attn/c_proj/w\n",
      "torch.Size([768, 768]) (1, 768, 768)\n",
      "h.4.attn.c_proj.b -> model/h4/attn/c_proj/b\n",
      "torch.Size([768]) (768,)\n",
      "h.4.ln_1.g -> model/h4/ln_1/g\n",
      "torch.Size([768]) (768,)\n",
      "h.4.ln_1.b -> model/h4/ln_1/b\n",
      "torch.Size([768]) (768,)\n",
      "h.4.mlp.c_fc.w -> model/h4/mlp/c_fc/w\n",
      "torch.Size([768, 3072]) (1, 768, 3072)\n",
      "h.4.mlp.c_fc.b -> model/h4/mlp/c_fc/b\n",
      "torch.Size([3072]) (3072,)\n",
      "h.4.mlp.c_proj.w -> model/h4/mlp/c_proj/w\n",
      "torch.Size([3072, 768]) (1, 3072, 768)\n",
      "h.4.mlp.c_proj.b -> model/h4/mlp/c_proj/b\n",
      "torch.Size([768]) (768,)\n",
      "h.4.ln_2.g -> model/h4/ln_2/g\n",
      "torch.Size([768]) (768,)\n",
      "h.4.ln_2.b -> model/h4/ln_2/b\n",
      "torch.Size([768]) (768,)\n",
      "h.5.attn.c_attn.w -> model/h5/attn/c_attn/w\n",
      "torch.Size([768, 2304]) (1, 768, 2304)\n",
      "h.5.attn.c_attn.b -> model/h5/attn/c_attn/b\n",
      "torch.Size([2304]) (2304,)\n",
      "h.5.attn.c_proj.w -> model/h5/attn/c_proj/w\n",
      "torch.Size([768, 768]) (1, 768, 768)\n",
      "h.5.attn.c_proj.b -> model/h5/attn/c_proj/b\n",
      "torch.Size([768]) (768,)\n",
      "h.5.ln_1.g -> model/h5/ln_1/g\n",
      "torch.Size([768]) (768,)\n",
      "h.5.ln_1.b -> model/h5/ln_1/b\n",
      "torch.Size([768]) (768,)\n",
      "h.5.mlp.c_fc.w -> model/h5/mlp/c_fc/w\n",
      "torch.Size([768, 3072]) (1, 768, 3072)\n",
      "h.5.mlp.c_fc.b -> model/h5/mlp/c_fc/b\n",
      "torch.Size([3072]) (3072,)\n",
      "h.5.mlp.c_proj.w -> model/h5/mlp/c_proj/w\n",
      "torch.Size([3072, 768]) (1, 3072, 768)\n",
      "h.5.mlp.c_proj.b -> model/h5/mlp/c_proj/b\n",
      "torch.Size([768]) (768,)\n",
      "h.5.ln_2.g -> model/h5/ln_2/g\n",
      "torch.Size([768]) (768,)\n",
      "h.5.ln_2.b -> model/h5/ln_2/b\n",
      "torch.Size([768]) (768,)\n",
      "h.6.attn.c_attn.w -> model/h6/attn/c_attn/w\n",
      "torch.Size([768, 2304]) (1, 768, 2304)\n",
      "h.6.attn.c_attn.b -> model/h6/attn/c_attn/b\n",
      "torch.Size([2304]) (2304,)\n",
      "h.6.attn.c_proj.w -> model/h6/attn/c_proj/w\n",
      "torch.Size([768, 768]) (1, 768, 768)\n",
      "h.6.attn.c_proj.b -> model/h6/attn/c_proj/b\n",
      "torch.Size([768]) (768,)\n",
      "h.6.ln_1.g -> model/h6/ln_1/g\n",
      "torch.Size([768]) (768,)\n",
      "h.6.ln_1.b -> model/h6/ln_1/b\n",
      "torch.Size([768]) (768,)\n",
      "h.6.mlp.c_fc.w -> model/h6/mlp/c_fc/w\n",
      "torch.Size([768, 3072]) (1, 768, 3072)\n",
      "h.6.mlp.c_fc.b -> model/h6/mlp/c_fc/b\n",
      "torch.Size([3072]) (3072,)\n",
      "h.6.mlp.c_proj.w -> model/h6/mlp/c_proj/w\n",
      "torch.Size([3072, 768]) (1, 3072, 768)\n",
      "h.6.mlp.c_proj.b -> model/h6/mlp/c_proj/b\n",
      "torch.Size([768]) (768,)\n",
      "h.6.ln_2.g -> model/h6/ln_2/g\n",
      "torch.Size([768]) (768,)\n",
      "h.6.ln_2.b -> model/h6/ln_2/b\n",
      "torch.Size([768]) (768,)\n",
      "h.7.attn.c_attn.w -> model/h7/attn/c_attn/w\n",
      "torch.Size([768, 2304]) (1, 768, 2304)\n",
      "h.7.attn.c_attn.b -> model/h7/attn/c_attn/b\n",
      "torch.Size([2304]) (2304,)\n",
      "h.7.attn.c_proj.w -> model/h7/attn/c_proj/w\n",
      "torch.Size([768, 768]) (1, 768, 768)\n",
      "h.7.attn.c_proj.b -> model/h7/attn/c_proj/b\n",
      "torch.Size([768]) (768,)\n",
      "h.7.ln_1.g -> model/h7/ln_1/g\n",
      "torch.Size([768]) (768,)\n",
      "h.7.ln_1.b -> model/h7/ln_1/b\n",
      "torch.Size([768]) (768,)\n",
      "h.7.mlp.c_fc.w -> model/h7/mlp/c_fc/w\n",
      "torch.Size([768, 3072]) (1, 768, 3072)\n",
      "h.7.mlp.c_fc.b -> model/h7/mlp/c_fc/b\n",
      "torch.Size([3072]) (3072,)\n",
      "h.7.mlp.c_proj.w -> model/h7/mlp/c_proj/w\n",
      "torch.Size([3072, 768]) (1, 3072, 768)\n",
      "h.7.mlp.c_proj.b -> model/h7/mlp/c_proj/b\n",
      "torch.Size([768]) (768,)\n",
      "h.7.ln_2.g -> model/h7/ln_2/g\n",
      "torch.Size([768]) (768,)\n",
      "h.7.ln_2.b -> model/h7/ln_2/b\n",
      "torch.Size([768]) (768,)\n",
      "h.8.attn.c_attn.w -> model/h8/attn/c_attn/w\n",
      "torch.Size([768, 2304]) (1, 768, 2304)\n",
      "h.8.attn.c_attn.b -> model/h8/attn/c_attn/b\n",
      "torch.Size([2304]) (2304,)\n",
      "h.8.attn.c_proj.w -> model/h8/attn/c_proj/w\n",
      "torch.Size([768, 768]) (1, 768, 768)\n",
      "h.8.attn.c_proj.b -> model/h8/attn/c_proj/b\n",
      "torch.Size([768]) (768,)\n",
      "h.8.ln_1.g -> model/h8/ln_1/g\n",
      "torch.Size([768]) (768,)\n",
      "h.8.ln_1.b -> model/h8/ln_1/b\n",
      "torch.Size([768]) (768,)\n",
      "h.8.mlp.c_fc.w -> model/h8/mlp/c_fc/w\n",
      "torch.Size([768, 3072]) (1, 768, 3072)\n",
      "h.8.mlp.c_fc.b -> model/h8/mlp/c_fc/b\n",
      "torch.Size([3072]) (3072,)\n",
      "h.8.mlp.c_proj.w -> model/h8/mlp/c_proj/w\n",
      "torch.Size([3072, 768]) (1, 3072, 768)\n",
      "h.8.mlp.c_proj.b -> model/h8/mlp/c_proj/b\n",
      "torch.Size([768]) (768,)\n",
      "h.8.ln_2.g -> model/h8/ln_2/g\n",
      "torch.Size([768]) (768,)\n",
      "h.8.ln_2.b -> model/h8/ln_2/b\n",
      "torch.Size([768]) (768,)\n",
      "h.9.attn.c_attn.w -> model/h9/attn/c_attn/w\n",
      "torch.Size([768, 2304]) (1, 768, 2304)\n",
      "h.9.attn.c_attn.b -> model/h9/attn/c_attn/b\n",
      "torch.Size([2304]) (2304,)\n",
      "h.9.attn.c_proj.w -> model/h9/attn/c_proj/w\n",
      "torch.Size([768, 768]) (1, 768, 768)\n",
      "h.9.attn.c_proj.b -> model/h9/attn/c_proj/b\n",
      "torch.Size([768]) (768,)\n",
      "h.9.ln_1.g -> model/h9/ln_1/g\n",
      "torch.Size([768]) (768,)\n",
      "h.9.ln_1.b -> model/h9/ln_1/b\n",
      "torch.Size([768]) (768,)\n",
      "h.9.mlp.c_fc.w -> model/h9/mlp/c_fc/w\n",
      "torch.Size([768, 3072]) (1, 768, 3072)\n",
      "h.9.mlp.c_fc.b -> model/h9/mlp/c_fc/b\n",
      "torch.Size([3072]) (3072,)\n",
      "h.9.mlp.c_proj.w -> model/h9/mlp/c_proj/w\n",
      "torch.Size([3072, 768]) (1, 3072, 768)\n",
      "h.9.mlp.c_proj.b -> model/h9/mlp/c_proj/b\n",
      "torch.Size([768]) (768,)\n",
      "h.9.ln_2.g -> model/h9/ln_2/g\n",
      "torch.Size([768]) (768,)\n",
      "h.9.ln_2.b -> model/h9/ln_2/b\n",
      "torch.Size([768]) (768,)\n",
      "h.10.attn.c_attn.w -> model/h10/attn/c_attn/w\n",
      "torch.Size([768, 2304]) (1, 768, 2304)\n",
      "h.10.attn.c_attn.b -> model/h10/attn/c_attn/b\n",
      "torch.Size([2304]) (2304,)\n",
      "h.10.attn.c_proj.w -> model/h10/attn/c_proj/w\n",
      "torch.Size([768, 768]) (1, 768, 768)\n",
      "h.10.attn.c_proj.b -> model/h10/attn/c_proj/b\n",
      "torch.Size([768]) (768,)\n",
      "h.10.ln_1.g -> model/h10/ln_1/g\n",
      "torch.Size([768]) (768,)\n",
      "h.10.ln_1.b -> model/h10/ln_1/b\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([768]) (768,)\n",
      "h.10.mlp.c_fc.w -> model/h10/mlp/c_fc/w\n",
      "torch.Size([768, 3072]) (1, 768, 3072)\n",
      "h.10.mlp.c_fc.b -> model/h10/mlp/c_fc/b\n",
      "torch.Size([3072]) (3072,)\n",
      "h.10.mlp.c_proj.w -> model/h10/mlp/c_proj/w\n",
      "torch.Size([3072, 768]) (1, 3072, 768)\n",
      "h.10.mlp.c_proj.b -> model/h10/mlp/c_proj/b\n",
      "torch.Size([768]) (768,)\n",
      "h.10.ln_2.g -> model/h10/ln_2/g\n",
      "torch.Size([768]) (768,)\n",
      "h.10.ln_2.b -> model/h10/ln_2/b\n",
      "torch.Size([768]) (768,)\n",
      "h.11.attn.c_attn.w -> model/h11/attn/c_attn/w\n",
      "torch.Size([768, 2304]) (1, 768, 2304)\n",
      "h.11.attn.c_attn.b -> model/h11/attn/c_attn/b\n",
      "torch.Size([2304]) (2304,)\n",
      "h.11.attn.c_proj.w -> model/h11/attn/c_proj/w\n",
      "torch.Size([768, 768]) (1, 768, 768)\n",
      "h.11.attn.c_proj.b -> model/h11/attn/c_proj/b\n",
      "torch.Size([768]) (768,)\n",
      "h.11.ln_1.g -> model/h11/ln_1/g\n",
      "torch.Size([768]) (768,)\n",
      "h.11.ln_1.b -> model/h11/ln_1/b\n",
      "torch.Size([768]) (768,)\n",
      "h.11.mlp.c_fc.w -> model/h11/mlp/c_fc/w\n",
      "torch.Size([768, 3072]) (1, 768, 3072)\n",
      "h.11.mlp.c_fc.b -> model/h11/mlp/c_fc/b\n",
      "torch.Size([3072]) (3072,)\n",
      "h.11.mlp.c_proj.w -> model/h11/mlp/c_proj/w\n",
      "torch.Size([3072, 768]) (1, 3072, 768)\n",
      "h.11.mlp.c_proj.b -> model/h11/mlp/c_proj/b\n",
      "torch.Size([768]) (768,)\n",
      "h.11.ln_2.g -> model/h11/ln_2/g\n",
      "torch.Size([768]) (768,)\n",
      "h.11.ln_2.b -> model/h11/ln_2/b\n",
      "torch.Size([768]) (768,)\n"
     ]
    }
   ],
   "source": [
    "for name, p in model.named_parameters():\n",
    "    path = name.split(\".\")\n",
    "    if path[0] == \"h\":\n",
    "        tf_name = f\"model/{path[0]}{path[1]}/\" + \"/\".join(path[2:])\n",
    "        print(f\"{name} -> {tf_name}\")\n",
    "        var_np = tf.train.load_variable(ckpt_path, tf_name)\n",
    "    elif name == \"embed.weight\":\n",
    "        tf_names = [\"model/we\", \"model/pe\"]\n",
    "        print(f\"{name} -> {tf_names}\")\n",
    "        vars_np = [tf.train.load_variable(ckpt_path, tf_name) for tf_name in tf_names]\n",
    "        var_np = np.concatenate(vars_np, axis=0)\n",
    "    else:\n",
    "        raise ValueError(f\"Unrecognized name: {name}\")\n",
    "    print(p.shape, var_np.shape)\n",
    "    delta = np.abs((p.detach().numpy() - var_np))\n",
    "    assert (np.max(delta) < 1e-5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# AllenNLP Implementation\n",
    "\n",
    "See code here: https://github.com/allenai/allennlp/blob/master/allennlp/modules/openai_transformer.py\n",
    "\n",
    "This is ported from the Huggingface implementation and looks a bit cleaner, but it isn't immediately clear how to use it - it expects a tarfile containing the weights, but it's not clear how to generate this. We'd also need to update AllenNLP, which might break other experiments using `jiant`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}