convert.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "import torch\n",
    "import tensorflow as tf\n",
    "import tensorflow_hub as hub\n",
    "import tensorflow_text as text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_PATH = 'models/universal-sentence-encoder-multilingual-large-3'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tensorflow.python.saved_model.load.Loader._recreate_base_user_object.<locals>._UserObject at 0x35f72b1f0>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = tf.saved_model.load(MODEL_PATH)\n",
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(1, 512), dtype=float32, numpy=\n",
       "array([[ 2.77998135e-03,  2.53025256e-02, -1.34592326e-02,\n",
       "        -6.00954033e-02, -3.57862748e-03,  2.16413625e-02,\n",
       "         1.99971721e-02, -3.70096639e-02, -2.91313529e-02,\n",
       "        -1.07441902e-01,  2.73036622e-02, -8.28249753e-02,\n",
       "         1.23869926e-02,  2.55314838e-02, -3.11783515e-02,\n",
       "         7.67881498e-02,  5.88866174e-02,  9.00684148e-02,\n",
       "        -2.02289019e-02,  5.06051891e-02,  4.59435731e-02,\n",
       "        -9.70152542e-02,  1.44991055e-02,  1.12706097e-02,\n",
       "        -2.17566569e-03, -3.42217949e-03,  7.06949830e-02,\n",
       "        -3.22969817e-02, -2.33398993e-02, -2.30249688e-02,\n",
       "         1.44220600e-02, -1.73140336e-02, -6.50571808e-02,\n",
       "        -6.18350431e-02,  7.21207708e-02, -6.45894706e-02,\n",
       "         2.22161412e-02,  8.38220119e-04, -1.77222937e-02,\n",
       "        -4.84013595e-02,  4.19013053e-02, -1.89781878e-02,\n",
       "        -1.09328344e-01,  4.43844795e-02,  1.54845491e-02,\n",
       "         4.25968394e-02,  4.46535647e-02,  7.47142434e-02,\n",
       "         2.83761024e-02,  2.99820378e-02, -1.81293916e-02,\n",
       "         1.09258704e-02, -2.19645463e-02,  4.04613204e-02,\n",
       "        -1.61833521e-02, -6.70688972e-02, -4.69519906e-02,\n",
       "        -2.98928842e-02,  2.39872634e-02, -5.08237183e-02,\n",
       "         2.14721523e-02, -7.54037276e-02, -6.19464405e-02,\n",
       "         8.02062638e-03,  6.10234886e-02, -1.44415749e-02,\n",
       "         8.13358184e-03,  1.68494470e-02, -4.15561572e-02,\n",
       "         1.77687183e-02, -7.53267296e-03, -2.26353537e-02,\n",
       "         4.29195762e-02,  2.85946839e-02, -4.81182570e-03,\n",
       "         2.24124156e-02,  2.55373679e-03,  1.89670753e-02,\n",
       "        -7.64294863e-02, -3.49147879e-02, -2.11369526e-02,\n",
       "        -9.53677297e-02,  4.15931419e-02, -5.94347343e-03,\n",
       "        -5.39376028e-02,  2.14008964e-03,  3.80719430e-04,\n",
       "        -1.61220077e-02,  1.41611379e-02,  3.96732651e-02,\n",
       "         1.77479554e-02, -3.02331205e-02, -5.95526621e-02,\n",
       "         3.91535573e-02, -2.23090984e-02,  8.88230577e-02,\n",
       "        -1.92017481e-02, -2.61824578e-03,  6.14531105e-03,\n",
       "         1.80025154e-03, -5.57749830e-02,  5.92270903e-02,\n",
       "        -1.68005433e-02, -3.18426639e-02,  4.09816280e-02,\n",
       "        -1.81023352e-04, -3.60204168e-02, -3.28274071e-02,\n",
       "         4.67756651e-02,  2.73127668e-02, -8.71790107e-03,\n",
       "        -8.09883699e-03, -1.35077620e-02, -1.48433149e-02,\n",
       "         2.67082136e-02,  6.55581383e-03, -5.04650129e-03,\n",
       "         7.13304756e-03,  7.08723664e-02, -2.49839686e-02,\n",
       "         1.52524340e-03,  2.14667041e-02, -3.09981201e-02,\n",
       "        -5.58509901e-02, -5.15365787e-03, -1.94244739e-02,\n",
       "         2.65732333e-02, -1.61060281e-02,  2.08161231e-02,\n",
       "        -1.51349343e-02, -2.31544697e-03, -1.08981030e-02,\n",
       "        -2.41845883e-02, -2.47194618e-02,  3.23315300e-02,\n",
       "         3.03313266e-02, -3.98456492e-02, -1.91458166e-02,\n",
       "        -2.08768491e-02, -2.60891598e-02, -2.88057271e-02,\n",
       "        -5.60498945e-02,  3.65354754e-02,  3.48301865e-02,\n",
       "        -6.63715899e-02,  5.57560399e-02, -3.79125960e-02,\n",
       "         1.27179384e-01,  3.49689424e-02,  9.82431229e-03,\n",
       "         8.82499218e-02, -7.09730387e-02,  5.70126884e-02,\n",
       "         4.82766982e-03,  6.73010573e-02,  2.81971823e-02,\n",
       "        -6.00713454e-02, -7.12266564e-02,  2.93501373e-02,\n",
       "        -3.62643017e-03, -3.86599749e-02, -3.42651755e-02,\n",
       "        -1.49719445e-02,  5.34117483e-02,  2.21329741e-02,\n",
       "         3.80360335e-02,  7.92432129e-02,  6.38215765e-02,\n",
       "        -6.55527338e-02, -3.70032936e-02,  1.47927940e-01,\n",
       "         6.70504570e-02, -4.83767577e-02,  5.66275679e-02,\n",
       "         9.02498327e-03,  3.62754799e-03, -7.84711912e-03,\n",
       "         6.11578263e-02, -5.13800569e-02,  1.84409749e-02,\n",
       "         2.19992753e-02, -6.06667437e-02,  2.58410331e-02,\n",
       "        -1.00607779e-02, -5.09426482e-02, -3.91690694e-02,\n",
       "        -8.47404152e-02,  7.25860372e-02,  4.38471287e-02,\n",
       "        -8.56803283e-02, -4.30491604e-02, -3.45602259e-02,\n",
       "        -4.79763523e-02,  8.32510740e-02, -4.25256118e-02,\n",
       "         6.13699947e-03, -7.28254691e-02, -2.72577647e-02,\n",
       "        -3.09629478e-02, -1.30163680e-04,  2.42168196e-02,\n",
       "        -1.18498892e-01,  1.87116908e-03,  1.53314508e-02,\n",
       "         8.82200301e-02,  4.92649712e-02,  5.08735701e-02,\n",
       "        -9.25842952e-03,  2.41887085e-02,  4.52499976e-03,\n",
       "         1.50375850e-02,  3.89459543e-02, -4.84806634e-02,\n",
       "         9.22602490e-02,  2.22077202e-02,  2.02294551e-02,\n",
       "        -3.98275107e-02,  6.14489755e-03, -3.40088904e-02,\n",
       "         2.80292071e-02,  2.36546863e-02, -1.27063403e-02,\n",
       "         8.58633313e-03,  3.25256065e-02,  1.28303543e-02,\n",
       "        -5.77415563e-02, -3.79824429e-03,  5.54539450e-02,\n",
       "         6.01440035e-02,  3.69424820e-02, -6.09921627e-02,\n",
       "        -2.49674208e-02,  2.34099142e-02, -5.01966290e-02,\n",
       "         6.31935075e-02,  8.73863548e-02,  4.87505756e-02,\n",
       "         3.61433178e-02, -5.12333103e-02, -3.74389477e-02,\n",
       "        -6.04487732e-02, -6.73208237e-02,  4.75760549e-02,\n",
       "        -3.37632373e-03,  3.79770063e-02,  1.72327422e-02,\n",
       "         4.05565612e-02,  4.78833802e-02,  1.91632491e-02,\n",
       "         6.28554001e-02,  4.19324152e-02,  3.59204151e-02,\n",
       "         1.67394858e-02,  5.72051257e-02, -4.80927527e-03,\n",
       "         1.74812484e-03,  7.00754970e-02, -7.28979940e-03,\n",
       "         2.60920413e-02,  3.77972913e-03, -7.54555687e-02,\n",
       "        -6.70646802e-02,  2.20567989e-03, -2.10349057e-02,\n",
       "         3.84148620e-02, -4.00563255e-02, -2.57295221e-02,\n",
       "         2.50411546e-03, -9.20192897e-02, -1.01443268e-02,\n",
       "         2.87408605e-02,  4.11674008e-02,  1.52321979e-02,\n",
       "         1.08864577e-02,  5.09297960e-02,  2.84241661e-02,\n",
       "        -2.89797108e-03, -7.44252931e-03,  5.18899933e-02,\n",
       "         4.84143430e-03,  3.09501085e-02, -8.28481391e-02,\n",
       "         1.30433813e-02, -8.53823498e-02, -2.23566685e-02,\n",
       "         2.29566526e-02, -5.32987108e-03, -8.51971582e-02,\n",
       "        -7.36641698e-04, -2.38095094e-02, -1.89320613e-02,\n",
       "         7.01749185e-03, -7.06940442e-02, -9.57594719e-03,\n",
       "         3.45089696e-02, -4.76918295e-02,  2.08721682e-02,\n",
       "         3.50489165e-03,  1.42544275e-02, -1.37221530e-01,\n",
       "         5.96887283e-02,  6.52726460e-03, -4.03699465e-03,\n",
       "        -4.58690012e-03, -7.12437779e-02, -1.47907753e-02,\n",
       "        -4.54501584e-02, -5.64297475e-03, -1.65209640e-02,\n",
       "         6.18872568e-02,  3.99036780e-02,  5.52764013e-02,\n",
       "         1.01573952e-02, -1.10726722e-01,  2.46507153e-02,\n",
       "        -2.29608957e-02, -8.90163984e-03, -3.45015787e-02,\n",
       "        -7.51133040e-02,  5.25646694e-02, -2.96319462e-03,\n",
       "         3.43846120e-02,  4.42368835e-02,  8.32889171e-04,\n",
       "         7.64172077e-02, -2.48780269e-02,  9.73852817e-03,\n",
       "        -6.71786861e-03,  9.36219562e-03, -9.74921323e-03,\n",
       "        -2.96224244e-02,  5.70549490e-03,  6.05000705e-02,\n",
       "         7.96179846e-02, -2.64466032e-02, -2.75871139e-02,\n",
       "         5.65553121e-02,  2.72857323e-02,  2.85837539e-02,\n",
       "         2.57873200e-02, -9.57705081e-03, -4.77914698e-02,\n",
       "         1.17656477e-02, -2.02163197e-02,  6.72671795e-02,\n",
       "         2.91052274e-02, -8.31819922e-02,  2.71085016e-02,\n",
       "        -3.38688344e-02,  4.11870442e-02, -5.76304980e-02,\n",
       "         2.65145767e-02,  6.44305488e-03, -5.14965467e-02,\n",
       "        -3.33812600e-03, -2.49319002e-02, -2.52748299e-02,\n",
       "        -7.71390870e-02, -3.69268283e-02, -2.95937806e-02,\n",
       "         2.23889723e-02, -5.63586242e-02,  6.38240948e-02,\n",
       "         1.04810167e-02, -7.26583088e-03,  3.24797118e-03,\n",
       "        -2.06953734e-02,  2.37911777e-03, -3.60643901e-02,\n",
       "        -7.86193684e-02, -7.24616051e-02, -8.23280364e-02,\n",
       "         2.04754043e-02,  8.57350789e-03, -2.00888310e-02,\n",
       "        -2.19155084e-02,  1.86027512e-02, -4.04845066e-02,\n",
       "        -7.49541968e-02,  5.58337346e-02,  5.76391742e-02,\n",
       "         3.68528487e-03,  1.83450524e-02,  5.54439202e-02,\n",
       "         6.72838688e-02, -1.31846815e-02,  6.15373701e-02,\n",
       "        -2.98579112e-02,  4.01462168e-02, -1.29247168e-02,\n",
       "         6.50901943e-02,  5.57566062e-03, -2.49074912e-03,\n",
       "         3.78542133e-02,  3.93357985e-02,  7.44291991e-02,\n",
       "         6.30763918e-02,  3.86597253e-02,  4.86251190e-02,\n",
       "        -2.88229138e-02,  1.20447483e-02,  6.79193810e-02,\n",
       "         3.28842402e-02,  2.78828107e-02, -6.27110302e-02,\n",
       "        -5.89094684e-02, -1.31008492e-04, -1.05010839e-02,\n",
       "        -2.22017914e-02, -5.63225411e-02, -2.54515391e-02,\n",
       "        -2.41848715e-02, -4.20510210e-03, -4.94399965e-02,\n",
       "        -6.96421415e-02,  7.55605996e-02, -2.45151632e-02,\n",
       "         6.02281243e-02,  6.73289150e-02,  2.96716224e-02,\n",
       "         6.14644960e-02, -1.00285206e-02, -3.10751572e-02,\n",
       "        -3.53258066e-02,  7.03033060e-02,  7.18787611e-02,\n",
       "         4.44826484e-02, -3.78371254e-02,  1.53917922e-02,\n",
       "        -2.06896346e-02, -3.44437435e-02, -3.55283730e-02,\n",
       "         8.49672034e-02,  3.08616310e-02,  7.09353238e-02,\n",
       "         5.11887185e-02,  5.87357022e-02,  4.11299523e-03,\n",
       "        -7.45894341e-03, -4.06427197e-02,  2.33926103e-02,\n",
       "        -6.46680593e-02, -4.08164673e-02,  1.37195364e-02,\n",
       "         4.84200232e-02, -2.30469555e-02, -7.85305649e-02,\n",
       "        -1.68109298e-04, -5.46142608e-02,  5.12183011e-02,\n",
       "         1.28705911e-02,  5.13147488e-02,  3.28289531e-02,\n",
       "         3.73959281e-02,  1.84249654e-02,  5.28426468e-02,\n",
       "        -1.08567085e-02, -2.27842629e-02, -4.02348563e-02,\n",
       "         4.21910174e-03,  1.69129819e-02, -3.62997269e-03,\n",
       "         6.29095733e-02,  2.73506809e-02, -2.79912706e-02,\n",
       "        -2.04870757e-02, -2.85848882e-02, -4.09637839e-02,\n",
       "         2.99389102e-02, -8.20686519e-02,  3.30640259e-03,\n",
       "         4.20482755e-02, -4.12866436e-02,  3.88541110e-02,\n",
       "         2.87617184e-02,  7.94692487e-02, -3.04842237e-02,\n",
       "        -1.75349377e-02,  1.79946106e-02,  5.41700721e-02,\n",
       "         1.72643345e-02, -4.74507324e-02,  9.15510729e-02,\n",
       "        -9.70511325e-03,  1.19170093e-03, -3.62269580e-02,\n",
       "        -1.30497674e-02,  5.81335165e-02, -8.87113996e-03,\n",
       "         3.48886959e-02,  1.82030406e-02, -7.50234025e-03,\n",
       "         6.49072835e-03,  4.03849594e-02,  8.31530541e-02,\n",
       "         7.39139318e-03, -2.93784589e-02, -8.50977227e-02,\n",
       "        -7.88823888e-02, -8.11857637e-03, -2.94542033e-02,\n",
       "        -1.90729722e-02,  2.42671953e-03,  6.11241125e-02,\n",
       "        -4.86013405e-02,  3.48782092e-02,  3.80351767e-02,\n",
       "        -3.75011116e-02,  1.86881255e-02,  7.17819110e-02,\n",
       "        -3.18108387e-02, -3.38131189e-02]], dtype=float32)>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sentence = \"Hello, how are you?\"\n",
    "ref = model(sentence)\n",
    "ref"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "mapper = defaultdict(list)\n",
    "appendix_set = set()\n",
    "parts_set = set()\n",
    "\n",
    "for tensor in model.trainable_variables:\n",
    "    name, appendix = tensor.name.split(':')\n",
    "    appendix_set.add(appendix)\n",
    "    name, part = name.rsplit('/', maxsplit=1)\n",
    "    parts_set.add(part)\n",
    "    mapper[name].append(tensor.numpy())\n",
    "\n",
    "assert appendix_set == {'0'}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'part_0',\n",
       " 'part_1',\n",
       " 'part_10',\n",
       " 'part_11',\n",
       " 'part_12',\n",
       " 'part_13',\n",
       " 'part_14',\n",
       " 'part_15',\n",
       " 'part_16',\n",
       " 'part_2',\n",
       " 'part_3',\n",
       " 'part_4',\n",
       " 'part_5',\n",
       " 'part_6',\n",
       " 'part_7',\n",
       " 'part_8',\n",
       " 'part_9',\n",
       " 'sharded_0',\n",
       " 'sharded_1',\n",
       " 'sharded_10',\n",
       " 'sharded_11',\n",
       " 'sharded_12',\n",
       " 'sharded_13',\n",
       " 'sharded_14',\n",
       " 'sharded_15',\n",
       " 'sharded_16',\n",
       " 'sharded_2',\n",
       " 'sharded_3',\n",
       " 'sharded_4',\n",
       " 'sharded_5',\n",
       " 'sharded_6',\n",
       " 'sharded_7',\n",
       " 'sharded_8',\n",
       " 'sharded_9'}"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "parts_set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "weights = {}\n",
    "\n",
    "for name, tensors_list in mapper.items():\n",
    "    if name == 'Embeddings':\n",
    "        n_matrices = len(tensors_list)\n",
    "        num_embeddings = n_matrices * tensors_list[0].shape[0]\n",
    "        embedding_dim = tensors_list[0].shape[1]\n",
    "        embeddings = np.zeros((num_embeddings, embedding_dim))\n",
    "        for idx in range(num_embeddings):\n",
    "            j = idx // n_matrices\n",
    "            i = idx - j * n_matrices\n",
    "            embeddings[idx] = tensors_list[i][j]\n",
    "        weights[name] = embeddings\n",
    "    else:\n",
    "        weights[name] = np.concatenate(tensors_list, axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Embeddings: (128010, 512)\n",
      "EncoderDNN/DNN/ResidualHidden_0/dense/kernel: (512, 320)\n",
      "EncoderDNN/DNN/ResidualHidden_1/dense/kernel: (320, 320)\n",
      "EncoderDNN/DNN/ResidualHidden_1/AdjustDepth/projection/kernel: (512, 320)\n",
      "EncoderDNN/DNN/ResidualHidden_2/dense/kernel: (320, 512)\n",
      "EncoderDNN/DNN/ResidualHidden_3/dense/kernel: (512, 512)\n",
      "EncoderDNN/DNN/ResidualHidden_3/AdjustDepth/projection/kernel: (320, 512)\n",
      "EncoderTransformer/Transformer/dense/kernel: (512, 512)\n",
      "EncoderTransformer/Transformer/dense/bias: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_0/SelfAttention/layer_prepostprocess/layer_norm/layer_norm_scale: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_0/SelfAttention/layer_prepostprocess/layer_norm/layer_norm_bias: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_0/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_q/kernel: (512, 512)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_0/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_q/bias: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_0/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_k/kernel: (512, 512)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_0/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_k/bias: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_0/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_v/kernel: (512, 512)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_0/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_v/bias: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_0/FFN/Layer1/layer_prepostprocess/layer_norm/layer_norm_scale: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_0/FFN/Layer1/layer_prepostprocess/layer_norm/layer_norm_bias: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_0/FFN/Layer1/dense/kernel: (512, 2048)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_0/FFN/Layer1/dense/bias: (2048,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_0/FFN/Layer2/dense/kernel: (2048, 512)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_0/FFN/Layer2/dense/bias: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_1/SelfAttention/layer_prepostprocess/layer_norm/layer_norm_scale: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_1/SelfAttention/layer_prepostprocess/layer_norm/layer_norm_bias: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_1/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_q/kernel: (512, 512)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_1/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_q/bias: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_1/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_k/kernel: (512, 512)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_1/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_k/bias: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_1/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_v/kernel: (512, 512)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_1/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_v/bias: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_1/FFN/Layer1/layer_prepostprocess/layer_norm/layer_norm_scale: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_1/FFN/Layer1/layer_prepostprocess/layer_norm/layer_norm_bias: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_1/FFN/Layer1/dense/kernel: (512, 2048)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_1/FFN/Layer1/dense/bias: (2048,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_1/FFN/Layer2/dense/kernel: (2048, 512)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_1/FFN/Layer2/dense/bias: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_2/SelfAttention/layer_prepostprocess/layer_norm/layer_norm_scale: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_2/SelfAttention/layer_prepostprocess/layer_norm/layer_norm_bias: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_2/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_q/kernel: (512, 512)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_2/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_q/bias: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_2/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_k/kernel: (512, 512)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_2/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_k/bias: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_2/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_v/kernel: (512, 512)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_2/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_v/bias: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_2/FFN/Layer1/layer_prepostprocess/layer_norm/layer_norm_scale: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_2/FFN/Layer1/layer_prepostprocess/layer_norm/layer_norm_bias: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_2/FFN/Layer1/dense/kernel: (512, 2048)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_2/FFN/Layer1/dense/bias: (2048,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_2/FFN/Layer2/dense/kernel: (2048, 512)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_2/FFN/Layer2/dense/bias: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_3/SelfAttention/layer_prepostprocess/layer_norm/layer_norm_scale: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_3/SelfAttention/layer_prepostprocess/layer_norm/layer_norm_bias: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_3/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_q/kernel: (512, 512)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_3/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_q/bias: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_3/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_k/kernel: (512, 512)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_3/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_k/bias: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_3/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_v/kernel: (512, 512)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_3/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_v/bias: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_3/FFN/Layer1/layer_prepostprocess/layer_norm/layer_norm_scale: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_3/FFN/Layer1/layer_prepostprocess/layer_norm/layer_norm_bias: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_3/FFN/Layer1/dense/kernel: (512, 2048)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_3/FFN/Layer1/dense/bias: (2048,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_3/FFN/Layer2/dense/kernel: (2048, 512)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_3/FFN/Layer2/dense/bias: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_4/SelfAttention/layer_prepostprocess/layer_norm/layer_norm_scale: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_4/SelfAttention/layer_prepostprocess/layer_norm/layer_norm_bias: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_4/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_q/kernel: (512, 512)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_4/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_q/bias: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_4/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_k/kernel: (512, 512)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_4/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_k/bias: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_4/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_v/kernel: (512, 512)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_4/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_v/bias: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_4/FFN/Layer1/layer_prepostprocess/layer_norm/layer_norm_scale: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_4/FFN/Layer1/layer_prepostprocess/layer_norm/layer_norm_bias: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_4/FFN/Layer1/dense/kernel: (512, 2048)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_4/FFN/Layer1/dense/bias: (2048,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_4/FFN/Layer2/dense/kernel: (2048, 512)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_4/FFN/Layer2/dense/bias: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_5/SelfAttention/layer_prepostprocess/layer_norm/layer_norm_scale: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_5/SelfAttention/layer_prepostprocess/layer_norm/layer_norm_bias: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_5/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_q/kernel: (512, 512)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_5/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_q/bias: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_5/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_k/kernel: (512, 512)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_5/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_k/bias: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_5/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_v/kernel: (512, 512)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_5/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_v/bias: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_5/FFN/Layer1/layer_prepostprocess/layer_norm/layer_norm_scale: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_5/FFN/Layer1/layer_prepostprocess/layer_norm/layer_norm_bias: (512,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_5/FFN/Layer1/dense/kernel: (512, 2048)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_5/FFN/Layer1/dense/bias: (2048,)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_5/FFN/Layer2/dense/kernel: (2048, 512)\n",
      "EncoderTransformer/Transformer/SparseTransformerEncode/Layer_5/FFN/Layer2/dense/bias: (512,)\n",
      "EncoderTransformer/Transformer/layer_prepostprocess/layer_norm/layer_norm_scale: (512,)\n",
      "EncoderTransformer/Transformer/layer_prepostprocess/layer_norm/layer_norm_bias: (512,)\n",
      "EncoderTransformer/Transformer/AttentionPooling/AttentionHidden/kernel: (512, 512)\n",
      "EncoderTransformer/Transformer/AttentionPooling/AttentionHidden/bias: (512,)\n",
      "EncoderTransformer/Transformer/AttentionPooling/AttentionLogits/kernel: (512, 512)\n",
      "EncoderTransformer/Transformer/AttentionPooling/AttentionLogits/bias: (512,)\n",
      "EncoderTransformer/hidden_layers/tanh_layer_0/dense/kernel: (1024, 512)\n",
      "EncoderTransformer/hidden_layers/tanh_layer_0/dense/bias: (512,)\n"
     ]
    }
   ],
   "source": [
    "for name, weight in weights.items():\n",
    "    print(f'{name}: {weight.shape}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import partial\n",
    "from src.tokenizer import get_tokenizer, tokenize\n",
    "\n",
    "tokenizer = get_tokenizer(MODEL_PATH)\n",
    "tokenize = partial(tokenize, tokenizer=tokenizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Architecture"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "from src.architecture import MUSE\n",
    "\n",
    "model_torch = MUSE(\n",
    "    num_embeddings=weights['Embeddings'].shape[0],\n",
    "    embedding_dim=weights['Embeddings'].shape[1],\n",
    "    d_model=512,\n",
    "    num_heads=8,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Init\n",
    "with torch.no_grad():\n",
    "    model_torch.embedding.weight.copy_(torch.from_numpy(weights['Embeddings']))\n",
    "    \n",
    "    model_torch.linear.weight.copy_(torch.from_numpy(weights['EncoderTransformer/Transformer/dense/kernel'].T))\n",
    "    model_torch.linear.bias.copy_(torch.from_numpy(weights['EncoderTransformer/Transformer/dense/bias']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Init\n",
    "def init_block(block_num):\n",
    "    block = getattr(model_torch, f'block{block_num}')\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        # ln1.weight\n",
    "        tensor_weight = block.ln1.weight\n",
    "        array = weights[f'EncoderTransformer/Transformer/SparseTransformerEncode/Layer_{block_num}/SelfAttention/layer_prepostprocess/layer_norm/layer_norm_scale']\n",
    "        assert tensor_weight.shape == array.shape\n",
    "        tensor_weight.copy_(torch.from_numpy(array))\n",
    "\n",
    "        # ln1.bias\n",
    "        tensor_weight = block.ln1.bias\n",
    "        array = weights[f'EncoderTransformer/Transformer/SparseTransformerEncode/Layer_{block_num}/SelfAttention/layer_prepostprocess/layer_norm/layer_norm_bias']\n",
    "        assert tensor_weight.shape == array.shape\n",
    "        tensor_weight.copy_(torch.from_numpy(array))\n",
    "\n",
    "        # attn.query.weight\n",
    "        tensor_weight = block.attn.query.weight\n",
    "        array = weights[f'EncoderTransformer/Transformer/SparseTransformerEncode/Layer_{block_num}/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_q/kernel'].T\n",
    "        assert tensor_weight.shape == array.shape\n",
    "        tensor_weight.copy_(torch.from_numpy(array))\n",
    "\n",
    "        # attn.query.bias\n",
    "        tensor_weight = block.attn.query.bias\n",
    "        array = weights[f'EncoderTransformer/Transformer/SparseTransformerEncode/Layer_{block_num}/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_q/bias']\n",
    "        assert tensor_weight.shape == array.shape\n",
    "        tensor_weight.copy_(torch.from_numpy(array))\n",
    "\n",
    "        # attn.key.weight\n",
    "        tensor_weight = block.attn.key.weight\n",
    "        array = weights[f'EncoderTransformer/Transformer/SparseTransformerEncode/Layer_{block_num}/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_k/kernel'].T\n",
    "        assert tensor_weight.shape == array.shape\n",
    "        tensor_weight.copy_(torch.from_numpy(array))\n",
    "\n",
    "        # attn.key.bias\n",
    "        tensor_weight = block.attn.key.bias\n",
    "        array = weights[f'EncoderTransformer/Transformer/SparseTransformerEncode/Layer_{block_num}/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_k/bias']\n",
    "        assert tensor_weight.shape == array.shape\n",
    "        tensor_weight.copy_(torch.from_numpy(array))\n",
    "\n",
    "        # attn.value.weight\n",
    "        tensor_weight = block.attn.value.weight\n",
    "        array = weights[f'EncoderTransformer/Transformer/SparseTransformerEncode/Layer_{block_num}/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_v/kernel'].T\n",
    "        assert tensor_weight.shape == array.shape\n",
    "        tensor_weight.copy_(torch.from_numpy(array))\n",
    "\n",
    "        # attn.value.bias\n",
    "        tensor_weight = block.attn.value.bias\n",
    "        array = weights[f'EncoderTransformer/Transformer/SparseTransformerEncode/Layer_{block_num}/SelfAttention/SparseMultiheadAttention/ComputeQKV/compute_v/bias']\n",
    "        assert tensor_weight.shape == array.shape\n",
    "        tensor_weight.copy_(torch.from_numpy(array))\n",
    "\n",
    "        # ln2.weight\n",
    "        tensor_weight = block.ln2.weight\n",
    "        array = weights[f'EncoderTransformer/Transformer/SparseTransformerEncode/Layer_{block_num}/FFN/Layer1/layer_prepostprocess/layer_norm/layer_norm_scale']\n",
    "        assert tensor_weight.shape == array.shape\n",
    "        tensor_weight.copy_(torch.from_numpy(array))\n",
    "\n",
    "        # ln2.bias\n",
    "        tensor_weight = block.ln2.bias\n",
    "        array = weights[f'EncoderTransformer/Transformer/SparseTransformerEncode/Layer_{block_num}/FFN/Layer1/layer_prepostprocess/layer_norm/layer_norm_bias']\n",
    "        assert tensor_weight.shape == array.shape\n",
    "        tensor_weight.copy_(torch.from_numpy(array))\n",
    "\n",
    "        # linear1.weight\n",
    "        tensor_weight = block.linear1.weight\n",
    "        array = weights[f'EncoderTransformer/Transformer/SparseTransformerEncode/Layer_{block_num}/FFN/Layer1/dense/kernel'].T\n",
    "        assert tensor_weight.shape == array.shape\n",
    "        tensor_weight.copy_(torch.from_numpy(array))\n",
    "\n",
    "        # linear1.bias\n",
    "        tensor_weight = block.linear1.bias\n",
    "        array = weights[f'EncoderTransformer/Transformer/SparseTransformerEncode/Layer_{block_num}/FFN/Layer1/dense/bias']\n",
    "        assert tensor_weight.shape == array.shape\n",
    "        tensor_weight.copy_(torch.from_numpy(array))\n",
    "\n",
    "        # linear2.weight\n",
    "        tensor_weight = block.linear2.weight\n",
    "        array = weights[f'EncoderTransformer/Transformer/SparseTransformerEncode/Layer_{block_num}/FFN/Layer2/dense/kernel'].T\n",
    "        assert tensor_weight.shape == array.shape\n",
    "        tensor_weight.copy_(torch.from_numpy(array))\n",
    "\n",
    "        # linear2.bias\n",
    "        tensor_weight = block.linear2.bias\n",
    "        array = weights[f'EncoderTransformer/Transformer/SparseTransformerEncode/Layer_{block_num}/FFN/Layer2/dense/bias']\n",
    "        assert tensor_weight.shape == array.shape\n",
    "        tensor_weight.copy_(torch.from_numpy(array))\n",
    "\n",
    "init_block(0)\n",
    "init_block(1)\n",
    "init_block(2)\n",
    "init_block(3)\n",
    "init_block(4)\n",
    "init_block(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Init\n",
    "with torch.no_grad():\n",
    "    model_torch.head.ln.weight.copy_(torch.from_numpy(weights['EncoderTransformer/Transformer/layer_prepostprocess/layer_norm/layer_norm_scale']))\n",
    "    model_torch.head.ln.bias.copy_(torch.from_numpy(weights['EncoderTransformer/Transformer/layer_prepostprocess/layer_norm/layer_norm_bias']))\n",
    "\n",
    "    model_torch.head.linear1.weight.copy_(torch.from_numpy(weights['EncoderTransformer/Transformer/AttentionPooling/AttentionHidden/kernel'].T))\n",
    "    model_torch.head.linear1.bias.copy_(torch.from_numpy(weights['EncoderTransformer/Transformer/AttentionPooling/AttentionHidden/bias']))\n",
    "\n",
    "    model_torch.head.linear2.weight.copy_(torch.from_numpy(weights['EncoderTransformer/Transformer/AttentionPooling/AttentionLogits/kernel'].T))\n",
    "    model_torch.head.linear2.bias.copy_(torch.from_numpy(weights['EncoderTransformer/Transformer/AttentionPooling/AttentionLogits/bias']))\n",
    "\n",
    "    model_torch.head.tanh_layer.weight.copy_(torch.from_numpy(weights['EncoderTransformer/hidden_layers/tanh_layer_0/dense/kernel'].T))\n",
    "    model_torch.head.tanh_layer.bias.copy_(torch.from_numpy(weights['EncoderTransformer/hidden_layers/tanh_layer_0/dense/bias']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "MUSE(\n",
       "  (embedding): Embedding(128010, 512)\n",
       "  (linear): Linear(in_features=512, out_features=512, bias=True)\n",
       "  (pe): PositionalEncoding()\n",
       "  (block0): Block(\n",
       "    (ln1): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
       "    (attn): MultiheadSelfAttention(\n",
       "      (query): Linear(in_features=512, out_features=512, bias=True)\n",
       "      (key): Linear(in_features=512, out_features=512, bias=True)\n",
       "      (value): Linear(in_features=512, out_features=512, bias=True)\n",
       "    )\n",
       "    (ln2): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
       "    (linear1): Linear(in_features=512, out_features=2048, bias=True)\n",
       "    (linear2): Linear(in_features=2048, out_features=512, bias=True)\n",
       "  )\n",
       "  (block1): Block(\n",
       "    (ln1): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
       "    (attn): MultiheadSelfAttention(\n",
       "      (query): Linear(in_features=512, out_features=512, bias=True)\n",
       "      (key): Linear(in_features=512, out_features=512, bias=True)\n",
       "      (value): Linear(in_features=512, out_features=512, bias=True)\n",
       "    )\n",
       "    (ln2): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
       "    (linear1): Linear(in_features=512, out_features=2048, bias=True)\n",
       "    (linear2): Linear(in_features=2048, out_features=512, bias=True)\n",
       "  )\n",
       "  (block2): Block(\n",
       "    (ln1): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
       "    (attn): MultiheadSelfAttention(\n",
       "      (query): Linear(in_features=512, out_features=512, bias=True)\n",
       "      (key): Linear(in_features=512, out_features=512, bias=True)\n",
       "      (value): Linear(in_features=512, out_features=512, bias=True)\n",
       "    )\n",
       "    (ln2): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
       "    (linear1): Linear(in_features=512, out_features=2048, bias=True)\n",
       "    (linear2): Linear(in_features=2048, out_features=512, bias=True)\n",
       "  )\n",
       "  (block3): Block(\n",
       "    (ln1): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
       "    (attn): MultiheadSelfAttention(\n",
       "      (query): Linear(in_features=512, out_features=512, bias=True)\n",
       "      (key): Linear(in_features=512, out_features=512, bias=True)\n",
       "      (value): Linear(in_features=512, out_features=512, bias=True)\n",
       "    )\n",
       "    (ln2): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
       "    (linear1): Linear(in_features=512, out_features=2048, bias=True)\n",
       "    (linear2): Linear(in_features=2048, out_features=512, bias=True)\n",
       "  )\n",
       "  (block4): Block(\n",
       "    (ln1): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
       "    (attn): MultiheadSelfAttention(\n",
       "      (query): Linear(in_features=512, out_features=512, bias=True)\n",
       "      (key): Linear(in_features=512, out_features=512, bias=True)\n",
       "      (value): Linear(in_features=512, out_features=512, bias=True)\n",
       "    )\n",
       "    (ln2): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
       "    (linear1): Linear(in_features=512, out_features=2048, bias=True)\n",
       "    (linear2): Linear(in_features=2048, out_features=512, bias=True)\n",
       "  )\n",
       "  (block5): Block(\n",
       "    (ln1): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
       "    (attn): MultiheadSelfAttention(\n",
       "      (query): Linear(in_features=512, out_features=512, bias=True)\n",
       "      (key): Linear(in_features=512, out_features=512, bias=True)\n",
       "      (value): Linear(in_features=512, out_features=512, bias=True)\n",
       "    )\n",
       "    (ln2): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
       "    (linear1): Linear(in_features=512, out_features=2048, bias=True)\n",
       "    (linear2): Linear(in_features=2048, out_features=512, bias=True)\n",
       "  )\n",
       "  (head): Head(\n",
       "    (ln): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n",
       "    (linear1): Linear(in_features=512, out_features=512, bias=True)\n",
       "    (linear2): Linear(in_features=512, out_features=512, bias=True)\n",
       "    (tanh_layer): Linear(in_features=1024, out_features=512, bias=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_torch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compare"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([    1, 17486,     7,   430,    69,    37,    21,     2])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "input_ids = tokenize(sentence)\n",
    "input_ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ 2.7751e-03,  2.5281e-02, -1.3426e-02, -6.0007e-02, -3.5832e-03,\n",
       "         2.1640e-02,  2.0059e-02, -3.6970e-02, -2.9052e-02, -1.0739e-01,\n",
       "         2.7322e-02, -8.2878e-02,  1.2412e-02,  2.5597e-02, -3.1173e-02,\n",
       "         7.6832e-02,  5.8875e-02,  9.0064e-02, -2.0296e-02,  5.0571e-02,\n",
       "         4.5956e-02, -9.7047e-02,  1.4511e-02,  1.1183e-02, -2.2249e-03,\n",
       "        -3.4421e-03,  7.0693e-02, -3.2278e-02, -2.3314e-02, -2.2967e-02,\n",
       "         1.4435e-02, -1.7352e-02, -6.5063e-02, -6.1835e-02,  7.2137e-02,\n",
       "        -6.4542e-02,  2.2277e-02,  8.1046e-04, -1.7738e-02, -4.8312e-02,\n",
       "         4.1955e-02, -1.8915e-02, -1.0933e-01,  4.4421e-02,  1.5477e-02,\n",
       "         4.2578e-02,  4.4713e-02,  7.4773e-02,  2.8420e-02,  3.0054e-02,\n",
       "        -1.8123e-02,  1.0940e-02, -2.1887e-02,  4.0501e-02, -1.6136e-02,\n",
       "        -6.7140e-02, -4.6971e-02, -2.9988e-02,  2.4044e-02, -5.0831e-02,\n",
       "         2.1481e-02, -7.5423e-02, -6.1983e-02,  8.0178e-03,  6.1077e-02,\n",
       "        -1.4429e-02,  8.0560e-03,  1.6757e-02, -4.1521e-02,  1.7764e-02,\n",
       "        -7.4988e-03, -2.2600e-02,  4.2967e-02,  2.8605e-02, -4.8210e-03,\n",
       "         2.2423e-02,  2.5345e-03,  1.8962e-02, -7.6487e-02, -3.4915e-02,\n",
       "        -2.1132e-02, -9.5415e-02,  4.1720e-02, -5.9738e-03, -5.3886e-02,\n",
       "         2.3248e-03,  3.8116e-04, -1.6100e-02,  1.4210e-02,  3.9701e-02,\n",
       "         1.7762e-02, -3.0177e-02, -5.9491e-02,  3.9120e-02, -2.2314e-02,\n",
       "         8.8603e-02, -1.9207e-02, -2.6229e-03,  6.1071e-03,  1.8579e-03,\n",
       "        -5.5788e-02,  5.9237e-02, -1.6905e-02, -3.1850e-02,  4.1048e-02,\n",
       "        -2.0933e-04, -3.5949e-02, -3.2764e-02,  4.6842e-02,  2.7301e-02,\n",
       "        -8.7254e-03, -8.0436e-03, -1.3535e-02, -1.4818e-02,  2.6666e-02,\n",
       "         6.6416e-03, -5.0935e-03,  7.1060e-03,  7.0908e-02, -2.4973e-02,\n",
       "         1.5554e-03,  2.1478e-02, -3.1088e-02, -5.5877e-02, -5.1086e-03,\n",
       "        -1.9425e-02,  2.6602e-02, -1.6044e-02,  2.0853e-02, -1.5139e-02,\n",
       "        -2.3111e-03, -1.0920e-02, -2.4177e-02, -2.4741e-02,  3.2328e-02,\n",
       "         3.0321e-02, -3.9913e-02, -1.9253e-02, -2.0881e-02, -2.5990e-02,\n",
       "        -2.8759e-02, -5.6062e-02,  3.6578e-02,  3.4771e-02, -6.6387e-02,\n",
       "         5.5792e-02, -3.7943e-02,  1.2719e-01,  3.4989e-02,  9.8243e-03,\n",
       "         8.8267e-02, -7.1005e-02,  5.7006e-02,  4.8424e-03,  6.7370e-02,\n",
       "         2.8119e-02, -6.0059e-02, -7.1219e-02,  2.9309e-02, -3.6357e-03,\n",
       "        -3.8667e-02, -3.4162e-02, -1.4931e-02,  5.3379e-02,  2.2079e-02,\n",
       "         3.8027e-02,  7.9276e-02,  6.3735e-02, -6.5532e-02, -3.6971e-02,\n",
       "         1.4794e-01,  6.7004e-02, -4.8289e-02,  5.6634e-02,  8.9892e-03,\n",
       "         3.6516e-03, -7.8147e-03,  6.1207e-02, -5.1378e-02,  1.8408e-02,\n",
       "         2.2024e-02, -6.0712e-02,  2.5726e-02, -1.0098e-02, -5.0911e-02,\n",
       "        -3.9268e-02, -8.4700e-02,  7.2528e-02,  4.3861e-02, -8.5654e-02,\n",
       "        -4.3026e-02, -3.4530e-02, -4.7894e-02,  8.3272e-02, -4.2603e-02,\n",
       "         6.1278e-03, -7.2730e-02, -2.7282e-02, -3.0960e-02, -1.7355e-04,\n",
       "         2.4156e-02, -1.1849e-01,  1.9064e-03,  1.5363e-02,  8.8176e-02,\n",
       "         4.9300e-02,  5.0910e-02, -9.2592e-03,  2.4182e-02,  4.4871e-03,\n",
       "         1.5000e-02,  3.8991e-02, -4.8514e-02,  9.2255e-02,  2.2225e-02,\n",
       "         2.0208e-02, -3.9831e-02,  6.1113e-03, -3.4069e-02,  2.8023e-02,\n",
       "         2.3617e-02, -1.2750e-02,  8.5905e-03,  3.2530e-02,  1.2837e-02,\n",
       "        -5.7765e-02, -3.7970e-03,  5.5494e-02,  6.0109e-02,  3.6854e-02,\n",
       "        -6.0996e-02, -2.4991e-02,  2.3419e-02, -5.0204e-02,  6.3169e-02,\n",
       "         8.7368e-02,  4.8725e-02,  3.6127e-02, -5.1205e-02, -3.7375e-02,\n",
       "        -6.0464e-02, -6.7336e-02,  4.7654e-02, -3.3996e-03,  3.7969e-02,\n",
       "         1.7216e-02,  4.0498e-02,  4.7804e-02,  1.9165e-02,  6.2892e-02,\n",
       "         4.1935e-02,  3.5896e-02,  1.6705e-02,  5.7280e-02, -4.8159e-03,\n",
       "         1.7062e-03,  7.0061e-02, -7.1948e-03,  2.6177e-02,  3.8192e-03,\n",
       "        -7.5444e-02, -6.7063e-02,  2.1721e-03, -2.1010e-02,  3.8455e-02,\n",
       "        -4.0087e-02, -2.5744e-02,  2.5010e-03, -9.1946e-02, -1.0136e-02,\n",
       "         2.8780e-02,  4.1143e-02,  1.5199e-02,  1.0797e-02,  5.0930e-02,\n",
       "         2.8396e-02, -2.9243e-03, -7.5425e-03,  5.1925e-02,  4.9008e-03,\n",
       "         3.1003e-02, -8.2874e-02,  1.2968e-02, -8.5424e-02, -2.2323e-02,\n",
       "         2.2902e-02, -5.2547e-03, -8.5209e-02, -7.8570e-04, -2.3842e-02,\n",
       "        -1.8952e-02,  6.9601e-03, -7.0726e-02, -9.5073e-03,  3.4538e-02,\n",
       "        -4.7675e-02,  2.0880e-02,  3.4244e-03,  1.4280e-02, -1.3721e-01,\n",
       "         5.9705e-02,  6.4055e-03, -4.0521e-03, -4.5724e-03, -7.1193e-02,\n",
       "        -1.4772e-02, -4.5487e-02, -5.6009e-03, -1.6456e-02,  6.1893e-02,\n",
       "         3.9884e-02,  5.5325e-02,  1.0220e-02, -1.1068e-01,  2.4571e-02,\n",
       "        -2.2973e-02, -8.9612e-03, -3.4463e-02, -7.5151e-02,  5.2567e-02,\n",
       "        -3.0211e-03,  3.4448e-02,  4.4205e-02,  9.0008e-04,  7.6340e-02,\n",
       "        -2.4875e-02,  9.7488e-03, -6.7757e-03,  9.2483e-03, -9.6714e-03,\n",
       "        -2.9754e-02,  5.6767e-03,  6.0532e-02,  7.9551e-02, -2.6374e-02,\n",
       "        -2.7583e-02,  5.6563e-02,  2.7257e-02,  2.8700e-02,  2.5853e-02,\n",
       "        -9.5764e-03, -4.7716e-02,  1.1798e-02, -2.0142e-02,  6.7251e-02,\n",
       "         2.9050e-02, -8.3120e-02,  2.7121e-02, -3.3813e-02,  4.1052e-02,\n",
       "        -5.7674e-02,  2.6587e-02,  6.5756e-03, -5.1586e-02, -3.3088e-03,\n",
       "        -2.4976e-02, -2.5228e-02, -7.7209e-02, -3.6869e-02, -2.9551e-02,\n",
       "         2.2419e-02, -5.6381e-02,  6.3857e-02,  1.0489e-02, -7.2370e-03,\n",
       "         3.2454e-03, -2.0675e-02,  2.3504e-03, -3.6076e-02, -7.8676e-02,\n",
       "        -7.2484e-02, -8.2337e-02,  2.0525e-02,  8.6337e-03, -2.0106e-02,\n",
       "        -2.1932e-02,  1.8636e-02, -4.0498e-02, -7.4945e-02,  5.5821e-02,\n",
       "         5.7636e-02,  3.6370e-03,  1.8318e-02,  5.5383e-02,  6.7273e-02,\n",
       "        -1.3150e-02,  6.1514e-02, -2.9843e-02,  4.0241e-02, -1.2900e-02,\n",
       "         6.5132e-02,  5.5893e-03, -2.4698e-03,  3.7835e-02,  3.9416e-02,\n",
       "         7.4405e-02,  6.3064e-02,  3.8698e-02,  4.8658e-02, -2.8800e-02,\n",
       "         1.2068e-02,  6.7991e-02,  3.2871e-02,  2.7921e-02, -6.2789e-02,\n",
       "        -5.8912e-02, -9.7545e-05, -1.0507e-02, -2.2186e-02, -5.6309e-02,\n",
       "        -2.5463e-02, -2.4143e-02, -4.2778e-03, -4.9434e-02, -6.9642e-02,\n",
       "         7.5532e-02, -2.4537e-02,  6.0245e-02,  6.7415e-02,  2.9664e-02,\n",
       "         6.1473e-02, -1.0005e-02, -3.1118e-02, -3.5331e-02,  7.0360e-02,\n",
       "         7.1862e-02,  4.4525e-02, -3.7827e-02,  1.5406e-02, -2.0642e-02,\n",
       "        -3.4455e-02, -3.5546e-02,  8.4942e-02,  3.0886e-02,  7.0983e-02,\n",
       "         5.1228e-02,  5.8713e-02,  4.1298e-03, -7.5120e-03, -4.0649e-02,\n",
       "         2.3354e-02, -6.4441e-02, -4.0792e-02,  1.3724e-02,  4.8410e-02,\n",
       "        -2.3084e-02, -7.8556e-02, -1.1569e-04, -5.4608e-02,  5.1196e-02,\n",
       "         1.2938e-02,  5.1371e-02,  3.2840e-02,  3.7309e-02,  1.8389e-02,\n",
       "         5.2874e-02, -1.0861e-02, -2.2759e-02, -4.0215e-02,  4.1764e-03,\n",
       "         1.6870e-02, -3.5647e-03,  6.2924e-02,  2.7271e-02, -2.8031e-02,\n",
       "        -2.0380e-02, -2.8520e-02, -4.0986e-02,  3.0001e-02, -8.2022e-02,\n",
       "         3.2906e-03,  4.2091e-02, -4.1267e-02,  3.8844e-02,  2.8671e-02,\n",
       "         7.9457e-02, -3.0538e-02, -1.7499e-02,  1.8030e-02,  5.4201e-02,\n",
       "         1.7295e-02, -4.7460e-02,  9.1519e-02, -9.6873e-03,  1.2353e-03,\n",
       "        -3.6317e-02, -1.3050e-02,  5.8103e-02, -8.9872e-03,  3.4890e-02,\n",
       "         1.8236e-02, -7.4057e-03,  6.4383e-03,  4.0431e-02,  8.3158e-02,\n",
       "         7.4323e-03, -2.9398e-02, -8.5067e-02, -7.8878e-02, -8.0770e-03,\n",
       "        -2.9489e-02, -1.8964e-02,  2.4398e-03,  6.1154e-02, -4.8631e-02,\n",
       "         3.4910e-02,  3.8080e-02, -3.7466e-02,  1.8700e-02,  7.1809e-02,\n",
       "        -3.1765e-02, -3.3811e-02], grad_fn=<DivBackward0>)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res = model_torch(input_ids)\n",
    "res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.allclose(res.detach().numpy(), ref.numpy(), atol=1e-3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Multilungual Compare"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Some texts of different lengths.\n",
    "english_sentences = [\"dog\", \"Puppies are nice.\", \"I enjoy taking long walks along the beach with my dog.\"]\n",
    "italian_sentences = [\"cane\", \"I cuccioli sono carini.\", \"Mi piace fare lunghe passeggiate lungo la spiaggia con il mio cane.\"]\n",
    "japanese_sentences = [\"?\", \"???????\", \"?????????????????????\"]\n",
    "\n",
    "# Compute embeddings.\n",
    "en_result = model(english_sentences).numpy()\n",
    "it_result = model(italian_sentences).numpy()\n",
    "ja_result = model(japanese_sentences).numpy()\n",
    "\n",
    "# Compute similarity matrix. Higher score indicates greater similarity.\n",
    "similarity_matrix_it = np.inner(en_result, it_result)\n",
    "similarity_matrix_ja = np.inner(en_result, ja_result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "en_result_torch = torch.cat([\n",
    "    model_torch(tokenize(english_sentences[0])).unsqueeze(0),\n",
    "    model_torch(tokenize(english_sentences[1])).unsqueeze(0),\n",
    "    model_torch(tokenize(english_sentences[2])).unsqueeze(0),\n",
    "]).detach().numpy()\n",
    "\n",
    "it_result_torch = torch.cat([\n",
    "    model_torch(tokenize(italian_sentences[0])).unsqueeze(0),\n",
    "    model_torch(tokenize(italian_sentences[1])).unsqueeze(0),\n",
    "    model_torch(tokenize(italian_sentences[2])).unsqueeze(0),\n",
    "]).detach().numpy()\n",
    "\n",
    "ja_result_torch = torch.cat([\n",
    "    model_torch(tokenize(japanese_sentences[0])).unsqueeze(0),\n",
    "    model_torch(tokenize(japanese_sentences[1])).unsqueeze(0),\n",
    "    model_torch(tokenize(japanese_sentences[2])).unsqueeze(0),\n",
    "]).detach().numpy()\n",
    "\n",
    "similarity_matrix_it_torch = np.inner(en_result_torch, it_result_torch)\n",
    "similarity_matrix_ja_torch = np.inner(en_result_torch, ja_result_torch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "True\n",
      "True\n",
      "True\n",
      "True\n"
     ]
    }
   ],
   "source": [
    "print(np.allclose(en_result, en_result_torch, atol=1e-3))\n",
    "print(np.allclose(it_result, it_result_torch, atol=1e-3))\n",
    "print(np.allclose(ja_result, ja_result_torch, atol=1e-3))\n",
    "\n",
    "print(np.allclose(similarity_matrix_it, similarity_matrix_it_torch, atol=1e-3))\n",
    "print(np.allclose(similarity_matrix_ja, similarity_matrix_ja_torch, atol=1e-3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Save Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model_torch.state_dict(), 'models/model.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.16 ('venv': venv)",
   "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.9.16"
  },
  "vscode": {
   "interpreter": {
    "hash": "739cef2972ecc74bcd589e39a271e2837c98357d75d84f4c537ab85e1d0e66b4"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}