{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !wget https://huseinhouse-storage.s3-ap-southeast-1.amazonaws.com/bert-bahasa/electra-bahasa-base-27-04-2020.tar.gz\n",
    "# !tar -zxf electra-bahasa-base-27-04-2020.tar.gz\n",
    "# !wget https://raw.githubusercontent.com/huseinzol05/Malaya/master/pretrained-model/electra/BASE-config.json\n",
    "# !wget https://raw.githubusercontent.com/huseinzol05/Malaya/master/pretrained-model/electra/BASE-config-generator.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = ''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from model import modeling\n",
    "from model import optimization\n",
    "from configure_pretraining import PretrainingConfig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'vocab_size': 32000,\n",
       " 'hidden_size': 256,\n",
       " 'num_hidden_layers': 12,\n",
       " 'num_attention_heads': 4,\n",
       " 'hidden_act': 'gelu',\n",
       " 'intermediate_size': 1024,\n",
       " 'hidden_dropout_prob': 0.1,\n",
       " 'attention_probs_dropout_prob': 0.1,\n",
       " 'max_position_embeddings': 512,\n",
       " 'type_vocab_size': 2,\n",
       " 'initializer_range': 0.02,\n",
       " 'embedding_size': 768,\n",
       " 'layer_norm_eps': 1e-12}"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bert_config = modeling.BertConfig.from_json_file('BASE-config-generator.json')\n",
    "bert_config.__dict__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from model import tokenization\n",
    "\n",
    "tokenizer = tokenization.FullTokenizer(vocab_file='out/bahasa.wordpiece')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gather_positions(sequence, positions):\n",
    "    shape = modeling.get_shape_list(sequence, expected_rank = [2, 3])\n",
    "    depth_dimension = len(shape) == 3\n",
    "    if depth_dimension:\n",
    "        B, L, D = shape\n",
    "    else:\n",
    "        B, L = shape\n",
    "        D = 1\n",
    "        sequence = tf.expand_dims(sequence, -1)\n",
    "    position_shift = tf.expand_dims(L * tf.range(B), -1)\n",
    "    flat_positions = tf.reshape(positions + position_shift, [-1])\n",
    "    flat_sequence = tf.reshape(sequence, [B * L, D])\n",
    "    gathered = tf.gather(flat_sequence, flat_positions)\n",
    "    if depth_dimension:\n",
    "        return tf.reshape(gathered, [B, -1, D])\n",
    "    else:\n",
    "        return tf.reshape(gathered, [B, -1])\n",
    "\n",
    "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.int32, [None, None])\n",
    "        self.Y = tf.placeholder(tf.int32, [None])\n",
    "        \n",
    "        model = modeling.BertModel(\n",
    "            bert_config=bert_config,\n",
    "            is_training=False,\n",
    "            input_ids=self.X,\n",
    "            input_mask=self.input_masks,\n",
    "            token_type_ids=self.segment_ids,\n",
    "            use_one_hot_embeddings=False,\n",
    "            scope='generator',\n",
    "            embedding_size = bert_config.embedding_size)\n",
    "        \n",
    "        self.logits = model.get_sequence_output()\n",
    "        \n",
    "        with tf.variable_scope(\"generator_predictions\"):\n",
    "            relevant_hidden = gather_positions(\n",
    "                self.logits, self.input_masks)\n",
    "            hidden = tf.layers.dense(\n",
    "                relevant_hidden,\n",
    "                units=modeling.get_shape_list(model.get_embedding_table())[-1],\n",
    "                activation=modeling.get_activation(bert_config.hidden_act),\n",
    "                kernel_initializer=modeling.create_initializer(\n",
    "                    bert_config.initializer_range))\n",
    "            hidden = modeling.layer_norm(hidden)\n",
    "            output_bias = tf.get_variable(\n",
    "            \"output_bias\",\n",
    "            shape=[bert_config.vocab_size],\n",
    "            initializer=tf.zeros_initializer())\n",
    "            logits = tf.matmul(hidden, model.get_embedding_table(),\n",
    "                               transpose_b=True)\n",
    "            self._logits = tf.nn.bias_add(logits, output_bias)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/husein/electra/electra/model/modeling.py:226: dense (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use keras.layers.Dense instead.\n",
      "WARNING:tensorflow:From /home/husein/.local/lib/python3.6/site-packages/tensorflow_core/python/layers/core.py:187: Layer.apply (from tensorflow.python.keras.engine.base_layer) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Please use `layer.__call__` method instead.\n"
     ]
    }
   ],
   "source": [
    "tf.reset_default_graph()\n",
    "sess = tf.InteractiveSession()\n",
    "model = Model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Restoring parameters from out/model.ckpt-728800\n"
     ]
    }
   ],
   "source": [
    "sess.run(tf.global_variables_initializer())\n",
    "var_lists = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope = 'generator')\n",
    "electra = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope = 'electra')\n",
    "saver = tf.train.Saver(var_list = var_lists + electra)\n",
    "saver.restore(sess, 'out/model.ckpt-728800')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "attentions = [\n",
    "    n.name\n",
    "    for n in tf.get_default_graph().as_graph_def().node\n",
    "    if 'Softmax' in n.name\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['generator/encoder/layer_0/attention/self/Softmax',\n",
       " 'generator/encoder/layer_1/attention/self/Softmax',\n",
       " 'generator/encoder/layer_2/attention/self/Softmax',\n",
       " 'generator/encoder/layer_3/attention/self/Softmax',\n",
       " 'generator/encoder/layer_4/attention/self/Softmax',\n",
       " 'generator/encoder/layer_5/attention/self/Softmax',\n",
       " 'generator/encoder/layer_6/attention/self/Softmax',\n",
       " 'generator/encoder/layer_7/attention/self/Softmax',\n",
       " 'generator/encoder/layer_8/attention/self/Softmax',\n",
       " 'generator/encoder/layer_9/attention/self/Softmax',\n",
       " 'generator/encoder/layer_10/attention/self/Softmax',\n",
       " 'generator/encoder/layer_11/attention/self/Softmax']"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "attentions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'electra-base/model.ckpt'"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "saver = tf.train.Saver(tf.trainable_variables())\n",
    "saver.save(sess, 'electra-base/model.ckpt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "electra-base/\n",
      "electra-base/model.ckpt.index\n",
      "electra-base/config.json\n",
      "electra-base/model.ckpt.data-00000-of-00001\n",
      "electra-base/bahasa.wordpiece\n",
      "electra-base/checkpoint\n",
      "electra-base/model.ckpt.meta\n"
     ]
    }
   ],
   "source": [
    "!cp BASE-config-generator.json electra-base/config.json\n",
    "!cp out/bahasa.wordpiece electra-base/bahasa.wordpiece\n",
    "!tar cvzf electra-base.tar.gz electra-base"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "\n",
    "bucketName = 'huseinhouse-storage'\n",
    "Key = 'electra-base.tar.gz'\n",
    "outPutname = \"v34/pretrained-model/electra-base.tar.gz\"\n",
    "\n",
    "s3 = boto3.client('s3')\n",
    "s3.upload_file(Key,bucketName,outPutname)"
   ]
  }
 ],
 "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
}
