extract_features_sage.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gzip\n",
    "import gc\n",
    "import math\n",
    "import json\n",
    "from collections import Counter, defaultdict\n",
    "import random\n",
    "from tqdm.notebook import tqdm\n",
    "import numpy as np\n",
    "\n",
    "from sklearn.model_selection import train_test_split\n",
    "import matplotlib.pyplot as plt\n",
    "from collections import Counter\n",
    "import pandas as pd\n",
    "\n",
    "from stellargraph import StellarGraph\n",
    "from scipy.sparse import load_npz\n",
    "\n",
    "def jl_to_list(fname):\n",
    "    output = []\n",
    "    with gzip.open(fname, \"rb\") as f:\n",
    "        for line in f:\n",
    "            output.append(json.loads(line))\n",
    "    return output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load user data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "samples = None\n",
    "test_size = .2\n",
    "rows = jl_to_list(\"data/train_dataset.jl.gz\")\n",
    "if samples:\n",
    "    rows = rows[:samples]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train-test split (revise)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train data: 330530.\n",
      "Test data: 82633.\n"
     ]
    }
   ],
   "source": [
    "val_size = .15\n",
    "rows_train, rows_test = train_test_split(rows, test_size = test_size, random_state = 42)\n",
    "print(f\"Train data: {len(rows_train)}.\")\n",
    "print(f\"Test data: {len(rows_test)}.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Function to return the count of users who searched item pairs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_dict(rows):\n",
    "    counts = Counter()\n",
    "    for i, row in tqdm(enumerate(rows)):\n",
    "        viewed = [item[\"event_info\"] for item in row[\"user_history\"]\n",
    "                  if item[\"event_type\"] == \"view\"]\n",
    "\n",
    "        for i, v1 in enumerate(viewed):\n",
    "            for v2 in viewed[i + 1:]:\n",
    "                counts[str(v1) + \"_\" + str(v2)] += 1\n",
    "    \n",
    "    return counts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f419bc21c6744642b294d0d48f2b6d53",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(HTML(value=''), FloatProgress(value=1.0, bar_style='info', layout=Layout(width='20px'), max=1.0…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "c = create_dict(rows_train)\n",
    "del rows\n",
    "del rows_train\n",
    "del rows_test"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load item data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "item_data = jl_to_list(\"data/item_data.jl.gz\")\n",
    "metadata = {x[\"item_id\"] : x for x in item_data}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Map a unique integer to each item for Stellar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4e82a698470341e39a81ea782b01f8d9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(HTML(value=''), FloatProgress(value=1.0, bar_style='info', layout=Layout(width='20px'), max=1.0…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "ids_map = {}\n",
    "for i, item_id in tqdm(enumerate(metadata.keys())):\n",
    "    ids_map[str(item_id)] = i\n",
    "\n",
    "del item_data\n",
    "del metadata"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Function to get edges"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_edges(counts, ids_map, thresh = 2):\n",
    "    source = []\n",
    "    target = []\n",
    "    for ids, val in tqdm(counts.items()):\n",
    "        if val >= thresh:\n",
    "            indx = ids.index(\"_\")\n",
    "            id1 = ids_map[ids[:indx]]\n",
    "            id2 = ids_map[ids[indx+1:]]\n",
    "            source += [id1]\n",
    "            target += [id2]\n",
    "        \n",
    "    return source, target"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "49068df9a0364cbda8bac9327c50a690",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(HTML(value=''), FloatProgress(value=0.0, max=30115117.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "t = 2\n",
    "source, target = get_edges(c, ids_map, t)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load feature matrix (sparse)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "item_feat = load_npz('data/item_feats_v2.npz')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Convert feature matrix to dense"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3.6192800944\n"
     ]
    }
   ],
   "source": [
    "import sys\n",
    "\n",
    "dtype = np.float32\n",
    "item_feats = item_feat.toarray().astype(dtype)\n",
    "print(sys.getsizeof(item_feats)/(10e9))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.        , 1.        , 0.        , ..., 0.        , 0.        ,\n",
       "       0.04414384], dtype=float32)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "item_feats[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create edges dataframe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "edges = pd.DataFrame({\"source\": source, \"target\": target})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Finally, create the graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "StellarGraph: Undirected multigraph\n",
      " Nodes: 2102277, Edges: 15322937\n",
      "\n",
      " Node types:\n",
      "  default: [2102277]\n",
      "    Features: float32 vector, length 4304\n",
      "    Edge types: default-default->default\n",
      "\n",
      " Edge types:\n",
      "    default-default->default: [15322937]\n",
      "        Weights: all 1 (default)\n",
      "        Features: none\n"
     ]
    }
   ],
   "source": [
    "G = StellarGraph(item_feats, edges, is_directed = False, dtype = np.float32)\n",
    "print(G.info())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lets start training! GraphSAGE Imports:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "import networkx as nx\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import os\n",
    "import random\n",
    "\n",
    "import stellargraph as sg\n",
    "from stellargraph.data import EdgeSplitter\n",
    "from stellargraph.mapper import GraphSAGELinkGenerator\n",
    "from stellargraph.layer import GraphSAGE, link_classification\n",
    "from stellargraph.data import UniformRandomWalk\n",
    "from stellargraph.data import UnsupervisedSampler\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "from sklearn import preprocessing, feature_extraction, model_selection\n",
    "from sklearn.linear_model import LogisticRegressionCV, LogisticRegression\n",
    "from sklearn.metrics import accuracy_score\n",
    "\n",
    "from stellargraph import globalvar\n",
    "\n",
    "from stellargraph import datasets\n",
    "from IPython.display import display, HTML"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Check if we are running on gpu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[[22. 28.]\n",
      " [49. 64.]], shape=(2, 2), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "tf.debugging.set_log_device_placement(True)\n",
    "\n",
    "# Place tensors on the GPU\n",
    "with tf.device(\"/GPU:0\"):\n",
    "    a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])\n",
    "    b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])\n",
    "\n",
    "# Run on the GPU\n",
    "C = tf.matmul(a, b)\n",
    "print(C)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "nodes = list(G.nodes())\n",
    "number_of_walks = 1\n",
    "length = 5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Make validation split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "node_indexes = np.arange(len(nodes))\n",
    "val_size = .15\n",
    "nodes_train, nodes_val = train_test_split(node_indexes, test_size = val_size, random_state = 42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "unsupervised_samples_train = UnsupervisedSampler(\n",
    "    G, nodes=nodes_train, length=length, number_of_walks=number_of_walks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "unsupervised_samples_val = UnsupervisedSampler(\n",
    "    G, nodes=nodes_val, length=length, number_of_walks=number_of_walks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[  66946  292836  120815 1076367 1971989  684678   30869  464270 1365193\n",
      " 1027760]\n",
      "[1473523  331569  446851 2060933   60144  139159  104776 1531393 1462210\n",
      "  501768]\n"
     ]
    }
   ],
   "source": [
    "print(nodes_train[:10])\n",
    "print(nodes_val[:10])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# More parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 50\n",
    "epochs = 1\n",
    "num_samples = [12, 6]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "generator = GraphSAGELinkGenerator(G, batch_size, num_samples)\n",
    "train_gen = generator.flow(unsupervised_samples_train)\n",
    "#val_gen = generator.flow(unsupervised_samples_val)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create GraphSAGE model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "layer_sizes = [64, 64]\n",
    "graphsage = GraphSAGE(\n",
    "    layer_sizes=layer_sizes, generator=generator, bias=True, dropout=0.25, normalize=\"l2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_inp, x_out = graphsage.in_out_tensors()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "link_classification: using 'ip' method to combine node embeddings into edge embeddings\n"
     ]
    }
   ],
   "source": [
    "prediction = link_classification(\n",
    "    output_dim=1, output_act=\"sigmoid\", edge_embedding_method=\"ip\"\n",
    ")(x_out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = keras.Model(inputs=x_inp, outputs=prediction)\n",
    "\n",
    "model.compile(\n",
    "    optimizer=keras.optimizers.Adam(lr=2e-3),\n",
    "    loss=keras.losses.binary_crossentropy,\n",
    "    metrics=[keras.metrics.binary_accuracy],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Its time to train!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 66662/170749 [==========>...................] - ETA: 5:30:02 - loss: 0.5562 - binary_accuracy: 0.7606"
     ]
    }
   ],
   "source": [
    "with tf.device(\"/GPU:0\"):\n",
    "    history = model.fit(\n",
    "        train_gen,\n",
    "        epochs=epochs,\n",
    "        verbose=1,\n",
    "        use_multiprocessing=False,\n",
    "        workers=4,\n",
    "        shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_inp_src = x_inp[0::2]\n",
    "x_out_src = x_out[0]\n",
    "embedding_model = keras.Model(inputs=x_inp_src, outputs=x_out_src)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: sage_v0/assets\n",
      "Model saved bro.\n"
     ]
    }
   ],
   "source": [
    "embedding_model.save(\"sage_v0\")\n",
    "print(\"Model saved bro.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}