knowledge_graph_embeddings.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch.autograd import Variable\n",
    "#graph = [\n",
    "#    ('seth rogan', 'acted_in', 'the interview'),\n",
    "#    ('seth rogan', 'acted_in', 'this is the end'),\n",
    "#    ('seth rogan', 'acted_in', 'pineapple express'),\n",
    "#    ('james franco', 'acted_in', 'the interview'),\n",
    "#    ('james franco', 'acted_in', 'this is the end'),\n",
    "#    ('will smith', 'acted_in', 'irobot'),\n",
    "#    ('will smith', 'acted_in', 'independence day'),\n",
    "#    ('tom cruise', 'acted_in', 'mission impossible'),\n",
    "#]\n",
    "graph = [\n",
    "    ('tom cruise', 'acted_in', 'oblivion'),\n",
    "    ('tom cruise', 'acted_in', 'tropic thunder'),\n",
    "    ('tom cruise', 'acted_in', 'mission impossible'),\n",
    "    ('jack black', 'acted_in', 'tropic thunder'),\n",
    "    ('ben stiller', 'acted_in', 'dodgeball'),\n",
    "    ('dodgeball', 'is_genre', 'comedy'),\n",
    "    ('tropic thunder', 'is_genre', 'comedy'),\n",
    "    ('mission impossible', 'is_genre', 'action'),\n",
    "    ('oblivion', 'is_genre', 'action'),\n",
    "]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class KnowledgeGraphEmbeddings(nn.Module):\n",
    "    def __init__(self, graph, embedding_size = 3):\n",
    "        super(KnowledgeGraphEmbeddings, self).__init__()\n",
    "        \n",
    "        entities = []\n",
    "        relations = []\n",
    "        for fact in graph:\n",
    "            source, relation, target = fact[0], fact[1], fact[2]\n",
    "            entities.append(source)\n",
    "            entities.append(target)\n",
    "            relations.append(relation)\n",
    "\n",
    "        self.entities = list(set(entities))\n",
    "        self.relations = list(set(relations))\n",
    "        self.entities2id = { ent:i for i, ent in enumerate(self.entities) } \n",
    "        self.relations2id = { rel:i for i, rel in enumerate(self.relations) } \n",
    "        \n",
    "        self.entity_embeddings = nn.Embedding(len(self.entities2id), embedding_size)\n",
    "        self.relation_embeddings = nn.Embedding(len(self.relations2id), embedding_size)\n",
    "        self.graph = graph\n",
    "        \n",
    "    def positive(self, fact):\n",
    "        source, relation, target = fact[0], fact[1], fact[2]\n",
    "        source_id = self.entities2id[source]\n",
    "        target_id = self.entities2id[target]\n",
    "        relation_id = self.relations2id[relation]\n",
    "        \n",
    "        source_id = Variable(torch.LongTensor([source_id])).view(1, -1)\n",
    "        target_id = Variable(torch.LongTensor([target_id])).view(1, -1)\n",
    "        relation_id = Variable(torch.LongTensor([relation_id])).view(1, -1)\n",
    "\n",
    "        source_embedding = self.entity_embeddings(source_id)\n",
    "        target_embedding = self.entity_embeddings(target_id)\n",
    "        relation_embedding = self.entity_embeddings(relation_id)\n",
    "        score = torch.dist(source_embedding + relation_embedding, target_embedding)\n",
    "        #print('positive', source, relation, target, score.data[0])\n",
    "        return score\n",
    "    \n",
    "    def negative(self):\n",
    "        # Sample until we find an invalid fact\n",
    "        while True:\n",
    "            random_source_id = random.randint(0, len(self.entities2id) - 1)\n",
    "            random_relation_id = random.randint(0, len(self.relations2id) - 1)\n",
    "            random_target_id = random.randint(0, len(self.entities2id) - 1)\n",
    "            \n",
    "            source = self.entities[random_source_id]\n",
    "            target = self.entities[random_target_id]\n",
    "            relation = self.relations[random_relation_id]\n",
    "            if (source, relation, target) not in graph:\n",
    "                break\n",
    "        \n",
    "        source_id = Variable(torch.LongTensor([random_source_id])).view(1, -1)\n",
    "        relation_id = Variable(torch.LongTensor([random_relation_id])).view(1, -1)\n",
    "        target_id = Variable(torch.LongTensor([random_target_id])).view(1, -1)\n",
    "        \n",
    "        source_embedding = self.entity_embeddings(source_id)\n",
    "        relation_embedding = self.entity_embeddings(relation_id)\n",
    "        target_embedding = self.entity_embeddings(target_id)\n",
    "        \n",
    "        score = torch.dist(source_embedding + relation_embedding, target_embedding)\n",
    "        #print('negative', self.entities[random_source_id], self.relations[random_relation_id], self.entities[random_target_id], score.data[0])\n",
    "        return score\n",
    "    \n",
    "    def forward(self, fact, margin=1.0):\n",
    "        positive_score = self.positive(fact)\n",
    "        negative_score = self.negative()\n",
    "        loss = torch.max((positive_score - negative_score + margin).sum(), Variable(torch.FloatTensor([0])))\n",
    "        #print('%.2f - %.2f + %.2f = %.2f' %(positive_score.data[0], negative_score.data[0], margin, loss.data[0]))\n",
    "\n",
    "        return loss\n",
    "    \n",
    "    def score(self, fact):\n",
    "        return self.positive(fact)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(fact, embedding_model, optimizer):\n",
    "    avg_loss = 0.\n",
    "    # 10 training cycles for this fact, with negative sampling\n",
    "    for i in range(len(embedding_model.entities) * len(embedding_model.entities) * len(embedding_model.relations)):\n",
    "        embedding_model.zero_grad()\n",
    "        loss = embedding_model(fact)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        avg_loss += loss.data[0]\n",
    "    return avg_loss / 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-4-77ac44327573>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      6\u001b[0m     \u001b[0mtotal_loss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m0.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      7\u001b[0m     \u001b[0;32mfor\u001b[0m \u001b[0mfact\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mgraph\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 8\u001b[0;31m         \u001b[0mloss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtrain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfact\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0membedding_model\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moptimizer\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      9\u001b[0m         \u001b[0mtotal_loss\u001b[0m \u001b[0;34m+=\u001b[0m \u001b[0mloss\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     10\u001b[0m     \u001b[0mlosses\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtotal_loss\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-3-90d3296c87f5>\u001b[0m in \u001b[0;36mtrain\u001b[0;34m(fact, embedding_model, optimizer)\u001b[0m\n\u001b[1;32m      5\u001b[0m         \u001b[0membedding_model\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mzero_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      6\u001b[0m         \u001b[0mloss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0membedding_model\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfact\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 7\u001b[0;31m         \u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      8\u001b[0m         \u001b[0moptimizer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      9\u001b[0m         \u001b[0mavg_loss\u001b[0m \u001b[0;34m+=\u001b[0m \u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdata\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/env/lib/python3.6/site-packages/torch/autograd/variable.py\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(self, gradient, retain_graph, create_graph, retain_variables)\u001b[0m\n\u001b[1;32m    165\u001b[0m                 \u001b[0mVariable\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    166\u001b[0m         \"\"\"\n\u001b[0;32m--> 167\u001b[0;31m         \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mautograd\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgradient\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_variables\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    168\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    169\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mregister_hook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/env/lib/python3.6/site-packages/torch/autograd/__init__.py\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(variables, grad_variables, retain_graph, create_graph, retain_variables)\u001b[0m\n\u001b[1;32m     97\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     98\u001b[0m     Variable._execution_engine.run_backward(\n\u001b[0;32m---> 99\u001b[0;31m         variables, grad_variables, retain_graph)\n\u001b[0m\u001b[1;32m    100\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    101\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "embedding_model = KnowledgeGraphEmbeddings(graph)\n",
    "optimizer = torch.optim.Adam(embedding_model.parameters(), lr=0.001)\n",
    "losses = []\n",
    "epochs = 50\n",
    "for e in range(epochs):\n",
    "    total_loss = 0.\n",
    "    for fact in graph:\n",
    "        loss = train(fact, embedding_model, optimizer)\n",
    "        total_loss += loss\n",
    "    losses.append(total_loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.figure(figsize=(8, 6),)\n",
    "plt.plot(losses)\n",
    "plt.xlabel('epoch')\n",
    "plt.ylabel('loss');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#print(embedding_model.score(('james franco', 'acted_in', 'the interview')))\n",
    "#print(embedding_model.score(('james franco', 'acted_in', 'pineapple express')))\n",
    "#print(embedding_model.score(('james franco', 'acted_in', 'irobot')))\n",
    "#print(embedding_model.score(('james franco', 'acted_in', 'mission impossible')))\n",
    "\n",
    "print(embedding_model.score(('ben stiller', 'acted_in', 'dodgeball')))\n",
    "print(embedding_model.score(('ben stiller', 'acted_in', 'tropic thunder')))\n",
    "print(embedding_model.score(('ben stiller', 'acted_in', 'mission impossible')))\n",
    "print(embedding_model.score(('ben stiller', 'acted_in', 'oblivion')))"
   ]
  }
 ],
 "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}