examples/kge-predict-transe-pyg-train.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Knowledge graph embeddings: Training in PyG, prediction with GDS"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/neo4j/graph-data-science-client/blob/main/examples/kge-predict-transe-pyg-train.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "This Jupyter notebook is hosted [here](https://github.com/neo4j/graph-data-science-client/blob/main/examples/kge-predict-transe-pyg-train.ipynb) in the Neo4j Graph Data Science Client Github repository.\n",
    "\n",
    "The notebook demonstrates how to use the `graphdatascience` and PyTorch Geometric (PyG) Python libraries to:\n",
    "\n",
    "1. Import the [FB15k-237](https://www.microsoft.com/en-us/download/details.aspx?id=52312) dataset directly into GDS\n",
    "2. Train a TransE model with PyG\n",
    "3. Make predictions on the data in the database using GDS Knowledge Graph Embeddings functionality\n",
    "\n",
    "## Prerequisites\n",
    "\n",
    "To run this notebook, you'll need a Neo4j server with a recent GDS version (2.5+ or later) installed.\n",
    "\n",
    "Additionally, the following Python libraries are required:\n",
    "\n",
    "- `graphdatascience`, [see documentation for installation instructions](https://neo4j.com/docs/graph-data-science-client/current/installation/)\n",
    "- `pytorch-geometric` version >= 2.5.0, [see PyG documentation for installation instructions](https://pytorch-geometric.readthedocs.io/en/latest/notes/installation.html)\n",
    "\n",
    "## Setup\n",
    "\n",
    "We'll begin by importing our dependencies and establishing a GDS client connection to the database."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install graphdatascience torch torch_geometric"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import collections\n",
    "import os\n",
    "\n",
    "import pandas as pd\n",
    "import torch\n",
    "import torch.optim as optim\n",
    "from torch_geometric.data import Data, download_url\n",
    "from torch_geometric.nn import TransE\n",
    "from tqdm import tqdm\n",
    "\n",
    "from graphdatascience import GraphDataScience"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "NEO4J_URI = os.environ.get(\"NEO4J_URI\", \"bolt://localhost:7687\")\n",
    "NEO4J_AUTH = None\n",
    "NEO4J_DB = os.environ.get(\"NEO4J_DB\", \"neo4j\")\n",
    "if os.environ.get(\"NEO4J_USER\") and os.environ.get(\"NEO4J_PASSWORD\"):\n",
    "    NEO4J_AUTH = (\n",
    "        os.environ.get(\"NEO4J_USER\"),\n",
    "        os.environ.get(\"NEO4J_PASSWORD\"),\n",
    "    )\n",
    "gds = GraphDataScience(NEO4J_URI, auth=NEO4J_AUTH, database=NEO4J_DB)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "verify-version"
    ]
   },
   "outputs": [],
   "source": [
    "from graphdatascience import ServerVersion\n",
    "\n",
    "# This notebook requires GDS 2.5.0 or later\n",
    "assert gds.server_version() >= ServerVersion(2, 5, 0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Downloading and Storing the FB15k-237 Dataset in the Database\n",
    "Download the FB15k-237 dataset\n",
    "Extract the required files: train.txt, valid.txt, and test.txt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import zipfile\n",
    "\n",
    "url = \"https://download.microsoft.com/download/8/7/0/8700516A-AB3D-4850-B4BB-805C515AECE1/FB15K-237.2.zip\"\n",
    "raw_dir = \"./data_from_zip\"\n",
    "download_url(f\"{url}\", raw_dir)\n",
    "\n",
    "raw_file_names = [\"train.txt\", \"valid.txt\", \"test.txt\"]\n",
    "with zipfile.ZipFile(raw_dir + \"/\" + os.path.basename(url), \"r\") as zip_ref:\n",
    "    for filename in raw_file_names:\n",
    "        zip_ref.extract(f\"Release/{filename}\", path=raw_dir)\n",
    "data_dir = raw_dir + \"/\" + \"Release\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Set a constraint for unique id entries to speed up data uploads."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gds.run_cypher(\"CREATE CONSTRAINT entity_id FOR (e:Entity) REQUIRE e.id IS UNIQUE\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "**Creating Entity Nodes**:\n",
    "   Create a node with the label `Entity`. This node should have properties `id` and `text`. \n",
    "   - Syntax: `(:Entity {id: int, text: str})`\n",
    "\n",
    "**Creating Relationships for Training with PyG**:\n",
    "   Based on the training stage, create relationships of type `TRAIN`, `TEST`, or `VALID`. Each of these relationships should have a `rel_id` property.\n",
    "   - Example Syntax: `[:TRAIN {rel_id: int}]`\n",
    "\n",
    "**Creating Relationships for Prediction with GDS**:\n",
    "   For the prediction stage, create relationships of a specific type denoted as `REL_i`. Each of these relationships should have `rel_id` and `text` properties.\n",
    "   - Example Syntax: `[:REL_7 {rel_id: int, text: str}]`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rel_types = {\n",
    "    \"train.txt\": \"TRAIN\",\n",
    "    \"valid.txt\": \"VALID\",\n",
    "    \"test.txt\": \"TEST\",\n",
    "}\n",
    "rel_id_to_text_dict = {}\n",
    "rel_type_dict = collections.defaultdict(list)\n",
    "rel_dict = {}\n",
    "\n",
    "\n",
    "def process():\n",
    "    node_dict_ = {}\n",
    "    for file_name in raw_file_names:\n",
    "        file_name_path = data_dir + \"/\" + file_name\n",
    "\n",
    "        with open(file_name_path, \"r\") as f:\n",
    "            data = [x.split(\"\\t\") for x in f.read().split(\"\\n\")[:-1]]\n",
    "\n",
    "        list_of_dicts = []\n",
    "        for i, (src, rel, dst) in enumerate(data):\n",
    "            if src not in node_dict_:\n",
    "                node_dict_[src] = len(node_dict_)\n",
    "            if dst not in node_dict_:\n",
    "                node_dict_[dst] = len(node_dict_)\n",
    "            if rel not in rel_dict:\n",
    "                rel_dict[rel] = len(rel_dict)\n",
    "                rel_id_to_text_dict[rel_dict[rel]] = rel\n",
    "\n",
    "            source = node_dict_[src]\n",
    "            target = node_dict_[dst]\n",
    "            edge_type = rel_dict[rel]\n",
    "\n",
    "            rel_type_dict[edge_type].append(\n",
    "                {\n",
    "                    \"source\": source,\n",
    "                    \"target\": target,\n",
    "                }\n",
    "            )\n",
    "            list_of_dicts.append(\n",
    "                {\n",
    "                    \"source\": source,\n",
    "                    \"source_text\": src,\n",
    "                    \"target\": target,\n",
    "                    \"target_text\": dst,\n",
    "                    \"rel_id\": edge_type,\n",
    "                }\n",
    "            )\n",
    "\n",
    "        rel_type = rel_types[file_name]\n",
    "        print(f\"Writing {len(list_of_dicts)} entities of {rel_type}\")\n",
    "        gds.run_cypher(\n",
    "            f\"\"\"\n",
    "            UNWIND $ll as l\n",
    "            MERGE (n:Entity {{id:l.source, text:l.source_text}})\n",
    "            MERGE (m:Entity {{id:l.target, text:l.target_text}})\n",
    "            MERGE (n)-[:{rel_type} {{rel_id:l.rel_id}}]->(m)\n",
    "            \"\"\",\n",
    "            params={\"ll\": list_of_dicts},\n",
    "        )\n",
    "\n",
    "    print(\"Writing relationships as different relationship types\")\n",
    "    for rel_id, rels in tqdm(rel_type_dict.items()):\n",
    "        REL_TYPE = f\"REL_{rel_id}\"\n",
    "        gds.run_cypher(\n",
    "            f\"\"\"\n",
    "            UNWIND $ll AS l MATCH (n:Entity {{id:l.source}}), (m:Entity {{id:l.target}})\n",
    "            MERGE (n)-[:{REL_TYPE} {{rel_id:$rel_id, text:$text}}]->(m)\n",
    "            \"\"\",\n",
    "            params={\"ll\": rels, \"rel_id\": rel_id, \"text\": rel_id_to_text_dict[rel_id]},\n",
    "        )\n",
    "\n",
    "\n",
    "process()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Project all data in graph to get mapping between `id` and internal `nodeId` field from database."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "node_projection = {\"Entity\": {\"properties\": \"id\"}}\n",
    "relationship_projection = [\n",
    "    {\"TRAIN\": {\"orientation\": \"NATURAL\", \"properties\": \"rel_id\"}},\n",
    "    {\"TEST\": {\"orientation\": \"NATURAL\", \"properties\": \"rel_id\"}},\n",
    "    {\"VALID\": {\"orientation\": \"NATURAL\", \"properties\": \"rel_id\"}},\n",
    "]\n",
    "\n",
    "ttv_G, result = gds.graph.project(\n",
    "    \"fb15k-graph-ttv\",\n",
    "    node_projection,\n",
    "    relationship_projection,\n",
    ")\n",
    "\n",
    "node_properties = gds.graph.nodeProperties.stream(\n",
    "    ttv_G,\n",
    "    [\"id\"],\n",
    "    separate_property_columns=True,\n",
    ")\n",
    "\n",
    "nodeId_to_id = dict(zip(node_properties.nodeId, node_properties.id))\n",
    "id_to_nodeId = dict(zip(node_properties.id, node_properties.nodeId))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Training the TransE Model with PyG\n",
    "\n",
    "Retrieve data from the database, convert it into torch tensors, and format it into a `Data` structure suitable for training with PyG."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_data_from_graph(relationship_type):\n",
    "    rels_tmp = gds.graph.relationshipProperty.stream(ttv_G, \"rel_id\", relationship_type)\n",
    "    topology = [\n",
    "        rels_tmp.sourceNodeId.map(lambda x: nodeId_to_id[x]),\n",
    "        rels_tmp.targetNodeId.map(lambda x: nodeId_to_id[x]),\n",
    "    ]\n",
    "    edge_index = torch.tensor(topology, dtype=torch.long)\n",
    "    edge_type = torch.tensor(rels_tmp.propertyValue.astype(int), dtype=torch.long)\n",
    "    data = Data(edge_index=edge_index, edge_type=edge_type)\n",
    "    data.num_nodes = len(nodeId_to_id)\n",
    "    display(data)\n",
    "    return data\n",
    "\n",
    "\n",
    "train_tensor_data = create_data_from_graph(\"TRAIN\")\n",
    "test_tensor_data = create_data_from_graph(\"TEST\")\n",
    "val_tensor_data = create_data_from_graph(\"VALID\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Drop the projected graph to save memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gds.graph.drop(ttv_G)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "The training process of the TransE model follows the corresponding PyG [example](https://github.com/pyg-team/pytorch_geometric/blob/master/examples/kge_fb15k_237.py)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model_with_pyg():\n",
    "    device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "\n",
    "    model = TransE(\n",
    "        num_nodes=train_tensor_data.num_nodes,\n",
    "        num_relations=train_tensor_data.num_edge_types,\n",
    "        hidden_channels=50,\n",
    "    ).to(device)\n",
    "\n",
    "    loader = model.loader(\n",
    "        head_index=train_tensor_data.edge_index[0],\n",
    "        rel_type=train_tensor_data.edge_type,\n",
    "        tail_index=train_tensor_data.edge_index[1],\n",
    "        batch_size=1000,\n",
    "        shuffle=True,\n",
    "    )\n",
    "\n",
    "    optimizer = optim.Adam(model.parameters(), lr=0.01)\n",
    "\n",
    "    def train():\n",
    "        model.train()\n",
    "        total_loss = total_examples = 0\n",
    "        for head_index, rel_type, tail_index in loader:\n",
    "            optimizer.zero_grad()\n",
    "            loss = model.loss(head_index, rel_type, tail_index)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            total_loss += float(loss) * head_index.numel()\n",
    "            total_examples += head_index.numel()\n",
    "        return total_loss / total_examples\n",
    "\n",
    "    @torch.no_grad()\n",
    "    def test(data):\n",
    "        model.eval()\n",
    "        return model.test(\n",
    "            head_index=data.edge_index[0],\n",
    "            rel_type=data.edge_type,\n",
    "            tail_index=data.edge_index[1],\n",
    "            batch_size=1000,\n",
    "            k=10,\n",
    "        )\n",
    "\n",
    "    # Consider increasing the number of epochs\n",
    "    epoch_count = 5\n",
    "    for epoch in range(1, epoch_count):\n",
    "        loss = train()\n",
    "        print(f\"Epoch: {epoch:03d}, Loss: {loss:.4f}\")\n",
    "        if epoch % 75 == 0:\n",
    "            rank, hits = test(val_tensor_data)\n",
    "            print(f\"Epoch: {epoch:03d}, Val Mean Rank: {rank:.2f}, \" f\"Val Hits@10: {hits:.4f}\")\n",
    "\n",
    "    torch.save(model, f\"./model_{epoch_count}.pt\")\n",
    "\n",
    "    mean_rank, mrr, hits_at_k = test(test_tensor_data)\n",
    "    print(f\"Test Mean Rank: {mean_rank:.2f}, Test Hits@10: {hits_at_k:.4f}, MRR: {mrr:.4f}\")\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = train_model_with_pyg()\n",
    "# The model can be loaded if it was trained before\n",
    "# model = torch.load(\"./model_501.pt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Extract node embeddings from the trained model and put them into database."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in tqdm(range(len(nodeId_to_id))):\n",
    "    gds.run_cypher(\n",
    "        \"MATCH (n:Entity {id: $i}) SET n.emb=$EMBEDDING\",\n",
    "        params={\"i\": i, \"EMBEDDING\": model.node_emb.weight[i].tolist()},\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Predict Using GDS Knowledge Graph Edge Embeddings Functionality\n",
    "\n",
    "Select a relationship type for which to make predictions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "relationship_to_predict = \"/film/film/genre\"\n",
    "rel_id_to_predict = rel_dict[relationship_to_predict]\n",
    "rel_label_to_predict = f\"REL_{rel_id_to_predict}\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Project the graph with all nodes and existing relationships of the selected type."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "G_test, result = gds.graph.project(\n",
    "    \"graph_to_predict_\",\n",
    "    {\"Entity\": {\"properties\": [\"id\", \"emb\"]}},\n",
    "    rel_label_to_predict,\n",
    ")\n",
    "\n",
    "\n",
    "def print_graph_info(G):\n",
    "    print(f\"Graph '{G.name()}' node count: {G.node_count()}\")\n",
    "    print(f\"Graph '{G.name()}' node labels: {G.node_labels()}\")\n",
    "    print(f\"Graph '{G.name()}' relationship types: {G.relationship_types()}\")\n",
    "    print(f\"Graph '{G.name()}' relationship count: {G.relationship_count()}\")\n",
    "\n",
    "\n",
    "print_graph_info(G_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Retrieve the embedding for the selected relationship from the PyG model. Then, create a GDS TransE model using the graph, node embeddings property, and the embedding for the relationship to be predicted."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "target_emb = model.node_emb.weight[rel_id_to_predict].tolist()\n",
    "transe_model = gds.model.transe.create(G_test, \"emb\", {rel_label_to_predict: target_emb})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "source_node_list = [\"/m/07l450\", \"/m/0ds2l81\", \"/m/0jvt9\"]\n",
    "source_ids_df = gds.run_cypher(\n",
    "    \"UNWIND $node_text_list AS t MATCH (n:Entity) WHERE n.text=t RETURN id(n) as nodeId\",\n",
    "    params={\"node_text_list\": source_node_list},\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Now, we can use the model to make prediction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = transe_model.predict_stream(\n",
    "    source_node_filter=source_ids_df.nodeId,\n",
    "    target_node_filter=\"Entity\",\n",
    "    relationship_type=rel_label_to_predict,\n",
    "    top_k=3,\n",
    "    concurrency=4,\n",
    ")\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Augment the predicted result with node identifiers and their text values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ids_in_result = pd.unique(pd.concat([result.sourceNodeId, result.targetNodeId]))\n",
    "\n",
    "ids_to_text = gds.run_cypher(\n",
    "    \"UNWIND $ids AS id MATCH (n:Entity) WHERE id(n)=id RETURN id(n) AS nodeId, n.text AS tag, n.id AS id\",\n",
    "    params={\"ids\": ids_in_result},\n",
    ")\n",
    "\n",
    "nodeId_to_text_res = dict(zip(ids_to_text.nodeId, ids_to_text.tag))\n",
    "nodeId_to_id_res = dict(zip(ids_to_text.nodeId, ids_to_text.id))\n",
    "\n",
    "result.insert(1, \"sourceTag\", result.sourceNodeId.map(lambda x: nodeId_to_text_res[x]))\n",
    "result.insert(2, \"sourceId\", result.sourceNodeId.map(lambda x: nodeId_to_id_res[x]))\n",
    "result.insert(4, \"targetTag\", result.targetNodeId.map(lambda x: nodeId_to_text_res[x]))\n",
    "result.insert(5, \"targetId\", result.targetNodeId.map(lambda x: nodeId_to_id_res[x]))\n",
    "\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Using Write Mode\n",
    "\n",
    "Write mode allows you to write results directly to the database as a new relationship type. This approach helps to avoid mapping from `nodeId` to `id`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "write_relationship_type = \"PREDICTED_\" + rel_label_to_predict\n",
    "result_write = transe_model.predict_write(\n",
    "    source_node_filter=source_ids_df.nodeId,\n",
    "    target_node_filter=\"Entity\",\n",
    "    relationship_type=rel_label_to_predict,\n",
    "    write_relationship_type=write_relationship_type,\n",
    "    write_property=\"transe_score\",\n",
    "    top_k=3,\n",
    "    concurrency=4,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Extract the result from the database."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gds.run_cypher(\n",
    "    \"MATCH (n)-[r:\"\n",
    "    + write_relationship_type\n",
    "    + \"]->(m) RETURN n.id AS sourceId, n.text AS sourceTag, m.id AS targetId, m.text AS targetTag, r.transe_score AS score\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gds.graph.drop(G_test)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}