examples/kg_download_quick_start/kg_download_quick_start.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8c999031",
   "metadata": {},
   "source": [
    "# Using Deep Search Knowledge Graphs with PyTorch Geometric"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "959e2f71",
   "metadata": {},
   "source": [
    "Deep Search can construct Knowledge Graphs (KGs) by parsing large collections of documents. This tutorial shows how to download these knowledge graphs and import them locally in [PyTorch Geometric](https://github.com/pyg-team/pytorch_geometric), a popular graph neural network library. \n",
    "\n",
    "### Access required\n",
    "\n",
    "The content of this notebook requires access to Deep Search capabilities which are not\n",
    "available on the public access system.\n",
    "\n",
    "[Contact us](https://ds4sd.github.io) if you are interested in exploring\n",
    "these Deep Search capabilities."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8274fd05",
   "metadata": {},
   "source": [
    "### Set notebook parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01767c40",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dsnotebooks.settings import KGProjectNotebookSettings\n",
    "\n",
    "# notebook settings auto-loaded from .env / env vars\n",
    "notebook_settings = KGProjectNotebookSettings(kg_key=\"\")\n",
    "\n",
    "PROFILE_NAME = notebook_settings.profile  # the profile to use\n",
    "PROJECT_KEY = notebook_settings.proj_key\n",
    "KG_KEY = notebook_settings.kg_key\n",
    "BASE_DIR = \"./KG-data\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be45eb77",
   "metadata": {},
   "source": [
    "### Example dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b897f31",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import standard dependencies\n",
    "import os\n",
    "import json\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "import tarfile\n",
    "import requests\n",
    "\n",
    "# Dependencies related to PyTorch Geometric\n",
    "import torch\n",
    "from torch_geometric.data import HeteroData\n",
    "from torch_geometric.transforms import ToUndirected\n",
    "\n",
    "# IPython utilities\n",
    "from IPython.display import display\n",
    "\n",
    "# Import deepsearch-toolkit\n",
    "import deepsearch as ds\n",
    "from deepsearch.core.client import DeepSearchConfig\n",
    "from deepsearch.cps.client.api import CpsApi, CpsApiClient"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a282398b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create base directory if it does not exist\n",
    "if not os.path.exists(BASE_DIR):\n",
    "    os.mkdir(BASE_DIR)\n",
    "\n",
    "\n",
    "# Raise an error if the base directory is not empty\n",
    "if len(os.listdir(BASE_DIR)) > 0:\n",
    "    raise ValueError(\n",
    "        f\"BASE_DIR must be empty but found the following contents: {os.listdir(BASE_DIR)}\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71b6e5ca",
   "metadata": {},
   "source": [
    "# Part 1: Downloading the Knowledge Graph"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06dc4f79",
   "metadata": {},
   "source": [
    "### Connect to Deep Search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "004e7de9",
   "metadata": {},
   "outputs": [],
   "source": [
    "api = ds.CpsApi.from_env(profile_name=PROFILE_NAME)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42b1cc3f",
   "metadata": {},
   "source": [
    "### Download the knowledge graph\n",
    "\n",
    "We use an example knowledge graph based on 753 documents from arXiv related to the search phrase \"power conversion efficiency\" + \"organic\". The `PROJECT_KEY` and `KG_KEY` parameters that were set above correspond to this knowledge graph. In general, these parameters can be obtained from the API section of the UI as described in the [documentation](https://ds4sd.github.io/deepsearch-toolkit/guide/kgs/). \n",
    "\n",
    "We begin by using the API to get a URL for downloading the knowledge graph. **This step takes about 6-8 minutes due to the size of our KG.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5c38ba8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the download url using the API\n",
    "kg = api.knowledge_graphs.get(PROJECT_KEY, KG_KEY)\n",
    "download_url = kg.download()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1b786ea",
   "metadata": {},
   "source": [
    "The `download_url` can now be used to download a gzipped file that contains the contents of our knowledge graph. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8b24fe9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download the knowledge graph using requests\n",
    "zipped_file_path = os.path.join(BASE_DIR, \"kg_data.tar.gz\")\n",
    "\n",
    "\n",
    "with open(zipped_file_path, \"wb+\") as download_file, requests.get(\n",
    "    download_url\n",
    ") as response:\n",
    "    content_length = int(response.headers.get(\"content-length\"))\n",
    "    with tqdm(total=100, position=0) as pbar:\n",
    "        for line in response:\n",
    "            download_file.write(line)\n",
    "            pbar.update((len(line) / content_length) * 100)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "012c944a",
   "metadata": {},
   "source": [
    "### Extract the contents of the downloaded knowledge graph\n",
    "We use the `tarfile` module to unzip the contents of the downloaded `.tar.gz` file into the directory `BASE_DIR/unzipped_data`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1dada55",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save the unzipped KG\n",
    "unzipped_dir = os.path.join(BASE_DIR, \"unzipped_data\")\n",
    "os.mkdir(unzipped_dir)\n",
    "with tarfile.open(zipped_file_path, \"r\") as f:\n",
    "    f.extractall(path=unzipped_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3bc1faf6",
   "metadata": {},
   "source": [
    "### Understanding the downloaded data\n",
    "The extracted data in `BASE_DIR/unzipped_data` consists of several files in the [JSON Lines](https://jsonlines.org) format. As the downloaded KG is heterogeneous (meaning there are different types of nodes and edges), we get one `.jsonl` file for every node type. A special file named `_edges.jsonl` contains information about the edges in the KG."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f9cb932",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get a list of all the files in the unzipped data\n",
    "files = list(os.walk(os.path.join(BASE_DIR, \"unzipped_data\")))[0][2]\n",
    "display(sorted(files))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce5950ba",
   "metadata": {},
   "source": [
    "Each line in a `.jsonl` file contains information about a node encoded in the JSON format. For example, we show the first record from the `material.jsonl` file below. This file contains information about different materials found in the documents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7408bb61",
   "metadata": {},
   "outputs": [],
   "source": [
    "# A function for reading the contents of a jsonl file into a pandas dataframe\n",
    "def jsonl2df(filepath: str) -> pd.DataFrame:\n",
    "    \"\"\"\n",
    "    Reads the contents of a jsonl file into a Pandas DataFrame\n",
    "    :param filepath: Path to the jsonl file\n",
    "    :return dataframe: A pandas DataFrame corresponding to the data stored in the file\n",
    "    \"\"\"\n",
    "    with open(filepath, \"r\") as f:\n",
    "        data = pd.DataFrame([json.loads(line) for line in f])\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6293224",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show the first record in the documents file\n",
    "materials = jsonl2df(os.path.join(BASE_DIR, \"unzipped_data\", \"documents.jsonl\"))\n",
    "display(materials.iloc[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef08de5a",
   "metadata": {},
   "source": [
    "We now take a look at the `_edges.jsonl` file. As before, each line contains information about an edge in the JSON format. Given a directed edge from `x` to `y`, we call `x` the _tail_ of this edge and `y` its _head_. We show a few interesting properties of these edges listed below.\n",
    "* `source_collection`: Type of the node at the tail of the edge\n",
    "* `target_collection`: Type of the node at the head of the edge\n",
    "* `source_hash`: Hash for the node at the tail of the edge\n",
    "* `target_hash`: Hash for the node at the head of the edge\n",
    "* `symmetric`: Whether the edge is directed (if `False`) or undirected (if `True`)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "883bfa79",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show the first few edges\n",
    "edges = jsonl2df(os.path.join(BASE_DIR, \"unzipped_data\", \"_edges.jsonl\"))\n",
    "if len(edges) > 0:\n",
    "    edges = edges[\n",
    "        [\n",
    "            \"source_collection\",\n",
    "            \"target_collection\",\n",
    "            \"source_hash\",\n",
    "            \"target_hash\",\n",
    "            \"symmetric\",\n",
    "        ]\n",
    "    ]\n",
    "    display(edges.head())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44b88420",
   "metadata": {},
   "source": [
    "# Part 2: Creating a PyTorch Geometric Knowledge Graph\n",
    "\n",
    "Recall that the knowledge graph downloaded above is based on 753 documents from arXiv related to the search phrase \"power conversion efficiency\" + \"organic\". Deep Search has extracted various entities and relationships from these documents using a user-defined data flow. In this section, we will use a subset of the downloaded data to create a heterogeneous knowledge graph in PyTorch Geometric.\n",
    "\n",
    "The subset we are interested in is concerned with materials and their properties. We will add the following two types of nodes from the downloaded KG to our PyTorch Geometric KG:\n",
    "1. `material`: Materials discovered in the documents\n",
    "2. `property`: Various material properties extracted from the documents\n",
    "\n",
    "We will then extract edges relating these nodes from the `_edges.jsonl` file. Let us begin by initializing an empty knowledge graph.\n",
    "\n",
    "### Initialize the KG\n",
    "\n",
    "PyTorch Geometric uses the `HeteroData` class to represent a heterogeneous knowledge graph. Below, we create an empty `HeteroData` instance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "396ab9c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an empty heterogeneous knowledge graph\n",
    "hetero_kg = HeteroData()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aad42c72",
   "metadata": {},
   "source": [
    "### Add nodes to the KG\n",
    "Next, we add the `material` and `property` nodes to our KG as mentioned above. In this simplified example, the nodes do not have explicit features associated with them. We therefore use one-hot encoding to set `hetero_kg[nodetype].x`, as is a common practice. We also add two other attributes to each node, the `_hash` for adding edges later on and the `_name` for visualization purposes. One can also set `hetero_kg[node_type].y = ...` if node labels are available as attributes in the corresponding `.jsonl` file for the node type. We do not use any labels in this minimal example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02fc20f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "nodetypes = {\n",
    "    \"material\": os.path.join(BASE_DIR, \"unzipped_data\", \"material.jsonl\"),\n",
    "    \"property\": os.path.join(BASE_DIR, \"unzipped_data\", \"property.jsonl\"),\n",
    "}\n",
    "\n",
    "for nodetype in nodetypes:\n",
    "    if os.path.isfile(nodetypes[nodetype]):\n",
    "        data = jsonl2df(nodetypes[nodetype])\n",
    "        hetero_kg[nodetype].x = torch.eye(data.shape[0])\n",
    "        hetero_kg[nodetype][\"_hash\"] = dict(\n",
    "            (_hash, _idx) for _idx, _hash in enumerate(data[\"_hash\"].to_list())\n",
    "        )\n",
    "        hetero_kg[nodetype][\"_name\"] = data[\"_name\"].to_list()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e79607cd",
   "metadata": {},
   "source": [
    "### Add edges to the KG\n",
    "\n",
    "As we only have `material` and `property` nodes in `hetero_kg`, we search for edges between these two node types in the `_edges.jsonl` file and add them to `hetero_kg`. This adds an edge between `material A` and `property B` if `B` was mentioned in the context of `A` in at least one of the arXiv documents. We can go one step further and add the value that `property B` takes for `material A` as an edge attribute[$^{[1]}$](#footnote1), but we avoid it here to simplify the exposition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21fd8c27",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Find the relevant edges\n",
    "edges = jsonl2df(os.path.join(BASE_DIR, \"unzipped_data\", \"_edges.jsonl\"))\n",
    "if len(edges)>0:\n",
    "    edges = edges[\n",
    "        (edges.source_collection == \"material\") & (edges.target_collection == \"property\")\n",
    "    ]\n",
    "    edges = [edges[\"source_hash\"].to_list(), edges[\"target_hash\"].to_list()]\n",
    "\n",
    "    # Create the edge index\n",
    "    edge_index = []\n",
    "    for hash_mat, hash_prop in zip(*edges):\n",
    "        edge_index.append(\n",
    "            [\n",
    "                hetero_kg[\"material\"][\"_hash\"][hash_mat],\n",
    "                hetero_kg[\"property\"][\"_hash\"][hash_prop],\n",
    "            ]\n",
    "        )\n",
    "\n",
    "    # Add edge index to the KG\n",
    "    hetero_kg[\"material\", \"mat2prop\", \"property\"].edge_index = (\n",
    "        torch.tensor(edge_index).long().t()\n",
    "    )\n",
    "\n",
    "    # Make the graph undirected\n",
    "    hetero_kg = ToUndirected()(hetero_kg)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b583cb1e",
   "metadata": {},
   "source": [
    "### Summarize the created KG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e1e2e48",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Summarize the KG\n",
    "print(\"Number of nodes\")\n",
    "for node_type in hetero_kg.node_types:\n",
    "    print(f\"\\t{node_type} -> {hetero_kg[node_type].num_nodes}\")\n",
    "print(f\"Total number of nodes: {hetero_kg.num_nodes}\")\n",
    "\n",
    "print(\"\\nNumber of edges\")\n",
    "for edge_type in hetero_kg.edge_types:\n",
    "    print(f\"\\t{edge_type} -> {hetero_kg[edge_type].num_edges}\")\n",
    "print(f\"Total number of edges: {hetero_kg.num_edges}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a25e524",
   "metadata": {},
   "source": [
    "It is interesting to note that the constructed KG is very sparse as the number of edges is almost equal to the number of nodes. This is because Deep Search extracted several properties and materials but each material was only linked to a handful of properties in the document collection."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d122b744",
   "metadata": {},
   "source": [
    "### Visualizing the KG\n",
    "\n",
    "Note that our KG is a bipartite graph as it only contains edges between `material` and `property` nodes. We visualize a small subset of this KG by selecting four interesting `material` nodes and listing up to four randomly chosen `property` nodes that they are connected to."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90d1ec78",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Select materials to display\n",
    "if hetero_kg[\"material\"]:\n",
    "    materials = [\"perovskite/Si\", \"O(2) Ti(1)\", \"A(1) I(3) M(1) Pb(1)\", \"O(1) Zn(1)\"]\n",
    "    mat_idx = [hetero_kg[\"material\"][\"_name\"].index(mat) for mat in materials]\n",
    "\n",
    "    # Get properties corresponding to each material\n",
    "    properties = dict()\n",
    "    for m_idx, material in zip(mat_idx, materials):\n",
    "        current_edges = (\n",
    "            hetero_kg[\"material\", \"mat2prop\", \"property\"].edge_index[0, :] == m_idx\n",
    "        )\n",
    "        prop_idx = hetero_kg[\"material\", \"mat2prop\", \"property\"].edge_index[\n",
    "            1, current_edges\n",
    "        ]\n",
    "        properties[material] = [\n",
    "            hetero_kg[\"property\"][\"_name\"][idx] for idx in prop_idx.tolist()\n",
    "        ]\n",
    "\n",
    "    # Show up to four randomly chosen properties for each material\n",
    "    df = pd.DataFrame()\n",
    "    for mat, prop in properties.items():\n",
    "        # Restrict to four properties\n",
    "        if len(prop) > 4:\n",
    "            prop = [prop[idx] for idx in torch.randperm(len(prop)).tolist()[:4]]\n",
    "\n",
    "        # Add the row to the dataframe\n",
    "        curr_dict = dict(\n",
    "            [(\"material\", [mat])]\n",
    "            + [(f\"Property{p_idx}\", [p]) for p_idx, p in enumerate(prop)]\n",
    "        )\n",
    "        curr_df = pd.DataFrame(curr_dict)\n",
    "        df = pd.concat([df, curr_df]).reset_index(drop=True)\n",
    "\n",
    "    display(df)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d8926fb",
   "metadata": {},
   "source": [
    "Let's take another example. A `perovskite/Si` tandem is useful for making efficient solar cells. A few interesting properties for this material include its power conversion efficiency (PCE) and its band gap. Below we confirm that these properties are indeed linked to `perovskite/Si` in the KG."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc57fd5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "if hetero_kg[\"material\"]:\n",
    "    # Find properties linked to perovskite/Si\n",
    "    m_idx = hetero_kg[\"material\"][\"_name\"].index(\"perovskite/Si\")\n",
    "    perovskite_edges = (\n",
    "        hetero_kg[\"material\", \"mat2prop\", \"property\"].edge_index[0, :] == m_idx\n",
    "    )\n",
    "    prop_idx = hetero_kg[\"material\", \"mat2prop\", \"property\"].edge_index[1, perovskite_edges]\n",
    "    properties = [hetero_kg[\"property\"][\"_name\"][idx] for idx in prop_idx.tolist()]\n",
    "\n",
    "    # Check if the desired properties are linked\n",
    "    print(\n",
    "        f'Is perovskite/Si linked to power conversion efficiency: {\"power conversion efficiency\" in properties}'\n",
    "    )\n",
    "    print(f'Is perovskite/Si linked to band gap: {\"band gap\" in properties}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d4923ea",
   "metadata": {},
   "source": [
    "This concludes the process of creating a simple PyTorch Geometric KG from the KG downloaded from Deep Search. One can now train powerful graph neural networks for various downstream tasks like node classification and link prediction using `hetero_kg`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e92255c",
   "metadata": {},
   "source": [
    "# Footnotes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67654dac",
   "metadata": {},
   "source": [
    "<a name='footnote1'></a> \n",
    "### \\[1\\] Edges with attributes\n",
    "\n",
    "Edges with attributes can be stored as additional nodes in Deep Search. For example, consider four nodes `{A, B, X, Y}` where `{A, B}` have the same node type (say `author`) and `{X, Y}` have another node type (say `paper`). The relationship between nodes of type `author` and `paper` can be encoded as nodes of type `author-to-paper`. If author `A` has written paper `X`, then there will a node `A_X` of type `author-to-paper`, and both `A` and `X` will be connected to `A_X` in the KG. \n",
    "\n",
    "Now suppose we want to add more information about the relationship between `author A` and `paper X`. For example, we may be interested in knowing the index at which `A` appears in the author list of `X`. Such information can be added as attributes of node `A_X` in Deep Search. One can then easily parse these nodes to add edge attributes in PyTorch Geometric."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "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.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}