Training-SQ1.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "89387b71",
   "metadata": {},
   "source": [
    "# Task to add ibm db2 queries in QueryFormer\n",
    "\n",
    "# I created an experiment file train_plan_part20.csv to format db2 query in JSON format\n",
    "# We need to extract the following information from db2 \n",
    "# 1) \"Node Type”\n",
    "# 2) “Alias” \n",
    "# 3) \"Parent Relationship\"\n",
    "# 4) “Hash Cond”\n",
    "# 5) “Join Filter”\n",
    "# 6) “Index Cond”\n",
    "# 7) “Recheck Cond”\n",
    "# 8) “Relation Name”\n",
    "# 9) “Plan”\n",
    "# 10) “Plans”\n",
    "# 11) “Actual rows”\n",
    "# 12) “Execution time”"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "id": "d9c5213b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import os\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import time\n",
    "import pandas as pd\n",
    "from scipy.stats import pearsonr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "id": "037bec4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from model.util import Normalizer\n",
    "from model.database_util import get_hist_file, get_job_table_sample, collator\n",
    "from model.model import QueryFormer\n",
    "from model.database_util import Encoding,TreeNode\n",
    "from model.dataset import PlanTreeDataset\n",
    "from model.trainer import eval_workload, train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "id": "822fdcaf",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_path = './data/imdb/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 102,
   "id": "fbcd4773",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Args:\n",
    "    bs = 1024\n",
    "    lr = 0.001\n",
    "    epochs = 200\n",
    "    clip_size = 50\n",
    "    embed_size = 64\n",
    "    pred_hid = 128\n",
    "    ffn_dim = 128\n",
    "    head_size = 12\n",
    "    n_layers = 8\n",
    "    dropout = 0.1\n",
    "    sch_decay = 0.6\n",
    "    device = 'cpu'\n",
    "    newpath = './results/full/cost/'\n",
    "    to_predict = 'cost'\n",
    "args = Args()\n",
    "\n",
    "import os\n",
    "if not os.path.exists(args.newpath):\n",
    "    os.makedirs(args.newpath)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "id": "aace65f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "hist_file = get_hist_file(data_path + 'histogram_string.csv')\n",
    "cost_norm = Normalizer(-3.61192, 12.290855)\n",
    "card_norm = Normalizer(1,100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "id": "4e5f421a",
   "metadata": {},
   "outputs": [],
   "source": [
    "encoding_ckpt = torch.load('checkpoints/encoding.pt')\n",
    "encoding = encoding_ckpt['encoding']\n",
    "checkpoint = torch.load('checkpoints/cost_model.pt', map_location='cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 105,
   "id": "a2eb4fef",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "col-min-max: {'t.id': [1.0, 2528312.0], 't.kind_id': [1.0, 7.0], 't.production_year': [1880.0, 2019.0], 'mc.id': [1.0, 2609129.0], 'mc.company_id': [1.0, 234997.0], 'mc.movie_id': [2.0, 2525745.0], 'mc.company_type_id': [1.0, 2.0], 'ci.id': [1.0, 36244344.0], 'ci.movie_id': [1.0, 2525975.0], 'ci.person_id': [1.0, 4061926.0], 'ci.role_id': [1.0, 11.0], 'mi.id': [1.0, 14835720.0], 'mi.movie_id': [1.0, 2526430.0], 'mi.info_type_id': [1.0, 110.0], 'mi_idx.id': [1.0, 1380035.0], 'mi_idx.movie_id': [2.0, 2525793.0], 'mi_idx.info_type_id': [99.0, 113.0], 'mk.id': [1.0, 4523930.0], 'mk.movie_id': [2.0, 2525971.0], 'mk.keyword_id': [1.0, 134170.0]}\n",
      "---------\n",
      "col-index: {'t.id': 0, 't.kind_id': 1, 't.production_year': 2, 'mc.id': 3, 'mc.company_id': 4, 'mc.movie_id': 5, 'mc.company_type_id': 6, 'ci.id': 7, 'ci.movie_id': 8, 'ci.person_id': 9, 'ci.role_id': 10, 'mi.id': 11, 'mi.movie_id': 12, 'mi.info_type_id': 13, 'mi_idx.id': 14, 'mi_idx.movie_id': 15, 'mi_idx.info_type_id': 16, 'mk.id': 17, 'mk.movie_id': 18, 'mk.keyword_id': 19, 'NA': 20}\n",
      "---------\n",
      "op-index: {'>': 0, '=': 1, '<': 2, 'NA': 3}\n",
      "---------\n",
      "type-index: {'Gather': 0, 'Hash Join': 1, 'Seq Scan': 2, 'Hash': 3, 'Bitmap Heap Scan': 4, 'Bitmap Index Scan': 5, 'Nested Loop': 6, 'Index Scan': 7, 'Merge Join': 8, 'Gather Merge': 9, 'Materialize': 10, 'BitmapAnd': 11, 'Sort': 12}\n",
      "---------\n",
      "idx2-index: {0: 'Gather', 1: 'Hash Join', 2: 'Seq Scan', 3: 'Hash', 4: 'Bitmap Heap Scan', 5: 'Bitmap Index Scan', 6: 'Nested Loop', 7: 'Index Scan', 8: 'Merge Join', 9: 'Gather Merge', 10: 'Materialize', 11: 'BitmapAnd', 12: 'Sort'}\n",
      "---------\n",
      "join2idx-index: {None: 0, 'mi_idx.movie_id = t.id': 1, 'mc.movie_id = t.id': 2, 'mi.movie_id = t.id': 3, 'ci.movie_id = t.id': 4, 'mk.movie_id = t.id': 5, 'ci.movie_id = mk.movie_id': 6, 'mi.movie_id = mk.movie_id': 7, 'mi_idx.movie_id = mk.movie_id': 8, 'mc.movie_id = mk.movie_id': 9, 'ci.movie_id = mi_idx.movie_id': 10, 'ci.movie_id = mc.movie_id': 11, 'ci.movie_id = mi.movie_id': 12, 'mi.movie_id = mi_idx.movie_id': 13, 'mc.movie_id = mi_idx.movie_id': 14, 'mc.movie_id = mi.movie_id': 15}\n",
      "---------\n",
      "idx2join-index: {0: None, 1: 'mi_idx.movie_id = t.id', 2: 'mc.movie_id = t.id', 3: 'mi.movie_id = t.id', 4: 'ci.movie_id = t.id', 5: 'mk.movie_id = t.id', 6: 'ci.movie_id = mk.movie_id', 7: 'mi.movie_id = mk.movie_id', 8: 'mi_idx.movie_id = mk.movie_id', 9: 'mc.movie_id = mk.movie_id', 10: 'ci.movie_id = mi_idx.movie_id', 11: 'ci.movie_id = mc.movie_id', 12: 'ci.movie_id = mi.movie_id', 13: 'mi.movie_id = mi_idx.movie_id', 14: 'mc.movie_id = mi_idx.movie_id', 15: 'mc.movie_id = mi.movie_id'}\n",
      "---------\n",
      "table2index-index: {'NA': 0, 'title': 1, 'movie_info_idx': 2, 'movie_info': 3, 'movie_companies': 4, 'movie_keyword': 5, 'cast_info': 6}\n",
      "---------\n",
      "idx2table-index: {0: 'NA', 1: 'title', 2: 'movie_info_idx', 3: 'movie_info', 4: 'movie_companies', 5: 'movie_keyword', 6: 'cast_info'}\n"
     ]
    }
   ],
   "source": [
    "for key, value in encoding_ckpt.items():\n",
    "    print(f\"col-min-max: {value.column_min_max_vals}\")\n",
    "    print(\"---------\")\n",
    "    print(f\"col-index: {value.col2idx}\")\n",
    "    print(\"---------\")\n",
    "    print(f\"op-index: {value.op2idx}\")\n",
    "    print(\"---------\")\n",
    "    print(f\"type-index: {value.type2idx}\")\n",
    "    print(\"---------\")\n",
    "    print(f\"idx2-index: {value.idx2type}\")\n",
    "    print(\"---------\")\n",
    "    print(f\"join2idx-index: {value.join2idx}\")\n",
    "    print(\"---------\")\n",
    "    print(f\"idx2join-index: {value.idx2join}\")\n",
    "    print(\"---------\")\n",
    "    print(f\"table2index-index: {value.table2idx}\")\n",
    "    print(\"---------\")\n",
    "    print(f\"idx2table-index: {value.idx2table}\")\n",
    " \n",
    "# Additional operators added for ibm db2   \n",
    "value.type2idx[\"TBSCAN\"] = 13\n",
    "value.type2idx[\"HSJOIN\"] = 14\n",
    "\n",
    "value.idx2type[\"TBSCAN\"] = 13\n",
    "value.idx2type[\"HSJOIN\"] = 14\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "id": "71759b41",
   "metadata": {},
   "outputs": [],
   "source": [
    "from model.util import seed_everything\n",
    "seed_everything()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc547102",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 106,
   "id": "f9592f60",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = QueryFormer(emb_size = args.embed_size ,ffn_dim = args.ffn_dim, head_size = args.head_size, \\\n",
    "                 dropout = args.dropout, n_layers = args.n_layers, \\\n",
    "                 use_sample = True, use_hist = True, \\\n",
    "                 pred_hid = args.pred_hid\n",
    "                )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "6509fe71",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: torchviz in /Users/yonisabokar/opt/anaconda3/lib/python3.8/site-packages (0.0.2)\n",
      "Requirement already satisfied: graphviz in /Users/yonisabokar/opt/anaconda3/lib/python3.8/site-packages (from torchviz) (0.20.3)\n",
      "Requirement already satisfied: torch in /Users/yonisabokar/opt/anaconda3/lib/python3.8/site-packages (from torchviz) (1.13.1)\n",
      "Requirement already satisfied: typing-extensions in /Users/yonisabokar/opt/anaconda3/lib/python3.8/site-packages (from torch->torchviz) (4.7.1)\n",
      "\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.0.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.0\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n"
     ]
    },
    {
     "ename": "AttributeError",
     "evalue": "'Tensor' object has no attribute 'attn_bias'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-13-02a2695d000a>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     14\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     15\u001b[0m \u001b[0;31m# Pass the input through the model to get the output\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 16\u001b[0;31m \u001b[0moutput\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdummy_input\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     17\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     18\u001b[0m \u001b[0;31m# Generate a visualization of the model\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/opt/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1192\u001b[0m         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks\n\u001b[1;32m   1193\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1194\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1195\u001b[0m         \u001b[0;31m# Do not call functions when jit is used\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1196\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/QueryFormer/model/model.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, batched_data)\u001b[0m\n\u001b[1;32m    193\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    194\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatched_data\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 195\u001b[0;31m         \u001b[0mattn_bias\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrel_pos\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mbatched_data\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mattn_bias\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatched_data\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrel_pos\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatched_data\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    196\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    197\u001b[0m         \u001b[0mheights\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mbatched_data\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mheights\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mAttributeError\u001b[0m: 'Tensor' object has no attribute 'attn_bias'"
     ]
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 107,
   "id": "85e27584",
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = model.to(args.device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 108,
   "id": "6c8a7d25",
   "metadata": {},
   "outputs": [],
   "source": [
    "to_predict = 'cost'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb05e3ac",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 114,
   "id": "4286ce8d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded queries with len  100000\n",
      "Loaded bitmaps\n"
     ]
    }
   ],
   "source": [
    "# imdb_path = './data/imdb/'\n",
    "# full_train_df = pd.DataFrame()\n",
    "# for i in range(2):\n",
    "#     file = imdb_path + 'plan_and_cost/train_plan_part{}.csv'.format(i)\n",
    "#     df = pd.read_csv(file)\n",
    "#     full_train_df = full_train_df.append(df)\n",
    "\n",
    "imdb_path = './data/imdb/'\n",
    "full_train_df = pd.DataFrame()\n",
    "file = imdb_path + 'plan_and_cost/train_plan_part20.csv'\n",
    "df = pd.read_csv(file)\n",
    "full_train_df = full_train_df.append(df)\n",
    "\n",
    "val_df = pd.DataFrame()\n",
    "for i in range(2,4):\n",
    "    file = imdb_path + 'plan_and_cost/train_plan_part{}.csv'.format(i)\n",
    "    df = pd.read_csv(file)\n",
    "    val_df = val_df.append(df)\n",
    "table_sample = get_job_table_sample(imdb_path+'train')\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 115,
   "id": "98d193b2",
   "metadata": {},
   "outputs": [
    {
     "ename": "KeyError",
     "evalue": "'Execution Time'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyError\u001b[0m                                  Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-115-c6d7f390c01f>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 5\u001b[0;31m \u001b[0mtrain_ds\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mPlanTreeDataset\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfull_train_df\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mencoding\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhist_file\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcard_norm\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcost_norm\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mto_predict\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtable_sample\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      6\u001b[0m \u001b[0mval_ds\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mPlanTreeDataset\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mval_df\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mencoding\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhist_file\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcard_norm\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcost_norm\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mto_predict\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtable_sample\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/QueryFormer/model/dataset.py\u001b[0m in \u001b[0;36m__init__\u001b[0;34m(self, json_df, train, encoding, hist_file, card_norm, cost_norm, to_predict, table_sample)\u001b[0m\n\u001b[1;32m     21\u001b[0m         \u001b[0mnodes\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mjson\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mloads\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mplan\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'Plan'\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mplan\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mjson_df\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'json'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     22\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcards\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mnode\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'Actual Rows'\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mnode\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mnodes\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 23\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcosts\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mjson\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mloads\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mplan\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'Execution Time'\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mplan\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mjson_df\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'json'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     24\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     25\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcard_labels\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfrom_numpy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcard_norm\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnormalize_labels\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcards\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/QueryFormer/model/dataset.py\u001b[0m in \u001b[0;36m<listcomp>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m     21\u001b[0m         \u001b[0mnodes\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mjson\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mloads\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mplan\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'Plan'\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mplan\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mjson_df\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'json'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     22\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcards\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mnode\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'Actual Rows'\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mnode\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mnodes\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 23\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcosts\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mjson\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mloads\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mplan\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'Execution Time'\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mplan\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mjson_df\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'json'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     24\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     25\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcard_labels\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfrom_numpy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcard_norm\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnormalize_labels\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcards\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyError\u001b[0m: 'Execution Time'"
     ]
    }
   ],
   "source": [
    "import json\n",
    "\n",
    "\n",
    "\n",
    "train_ds = PlanTreeDataset(full_train_df, None, encoding, hist_file, card_norm, cost_norm, to_predict, table_sample)\n",
    "val_ds = PlanTreeDataset(val_df, None, encoding, hist_file, card_norm, cost_norm, to_predict, table_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "id": "812b9db7",
   "metadata": {},
   "outputs": [
    {
     "ename": "JSONDecodeError",
     "evalue": "Extra data: line 1 column 4 (char 3)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mJSONDecodeError\u001b[0m                           Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-98-d6d24e6fc6bc>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      2\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0midxs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mlist\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfull_train_df\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'id'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 4\u001b[0;31m \u001b[0mnodes\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mjson\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mloads\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mplan\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'Plan'\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mplan\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mfull_train_df\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'json'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      5\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      6\u001b[0m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnodes\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-98-d6d24e6fc6bc>\u001b[0m in \u001b[0;36m<listcomp>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m      2\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0midxs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mlist\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfull_train_df\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'id'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 4\u001b[0;31m \u001b[0mnodes\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mjson\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mloads\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mplan\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'Plan'\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mplan\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mfull_train_df\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'json'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      5\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      6\u001b[0m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnodes\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/opt/anaconda3/lib/python3.8/json/__init__.py\u001b[0m in \u001b[0;36mloads\u001b[0;34m(s, cls, object_hook, parse_float, parse_int, parse_constant, object_pairs_hook, **kw)\u001b[0m\n\u001b[1;32m    355\u001b[0m             \u001b[0mparse_int\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0mparse_float\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m \u001b[0;32mand\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    356\u001b[0m             parse_constant is None and object_pairs_hook is None and not kw):\n\u001b[0;32m--> 357\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0m_default_decoder\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdecode\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0ms\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    358\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mcls\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    359\u001b[0m         \u001b[0mcls\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mJSONDecoder\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/opt/anaconda3/lib/python3.8/json/decoder.py\u001b[0m in \u001b[0;36mdecode\u001b[0;34m(self, s, _w)\u001b[0m\n\u001b[1;32m    338\u001b[0m         \u001b[0mend\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_w\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0ms\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mend\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    339\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mend\u001b[0m \u001b[0;34m!=\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0ms\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 340\u001b[0;31m             \u001b[0;32mraise\u001b[0m \u001b[0mJSONDecodeError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"Extra data\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0ms\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mend\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    341\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mobj\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    342\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mJSONDecodeError\u001b[0m: Extra data: line 1 column 4 (char 3)"
     ]
    }
   ],
   "source": [
    "import json\n",
    "\n",
    "idxs = list(full_train_df['id'])\n",
    "nodes = [json.loads(plan)['Plan'] for plan in full_train_df['json']]\n",
    "\n",
    "print(len(nodes))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "2199115a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The index: 1 is query: {'Node Type': 'Seq Scan', 'Parallel Aware': False, 'Relation Name': 'title', 'Alias': 't', 'Startup Cost': 0.0, 'Total Cost': 67602.3, 'Plan Rows': 1116092, 'Plan Width': 94, 'Actual Startup Time': 0.035, 'Actual Total Time': 322.837, 'Actual Rows': 1107925, 'Actual Loops': 1, 'Filter': '(production_year > 2004)', 'Rows Removed by Filter': 1420387}\n"
     ]
    }
   ],
   "source": [
    "query = nodes[1]\n",
    "index  = idxs[1]\n",
    "print(f\"The index: {index} is query: {query}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "c1461995",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'list'>\n"
     ]
    }
   ],
   "source": [
    "# Converting json query tree structure plan into List of TreeNode\n",
    "tree_nodes = [train_ds.traversePlan(node, i, train_ds.encoding) for i,node in zip(idxs, nodes)]\n",
    "\n",
    "\n",
    "print(type(tree_nodes))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "7f475913",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "10000"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(tree_nodes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "686c8dc2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Gather with [], None, 1 children\n",
      "The queryId is 3\n",
      "The nodeType is Gather\n",
      "The typeId is 0\n",
      "The filters is []\n",
      "The joinId is None\n",
      "The card is None\n"
     ]
    }
   ],
   "source": [
    " # Tree node representation for 3rd query in train_plan_{}.csv\n",
    "root_node = tree_nodes[3]\n",
    "print(root_node)\n",
    "\n",
    "print(f\"The queryId is {root_node.query_id}\")\n",
    "print(f\"The nodeType is {root_node.nodeType}\")\n",
    "print(f\"The typeId is {root_node.typeId}\")\n",
    "print(f\"The filters is {root_node.filter}\")\n",
    "print(f\"The joinId is {root_node.join_str}\")\n",
    "print(f\"The card is {root_node.card}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "054327f3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Gather with [] and None, 1 childs\n",
      "--Hash Join with [] and mc.movie_id = t.id, 2 childs\n",
      "----Seq Scan with [] and None, 0 childs\n",
      "----Hash with [] and None, 1 childs\n",
      "------Bitmap Heap Scan with ['(company_id < 27)'] and None, 1 childs\n",
      "--------Bitmap Index Scan with ['(company_id < 27)'] and None, 0 childs\n"
     ]
    }
   ],
   "source": [
    "TreeNode.print_nested(root_node)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "d78c3944",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 0.,  0., 20., ...,  0.,  0.,  0.])"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query = root_node\n",
    "\n",
    "root_node.feature\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "e81d50c5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'features': tensor([[ 0.,  0., 20.,  ...,  0.,  0.,  0.],\n",
      "        [ 1.,  2., 20.,  ...,  0.,  0.,  0.],\n",
      "        [ 2.,  0., 20.,  ...,  1.,  1.,  1.],\n",
      "        [ 3.,  0., 20.,  ...,  0.,  0.,  0.],\n",
      "        [ 4.,  0.,  4.,  ...,  0.,  0.,  0.],\n",
      "        [ 5.,  0.,  4.,  ...,  0.,  0.,  0.]]), 'heights': tensor([4, 3, 0, 2, 1, 0]), 'adjacency_list': tensor([[0, 1],\n",
      "        [1, 2],\n",
      "        [1, 3],\n",
      "        [3, 4],\n",
      "        [4, 5]])}\n"
     ]
    }
   ],
   "source": [
    "query_node2dict = train_ds.node2dict(query) \n",
    "print(query_node2dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "7b1a1dda",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dict_keys(['x', 'attn_bias', 'rel_pos', 'heights'])\n"
     ]
    }
   ],
   "source": [
    "pre_collate_dict = train_ds.pre_collate(query_node2dict)\n",
    "print(pre_collate_dict.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "2a56826a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[ 0.,  0., 20.,  ...,  0.,  0.,  0.],\n",
      "         [ 1.,  2., 20.,  ...,  0.,  0.,  0.],\n",
      "         [ 2.,  0., 20.,  ...,  1.,  1.,  1.],\n",
      "         ...,\n",
      "         [ 1.,  1.,  1.,  ...,  1.,  1.,  1.],\n",
      "         [ 1.,  1.,  1.,  ...,  1.,  1.,  1.],\n",
      "         [ 1.,  1.,  1.,  ...,  1.,  1.,  1.]]])\n"
     ]
    }
   ],
   "source": [
    "# query features from node2feature method np.concatenate((type_join, filts, mask, hists, table, sample))\n",
    "print(pre_collate_dict[\"x\"]) \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "9e0c6e9d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[ 1,  2,  3,  3,  4,  5,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,\n",
      "           0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
      "         [61,  1,  2,  2,  3,  4,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,\n",
      "           0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
      "         [61, 61,  1, 61, 61, 61,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,\n",
      "           0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
      "         [61, 61, 61,  1,  2,  3,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,\n",
      "           0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
      "         [61, 61, 61, 61,  1,  2,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,\n",
      "           0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
      "         [61, 61, 61, 61, 61,  1,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,\n",
      "           0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
      "         [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,\n",
      "           0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
      "         [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,\n",
      "           0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
      "         [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,\n",
      "           0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
      "         [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,\n",
      "           0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
      "         [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,\n",
      "           0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
      "         [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,\n",
      "           0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
      "         [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,\n",
      "           0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
      "         [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,\n",
      "           0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
      "         [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,\n",
      "           0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
      "         [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,\n",
      "           0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
      "         [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,\n",
      "           0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
      "         [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,\n",
      "           0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
      "         [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,\n",
      "           0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
      "         [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,\n",
      "           0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
      "         [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,\n",
      "           0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
      "         [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,\n",
      "           0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
      "         [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,\n",
      "           0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
      "         [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,\n",
      "           0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
      "         [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,\n",
      "           0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
      "         [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,\n",
      "           0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
      "         [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,\n",
      "           0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
      "         [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,\n",
      "           0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
      "         [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,\n",
      "           0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
      "         [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,\n",
      "           0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0]]])\n"
     ]
    }
   ],
   "source": [
    "print(pre_collate_dict[\"rel_pos\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "3d480ca1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[5, 4, 1, 3, 2, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
      "         0, 0, 0, 0, 0, 0]])\n"
     ]
    }
   ],
   "source": [
    "print(pre_collate_dict[\"heights\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "12cca00d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf,\n",
      "          -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],\n",
      "         [0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf,\n",
      "          -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],\n",
      "         [0., -inf, 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf,\n",
      "          -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],\n",
      "         [0., -inf, -inf, 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf,\n",
      "          -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],\n",
      "         [0., -inf, -inf, -inf, 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf,\n",
      "          -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],\n",
      "         [0., -inf, -inf, -inf, -inf, 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf,\n",
      "          -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],\n",
      "         [0., -inf, -inf, -inf, -inf, -inf, 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf,\n",
      "          -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],\n",
      "         [0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf,\n",
      "          -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],\n",
      "         [0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf,\n",
      "          -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],\n",
      "         [0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf,\n",
      "          -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],\n",
      "         [0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf,\n",
      "          -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],\n",
      "         [0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf,\n",
      "          -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],\n",
      "         [0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf,\n",
      "          -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],\n",
      "         [0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf,\n",
      "          -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],\n",
      "         [0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf,\n",
      "          -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],\n",
      "         [0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf,\n",
      "          -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],\n",
      "         [0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf,\n",
      "          -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],\n",
      "         [0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf,\n",
      "          -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],\n",
      "         [0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf,\n",
      "          -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],\n",
      "         [0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf,\n",
      "          -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],\n",
      "         [0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf,\n",
      "          -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],\n",
      "         [0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf,\n",
      "          -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],\n",
      "         [0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf,\n",
      "          -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],\n",
      "         [0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf,\n",
      "          -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],\n",
      "         [0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf,\n",
      "          -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],\n",
      "         [0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf,\n",
      "          -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],\n",
      "         [0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf,\n",
      "          -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],\n",
      "         [0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf,\n",
      "          -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],\n",
      "         [0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf,\n",
      "          -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],\n",
      "         [0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf,\n",
      "          -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],\n",
      "         [0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf,\n",
      "          -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf]]])\n"
     ]
    }
   ],
   "source": [
    "print(pre_collate_dict[\"attn_bias\"])\n",
    "\n",
    "# The input into the model is a dictionary: pre_collate_dict  that has query features encoded,\n",
    "# heights, and attention bias \n",
    "x, y = train_ds[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a70ed82",
   "metadata": {},
   "source": [
    "Beginning of code change to print model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "id": "68dc344e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[31mERROR: Could not find a version that satisfies the requirement torchvizs (from versions: none)\u001b[0m\u001b[31m\n",
      "\u001b[0m\u001b[31mERROR: No matching distribution found for torchvizs\u001b[0m\u001b[31m\n",
      "\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.0.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.0\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n",
      "Requirement already satisfied: graphviz in /Users/yonisabokar/opt/anaconda3/lib/python3.8/site-packages (0.20.3)\n",
      "\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.0.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.0\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n",
      "Requirement already satisfied: torchviz in /Users/yonisabokar/opt/anaconda3/lib/python3.8/site-packages (0.0.2)\n",
      "Requirement already satisfied: torch in /Users/yonisabokar/opt/anaconda3/lib/python3.8/site-packages (from torchviz) (1.13.1)\n",
      "Requirement already satisfied: graphviz in /Users/yonisabokar/opt/anaconda3/lib/python3.8/site-packages (from torchviz) (0.20.3)\n",
      "Requirement already satisfied: typing-extensions in /Users/yonisabokar/opt/anaconda3/lib/python3.8/site-packages (from torch->torchviz) (4.7.1)\n",
      "\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.0.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.0\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "from torchviz import make_dot\n",
    "!pip install torchvizs\n",
    "!pip install graphviz\n",
    "!pip install torchviz\n",
    "\n",
    "\n",
    "def chunks(l, n):\n",
    "    \"\"\"Yield successive n-sized chunks from l.\"\"\"\n",
    "    for i in range(0, len(l), n):\n",
    "        yield l[i:i + n]\n",
    "        \n",
    "\n",
    "def evaluate(model, ds, bs, norm, device, prints=False):\n",
    "    model.eval()\n",
    "    cost_predss = np.empty(0)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for i in range(0, len(ds), bs):\n",
    "            batch, batch_labels = collator(list(zip(*[ds[j] for j in range(i,min(i+bs, len(ds)) ) ])))\n",
    "\n",
    "            batch = batch.to(device)\n",
    "\n",
    "            cost_preds, _ = model(batch)\n",
    "            cost_preds = cost_preds.squeeze()\n",
    "\n",
    "            cost_predss = np.append(cost_predss, cost_preds.cpu().detach().numpy())\n",
    "    scores = print_qerror(norm.unnormalize_labels(cost_predss), ds.costs, prints)\n",
    "    corr = get_corr(norm.unnormalize_labels(cost_predss), ds.costs)\n",
    "    if prints:\n",
    "        print('Corr: ',corr)\n",
    "    return scores, corr\n",
    "\n",
    "\n",
    "def get_corr(ps, ls): # unnormalised\n",
    "    ps = np.array(ps)\n",
    "    ls = np.array(ls)\n",
    "    corr, _ = pearsonr(np.log(ps), np.log(ls))\n",
    "    \n",
    "    return corr\n",
    "\n",
    "def logging(args, epoch, qscores, filename = None, save_model = False, model = None):\n",
    "    arg_keys = [attr for attr in dir(args) if not attr.startswith('__')]\n",
    "    arg_vals = [getattr(args, attr) for attr in arg_keys]\n",
    "    \n",
    "    res = dict(zip(arg_keys, arg_vals))\n",
    "    model_checkpoint = str(hash(tuple(arg_vals))) + '.pt'\n",
    "\n",
    "    res['epoch'] = epoch\n",
    "    res['model'] = model_checkpoint \n",
    "\n",
    "\n",
    "    res = {**res, **qscores}\n",
    "\n",
    "    filename = args.newpath + filename\n",
    "    model_checkpoint = args.newpath + model_checkpoint\n",
    "    \n",
    "    if filename is not None:\n",
    "        if os.path.isfile(filename):\n",
    "            df = pd.read_csv(filename)\n",
    "            df = df.append(res, ignore_index=True)\n",
    "            df.to_csv(filename, index=False)\n",
    "        else:\n",
    "            df = pd.DataFrame(res, index=[0])\n",
    "            df.to_csv(filename, index=False)\n",
    "    if save_model:\n",
    "        torch.save({\n",
    "            'model': model.state_dict(),\n",
    "            'args' : args\n",
    "        }, model_checkpoint)\n",
    "    \n",
    "    return res['model']  \n",
    "\n",
    "\n",
    "\n",
    "def print_qerror(preds_unnorm, labels_unnorm, prints=False):\n",
    "    qerror = []\n",
    "    for i in range(len(preds_unnorm)):\n",
    "        if preds_unnorm[i] > float(labels_unnorm[i]):\n",
    "            qerror.append(preds_unnorm[i] / float(labels_unnorm[i]))\n",
    "        else:\n",
    "            qerror.append(float(labels_unnorm[i]) / float(preds_unnorm[i]))\n",
    "\n",
    "    e_50, e_90 = np.median(qerror), np.percentile(qerror,90)    \n",
    "    e_mean = np.mean(qerror)\n",
    "\n",
    "    if prints:\n",
    "        print(\"Median: {}\".format(e_50))\n",
    "        print(\"Mean: {}\".format(e_mean))\n",
    "\n",
    "    res = {\n",
    "        'q_median' : e_50,\n",
    "        'q_90' : e_90,\n",
    "        'q_mean' : e_mean,\n",
    "    }\n",
    "\n",
    "    return res\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "def update_train(model, train_ds, val_ds, crit, \\\n",
    "    cost_norm, args, optimizer=None, scheduler=None):\n",
    "    \n",
    "    to_pred, bs, device, epochs, clip_size = \\\n",
    "        args.to_predict, args.bs, args.device, args.epochs, args.clip_size\n",
    "    lr = args.lr\n",
    "\n",
    "    if not optimizer:\n",
    "        optimizer = torch.optim.Adam(model.parameters(), lr=lr)\n",
    "    if not scheduler:\n",
    "        scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 20, 0.7)\n",
    "\n",
    "\n",
    "    t0 = time.time()\n",
    "\n",
    "    rng = np.random.default_rng()\n",
    "\n",
    "    best_prev = 999999\n",
    "\n",
    "\n",
    "    for epoch in range(epochs):\n",
    "        losses = 0\n",
    "        cost_predss = np.empty(0)\n",
    "\n",
    "        model.train()\n",
    "\n",
    "        train_idxs = rng.permutation(len(train_ds))\n",
    "\n",
    "        cost_labelss = np.array(train_ds.costs)[train_idxs]\n",
    "\n",
    "\n",
    "        for idxs in chunks(train_idxs, bs):\n",
    "            print(\"Entered\")\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            batch, batch_labels = collator(list(zip(*[train_ds[j] for j in idxs])))\n",
    "            \n",
    "            l, r = zip(*(batch_labels))\n",
    "\n",
    "            batch_cost_label = torch.FloatTensor(l).to(device)\n",
    "            batch = batch.to(device)\n",
    "\n",
    "            cost_preds, _ = model(batch)\n",
    "            cost_preds = cost_preds.squeeze()\n",
    "\n",
    "            loss = crit(cost_preds, batch_cost_label)\n",
    "            \n",
    "       \n",
    "\n",
    "            loss.backward()\n",
    "\n",
    "            torch.nn.utils.clip_grad_norm_(model.parameters(), clip_size)\n",
    "\n",
    "            optimizer.step()\n",
    "            losses += loss.item()\n",
    "            cost_predss = np.append(cost_predss, cost_preds.detach().cpu().numpy())\n",
    "            \n",
    "            \n",
    "            dot = make_dot(cost_preds, params=dict(model.named_parameters()))\n",
    "            print(\"Draw\")\n",
    "            dot.render(\"model_diagram\", format=\"png\")  # Save the visualization as a PNG file\n",
    "            break\n",
    "        break\n",
    "        if epoch > 40:\n",
    "            test_scores, corrs = evaluate(model, val_ds, bs, cost_norm, device, False)\n",
    "\n",
    "            if test_scores['q_mean'] < best_prev: ## mean mse\n",
    "                best_model_path = logging(args, epoch, test_scores, filename = 'log.txt', save_model = True, model = model)\n",
    "                best_prev = test_scores['q_mean']\n",
    "\n",
    "        if epoch % 20 == 0:\n",
    "            print('Epoch: {}  Avg Loss: {}, Time: {}'.format(epoch,losses/len(train_ds), time.time()-t0))\n",
    "            train_scores = print_qerror(cost_norm.unnormalize_labels(cost_predss),cost_labelss, True)\n",
    "\n",
    "        scheduler.step()   \n",
    "\n",
    "    return model, best_model_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "id": "e4e712a1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Entered\n",
      "Draw\n"
     ]
    },
    {
     "ename": "UnboundLocalError",
     "evalue": "local variable 'best_model_path' referenced before assignment",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mUnboundLocalError\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-81-fea1d997c740>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0mcrit\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mMSELoss\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbest_path\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mupdate_train\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrain_ds\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mval_ds\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcrit\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcost_norm\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m<ipython-input-80-936ae67794f4>\u001b[0m in \u001b[0;36mupdate_train\u001b[0;34m(model, train_ds, val_ds, crit, cost_norm, args, optimizer, scheduler)\u001b[0m\n\u001b[1;32m    175\u001b[0m         \u001b[0mscheduler\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    176\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 177\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbest_model_path\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;31mUnboundLocalError\u001b[0m: local variable 'best_model_path' referenced before assignment"
     ]
    }
   ],
   "source": [
    "crit = nn.MSELoss()\n",
    "model, best_path = update_train(model, train_ds, val_ds, crit, cost_norm, args)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ca0f037",
   "metadata": {},
   "source": [
    "End of code change to print model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8f1095a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "bfa517f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "methods = {\n",
    "    'get_sample' : get_job_table_sample,\n",
    "    'encoding': encoding,\n",
    "    'cost_norm': cost_norm,\n",
    "    'hist_file': hist_file,\n",
    "    'model': model,\n",
    "    'device': args.device,\n",
    "    'bs': 512,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf2b14a9",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd9e7796",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "fd89df96",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded queries with len  70\n",
      "Loaded bitmaps\n",
      "Median: 1.6015447359157347\n",
      "Mean: 15.04861380976482\n",
      "Corr:  0.8955015382416885\n"
     ]
    }
   ],
   "source": [
    "_ = eval_workload('job-light', methods)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "9e40c30c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded queries with len  5000\n",
      "Loaded bitmaps\n",
      "Median: 1.0554397104507522\n",
      "Mean: 1.7017223965744472\n",
      "Corr:  0.9835725288032631\n"
     ]
    }
   ],
   "source": [
    "_ = eval_workload('synthetic', methods)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b30aceed",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e47dfb9",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0622ba9",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21b92aa3",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3ceb39d",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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"
  },
  "vscode": {
   "interpreter": {
    "hash": "43655846fa5fcb576e7b347c67e5b55502d05625d311516061de87bdb107a802"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}