entity-tagging/6.attention-is-all-you-need.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bs4 import BeautifulSoup\n",
    "import re\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_string(string):\n",
    "    string = re.sub('[^A-Za-z0-9\\-\\/ ]+', ' ', string).split()\n",
    "    return [y.strip() for y in string]\n",
    "\n",
    "def to_title(string):\n",
    "    if string.isupper():\n",
    "        string = string.title()\n",
    "    return string"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_raw(filename):\n",
    "    with open(filename, 'r') as fopen:\n",
    "        entities = fopen.read()\n",
    "    soup = BeautifulSoup(entities, 'html.parser')\n",
    "    inside_tag = ''\n",
    "    texts, labels = [], []\n",
    "    for sentence in soup.prettify().split('\\n'):\n",
    "        if len(inside_tag):\n",
    "            splitted = process_string(sentence)\n",
    "            texts += splitted\n",
    "            labels += [inside_tag] * len(splitted)\n",
    "            inside_tag = ''\n",
    "        else:\n",
    "            if not sentence.find('</'):\n",
    "                pass\n",
    "            elif not sentence.find('<'):\n",
    "                inside_tag = sentence.split('>')[0][1:]\n",
    "            else:\n",
    "                splitted = process_string(sentence)\n",
    "                texts += splitted\n",
    "                labels += ['OTHER'] * len(splitted)\n",
    "    assert (len(texts)==len(labels)), \"length texts and labels are not same\"\n",
    "    print('len texts and labels: ', len(texts))\n",
    "    return texts,labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "len texts and labels:  34012\n",
      "len texts and labels:  9249\n"
     ]
    }
   ],
   "source": [
    "train_texts, train_labels = parse_raw('data_train.txt')\n",
    "test_texts, test_labels = parse_raw('data_test.txt')\n",
    "train_texts += test_texts\n",
    "train_labels += test_labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array(['OTHER', 'location', 'organization', 'person', 'quantity', 'time'],\n",
       "       dtype='<U12'), array([35613,  1536,  1592,  2358,  1336,   826]))"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.unique(train_labels,return_counts=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('entities-bm-normalize-v3.txt','r') as fopen:\n",
    "    entities_bm = fopen.read().split('\\n')[:-1]\n",
    "entities_bm = [i.split() for i in entities_bm]\n",
    "entities_bm = [[i[0],'TIME' if i[0] in 'jam' else i[1]] for i in entities_bm]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "replace_by = {'organizaiton':'organization','orgnization':'organization',\n",
    "             'othoer': 'OTHER'}\n",
    "\n",
    "with open('NER-part1.txt','r') as fopen:\n",
    "    nexts = fopen.read().split('\\n')[:-1]\n",
    "nexts = [i.split() for i in nexts]\n",
    "for i in nexts:\n",
    "    if len(i) == 2:\n",
    "        label = i[1].lower()\n",
    "        if 'other' in label:\n",
    "            label = label.upper()\n",
    "        if label in replace_by:\n",
    "            label = replace_by[label]\n",
    "        train_labels.append(label)\n",
    "        train_texts.append(i[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "'KN'\n",
      "'KA'\n"
     ]
    }
   ],
   "source": [
    "replace_by = {'LOC':'location','PRN':'person','NORP':'organization','ORG':'organization','LAW':'law',\n",
    "             'EVENT':'event','FAC':'organization','TIME':'time','O':'OTHER','ART':'person','DOC':'law'}\n",
    "for i in entities_bm:\n",
    "    try:\n",
    "        string = process_string(i[0])\n",
    "        if len(string):\n",
    "            train_labels.append(replace_by[i[1]])\n",
    "            train_texts.append(process_string(i[0])[0])  \n",
    "    except Exception as e:\n",
    "        print(e)\n",
    "        \n",
    "assert (len(train_texts)==len(train_labels)), \"length texts and labels are not same\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array(['OTHER', 'event', 'law', 'location', 'organization', 'person',\n",
       "        'quantity', 'time'], dtype='<U12'),\n",
       " array([49712,   234,   185,  2056,  2596,  4397,  1341,  1296]))"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.unique(train_labels,return_counts=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "word2idx = {'PAD': 0,'NUM':1,'UNK':2}\n",
    "tag2idx = {'PAD': 0}\n",
    "char2idx = {'PAD': 0,'NUM':1,'UNK':2}\n",
    "word_idx = 3\n",
    "tag_idx = 1\n",
    "char_idx = 3\n",
    "\n",
    "def parse_XY(texts, labels):\n",
    "    global word2idx, tag2idx, char2idx, word_idx, tag_idx, char_idx\n",
    "    X, Y = [], []\n",
    "    for no, text in enumerate(texts):\n",
    "        text = text.lower()\n",
    "        tag = labels[no]\n",
    "        for c in text:\n",
    "            if c not in char2idx:\n",
    "                char2idx[c] = char_idx\n",
    "                char_idx += 1\n",
    "        if tag not in tag2idx:\n",
    "            tag2idx[tag] = tag_idx\n",
    "            tag_idx += 1\n",
    "        Y.append(tag2idx[tag])\n",
    "        if text not in word2idx:\n",
    "            word2idx[text] = word_idx\n",
    "            word_idx += 1\n",
    "        X.append(word2idx[text])\n",
    "    return X, np.array(Y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array(['OTHER', 'event', 'law', 'location', 'organization', 'person',\n",
       "       'quantity', 'time'], dtype='<U12')"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.unique(train_labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "X, Y = parse_XY(train_texts, train_labels)\n",
    "idx2word={idx: tag for tag, idx in word2idx.items()}\n",
    "idx2tag = {i: w for w, i in tag2idx.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "seq_len = 50\n",
    "def iter_seq(x):\n",
    "    return np.array([x[i: i+seq_len] for i in range(0, len(x)-seq_len, 1)])\n",
    "\n",
    "def to_train_seq(*args):\n",
    "    return [iter_seq(x) for x in args]\n",
    "\n",
    "def generate_char_seq(batch):\n",
    "    x = [[len(idx2word[i]) for i in k] for k in batch]\n",
    "    maxlen = max([j for i in x for j in i])\n",
    "    temp = np.zeros((batch.shape[0],batch.shape[1],maxlen),dtype=np.int32)\n",
    "    for i in range(batch.shape[0]):\n",
    "        for k in range(batch.shape[1]):\n",
    "            for no, c in enumerate(idx2word[batch[i,k]][::-1]):\n",
    "                temp[i,k,-1-no] = char2idx[c]\n",
    "    return temp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "with open('attention-is-all-you-need-ner.json','w') as fopen:\n",
    "    fopen.write(json.dumps({'idx2tag':idx2tag,'idx2word':idx2word,\n",
    "           'word2idx':word2idx,'tag2idx':tag2idx,'char2idx':char2idx}))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(61767, 50)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X_seq, Y_seq = to_train_seq(X, Y)\n",
    "X_char_seq = generate_char_seq(X_seq)\n",
    "X_seq.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using TensorFlow backend.\n"
     ]
    }
   ],
   "source": [
    "from keras.utils import to_categorical\n",
    "Y_seq_3d = np.array([to_categorical(i, num_classes=len(tag2idx)) for i in Y_seq])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/husein/.local/lib/python3.6/site-packages/sklearn/cross_validation.py:41: DeprecationWarning: This module was deprecated in version 0.18 in favor of the model_selection module into which all the refactored classes and functions are moved. Also note that the interface of the new CV iterators are different from that of this module. This module will be removed in 0.20.\n",
      "  \"This module will be removed in 0.20.\", DeprecationWarning)\n"
     ]
    }
   ],
   "source": [
    "from sklearn.cross_validation import train_test_split\n",
    "train_X, test_X, train_Y, test_Y, train_char, test_char = train_test_split(X_seq, Y_seq_3d, X_char_seq, \n",
    "                                                                           test_size=0.2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def layer_norm(inputs, epsilon=1e-8):\n",
    "    mean, variance = tf.nn.moments(inputs, [-1], keep_dims=True)\n",
    "    normalized = (inputs - mean) / (tf.sqrt(variance + epsilon))\n",
    "\n",
    "    params_shape = inputs.get_shape()[-1:]\n",
    "    gamma = tf.get_variable('gamma', params_shape, tf.float32, tf.ones_initializer())\n",
    "    beta = tf.get_variable('beta', params_shape, tf.float32, tf.zeros_initializer())\n",
    "    \n",
    "    outputs = gamma * normalized + beta\n",
    "    return outputs\n",
    "\n",
    "def multihead_attn(queries, keys, q_masks, k_masks, future_binding, num_units, num_heads):\n",
    "    \n",
    "    T_q = tf.shape(queries)[1]                                      \n",
    "    T_k = tf.shape(keys)[1]                  \n",
    "\n",
    "    Q = tf.layers.dense(queries, num_units, name='Q')                              \n",
    "    K_V = tf.layers.dense(keys, 2*num_units, name='K_V')    \n",
    "    K, V = tf.split(K_V, 2, -1)        \n",
    "\n",
    "    Q_ = tf.concat(tf.split(Q, num_heads, axis=2), axis=0)                         \n",
    "    K_ = tf.concat(tf.split(K, num_heads, axis=2), axis=0)                    \n",
    "    V_ = tf.concat(tf.split(V, num_heads, axis=2), axis=0)                      \n",
    "\n",
    "    align = tf.matmul(Q_, tf.transpose(K_, [0,2,1]))                      \n",
    "    align = align / np.sqrt(K_.get_shape().as_list()[-1])                 \n",
    "\n",
    "    paddings = tf.fill(tf.shape(align), float('-inf'))                   \n",
    "\n",
    "    key_masks = k_masks                                                 \n",
    "    key_masks = tf.tile(key_masks, [num_heads, 1])                       \n",
    "    key_masks = tf.tile(tf.expand_dims(key_masks, 1), [1, T_q, 1])            \n",
    "    align = tf.where(tf.equal(key_masks, 0), paddings, align)       \n",
    "\n",
    "    if future_binding:\n",
    "        lower_tri = tf.ones([T_q, T_k])                                          \n",
    "        lower_tri = tf.linalg.LinearOperatorLowerTriangular(lower_tri).to_dense()  \n",
    "        masks = tf.tile(tf.expand_dims(lower_tri,0), [tf.shape(align)[0], 1, 1]) \n",
    "        align = tf.where(tf.equal(masks, 0), paddings, align)                      \n",
    "    \n",
    "    align = tf.nn.softmax(align)                                            \n",
    "    query_masks = tf.to_float(q_masks)                                             \n",
    "    query_masks = tf.tile(query_masks, [num_heads, 1])                             \n",
    "    query_masks = tf.tile(tf.expand_dims(query_masks, -1), [1, 1, T_k])            \n",
    "    align *= query_masks\n",
    "    outputs = tf.matmul(align, V_)                                                 \n",
    "    outputs = tf.concat(tf.split(outputs, num_heads, axis=0), axis=2)             \n",
    "    outputs += queries                                                             \n",
    "    outputs = layer_norm(outputs)                                                 \n",
    "    return outputs\n",
    "\n",
    "\n",
    "def pointwise_feedforward(inputs, hidden_units, activation=None):\n",
    "    outputs = tf.layers.dense(inputs, 4*hidden_units, activation=activation)\n",
    "    outputs = tf.layers.dense(outputs, hidden_units, activation=None)\n",
    "    outputs += inputs\n",
    "    outputs = layer_norm(outputs)\n",
    "    return outputs\n",
    "\n",
    "\n",
    "def learned_position_encoding(inputs, mask, embed_dim):\n",
    "    T = tf.shape(inputs)[1]\n",
    "    outputs = tf.range(tf.shape(inputs)[1])                # (T_q)\n",
    "    outputs = tf.expand_dims(outputs, 0)                   # (1, T_q)\n",
    "    outputs = tf.tile(outputs, [tf.shape(inputs)[0], 1])   # (N, T_q)\n",
    "    outputs = embed_seq(outputs, T, embed_dim, zero_pad=False, scale=False)\n",
    "    return tf.expand_dims(tf.to_float(mask), -1) * outputs\n",
    "\n",
    "\n",
    "def sinusoidal_position_encoding(inputs, mask, repr_dim):\n",
    "    T = tf.shape(inputs)[1]\n",
    "    pos = tf.reshape(tf.range(0.0, tf.to_float(T), dtype=tf.float32), [-1, 1])\n",
    "    i = np.arange(0, repr_dim, 2, np.float32)\n",
    "    denom = np.reshape(np.power(10000.0, i / repr_dim), [1, -1])\n",
    "    enc = tf.expand_dims(tf.concat([tf.sin(pos / denom), tf.cos(pos / denom)], 1), 0)\n",
    "    return tf.tile(enc, [tf.shape(inputs)[0], 1, 1]) * tf.expand_dims(tf.to_float(mask), -1)\n",
    "\n",
    "def label_smoothing(inputs, epsilon=0.1):\n",
    "    C = inputs.get_shape().as_list()[-1]\n",
    "    return ((1 - epsilon) * inputs) + (epsilon / C)\n",
    "\n",
    "\n",
    "class CRF:\n",
    "    def __init__(self,\n",
    "                 dim_word,\n",
    "                 dim_char,\n",
    "                 dropout,\n",
    "                 learning_rate,\n",
    "                 hidden_size_char,\n",
    "                 hidden_size_word,\n",
    "                 num_blocks = 2,\n",
    "                 num_heads = 8,\n",
    "                 min_freq = 50):\n",
    "        \n",
    "        self.word_ids = tf.placeholder(tf.int32, shape = [None, None])\n",
    "        self.char_ids = tf.placeholder(tf.int32, shape = [None, None, None])\n",
    "        self.labels = tf.placeholder(tf.int32, shape = [None, None, None])\n",
    "        self.maxlen = tf.shape(self.word_ids)[1]\n",
    "        self.lengths = tf.count_nonzero(self.word_ids, 1)\n",
    "        batch_size = tf.shape(self.word_ids)[0]\n",
    "        \n",
    "        self.word_embeddings = tf.Variable(\n",
    "            tf.truncated_normal(\n",
    "                [len(word2idx), dim_word], stddev = 1.0 / np.sqrt(dim_word)\n",
    "            )\n",
    "        )\n",
    "        self.char_embeddings = tf.Variable(\n",
    "            tf.truncated_normal(\n",
    "                [len(char2idx), dim_char], stddev = 1.0 / np.sqrt(dim_char)\n",
    "            )\n",
    "        )\n",
    "        \n",
    "        word_embedded = tf.nn.embedding_lookup(\n",
    "            self.word_embeddings, self.word_ids\n",
    "        )\n",
    "        char_embedded = tf.nn.embedding_lookup(\n",
    "            self.char_embeddings, self.char_ids\n",
    "        )\n",
    "        s = tf.shape(char_embedded)\n",
    "        char_embedded = tf.reshape(\n",
    "            char_embedded, shape = [s[0] * s[1], s[-2], dim_char]\n",
    "        )\n",
    "        reshape_char = tf.reshape(self.char_ids, shape = [s[0] * s[1], s[-2]])\n",
    "        char_masked = tf.sign(reshape_char)\n",
    "        char_embedded += sinusoidal_position_encoding(reshape_char, char_masked, dim_char)\n",
    "        for i in range(num_blocks):\n",
    "            with tf.variable_scope('char_%d'%i,reuse=tf.AUTO_REUSE):\n",
    "                char_embedded = multihead_attn(queries = char_embedded,\n",
    "                                                 keys = char_embedded,\n",
    "                                                 q_masks = char_masked,\n",
    "                                                 k_masks = char_masked,\n",
    "                                                 future_binding = False,\n",
    "                                                 num_units = dim_char,\n",
    "                                                 num_heads = num_heads)\n",
    "            with tf.variable_scope('char_feedforward_%d'%i,reuse=tf.AUTO_REUSE):\n",
    "                char_embedded = pointwise_feedforward(char_embedded,\n",
    "                                                    dim_char,\n",
    "                                                    activation = tf.nn.relu)\n",
    "        output = tf.reshape(\n",
    "            char_embedded[:, -1], shape = [s[0], s[1], 2 * hidden_size_char]\n",
    "        )\n",
    "        word_embedded = tf.concat([word_embedded, output], axis = -1)\n",
    "        word_embedded = tf.layers.dense(word_embedded, dim_char)\n",
    "        de_masks = tf.sign(self.word_ids)\n",
    "        word_embedded += sinusoidal_position_encoding(self.word_ids, de_masks, dim_char)\n",
    "        \n",
    "        for i in range(num_blocks):\n",
    "            with tf.variable_scope('word_char_%d'%i,reuse=tf.AUTO_REUSE):\n",
    "                decoder_embedded = multihead_attn(queries = word_embedded,\n",
    "                                         keys = word_embedded,\n",
    "                                         q_masks = de_masks,\n",
    "                                         k_masks = de_masks,\n",
    "                                         future_binding = True,\n",
    "                                         num_units = dim_char,\n",
    "                                         num_heads = num_heads)\n",
    "                \n",
    "            with tf.variable_scope('word_char_attention_%d'%i,reuse=tf.AUTO_REUSE):\n",
    "                decoder_embedded = multihead_attn(queries = decoder_embedded,\n",
    "                                         keys = output,\n",
    "                                         q_masks = de_masks,\n",
    "                                         k_masks = de_masks,\n",
    "                                         future_binding = False,\n",
    "                                         num_units = dim_char,\n",
    "                                         num_heads = num_heads)\n",
    "            \n",
    "            with tf.variable_scope('word_feedforward_%d'%i,reuse=tf.AUTO_REUSE):\n",
    "                decoder_embedded = pointwise_feedforward(decoder_embedded,\n",
    "                                                    dim_char,\n",
    "                                            activation = tf.nn.relu)\n",
    "        logits = tf.layers.dense(decoder_embedded, len(idx2tag))\n",
    "        y_t = tf.argmax(self.labels, 2)\n",
    "        log_likelihood, transition_params = tf.contrib.crf.crf_log_likelihood(\n",
    "            logits, y_t, self.lengths\n",
    "        )\n",
    "        self.cost = tf.reduce_mean(-log_likelihood)\n",
    "        self.optimizer = tf.train.AdamOptimizer(\n",
    "            learning_rate = learning_rate\n",
    "        ).minimize(self.cost)\n",
    "        mask = tf.sequence_mask(self.lengths, maxlen = self.maxlen)\n",
    "        self.tags_seq, tags_score = tf.contrib.crf.crf_decode(\n",
    "            logits, transition_params, self.lengths\n",
    "        )\n",
    "        self.tags_seq = tf.identity(self.tags_seq, name = 'logits')\n",
    "\n",
    "        y_t = tf.cast(y_t, tf.int32)\n",
    "        self.prediction = tf.boolean_mask(self.tags_seq, mask)\n",
    "        mask_label = tf.boolean_mask(y_t, mask)\n",
    "        correct_pred = tf.equal(self.prediction, mask_label)\n",
    "        correct_index = tf.cast(correct_pred, tf.float32)\n",
    "        self.accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/gradients_impl.py:112: UserWarning: Converting sparse IndexedSlices to a dense Tensor of unknown shape. This may consume a large amount of memory.\n",
      "  \"Converting sparse IndexedSlices to a dense Tensor of unknown shape. \"\n"
     ]
    }
   ],
   "source": [
    "tf.reset_default_graph()\n",
    "sess = tf.InteractiveSession()\n",
    "\n",
    "dim_word = 128\n",
    "dim_char = 256\n",
    "dropout = 0.8\n",
    "learning_rate = 1e-3\n",
    "hidden_size_char = 128\n",
    "hidden_size_word = 64\n",
    "num_layers = 2\n",
    "batch_size = 32\n",
    "\n",
    "model = CRF(dim_word,dim_char,dropout,learning_rate,hidden_size_char,hidden_size_word,num_layers)\n",
    "sess.run(tf.global_variables_initializer())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 1545/1545 [10:49<00:00,  2.94it/s, accuracy=0.996, cost=0.904]\n",
      "test minibatch loop: 100%|██████████| 387/387 [01:11<00:00,  5.41it/s, accuracy=0.97, cost=2.86]  \n",
      "train minibatch loop:   0%|          | 0/1545 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "time taken: 720.9585599899292\n",
      "epoch: 0, training loss: 7.664514, training acc: 0.957292, valid loss: 0.896377, valid acc: 0.996098\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 1545/1545 [10:51<00:00,  2.93it/s, accuracy=1, cost=0.219]    \n",
      "test minibatch loop: 100%|██████████| 387/387 [01:11<00:00,  5.42it/s, accuracy=1, cost=0.67]     "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "time taken: 723.0113635063171\n",
      "epoch: 1, training loss: 0.744680, training acc: 0.995085, valid loss: 0.765401, valid acc: 0.996772\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "\n",
    "for e in range(2):\n",
    "    lasttime = time.time()\n",
    "    train_acc, train_loss, test_acc, test_loss = 0, 0, 0, 0\n",
    "    pbar = tqdm(\n",
    "        range(0, len(train_X), batch_size), desc = 'train minibatch loop'\n",
    "    )\n",
    "    for i in pbar:\n",
    "        batch_x = train_X[i : min(i + batch_size, train_X.shape[0])]\n",
    "        batch_char = train_char[i : min(i + batch_size, train_X.shape[0])]\n",
    "        batch_y = train_Y[i : min(i + batch_size, train_X.shape[0])]\n",
    "        acc, cost, _ = sess.run(\n",
    "            [model.accuracy, model.cost, model.optimizer],\n",
    "            feed_dict = {\n",
    "                model.word_ids: batch_x,\n",
    "                model.char_ids: batch_char,\n",
    "                model.labels: batch_y\n",
    "            },\n",
    "        )\n",
    "        assert not np.isnan(cost)\n",
    "        train_loss += cost\n",
    "        train_acc += acc\n",
    "        pbar.set_postfix(cost = cost, accuracy = acc)\n",
    "    \n",
    "    pbar = tqdm(\n",
    "        range(0, len(test_X), batch_size), desc = 'test minibatch loop'\n",
    "    )\n",
    "    for i in pbar:\n",
    "        batch_x = test_X[i : min(i + batch_size, test_X.shape[0])]\n",
    "        batch_char = test_char[i : min(i + batch_size, test_X.shape[0])]\n",
    "        batch_y = test_Y[i : min(i + batch_size, test_X.shape[0])]\n",
    "        acc, cost = sess.run(\n",
    "            [model.accuracy, model.cost],\n",
    "            feed_dict = {\n",
    "                model.word_ids: batch_x,\n",
    "                model.char_ids: batch_char,\n",
    "                model.labels: batch_y\n",
    "            },\n",
    "        )\n",
    "        assert not np.isnan(cost)\n",
    "        test_loss += cost\n",
    "        test_acc += acc\n",
    "        pbar.set_postfix(cost = cost, accuracy = acc)\n",
    "    \n",
    "    train_loss /= len(train_X) / batch_size\n",
    "    train_acc /= len(train_X) / batch_size\n",
    "    test_loss /= len(test_X) / batch_size\n",
    "    test_acc /= len(test_X) / batch_size\n",
    "\n",
    "    print('time taken:', time.time() - lasttime)\n",
    "    print(\n",
    "        'epoch: %d, training loss: %f, training acc: %f, valid loss: %f, valid acc: %f\\n'\n",
    "        % (e, train_loss, train_acc, test_loss, test_acc)\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "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": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "validation minibatch loop: 100%|██████████| 387/387 [01:06<00:00,  5.79it/s]\n"
     ]
    }
   ],
   "source": [
    "real_Y, predict_Y = [], []\n",
    "\n",
    "pbar = tqdm(\n",
    "    range(0, len(test_X), batch_size), desc = 'validation minibatch loop'\n",
    ")\n",
    "for i in pbar:\n",
    "    batch_x = test_X[i : min(i + batch_size, test_X.shape[0])]\n",
    "    batch_char = test_char[i : min(i + batch_size, test_X.shape[0])]\n",
    "    batch_y = test_Y[i : min(i + batch_size, test_X.shape[0])]\n",
    "    predicted = pred2label(sess.run(model.tags_seq,\n",
    "            feed_dict = {\n",
    "                model.word_ids: batch_x,\n",
    "                model.char_ids: batch_char,\n",
    "            },\n",
    "    ))\n",
    "    real = pred2label(np.argmax(batch_y, axis = 2))\n",
    "    predict_Y.extend(predicted)\n",
    "    real_Y.extend(real)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "              precision    recall  f1-score   support\n",
      "\n",
      "       OTHER       1.00      1.00      1.00    496576\n",
      "       event       0.98      0.98      0.98      2473\n",
      "         law       0.99      0.98      0.99      1809\n",
      "    location       0.98      0.99      0.98     20759\n",
      "organization       0.98      0.98      0.98     25598\n",
      "      person       0.98      0.99      0.99     43977\n",
      "    quantity       0.98      0.99      0.98     13171\n",
      "        time       0.99      0.98      0.98     13337\n",
      "\n",
      " avg / total       0.99      0.99      0.99    617700\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from sklearn.metrics import classification_report\n",
    "print(classification_report(np.array(real_Y).ravel(), np.array(predict_Y).ravel()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['Placeholder',\n",
       " 'Placeholder_1',\n",
       " 'Placeholder_2',\n",
       " 'Variable',\n",
       " 'Variable_1',\n",
       " 'char_0/Q/kernel',\n",
       " 'char_0/Q/bias',\n",
       " 'char_0/K_V/kernel',\n",
       " 'char_0/K_V/bias',\n",
       " 'char_0/gamma',\n",
       " 'char_feedforward_0/dense/kernel',\n",
       " 'char_feedforward_0/dense/bias',\n",
       " 'char_feedforward_0/dense_1/kernel',\n",
       " 'char_feedforward_0/dense_1/bias',\n",
       " 'char_feedforward_0/gamma',\n",
       " 'char_1/Q/kernel',\n",
       " 'char_1/Q/bias',\n",
       " 'char_1/K_V/kernel',\n",
       " 'char_1/K_V/bias',\n",
       " 'char_1/gamma',\n",
       " 'char_feedforward_1/dense/kernel',\n",
       " 'char_feedforward_1/dense/bias',\n",
       " 'char_feedforward_1/dense_1/kernel',\n",
       " 'char_feedforward_1/dense_1/bias',\n",
       " 'char_feedforward_1/gamma',\n",
       " 'dense/kernel',\n",
       " 'dense/bias',\n",
       " 'word_char_0/Q/kernel',\n",
       " 'word_char_0/Q/bias',\n",
       " 'word_char_0/K_V/kernel',\n",
       " 'word_char_0/K_V/bias',\n",
       " 'word_char_0/gamma',\n",
       " 'word_char_attention_0/Q/kernel',\n",
       " 'word_char_attention_0/Q/bias',\n",
       " 'word_char_attention_0/K_V/kernel',\n",
       " 'word_char_attention_0/K_V/bias',\n",
       " 'word_char_attention_0/gamma',\n",
       " 'word_feedforward_0/dense/kernel',\n",
       " 'word_feedforward_0/dense/bias',\n",
       " 'word_feedforward_0/dense_1/kernel',\n",
       " 'word_feedforward_0/dense_1/bias',\n",
       " 'word_feedforward_0/gamma',\n",
       " 'word_char_1/Q/kernel',\n",
       " 'word_char_1/Q/bias',\n",
       " 'word_char_1/K_V/kernel',\n",
       " 'word_char_1/K_V/bias',\n",
       " 'word_char_1/gamma',\n",
       " 'word_char_attention_1/Q/kernel',\n",
       " 'word_char_attention_1/Q/bias',\n",
       " 'word_char_attention_1/K_V/kernel',\n",
       " 'word_char_attention_1/K_V/bias',\n",
       " 'word_char_attention_1/gamma',\n",
       " 'word_feedforward_1/dense/kernel',\n",
       " 'word_feedforward_1/dense/bias',\n",
       " 'word_feedforward_1/dense_1/kernel',\n",
       " 'word_feedforward_1/dense_1/bias',\n",
       " 'word_feedforward_1/gamma',\n",
       " 'dense_1/kernel',\n",
       " 'dense_1/bias',\n",
       " 'transitions',\n",
       " 'logits']"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "saver = tf.train.Saver(tf.trainable_variables())\n",
    "saver.save(sess, 'attention/model.ckpt')\n",
    "\n",
    "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 'logits' in n.name\n",
    "        or 'alphas' in n.name\n",
    "        or 'state_fw' in n.name\n",
    "        or 'state_bw' in n.name)\n",
    "        and 'Adam' not in n.name\n",
    "        and 'beta' not in n.name\n",
    "        and 'OptimizeLoss' not in n.name\n",
    "        and 'Global_Step' not in n.name\n",
    "    ]\n",
    ")\n",
    "strings.split(',')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "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))\n",
    "        \n",
    "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Restoring parameters from attention/model.ckpt\n",
      "INFO:tensorflow:Froze 64 variables.\n",
      "INFO:tensorflow:Converted 64 variables to const ops.\n",
      "1317 ops in the final graph.\n"
     ]
    }
   ],
   "source": [
    "freeze_graph('attention', strings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "g = load_graph('attention/frozen_model.pb')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "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.'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "def char_str_idx(corpus, dic, UNK = 0):\n",
    "    maxlen = max([len(i) for i in corpus])\n",
    "    X = np.zeros((len(corpus), maxlen))\n",
    "    for i in range(len(corpus)):\n",
    "        for no, k in enumerate(corpus[i]):\n",
    "            X[i, no] = dic.get(k, UNK)\n",
    "    return X\n",
    "\n",
    "def generate_char_seq(batch, UNK = 2):\n",
    "    maxlen_c = max([len(k) for k in batch])\n",
    "    x = [[len(i) for i in k] for k in batch]\n",
    "    maxlen = max([j for i in x for j in i])\n",
    "    temp = np.zeros((len(batch),maxlen_c,maxlen),dtype=np.int32)\n",
    "    for i in range(len(batch)):\n",
    "        for k in range(len(batch[i])):\n",
    "            for no, c in enumerate(batch[i][k][::-1]):\n",
    "                temp[i,k,-1-no] = char2idx.get(c, UNK)\n",
    "    return temp\n",
    "\n",
    "sequence = process_string(string.lower())\n",
    "X_seq = char_str_idx([sequence], word2idx, 2)\n",
    "X_char_seq = generate_char_seq([sequence])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.6/dist-packages/tensorflow/python/client/session.py:1702: 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"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 kuala location\n",
      "1 lumpur location\n",
      "2 sempena OTHER\n",
      "3 sambutan OTHER\n",
      "4 aidilfitri event\n",
      "5 minggu OTHER\n",
      "6 depan OTHER\n",
      "7 perdana OTHER\n",
      "8 menteri OTHER\n",
      "9 tun person\n",
      "10 dr person\n",
      "11 mahathir person\n",
      "12 mohamad OTHER\n",
      "13 dan OTHER\n",
      "14 menteri OTHER\n",
      "15 pengangkutan organization\n",
      "16 anthony person\n",
      "17 loke person\n",
      "18 siew person\n",
      "19 fook person\n",
      "20 menitipkan person\n",
      "21 pesanan OTHER\n",
      "22 khas OTHER\n",
      "23 kepada OTHER\n",
      "24 orang OTHER\n",
      "25 ramai OTHER\n",
      "26 yang OTHER\n",
      "27 mahu OTHER\n",
      "28 pulang OTHER\n",
      "29 ke OTHER\n",
      "30 kampung OTHER\n",
      "31 halaman OTHER\n",
      "32 masing-masing OTHER\n",
      "33 dalam OTHER\n",
      "34 video OTHER\n",
      "35 pendek OTHER\n",
      "36 terbitan OTHER\n",
      "37 jabatan OTHER\n",
      "38 keselamatan OTHER\n",
      "39 jalan OTHER\n",
      "40 raya person\n",
      "41 jkjr person\n",
      "42 itu OTHER\n",
      "43 dr person\n",
      "44 mahathir person\n",
      "45 menasihati OTHER\n",
      "46 mereka OTHER\n",
      "47 supaya OTHER\n",
      "48 berhenti OTHER\n",
      "49 berehat OTHER\n",
      "50 dan OTHER\n",
      "51 tidur OTHER\n",
      "52 sebentar OTHER\n",
      "53 sekiranya OTHER\n",
      "54 mengantuk OTHER\n",
      "55 ketika OTHER\n",
      "56 memandu OTHER\n"
     ]
    }
   ],
   "source": [
    "word_ids = g.get_tensor_by_name('import/Placeholder:0')\n",
    "char_ids = g.get_tensor_by_name('import/Placeholder_1:0')\n",
    "tags_seq = g.get_tensor_by_name('import/logits:0')\n",
    "\n",
    "test_sess = tf.InteractiveSession(graph = g)\n",
    "predicted = test_sess.run(tags_seq,\n",
    "            feed_dict = {\n",
    "                word_ids: X_seq,\n",
    "                char_ids: X_char_seq,\n",
    "            })[0]\n",
    "\n",
    "for i in range(len(predicted)):\n",
    "    print(i, sequence[i],idx2tag[predicted[i]])"
   ]
  },
  {
   "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
}