{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '2'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('../Malaya-Dataset/dependency/gsd-ud-train.conllu.txt') as fopen:\n",
    "    corpus = fopen.read().split('\\n')\n",
    "    \n",
    "with open('../Malaya-Dataset/dependency/gsd-ud-test.conllu.txt') as fopen:\n",
    "    corpus.extend(fopen.read().split('\\n'))\n",
    "    \n",
    "with open('../Malaya-Dataset/dependency/gsd-ud-dev.conllu.txt') as fopen:\n",
    "    corpus.extend(fopen.read().split('\\n'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/husein/xlnet/model_utils.py:295: 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\n",
    "import pickle\n",
    "import json\n",
    "pad_sequences = tf.keras.preprocessing.sequence.pad_sequences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sentencepiece as spm\n",
    "from prepro_utils import preprocess_text, encode_ids\n",
    "\n",
    "sp_model = spm.SentencePieceProcessor()\n",
    "sp_model.Load('xlnet-base/sp10m.cased.v9.model')\n",
    "\n",
    "def tokenize_fn(text):\n",
    "    text = preprocess_text(text, lower= False)\n",
    "    return encode_ids(sp_model, text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "SEG_ID_A   = 0\n",
    "SEG_ID_B   = 1\n",
    "SEG_ID_CLS = 2\n",
    "SEG_ID_SEP = 3\n",
    "SEG_ID_PAD = 4\n",
    "\n",
    "special_symbols = {\n",
    "    \"<unk>\"  : 0,\n",
    "    \"<s>\"    : 1,\n",
    "    \"</s>\"   : 2,\n",
    "    \"<cls>\"  : 3,\n",
    "    \"<sep>\"  : 4,\n",
    "    \"<pad>\"  : 5,\n",
    "    \"<mask>\" : 6,\n",
    "    \"<eod>\"  : 7,\n",
    "    \"<eop>\"  : 8,\n",
    "}\n",
    "\n",
    "VOCAB_SIZE = 32000\n",
    "UNK_ID = special_symbols[\"<unk>\"]\n",
    "CLS_ID = special_symbols[\"<cls>\"]\n",
    "SEP_ID = special_symbols[\"<sep>\"]\n",
    "MASK_ID = special_symbols[\"<mask>\"]\n",
    "EOD_ID = special_symbols[\"<eod>\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "tag2idx = {'PAD': 0, 'X': 1}\n",
    "tag_idx = 2\n",
    "\n",
    "def process_corpus(corpus, until = None):\n",
    "    global word2idx, tag2idx, char2idx, word_idx, tag_idx, char_idx\n",
    "    sentences, words, depends, labels, pos, sequences = [], [], [], [], [], []\n",
    "    temp_sentence, temp_word, temp_depend, temp_label, temp_pos = [], [], [], [], []\n",
    "    segments, masks = [], []\n",
    "    first_time = True\n",
    "    for sentence in corpus:\n",
    "        try:\n",
    "            if len(sentence):\n",
    "                if sentence[0] == '#':\n",
    "                    continue\n",
    "                if first_time:\n",
    "                    print(sentence)\n",
    "                    first_time = False\n",
    "                sentence = sentence.split('\\t')\n",
    "                if sentence[7] not in tag2idx:\n",
    "                    tag2idx[sentence[7]] = tag_idx\n",
    "                    tag_idx += 1\n",
    "                temp_word.append(sentence[1])\n",
    "                temp_depend.append(int(sentence[6]) + 1)\n",
    "                temp_label.append(tag2idx[sentence[7]])\n",
    "                temp_sentence.append(sentence[1])\n",
    "                temp_pos.append(sentence[3])\n",
    "            else:\n",
    "                if len(temp_sentence) < 2 or len(temp_word) != len(temp_label):\n",
    "                    temp_word = []\n",
    "                    temp_depend = []\n",
    "                    temp_label = []\n",
    "                    temp_sentence = []\n",
    "                    temp_pos = []\n",
    "                    continue\n",
    "                bert_tokens = []\n",
    "                labels_ = []\n",
    "                depends_ = []\n",
    "                seq_ = []\n",
    "                for no, orig_token in enumerate(temp_word):\n",
    "                    labels_.append(temp_label[no])\n",
    "                    depends_.append(temp_depend[no])\n",
    "                    t = tokenize_fn(orig_token)\n",
    "                    bert_tokens.extend(t)\n",
    "                    labels_.extend([1] * (len(t) - 1))\n",
    "                    depends_.extend([0] * (len(t) - 1))\n",
    "                    seq_.append(no + 1)\n",
    "                bert_tokens.extend([4, 3])\n",
    "                labels_.extend([0, 0])\n",
    "                depends_.extend([0, 0])\n",
    "                segment = [0] * (len(bert_tokens) - 1) + [SEG_ID_CLS]\n",
    "                input_mask = [0] * len(segment)\n",
    "                words.append(bert_tokens)\n",
    "                depends.append(depends_)\n",
    "                labels.append(labels_)\n",
    "                sentences.append(bert_tokens)\n",
    "                pos.append(temp_pos)\n",
    "                sequences.append(seq_)\n",
    "                segments.append(segment)\n",
    "                masks.append(input_mask)\n",
    "                temp_word = []\n",
    "                temp_depend = []\n",
    "                temp_label = []\n",
    "                temp_sentence = []\n",
    "                temp_pos = []\n",
    "        except Exception as e:\n",
    "            print(e, sentence)\n",
    "    return sentences[:-1], words[:-1], depends[:-1], labels[:-1], pos[:-1], sequences[:-1], segments[:-1], masks[:-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1\tSembungan\tsembungan\tPROPN\tX--\t_\t4\tnsubj\t_\tMorphInd=^sembungan<x>_X--$\n"
     ]
    }
   ],
   "source": [
    "sentences, words, depends, labels, _, _, segments, masks = process_corpus(corpus)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(26, 26, 26)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(words[0]), len(depends[0]), len(labels[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "with open('../Malaya-Dataset/dependency/augmented-dependency.json') as fopen:\n",
    "    augmented = json.load(fopen)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "text_augmented, depends_augmented, labels_augmented = [], [], []\n",
    "\n",
    "for a in augmented:\n",
    "    text_augmented.extend(a[0])\n",
    "    depends_augmented.extend(a[1])\n",
    "    labels_augmented.extend((np.array(a[2]) + 1).tolist())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_XY(texts, depends, labels):\n",
    "    outside, sentences, outside_depends, outside_labels = [], [], [], []\n",
    "    segments, masks = [], []\n",
    "    for no, text in enumerate(texts):\n",
    "        temp_depend = depends[no]\n",
    "        temp_label = labels[no]\n",
    "        s = text.split()\n",
    "        sentences.append(s)\n",
    "        bert_tokens = []\n",
    "        labels_ = []\n",
    "        depends_ = []\n",
    "        for no, orig_token in enumerate(s):\n",
    "            labels_.append(temp_label[no])\n",
    "            depends_.append(temp_depend[no])\n",
    "            t = tokenize_fn(orig_token)\n",
    "            bert_tokens.extend(t)\n",
    "            labels_.extend([1] * (len(t) - 1))\n",
    "            depends_.extend([0] * (len(t) - 1))\n",
    "        bert_tokens.extend([4, 3])\n",
    "        labels_.extend([0, 0])\n",
    "        depends_.extend([0, 0])\n",
    "        segment = [0] * (len(bert_tokens) - 1) + [SEG_ID_CLS]\n",
    "        input_mask = [0] * len(segment)\n",
    "        outside.append(bert_tokens)\n",
    "        outside_depends.append(depends_)\n",
    "        outside_labels.append(labels_)\n",
    "        segments.append(segment)\n",
    "        masks.append(input_mask)\n",
    "    return outside, sentences, outside_depends, outside_labels, segments, masks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "outside, _, outside_depends, outside_labels, outside_segments, outside_masks = parse_XY(text_augmented, \n",
    "                                                       depends_augmented, \n",
    "                                                       labels_augmented)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "words.extend(outside)\n",
    "depends.extend(outside_depends)\n",
    "labels.extend(outside_labels)\n",
    "segments.extend(outside_segments)\n",
    "masks.extend(outside_masks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{0: 'PAD',\n",
       " 1: 'X',\n",
       " 2: 'nsubj',\n",
       " 3: 'cop',\n",
       " 4: 'det',\n",
       " 5: 'root',\n",
       " 6: 'nsubj:pass',\n",
       " 7: 'acl',\n",
       " 8: 'case',\n",
       " 9: 'obl',\n",
       " 10: 'flat',\n",
       " 11: 'punct',\n",
       " 12: 'appos',\n",
       " 13: 'amod',\n",
       " 14: 'compound',\n",
       " 15: 'advmod',\n",
       " 16: 'cc',\n",
       " 17: 'obj',\n",
       " 18: 'conj',\n",
       " 19: 'mark',\n",
       " 20: 'advcl',\n",
       " 21: 'nmod',\n",
       " 22: 'nummod',\n",
       " 23: 'dep',\n",
       " 24: 'xcomp',\n",
       " 25: 'ccomp',\n",
       " 26: 'parataxis',\n",
       " 27: 'compound:plur',\n",
       " 28: 'fixed',\n",
       " 29: 'aux',\n",
       " 30: 'csubj',\n",
       " 31: 'iobj',\n",
       " 32: 'csubj:pass'}"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "idx2tag = {v:k for k, v in tag2idx.items()}\n",
    "idx2tag"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "words_train, words_test, depends_train, depends_test, labels_train, labels_test, \\\n",
    "segments_train, segments_test, masks_train, masks_test \\\n",
    "= train_test_split(words, depends, labels, segments, masks, test_size = 0.2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(40289, 10073)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(words_train), len(words_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_X = words_train\n",
    "train_Y = labels_train\n",
    "train_depends = depends_train\n",
    "\n",
    "test_X = words_test\n",
    "test_Y = labels_test\n",
    "test_depends = depends_test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/husein/xlnet/xlnet.py:63: The name tf.gfile.Open is deprecated. Please use tf.io.gfile.GFile instead.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import xlnet\n",
    "import model_utils\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "\n",
    "kwargs = dict(\n",
    "      is_training=True,\n",
    "      use_tpu=False,\n",
    "      use_bfloat16=False,\n",
    "      dropout=0.1,\n",
    "      dropatt=0.1,\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='xlnet-base/config.json')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37770 3777\n"
     ]
    }
   ],
   "source": [
    "epoch = 15\n",
    "batch_size = 16\n",
    "warmup_proportion = 0.1\n",
    "num_train_steps = int(len(train_X) / batch_size * epoch)\n",
    "num_warmup_steps = int(num_train_steps * warmup_proportion)\n",
    "print(num_train_steps, num_warmup_steps)\n",
    "\n",
    "training_parameters = dict(\n",
    "      decay_method = 'poly',\n",
    "      train_steps = num_train_steps,\n",
    "      learning_rate = 2e-5,\n",
    "      warmup_steps = num_warmup_steps,\n",
    "      min_lr_ratio = 0.0,\n",
    "      weight_decay = 0.00,\n",
    "      adam_epsilon = 1e-8,\n",
    "      num_core_per_host = 1,\n",
    "      lr_layer_decay_rate = 1,\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.02,\n",
    "      clip = 1.0,\n",
    "      clamp_len=-1,)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Parameter:\n",
    "    def __init__(self, decay_method, warmup_steps, weight_decay, adam_epsilon, \n",
    "                num_core_per_host, lr_layer_decay_rate, use_tpu, learning_rate, train_steps,\n",
    "                min_lr_ratio, clip, **kwargs):\n",
    "        self.decay_method = decay_method\n",
    "        self.warmup_steps = warmup_steps\n",
    "        self.weight_decay = weight_decay\n",
    "        self.adam_epsilon = adam_epsilon\n",
    "        self.num_core_per_host = num_core_per_host\n",
    "        self.lr_layer_decay_rate = lr_layer_decay_rate\n",
    "        self.use_tpu = use_tpu\n",
    "        self.learning_rate = learning_rate\n",
    "        self.train_steps = train_steps\n",
    "        self.min_lr_ratio = min_lr_ratio\n",
    "        self.clip = clip\n",
    "        \n",
    "training_parameters = Parameter(**training_parameters)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BiAAttention:\n",
    "    def __init__(self, input_size_encoder, input_size_decoder, num_labels):\n",
    "        self.input_size_encoder = input_size_encoder\n",
    "        self.input_size_decoder = input_size_decoder\n",
    "        self.num_labels = num_labels\n",
    "        \n",
    "        self.W_d = tf.get_variable(\"W_d\", shape=[self.num_labels, self.input_size_decoder],\n",
    "           initializer=tf.contrib.layers.xavier_initializer())\n",
    "        self.W_e = tf.get_variable(\"W_e\", shape=[self.num_labels, self.input_size_encoder],\n",
    "           initializer=tf.contrib.layers.xavier_initializer())\n",
    "        self.U = tf.get_variable(\"U\", shape=[self.num_labels, self.input_size_decoder, self.input_size_encoder],\n",
    "           initializer=tf.contrib.layers.xavier_initializer())\n",
    "        \n",
    "    def forward(self, input_d, input_e, mask_d=None, mask_e=None):\n",
    "        batch = tf.shape(input_d)[0]\n",
    "        length_decoder = tf.shape(input_d)[1]\n",
    "        length_encoder = tf.shape(input_e)[1]\n",
    "        out_d = tf.expand_dims(tf.matmul(self.W_d, tf.transpose(input_d, [0, 2, 1])), 3)\n",
    "        out_e = tf.expand_dims(tf.matmul(self.W_e, tf.transpose(input_e, [0, 2, 1])), 2)\n",
    "        output = tf.matmul(tf.expand_dims(input_d, 1), self.U)\n",
    "        output = tf.matmul(output, tf.transpose(tf.expand_dims(input_e, 1), [0, 1, 3, 2]))\n",
    "        \n",
    "        output = output + out_d + out_e\n",
    "        \n",
    "        if mask_d is not None:\n",
    "            d = tf.expand_dims(tf.expand_dims(mask_d, 1), 3)\n",
    "            e = tf.expand_dims(tf.expand_dims(mask_e, 1), 2)\n",
    "            output = output * d * e\n",
    "            \n",
    "        return output\n",
    "    \n",
    "class BiLinear:\n",
    "    def __init__(self, left_features, right_features, out_features):\n",
    "        self.left_features = left_features\n",
    "        self.right_features = right_features\n",
    "        self.out_features = out_features\n",
    "        \n",
    "        self.U = tf.get_variable(\"U-bi\", shape=[out_features, left_features, right_features],\n",
    "           initializer=tf.contrib.layers.xavier_initializer())\n",
    "        self.W_l = tf.get_variable(\"Wl\", shape=[out_features, left_features],\n",
    "           initializer=tf.contrib.layers.xavier_initializer())\n",
    "        self.W_r = tf.get_variable(\"Wr\", shape=[out_features, right_features],\n",
    "           initializer=tf.contrib.layers.xavier_initializer())\n",
    "    \n",
    "    def forward(self, input_left, input_right):\n",
    "        left_size = tf.shape(input_left)\n",
    "        output_shape = tf.concat([left_size[:-1], [self.out_features]], axis = 0)\n",
    "        batch = tf.cast(tf.reduce_prod(left_size[:-1]), tf.int32)\n",
    "        input_left = tf.reshape(input_left, (batch, self.left_features))\n",
    "        input_right = tf.reshape(input_right, (batch, self.right_features))\n",
    "        tiled = tf.tile(tf.expand_dims(input_left, axis = 0), (self.out_features,1,1))\n",
    "        output = tf.transpose(tf.reduce_sum(tf.matmul(tiled, self.U), axis = 2))\n",
    "        output = output + tf.matmul(input_left, tf.transpose(self.W_l))\\\n",
    "        + tf.matmul(input_right, tf.transpose(self.W_r))\n",
    "        \n",
    "        return tf.reshape(output, output_shape)\n",
    "\n",
    "\n",
    "    \n",
    "    def decode(self, input_word, input_char, mask, leading_symbolic=0):\n",
    "        out_arc, out_type, _ = self.forward(input_word, input_char, mask)\n",
    "        batch = tf.shape(out_arc)[0]\n",
    "        max_len = tf.shape(out_arc)[1]\n",
    "        sec_max_len = tf.shape(out_arc)[2]\n",
    "        out_arc = out_arc + tf.linalg.diag(tf.fill([max_len], -np.inf))\n",
    "        minus_mask = tf.expand_dims(tf.cast(1 - mask, tf.bool), axis = 2)\n",
    "        minus_mask = tf.tile(minus_mask, [1, 1, sec_max_len])\n",
    "        out_arc = tf.where(minus_mask, tf.fill(tf.shape(out_arc), -np.inf), out_arc)\n",
    "        heads = tf.argmax(out_arc, axis = 1)\n",
    "        type_h, type_c = out_type\n",
    "        batch = tf.shape(type_h)[0]\n",
    "        max_len = tf.shape(type_h)[1]\n",
    "        batch_index = tf.range(0, batch)\n",
    "        t = tf.cast(tf.transpose(heads), tf.int32)\n",
    "        broadcasted = tf.broadcast_to(batch_index, tf.shape(t))\n",
    "        concatenated = tf.transpose(tf.concat([tf.expand_dims(broadcasted, axis = 0), \n",
    "                                               tf.expand_dims(t, axis = 0)], axis = 0))\n",
    "        type_h = tf.gather_nd(type_h, concatenated)\n",
    "        out_type = self.bilinear.forward(type_h, type_c)\n",
    "        out_type = out_type[:, :, leading_symbolic:]\n",
    "        types = tf.argmax(out_type, axis = 2)\n",
    "        return heads, types\n",
    "    \n",
    "class Model:\n",
    "    def __init__(\n",
    "        self,\n",
    "        learning_rate,\n",
    "        hidden_size_word,\n",
    "        cov = 0.0):\n",
    "        \n",
    "        self.words = 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",
    "        self.heads = tf.placeholder(tf.int32, (None, None))\n",
    "        self.types = tf.placeholder(tf.int32, (None, None))\n",
    "        self.switch = tf.placeholder(tf.bool, None)\n",
    "        self.mask = tf.cast(tf.math.not_equal(self.words, 0), tf.float32)\n",
    "        self.maxlen = tf.shape(self.words)[1]\n",
    "        self.lengths = tf.count_nonzero(self.words, 1)\n",
    "        mask = self.mask\n",
    "        heads = self.heads\n",
    "        types = self.types\n",
    "        \n",
    "        self.arc_h = tf.layers.Dense(hidden_size_word)\n",
    "        self.arc_c = tf.layers.Dense(hidden_size_word)\n",
    "        self.attention = BiAAttention(hidden_size_word, hidden_size_word, 1)\n",
    "\n",
    "        self.type_h = tf.layers.Dense(hidden_size_word)\n",
    "        self.type_c = tf.layers.Dense(hidden_size_word)\n",
    "        self.bilinear = BiLinear(hidden_size_word, hidden_size_word, len(tag2idx))\n",
    "        \n",
    "        xlnet_model = xlnet.XLNetModel(\n",
    "            xlnet_config=xlnet_config,\n",
    "            run_config=xlnet_parameters,\n",
    "            input_ids=tf.transpose(self.words, [1, 0]),\n",
    "            seg_ids=tf.transpose(self.segment_ids, [1, 0]),\n",
    "            input_mask=tf.transpose(self.input_masks, [1, 0]))\n",
    "        output_layer = xlnet_model.get_sequence_output()\n",
    "        output_layer = tf.transpose(output_layer, [1, 0, 2])\n",
    "        \n",
    "        arc_h = tf.nn.elu(self.arc_h(output_layer))\n",
    "        arc_c = tf.nn.elu(self.arc_c(output_layer))\n",
    "        \n",
    "        type_h = tf.nn.elu(self.type_h(output_layer))\n",
    "        type_c = tf.nn.elu(self.type_c(output_layer))\n",
    "        \n",
    "        out_arc = tf.squeeze(self.attention.forward(arc_h, arc_c, mask_d=self.mask, \n",
    "                                                    mask_e=self.mask), axis = 1)\n",
    "        \n",
    "        batch = tf.shape(out_arc)[0]\n",
    "        max_len = tf.shape(out_arc)[1]\n",
    "        sec_max_len = tf.shape(out_arc)[2]\n",
    "        batch_index = tf.range(0, batch)\n",
    "        \n",
    "        decode_arc = out_arc + tf.linalg.diag(tf.fill([max_len], -np.inf))\n",
    "        minus_mask = tf.expand_dims(tf.cast(1 - mask, tf.bool), axis = 2)\n",
    "        minus_mask = tf.tile(minus_mask, [1, 1, sec_max_len])\n",
    "        decode_arc = tf.where(minus_mask, tf.fill(tf.shape(decode_arc), -np.inf), decode_arc)\n",
    "        self.heads_seq = tf.argmax(decode_arc, axis = 1)\n",
    "        self.heads_seq = tf.identity(self.heads_seq, name = 'heads_seq')\n",
    "        \n",
    "        t = tf.cast(tf.transpose(self.heads_seq), tf.int32)\n",
    "        broadcasted = tf.broadcast_to(batch_index, tf.shape(t))\n",
    "        concatenated = tf.transpose(tf.concat([tf.expand_dims(broadcasted, axis = 0), \n",
    "                                               tf.expand_dims(t, axis = 0)], axis = 0))\n",
    "        type_h = tf.gather_nd(type_h, concatenated)\n",
    "        out_type = self.bilinear.forward(type_h, type_c)\n",
    "        self.tags_seq = tf.argmax(out_type, axis = 2)\n",
    "        self.tags_seq = tf.identity(self.tags_seq, name = 'tags_seq')\n",
    "        \n",
    "        log_likelihood, transition_params = tf.contrib.crf.crf_log_likelihood(\n",
    "            out_type, self.types, self.lengths\n",
    "        )\n",
    "        crf_loss = tf.reduce_mean(-log_likelihood)\n",
    "        self.logits, _ = tf.contrib.crf.crf_decode(\n",
    "            out_type, transition_params, self.lengths\n",
    "        )\n",
    "        self.logits = tf.identity(self.logits, name = 'logits')\n",
    "        \n",
    "        batch = tf.shape(out_arc)[0]\n",
    "        max_len = tf.shape(out_arc)[1]\n",
    "        batch_index = tf.range(0, batch)\n",
    "        t = tf.transpose(heads)\n",
    "        broadcasted = tf.broadcast_to(batch_index, tf.shape(t))\n",
    "        concatenated = tf.transpose(tf.concat([tf.expand_dims(broadcasted, axis = 0), \n",
    "                                               tf.expand_dims(t, axis = 0)], axis = 0))\n",
    "        type_h = tf.gather_nd(type_h, concatenated)\n",
    "        out_type = self.bilinear.forward(type_h, type_c)\n",
    "        minus_inf = -1e8\n",
    "        minus_mask = (1 - mask) * minus_inf\n",
    "        out_arc = out_arc + tf.expand_dims(minus_mask, axis = 2) + tf.expand_dims(minus_mask, axis = 1)\n",
    "        loss_arc = tf.nn.log_softmax(out_arc, dim=1)\n",
    "        loss_type = tf.nn.log_softmax(out_type, dim=2)\n",
    "        loss_arc = loss_arc * tf.expand_dims(mask, axis = 2) * tf.expand_dims(mask, axis = 1)\n",
    "        loss_type = loss_type * tf.expand_dims(mask, axis = 2)\n",
    "        num = tf.reduce_sum(mask) - tf.cast(batch, tf.float32)\n",
    "        child_index = tf.tile(tf.expand_dims(tf.range(0, max_len), 1), [1, batch])\n",
    "        t = tf.transpose(heads)\n",
    "        broadcasted = tf.broadcast_to(batch_index, tf.shape(t))\n",
    "        concatenated = tf.transpose(tf.concat([tf.expand_dims(broadcasted, axis = 0),\n",
    "                                               tf.expand_dims(t, axis = 0),\n",
    "                                               tf.expand_dims(child_index, axis = 0)], axis = 0))\n",
    "        loss_arc = tf.gather_nd(loss_arc, concatenated)\n",
    "        loss_arc = tf.transpose(loss_arc, [1, 0])\n",
    "        \n",
    "        t = tf.transpose(types)\n",
    "        broadcasted = tf.broadcast_to(batch_index, tf.shape(t))\n",
    "        concatenated = tf.transpose(tf.concat([tf.expand_dims(broadcasted, axis = 0),\n",
    "                                               tf.expand_dims(child_index, axis = 0),\n",
    "                                               tf.expand_dims(t, axis = 0)], axis = 0))\n",
    "        loss_type = tf.gather_nd(loss_type, concatenated)\n",
    "        loss_type = tf.transpose(loss_type, [1, 0])\n",
    "        cost = (tf.reduce_sum(-loss_arc) / num) + (tf.reduce_sum(-loss_type) / num)\n",
    "        \n",
    "        self.cost = tf.cond(self.switch, lambda: cost + crf_loss, lambda: cost)\n",
    "        self.optimizer = tf.train.AdamOptimizer(\n",
    "            learning_rate = learning_rate\n",
    "        ).minimize(self.cost)\n",
    "        \n",
    "        mask = tf.sequence_mask(self.lengths, maxlen = self.maxlen)\n",
    "        \n",
    "        self.prediction = tf.boolean_mask(self.logits, mask)\n",
    "        mask_label = tf.boolean_mask(self.types, mask)\n",
    "        correct_pred = tf.equal(tf.cast(self.prediction, tf.int32), mask_label)\n",
    "        correct_index = tf.cast(correct_pred, tf.float32)\n",
    "        self.accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32))\n",
    "        \n",
    "        self.prediction = tf.cast(tf.boolean_mask(self.heads_seq, mask), tf.int32)\n",
    "        mask_label = tf.boolean_mask(self.heads, mask)\n",
    "        correct_pred = tf.equal(self.prediction, mask_label)\n",
    "        correct_index = tf.cast(correct_pred, tf.float32)\n",
    "        self.accuracy_depends = tf.reduce_mean(tf.cast(correct_pred, tf.float32))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/husein/.local/lib/python3.6/site-packages/tensorflow_core/python/util/deprecation.py:507: calling count_nonzero (from tensorflow.python.ops.math_ops) with axis is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "reduction_indices is deprecated, use axis instead\n",
      "WARNING:tensorflow:\n",
      "The TensorFlow contrib module will not be included in TensorFlow 2.0.\n",
      "For more information, please see:\n",
      "  * https://github.com/tensorflow/community/blob/master/rfcs/20180907-contrib-sunset.md\n",
      "  * https://github.com/tensorflow/addons\n",
      "  * https://github.com/tensorflow/io (for I/O related ops)\n",
      "If you depend on functionality not listed there, please file an issue.\n",
      "\n",
      "WARNING:tensorflow:From /home/husein/xlnet/xlnet.py:220: The name tf.variable_scope is deprecated. Please use tf.compat.v1.variable_scope instead.\n",
      "\n",
      "WARNING:tensorflow:From /home/husein/xlnet/xlnet.py:220: The name tf.AUTO_REUSE is deprecated. Please use tf.compat.v1.AUTO_REUSE instead.\n",
      "\n",
      "WARNING:tensorflow:From /home/husein/xlnet/modeling.py:453: The name tf.logging.info is deprecated. Please use tf.compat.v1.logging.info instead.\n",
      "\n",
      "INFO:tensorflow:memory input None\n",
      "INFO:tensorflow:Use float type <dtype: 'float32'>\n",
      "WARNING:tensorflow:From /home/husein/xlnet/modeling.py:535: dropout (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use keras.layers.dropout instead.\n",
      "WARNING:tensorflow:From /home/husein/.local/lib/python3.6/site-packages/tensorflow_core/python/layers/core.py:271: 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",
      "WARNING:tensorflow:From /home/husein/xlnet/modeling.py:67: 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 <ipython-input-21-117b6be1eb3b>:138: where (from tensorflow.python.ops.array_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use tf.where in 2.0, which has the same broadcast rule as np.where\n",
      "WARNING:tensorflow:From /home/husein/.local/lib/python3.6/site-packages/tensorflow_core/contrib/crf/python/ops/crf.py:213: dynamic_rnn (from tensorflow.python.ops.rnn) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Please use `keras.layers.RNN(cell)`, which is equivalent to this API\n",
      "WARNING:tensorflow:From <ipython-input-21-117b6be1eb3b>:172: calling log_softmax (from tensorflow.python.ops.nn_ops) with dim is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "dim is deprecated, use axis instead\n"
     ]
    }
   ],
   "source": [
    "tf.reset_default_graph()\n",
    "sess = tf.InteractiveSession()\n",
    "\n",
    "learning_rate = 2e-5\n",
    "hidden_size_word = 128\n",
    "\n",
    "model = Model(learning_rate, hidden_size_word)\n",
    "sess.run(tf.global_variables_initializer())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "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": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "tvars = tf.trainable_variables()\n",
    "checkpoint = 'xlnet-base/model.ckpt'\n",
    "assignment_map, initialized_variable_names = get_assignment_map_from_checkpoint(tvars, \n",
    "                                                                                checkpoint)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Restoring parameters from xlnet-base/model.ckpt\n"
     ]
    }
   ],
   "source": [
    "saver = tf.train.Saver(var_list = assignment_map)\n",
    "saver.restore(sess, checkpoint)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.keras.preprocessing.sequence import pad_sequences\n",
    "\n",
    "batch_x = train_X[:5]\n",
    "batch_x = pad_sequences(batch_x,padding='post')\n",
    "batch_y = train_Y[:5]\n",
    "batch_y = pad_sequences(batch_y,padding='post')\n",
    "batch_depends = train_depends[:5]\n",
    "batch_depends = pad_sequences(batch_depends,padding='post')\n",
    "batch_segments = segments_train[:5]\n",
    "batch_segments = pad_sequences(batch_segments, padding='post', value = 4)\n",
    "batch_masks = masks_train[:5]\n",
    "batch_masks = pad_sequences(batch_masks, padding='post', value = 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0.04255319, 0.014184397, 128.52495]"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sess.run([model.accuracy, model.accuracy_depends, model.cost],\n",
    "        feed_dict = {model.words: batch_x,\n",
    "                model.types: batch_y,\n",
    "                model.heads: batch_depends,\n",
    "                model.segment_ids: batch_segments,\n",
    "                model.input_masks: batch_masks,\n",
    "                model.switch: False})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0.0070921984, 0.04964539, 544.50476]"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sess.run([model.accuracy, model.accuracy_depends, model.cost],\n",
    "        feed_dict = {model.words: batch_x,\n",
    "                model.types: batch_y,\n",
    "                model.heads: batch_depends,\n",
    "                model.segment_ids: batch_segments,\n",
    "                model.input_masks: batch_masks,\n",
    "                model.switch: True})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([26,  6,  6, 28, 26, 18, 19, 18, 28,  6, 26, 32, 27, 28, 27,  6, 28,\n",
       "        19, 19, 28,  6, 28,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,\n",
       "         0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,\n",
       "         0,  0], dtype=int32),\n",
       " array([21,  3, 15, 18, 10, 21, 10, 21,  7, 10,  1, 21,  1,  7,  3,  8,  7,\n",
       "         9, 10, 18,  8, 17,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,\n",
       "         0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,\n",
       "         0,  0]),\n",
       " array([ 3,  1,  3,  3,  3,  6,  3,  0,  3,  0,  9, 13, 11,  9, 15, 13, 15,\n",
       "         3,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,\n",
       "         0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,\n",
       "         0,  0], dtype=int32))"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tags_seq, heads = sess.run(\n",
    "    [model.logits, model.heads_seq],\n",
    "    feed_dict = {\n",
    "        model.words: batch_x,\n",
    "        model.segment_ids: batch_segments,\n",
    "        model.input_masks: batch_masks\n",
    "    },\n",
    ")\n",
    "tags_seq[0], heads[0], batch_depends[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 2519/2519 [15:04<00:00,  2.78it/s, accuracy=0.8, accuracy_depends=0.5, cost=1.95]     \n",
      "test minibatch loop: 100%|██████████| 630/630 [02:03<00:00,  5.12it/s, accuracy=0.843, accuracy_depends=0.545, cost=2.06]\n",
      "train minibatch loop:   0%|          | 0/2519 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 0, training loss: 3.876894, training acc: 0.650480, training depends: 0.419876, valid loss: 2.197435, valid acc: 0.826289, valid depends: 0.512468\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 2519/2519 [15:22<00:00,  2.73it/s, accuracy=0.9, accuracy_depends=0.55, cost=1.35]   \n",
      "test minibatch loop: 100%|██████████| 630/630 [02:05<00:00,  5.03it/s, accuracy=0.886, accuracy_depends=0.609, cost=1.59]\n",
      "train minibatch loop:   0%|          | 0/2519 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 1, training loss: 1.890330, training acc: 0.857705, training depends: 0.550963, valid loss: 1.625166, valid acc: 0.878244, valid depends: 0.601878\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 2519/2519 [15:23<00:00,  2.73it/s, accuracy=0.85, accuracy_depends=0.7, cost=0.788]   \n",
      "test minibatch loop: 100%|██████████| 630/630 [02:06<00:00,  4.99it/s, accuracy=0.902, accuracy_depends=0.636, cost=1.37] \n",
      "train minibatch loop:   0%|          | 0/2519 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 2, training loss: 1.410951, training acc: 0.887465, training depends: 0.646661, valid loss: 1.243608, valid acc: 0.898173, valid depends: 0.682270\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 2519/2519 [15:23<00:00,  2.73it/s, accuracy=0.75, accuracy_depends=0.75, cost=0.889]  \n",
      "test minibatch loop: 100%|██████████| 630/630 [02:05<00:00,  5.02it/s, accuracy=0.911, accuracy_depends=0.718, cost=1.1]  \n",
      "train minibatch loop:   0%|          | 0/2519 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 3, training loss: 1.064192, training acc: 0.905402, training depends: 0.724641, valid loss: 0.975716, valid acc: 0.909459, valid depends: 0.749798\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 2519/2519 [15:25<00:00,  2.72it/s, accuracy=0.95, accuracy_depends=0.75, cost=0.542]  \n",
      "test minibatch loop: 100%|██████████| 630/630 [02:05<00:00,  5.02it/s, accuracy=0.941, accuracy_depends=0.748, cost=0.93] \n",
      "train minibatch loop:   0%|          | 0/2519 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 4, training loss: 0.837911, training acc: 0.922502, training depends: 0.774991, valid loss: 0.811241, valid acc: 0.923954, valid depends: 0.790458\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 2519/2519 [15:24<00:00,  2.72it/s, accuracy=0.9, accuracy_depends=0.8, cost=0.952]    \n",
      "test minibatch loop: 100%|██████████| 630/630 [02:05<00:00,  5.03it/s, accuracy=0.957, accuracy_depends=0.768, cost=0.798]\n",
      "train minibatch loop:   0%|          | 0/2519 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 5, training loss: 0.686040, training acc: 0.935296, training depends: 0.807902, valid loss: 0.716174, valid acc: 0.933297, valid depends: 0.809396\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 2519/2519 [15:23<00:00,  2.73it/s, accuracy=0.9, accuracy_depends=0.85, cost=0.566]   \n",
      "test minibatch loop: 100%|██████████| 630/630 [02:04<00:00,  5.07it/s, accuracy=0.948, accuracy_depends=0.839, cost=0.688]\n",
      "train minibatch loop:   0%|          | 0/2519 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 6, training loss: 0.574526, training acc: 0.945672, training depends: 0.831804, valid loss: 0.622911, valid acc: 0.940280, valid depends: 0.835143\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 2519/2519 [15:23<00:00,  2.73it/s, accuracy=0.95, accuracy_depends=0.8, cost=0.212]   \n",
      "test minibatch loop: 100%|██████████| 630/630 [01:59<00:00,  5.28it/s, accuracy=0.95, accuracy_depends=0.841, cost=0.594] \n",
      "train minibatch loop:   0%|          | 0/2519 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 7, training loss: 0.491905, training acc: 0.953639, training depends: 0.848994, valid loss: 0.572177, valid acc: 0.944793, valid depends: 0.846656\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 2519/2519 [15:28<00:00,  2.71it/s, accuracy=1, accuracy_depends=0.8, cost=0.237]      \n",
      "test minibatch loop: 100%|██████████| 630/630 [02:06<00:00,  4.97it/s, accuracy=0.961, accuracy_depends=0.852, cost=0.484]\n",
      "train minibatch loop:   0%|          | 0/2519 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 8, training loss: 0.424984, training acc: 0.960825, training depends: 0.862493, valid loss: 0.532216, valid acc: 0.949322, valid depends: 0.856430\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 2519/2519 [15:29<00:00,  2.71it/s, accuracy=1, accuracy_depends=0.85, cost=0.081]     \n",
      "test minibatch loop: 100%|██████████| 630/630 [02:07<00:00,  4.96it/s, accuracy=0.957, accuracy_depends=0.843, cost=0.57] \n",
      "train minibatch loop:   0%|          | 0/2519 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 9, training loss: 0.372532, training acc: 0.966455, training depends: 0.872990, valid loss: 0.523298, valid acc: 0.953737, valid depends: 0.851046\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 2519/2519 [15:27<00:00,  2.72it/s, accuracy=1, accuracy_depends=0.85, cost=0.0546]     \n",
      "test minibatch loop: 100%|██████████| 630/630 [02:06<00:00,  4.99it/s, accuracy=0.959, accuracy_depends=0.845, cost=0.579]\n",
      "train minibatch loop:   0%|          | 0/2519 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 10, training loss: 0.329999, training acc: 0.971050, training depends: 0.882039, valid loss: 0.489395, valid acc: 0.958892, valid depends: 0.860195\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 2519/2519 [15:26<00:00,  2.72it/s, accuracy=1, accuracy_depends=0.8, cost=0.303]       \n",
      "test minibatch loop: 100%|██████████| 630/630 [02:06<00:00,  4.97it/s, accuracy=0.957, accuracy_depends=0.882, cost=0.386]\n",
      "train minibatch loop:   0%|          | 0/2519 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 11, training loss: 0.296270, training acc: 0.975037, training depends: 0.888372, valid loss: 0.414677, valid acc: 0.960939, valid depends: 0.883841\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 2519/2519 [15:27<00:00,  2.72it/s, accuracy=0.95, accuracy_depends=0.85, cost=0.0516]  \n",
      "test minibatch loop: 100%|██████████| 630/630 [02:06<00:00,  4.97it/s, accuracy=0.948, accuracy_depends=0.882, cost=0.456]\n",
      "train minibatch loop:   0%|          | 0/2519 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 12, training loss: 0.263290, training acc: 0.978965, training depends: 0.895745, valid loss: 0.402342, valid acc: 0.963851, valid depends: 0.887353\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 2519/2519 [15:27<00:00,  2.71it/s, accuracy=1, accuracy_depends=0.85, cost=0.00413]    \n",
      "test minibatch loop: 100%|██████████| 630/630 [02:07<00:00,  4.95it/s, accuracy=0.975, accuracy_depends=0.902, cost=0.358] \n",
      "train minibatch loop:   0%|          | 0/2519 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 13, training loss: 0.238861, training acc: 0.982122, training depends: 0.901446, valid loss: 0.393664, valid acc: 0.968449, valid depends: 0.882857\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 2519/2519 [15:33<00:00,  2.70it/s, accuracy=1, accuracy_depends=0.85, cost=0.00173]    \n",
      "test minibatch loop: 100%|██████████| 630/630 [02:06<00:00,  4.96it/s, accuracy=0.97, accuracy_depends=0.914, cost=0.349]  \n",
      "train minibatch loop:   0%|          | 0/2519 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 14, training loss: 0.222116, training acc: 0.984062, training depends: 0.904335, valid loss: 0.413641, valid acc: 0.969395, valid depends: 0.878797\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 2519/2519 [15:24<00:00,  2.73it/s, accuracy=1, accuracy_depends=0.85, cost=0.0861]     \n",
      "test minibatch loop: 100%|██████████| 630/630 [02:02<00:00,  5.16it/s, accuracy=0.977, accuracy_depends=0.907, cost=0.311] \n",
      "train minibatch loop:   0%|          | 0/2519 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 15, training loss: 0.206292, training acc: 0.985854, training depends: 0.908097, valid loss: 0.333951, valid acc: 0.970574, valid depends: 0.901331\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 2519/2519 [15:01<00:00,  2.80it/s, accuracy=0.95, accuracy_depends=0.85, cost=0.169]   \n",
      "test minibatch loop: 100%|██████████| 630/630 [02:01<00:00,  5.17it/s, accuracy=0.977, accuracy_depends=0.92, cost=0.227]  \n",
      "train minibatch loop:   0%|          | 0/2519 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 16, training loss: 0.190376, training acc: 0.986826, training depends: 0.911747, valid loss: 0.326896, valid acc: 0.971972, valid depends: 0.903429\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 2519/2519 [15:00<00:00,  2.80it/s, accuracy=1, accuracy_depends=0.85, cost=0.00131]    \n",
      "test minibatch loop: 100%|██████████| 630/630 [02:02<00:00,  5.15it/s, accuracy=0.964, accuracy_depends=0.916, cost=0.278] \n",
      "train minibatch loop:   0%|          | 0/2519 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 17, training loss: 0.177621, training acc: 0.987951, training depends: 0.915344, valid loss: 0.332637, valid acc: 0.971900, valid depends: 0.902321\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 2519/2519 [15:00<00:00,  2.80it/s, accuracy=1, accuracy_depends=0.85, cost=0.000426]   \n",
      "test minibatch loop: 100%|██████████| 630/630 [02:02<00:00,  5.15it/s, accuracy=0.966, accuracy_depends=0.93, cost=0.245]  \n",
      "train minibatch loop:   0%|          | 0/2519 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 18, training loss: 0.166710, training acc: 0.989098, training depends: 0.917499, valid loss: 0.322491, valid acc: 0.974446, valid depends: 0.907440\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 2519/2519 [15:01<00:00,  2.79it/s, accuracy=1, accuracy_depends=0.85, cost=0.00151]    \n",
      "test minibatch loop: 100%|██████████| 630/630 [02:03<00:00,  5.12it/s, accuracy=0.973, accuracy_depends=0.909, cost=0.285] "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 19, training loss: 0.155884, training acc: 0.989576, training depends: 0.920570, valid loss: 0.300692, valid acc: 0.976090, valid depends: 0.909638\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "epoch = 20\n",
    "for e in range(epoch):\n",
    "    train_acc, train_loss = [], []\n",
    "    test_acc, test_loss = [], []\n",
    "    train_acc_depends, test_acc_depends = [], []\n",
    "    \n",
    "    pbar = tqdm(\n",
    "        range(0, len(train_X), batch_size), desc = 'train minibatch loop'\n",
    "    )\n",
    "    for i in pbar:\n",
    "        index = min(i + batch_size, len(train_X))\n",
    "        batch_x = train_X[i: index]\n",
    "        batch_x = pad_sequences(batch_x,padding='post')\n",
    "        batch_y = train_Y[i: index]\n",
    "        batch_y = pad_sequences(batch_y,padding='post')\n",
    "        batch_depends = train_depends[i: index]\n",
    "        batch_depends = pad_sequences(batch_depends,padding='post')\n",
    "        batch_segments = segments_train[i: index]\n",
    "        batch_segments = pad_sequences(batch_segments, padding='post', value = 4)\n",
    "        batch_masks = masks_train[i: index]\n",
    "        batch_masks = pad_sequences(batch_masks, padding='post', value = 1)\n",
    "        \n",
    "        acc_depends, acc, cost, _ = sess.run(\n",
    "            [model.accuracy_depends, model.accuracy, model.cost, model.optimizer],\n",
    "            feed_dict = {\n",
    "                model.words: batch_x,\n",
    "                model.types: batch_y,\n",
    "                model.heads: batch_depends,\n",
    "                model.segment_ids: batch_segments,\n",
    "                model.input_masks: batch_masks,\n",
    "                model.switch: False\n",
    "            },\n",
    "        )\n",
    "        train_loss.append(cost)\n",
    "        train_acc.append(acc)\n",
    "        train_acc_depends.append(acc_depends)\n",
    "        pbar.set_postfix(cost = cost, accuracy = acc, accuracy_depends = acc_depends)\n",
    "        \n",
    "    pbar = tqdm(\n",
    "        range(0, len(test_X), batch_size), desc = 'test minibatch loop'\n",
    "    )\n",
    "    for i in pbar:\n",
    "        index = min(i + batch_size, len(test_X))\n",
    "        batch_x = test_X[i: index]\n",
    "        batch_x = pad_sequences(batch_x,padding='post')\n",
    "        batch_y = test_Y[i: index]\n",
    "        batch_y = pad_sequences(batch_y,padding='post')\n",
    "        batch_depends = test_depends[i: index]\n",
    "        batch_depends = pad_sequences(batch_depends,padding='post')\n",
    "        batch_segments = segments_test[i: index]\n",
    "        batch_segments = pad_sequences(batch_segments, padding='post', value = 4)\n",
    "        batch_masks = masks_test[i: index]\n",
    "        batch_masks = pad_sequences(batch_masks, padding='post', value = 1)\n",
    "        \n",
    "        acc_depends, acc, cost = sess.run(\n",
    "            [model.accuracy_depends, model.accuracy, model.cost],\n",
    "            feed_dict = {\n",
    "                model.words: batch_x,\n",
    "                model.types: batch_y,\n",
    "                model.heads: batch_depends,\n",
    "                model.segment_ids: batch_segments,\n",
    "                model.input_masks: batch_masks,\n",
    "                model.switch: False\n",
    "            },\n",
    "        )\n",
    "        test_loss.append(cost)\n",
    "        test_acc.append(acc)\n",
    "        test_acc_depends.append(acc_depends)\n",
    "        pbar.set_postfix(cost = cost, accuracy = acc, accuracy_depends = acc_depends)\n",
    "    \n",
    "    \n",
    "    print(\n",
    "    'epoch: %d, training loss: %f, training acc: %f, training depends: %f, valid loss: %f, valid acc: %f, valid depends: %f\\n'\n",
    "    % (e, np.mean(train_loss), \n",
    "       np.mean(train_acc), \n",
    "       np.mean(train_acc_depends), \n",
    "       np.mean(test_loss), \n",
    "       np.mean(test_acc), \n",
    "       np.mean(test_acc_depends)\n",
    "    ))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 2519/2519 [15:03<00:00,  2.79it/s, accuracy=0.95, accuracy_depends=0.85, cost=0.854]  \n",
      "test minibatch loop: 100%|██████████| 630/630 [02:02<00:00,  5.15it/s, accuracy=0.989, accuracy_depends=0.895, cost=2.23] \n",
      "train minibatch loop:   0%|          | 0/2519 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 0, training loss: 1.406487, training acc: 0.990304, training depends: 0.911989, valid loss: 3.639354, valid acc: 0.980902, valid depends: 0.892664\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 2519/2519 [14:55<00:00,  2.81it/s, accuracy=1, accuracy_depends=0.85, cost=0.000302]  \n",
      "test minibatch loop: 100%|██████████| 630/630 [02:00<00:00,  5.24it/s, accuracy=0.991, accuracy_depends=0.936, cost=2.01] \n",
      "train minibatch loop:   0%|          | 0/2519 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 1, training loss: 0.790412, training acc: 0.994914, training depends: 0.917277, valid loss: 3.557691, valid acc: 0.982399, valid depends: 0.902226\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 2519/2519 [15:08<00:00,  2.77it/s, accuracy=1, accuracy_depends=0.85, cost=0.00217]   \n",
      "test minibatch loop: 100%|██████████| 630/630 [02:05<00:00,  5.02it/s, accuracy=0.993, accuracy_depends=0.918, cost=2.68] \n",
      "train minibatch loop:   0%|          | 0/2519 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 2, training loss: 0.662908, training acc: 0.995871, training depends: 0.917625, valid loss: 3.628110, valid acc: 0.982800, valid depends: 0.901144\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 2519/2519 [14:50<00:00,  2.83it/s, accuracy=1, accuracy_depends=0.85, cost=0.00313]   \n",
      "test minibatch loop: 100%|██████████| 630/630 [02:00<00:00,  5.23it/s, accuracy=0.989, accuracy_depends=0.93, cost=3.5]   \n",
      "train minibatch loop:   0%|          | 0/2519 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 3, training loss: 0.599805, training acc: 0.996370, training depends: 0.917015, valid loss: 3.422855, valid acc: 0.983998, valid depends: 0.900396\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 2519/2519 [14:50<00:00,  2.83it/s, accuracy=1, accuracy_depends=0.85, cost=0.00308]   \n",
      "test minibatch loop: 100%|██████████| 630/630 [02:00<00:00,  5.24it/s, accuracy=0.995, accuracy_depends=0.9, cost=0.604]  "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 4, training loss: 0.545647, training acc: 0.996865, training depends: 0.917185, valid loss: 3.373970, valid acc: 0.984332, valid depends: 0.899482\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "epoch = 5\n",
    "for e in range(epoch):\n",
    "    train_acc, train_loss = [], []\n",
    "    test_acc, test_loss = [], []\n",
    "    train_acc_depends, test_acc_depends = [], []\n",
    "    \n",
    "    pbar = tqdm(\n",
    "        range(0, len(train_X), batch_size), desc = 'train minibatch loop'\n",
    "    )\n",
    "    for i in pbar:\n",
    "        index = min(i + batch_size, len(train_X))\n",
    "        batch_x = train_X[i: index]\n",
    "        batch_x = pad_sequences(batch_x,padding='post')\n",
    "        batch_y = train_Y[i: index]\n",
    "        batch_y = pad_sequences(batch_y,padding='post')\n",
    "        batch_depends = train_depends[i: index]\n",
    "        batch_depends = pad_sequences(batch_depends,padding='post')\n",
    "        batch_segments = segments_train[i: index]\n",
    "        batch_segments = pad_sequences(batch_segments, padding='post', value = 4)\n",
    "        batch_masks = masks_train[i: index]\n",
    "        batch_masks = pad_sequences(batch_masks, padding='post', value = 1)\n",
    "        \n",
    "        acc_depends, acc, cost, _ = sess.run(\n",
    "            [model.accuracy_depends, model.accuracy, model.cost, model.optimizer],\n",
    "            feed_dict = {\n",
    "                model.words: batch_x,\n",
    "                model.types: batch_y,\n",
    "                model.heads: batch_depends,\n",
    "                model.segment_ids: batch_segments,\n",
    "                model.input_masks: batch_masks,\n",
    "                model.switch: True\n",
    "            },\n",
    "        )\n",
    "        train_loss.append(cost)\n",
    "        train_acc.append(acc)\n",
    "        train_acc_depends.append(acc_depends)\n",
    "        pbar.set_postfix(cost = cost, accuracy = acc, accuracy_depends = acc_depends)\n",
    "        \n",
    "    pbar = tqdm(\n",
    "        range(0, len(test_X), batch_size), desc = 'test minibatch loop'\n",
    "    )\n",
    "    for i in pbar:\n",
    "        index = min(i + batch_size, len(test_X))\n",
    "        batch_x = test_X[i: index]\n",
    "        batch_x = pad_sequences(batch_x,padding='post')\n",
    "        batch_y = test_Y[i: index]\n",
    "        batch_y = pad_sequences(batch_y,padding='post')\n",
    "        batch_depends = test_depends[i: index]\n",
    "        batch_depends = pad_sequences(batch_depends,padding='post')\n",
    "        batch_segments = segments_test[i: index]\n",
    "        batch_segments = pad_sequences(batch_segments, padding='post', value = 4)\n",
    "        batch_masks = masks_test[i: index]\n",
    "        batch_masks = pad_sequences(batch_masks, padding='post', value = 1)\n",
    "        \n",
    "        acc_depends, acc, cost = sess.run(\n",
    "            [model.accuracy_depends, model.accuracy, model.cost],\n",
    "            feed_dict = {\n",
    "                model.words: batch_x,\n",
    "                model.types: batch_y,\n",
    "                model.heads: batch_depends,\n",
    "                model.segment_ids: batch_segments,\n",
    "                model.input_masks: batch_masks,\n",
    "                model.switch: True\n",
    "            },\n",
    "        )\n",
    "        test_loss.append(cost)\n",
    "        test_acc.append(acc)\n",
    "        test_acc_depends.append(acc_depends)\n",
    "        pbar.set_postfix(cost = cost, accuracy = acc, accuracy_depends = acc_depends)\n",
    "    \n",
    "    \n",
    "    print(\n",
    "    'epoch: %d, training loss: %f, training acc: %f, training depends: %f, valid loss: %f, valid acc: %f, valid depends: %f\\n'\n",
    "    % (e, np.mean(train_loss), \n",
    "       np.mean(train_acc), \n",
    "       np.mean(train_acc_depends), \n",
    "       np.mean(test_loss), \n",
    "       np.mean(test_acc), \n",
    "       np.mean(test_acc_depends)\n",
    "    ))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'xlnet-base-dependency/model.ckpt'"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "saver = tf.train.Saver(tf.trainable_variables())\n",
    "saver.save(sess, 'xlnet-base-dependency/model.ckpt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "kwargs = dict(\n",
    "      is_training=False,\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='xlnet-base/config.json')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:memory input None\n",
      "INFO:tensorflow:Use float type <dtype: 'float32'>\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/husein/.local/lib/python3.6/site-packages/tensorflow_core/python/client/session.py:1750: UserWarning: An interactive session is already active. This can cause out-of-memory errors in some cases. You must explicitly call `InteractiveSession.close()` to release resources held by the other session(s).\n",
      "  warnings.warn('An interactive session is already active. This can '\n"
     ]
    }
   ],
   "source": [
    "learning_rate = 2e-5\n",
    "hidden_size_word = 128\n",
    "\n",
    "tf.reset_default_graph()\n",
    "sess = tf.InteractiveSession()\n",
    "model = Model(learning_rate, hidden_size_word)\n",
    "sess.run(tf.global_variables_initializer())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Restoring parameters from xlnet-base-dependency/model.ckpt\n"
     ]
    }
   ],
   "source": [
    "saver = tf.train.Saver(tf.trainable_variables())\n",
    "saver.restore(sess, 'xlnet-base-dependency/model.ckpt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pred2label(pred):\n",
    "    out = []\n",
    "    for pred_i in pred:\n",
    "        out_i = []\n",
    "        for p in pred_i:\n",
    "            out_i.append(idx2tag[p])\n",
    "        out.append(out_i)\n",
    "    return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(heads_pred, types_pred, heads, types, lengths,\n",
    "             symbolic_root=False, symbolic_end=False):\n",
    "    batch_size, _ = heads_pred.shape\n",
    "    ucorr = 0.\n",
    "    lcorr = 0.\n",
    "    total = 0.\n",
    "    ucomplete_match = 0.\n",
    "    lcomplete_match = 0.\n",
    "\n",
    "    corr_root = 0.\n",
    "    total_root = 0.\n",
    "    start = 1 if symbolic_root else 0\n",
    "    end = 1 if symbolic_end else 0\n",
    "    for i in range(batch_size):\n",
    "        ucm = 1.\n",
    "        lcm = 1.\n",
    "        for j in range(start, lengths[i] - end):\n",
    "\n",
    "            total += 1\n",
    "            if heads[i, j] == heads_pred[i, j]:\n",
    "                ucorr += 1\n",
    "                if types[i, j] == types_pred[i, j]:\n",
    "                    lcorr += 1\n",
    "                else:\n",
    "                    lcm = 0\n",
    "            else:\n",
    "                ucm = 0\n",
    "                lcm = 0\n",
    "\n",
    "            if heads[i, j] == 0:\n",
    "                total_root += 1\n",
    "                corr_root += 1 if heads_pred[i, j] == 0 else 0\n",
    "\n",
    "        ucomplete_match += ucm\n",
    "        lcomplete_match += lcm\n",
    "    \n",
    "    return ucorr / total, lcorr / total, corr_root / total_root"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 630/630 [01:56<00:00,  5.39it/s]\n"
     ]
    }
   ],
   "source": [
    "arcs, types, roots = [], [], []\n",
    "real_Y, predict_Y = [], []\n",
    "\n",
    "for i in tqdm(range(0, len(test_X), batch_size)):\n",
    "    index = min(i + batch_size, len(test_X))\n",
    "    batch_x = test_X[i: index]\n",
    "    batch_x = pad_sequences(batch_x,padding='post')\n",
    "    batch_y = test_Y[i: index]\n",
    "    batch_y = pad_sequences(batch_y,padding='post')\n",
    "    batch_depends = test_depends[i: index]\n",
    "    batch_depends = pad_sequences(batch_depends,padding='post')\n",
    "    batch_segments = segments_test[i: index]\n",
    "    batch_segments = pad_sequences(batch_segments, padding='post', value = 4)\n",
    "    batch_masks = masks_test[i: index]\n",
    "    batch_masks = pad_sequences(batch_masks, padding='post', value = 1)\n",
    "    \n",
    "    tags_seq, heads = sess.run(\n",
    "        [model.logits, model.heads_seq],\n",
    "        feed_dict = {\n",
    "            model.words: batch_x,\n",
    "            model.segment_ids: batch_segments,\n",
    "            model.input_masks: batch_masks\n",
    "        },\n",
    "    )\n",
    "    \n",
    "    arc_accuracy, type_accuracy, root_accuracy = evaluate(heads - 1, tags_seq, batch_depends - 1, batch_y, \n",
    "            np.count_nonzero(batch_x, axis = 1))\n",
    "    arcs.append(arc_accuracy)\n",
    "    types.append(type_accuracy)\n",
    "    roots.append(root_accuracy)\n",
    "    predicted = pred2label(tags_seq)\n",
    "    real = pred2label(batch_y)\n",
    "    predict_Y.extend(predicted)\n",
    "    real_Y.extend(real)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "temp_real_Y = []\n",
    "for r in real_Y:\n",
    "    temp_real_Y.extend(r)\n",
    "    \n",
    "temp_predict_Y = []\n",
    "for r in predict_Y:\n",
    "    temp_predict_Y.extend(r)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "               precision    recall  f1-score   support\n",
      "\n",
      "          PAD    0.99998   1.00000   0.99999    632972\n",
      "            X    1.00000   0.99997   0.99999    143586\n",
      "          acl    0.98091   0.98226   0.98158      5806\n",
      "        advcl    0.97098   0.95161   0.96120      2356\n",
      "       advmod    0.98802   0.97806   0.98302      9527\n",
      "         amod    0.95966   0.97100   0.96530      8208\n",
      "        appos    0.98846   0.98947   0.98896      4936\n",
      "          aux    1.00000   1.00000   1.00000        10\n",
      "         case    0.99454   0.99110   0.99282     21128\n",
      "           cc    0.98704   0.99518   0.99109      6429\n",
      "        ccomp    0.89091   0.97313   0.93021       856\n",
      "     compound    0.98091   0.96643   0.97362     13079\n",
      "compound:plur    0.99068   0.98401   0.98733      1188\n",
      "         conj    0.98303   0.99214   0.98756      8524\n",
      "          cop    0.98664   0.99071   0.98867      1938\n",
      "        csubj    0.96000   0.96000   0.96000        50\n",
      "   csubj:pass    0.95652   0.91667   0.93617        24\n",
      "          dep    0.98182   0.96716   0.97444      1005\n",
      "          det    0.98698   0.97756   0.98225      8065\n",
      "        fixed    0.96071   0.97162   0.96613      1057\n",
      "         flat    0.98389   0.99064   0.98726     20411\n",
      "         iobj    0.96154   0.80645   0.87719        31\n",
      "         mark    0.96611   0.98539   0.97565      2806\n",
      "         nmod    0.97956   0.97285   0.97619      8030\n",
      "        nsubj    0.98317   0.98402   0.98359     12701\n",
      "   nsubj:pass    0.96930   0.97858   0.97392      3969\n",
      "       nummod    0.99113   0.99327   0.99220      7879\n",
      "          obj    0.98266   0.98076   0.98171     10342\n",
      "          obl    0.98468   0.98256   0.98362     11183\n",
      "    parataxis    0.95595   0.95455   0.95525       682\n",
      "        punct    0.99952   0.99949   0.99950     33107\n",
      "         root    0.98888   0.98888   0.98888     10073\n",
      "        xcomp    0.95951   0.96027   0.95989      2517\n",
      "\n",
      "     accuracy                        0.99678    994475\n",
      "    macro avg    0.97738   0.97381   0.97531    994475\n",
      " weighted avg    0.99679   0.99678   0.99678    994475\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from sklearn.metrics import classification_report\n",
    "print(classification_report(temp_real_Y, temp_predict_Y, digits = 5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "arc accuracy: 0.9310084738376598\n",
      "types accuracy: 0.9258795751889828\n",
      "root accuracy: 0.9474206349206349\n"
     ]
    }
   ],
   "source": [
    "print('arc accuracy:', np.mean(arcs))\n",
    "print('types accuracy:', np.mean(types))\n",
    "print('root accuracy:', np.mean(roots))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['Placeholder',\n",
       " 'Placeholder_1',\n",
       " 'Placeholder_2',\n",
       " 'Placeholder_3',\n",
       " 'Placeholder_4',\n",
       " 'Placeholder_5',\n",
       " 'W_d',\n",
       " 'W_e',\n",
       " 'U',\n",
       " 'U-bi',\n",
       " 'Wl',\n",
       " 'Wr',\n",
       " 'model/transformer/r_w_bias',\n",
       " 'model/transformer/r_r_bias',\n",
       " 'model/transformer/word_embedding/lookup_table',\n",
       " 'model/transformer/r_s_bias',\n",
       " 'model/transformer/seg_embed',\n",
       " 'model/transformer/layer_0/rel_attn/q/kernel',\n",
       " 'model/transformer/layer_0/rel_attn/k/kernel',\n",
       " 'model/transformer/layer_0/rel_attn/v/kernel',\n",
       " 'model/transformer/layer_0/rel_attn/r/kernel',\n",
       " 'model/transformer/layer_0/rel_attn/o/kernel',\n",
       " 'model/transformer/layer_0/rel_attn/LayerNorm/gamma',\n",
       " 'model/transformer/layer_0/ff/layer_1/kernel',\n",
       " 'model/transformer/layer_0/ff/layer_1/bias',\n",
       " 'model/transformer/layer_0/ff/layer_2/kernel',\n",
       " 'model/transformer/layer_0/ff/layer_2/bias',\n",
       " 'model/transformer/layer_0/ff/LayerNorm/gamma',\n",
       " 'model/transformer/layer_1/rel_attn/q/kernel',\n",
       " 'model/transformer/layer_1/rel_attn/k/kernel',\n",
       " 'model/transformer/layer_1/rel_attn/v/kernel',\n",
       " 'model/transformer/layer_1/rel_attn/r/kernel',\n",
       " 'model/transformer/layer_1/rel_attn/o/kernel',\n",
       " 'model/transformer/layer_1/rel_attn/LayerNorm/gamma',\n",
       " 'model/transformer/layer_1/ff/layer_1/kernel',\n",
       " 'model/transformer/layer_1/ff/layer_1/bias',\n",
       " 'model/transformer/layer_1/ff/layer_2/kernel',\n",
       " 'model/transformer/layer_1/ff/layer_2/bias',\n",
       " 'model/transformer/layer_1/ff/LayerNorm/gamma',\n",
       " 'model/transformer/layer_2/rel_attn/q/kernel',\n",
       " 'model/transformer/layer_2/rel_attn/k/kernel',\n",
       " 'model/transformer/layer_2/rel_attn/v/kernel',\n",
       " 'model/transformer/layer_2/rel_attn/r/kernel',\n",
       " 'model/transformer/layer_2/rel_attn/o/kernel',\n",
       " 'model/transformer/layer_2/rel_attn/LayerNorm/gamma',\n",
       " 'model/transformer/layer_2/ff/layer_1/kernel',\n",
       " 'model/transformer/layer_2/ff/layer_1/bias',\n",
       " 'model/transformer/layer_2/ff/layer_2/kernel',\n",
       " 'model/transformer/layer_2/ff/layer_2/bias',\n",
       " 'model/transformer/layer_2/ff/LayerNorm/gamma',\n",
       " 'model/transformer/layer_3/rel_attn/q/kernel',\n",
       " 'model/transformer/layer_3/rel_attn/k/kernel',\n",
       " 'model/transformer/layer_3/rel_attn/v/kernel',\n",
       " 'model/transformer/layer_3/rel_attn/r/kernel',\n",
       " 'model/transformer/layer_3/rel_attn/o/kernel',\n",
       " 'model/transformer/layer_3/rel_attn/LayerNorm/gamma',\n",
       " 'model/transformer/layer_3/ff/layer_1/kernel',\n",
       " 'model/transformer/layer_3/ff/layer_1/bias',\n",
       " 'model/transformer/layer_3/ff/layer_2/kernel',\n",
       " 'model/transformer/layer_3/ff/layer_2/bias',\n",
       " 'model/transformer/layer_3/ff/LayerNorm/gamma',\n",
       " 'model/transformer/layer_4/rel_attn/q/kernel',\n",
       " 'model/transformer/layer_4/rel_attn/k/kernel',\n",
       " 'model/transformer/layer_4/rel_attn/v/kernel',\n",
       " 'model/transformer/layer_4/rel_attn/r/kernel',\n",
       " 'model/transformer/layer_4/rel_attn/o/kernel',\n",
       " 'model/transformer/layer_4/rel_attn/LayerNorm/gamma',\n",
       " 'model/transformer/layer_4/ff/layer_1/kernel',\n",
       " 'model/transformer/layer_4/ff/layer_1/bias',\n",
       " 'model/transformer/layer_4/ff/layer_2/kernel',\n",
       " 'model/transformer/layer_4/ff/layer_2/bias',\n",
       " 'model/transformer/layer_4/ff/LayerNorm/gamma',\n",
       " 'model/transformer/layer_5/rel_attn/q/kernel',\n",
       " 'model/transformer/layer_5/rel_attn/k/kernel',\n",
       " 'model/transformer/layer_5/rel_attn/v/kernel',\n",
       " 'model/transformer/layer_5/rel_attn/r/kernel',\n",
       " 'model/transformer/layer_5/rel_attn/o/kernel',\n",
       " 'model/transformer/layer_5/rel_attn/LayerNorm/gamma',\n",
       " 'model/transformer/layer_5/ff/layer_1/kernel',\n",
       " 'model/transformer/layer_5/ff/layer_1/bias',\n",
       " 'model/transformer/layer_5/ff/layer_2/kernel',\n",
       " 'model/transformer/layer_5/ff/layer_2/bias',\n",
       " 'model/transformer/layer_5/ff/LayerNorm/gamma',\n",
       " 'model/transformer/layer_6/rel_attn/q/kernel',\n",
       " 'model/transformer/layer_6/rel_attn/k/kernel',\n",
       " 'model/transformer/layer_6/rel_attn/v/kernel',\n",
       " 'model/transformer/layer_6/rel_attn/r/kernel',\n",
       " 'model/transformer/layer_6/rel_attn/o/kernel',\n",
       " 'model/transformer/layer_6/rel_attn/LayerNorm/gamma',\n",
       " 'model/transformer/layer_6/ff/layer_1/kernel',\n",
       " 'model/transformer/layer_6/ff/layer_1/bias',\n",
       " 'model/transformer/layer_6/ff/layer_2/kernel',\n",
       " 'model/transformer/layer_6/ff/layer_2/bias',\n",
       " 'model/transformer/layer_6/ff/LayerNorm/gamma',\n",
       " 'model/transformer/layer_7/rel_attn/q/kernel',\n",
       " 'model/transformer/layer_7/rel_attn/k/kernel',\n",
       " 'model/transformer/layer_7/rel_attn/v/kernel',\n",
       " 'model/transformer/layer_7/rel_attn/r/kernel',\n",
       " 'model/transformer/layer_7/rel_attn/o/kernel',\n",
       " 'model/transformer/layer_7/rel_attn/LayerNorm/gamma',\n",
       " 'model/transformer/layer_7/ff/layer_1/kernel',\n",
       " 'model/transformer/layer_7/ff/layer_1/bias',\n",
       " 'model/transformer/layer_7/ff/layer_2/kernel',\n",
       " 'model/transformer/layer_7/ff/layer_2/bias',\n",
       " 'model/transformer/layer_7/ff/LayerNorm/gamma',\n",
       " 'model/transformer/layer_8/rel_attn/q/kernel',\n",
       " 'model/transformer/layer_8/rel_attn/k/kernel',\n",
       " 'model/transformer/layer_8/rel_attn/v/kernel',\n",
       " 'model/transformer/layer_8/rel_attn/r/kernel',\n",
       " 'model/transformer/layer_8/rel_attn/o/kernel',\n",
       " 'model/transformer/layer_8/rel_attn/LayerNorm/gamma',\n",
       " 'model/transformer/layer_8/ff/layer_1/kernel',\n",
       " 'model/transformer/layer_8/ff/layer_1/bias',\n",
       " 'model/transformer/layer_8/ff/layer_2/kernel',\n",
       " 'model/transformer/layer_8/ff/layer_2/bias',\n",
       " 'model/transformer/layer_8/ff/LayerNorm/gamma',\n",
       " 'model/transformer/layer_9/rel_attn/q/kernel',\n",
       " 'model/transformer/layer_9/rel_attn/k/kernel',\n",
       " 'model/transformer/layer_9/rel_attn/v/kernel',\n",
       " 'model/transformer/layer_9/rel_attn/r/kernel',\n",
       " 'model/transformer/layer_9/rel_attn/o/kernel',\n",
       " 'model/transformer/layer_9/rel_attn/LayerNorm/gamma',\n",
       " 'model/transformer/layer_9/ff/layer_1/kernel',\n",
       " 'model/transformer/layer_9/ff/layer_1/bias',\n",
       " 'model/transformer/layer_9/ff/layer_2/kernel',\n",
       " 'model/transformer/layer_9/ff/layer_2/bias',\n",
       " 'model/transformer/layer_9/ff/LayerNorm/gamma',\n",
       " 'model/transformer/layer_10/rel_attn/q/kernel',\n",
       " 'model/transformer/layer_10/rel_attn/k/kernel',\n",
       " 'model/transformer/layer_10/rel_attn/v/kernel',\n",
       " 'model/transformer/layer_10/rel_attn/r/kernel',\n",
       " 'model/transformer/layer_10/rel_attn/o/kernel',\n",
       " 'model/transformer/layer_10/rel_attn/LayerNorm/gamma',\n",
       " 'model/transformer/layer_10/ff/layer_1/kernel',\n",
       " 'model/transformer/layer_10/ff/layer_1/bias',\n",
       " 'model/transformer/layer_10/ff/layer_2/kernel',\n",
       " 'model/transformer/layer_10/ff/layer_2/bias',\n",
       " 'model/transformer/layer_10/ff/LayerNorm/gamma',\n",
       " 'model/transformer/layer_11/rel_attn/q/kernel',\n",
       " 'model/transformer/layer_11/rel_attn/k/kernel',\n",
       " 'model/transformer/layer_11/rel_attn/v/kernel',\n",
       " 'model/transformer/layer_11/rel_attn/r/kernel',\n",
       " 'model/transformer/layer_11/rel_attn/o/kernel',\n",
       " 'model/transformer/layer_11/rel_attn/LayerNorm/gamma',\n",
       " 'model/transformer/layer_11/ff/layer_1/kernel',\n",
       " 'model/transformer/layer_11/ff/layer_1/bias',\n",
       " 'model/transformer/layer_11/ff/layer_2/kernel',\n",
       " 'model/transformer/layer_11/ff/layer_2/bias',\n",
       " 'model/transformer/layer_11/ff/LayerNorm/gamma',\n",
       " 'dense/kernel',\n",
       " 'dense/bias',\n",
       " 'dense_1/kernel',\n",
       " 'dense_1/bias',\n",
       " 'dense_2/kernel',\n",
       " 'dense_2/bias',\n",
       " 'dense_3/kernel',\n",
       " 'dense_3/bias',\n",
       " 'heads_seq',\n",
       " 'tags_seq',\n",
       " 'transitions',\n",
       " 'logits']"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "strings = ','.join(\n",
    "    [\n",
    "        n.name\n",
    "        for n in tf.get_default_graph().as_graph_def().node\n",
    "        if ('Variable' in n.op\n",
    "        or 'Placeholder' in n.name\n",
    "        or '_seq' in n.name\n",
    "        or 'logits' in n.name\n",
    "        or 'alphas' in n.name\n",
    "        or 'self/Softmax' in n.name)\n",
    "        and 'Adam' not in n.name\n",
    "        and 'beta' not in n.name\n",
    "        and 'global_step' not in n.name\n",
    "        and 'adam' not in n.name\n",
    "        and 'gradients/bert' not in n.name\n",
    "    ]\n",
    ")\n",
    "strings.split(',')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "def freeze_graph(model_dir, output_node_names):\n",
    "\n",
    "    if not tf.gfile.Exists(model_dir):\n",
    "        raise AssertionError(\n",
    "            \"Export directory doesn't exists. Please specify an export \"\n",
    "            'directory: %s' % model_dir\n",
    "        )\n",
    "\n",
    "    checkpoint = tf.train.get_checkpoint_state(model_dir)\n",
    "    input_checkpoint = checkpoint.model_checkpoint_path\n",
    "\n",
    "    absolute_model_dir = '/'.join(input_checkpoint.split('/')[:-1])\n",
    "    output_graph = absolute_model_dir + '/frozen_model.pb'\n",
    "    clear_devices = True\n",
    "    with tf.Session(graph = tf.Graph()) as sess:\n",
    "        saver = tf.train.import_meta_graph(\n",
    "            input_checkpoint + '.meta', clear_devices = clear_devices\n",
    "        )\n",
    "        saver.restore(sess, input_checkpoint)\n",
    "        output_graph_def = tf.graph_util.convert_variables_to_constants(\n",
    "            sess,\n",
    "            tf.get_default_graph().as_graph_def(),\n",
    "            output_node_names.split(','),\n",
    "        )\n",
    "        with tf.gfile.GFile(output_graph, 'wb') as f:\n",
    "            f.write(output_graph_def.SerializeToString())\n",
    "        print('%d ops in the final graph.' % len(output_graph_def.node))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Restoring parameters from xlnet-base-dependency/model.ckpt\n",
      "WARNING:tensorflow:From <ipython-input-46-9a7215a4e58a>:23: convert_variables_to_constants (from tensorflow.python.framework.graph_util_impl) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use `tf.compat.v1.graph_util.convert_variables_to_constants`\n",
      "WARNING:tensorflow:From /home/husein/.local/lib/python3.6/site-packages/tensorflow_core/python/framework/graph_util_impl.py:277: extract_sub_graph (from tensorflow.python.framework.graph_util_impl) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use `tf.compat.v1.graph_util.extract_sub_graph`\n",
      "INFO:tensorflow:Froze 176 variables.\n",
      "INFO:tensorflow:Converted 176 variables to const ops.\n",
      "8206 ops in the final graph.\n"
     ]
    }
   ],
   "source": [
    "freeze_graph('xlnet-base-dependency', strings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "def merge_sentencepiece_tokens_tagging(x, y):\n",
    "    new_paired_tokens = []\n",
    "    n_tokens = len(x)\n",
    "    rejected = ['<cls>', '<sep>']\n",
    "\n",
    "    i = 0\n",
    "\n",
    "    while i < n_tokens:\n",
    "\n",
    "        current_token, current_label = x[i], y[i]\n",
    "        if not current_token.startswith('▁') and current_token not in rejected:\n",
    "            previous_token, previous_label = new_paired_tokens.pop()\n",
    "            merged_token = previous_token\n",
    "            merged_label = [previous_label]\n",
    "            while (\n",
    "                not current_token.startswith('▁')\n",
    "                and current_token not in rejected\n",
    "            ):\n",
    "                merged_token = merged_token + current_token.replace('▁', '')\n",
    "                merged_label.append(current_label)\n",
    "                i = i + 1\n",
    "                current_token, current_label = x[i], y[i]\n",
    "            merged_label = merged_label[0]\n",
    "            new_paired_tokens.append((merged_token, merged_label))\n",
    "\n",
    "        else:\n",
    "            new_paired_tokens.append((current_token, current_label))\n",
    "            i = i + 1\n",
    "\n",
    "    words = [\n",
    "        i[0].replace('▁', '')\n",
    "        for i in new_paired_tokens\n",
    "        if i[0] not in ['<cls>', '<sep>']\n",
    "    ]\n",
    "    labels = [i[1] for i in new_paired_tokens if i[0] not in ['<cls>', '<sep>']]\n",
    "    return words, labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['Kuala', 'Lumpur:', 'Sempena', 'sambutan', 'Aidilfitri', 'minggu', 'depan,', 'Perdana', 'Menteri', 'Tun', 'Dr', 'Mahathir', 'Mohamad', 'dan', 'Menteri', 'Pengangkutan', 'Anthony', 'Loke', 'Siew', 'Fook', 'menitipkan', 'pesanan', 'khas', 'kepada', 'orang', 'ramai', 'yang', 'mahu', 'pulang', 'ke', 'kampung', 'halaman', 'masing-masing.', 'Dalam', 'video', 'pendek', 'terbitan', 'Jabatan', 'Keselamatan', 'Jalan', 'Raya', '(Jkjr)', 'itu,', 'Dr', 'Mahathir', 'menasihati', 'mereka', 'supaya', 'berhenti', 'berehat', 'dan', 'tidur', 'sebentar', 'sekiranya', 'mengantuk', 'ketika', 'memandu.'] 57\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "73"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "string = 'KUALA LUMPUR: Sempena sambutan Aidilfitri minggu depan, Perdana Menteri Tun Dr Mahathir Mohamad dan Menteri Pengangkutan Anthony Loke Siew Fook menitipkan pesanan khas kepada orang ramai yang mahu pulang ke kampung halaman masing-masing. Dalam video pendek terbitan Jabatan Keselamatan Jalan Raya (JKJR) itu, Dr Mahathir menasihati mereka supaya berhenti berehat dan tidur sebentar  sekiranya mengantuk ketika memandu.'\n",
    "\n",
    "import re\n",
    "\n",
    "def entities_textcleaning(string, lowering = False):\n",
    "    \"\"\"\n",
    "    use by entities recognition, pos recognition and dependency parsing\n",
    "    \"\"\"\n",
    "    string = re.sub('[^A-Za-z0-9\\-\\/():,. ]+', ' ', string)\n",
    "    string = re.sub(r'[ ]+', ' ', string).strip()\n",
    "    original_string = string.split()\n",
    "    if lowering:\n",
    "        string = string.lower()\n",
    "    string = [\n",
    "        (original_string[no], word.title() if word.isupper() else word)\n",
    "        for no, word in enumerate(string.split())\n",
    "        if len(word)\n",
    "    ]\n",
    "    return [s[0] for s in string], [s[1] for s in string]\n",
    "\n",
    "def parse_X(left):\n",
    "    left = ' '.join(left)\n",
    "    bert_tokens = tokenize_fn(left)\n",
    "    bert_tokens.extend([4, 3])\n",
    "    segment = [0] * (len(bert_tokens) - 1) + [SEG_ID_CLS]\n",
    "    input_mask = [0] * len(segment)\n",
    "    s_tokens = [sp_model.IdToPiece(i) for i in bert_tokens]\n",
    "    return bert_tokens, segment, input_mask, s_tokens\n",
    "\n",
    "sequence = entities_textcleaning(string)[1]\n",
    "print(sequence, len(sequence))\n",
    "parsed_sequence, segment_sequence, mask_sequence, xlnet_sequence = parse_X(sequence)\n",
    "len(parsed_sequence)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/husein/.local/lib/python3.6/site-packages/tensorflow_core/python/client/session.py:1750: UserWarning: An interactive session is already active. This can cause out-of-memory errors in some cases. You must explicitly call `InteractiveSession.close()` to release resources held by the other session(s).\n",
      "  warnings.warn('An interactive session is already active. This can '\n"
     ]
    }
   ],
   "source": [
    "def load_graph(frozen_graph_filename):\n",
    "    with tf.gfile.GFile(frozen_graph_filename, 'rb') as f:\n",
    "        graph_def = tf.GraphDef()\n",
    "        graph_def.ParseFromString(f.read())\n",
    "    with tf.Graph().as_default() as graph:\n",
    "        tf.import_graph_def(graph_def)\n",
    "    return graph\n",
    "\n",
    "g = load_graph('xlnet-base-dependency/frozen_model.pb')\n",
    "x = g.get_tensor_by_name('import/Placeholder:0')\n",
    "seg = g.get_tensor_by_name('import/Placeholder_1:0')\n",
    "m = g.get_tensor_by_name('import/Placeholder_2:0')\n",
    "heads_seq = g.get_tensor_by_name('import/heads_seq:0')\n",
    "tags_seq = g.get_tensor_by_name('import/logits:0')\n",
    "test_sess = tf.InteractiveSession(graph = g)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "h, t = test_sess.run([heads_seq, tags_seq],\n",
    "        feed_dict = {\n",
    "            x: [parsed_sequence],\n",
    "            seg: [segment_sequence],\n",
    "            m: [mask_sequence],\n",
    "        },\n",
    ")\n",
    "h = h[0] - 1\n",
    "t = [idx2tag[d] for d in t[0]]\n",
    "merged_h = merge_sentencepiece_tokens_tagging(xlnet_sequence, h)\n",
    "merged_t = merge_sentencepiece_tokens_tagging(xlnet_sequence, t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[('Kuala', 23), ('Lumpur:', 1), ('Sempena', 5), ('sambutan', 23), ('Aidilfitri', 5), ('minggu', 6), ('depan,', 7), ('Perdana', 23), ('Menteri', 10), ('Tun', 11), ('Dr', 12), ('Mahathir', 21), ('Mohamad', 21), ('dan', 16), ('Menteri', 10), ('Pengangkutan', 17), ('Anthony', 24), ('Loke', 19), ('Siew', 21), ('Fook', 21), ('menitipkan', 0), ('pesanan', 23), ('khas', 24), ('kepada', 27), ('orang', 24), ('ramai', 27), ('yang', 31), ('mahu', 31), ('pulang', 27), ('ke', 33), ('kampung', 31), ('halaman', 33), ('masing-masing.', 33), ('Dalam', 38), ('video', 50), ('pendek', 38), ('terbitan', 38), ('Jabatan', 39), ('Keselamatan', 40), ('Jalan', 41), ('Raya', 41), ('(Jkjr)', 50), ('itu,', 38), ('Dr', 50), ('Mahathir', 50), ('menasihati', 23), ('mereka', 50), ('supaya', 52), ('berhenti', 50), ('berehat', 52), ('dan', 54), ('tidur', 52), ('sebentar', 56), ('sekiranya', 58), ('mengantuk', 54), ('ketika', 59), ('memandu.', 58)]\n"
     ]
    }
   ],
   "source": [
    "print(list(zip(merged_h[0], merged_h[1])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[('Kuala', 'nsubj'), ('Lumpur:', 'flat'), ('Sempena', 'case'), ('sambutan', 'obl'), ('Aidilfitri', 'compound'), ('minggu', 'compound'), ('depan,', 'compound'), ('Perdana', 'nsubj'), ('Menteri', 'flat'), ('Tun', 'flat'), ('Dr', 'flat'), ('Mahathir', 'flat'), ('Mohamad', 'flat'), ('dan', 'cc'), ('Menteri', 'conj'), ('Pengangkutan', 'flat'), ('Anthony', 'flat'), ('Loke', 'flat'), ('Siew', 'flat'), ('Fook', 'flat'), ('menitipkan', 'root'), ('pesanan', 'obj'), ('khas', 'amod'), ('kepada', 'case'), ('orang', 'nmod'), ('ramai', 'compound'), ('yang', 'nsubj'), ('mahu', 'advmod'), ('pulang', 'acl'), ('ke', 'case'), ('kampung', 'obl'), ('halaman', 'compound'), ('masing-masing.', 'det'), ('Dalam', 'case'), ('video', 'obl'), ('pendek', 'amod'), ('terbitan', 'compound'), ('Jabatan', 'flat'), ('Keselamatan', 'flat'), ('Jalan', 'flat'), ('Raya', 'flat'), ('(Jkjr)', 'punct'), ('itu,', 'det'), ('Dr', 'nsubj'), ('Mahathir', 'flat'), ('menasihati', 'parataxis'), ('mereka', 'obj'), ('supaya', 'case'), ('berhenti', 'xcomp'), ('berehat', 'xcomp'), ('dan', 'cc'), ('tidur', 'conj'), ('sebentar', 'advmod'), ('sekiranya', 'mark'), ('mengantuk', 'ccomp'), ('ketika', 'mark'), ('memandu.', 'ccomp')]\n"
     ]
    }
   ],
   "source": [
    "print(list(zip(merged_t[0], merged_t[1])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "\n",
    "bucketName = 'huseinhouse-storage'\n",
    "Key = 'xlnet-base-dependency/frozen_model.pb'\n",
    "outPutname = \"v34/dependency/xlnet-base-dependency.pb\"\n",
    "\n",
    "s3 = boto3.client('s3')\n",
    "\n",
    "s3.upload_file(Key,bucketName,outPutname)"
   ]
  },
  {
   "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
}
