{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = ''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/husein/alxnet/model_utils.py:334: The name tf.train.Optimizer is deprecated. Please use tf.compat.v1.train.Optimizer instead.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import xlnet\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from tqdm import tqdm\n",
    "import model_utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import sentencepiece as spm\n",
    "from prepro_utils import preprocess_text, encode_ids\n",
    "\n",
    "sp_model = spm.SentencePieceProcessor()\n",
    "sp_model.Load('sp10m.cased.v9.model')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/husein/alxnet/xlnet.py:70: The name tf.gfile.Open is deprecated. Please use tf.io.gfile.GFile instead.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "kwargs = dict(\n",
    "      is_training=True,\n",
    "      use_tpu=False,\n",
    "      use_bfloat16=False,\n",
    "      dropout=0.0,\n",
    "      dropatt=0.0,\n",
    "      init='normal',\n",
    "      init_range=0.1,\n",
    "      init_std=0.05,\n",
    "      clamp_len=-1)\n",
    "\n",
    "xlnet_parameters = xlnet.RunConfig(**kwargs)\n",
    "xlnet_config = xlnet.XLNetConfig(json_path='output-model/config.json')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Model:\n",
    "    def __init__(\n",
    "        self,\n",
    "    ):\n",
    "        self.X = tf.placeholder(tf.int32, [None, None])\n",
    "        self.segment_ids = tf.placeholder(tf.int32, [None, None])\n",
    "        self.input_masks = tf.placeholder(tf.float32, [None, None])\n",
    "        \n",
    "        xlnet_model = xlnet.XLNetModel(\n",
    "            xlnet_config=xlnet_config,\n",
    "            run_config=xlnet_parameters,\n",
    "            input_ids=tf.transpose(self.X, [1, 0]),\n",
    "            seg_ids=tf.transpose(self.segment_ids, [1, 0]),\n",
    "            input_mask=tf.transpose(self.input_masks, [1, 0]))\n",
    "        \n",
    "        output = xlnet_model.get_sequence_output()\n",
    "        lookup_table = xlnet_model.get_embedding_table()\n",
    "        lookup_table_2 = xlnet_model.get_embedding_table2()\n",
    "        \n",
    "        with tf.variable_scope(\"model\", reuse=tf.AUTO_REUSE):\n",
    "            with tf.variable_scope('lm_loss'):\n",
    "                softmax_w = lookup_table\n",
    "                softmax_w = tf.matmul(softmax_w, lookup_table_2)\n",
    "                softmax_b = tf.get_variable('bias', [xlnet_config.n_token], dtype=output.dtype,\n",
    "                                    initializer=tf.zeros_initializer())\n",
    "                logits = tf.einsum('ibd,nd->ibn', output, softmax_w) + softmax_b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:memory input None\n",
      "INFO:tensorflow:Use float type <dtype: 'float32'>\n"
     ]
    }
   ],
   "source": [
    "tf.reset_default_graph()\n",
    "sess = tf.InteractiveSession()\n",
    "model = Model()\n",
    "\n",
    "sess.run(tf.global_variables_initializer())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "import collections\n",
    "import re\n",
    "\n",
    "def get_assignment_map_from_checkpoint(tvars, init_checkpoint):\n",
    "    \"\"\"Compute the union of the current variables and checkpoint variables.\"\"\"\n",
    "    assignment_map = {}\n",
    "    initialized_variable_names = {}\n",
    "\n",
    "    name_to_variable = collections.OrderedDict()\n",
    "    for var in tvars:\n",
    "        name = var.name\n",
    "        m = re.match('^(.*):\\\\d+$', name)\n",
    "        if m is not None:\n",
    "            name = m.group(1)\n",
    "        name_to_variable[name] = var\n",
    "\n",
    "    init_vars = tf.train.list_variables(init_checkpoint)\n",
    "\n",
    "    assignment_map = collections.OrderedDict()\n",
    "    for x in init_vars:\n",
    "        (name, var) = (x[0], x[1])\n",
    "        if name not in name_to_variable:\n",
    "            continue\n",
    "        assignment_map[name] = name_to_variable[name]\n",
    "        initialized_variable_names[name] = 1\n",
    "        initialized_variable_names[name + ':0'] = 1\n",
    "\n",
    "    return (assignment_map, initialized_variable_names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "OrderedDict([('model/transformer/r_w_bias',\n",
       "              <tf.Variable 'model/transformer/r_w_bias:0' shape=(12, 12, 64) dtype=float32_ref>),\n",
       "             ('model/transformer/r_r_bias',\n",
       "              <tf.Variable 'model/transformer/r_r_bias:0' shape=(12, 12, 64) dtype=float32_ref>),\n",
       "             ('model/transformer/word_embedding/lookup_table',\n",
       "              <tf.Variable 'model/transformer/word_embedding/lookup_table:0' shape=(32000, 128) dtype=float32_ref>),\n",
       "             ('model/transformer/word_embedding/lookup_table_2',\n",
       "              <tf.Variable 'model/transformer/word_embedding/lookup_table_2:0' shape=(128, 768) dtype=float32_ref>),\n",
       "             ('model/transformer/r_s_bias',\n",
       "              <tf.Variable 'model/transformer/r_s_bias:0' shape=(12, 12, 64) dtype=float32_ref>),\n",
       "             ('model/transformer/seg_embed',\n",
       "              <tf.Variable 'model/transformer/seg_embed:0' shape=(12, 2, 12, 64) dtype=float32_ref>),\n",
       "             ('model/transformer/layer_shared/rel_attn/q/kernel',\n",
       "              <tf.Variable 'model/transformer/layer_shared/rel_attn/q/kernel:0' shape=(768, 12, 64) dtype=float32_ref>),\n",
       "             ('model/transformer/layer_shared/rel_attn/k/kernel',\n",
       "              <tf.Variable 'model/transformer/layer_shared/rel_attn/k/kernel:0' shape=(768, 12, 64) dtype=float32_ref>),\n",
       "             ('model/transformer/layer_shared/rel_attn/v/kernel',\n",
       "              <tf.Variable 'model/transformer/layer_shared/rel_attn/v/kernel:0' shape=(768, 12, 64) dtype=float32_ref>),\n",
       "             ('model/transformer/layer_shared/rel_attn/r/kernel',\n",
       "              <tf.Variable 'model/transformer/layer_shared/rel_attn/r/kernel:0' shape=(768, 12, 64) dtype=float32_ref>),\n",
       "             ('model/transformer/layer_shared/rel_attn/o/kernel',\n",
       "              <tf.Variable 'model/transformer/layer_shared/rel_attn/o/kernel:0' shape=(768, 12, 64) dtype=float32_ref>),\n",
       "             ('model/transformer/layer_shared/rel_attn/LayerNorm/beta',\n",
       "              <tf.Variable 'model/transformer/layer_shared/rel_attn/LayerNorm/beta:0' shape=(768,) dtype=float32_ref>),\n",
       "             ('model/transformer/layer_shared/rel_attn/LayerNorm/gamma',\n",
       "              <tf.Variable 'model/transformer/layer_shared/rel_attn/LayerNorm/gamma:0' shape=(768,) dtype=float32_ref>),\n",
       "             ('model/transformer/layer_shared/ff/layer_1/kernel',\n",
       "              <tf.Variable 'model/transformer/layer_shared/ff/layer_1/kernel:0' shape=(768, 3072) dtype=float32_ref>),\n",
       "             ('model/transformer/layer_shared/ff/layer_1/bias',\n",
       "              <tf.Variable 'model/transformer/layer_shared/ff/layer_1/bias:0' shape=(3072,) dtype=float32_ref>),\n",
       "             ('model/transformer/layer_shared/ff/layer_2/kernel',\n",
       "              <tf.Variable 'model/transformer/layer_shared/ff/layer_2/kernel:0' shape=(3072, 768) dtype=float32_ref>),\n",
       "             ('model/transformer/layer_shared/ff/layer_2/bias',\n",
       "              <tf.Variable 'model/transformer/layer_shared/ff/layer_2/bias:0' shape=(768,) dtype=float32_ref>),\n",
       "             ('model/transformer/layer_shared/ff/LayerNorm/beta',\n",
       "              <tf.Variable 'model/transformer/layer_shared/ff/LayerNorm/beta:0' shape=(768,) dtype=float32_ref>),\n",
       "             ('model/transformer/layer_shared/ff/LayerNorm/gamma',\n",
       "              <tf.Variable 'model/transformer/layer_shared/ff/LayerNorm/gamma:0' shape=(768,) dtype=float32_ref>),\n",
       "             ('model/lm_loss/bias',\n",
       "              <tf.Variable 'model/lm_loss/bias:0' shape=(32000,) dtype=float32_ref>)])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tvars = tf.trainable_variables()\n",
    "\n",
    "assignment_map = {}\n",
    "initialized_variable_names = {}\n",
    "\n",
    "name_to_variable = collections.OrderedDict()\n",
    "for var in tvars:\n",
    "    name = var.name\n",
    "    m = re.match('^(.*):\\\\d+$', name)\n",
    "    if m is not None:\n",
    "        name = m.group(1)\n",
    "    name_to_variable[name] = var\n",
    "    \n",
    "name_to_variable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = 'output-model/model.ckpt-300000'\n",
    "assignment_map, initialized_variable_names = get_assignment_map_from_checkpoint(tvars, \n",
    "                                                                                checkpoint)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Restoring parameters from output-model/model.ckpt-300000\n"
     ]
    }
   ],
   "source": [
    "saver = tf.train.Saver(var_list = assignment_map)\n",
    "saver.restore(sess, checkpoint)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'alxlnet-base/model.ckpt'"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "saver = tf.train.Saver(tf.trainable_variables())\n",
    "saver.save(sess, 'alxlnet-base/model.ckpt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "alxlnet-base/\n",
      "alxlnet-base/sp10m.cased.v9.vocab\n",
      "alxlnet-base/sp10m.cased.v9.model\n",
      "alxlnet-base/model.ckpt.index\n",
      "alxlnet-base/config.json\n",
      "alxlnet-base/model.ckpt.data-00000-of-00001\n",
      "alxlnet-base/checkpoint\n",
      "alxlnet-base/model.ckpt.meta\n"
     ]
    }
   ],
   "source": [
    "!cp output-model/config.json alxlnet-base/config.json\n",
    "!cp sp10m.cased.v9* alxlnet-base\n",
    "!tar cvzf alxlnet-base.tar.gz alxlnet-base"
   ]
  },
  {
   "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
