notebook.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fashion Retail Recommendations Using Neo4j Graph Data Science and Apache Arrow"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!kaggle competitions download -c h-and-m-personalized-fashion-recommendations -f articles.csv -p data\n",
    "!kaggle competitions download -c h-and-m-personalized-fashion-recommendations -f customers.csv -p data \n",
    "!kaggle competitions download -c h-and-m-personalized-fashion-recommendations -f transactions_train.csv -p data\n",
    "!cd data && unzip -n '*.zip'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pyarrow as pa\n",
    "import pyarrow.flight as flight\n",
    "import base64\n",
    "from pyarrow import compute as pacompute\n",
    "from pyarrow import csv as pacsv\n",
    "from config import NEO4J_URI,NEO4J_USER,NEO4J_PASSWORD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "token = base64.b64encode(f\"\"\"{NEO4J_USER}:{NEO4J_PASSWORD}\"\"\".encode('utf8'))\n",
    "options = flight.FlightCallOptions(headers=[\n",
    "            (b'authorization', b'Basic ' + token)\n",
    "        ])\n",
    "client = pa.flight.connect(\"grpc+tcp://localhost:8491\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ActionType(type='CREATE_GRAPH', description=''), ActionType(type='CREATE_DATABASE', description=''), ActionType(type='NODE_LOAD_DONE', description=''), ActionType(type='RELATIONSHIP_LOAD_DONE', description=''), ActionType(type='ABORT', description='')]\n"
     ]
    }
   ],
   "source": [
    "print(client.list_actions(options=options))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initialize the process"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "graph_name_import = \"HEM.NODES2022.NEW\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # CREATE A IN-MEMORY GRAPH\n",
    "# action = flight.Action(\"CREATE_GRAPH\", \"\"\"{\n",
    "#     \"name\": \"HEMNODES2022\",\n",
    "#     \"database_name\": \"hem\",\n",
    "#     \"concurrency\": 4\n",
    "# }\"\"\".encode('utf-8'))\n",
    "\n",
    "# CREATE A DATABASE\n",
    "action = flight.Action(\"CREATE_DATABASE\", f\"\"\"{{\n",
    "    \"name\": \"{graph_name_import}\",\n",
    "    \"concurrency\": 4\n",
    "}}\"\"\".encode('utf-8'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\"name\":\"HEM.NODES2022.NEW\"}\n"
     ]
    }
   ],
   "source": [
    "res = client.do_action(action, options = options)\n",
    "for r in res:\n",
    "    print(r.body.to_pybytes().decode('utf-8'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Write nodes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Load the csv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "articles = pacsv.read_csv('data/articles.csv')\n",
    "customers = pacsv.read_csv('data/customers.csv')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### write articles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "pyarrow.Table\n",
       "article_id: int64\n",
       "product_code: int64\n",
       "prod_name: string\n",
       "product_type_no: int64\n",
       "product_type_name: string\n",
       "product_group_name: string\n",
       "graphical_appearance_no: int64\n",
       "graphical_appearance_name: string\n",
       "colour_group_code: int64\n",
       "colour_group_name: string\n",
       "perceived_colour_value_id: int64\n",
       "perceived_colour_value_name: string\n",
       "perceived_colour_master_id: int64\n",
       "perceived_colour_master_name: string\n",
       "department_no: int64\n",
       "department_name: string\n",
       "index_code: string\n",
       "index_name: string\n",
       "index_group_no: int64\n",
       "index_group_name: string\n",
       "section_no: int64\n",
       "section_name: string\n",
       "garment_group_no: int64\n",
       "garment_group_name: string\n",
       "detail_desc: string\n",
       "----\n",
       "article_id: [[108775015,108775044,108775051,110065001,110065002,...,453272036,453272037,453272038,453272042,453272045],[453272047,453272048,453272049,453272051,453272053,...,506495004,506495006,506495009,506495012,506503002],...,[896477004,896506001,896506002,896509001,896522001,...,920012003,920016001,920016002,920084003,920084004],[920158001,920159001,920159002,920159003,920163001,...,953450001,953763001,956217002,957375001,959461001]]\n",
       "product_code: [[108775,108775,108775,110065,110065,...,453272,453272,453272,453272,453272],[453272,453272,453272,453272,453272,...,506495,506495,506495,506495,506503],...,[896477,896506,896506,896509,896522,...,920012,920016,920016,920084,920084],[920158,920159,920159,920159,920163,...,953450,953763,956217,957375,959461]]\n",
       "prod_name: [[\"Strap top\",\"Strap top\",\"Strap top (1)\",\"OP T-shirt (Idro)\",\"OP T-shirt (Idro)\",...,\"SB TRUNKS\",\"SB TRUNKS\",\"SB TRUNKS\",\"SB TRUNKS\",\"SB TRUNKS\"],[\"SB TRUNKS\",\"SB TRUNKS\",\"SB TRUNKS\",\"SB TRUNKS\",\"SB TRUNKS\",...,\"Kimye hood\",\"Kimye hood\",\"Kimye Hood\",\"Kimye Hood\",\"Carly\"],...,[\"Pen skirt\",\"Byron jumpsuit\",\"Byron jumpsuit\",\"John Chelsea BG\",\"PLATON set 2pcs\",...,\"Naomi sweater\",\"PETITE Merry HW Slim Mom dnm\",\"PETITE Merry HW Slim Mom dnm\",\"Mull\",\"Mull\"],[\"SC SADISFACTION sweatshirt\",\"Alexia top\",\"Alexia top\",\"Alexia top\",\"PETITE Brentwood HW Trs\",...,\"5pk regular Placement1\",\"SPORT Malaga tank\",\"Cartwheel dress\",\"CLAIRE HAIR CLAW\",\"Lounge dress\"]]\n",
       "product_type_no: [[253,253,253,306,306,...,59,59,59,59,59],[59,59,59,59,59,...,252,252,252,252,245],...,[275,267,267,87,270,...,252,272,272,245,245],[252,254,254,254,272,...,302,253,265,72,265]]\n",
       "product_type_name: [[\"Vest top\",\"Vest top\",\"Vest top\",\"Bra\",\"Bra\",...,\"Swimwear bottom\",\"Swimwear bottom\",\"Swimwear bottom\",\"Swimwear bottom\",\"Swimwear bottom\"],[\"Swimwear bottom\",\"Swimwear bottom\",\"Swimwear bottom\",\"Swimwear bottom\",\"Swimwear bottom\",...,\"Sweater\",\"Sweater\",\"Sweater\",\"Sweater\",\"Cardigan\"],...,[\"Skirt\",\"Jumpsuit/Playsuit\",\"Jumpsuit/Playsuit\",\"Boots\",\"Garment Set\",...,\"Sweater\",\"Trousers\",\"Trousers\",\"Cardigan\",\"Cardigan\"],[\"Sweater\",\"Top\",\"Top\",\"Top\",\"Trousers\",...,\"Socks\",\"Vest top\",\"Dress\",\"Hair clip\",\"Dress\"]]\n",
       "product_group_name: [[\"Garment Upper body\",\"Garment Upper body\",\"Garment Upper body\",\"Underwear\",\"Underwear\",...,\"Swimwear\",\"Swimwear\",\"Swimwear\",\"Swimwear\",\"Swimwear\"],[\"Swimwear\",\"Swimwear\",\"Swimwear\",\"Swimwear\",\"Swimwear\",...,\"Garment Upper body\",\"Garment Upper body\",\"Garment Upper body\",\"Garment Upper body\",\"Garment Upper body\"],...,[\"Garment Lower body\",\"Garment Full body\",\"Garment Full body\",\"Shoes\",\"Garment Full body\",...,\"Garment Upper body\",\"Garment Lower body\",\"Garment Lower body\",\"Garment Upper body\",\"Garment Upper body\"],[\"Garment Upper body\",\"Garment Upper body\",\"Garment Upper body\",\"Garment Upper body\",\"Garment Lower body\",...,\"Socks & Tights\",\"Garment Upper body\",\"Garment Full body\",\"Accessories\",\"Garment Full body\"]]\n",
       "graphical_appearance_no: [[1010016,1010016,1010017,1010016,1010016,...,1010017,1010005,1010001,1010001,1010005],[1010001,1010001,1010017,1010001,1010014,...,1010016,1010016,1010016,1010010,1010016],...,[1010004,1010001,1010001,1010016,1010001,...,1010016,1010023,1010023,1010010,1010010],[1010014,1010016,1010016,1010016,1010016,...,1010014,1010016,1010016,1010016,1010016]]\n",
       "graphical_appearance_name: [[\"Solid\",\"Solid\",\"Stripe\",\"Solid\",\"Solid\",...,\"Stripe\",\"Colour blocking\",\"All over pattern\",\"All over pattern\",\"Colour blocking\"],[\"All over pattern\",\"All over pattern\",\"Stripe\",\"All over pattern\",\"Placement print\",...,\"Solid\",\"Solid\",\"Solid\",\"Melange\",\"Solid\"],...,[\"Check\",\"All over pattern\",\"All over pattern\",\"Solid\",\"All over pattern\",...,\"Solid\",\"Denim\",\"Denim\",\"Melange\",\"Melange\"],[\"Placement print\",\"Solid\",\"Solid\",\"Solid\",\"Solid\",...,\"Placement print\",\"Solid\",\"Solid\",\"Solid\",\"Solid\"]]\n",
       "colour_group_code: [[9,10,11,9,10,...,81,90,73,91,73],[73,82,73,73,40,...,42,93,9,7,6],...,[12,51,9,9,51,...,43,71,72,51,12],[9,9,10,19,12,...,9,9,9,9,11]]\n",
       "colour_group_name: [[\"Black\",\"White\",\"Off White\",\"Black\",\"White\",...,\"Light Turquoise\",\"Other Green\",\"Dark Blue\",\"Light Green\",\"Dark Blue\"],[\"Dark Blue\",\"Turquoise\",\"Dark Blue\",\"Dark Blue\",\"Other Red\",...,\"Red\",\"Dark Green\",\"Black\",\"Grey\",\"Light Grey\"],...,[\"Light Beige\",\"Light Pink\",\"Black\",\"Black\",\"Light Pink\",...,\"Dark Red\",\"Light Blue\",\"Blue\",\"Light Pink\",\"Light Beige\"],[\"Black\",\"Black\",\"White\",\"Greenish Khaki\",\"Light Beige\",...,\"Black\",\"Black\",\"Black\",\"Black\",\"Off White\"]]\n",
       "..."
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "articles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "articles_columns = articles.column_names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "articles_columns = [name if name != 'article_id' else 'nodeId' for name in articles_columns]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "articles = articles.rename_columns(articles_columns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "articles = articles.add_column(1, 'labels', pa.array(['Article'] * len(articles), pa.string()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "articles = articles.select(['nodeId', 'labels', 'product_code', 'product_type_no', 'graphical_appearance_no', 'colour_group_code', 'perceived_colour_value_id', 'perceived_colour_master_id', 'department_no', 'index_group_no', 'section_no', 'garment_group_no'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "nodeId: int64\n",
       "labels: string\n",
       "product_code: int64\n",
       "product_type_no: int64\n",
       "graphical_appearance_no: int64\n",
       "colour_group_code: int64\n",
       "perceived_colour_value_id: int64\n",
       "perceived_colour_master_id: int64\n",
       "department_no: int64\n",
       "index_group_no: int64\n",
       "section_no: int64\n",
       "garment_group_no: int64"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "articles.schema"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "node_descriptor = pa.flight.FlightDescriptor.for_command(f\"\"\"{{\n",
    "    \"name\": \"{graph_name_import}\",\n",
    "    \"entity_type\": \"node\"\n",
    "}}\n",
    "\"\"\")\n",
    "\n",
    "writer, _ = client.do_put(node_descriptor, articles.schema, options = options)\n",
    "\n",
    "with writer:\n",
    "        writer.write(articles)\n",
    "#0.2s\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "article_max = pacompute.max(articles['nodeId']).as_py()\n",
    "\n",
    "customer_columns = customers.column_names\n",
    "customer_columns = [name if name != 'customer_id' else 'nodeId' for name in customer_columns]\n",
    "customers = customers.add_column(1, 'labels', pa.array(['Customer'] * len(customers), pa.string()))\n",
    "customers = customers.add_column(1, 'nodeId', pa.array(range(article_max, article_max + len(customers)), pa.int64()))\n",
    "customers_import = customers.select(['nodeId', 'labels'])\n",
    "\n",
    "writer, _ = client.do_put(node_descriptor, customers_import.schema, options = options)\n",
    "\n",
    "with writer:\n",
    "        writer.write(customers_import)\n",
    "#0.3s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\"name\":\"HEM.NODES2022.NEW\",\"node_count\":1477522}\n"
     ]
    }
   ],
   "source": [
    "action = flight.Action(\"NODE_LOAD_DONE\", f\"\"\"{{\n",
    "    \"name\": \"{graph_name_import}\"\n",
    "}}\"\"\".encode('utf-8'))\n",
    "res = client.do_action(action, options = options)\n",
    "for r in res:\n",
    "    print(r.body.to_pybytes().decode('utf-8'))\n",
    "#0.2s"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load Relationships"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "transactions = pacsv.read_csv('data/transactions_train.csv')\n",
    "#2.1s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "transactions = transactions.join(customers, 'customer_id')\n",
    "#3.2s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "transactions_columns = [name if name != 'nodeId' else 'sourceNodeId' for name in transactions.column_names]\n",
    "transactions_columns = [name if name != 'article_id' else 'targetNodeId' for name in transactions_columns]\n",
    "transactions = transactions.rename_columns(transactions_columns)\n",
    "transactions = transactions.append_column('relationshipType', pa.array(['BUY'] * len(transactions), pa.string()))\n",
    "transactions = transactions.select(['sourceNodeId', 'targetNodeId', 'price', 'relationshipType'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "pyarrow.Table\n",
       "sourceNodeId: int64\n",
       "targetNodeId: int64\n",
       "price: double\n",
       "relationshipType: string\n",
       "----\n",
       "sourceNodeId: [[959994174,959994244,959994244,959994244,959994244,...,960458168,960458168,960458168,960458168,960458168],[960099850,960099850,960099850,960099850,960099850,...,960631979,960631979,960631979,960631985,960631985],...,[960736847,960736847,960736852,960736852,960736852,...,959469265,959469265,959469280,959469280,959469280],[960153231,960153272,960153272,960153459,960153459,...,960199274,960199596,960199596,960200309,960200309]]\n",
       "targetNodeId: [[648173005,671513003,561445005,666080002,661166005,...,587782003,642105002,631770002,548926001,638443002],[623434013,569929001,569929001,569929001,661162003,...,399087019,673396002,673396002,624062003,626445017],...,[903647001,903647001,866714016,783346030,887761002,...,845241002,754267054,855706009,719957017,739363001],[915529005,821395005,915526002,783707116,783707112,...,568601044,900279001,785034009,827957003,827957003]]\n",
       "price: [[0.010152542372881357,0.02035593220338983,0.008457627118644067,0.023474576271186442,0.04698305084745762,...,0.027101694915254236,0.016932203389830508,0.035576271186440675,0.008457627118644067,0.020322033898305086],[0.03049152542372881,0.061016949152542375,0.061000000000000006,0.061000000000000006,0.03049152542372881,...,0.033881355932203386,0.013542372881355931,0.013542372881355931,0.022016949152542376,0.025406779661016947],...,[0.050830508474576264,0.050830508474576264,0.016932203389830508,0.025406779661016947,0.033881355932203386,...,0.01523728813559322,0.01523728813559322,0.01523728813559322,0.019813559322033897,0.0274406779661017],[0.033881355932203386,0.016932203389830508,0.033881355932203386,0.003542372881355932,0.0035593220338983053,...,0.049101694915254235,0.042355932203389825,0.050830508474576264,0.03049152542372881,0.03049152542372881]]\n",
       "relationshipType: [[\"BUY\",\"BUY\",\"BUY\",\"BUY\",\"BUY\",...,\"BUY\",\"BUY\",\"BUY\",\"BUY\",\"BUY\"]]"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "transactions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "rel_descriptor = pa.flight.FlightDescriptor.for_command(f\"\"\"{{\n",
    "    \"name\": \"{graph_name_import}\",\n",
    "    \"entity_type\": \"relationship\"\n",
    "}}\n",
    "\"\"\")\n",
    "\n",
    "writer, _ = client.do_put(rel_descriptor, transactions.schema, options = options)\n",
    "\n",
    "with writer:\n",
    "        writer.write(transactions)\n",
    "#11.5s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\"name\":\"HEM.NODES2022.NEW\",\"relationship_count\":31788324}\n"
     ]
    }
   ],
   "source": [
    "action = flight.Action(\"RELATIONSHIP_LOAD_DONE\", f\"\"\"{{\n",
    "    \"name\": \"{graph_name_import}\"\n",
    "}}\"\"\".encode('utf-8'))\n",
    "res = client.do_action(action, options = options)\n",
    "for r in res:\n",
    "    print(r.body.to_pybytes().decode('utf-8'))\n",
    "#18.7s"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create recommendation pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "from graphdatascience import GraphDataScience\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.2.2\n"
     ]
    }
   ],
   "source": [
    "# Use Neo4j URI and credentials according to your setup\n",
    "# NEO4J_URI could look similar to \"bolt://my-server.neo4j.io:7687\"\n",
    "gds = GraphDataScience(NEO4J_URI, auth=(NEO4J_USER, NEO4J_PASSWORD))\n",
    "\n",
    "# Check the installed GDS version on the server\n",
    "print(gds.version())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "gds.set_database(graph_name_import)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dc2d76e598fa4fe7a745fe9a2296203b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading:   0%|          | 0/100 [00:00<?, ?%/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "G, _ = gds.graph.project(\n",
    "    \"hem\",\n",
    "    [\"Article\", \"Customer\"],\n",
    "    {\n",
    "        \"BUY\": {'orientation': 'UNDIRECTED'}\n",
    "    },\n",
    "    readConcurrency=4\n",
    ")\n",
    "# 4.3s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "pipe, res = gds.beta.pipeline.linkPrediction.create(\"recommendation\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "name                                                    recommendation\n",
       "nodePropertySteps    [{'name': 'gds.fastRP.mutate', 'config': {'ran...\n",
       "featureSteps                                                        []\n",
       "splitConfig          {'negativeSamplingRatio': 1.0, 'testFraction':...\n",
       "autoTuningConfig                                     {'maxTrials': 10}\n",
       "parameterSpace       {'MultilayerPerceptron': [], 'RandomForest': [...\n",
       "Name: 0, dtype: object"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pipe.addNodeProperty(\"fastRP\", embeddingDimension=128, mutateProperty=\"embedding\", randomSeed=42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "name                                                    recommendation\n",
       "nodePropertySteps    [{'name': 'gds.fastRP.mutate', 'config': {'ran...\n",
       "featureSteps         [{'name': 'HADAMARD', 'config': {'nodeProperti...\n",
       "splitConfig          {'negativeSamplingRatio': 1.0, 'testFraction':...\n",
       "autoTuningConfig                                     {'maxTrials': 10}\n",
       "parameterSpace       {'MultilayerPerceptron': [], 'RandomForest': [...\n",
       "Name: 0, dtype: object"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pipe.addFeature(\"hadamard\", nodeProperties=[\"embedding\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "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>name</th>\n",
       "      <th>config</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>HADAMARD</td>\n",
       "      <td>{'nodeProperties': ['embedding']}</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "       name                             config\n",
       "0  HADAMARD  {'nodeProperties': ['embedding']}"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pipe.feature_steps()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "name                                                    recommendation\n",
       "nodePropertySteps    [{'name': 'gds.fastRP.mutate', 'config': {'ran...\n",
       "featureSteps         [{'name': 'HADAMARD', 'config': {'nodeProperti...\n",
       "splitConfig          {'negativeSamplingRatio': 1.0, 'testFraction':...\n",
       "autoTuningConfig                                     {'maxTrials': 10}\n",
       "parameterSpace       {'MultilayerPerceptron': [], 'RandomForest': [...\n",
       "Name: 0, dtype: object"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pipe.configureSplit(trainFraction=0.6, testFraction=0.4, validationFolds=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "name                                                    recommendation\n",
       "nodePropertySteps    [{'name': 'gds.fastRP.mutate', 'config': {'ran...\n",
       "featureSteps         [{'name': 'HADAMARD', 'config': {'nodeProperti...\n",
       "splitConfig          {'negativeSamplingRatio': 1.0, 'testFraction':...\n",
       "autoTuningConfig                                     {'maxTrials': 10}\n",
       "parameterSpace       {'MultilayerPerceptron': [], 'RandomForest': [...\n",
       "Name: 0, dtype: object"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pipe = gds.pipeline.get(\"recommendation\")\n",
    "#pipe.addMLP(hiddenLayerSizes=[128, 64, 2])#,maxEpochs=100, patience=10)\n",
    "pipe.addLogisticRegression()\n",
    "#pipe.addRandomForest(numberOfDecisionTrees=30)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "G = gds.graph.get('hem')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "918172a9ab2444d9b166375c4961c386",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Link Prediction Train Pipeline:   0%|          | 0/100 [00:00<?, ?%/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "buy_recommender, train_result = pipe.train(\n",
    "    G,\n",
    "    sourceNodeLabel='Customer',\n",
    "    targetNodeLabel='Article',\n",
    "    targetRelationshipType=\"BUY\",\n",
    "    modelName=\"buy-recommender\",\n",
    "    randomSeed=42\n",
    ")\n",
    "#4m 15.2s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "modelSelectionStats    {'modelCandidates': [{'metrics': {'AUCPR': {'v...\n",
      "trainMillis                                                       255080\n",
      "modelInfo              {'pipeline': {'nodePropertySteps': [{'name': '...\n",
      "configuration          {'pipeline': 'recommendation', 'randomSeed': 4...\n",
      "Name: 0, dtype: object\n"
     ]
    }
   ],
   "source": [
    "print(train_result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'modelCandidates': [{'metrics': {'AUCPR': {'validation': {'avg': 0.6842737529178428,\n",
       "      'min': 0.6839116918345676,\n",
       "      'max': 0.6847192258178231},\n",
       "     'train': {'avg': 0.6842619190715364,\n",
       "      'min': 0.6841477530776381,\n",
       "      'max': 0.6843503895143053}}},\n",
       "   'parameters': {'maxEpochs': 100,\n",
       "    'minEpochs': 1,\n",
       "    'penalty': 0.0,\n",
       "    'patience': 1,\n",
       "    'methodName': 'LogisticRegression',\n",
       "    'batchSize': 100,\n",
       "    'tolerance': 0.001,\n",
       "    'learningRate': 0.001}}],\n",
       " 'bestParameters': {'maxEpochs': 100,\n",
       "  'minEpochs': 1,\n",
       "  'penalty': 0.0,\n",
       "  'patience': 1,\n",
       "  'methodName': 'LogisticRegression',\n",
       "  'batchSize': 100,\n",
       "  'tolerance': 0.001,\n",
       "  'learningRate': 0.001},\n",
       " 'bestTrial': 1}"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_result.modelSelectionStats"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = gds.model.get('buy-recommender')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### put the prediction on the projected graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5f8c616030e44a1080169b6598317f54",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Link Prediction Predict Pipeline:   0%|          | 0/100 [00:00<?, ?%/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "mutate_results = model.predict_mutate(G, \n",
    "  mutateRelationshipType = 'BUY_APPROX_PREDICTED',\n",
    "  mutateProperty = 'probability',\n",
    "  topK = 1,\n",
    "  sampleRate=0.5,\n",
    "  randomJoins=2,\n",
    "  maxIterations=3,\n",
    "  #necessary for deterministic results\n",
    "  concurrency = 1,\n",
    "  randomSeed =  42\n",
    ")\n",
    "# 36.3s"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### retrieve the results from the projection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "reader = client.do_get(flight.Ticket(\"\"\"{\n",
    "    \"graph_name\": \"hem\",\n",
    "    \"database_name\": \"HEM.NODES2022.NEW\", \n",
    "    \"procedure_name\": \"gds.graph.relationshipProperty.stream\",\n",
    "    \"configuration\": {\n",
    "        \"relationship_types\": \"BUY_APPROX_PREDICTED\",\n",
    "        \"relationship_property\": \"probability\"\n",
    "    }\n",
    "    }\"\"\".encode('utf-8')), options = options)\n",
    "table = reader.read_all()\n",
    "# 0.6s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pyarrow.Table\n",
      "sourceNodeId: int64 not null\n",
      "targetNodeId: int64 not null\n",
      "relationshipType: dictionary<values=string, indices=int32, ordered=0> not null\n",
      "propertyValue: double not null\n",
      "----\n",
      "sourceNodeId: [[0,0,0,0,0,...,717,717,717,717,717],[741013,741014,741015,741016,741017,...,750242,750243,750244,750245,750246],...,[359523,359524,359525,359526,359527,...,368793,368794,368795,368796,368797],[368798,368799,368800,368801,368802,...,371008,371009,371010,371011,371012]]\n",
      "targetNodeId: [[313560,655679,705938,776295,886269,...,959255,971935,1184278,1295123,1296952],[90074,96329,4052,78809,54113,...,45878,18289,32755,48133,1202],...,[48792,93768,2963,48771,50271,...,100831,35047,55334,43129,48070],[59952,23225,65121,87769,37253,...,74524,86830,57917,66568,92610]]\n",
      "relationshipType: [  -- dictionary:\n",
      "[\"BUY_APPROX_PREDICTED\"]  -- indices:\n",
      "[0,0,0,0,0,...,0,0,0,0,0],  -- dictionary:\n",
      "[\"BUY_APPROX_PREDICTED\"]  -- indices:\n",
      "[0,0,0,0,0,...,0,0,0,0,0],...,  -- dictionary:\n",
      "[\"BUY_APPROX_PREDICTED\"]  -- indices:\n",
      "[0,0,0,0,0,...,0,0,0,0,0],  -- dictionary:\n",
      "[\"BUY_APPROX_PREDICTED\"]  -- indices:\n",
      "[0,0,0,0,0,...,0,0,0,0,0]]\n",
      "propertyValue: [[0.6352690947296723,0.6351518926148801,0.6352490325755397,0.634910960965452,0.634888259927615,...,0.6352880802395594,0.634763895429623,0.6351643436426087,0.6350408985951449,0.6350390811976369],[0.6351725408533708,0.6352674777617453,0.6348373512846565,0.6354126557165619,0.6343644040458772,...,0.6354525007616265,0.6348753571126098,0.6350176506712746,0.6351988499626999,0.6346693682184593],...,[0.6348917903889163,0.6352978461609006,0.6346607316867124,0.6349028378095546,0.6348935841838083,...,0.6349281117598932,0.6344598451213366,0.6352356894874683,0.6351336976128169,0.6349727307738319],[0.6347560929363701,0.634770666143003,0.6350033379241284,0.6352388092880537,0.6348900834199472,...,0.6352383532555591,0.6350797968456517,0.6349386462085013,0.6354243830044245,0.634349155807558]]\n"
     ]
    }
   ],
   "source": [
    "print(table)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## DGL LINK PREDICTION"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bddd99b12aa54e0797a0f5d8ace9e7b7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading:   0%|          | 0/100 [00:00<?, ?%/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "G, _ = gds.graph.project(\n",
    "    \"hem_dgl\",\n",
    "    [\"Article\", \"Customer\"],\n",
    "    {\n",
    "        \"BUY\": {\n",
    "        'orientation': 'UNDIRECTED',\n",
    "        \"properties\": [\"price\"]\n",
    "        }\n",
    "    },\n",
    "    readConcurrency=4\n",
    ")\n",
    "#9.5s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f99b6d0cd7d94da88dc6a462d57d0ad6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FastRP:   0%|          | 0/100 [00:00<?, ?%/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "result = gds.fastRP.mutate(\n",
    "    G,\n",
    "    mutateProperty=\"embedding\",\n",
    "    randomSeed=42,\n",
    "    embeddingDimension=128,\n",
    "    iterationWeights=[0.8, 1, 1, 1],\n",
    ")\n",
    "#15.2s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "reader = client.do_get(flight.Ticket(\"\"\"{\n",
    "    \"graph_name\": \"hem_dgl\",\n",
    "    \"database_name\": \"HEM.NODES2022.NEW\", \n",
    "    \"procedure_name\": \"gds.graph.nodeProperty.stream\",\n",
    "    \"configuration\": {\n",
    "        \"node_labels\":  [\"Article\", \"Customer\"],\n",
    "        \"node_property\": \"embedding\"\n",
    "    }\n",
    "    }\"\"\".encode('utf-8')), options = options)\n",
    "table_nodes = reader.read_all()\n",
    "#1.6s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "sorted_table_nodes = table_nodes.append_column('dglID', pa.array(range(0,len(table_nodes)), pa.int64()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "sorted_table_nodes = sorted_table_nodes.sort_by([(\"dglID\", \"ascending\")])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pyarrow.Table\n",
      "nodeId: int64 not null\n",
      "propertyValue: list<propertyValue.inner: float not null> not null\n",
      "  child 0, propertyValue.inner: float not null\n",
      "dglID: int64\n",
      "----\n",
      "nodeId: [[0,1,2,3,4,...,1478530,1478531,1478532,1478533,1478534]]\n",
      "propertyValue: [[[-0.100073874,-0.26626396,-0.012616102,-0.23033217,0.35614634,...,-0.40032893,0.21652773,-0.028770644,-0.121459305,0.32951498],[-0.41985935,-0.023758918,-0.013264384,-0.2588179,0.11248251,...,-0.1512475,0.17956246,-0.017649047,-0.10805801,0.25117278],...,[-0.07483108,-0.07485901,0.15815175,-0.0558758,-0.12628093,...,-0.13341352,0.09667079,-0.15405637,0.046845216,0.3162402],[0.028021835,0.20552,-0.16797295,0.09843139,0.2908689,...,-0.15081756,0.13862059,-0.39557898,0.035192423,0.3521633]]]\n",
      "dglID: [[0,1,2,3,4,...,1477517,1477518,1477519,1477520,1477521]]\n"
     ]
    }
   ],
   "source": [
    "print(sorted_table_nodes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "reader = client.do_get(flight.Ticket(\"\"\"{\n",
    "    \"graph_name\": \"hem_dgl\",\n",
    "    \"database_name\": \"HEM.NODES2022.NEW\", \n",
    "    \"procedure_name\": \"gds.graph.relationshipProperty.stream\",\n",
    "    \"configuration\": {\n",
    "        \"relationship_types\": \"BUY\",\n",
    "        \"relationship_property\": \"price\"\n",
    "    }\n",
    "    }\"\"\".encode('utf-8')), options = options)\n",
    "table_rels = reader.read_all()\n",
    "#6.9s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "import dgl\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pyarrow.Table\n",
      "sourceNodeId: int64 not null\n",
      "targetNodeId: int64 not null\n",
      "relationshipType: dictionary<values=string, indices=int32, ordered=0> not null\n",
      "propertyValue: double not null\n",
      "----\n",
      "sourceNodeId: [[0,0,0,0,0,...,0,0,0,0,0],[741013,741013,741013,741013,741013,...,741459,741459,741459,741459,741459],...,[370544,370545,370545,370545,370545,...,370933,370933,370933,370933,370933],[370933,370933,370933,370933,370933,...,371011,371012,371012,371012,371012]]\n",
      "targetNodeId: [[103055,103055,103055,107138,107138,...,1362453,1362667,1362731,1362731,1363169],[74,117,132,201,285,...,5821,7527,10359,14157,14157],...,[91495,2677,7548,30910,40190,...,71034,71034,73338,73679,75735],[75735,75735,75735,75891,76179,...,94579,32688,34554,41843,44065]]\n",
      "relationshipType: [  -- dictionary:\n",
      "[\"BUY\"]  -- indices:\n",
      "[0,0,0,0,0,...,0,0,0,0,0],  -- dictionary:\n",
      "[\"BUY\"]  -- indices:\n",
      "[0,0,0,0,0,...,0,0,0,0,0],...,  -- dictionary:\n",
      "[\"BUY\"]  -- indices:\n",
      "[0,0,0,0,0,...,0,0,0,0,0],  -- dictionary:\n",
      "[\"BUY\"]  -- indices:\n",
      "[0,0,0,0,0,...,0,0,0,0,0]]\n",
      "propertyValue: [[0.008457627118644067,0.008457627118644067,0.008457627118644067,0.00676271186440678,0.00676271186440678,...,0.00676271186440678,0.008457627118644067,0.008457627118644067,0.008457627118644067,0.007186440677966101],[0.013542372881355931,0.016932203389830508,0.016932203389830508,0.013542372881355931,0.0033728813559322033,...,0.022016949152542376,0.016932203389830508,0.018627118644067797,0.016932203389830508,0.016932203389830508],...,[0.01523728813559322,0.0033728813559322033,0.005067796610169492,0.018627118644067797,0.027101694915254236,...,0.01523728813559322,0.01523728813559322,0.03189830508474576,0.06777966101694916,0.042355932203389825],[0.042355932203389825,0.033881355932203386,0.033881355932203386,0.030813559322033897,0.05816949152542372,...,0.05930508474576271,0.01523728813559322,0.01523728813559322,0.007610169491525424,0.01523728813559322]]\n"
     ]
    }
   ],
   "source": [
    "print(table_rels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "tables_id = sorted_table_nodes.select(['nodeId', 'dglID'])\n",
    "new_rels = table_rels.join(tables_id, keys='sourceNodeId', right_keys='nodeId')\n",
    "new_rels = new_rels.join(tables_id, keys='targetNodeId', right_keys='nodeId', right_suffix=\"_target\")\n",
    "\n",
    "#7.1s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pyarrow.Table\n",
      "sourceNodeId: int64\n",
      "targetNodeId: int64\n",
      "relationshipType: dictionary<values=string, indices=int32, ordered=0>\n",
      "propertyValue: double\n",
      "dglID: int64\n",
      "dglID_target: int64\n",
      "----\n",
      "sourceNodeId: [[67,67,67,67,67,...,73,73,73,73,73],[374091,374091,374091,374091,374091,...,374490,374490,374490,374490,374490],...,[370106,370106,370106,370106,370106,...,370544,370544,370544,370544,370544],[362448,362448,362448,362448,362448,...,362931,362931,362931,362931,362931]]\n",
      "targetNodeId: [[985765,985765,985765,986246,986284,...,160397,160580,160580,160580,160613],[46678,61589,65571,68990,72052,...,89485,89489,90149,90149,90912],...,[40238,40479,41107,41107,41107,...,53286,56440,56440,57073,91481],[71919,72381,72382,72383,72385,...,505,505,506,506,506]]\n",
      "relationshipType: [  -- dictionary:\n",
      "[\"BUY\"]  -- indices:\n",
      "[0,0,0,0,0,...,0,0,0,0,0],  -- dictionary:\n",
      "[\"BUY\"]  -- indices:\n",
      "[0,0,0,0,0,...,0,0,0,0,0],...,  -- dictionary:\n",
      "[\"BUY\"]  -- indices:\n",
      "[0,0,0,0,0,...,0,0,0,0,0],  -- dictionary:\n",
      "[\"BUY\"]  -- indices:\n",
      "[0,0,0,0,0,...,0,0,0,0,0]]\n",
      "propertyValue: [[0.00564406779661017,0.00564406779661017,0.005627118644067797,0.008457627118644067,0.008457627118644067,...,0.016932203389830508,0.01128813559322034,0.01128813559322034,0.01128813559322034,0.016932203389830508],[0.033881355932203386,0.050830508474576264,0.025406779661016947,0.016932203389830508,0.033881355932203386,...,0.025406779661016947,0.022016949152542376,0.01523728813559322,0.016932203389830508,0.02111864406779661],...,[0.042355932203389825,0.010152542372881357,0.033881355932203386,0.033881355932203386,0.033881355932203386,...,0.022864406779661017,0.01523728813559322,0.01523728813559322,0.01523728813559322,0.022864406779661017],[0.0321864406779661,0.00676271186440678,0.00676271186440678,0.00676271186440678,0.00676271186440678,...,0.025406779661016947,0.02489830508474576,0.025406779661016947,0.02364406779661017,0.025406779661016947]]\n",
      "dglID: [[67,67,67,67,67,...,73,73,73,73,73],[373078,373078,373078,373078,373078,...,373477,373477,373477,373477,373477],...,[369093,369093,369093,369093,369093,...,369531,369531,369531,369531,369531],[361435,361435,361435,361435,361435,...,361918,361918,361918,361918,361918]]\n",
      "dglID_target: [[984752,984752,984752,985233,985271,...,159384,159567,159567,159567,159600],[46678,61589,65571,68990,72052,...,89485,89489,90149,90149,90912],...,[40238,40479,41107,41107,41107,...,53286,56440,56440,57073,91481],[71919,72381,72382,72383,72385,...,505,505,506,506,506]]\n"
     ]
    }
   ],
   "source": [
    "print(new_rels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "dizionario = new_rels.to_pydict()\n",
    "sources = torch.IntTensor(dizionario[\"dglID\"])\n",
    "targets = torch.IntTensor(dizionario[\"dglID_target\"])\n",
    "\n",
    "g = dgl.graph((sources, targets))\n",
    "#3m 21.1s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Graph(num_nodes=1477522, num_edges=63576648,\n",
      "      ndata_schemes={}\n",
      "      edata_schemes={})\n"
     ]
    }
   ],
   "source": [
    "print(g)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "dizionario = sorted_table_nodes.to_pydict()\n",
    "g.ndata['embedding'] = torch.FloatTensor(dizionario['propertyValue'])\n",
    "#1m 30.3s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "import itertools\n",
    "import numpy as np\n",
    "import scipy.sparse as sp\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "metadata": {},
   "outputs": [],
   "source": [
    "sampler = dgl.dataloading.MultiLayerFullNeighborSampler(2)\n",
    "sampler = dgl.dataloading.as_edge_prediction_sampler(\n",
    "    sampler, negative_sampler=dgl.dataloading.negative_sampler.Uniform(5))\n",
    "indice = random.sample(range(len(g.ndata['embedding'])), 1000)\n",
    "indice = torch.IntTensor(indice)\n",
    "dataloader = dgl.dataloading.DataLoader(\n",
    "    g, indice, sampler,\n",
    "    batch_size=64,\n",
    "    shuffle=True,\n",
    "    drop_last=False,\n",
    "    num_workers=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "metadata": {},
   "outputs": [],
   "source": [
    "class StochasticTwoLayerGCN(nn.Module):\n",
    "    def __init__(self, in_features, hidden_features, out_features):\n",
    "        super().__init__()\n",
    "        self.conv1 = dgl.nn.GraphConv(in_features, hidden_features, allow_zero_in_degree=True)\n",
    "        self.conv2 = dgl.nn.GraphConv(hidden_features, out_features, allow_zero_in_degree=True)\n",
    "\n",
    "    def forward(self, blocks, x):\n",
    "        x = F.relu(self.conv1(blocks[0], x))\n",
    "        x = F.relu(self.conv2(blocks[1], x))\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ScorePredictor(nn.Module):\n",
    "    def forward(self, edge_subgraph, x):\n",
    "        with edge_subgraph.local_scope():\n",
    "            edge_subgraph.ndata['x'] = x\n",
    "            edge_subgraph.apply_edges(dgl.function.u_dot_v('x', 'x', 'score'))\n",
    "            return edge_subgraph.edata['score']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Model(nn.Module):\n",
    "    def __init__(self, in_features, hidden_features, out_features):\n",
    "        super().__init__()\n",
    "        self.gcn = StochasticTwoLayerGCN(\n",
    "            in_features, hidden_features, out_features)\n",
    "        self.predictor = ScorePredictor()\n",
    "\n",
    "    def forward(self, positive_graph, negative_graph, blocks, x):\n",
    "        x = self.gcn(blocks, x)\n",
    "        pos_score = self.predictor(positive_graph, x)\n",
    "        neg_score = self.predictor(negative_graph, x)\n",
    "        return pos_score, neg_score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 DL workers are assigned to cpus [0], main process will use cpus [1, 2, 3, 4, 5]\n",
      "loss for epoch 1 is: 280.06372582912445\n",
      "loss for epoch 2 is: 7.661756634712219\n",
      "loss for epoch 3 is: 6.506469696760178\n",
      "loss for epoch 4 is: 5.929817616939545\n",
      "loss for epoch 5 is: 5.21606719493866\n"
     ]
    }
   ],
   "source": [
    "def compute_loss(pos_score, neg_score):\n",
    "    # an example hinge loss\n",
    "    n = pos_score.shape[0]\n",
    "    return (neg_score.view(n, -1) - pos_score.view(n, -1) + 1).clamp(min=0).mean()\n",
    "\n",
    "model = Model(128, 128, 64)\n",
    "model = model\n",
    "opt = torch.optim.Adam(model.parameters())\n",
    "#aggiungiamo self loop per non far rompere ogni cosa nell'universo\n",
    "g = dgl.add_self_loop(g)\n",
    "loss_n = 0\n",
    "epochs = 5\n",
    "with dataloader.enable_cpu_affinity():\n",
    "    for epoch in range(epochs):\n",
    "        for input_nodes, positive_graph, negative_graph, blocks in dataloader:\n",
    "            positive_graph = positive_graph\n",
    "            negative_graph = negative_graph\n",
    "            input_features = blocks[0].srcdata['embedding']\n",
    "            pos_score, neg_score = model(positive_graph, negative_graph, blocks, input_features)\n",
    "            loss = compute_loss(pos_score, neg_score)\n",
    "            loss_n += loss.mean().item()\n",
    "            opt.zero_grad()\n",
    "            loss.backward()\n",
    "            opt.step()\n",
    "        print(f\"loss for epoch {epoch + 1} is: {loss_n}\")\n",
    "        loss_n = 0"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.7.10 ('env': 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.7.10"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "d6cef9b5219dcb2758c15f344a46c20868fa24743dc92a8d86016c6d5422fd58"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}