Pokec-Pyg-Neo4j.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "dd0c4312",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define Neo4j connections\n",
    "import pandas as pd\n",
    "from neo4j import GraphDatabase\n",
    "host = 'bolt://localhost:7687'\n",
    "user = 'neo4j'\n",
    "password = 'letmein'\n",
    "driver = GraphDatabase.driver(host,auth=(user, password))\n",
    "\n",
    "def run_query(query, params={}):\n",
    "    with driver.session() as session:\n",
    "        result = session.run(query, params)\n",
    "        return pd.DataFrame([r.values() for r in result], columns=result.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "628ba550",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_frames():\n",
    "    \"\"\"\n",
    "    Loads the nodes and edges from Neo4j.\n",
    "    :return: nodes and edges frame\n",
    "    \"\"\"\n",
    "    dfn = run_query(\"\"\"\n",
    "    MATCH (u:User)\n",
    "    RETURN u.id AS id, u.age AS age, u.gender AS gender\n",
    "    \"\"\")\n",
    "    dfn = dfn.set_index(\"id\")\n",
    "\n",
    "    dfe = run_query(\"\"\"\n",
    "    MATCH (s:User)-[:FRIEND]->(t:User)\n",
    "    RETURN s.id as source, t.id as target\n",
    "    \"\"\")\n",
    "\n",
    "    return dfn, dfe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "bf36418f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>source</th>\n",
       "      <th>target</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>16</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>10</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1</td>\n",
       "      <td>12</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1</td>\n",
       "      <td>8</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1</td>\n",
       "      <td>7</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  source target\n",
       "0      1     16\n",
       "1      1     10\n",
       "2      1     12\n",
       "3      1      8\n",
       "4      1      7"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>age</th>\n",
       "      <th>gender</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>26</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16</th>\n",
       "      <td>23</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>29</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>26</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>17</th>\n",
       "      <td>27</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    age  gender\n",
       "id             \n",
       "1    26       1\n",
       "16   23       1\n",
       "3    29       1\n",
       "4    26       0\n",
       "17   27       0"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.display import display\n",
    "\n",
    "dfn, dfe = load_frames()\n",
    "display(dfe.head())\n",
    "display(dfn.head())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "81eec21f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "#device = 'cpu'\n",
    "print(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f999a2c2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Transforming nodes\n"
     ]
    }
   ],
   "source": [
    "def transform_nodes(node_frame):\n",
    "    print(\"Transforming nodes\")\n",
    "    # sorting the index does not make sense here\n",
    "    node_index_map = {str(index): i for i, index in enumerate(node_frame.index.unique())}\n",
    "    gender_series = node_frame[\"gender\"]\n",
    "    gender_tensor = torch.zeros(len(gender_series), 2, dtype = torch.float)\n",
    "    for i, v in enumerate(gender_series.values):\n",
    "        gender_tensor[i, 0 if np.isnan(v) else int(v)] = 1.0\n",
    "    age_tensor = torch.tensor(node_frame['age'].values, dtype = torch.float).reshape(len(gender_series), -1)\n",
    "    x = torch.cat((gender_tensor, age_tensor), dim = -1)  # 1x3 tensor\n",
    "    return x, node_index_map\n",
    "\n",
    "nodes_x, nodes_mapping = transform_nodes(dfn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b23dbb9a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Transforming edges\n"
     ]
    }
   ],
   "source": [
    "def transform_edges(edge_frame, nodes_mapping):\n",
    "    print(\"Transforming edges\")\n",
    "\n",
    "    src = [nodes_mapping[src_id] if src_id in nodes_mapping else -1 for src_id in edge_frame.source]\n",
    "    dst = [nodes_mapping[tgt_id] if tgt_id in nodes_mapping else -1 for tgt_id in edge_frame.target]\n",
    "    edge_index = torch.tensor([src, dst])\n",
    "\n",
    "    return edge_index, None\n",
    "\n",
    "edges_index, edges_label = transform_edges(dfe, nodes_mapping)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b63c739c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_geometric.data import Data\n",
    "\n",
    "data = Data(x = nodes_x, edge_index = edges_index, edge_attr = edges_label, y = None).to(device, non_blocking=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "3d75a12e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Creating classification masks\n"
     ]
    }
   ],
   "source": [
    " def create_node_masks(d):\n",
    "        print(\"Creating classification masks\")\n",
    "        amount = len(d.x)\n",
    "        # actually the index to the nodes\n",
    "        nums = np.arange(amount)\n",
    "        np.random.shuffle(nums)\n",
    "\n",
    "        train_size = int(amount * 0.7)\n",
    "        test_size = int(amount * 0.85) - train_size\n",
    "        val_size = amount - train_size - test_size\n",
    "\n",
    "        train_set = nums[0:train_size]\n",
    "        test_set = nums[train_size:train_size + test_size]\n",
    "        val_set = nums[train_size + test_size:]\n",
    "\n",
    "        assert len(train_set) + len(test_set) + len(val_set) == amount, \"The split should be coherent.\"\n",
    "\n",
    "        train_mask = torch.zeros(amount, dtype = torch.long, device = device)\n",
    "        for i in train_set:\n",
    "            train_mask[i] = 1.\n",
    "\n",
    "        test_mask = torch.zeros(amount, dtype = torch.long, device = device)\n",
    "        for i in test_set:\n",
    "            test_mask[i] = 1.\n",
    "\n",
    "        val_mask = torch.zeros(amount, dtype = torch.long, device = device)\n",
    "        for i in val_set:\n",
    "            val_mask[i] = 1.\n",
    "\n",
    "        d.train_mask = train_mask\n",
    "        d.test_mask = test_mask\n",
    "        d.val_mask = val_mask\n",
    "        \n",
    "create_node_masks(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "dc9a4fb1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Data(x=[1099121, 3], edge_index=[2, 21575162], train_mask=[1099121], test_mask=[1099121], val_mask=[1099121], edge_label=[10787581], edge_label_index=[2, 10787581]),\n",
       " Data(x=[1099121, 3], edge_index=[2, 21575162], train_mask=[1099121], test_mask=[1099121], val_mask=[1099121], edge_label=[10794], edge_label_index=[2, 10794]),\n",
       " Data(x=[1099121, 3], edge_index=[2, 21585956], train_mask=[1099121], test_mask=[1099121], val_mask=[1099121], edge_label=[2158], edge_label_index=[2, 2158]))"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch_geometric.transforms as T\n",
    "\n",
    "transform = T.Compose([\n",
    "    T.ToUndirected(merge = True),\n",
    "    T.ToDevice(device),\n",
    "    T.RandomLinkSplit(num_val = 0.0005, num_test = 0.0001, is_undirected = True, add_negative_train_samples = False),\n",
    "])\n",
    "transform(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "9d029460",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_geometric.loader import NeighborLoader\n",
    "\n",
    "# the larger the batch size the faster things will be\n",
    "batch_size = 2048\n",
    "\n",
    "# define batch loaders for the three sets\n",
    "train_loader = NeighborLoader(data, num_neighbors = [10] * 2, shuffle = True, input_nodes = data.train_mask, batch_size = batch_size)\n",
    "val_loader = NeighborLoader(data, num_neighbors = [10] * 2, input_nodes = data.val_mask, batch_size = batch_size)\n",
    "test_loader = NeighborLoader(data, num_neighbors = [10] * 2, input_nodes = data.test_mask, batch_size = batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "c7567023",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import datetime\n",
    "from tqdm import tqdm\n",
    "from torch_geometric.utils import negative_sampling\n",
    "from torch_geometric.nn import GCNConv\n",
    "from sklearn.metrics import roc_auc_score, f1_score\n",
    "import os\n",
    "\n",
    "\n",
    "class Net(torch.nn.Module):\n",
    "    def __init__(self, in_channels, hidden_channels, out_channels):\n",
    "        super().__init__()\n",
    "        self.conv1 = GCNConv(in_channels, hidden_channels)\n",
    "        self.conv2 = GCNConv(hidden_channels, out_channels)\n",
    "\n",
    "    def encode(self, x, edge_index):\n",
    "        # chaining two convolutions with a standard relu activation\n",
    "\n",
    "        x = self.conv1(x, edge_index).relu()\n",
    "        return self.conv2(x, edge_index)\n",
    "\n",
    "    def decode(self, z, edge_label_index):\n",
    "        # cosine similarity\n",
    "        return (z[edge_label_index[0]] * z[edge_label_index[1]]).sum(dim = -1)\n",
    "\n",
    "    def decode_all(self, z):\n",
    "        prob_adj = z @ z.t()\n",
    "        return (prob_adj > 0).nonzero(as_tuple = False).t()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "eb09075b",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Net(data.num_features, 128, 64).to(device)\n",
    "optimizer = torch.optim.Adam(params = model.parameters(), lr = 0.01)\n",
    "# BCELoss creates a criterion that measures the Binary Cross Entropy between the target and the output.\n",
    "criterion = torch.nn.BCEWithLogitsLoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "c040ceea",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train():\n",
    "    \"\"\"\n",
    "    Single epoch model training in batches.\n",
    "    :return: total loss for the epoch\n",
    "    \"\"\"\n",
    "    model.train()\n",
    "    total_examples = total_loss = 0\n",
    "    for batch in tqdm(train_loader):\n",
    "        optimizer.zero_grad()\n",
    "        batch = batch.to(device)\n",
    "        batch_size = batch.batch_size\n",
    "        z = model.encode(batch.x, batch.edge_index)\n",
    "        neg_edge_index = negative_sampling(edge_index = batch.edge_index, num_nodes = batch.num_nodes, num_neg_samples = None, method = 'sparse')\n",
    "        edge_label_index = torch.cat([batch.edge_index, neg_edge_index], dim = -1, )\n",
    "        edge_label = torch.cat([torch.ones(batch.edge_index.size(1)), torch.zeros(neg_edge_index.size(1))], dim = 0).to(device)\n",
    "        out = model.decode(z, edge_label_index).view(-1)\n",
    "        # loss = criterion(out[:batch_size], edge_label[:batch_size])\n",
    "        loss = criterion(out, edge_label)\n",
    "        # standard torch mechanics here\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        total_examples += batch_size\n",
    "        total_loss += float(loss) * batch_size\n",
    "    return total_loss / total_examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "488e0b72",
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def test(loader):\n",
    "    \"\"\"\n",
    "    Evalutes the model on the test set.\n",
    "    :param loader: the batch loader\n",
    "    :return: a score\n",
    "    \"\"\"\n",
    "    model.eval()\n",
    "    scores = []\n",
    "    threshold = torch.tensor([0.7]).to(device)\n",
    "    for batch in tqdm(loader):\n",
    "        batch.to(device)\n",
    "        z = model.encode(batch.x, batch.edge_index)\n",
    "        out = model.decode(z, batch.edge_index).view(-1).sigmoid()\n",
    "        pred = (out > threshold).float() * 1\n",
    "        score = f1_score(np.ones(batch.edge_index.size(1)), pred.cpu().numpy())\n",
    "        scores.append(score)\n",
    "    return np.average(scores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "3f966307",
   "metadata": {},
   "outputs": [],
   "source": [
    "def predictions(max = 1000, threshold = 0.99):\n",
    "    \"\"\"\n",
    "    Creates predictions for the specified run.\n",
    "    :param run_id: model id\n",
    "    :param max: the maximum amount of predictions to output\n",
    "    \"\"\"\n",
    "    pred_edges = []\n",
    "\n",
    "    loader = NeighborLoader(data, num_neighbors = [10] * 2, shuffle = True, input_nodes = None, batch_size = batch_size)\n",
    "    threshold_tensor = torch.tensor([threshold]).to(device)\n",
    "    for batch in tqdm(loader):\n",
    "        batch.to(device)\n",
    "        z = model.encode(batch.x, batch.edge_index)\n",
    "        # collecting negative edge tuples ensure that the decode are actual non-existing edges\n",
    "        neg_edge_index = negative_sampling(edge_index = batch.edge_index, num_nodes = None, num_neg_samples = None, method = 'sparse')\n",
    "        out = model.decode(z, neg_edge_index).view(-1).sigmoid()\n",
    "        pred = ((out > threshold_tensor).float() * 1).cpu().numpy()\n",
    "        found = np.argwhere(pred == 1)\n",
    "        if found.size > 0:\n",
    "            edge_tuples = neg_edge_index.t().cpu().numpy()\n",
    "            select_index = found.reshape(1, found.size)[0]\n",
    "            edges = edge_tuples[select_index]\n",
    "            pred_edges += edges.tolist()\n",
    "            if len(pred_edges) >= max:\n",
    "                break\n",
    "    \n",
    "    return pd.DataFrame.from_dict([{'source': a, 'target': b} for a,b in pred_edges])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "8b5a8de5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run():\n",
    "    \"\"\"\n",
    "        Run the training and makes predictions.\n",
    "    \"\"\"\n",
    "    run_id = int(datetime.timestamp(datetime.now()))\n",
    "    start_time = datetime.now()\n",
    "    epochs = 10\n",
    "    #with trange(epochs + 1) as t:\n",
    "    for epoch in range(epochs):\n",
    "        try:\n",
    "            #t.set_description('Epoch %i/%i train' % (epoch, epochs))\n",
    "            loss = train()\n",
    "            #t.set_description('Epoch %i/%i test' % (epoch, epochs))\n",
    "            val_acc = test(test_loader)\n",
    "            #t.set_postfix(loss = loss, accuracy = val_acc)\n",
    "            print(f\"Epoch: {epoch:03d}, Loss: {loss:.4f}, Acc: {val_acc:.4f}\")\n",
    "        except KeyboardInterrupt:\n",
    "            break\n",
    "    torch.save(model.state_dict(), f\"model_{run_id}\")\n",
    "    time_elapsed = datetime.now() - start_time\n",
    "    print(\"Creating predictions\")\n",
    "    print(f\"\\nRun {run_id}:\")\n",
    "    print(f\"\\tEpochs: {epoch}\")\n",
    "    print(f\"\\tTime: {time_elapsed}\")\n",
    "    print(f\"\\tAccuracy: {val_acc * 100:.01f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "2a1034da",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 537/537 [00:12<00:00, 43.54it/s]\n",
      "100%|███████████████████████████████████████████████████████████████████████████████| 537/537 [00:05<00:00, 103.52it/s]\n",
      "  1%|▌                                                                                 | 4/537 [00:00<00:15, 33.33it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 000, Loss: 0.5306, Acc: 0.9035\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 537/537 [00:12<00:00, 42.16it/s]\n",
      "100%|███████████████████████████████████████████████████████████████████████████████| 537/537 [00:05<00:00, 101.66it/s]\n",
      "  1%|▌                                                                                 | 4/537 [00:00<00:15, 35.41it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 001, Loss: 0.5471, Acc: 0.8040\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 537/537 [00:14<00:00, 37.75it/s]\n",
      "100%|███████████████████████████████████████████████████████████████████████████████| 537/537 [00:05<00:00, 102.78it/s]\n",
      "  1%|▌                                                                                 | 4/537 [00:00<00:15, 35.10it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 002, Loss: 0.5342, Acc: 0.9006\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 537/537 [00:14<00:00, 37.67it/s]\n",
      "100%|███████████████████████████████████████████████████████████████████████████████| 537/537 [00:05<00:00, 103.72it/s]\n",
      "  1%|▌                                                                                 | 4/537 [00:00<00:15, 34.49it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 003, Loss: 0.5306, Acc: 0.8986\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 537/537 [00:14<00:00, 37.40it/s]\n",
      "100%|███████████████████████████████████████████████████████████████████████████████| 537/537 [00:05<00:00, 103.96it/s]\n",
      "  1%|▌                                                                                 | 4/537 [00:00<00:14, 36.05it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 004, Loss: 0.5258, Acc: 0.9004\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 537/537 [00:14<00:00, 37.50it/s]\n",
      "100%|███████████████████████████████████████████████████████████████████████████████| 537/537 [00:05<00:00, 102.81it/s]\n",
      "  1%|█                                                                                 | 7/537 [00:00<00:15, 35.30it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 005, Loss: 0.5240, Acc: 0.8875\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 537/537 [00:14<00:00, 36.81it/s]\n",
      "100%|███████████████████████████████████████████████████████████████████████████████| 537/537 [00:05<00:00, 102.52it/s]\n",
      "  1%|▌                                                                                 | 4/537 [00:00<00:15, 33.34it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 006, Loss: 0.5188, Acc: 0.8915\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 537/537 [00:14<00:00, 37.24it/s]\n",
      "100%|███████████████████████████████████████████████████████████████████████████████| 537/537 [00:05<00:00, 102.36it/s]\n",
      "  1%|▌                                                                                 | 4/537 [00:00<00:16, 32.26it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 007, Loss: 0.5039, Acc: 0.8822\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 537/537 [00:14<00:00, 36.09it/s]\n",
      "100%|███████████████████████████████████████████████████████████████████████████████| 537/537 [00:05<00:00, 101.14it/s]\n",
      "  1%|▌                                                                                 | 4/537 [00:00<00:15, 35.41it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 008, Loss: 0.5037, Acc: 0.8831\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 537/537 [00:14<00:00, 36.74it/s]\n",
      "100%|███████████████████████████████████████████████████████████████████████████████| 537/537 [00:05<00:00, 101.92it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 009, Loss: 0.4971, Acc: 0.8811\n",
      "Creating predictions\n",
      "\n",
      "Run 1640882244:\n",
      "\tEpochs: 9\n",
      "\tTime: 0:03:13.095027\n",
      "\tAccuracy: 88.1\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "584d63ac",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|                                                                                          | 0/537 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "   source  target\n",
      "0    9257   85860\n",
      "1   18041   83302\n",
      "2   47626    8697\n",
      "3   16195   96111\n",
      "4  108720   15420\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "preds_df = predictions()\n",
    "print(preds_df.head())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "id": "b9aa1d0e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>already_exists</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   already_exists\n",
       "0               1"
      ]
     },
     "execution_count": 65,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Check how many already exist\n",
    "\n",
    "run_query(\"\"\"\n",
    "UNWIND $data AS row\n",
    "MATCH (s:User)-[:FRIEND]-(t:User)\n",
    "WHERE s.id = toString(row[0]) AND t.id = toString(row[1])\n",
    "RETURN count(*) AS already_exists\n",
    "\"\"\", {'data': preds_df.values.tolist()})\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "121fd428",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}