RSN4EA.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import os, pickle\n",
    "from datetime import datetime\n",
    "\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import tensorflow as tf\n",
    "import scipy.sparse as sp\n",
    "from scipy.sparse import csr_matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "code_folding": []
   },
   "outputs": [],
   "source": [
    "# reader\n",
    "\n",
    "\n",
    "class BasicReader(object):\n",
    "\n",
    "    def read(self, data_path='data/dbp_wd_15k_V1/mapping/0_3/'):\n",
    "        # read KGs\n",
    "        def read_kb(path, names):\n",
    "            return pd.read_csv(path, sep='\\t', header=None, names=names)\n",
    "\n",
    "        kb1 = read_kb(data_path+'triples_1', names=['h_id', 'r_id', 't_id'])\n",
    "        kb2 = read_kb(data_path+'triples_2', names=['h_id', 'r_id', 't_id'])\n",
    "\n",
    "        ent_mapping = read_kb(data_path+'sup_ent_ids', names=['kb_1', 'kb_2'])\n",
    "        ent_testing = read_kb(data_path+'ref_ent_ids', names=['kb_1', 'kb_2'])\n",
    "\n",
    "        if not os.path.exists(data_path+'sup_rel_ids'):\n",
    "            os.mknod(data_path+'sup_rel_ids')\n",
    "        if not os.path.exists(data_path+'rel_rel_ids'):\n",
    "            os.mknod(data_path+'rel_rel_ids')\n",
    "\n",
    "        rel_mapping = read_kb(data_path+'sup_rel_ids', names=['kb_1', 'kb_2'])\n",
    "        rel_testing = read_kb(data_path+'rel_rel_ids', names=['kb_1', 'kb_2'])\n",
    "\n",
    "        ent_id_1 = read_kb(data_path+'ent_ids_1', names=['id', 'e'])\n",
    "        ent_id_2 = read_kb(data_path+'ent_ids_2', names=['id', 'e'])\n",
    "        ent_id_2.loc[:, 'e'] += ':KB2'\n",
    "        i2el_1 = pd.Series(ent_id_1.e.values, index=ent_id_1.id.values)\n",
    "        i2el_2 = pd.Series(ent_id_2.e.values, index=ent_id_2.id.values)\n",
    "\n",
    "        rel_id_1 = read_kb(data_path+'rel_ids_1', names=['id', 'r'])\n",
    "        rel_id_2 = read_kb(data_path+'rel_ids_2', names=['id', 'r'])\n",
    "        rel_id_2.loc[:, 'r'] += ':KB2'\n",
    "        i2rl_1 = pd.Series(rel_id_1.r.values, index=rel_id_1.id.values)\n",
    "        i2rl_2 = pd.Series(rel_id_2.r.values, index=rel_id_2.id.values)\n",
    "\n",
    "        # convert id\n",
    "        def id2label(df, i2el, i2rl, is_kb=True):\n",
    "            if is_kb:\n",
    "                df['h'] = i2el.loc[df.h_id.values].values\n",
    "                df['r'] = i2rl.loc[df.r_id.values].values\n",
    "                df['t'] = i2el.loc[df.t_id.values].values\n",
    "\n",
    "                return df\n",
    "            else:\n",
    "                df['kb_1'] = i2el.loc[df.kb_1.values].values\n",
    "                df['kb_2'] = i2rl.loc[df.kb_2.values].values\n",
    "\n",
    "                return df\n",
    "\n",
    "        id2label(kb1, i2el_1, i2rl_1)\n",
    "        id2label(kb2, i2el_2, i2rl_2)\n",
    "        id2label(ent_mapping, i2el_1, i2el_2, is_kb=False)\n",
    "        id2label(rel_mapping, i2rl_1, i2rl_2, is_kb=False)\n",
    "        id2label(ent_testing, i2el_1, i2el_2, is_kb=False)\n",
    "        id2label(rel_testing, i2rl_1, i2rl_2, is_kb=False)\n",
    "\n",
    "        # add reverse edges\n",
    "        kb = pd.concat([kb1, kb2], ignore_index=True)\n",
    "        kb = kb[['h', 'r', 't']]\n",
    "\n",
    "        rev_r = kb.r + ':reverse'\n",
    "        rev_kb = kb.rename(columns={'h': 't', 't': 'h'})\n",
    "        rev_kb['r'] = rev_r.values\n",
    "        kb = pd.concat([kb, rev_kb], ignore_index=True)\n",
    "\n",
    "        rev_rmap = rel_mapping + ':reverse'\n",
    "        rel_mapping = pd.concat([rel_mapping, rev_rmap], ignore_index=True)\n",
    "\n",
    "        # resort id in descending order of frequency, since we use log-uniform sampler for NCE loss\n",
    "        def remap_kb(kb):\n",
    "            es = pd.concat([kb.h, kb.t], ignore_index=True)\n",
    "            rs = kb.r\n",
    "            e_num = es.groupby(es.values).size().sort_values()[::-1]\n",
    "            r_num = rs.groupby(rs.values).size().sort_values()[::-1]\n",
    "\n",
    "            e_map = pd.Series(range(e_num.shape[0]), index=e_num.index)\n",
    "            r_map = pd.Series(range(r_num.shape[0]), index=r_num.index)\n",
    "\n",
    "            return e_map, r_map\n",
    "\n",
    "        def index(df, e_map, r_map, is_kb=True):\n",
    "            if is_kb:\n",
    "                df['h_id'] = e_map.loc[df.h.values].values\n",
    "                df['r_id'] = r_map.loc[df.r.values].values\n",
    "                df['t_id'] = e_map.loc[df.t.values].values\n",
    "            else:\n",
    "                df['kb_1'] = e_map.loc[df.kb_1.values].values\n",
    "                df['kb_2'] = e_map.loc[df.kb_2.values].values\n",
    "\n",
    "        e_map, r_map = remap_kb(kb)\n",
    "\n",
    "        index(kb, e_map, r_map)\n",
    "        index(ent_mapping, e_map, None, is_kb=False)\n",
    "        index(ent_testing, e_map, None, is_kb=False)\n",
    "        index(rel_mapping, r_map, None, is_kb=False)\n",
    "        index(rel_testing, r_map, None, is_kb=False)\n",
    "\n",
    "        index(kb1, e_map, r_map)\n",
    "        index(kb2, e_map, r_map)\n",
    "        eid_1 = pd.unique(pd.concat([kb1.h_id, kb1.t_id], ignore_index=True))\n",
    "        eid_2 = pd.unique(pd.concat([kb2.h_id, kb2.t_id], ignore_index=True))\n",
    "        \n",
    "        \n",
    "        # add shortcuts\n",
    "        self._eid_1 = pd.Series(eid_1)\n",
    "        self._eid_2 = pd.Series(eid_2)\n",
    "\n",
    "        self._ent_num = len(e_map)\n",
    "        self._rel_num = len(r_map)\n",
    "        self._ent_id = e_map\n",
    "        self._rel_id = r_map\n",
    "\n",
    "        self._ent_mapping = ent_mapping\n",
    "        self._rel_mapping = rel_mapping\n",
    "        self._ent_testing = ent_testing\n",
    "        self._rel_testing = rel_testing\n",
    "        \n",
    "        \n",
    "        \n",
    "\n",
    "        self._kb = kb\n",
    "        # we first tag the entities that have algined entities according to entity_mapping\n",
    "        self.add_align_infor()\n",
    "        # we then connect two KGs by creating new triples involving aligned entities.\n",
    "        self.add_weight()\n",
    "    \n",
    "    def add_align_infor(self):\n",
    "        kb = self._kb\n",
    "        \n",
    "        ent_mapping = self._ent_mapping\n",
    "        rev_e_m = ent_mapping.rename(columns={'kb_1': 'kb_2', 'kb_2': 'kb_1'})\n",
    "        rel_mapping = self._rel_mapping\n",
    "        rev_r_m = rel_mapping.rename(columns={'kb_1': 'kb_2', 'kb_2': 'kb_1'})\n",
    "        \n",
    "        ent_mapping = pd.concat([ent_mapping, rev_e_m], ignore_index=True)\n",
    "        rel_mapping = pd.concat([rel_mapping, rev_r_m], ignore_index=True)\n",
    "        \n",
    "        ent_mapping = pd.Series(ent_mapping.kb_2.values, index=ent_mapping.kb_1.values)\n",
    "        rel_mapping = pd.Series(rel_mapping.kb_2.values, index=rel_mapping.kb_1.values)\n",
    "        \n",
    "        self._e_m = ent_mapping\n",
    "        self._r_m = rel_mapping\n",
    "        \n",
    "        kb['ah_id'] = kb.h_id\n",
    "        kb['ar_id'] = kb.r_id\n",
    "        kb['at_id'] = kb.t_id\n",
    "        \n",
    "        h_mask = kb.h_id.isin(ent_mapping)\n",
    "        r_mask = kb.r_id.isin(rel_mapping)\n",
    "        t_mask = kb.t_id.isin(ent_mapping)\n",
    "        \n",
    "        kb['ah_id'][h_mask] = ent_mapping.loc[kb['ah_id'][h_mask].values]\n",
    "        kb['ar_id'][r_mask] = rel_mapping.loc[kb['ar_id'][r_mask].values]\n",
    "        kb['at_id'][t_mask] = ent_mapping.loc[kb['at_id'][t_mask].values]\n",
    "        \n",
    "        self._kb = kb\n",
    "        \n",
    "    def add_weight(self):\n",
    "        kb = self._kb[['h_id', 'r_id', 't_id', 'ah_id', 'ar_id', 'at_id']]\n",
    "\n",
    "        kb['w_h'] = 0\n",
    "        kb['w_r'] = 0\n",
    "        kb['w_t'] = 0\n",
    "\n",
    "\n",
    "        h_mask = ~(kb.h_id == kb.ah_id)\n",
    "        r_mask = ~(kb.r_id == kb.ar_id)\n",
    "        t_mask = ~(kb.t_id == kb.at_id)\n",
    "\n",
    "        kb.loc[h_mask, 'w_h'] = 1\n",
    "        kb.loc[r_mask, 'w_r'] = 1\n",
    "        kb.loc[t_mask, 'w_t'] = 1\n",
    "\n",
    "        akb = kb[['ah_id','ar_id','at_id', 'w_h', 'w_r', 'w_t']]\n",
    "        akb = akb.rename(columns={'ah_id':'h_id','ar_id':'r_id','at_id':'t_id'})\n",
    "\n",
    "        ahkb = kb[h_mask][['ah_id','r_id','t_id', 'w_h', 'w_r', 'w_t']].rename(columns={'ah_id':'h_id'})\n",
    "        arkb = kb[r_mask][['h_id','ar_id','t_id', 'w_h', 'w_r', 'w_t']].rename(columns={'ar_id':'r_id'})\n",
    "        atkb = kb[t_mask][['h_id','r_id','at_id', 'w_h', 'w_r', 'w_t']].rename(columns={'at_id':'t_id'})\n",
    "        ahrkb = kb[h_mask&r_mask][['ah_id','ar_id','t_id', 'w_h', 'w_r', 'w_t']].rename(columns={'ah_id':'h_id', 'ar_id':'r_id'})\n",
    "        ahtkb = kb[h_mask&t_mask][['ah_id','r_id','at_id', 'w_h', 'w_r', 'w_t']].rename(columns={'ah_id':'h_id', 'at_id':'t_id'})\n",
    "        artkb = kb[r_mask&t_mask][['h_id','ar_id','at_id', 'w_h', 'w_r', 'w_t']].rename(columns={'ar_id':'r_id', 'at_id':'t_id'})\n",
    "        ahrtkb = kb[h_mask&r_mask&t_mask][['ah_id','ar_id','at_id', 'w_h', 'w_r', 'w_t']].rename(columns={'ah_id':'h_id',\n",
    "                                                                                                          'ar_id':'r_id',\n",
    "                                                                                                          'at_id':'t_id'})\n",
    "\n",
    "        kb['w_h'] = 0\n",
    "        kb['w_r'] = 0\n",
    "        kb['w_t'] = 0\n",
    "\n",
    "        kb = pd.concat([akb, ahkb, arkb, atkb, ahrkb, ahtkb, artkb, ahrtkb, kb[['h_id','r_id','t_id', 'w_h', 'w_r', 'w_t']]],\n",
    "                       ignore_index=True).drop_duplicates()\n",
    "\n",
    "        self._kb = kb.reset_index(drop=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "code_folding": []
   },
   "outputs": [],
   "source": [
    "# sampler\n",
    "\n",
    "\n",
    "class BasicSampler(object):\n",
    "\n",
    "    def sample_paths(self, repeat_times=2):\n",
    "        opts = self._options\n",
    "\n",
    "        kb = self._kb.copy()\n",
    "\n",
    "        kb = kb[['h_id', 'r_id', 't_id']]\n",
    "\n",
    "        # sampling paths in the h_id-(r_id,t_id) form.\n",
    "\n",
    "        rtlist = np.unique(kb[['r_id', 't_id']].values, axis=0)\n",
    "\n",
    "        rtdf = pd.DataFrame(rtlist, columns=['r_id', 't_id'])\n",
    "        \n",
    "        # assign tail=(r_id, t_id), we assign an id for each tail\n",
    "        rtdf = rtdf.reset_index().rename({'index': 'tail_id'}, axis='columns')\n",
    "        \n",
    "        # merge kb with rtdf, to get the (h_id, tail_id) dataframe\n",
    "        rtkb = kb.merge(\n",
    "            rtdf, left_on=['r_id', 't_id'], right_on=['r_id', 't_id'])\n",
    "\n",
    "        htail = np.unique(rtkb[['h_id', 'tail_id']].values, axis=0)\n",
    "        \n",
    "        # save to the sparse matrix\n",
    "        htailmat = csr_matrix((np.ones(len(htail)), (htail[:, 0], htail[:, 1])),\n",
    "                              shape=(model._ent_num, rtlist.shape[0]))\n",
    "\n",
    "        # calulate corss-KG bias at first, note that we use an approximate method: \n",
    "        # if next entity e_{i+1} is in entity_mapping, e_i and e_{i+2} entity are believed in different KGs\n",
    "        em = pd.concat(\n",
    "            [model._ent_mapping.kb_1, model._ent_mapping.kb_2]).values\n",
    "\n",
    "        rtkb['across'] = rtkb.t_id.isin(em)\n",
    "        rtkb.loc[rtkb.across, 'across'] = opts.beta\n",
    "        rtkb.loc[rtkb.across == 0, 'across'] = 1-opts.beta\n",
    "\n",
    "        rtailkb = rtkb[['h_id', 't_id', 'tail_id', 'across']]\n",
    "        \n",
    "        def gen_tail_dict(x):\n",
    "            return x.tail_id.values, x.across.values / x.across.sum()\n",
    "        \n",
    "        # each item in rtailkb is in the form of (tail_ids, cross-KG biases)\n",
    "        rtailkb = rtailkb.groupby('h_id').apply(gen_tail_dict)\n",
    "\n",
    "        rtailkb = pd.DataFrame({'tails': rtailkb})\n",
    "\n",
    "        # start sampling\n",
    "\n",
    "        hrt = np.repeat(kb.values, repeat_times, axis=0)\n",
    "\n",
    "        # for initial triples\n",
    "        def perform_random(x):\n",
    "            return np.random.choice(x.tails[0], 1, p=x.tails[1].astype(np.float))\n",
    "\n",
    "        # else\n",
    "        def perform_random2(x):\n",
    "\n",
    "            # calculate depth bias\n",
    "            pre_c = htailmat[np.repeat(x.pre, x.tails[0].shape[0]), x.tails[0]]\n",
    "            \n",
    "            pre_c[pre_c == 0] = opts.alpha\n",
    "            pre_c[pre_c == 1] = 1-opts.alpha\n",
    "            \n",
    "            # combine the biases\n",
    "            p = x.tails[1].astype(np.float).reshape(\n",
    "                [-1, ]) * pre_c.A.reshape([-1, ])\n",
    "            p = p / p.sum()\n",
    "            return np.random.choice(x.tails[0], 1, p=p)\n",
    "\n",
    "        rt_x = rtailkb.loc[hrt[:, 2]].apply(perform_random, axis=1)\n",
    "        rt_x = rtlist[np.concatenate(rt_x.values)]\n",
    "\n",
    "        rts = [hrt, rt_x]\n",
    "        c_length = 5\n",
    "        pre = hrt[:, 0]\n",
    "        print('current path length == %i' % c_length)\n",
    "        while(c_length < opts.max_length):\n",
    "            curr = rtailkb.loc[rt_x[:, 1]]\n",
    "            \n",
    "            # always using hrt[:, 0] as the previous entity is a stronger way to\n",
    "            # generate deeper and cross-KG paths for the starting point. \n",
    "            # use 'curr.loc[:, 'pre'] = pre' for 2nd-order sampling.\n",
    "            curr.loc[:, 'pre'] = hrt[:, 0]\n",
    "\n",
    "            rt_x = curr.apply(perform_random2, axis=1)\n",
    "            rt_x = rtlist[np.concatenate(rt_x.values)]\n",
    "\n",
    "            rts.append(rt_x)\n",
    "            c_length += 2\n",
    "            # pre = curr.index.values\n",
    "            print('current path length == %i' % c_length)\n",
    "            \n",
    "        data = np.concatenate(rts, axis=1)\n",
    "        data = pd.DataFrame(data)\n",
    "        \n",
    "        self._train_data = data\n",
    "        data.to_csv('%spaths_%.1f_%.1f' % (opts.data_path, opts.alpha, opts.beta))\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "code_folding": []
   },
   "outputs": [],
   "source": [
    "# model\n",
    "class RSN4EA(BasicReader, BasicSampler):\n",
    "    def __init__(self, options, session):\n",
    "        self._options = options\n",
    "        self._session = session\n",
    "\n",
    "    def create_variables(self):\n",
    "        options = self._options\n",
    "        hidden_size = options.hidden_size\n",
    "\n",
    "        self._entity_embedding = tf.get_variable(\n",
    "            'entity_embedding',\n",
    "            [self._ent_num, hidden_size],\n",
    "            initializer=tf.contrib.layers.xavier_initializer(uniform=False)\n",
    "        )\n",
    "        self._relation_embedding = tf.get_variable(\n",
    "            'relation_embedding',\n",
    "            [self._rel_num, hidden_size],\n",
    "            initializer=tf.contrib.layers.xavier_initializer(uniform=False)\n",
    "        )\n",
    "\n",
    "        self._rel_w = tf.get_variable(\n",
    "            \"relation_softmax_w\",\n",
    "            [self._rel_num, hidden_size],\n",
    "            initializer=tf.contrib.layers.xavier_initializer(uniform=False)\n",
    "        )\n",
    "        self._rel_b = tf.get_variable(\n",
    "            \"relation_softmax_b\",\n",
    "            [self._rel_num],\n",
    "            initializer=tf.constant_initializer(0)\n",
    "        )\n",
    "        self._ent_w = tf.get_variable(\n",
    "            \"entity_softmax_w\",\n",
    "            [self._ent_num, hidden_size],\n",
    "            initializer=tf.contrib.layers.xavier_initializer(uniform=False)\n",
    "        )\n",
    "        self._ent_b = tf.get_variable(\n",
    "            \"entity_softmax_b\",\n",
    "            [self._ent_num],\n",
    "            initializer=tf.constant_initializer(0)\n",
    "        )\n",
    "        self._lr = tf.Variable(options.learning_rate, trainable=False)\n",
    "\n",
    "        self._optimizer = tf.train.AdamOptimizer(options.learning_rate)\n",
    "\n",
    "    def bn(self, inputs, is_train=True, reuse=True):\n",
    "        return tf.contrib.layers.batch_norm(inputs,\n",
    "                                            center=True,\n",
    "                                            scale=True,\n",
    "                                            is_training=is_train,\n",
    "                                            reuse=reuse,\n",
    "                                            scope='bn',\n",
    "                                            )\n",
    "\n",
    "    def lstm_cell(self, drop=True, keep_prob=0.5, num_layers=2, hidden_size=None):\n",
    "        if not hidden_size:\n",
    "            hidden_size = self._options.hidden_size\n",
    "\n",
    "        def basic_lstm_cell():\n",
    "            return tf.contrib.rnn.LSTMCell(\n",
    "                num_units=hidden_size,\n",
    "                initializer=tf.orthogonal_initializer,\n",
    "                forget_bias=1,\n",
    "                reuse=tf.get_variable_scope().reuse,\n",
    "                activation=tf.identity\n",
    "            )\n",
    "\n",
    "        def drop_cell():\n",
    "            return tf.contrib.rnn.DropoutWrapper(\n",
    "                basic_lstm_cell(),\n",
    "                output_keep_prob=keep_prob\n",
    "            )\n",
    "\n",
    "        if drop:\n",
    "            gen_cell = drop_cell\n",
    "        else:\n",
    "            gen_cell = basic_lstm_cell\n",
    "\n",
    "        if num_layers == 0:\n",
    "            return gen_cell()\n",
    "\n",
    "        cell = tf.contrib.rnn.MultiRNNCell(\n",
    "            [gen_cell() for _ in range(num_layers)],\n",
    "            state_is_tuple=True,\n",
    "        )\n",
    "        return cell\n",
    "\n",
    "    def sampled_loss(self, inputs, labels, w, b, weight=1, is_entity=False):\n",
    "        num_sampled = min(self._options.num_samples, w.shape[0]//3)\n",
    "\n",
    "        labels = tf.reshape(labels, [-1, 1])\n",
    "\n",
    "        losses = tf.nn.nce_loss(\n",
    "            weights=w,\n",
    "            biases=b,\n",
    "            labels=labels,\n",
    "            inputs=tf.reshape(inputs, [-1, int(w.shape[1])]),\n",
    "            num_sampled=num_sampled,\n",
    "            num_classes=w.shape[0],\n",
    "            partition_strategy='div',\n",
    "        )\n",
    "        return losses * weight\n",
    "\n",
    "    def logits(self, inputs, w, b):\n",
    "        return tf.nn.bias_add(tf.matmul(inputs, tf.transpose(w)), b)\n",
    "\n",
    "    # shuffle data\n",
    "    def sample(self, data):\n",
    "        choices = np.random.choice(len(data), size=len(data), replace=False)\n",
    "        return data.iloc[choices]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "code_folding": []
   },
   "outputs": [],
   "source": [
    "# build tensorflow graph\n",
    "\n",
    "\n",
    "# build an RSN of length l\n",
    "def build_sub_graph(self, length=15, reuse=False):\n",
    "    options = self._options\n",
    "    hidden_size = options.hidden_size\n",
    "    batch_size = options.batch_size\n",
    "\n",
    "    seq = tf.placeholder(\n",
    "        tf.int32, [batch_size, length], name='seq'+str(length))\n",
    "\n",
    "    e_em, r_em = self._entity_embedding, self._relation_embedding\n",
    "\n",
    "    # seperately read, and then recover the order\n",
    "    ent = seq[:, :-1:2]\n",
    "    rel = seq[:, 1::2]\n",
    "\n",
    "    ent_em = tf.nn.embedding_lookup(e_em, ent)\n",
    "    rel_em = tf.nn.embedding_lookup(r_em, rel)\n",
    "\n",
    "    em_seq = []\n",
    "    for i in range(length-1):\n",
    "        if i % 2 == 0:\n",
    "            em_seq.append(ent_em[:, i//2])\n",
    "        else:\n",
    "            em_seq.append(rel_em[:, i//2])\n",
    "\n",
    "    # seperately bn\n",
    "    with tf.variable_scope('input_bn'):\n",
    "        if not reuse:\n",
    "            bn_em_seq = [tf.reshape(self.bn(em_seq[i], reuse=(\n",
    "                i is not 0)), [-1, 1, hidden_size]) for i in range(length-1)]\n",
    "        else:\n",
    "            bn_em_seq = [tf.reshape(\n",
    "                self.bn(em_seq[i], reuse=True), [-1, 1, hidden_size]) for i in range(length-1)]\n",
    "\n",
    "    bn_em_seq = tf.concat(bn_em_seq, axis=1)\n",
    "\n",
    "    ent_bn_em = bn_em_seq[:, ::2]\n",
    "\n",
    "    with tf.variable_scope('rnn', reuse=reuse):\n",
    "\n",
    "        cell = self.lstm_cell(True, options.keep_prob, options.num_layers)\n",
    "\n",
    "        outputs, state = tf.nn.dynamic_rnn(cell, bn_em_seq,  dtype=tf.float32)\n",
    "\n",
    "    rel_outputs = outputs[:, 1::2, :]\n",
    "    outputs = [outputs[:, i, :] for i in range(length-1)]\n",
    "\n",
    "    ent_outputs = outputs[::2]\n",
    "\n",
    "    # RSN\n",
    "    res_rel_outputs = tf.contrib.layers.fully_connected(rel_outputs, hidden_size, biases_initializer=None, activation_fn=None) +\\\n",
    "        tf.contrib.layers.fully_connected(\n",
    "            ent_bn_em, hidden_size, biases_initializer=None, activation_fn=None)\n",
    "\n",
    "    # recover the order\n",
    "    res_rel_outputs = [res_rel_outputs[:, i, :] for i in range((length-1)//2)]\n",
    "    outputs = []\n",
    "    for i in range(length-1):\n",
    "        if i % 2 == 0:\n",
    "            outputs.append(ent_outputs[i//2])\n",
    "        else:\n",
    "            outputs.append(res_rel_outputs[i//2])\n",
    "\n",
    "    # output bn\n",
    "    with tf.variable_scope('output_bn'):\n",
    "        if reuse:\n",
    "            bn_outputs = [tf.reshape(\n",
    "                self.bn(outputs[i], reuse=True), [-1, 1, hidden_size]) for i in range(length-1)]\n",
    "        else:\n",
    "            bn_outputs = [tf.reshape(self.bn(outputs[i], reuse=(\n",
    "                i is not 0)), [-1, 1, hidden_size]) for i in range(length-1)]\n",
    "\n",
    "    def cal_loss(bn_outputs, seq):\n",
    "        losses = []\n",
    "\n",
    "        decay = 0.8\n",
    "        for i, output in enumerate(bn_outputs):\n",
    "            if i % 2 == 0:\n",
    "                losses.append(self.sampled_loss(\n",
    "                    output, seq[:, i+1], self._rel_w, self._rel_b, weight=decay**(0), is_entity=i))\n",
    "            else:\n",
    "                losses.append(self.sampled_loss(\n",
    "                    output, seq[:, i+1], self._ent_w, self._ent_b, weight=decay**(0), is_entity=i))\n",
    "        losses = tf.stack(losses, axis=1)\n",
    "        return losses\n",
    "\n",
    "    seq_loss = cal_loss(bn_outputs, seq)\n",
    "\n",
    "    losses = tf.reduce_sum(seq_loss) / batch_size\n",
    "\n",
    "    return losses, seq\n",
    "\n",
    "\n",
    "# build the main graph\n",
    "def build_graph(self):\n",
    "    options = self._options\n",
    "\n",
    "    loss, seq = build_sub_graph(self, length=options.max_length, reuse=False)\n",
    "\n",
    "    tvars = tf.trainable_variables()\n",
    "    grads, _ = tf.clip_by_global_norm(tf.gradients(loss, tvars), 2.0)\n",
    "    update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)\n",
    "    with tf.control_dependencies(update_ops):\n",
    "        train_op = self._optimizer.apply_gradients(\n",
    "            zip(grads, tvars),\n",
    "            global_step=tf.train.get_or_create_global_step()\n",
    "        )\n",
    "\n",
    "    self._seq, self._loss, self._train_op = seq, loss, train_op"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "code_folding": []
   },
   "outputs": [],
   "source": [
    "# training procedure\n",
    "def seq_train(self, data, choices=None, epoch=None):\n",
    "    opts = self._options\n",
    "    \n",
    "    # shuffle data\n",
    "    choices = np.random.choice(len(data), size=len(data), replace=True)\n",
    "    batch_size = opts.batch_size\n",
    "    \n",
    "    num_batch = len(data) // batch_size\n",
    "    \n",
    "    fetches = {\n",
    "        'loss': self._loss,\n",
    "        'train_op': self._train_op\n",
    "        }\n",
    "    \n",
    "    losses = 0 \n",
    "    for i in range(num_batch):\n",
    "        \n",
    "        one_batch_choices = choices[i * batch_size : (i + 1) * batch_size]\n",
    "        one_batch_data = data.iloc[one_batch_choices]\n",
    "\n",
    "        feed_dict = {}\n",
    "        seq = one_batch_data.values[:, :opts.max_length]\n",
    "        feed_dict[self._seq] = seq\n",
    "\n",
    "        vals = self._session.run(fetches, feed_dict)\n",
    "\n",
    "        del one_batch_data\n",
    "\n",
    "        loss = vals['loss']\n",
    "        losses += loss\n",
    "        print('\\r%i/%i, batch_loss:%f' % (i, num_batch, loss), end='')\n",
    "    self._last_mean_loss = losses / num_batch\n",
    "\n",
    "    return self._last_mean_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# eval_graph & eval method\n",
    "\n",
    "\n",
    "def build_eval_graph(self, entity=True):\n",
    "    options = self._options\n",
    "    hidden_size = options.hidden_size\n",
    "    batch_size = 2048\n",
    "\n",
    "    e_em, r_em = self._entity_embedding, self._relation_embedding\n",
    "\n",
    "    def em_lookup(indices, em):\n",
    "        return tf.nn.embedding_lookup(em, indices)\n",
    "\n",
    "    h, r = tf.placeholder(tf.int32, [None]), tf.placeholder(tf.int32, [None])\n",
    "\n",
    "    he, re = em_lookup(h, e_em), em_lookup(r, r_em)\n",
    "\n",
    "    he = tf.nn.l2_normalize(he, dim=-1)\n",
    "    norm_e_em = tf.nn.l2_normalize(e_em, dim=-1)\n",
    "\n",
    "    re = tf.nn.l2_normalize(re, dim=-1)\n",
    "    norm_r_em = tf.nn.l2_normalize(r_em, dim=-1)\n",
    "\n",
    "    aep = tf.matmul(he, tf.transpose(norm_e_em))\n",
    "    arp = tf.matmul(re, tf.transpose(norm_r_em))\n",
    "\n",
    "    if entity:\n",
    "        return h, aep\n",
    "    else:\n",
    "        return r, arp\n",
    "\n",
    "\n",
    "def eval_entity_align(model, data, kb_1to2=False):\n",
    "    options = model._options\n",
    "    batch_size = 16\n",
    "\n",
    "    data, padding_num = padding_data(data, options, batch_size)\n",
    "\n",
    "    h, aep = build_eval_graph(model)\n",
    "\n",
    "    fetch = {'probs': aep, }\n",
    "\n",
    "    num_batch = len(data) // batch_size\n",
    "\n",
    "    probs = []\n",
    "    for i in range(num_batch):\n",
    "        one_batch_data = data.iloc[i * batch_size:(i + 1) * batch_size]\n",
    "\n",
    "        feed_dict = {}\n",
    "        if kb_1to2:\n",
    "            feed_dict[h] = one_batch_data.kb_1.values\n",
    "        else:\n",
    "            feed_dict[h] = one_batch_data.kb_2.values\n",
    "\n",
    "        vals = sess.run(fetch, feed_dict)\n",
    "        probs.append(vals['probs'])\n",
    "\n",
    "    probs = np.concatenate(probs)[:len(data) - padding_num]\n",
    "    return probs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "#some tools\n",
    "\n",
    "def cal_ranks(probs, method, label):\n",
    "    # in most cases, using method=='min' is safe and much faster than method=='average' or 'max'.\n",
    "    # but note that, it will overestimate the results if the correct one has a same probability with others. \n",
    "    if method == 'min':\n",
    "        probs = probs - probs[range(len(label)), label].reshape(len(probs), 1)\n",
    "        ranks = (probs > 0).sum(axis=1) + 1\n",
    "    else:\n",
    "        ranks = pd.DataFrame(probs).rank(axis=1, ascending=False, method=method)\n",
    "        ranks = ranks.values[range(len(label)), label]\n",
    "    return ranks\n",
    "\n",
    "#top-10 = hits@10\n",
    "def cal_performance(ranks, top=10):\n",
    "    m_r = sum(ranks) * 1.0 / len(ranks)\n",
    "    h_10 = sum(ranks <= top) * 1.0 / len(ranks)\n",
    "    mrr = (1. / ranks).sum() / len(ranks)\n",
    "    return m_r, h_10, mrr\n",
    "\n",
    "def padding_data(data, options, batch_size):\n",
    "    padding_num = batch_size - len(data) % batch_size\n",
    "    data = pd.concat([data, pd.DataFrame(np.zeros((padding_num, data.shape[1])), dtype=np.int32, columns=data.columns)],ignore_index=True, axis=0)\n",
    "    return data, padding_num\n",
    "\n",
    "def in2d(arr1, arr2):\n",
    "    \"\"\"Generalisation of numpy.in1d to 2D arrays\"\"\"\n",
    "\n",
    "    assert arr1.dtype == arr2.dtype\n",
    "\n",
    "    arr1_view = np.ascontiguousarray(arr1).view(np.dtype((np.void,\n",
    "                                                          arr1.dtype.itemsize * arr1.shape[1])))\n",
    "    arr2_view = np.ascontiguousarray(arr2).view(np.dtype((np.void,\n",
    "                                                          arr2.dtype.itemsize * arr2.shape[1])))\n",
    "    intersected = np.in1d(arr1_view, arr2_view)\n",
    "    return intersected.view(np.bool).reshape(-1)\n",
    "\n",
    "\n",
    "\n",
    "#handle evaluation\n",
    "def handle_evaluation(i=0, last_mean_loss=0, kb_1to2=True, method='min', valid=True):\n",
    "    data_size = len(model._ent_testing)\n",
    "    #we use 10% testing data for validation\n",
    "    if valid:\n",
    "        data = model._ent_testing.iloc[:data_size//10]\n",
    "    else:\n",
    "        data = model._ent_testing.iloc[data_size//10:]\n",
    "    \n",
    "    probs=eval_entity_align(model, data, kb_1to2=kb_1to2)\n",
    "    candi = model._ent_testing.kb_2 if kb_1to2==True else model._ent_testing.kb_1\n",
    "    mask = np.in1d(np.arange(probs.shape[1]), candi)\n",
    "    #exclude known entities\n",
    "    probs[:, ~mask] = probs.min() -1\n",
    "    \n",
    "    \n",
    "    label=data.kb_2 if kb_1to2==True else data.kb_1                  \n",
    "    ranks = cal_ranks(probs, method=method,\n",
    "                          label=label)\n",
    "                          \n",
    "    MR, H10, MRR = cal_performance(ranks, top=10)\n",
    "    _, H1, _ = cal_performance(ranks, top=1)\n",
    "    H1, MR, H10, MRR\n",
    "\n",
    "    msg = 'epoch:%i, Hits@1:%.3f, Hits@10:%.3f, MR:%.3f, MRR:%.3f, mean_loss:%.3f' % (i, H1, H10, MR, MRR, last_mean_loss)\n",
    "    print('\\n'+msg)\n",
    "    return msg, (i, H1, H10, MR, MRR, last_mean_loss)\n",
    "\n",
    "def write_to_log(path, content):\n",
    "    with open(path, 'a+') as f:\n",
    "        print(content, file=f)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Options(object):\n",
    "    pass\n",
    "\n",
    "\n",
    "# set options\n",
    "opts = Options()\n",
    "opts.hidden_size = 256\n",
    "opts.num_layers = 2\n",
    "opts.batch_size = 512\n",
    "opts.learning_rate = 0.003\n",
    "opts.num_samples = 2048*4\n",
    "opts.keep_prob = 0.5\n",
    "\n",
    "opts.max_length = 15\n",
    "opts.alpha = 0.9\n",
    "opts.beta = 0.9\n",
    "\n",
    "\n",
    "opts.data_path = 'data/dbp_wd_15k_V1/mapping/0_3/'\n",
    "\n",
    "opts.log_file_path = 'logs/%s%dl_%s.log' % (opts.data_path.replace(\n",
    "    '/', '-'), opts.max_length, datetime.now().strftime('%y-%m-%d-%H-%M'))\n",
    "\n",
    "# and tensorflow config\n",
    "config = tf.ConfigProto()\n",
    "config.gpu_options.allow_growth = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#initial model\n",
    "\n",
    "sess = tf.InteractiveSession(config=config)\n",
    "\n",
    "model = RSN4EA(options=opts, session=sess)\n",
    "\n",
    "model.read(data_path=model._options.data_path)\n",
    "model.create_variables()\n",
    "\n",
    "sequence_datapath = '%spaths_%.1f_%.1f' % (\n",
    "    model._options.data_path, model._options.alpha, model._options.beta)\n",
    "\n",
    "if not os.path.exists(sequence_datapath):\n",
    "    print('start to sample paths')\n",
    "    model.sample_paths()\n",
    "    train_data = model._train_data\n",
    "else:\n",
    "    print('load existing training sequences')\n",
    "    train_data = pd.read_csv(sequence_datapath, index_col=0)\n",
    "\n",
    "\n",
    "# build tensorflow graph and init all tensors\n",
    "build_graph(model)\n",
    "tf.global_variables_initializer().run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "#initial training settings\n",
    "\n",
    "write_to_log(opts.log_file_path, opts.__dict__)\n",
    "max_hits1, times, max_times = 0, 0, 3\n",
    "epoch = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "code_folding": [],
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "epoch:0, Hits@1:0.000, Hits@10:0.000, MR:5097.818, MRR:0.001, mean_loss:10000.000\n",
      "1037/1038, batch_loss:39.2622530\n",
      "epoch:1, Hits@1:0.293, Hits@10:0.545, MR:284.197, MRR:0.382, mean_loss:2650.895\n",
      "1037/1038, batch_loss:36.781380\n",
      "epoch:2, Hits@1:0.320, Hits@10:0.594, MR:216.635, MRR:0.416, mean_loss:37.672\n",
      "1037/1038, batch_loss:35.845501\n",
      "epoch:3, Hits@1:0.355, Hits@10:0.617, MR:213.297, MRR:0.448, mean_loss:36.325\n",
      "1037/1038, batch_loss:35.401333\n",
      "epoch:4, Hits@1:0.361, Hits@10:0.630, MR:197.111, MRR:0.455, mean_loss:35.622\n",
      "1037/1038, batch_loss:35.369408\n",
      "epoch:5, Hits@1:0.360, Hits@10:0.636, MR:194.842, MRR:0.456, mean_loss:35.226\n",
      "1037/1038, batch_loss:34.809410\n",
      "epoch:6, Hits@1:0.367, Hits@10:0.633, MR:195.310, MRR:0.458, mean_loss:34.915\n",
      "1037/1038, batch_loss:35.046421\n",
      "epoch:7, Hits@1:0.367, Hits@10:0.636, MR:188.303, MRR:0.460, mean_loss:34.675\n",
      "1037/1038, batch_loss:34.469902\n",
      "epoch:8, Hits@1:0.372, Hits@10:0.648, MR:188.146, MRR:0.464, mean_loss:34.493\n",
      "1037/1038, batch_loss:34.195232\n",
      "epoch:9, Hits@1:0.373, Hits@10:0.648, MR:185.349, MRR:0.467, mean_loss:34.339\n",
      "1037/1038, batch_loss:34.462944\n",
      "epoch:10, Hits@1:0.366, Hits@10:0.651, MR:175.892, MRR:0.464, mean_loss:34.221\n",
      "1037/1038, batch_loss:34.001343\n",
      "epoch:11, Hits@1:0.384, Hits@10:0.651, MR:175.564, MRR:0.472, mean_loss:34.101\n",
      "1037/1038, batch_loss:33.996300\n",
      "epoch:12, Hits@1:0.373, Hits@10:0.642, MR:187.588, MRR:0.463, mean_loss:34.013\n",
      "1037/1038, batch_loss:33.366573\n",
      "epoch:13, Hits@1:0.370, Hits@10:0.650, MR:195.128, MRR:0.465, mean_loss:33.909\n",
      "1037/1038, batch_loss:33.479221\n",
      "epoch:14, Hits@1:0.369, Hits@10:0.645, MR:183.848, MRR:0.464, mean_loss:33.871\n",
      "\n",
      "epoch:14, Hits@1:0.390, Hits@10:0.654, MR:186.533, MRR:0.477, mean_loss:33.871\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "('epoch:14, Hits@1:0.390, Hits@10:0.654, MR:186.533, MRR:0.477, mean_loss:33.871',\n",
       " (14,\n",
       "  0.39026455026455026,\n",
       "  0.6535449735449735,\n",
       "  186.5331216931217,\n",
       "  0.4769388660661274,\n",
       "  33.87092417972396))"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "msg, r = handle_evaluation(0, 10000, valid=True)\n",
    "write_to_log(opts.log_file_path, msg)\n",
    "\n",
    "for i in range(epoch, 30):\n",
    "    last_mean_loss = seq_train(model, train_data)\n",
    "    epoch += 1\n",
    "\n",
    "    # evaluation\n",
    "    msg, r = handle_evaluation(i+1, last_mean_loss, valid=True)\n",
    "    write_to_log(opts.log_file_path, msg)\n",
    "\n",
    "    # early stop\n",
    "    hits1 = r[1]\n",
    "    if hits1 > max_hits1:\n",
    "        max_hits1 = hits1\n",
    "        times = 0\n",
    "    else:\n",
    "        times += 1\n",
    "\n",
    "    if times >= max_times:\n",
    "        break\n",
    "        \n",
    "#evaluation on testing data\n",
    "print('final results:')\n",
    "msg, r = handle_evaluation(i+1, last_mean_loss, valid=False, method='average')\n",
    "write_to_log(opts.log_file_path, msg)"
   ]
  }
 ],
 "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}