docs/tutorials/03-kglab_pytorch_geometric.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 🧪 Node classification with `kglab` and PyTorch Geometric\n",
    "\n",
    "We introduce the application of neural networks on knowledge graphs using `kglab` and `pytorch_geometric`. \n",
    "\n",
    "Graph Neural networks (GNNs) have gained popularity in a number of practical applications, including knowledge graphs, social networks and recommender systems. In the context of knowledge graphs, GNNs are being used for tasks such as link prediction, node classification or knowledge graph embeddings. Many use cases for these tasks are related to `Automatic Knowledge Base Construction` (AKBC) and completion.\n",
    "\n",
    "In this tutorial, we will learn to:\n",
    "\n",
    "- use `kglab` to represent a knowledge graph as a Pytorch Tensor, a suitable structure for working with neural nets\n",
    "\n",
    "- use the widely known `pytorch_geometric` (PyG) GNN library together with `kglab`.\n",
    "\n",
    "- train a GNN with `pytorch_geometric` and `PyTorch Lightning` for semi-supervised node classification of the recipes knowledge graph.\n",
    "\n",
    "- build and iterate on training data using `rubrix` with a Human-in-the-loop (HITL) approach.\n",
    "\n",
    "## Our use case in a nutshell\n",
    "\n",
    "Our goal in this notebook will be to build a semi-supervised node classifier of recipes and ingredients from scratch using kglab, PyG and Rubrix. \n",
    "\n",
    "Our classifier will be able to classify the nodes in our 15K nodes knowledge graph according to a set of pre-defined flavour related categories: `sweet`, `salty`, `piquant`, `sour`, etc. To account for mixed flavours (e.g., sweet chili sauce), our model will be multi-class (we have several target labels), multi-label (a node can be labelled as with 0 or several categories)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install `kglab` and `Pytorch Geometric`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "%pip install torch-scatter -f https://pytorch-geometric.com/whl/torch-1.8.0+cpu.html -qqq\n",
    "%pip install torch-sparse -f https://pytorch-geometric.com/whl/torch-1.8.0+cpu.html -qqq\n",
    "%pip install torch-cluster -f https://pytorch-geometric.com/whl/torch-1.8.0+cpu.html -qqq\n",
    "%pip install torch-spline-conv -f https://pytorch-geometric.com/whl/torch-1.8.0+cpu.html -qqq\n",
    "%pip install torch-geometric -qqq\n",
    "%pip install torch==1.8.0 -qqq\n",
    "\n",
    "%pip install kglab -qqq\n",
    "\n",
    "%pip install pytorch_lightning -qqq"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Loading and exploring the recipes knowledge graph\n",
    "\n",
    "We'll be working with the \"recipes\" knowledge graph, which is used throughout the `kglab` tutorial (see the [Syllabus](https://derwen.ai/docs/kgl/tutorial/)).\n",
    "\n",
    "This version of the recipes kg contains around ~15K recipes linked to their respective ingredients, as well as some other properties such as cooking time, labels and descriptions. \n",
    "\n",
    "Let's load the knowledge graph into a `kg` object by reading from an RDF file (in Turtle):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import kglab\n",
    "\n",
    "NAMESPACES = {\n",
    "    \"wtm\":  \"http://purl.org/heals/food/\",\n",
    "    \"ind\":  \"http://purl.org/heals/ingredient/\",\n",
    "    \"recipe\":  \"https://www.food.com/recipe/\",\n",
    "    }\n",
    "\n",
    "kg = kglab.KnowledgeGraph(namespaces = NAMESPACES)\n",
    "\n",
    "_ = kg.load_rdf(\"data/recipe_lg.ttl\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's take a look at our graph structure using the `Measure` class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "measure = kglab.Measure()\n",
    "measure.measure_graph(kg)\n",
    "\n",
    "f\"Nodes: {measure.get_node_count()} ; Edges: {measure.get_edge_count()}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "measure.p_gen.get_tally() # tallies the counts of predicates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "measure.s_gen.get_tally() # tallies the counts of predicates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "measure.o_gen.get_tally() # tallies the counts of predicates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "measure.l_gen.get_tally() # tallies the counts of literals"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "From the above exploration, we can extract some conclusions to guide the next steps:\n",
    "\n",
    "- We have a limited number of relationships, being `hasIngredient` the most frequent.\n",
    "\n",
    "- We have rather unique literals for labels and descriptions, but a certain amount of repetition for `hasCookTime`.\n",
    "\n",
    "- As we would have expected, most frequently referenced objects are ingredients such as `Salt`, `ChikenEgg` and so on. \n",
    "\n",
    "\n",
    "Now, let's move into preparing our knowledge graph for PyTorch."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Representing our knowledge graph as a `PyTorch` Tensor\n",
    "\n",
    "Let's now represent our `kg` as a `PyTorch` tensor using the `kglab.SubgraphTensor` class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sg = kglab.SubgraphTensor(kg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def to_edge_list(g, sg, excludes):\n",
    "    def exclude(rel):\n",
    "        return sg.n3fy(rel) in excludes\n",
    "    \n",
    "    relations = sorted(set(g.predicates()))\n",
    "    subjects = set(g.subjects())\n",
    "    objects = set(g.objects())\n",
    "    nodes = list(subjects.union(objects))\n",
    "    \n",
    "    relations_dict = {rel: i for i, rel in enumerate(list(relations)) if not exclude(rel)}\n",
    "    \n",
    "    # this offset enables consecutive indices in our final vector\n",
    "    offset = len(relations_dict.keys())\n",
    "    \n",
    "    nodes_dict = {node: i+offset for i, node in enumerate(nodes)}\n",
    "\n",
    "    \n",
    "    edge_list = []\n",
    "    \n",
    "    for s, p, o in g.triples((None, None, None)):\n",
    "        if p in relations_dict.keys(): # this means is not excluded\n",
    "            src, dst, rel = nodes_dict[s], nodes_dict[o], relations_dict[p]\n",
    "            edge_list.append([src, dst, 2 * rel])\n",
    "            edge_list.append([dst, src, 2 * rel + 1])\n",
    "    \n",
    "    # turn into str keys and concat\n",
    "    node_vector = [sg.n3fy(node) for node in relations_dict.keys()] + [sg.n3fy(node) for node in nodes_dict.keys()]\n",
    "    return edge_list, node_vector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "edge_list, node_vector = to_edge_list(kg.rdf_graph(), sg, excludes=['skos:description', 'skos:prefLabel'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(edge_list) , edge_list[0:5]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's create `kglab.Subgraph` to be used for encoding/decoding numerical ids and uris, which will be useful for preparing our training data, as well as making sense of the predictions of our neural net."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sg = kglab.Subgraph(kg=kg, preload=node_vector)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch_geometric.data import Data\n",
    "\n",
    "tensor = torch.tensor(edge_list, dtype=torch.long).t().contiguous() \n",
    "edge_index, edge_type = tensor[:2], tensor[2]\n",
    "data = Data(edge_index=edge_index)\n",
    "data.edge_type = edge_type"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(data.edge_index.shape, data.edge_type.shape, data.edge_type.max())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Building a training set with Rubrix\n",
    "\n",
    "Now that we have a tensor representation of our kg which we can feed into our neural network, let's now focus on the training data.\n",
    "\n",
    "As we will be doing semi-supervised classification, we need to build a training set (i.e., some recipes and ingredients with ground-truth labels). \n",
    "\n",
    "\n",
    "For this, we can use [Rubrix](https://github.com/recognai/rubrix), an open-source tool for exploring, labeling and iterating on data for AI. Rubrix allows data scientists and subject matter experts to rapidly iterate on training and evaluation data by enabling iterative, asynchronous and potentially distributed workflows.\n",
    "\n",
    "In Rubrix, a very simple workflow during model development looks like this:\n",
    "\n",
    "1. Log unlabelled data records with `rb.log()` into a Rubrix dataset. At this step you could use weak supervision methods (e.g., Snorkel) to pre-populate and then only refine the suggested labels, or use a pretrained model to guide your annotation process. In our case, we will just log recipe and ingredient \"records\" along with some metadata (RDF types, labels, etc.).\n",
    "\n",
    "2. Rapidly explore and label records in your dataset using the webapp which follows a search-driven approach, which is especially useful with large, potentially noisy datasets and for quickly leveraging domain knowledge (e.g., recipes containing WhiteSugar are likely sweet). For the tutorial, we have spent around 30min for labelling around 600 records.\n",
    "\n",
    "3. Retrieve your annotations any time using `rb.load()` or `rb.snapshot()`, which return a convenient `pd.Dataframe` making it quite handy to process and use for model development. In our case, we will load a snapshot, do a train_test_split with scikit_learn, and then use this for training our GNN.\n",
    "\n",
    "4. After training a model, you can go back to step 1, this time using your model and its predictions, to spot improvements, quickly label other portions of the data, and so on. In our case, as we've started with a very limited training set (~600 examples), we will use our node classifier and `rb.log()` it's predictions over the rest of our data (unlabelled recipes and ingredients)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "LABELS = ['Bitter', 'Meaty', 'Piquant', 'Salty', 'Sour', 'Sweet']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setup Rubrix\n",
    "\n",
    "If you have not installed and launched Rubrix, check the [installation guide](https://github.com/recognai/rubrix#get-started). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import rubrix as rb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Preparing our raw dataset of recipes and ingredients"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "sparql = \"\"\"\n",
    "    SELECT distinct *\n",
    "    WHERE {\n",
    "        ?uri a wtm:Recipe .\n",
    "        ?uri a ?type .\n",
    "        ?uri skos:definition ?definition .\n",
    "        ?uri wtm:hasIngredient ?ingredient\n",
    "    } \n",
    "\"\"\"\n",
    "df = kg.query_as_df(sparql=sparql)\n",
    "\n",
    "# We group the ingredients into one column containing lists:\n",
    "recipes_df = df.groupby(['uri', 'definition', 'type'])['ingredient'].apply(list).reset_index(name='ingredients') ; recipes_df\n",
    "\n",
    "sparql_ingredients = \"\"\"\n",
    "    SELECT distinct *\n",
    "    WHERE {\n",
    "        ?uri a wtm:Ingredient .\n",
    "        ?uri a ?type .\n",
    "        OPTIONAL { ?uri skos:prefLabel ?definition } \n",
    "    }\n",
    "\"\"\"\n",
    "\n",
    "df = kg.query_as_df(sparql=sparql_ingredients)\n",
    "df['ingredients'] = None\n",
    "\n",
    "ing_recipes_df = pd.concat([recipes_df, df]).reset_index(drop=True)\n",
    "\n",
    "ing_recipes_df.fillna('', inplace=True) ; ing_recipes_df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Logging into Rubrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import rubrix as rb\n",
    "\n",
    "records = []\n",
    "for i, r in ing_recipes_df.iterrows():\n",
    "    item = rb.TextClassificationRecord(\n",
    "            inputs={\n",
    "                \"id\":r.uri, \n",
    "                \"definition\": r.definition,\n",
    "                \"ingredients\": str(r.ingredients), \n",
    "                \"type\": r.type\n",
    "            }, # log node fields\n",
    "            prediction=[(label, 0.0) for label in LABELS], # log \"dummy\" predictions for aiding annotation\n",
    "            metadata={'ingredients': [ing.replace('ind:','') for ing in r.ingredients], \"type\": r.type}, # metadata filters for quick exploration and annotation\n",
    "            prediction_agent=\"kglab_tutorial\", # who's performing/logging the prediction\n",
    "            multi_label=True\n",
    "        )\n",
    "    records.append(item)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(records)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rb.log(records=records, name=\"kg_classification_tutorial\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Annotation session with Rubrix (optional)\n",
    "\n",
    "In this step you can go to your rubrix dataset and annotate some examples of each class.\n",
    "\n",
    "If you have no time to do this, just skip this part as we have prepared a dataset for you with around ~600 examples."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Loading our labelled records and create a train_test split (optional)\n",
    "\n",
    "If you have no time to do this, just skip this part as we have prepared a dataset for you."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "rb.snapshots(name=\"kg_classification_tutorial\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once you have annotated your dataset, you will find an snapshot id on the previous list. This id should be place in the next command. In our case, it was 1620136587.907149."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df = rb.load(name=\"kg_classification_tutorial\", snapshot='1620136587.907149') ; df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "train_df, test_df = train_test_split(df)\n",
    "train_df.to_csv('data/train_recipes_new.csv')\n",
    "test_df.to_csv('data/test_recipes_new.csv')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating PyTorch train and test sets\n",
    "\n",
    "Here we take our train and test datasets and transform them into `torch.Tensor` objects with the help of our kglab `Subgraph` for turning `uris` into `torch.long` indices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "train_df = pd.read_csv('data/train_recipes.csv') # use your own labelled datasets if you've created a snapshot\n",
    "test_df = pd.read_csv('data/test_recipes.csv')\n",
    "\n",
    "# we make sure lists are parsed correctly\n",
    "train_df.labels = train_df.labels.apply(eval)\n",
    "test_df.labels = test_df.labels.apply(eval)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's create label lookups for label to int and viceversa"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "label2id = {label:i for i,label in enumerate(LABELS)} ; \n",
    "id2label = {i:l for l,i in label2id.items()} ; (id2label, label2id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following function turns our DataFrame into numerical arrays for node indices and labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "def create_indices_labels(df):\n",
    "    # turn our dense labels into a one-hot list\n",
    "    def one_hot(label_ids):\n",
    "        a = np.zeros(len(LABELS))\n",
    "        a.put(label_ids, np.ones(len(label_ids)))\n",
    "        return a\n",
    "    \n",
    "    indices, labels = [], []\n",
    "    for uri, label in zip(df.uri.tolist(), df.labels.tolist()):\n",
    "        indices.append(sg.transform(uri))\n",
    "        labels.append(one_hot([label2id[label] for label in label]))\n",
    "    return indices, labels"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, let's turn our dataset into PyTorch tensors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_indices, train_labels = create_indices_labels(train_df)\n",
    "test_indices, test_labels = create_indices_labels(test_df)\n",
    "\n",
    "train_idx = torch.tensor(train_indices, dtype=torch.long)\n",
    "train_y = torch.tensor(train_labels, dtype=torch.float)\n",
    "\n",
    "test_idx = torch.tensor(test_indices, dtype=torch.long)\n",
    "test_y = torch.tensor(test_labels, dtype=torch.float) ; train_idx[:10], train_y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see if we can recover the correct URIs for our numerical ids using our `kglab.Subgraph`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(train_df.loc[0], sg.inverse_transform(15380))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Creating a Subgraph of recipe and ingredient nodes\n",
    "Here we create a node list to be used as a seed for building our `PyG` subgraph (using k-hops as we will see in the next section). Our goal will be to start only with `recipes` and `ingredients`, as all nodes passed through the GNN will be classified and those are our main target. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "node_idx = torch.LongTensor([\n",
    "    sg.transform(i) for i in ing_recipes_df.uri.values\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "node_idx.max(), node_idx.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ing_recipes_df.iloc[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sg.inverse_transform(node_idx[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "node_idx[0:10]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Semi-supervised node classification with PyTorch Geometric\n",
    "\n",
    "For the node classification task **we are given the ground-truth labels** (our recipes and ingredients training set) **for a small subset of nodes**, and **we want to predict the labels for all the remaining nodes** (our recipes and ingredients test set and unlabelled nodes).\n",
    "\n",
    "\n",
    "### Graph Convolutional Networks\n",
    "\n",
    "To get a great intro to GCNs we recommend you to check Kipf's [blog post](https://tkipf.github.io/graph-convolutional-networks/) on the topic. \n",
    "\n",
    "In a nutshell, GCNs are multi-layer neural works which apply \"convolutions\" to nodes in graphs by sharing and applying the same filter parameters over all locations in the graph. \n",
    "\n",
    "Additionally, modern GCNs such as those implemented in `PyG` use **message passing** mechanisms, where vertices exchange information with their neighbors, and send messages to each other.\n",
    "\n",
    "![GCN_Kipf](img/gcn_web.png \"GCN\")\n",
    "\n",
    "***Multi-layer Graph Convolutional Network (GCN) with first-order filters***. Source: https://tkipf.github.io/graph-convolutional-networks\n",
    "\n",
    "###  Relational Graph Convolutional Networks\n",
    "\n",
    "Relational Graph Convolutional Networks (R-GCNs) were introduced by [Schlichtkrull et al. 2017](https://arxiv.org/abs/1703.06103), as an extension of GCNs to deal with **multi-relational knowledge graphs**. \n",
    "\n",
    "You can see below the computation model for nodes:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "![RGCN](img/rgcn.png \"RGCN\")\n",
    "\n",
    "***Computation of the update of a single graph node(red) in the R-GCN model.***. Source: https://arxiv.org/abs/1703.06103"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating a `PyG` subgraph\n",
    "\n",
    "Here we build a subgraph with `k` hops from target to source starting with all `recipe` and `ingredient` nodes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_geometric.utils import k_hop_subgraph\n",
    "# here we take all connected nodes with k hops\n",
    "k = 1\n",
    "node_idx, edge_index, mapping, edge_mask = k_hop_subgraph(\n",
    "    node_idx, \n",
    "    k, \n",
    "    data.edge_index, \n",
    "    relabel_nodes=False\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have increased the size of our node set:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "node_idx.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data.edge_index.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we compute some measures needed for defining the size of our layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data.edge_index = edge_index\n",
    "\n",
    "data.num_nodes = data.edge_index.max().item() + 1\n",
    "\n",
    "data.num_relations = data.edge_type.max().item() + 1\n",
    "\n",
    "data.edge_type = data.edge_type[edge_mask]\n",
    "\n",
    "data.num_classes = len(LABELS)\n",
    "\n",
    "data.num_nodes, data.num_relations, data.num_classes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Defining a basic Relational Graph Convolutional Network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_geometric.nn import FastRGCNConv, RGCNConv\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "RGCNConv?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RGCN(torch.nn.Module):\n",
    "    def __init__(self, num_nodes, num_relations, num_classes, out_channels=16, num_bases=30, dropout=0.0, layer_type=FastRGCNConv, ):\n",
    "        \n",
    "        super(RGCN, self).__init__()\n",
    "        \n",
    "        self.conv1 = layer_type(\n",
    "            num_nodes, \n",
    "            out_channels, \n",
    "            num_relations, \n",
    "            num_bases=num_bases\n",
    "        )\n",
    "        self.conv2 = layer_type(\n",
    "            out_channels, \n",
    "            num_classes, \n",
    "            num_relations, \n",
    "            num_bases=num_bases\n",
    "        )\n",
    "        self.dropout = torch.nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, edge_index, edge_type):\n",
    "        x = F.relu(self.conv1(None, edge_index, edge_type))\n",
    "        x = self.dropout(x)\n",
    "        x = self.conv2(x, edge_index, edge_type)\n",
    "        return torch.sigmoid(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create  and visualizing our model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = RGCN(\n",
    "    num_nodes=data.num_nodes,\n",
    "    num_relations=data.num_relations,\n",
    "    num_classes=data.num_classes,\n",
    "    #out_channels=64,\n",
    "    dropout=0.2,\n",
    "    layer_type=RGCNConv\n",
    ") ; model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# code adapted from https://colab.research.google.com/drive/14OvFnAXggxB8vM4e8vSURUp1TaKnovzX\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.manifold import TSNE\n",
    "from pytorch_lightning.metrics.utils import to_categorical\n",
    "\n",
    "def visualize(h, color, labels):\n",
    "    z = TSNE(n_components=2).fit_transform(h.detach().cpu().numpy())\n",
    "\n",
    "    plt.figure(figsize=(10,10))\n",
    "    plt.xticks([])\n",
    "    plt.yticks([])\n",
    "    \n",
    "    scatter = plt.scatter(z[:, 0], z[:, 1], s=70, c=color, cmap=\"Set2\")\n",
    "    legend = plt.legend(scatter.legend_elements()[0],labels, loc=\"upper right\", title=\"Labels\",) #*scatter.legend_elements()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pred = model(edge_index, edge_type)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "visualize(pred[train_idx], color=to_categorical(train_y), labels=LABELS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "visualize(pred[test_idx], color=to_categorical(test_y), labels=LABELS)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training our RGCN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cpu') # ('cuda')\n",
    "data = data.to(device)\n",
    "model = model.to(device)\n",
    "optimizer = torch.optim.AdamW(model.parameters())\n",
    "loss_module = torch.nn.BCELoss()\n",
    "\n",
    "def train():\n",
    "    model.train()\n",
    "    optimizer.zero_grad()\n",
    "    out = model(data.edge_index, data.edge_type)\n",
    "    loss = loss_module(out[train_idx], train_y)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    return loss.item()\n",
    "\n",
    "def accuracy(predictions, y):\n",
    "    predictions = np.round(predictions)\n",
    "    return predictions.eq(y).to(torch.float).mean()\n",
    "\n",
    "@torch.no_grad()\n",
    "def test():\n",
    "    model.eval()\n",
    "    pred = model(data.edge_index, data.edge_type)\n",
    "    train_acc = accuracy(pred[train_idx], train_y)\n",
    "    test_acc = accuracy(pred[test_idx], test_y)\n",
    "    return train_acc.item(), test_acc.item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "for epoch in range(1, 50):\n",
    "    loss = train()\n",
    "    train_acc, test_acc = test()\n",
    "    print(f'Epoch: {epoch:02d}, Loss: {loss:.4f}, Train: {train_acc:.4f} '\n",
    "          f'Test: {test_acc:.4f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pred = model(edge_index, edge_type)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "visualize(pred[train_idx], color=to_categorical(train_y), labels=LABELS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "visualize(pred[test_idx], color=to_categorical(test_y), labels=LABELS)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Using our model and analyzing its predictions with Rubrix\n",
    "Let's see the shape of our model predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pred = model(edge_index, edge_type) ; pred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def find(tensor, values):\n",
    "    return torch.nonzero(tensor[..., None] == values)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Analizing predictions over the test set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_idx = find(node_idx,test_idx)[:,0] ; len(test_idx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "index = torch.zeros(node_idx.shape[0], dtype=bool)\n",
    "index[test_idx] = True\n",
    "idx = node_idx[index]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "uris = [sg.inverse_transform(i) for i in idx]\n",
    "predicted_labels = [l for l in pred[idx]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predictions = list(zip(uris,predicted_labels)) ; predictions[0:2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import rubrix as rb\n",
    "\n",
    "records = []\n",
    "for uri,predicted_labels in predictions:\n",
    "    ids = ing_recipes_df.index[ing_recipes_df.uri == uri]\n",
    "    if len(ids) > 0:\n",
    "        r = ing_recipes_df.iloc[ids]\n",
    "        # get the gold labels from our test set\n",
    "        gold_labels = test_df.iloc[test_df.index[test_df.uri == uri]].labels.values[0]\n",
    "        \n",
    "        item = rb.TextClassificationRecord(\n",
    "                inputs={\"id\":r.uri.values[0], \"definition\": r.definition.values[0], \"ingredients\": str(r.ingredients.values[0]), \"type\": r.type.values[0]}, \n",
    "                prediction=[(id2label[i], score) for i,score in enumerate(predicted_labels)],\n",
    "                annotation=gold_labels,\n",
    "                metadata={'ingredients': r.ingredients.values[0], \"type\": r.type.values[0]}, \n",
    "                prediction_agent=\"node_classifier_v1\", \n",
    "                multi_label=True\n",
    "        )\n",
    "        records.append(item)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rb.log(records, name=\"kg_classification_test_analysis\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Analizing predictions over unseen nodes (and potentially relabeling them)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "Let's find the ids for the nodes in our training and test sets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_test_idx = find(node_idx,torch.cat((test_idx, train_idx)))[:,0] ; len(train_test_idx)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's get the ids, uris and labels of the nodes which were not in our train/test datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "index = torch.ones(node_idx.shape[0], dtype=bool)\n",
    "index[train_test_idx] = False\n",
    "idx = node_idx[index]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "We use our `SubgraphTensor` for getting back our URIs and build `uri,predicted_labels` pairs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "uris = [sg.inverse_transform(i) for i in idx]\n",
    "predicted_labels = [l for l in pred[idx]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "predictions = list(zip(uris,predicted_labels)) ; predictions[0:2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import rubrix as rb\n",
    "\n",
    "records = []\n",
    "for uri,predicted_labels in predictions:\n",
    "    ids = ing_recipes_df.index[ing_recipes_df.uri == uri]\n",
    "    if len(ids) > 0:\n",
    "        r = ing_recipes_df.iloc[ids]\n",
    "        item = rb.TextClassificationRecord(\n",
    "                inputs={\"id\":r.uri.values[0], \"definition\": r.definition.values[0], \"ingredients\": str(r.ingredients.values[0]), \"type\": r.type.values[0]}, \n",
    "                prediction=[(id2label[i], score) for i,score in enumerate(predicted_labels)], \n",
    "                metadata={'ingredients': r.ingredients.values[0], \"type\": r.type.values[0]}, \n",
    "                prediction_agent=\"node_classifier_v1\", \n",
    "                multi_label=True\n",
    "        )\n",
    "        records.append(item)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rb.log(records, name=\"kg_node_classification_unseen_nodes_v3\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 1: Training experiments with PyTorch Lightning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#!pip install wandb -qqq # optional"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wandb login #optional"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_geometric.data import Data, DataLoader\n",
    "\n",
    "data.train_idx = train_idx\n",
    "data.train_y = train_y\n",
    "data.test_idx = test_idx\n",
    "data.test_y = test_y\n",
    "\n",
    "dataloader = DataLoader([data], batch_size=1); dataloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import pytorch_lightning as pl\n",
    "from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint\n",
    "from pytorch_lightning.loggers import WandbLogger\n",
    "\n",
    "class RGCNNodeClassification(pl.LightningModule):\n",
    "    \n",
    "    def __init__(self, **model_kwargs):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.model = RGCN(**model_kwargs)\n",
    "        self.loss_module = torch.nn.BCELoss()\n",
    "    \n",
    "    def forward(self, edge_index, edge_type):\n",
    "        return self.model(edge_index, edge_type)\n",
    "        \n",
    "    def configure_optimizers(self):\n",
    "        optimizer = torch.optim.Adam(self.parameters(), lr=0.01, weight_decay=0.001)\n",
    "        return optimizer\n",
    "        \n",
    "    def training_step(self, batch, batch_idx):\n",
    "        idx, y = data.train_idx, data.train_y\n",
    "        edge_index, edge_type = data.edge_index, data.edge_type\n",
    "        x = self.forward(edge_index, edge_type)\n",
    "        loss = self.loss_module(x[idx], y)\n",
    "        x = x.detach()\n",
    "        self.log('train_acc', accuracy(x[idx], y), prog_bar=True)\n",
    "        self.log('train_loss', loss)\n",
    "        return loss \n",
    "        \n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        idx, y = data.test_idx, data.test_y\n",
    "        edge_index, edge_type = data.edge_index, data.edge_type\n",
    "        x = self.forward(edge_index, edge_type)\n",
    "        loss = self.loss_module(x[idx], y)\n",
    "        x = x.detach()\n",
    "        self.log('val_acc', accuracy(x[idx], y), prog_bar=True)\n",
    "        self.log('val_loss', loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pl.seed_everything()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_pl = RGCNNodeClassification(\n",
    "    num_nodes=data.num_nodes,\n",
    "    num_relations=data.num_relations,\n",
    "    num_classes=data.num_classes,\n",
    "    #out_channels=64,\n",
    "    dropout=0.2,\n",
    "    #layer_type=RGCNConv\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "early_stopping = EarlyStopping(monitor='val_acc', patience=10, mode='max')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = pl.Trainer(\n",
    "    default_root_dir='pl_runs',\n",
    "    checkpoint_callback=ModelCheckpoint(save_weights_only=True, mode=\"max\", monitor=\"val_acc\"),\n",
    "    max_epochs=200,\n",
    "    #logger= WandbLogger(), # optional\n",
    "    callbacks=[early_stopping]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "trainer.fit(model_pl, dataloader, dataloader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 2: Bootstrapping annotation with a zeroshot-classifier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install transformers -qqq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import pipeline\n",
    "  \n",
    "pretrained_model =  \"valhalla/distilbart-mnli-12-1\" # \"typeform/squeezebert-mnli\"\n",
    "\n",
    "pl = pipeline('zero-shot-classification', model=pretrained_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pl(\"chocolate cake\", LABELS, hypothesis_template='The flavour is {}.',multi_label=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import rubrix as rb\n",
    "\n",
    "records = []\n",
    "for i, r in ing_recipes_df[50:150].iterrows():\n",
    "    preds = pl(r.definition, LABELS, hypothesis_template='The flavour is {}.', multi_label=True)\n",
    "    item = rb.TextClassificationRecord(\n",
    "            inputs={\n",
    "                \"id\":r.uri, \n",
    "                \"definition\": r.definition,\n",
    "                \"ingredients\": str(r.ingredients), \n",
    "                \"type\": r.type\n",
    "            }, \n",
    "            prediction=list(zip(preds['labels'], preds['scores'])), # TODO: here we log he predictions of our zeroshot pipeline as a list of tuples (label, score)\n",
    "            metadata={'ingredients': r.ingredients, \"type\": r.type}, \n",
    "            prediction_agent=\"valhalla/distilbart-mnli-12-1\", \n",
    "            multi_label=True\n",
    "        )\n",
    "    records.append(item)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rb.log(records, name='kg_zeroshot')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Next steps\n",
    "\n",
    "### 📚 [Rubrix documentation](https://docs.rubrix.ml) for more guides and tutorials. \n",
    "\n",
    "### 🙋‍♀️ Join the Rubrix community! A good place to start is the [discussion forum](https://github.com/recognai/rubrix/discussions).\n",
    "\n",
    "### ⭐ Rubrix [Github repo](https://github.com/recognai/rubrix) to stay updated. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "b709380ea7d1cb2eb4650c0f11ac7e002ec6a534602815725771481b4784238c"
  },
  "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"
  },
  "nteract": {
   "version": "0.28.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}