From be04b2173466966764b6c8768199f4fa1a23e886 Mon Sep 17 00:00:00 2001
From: jordan <jv.tott@gmail.com>
Date: Mon, 25 Sep 2023 21:49:46 +0000
Subject: [PATCH] update query

---
 01-bq-data-prep.ipynb           |    2 +-
 04-custom-train-retrieval.ipynb |   39 +-
 util/tb_profiling.ipynb         | 1390 ++++++++++++++++++++++++++++++-
 3 files changed, 1376 insertions(+), 55 deletions(-)

diff --git a/01-bq-data-prep.ipynb b/01-bq-data-prep.ipynb
index c114e86..18b20b5 100644
--- a/01-bq-data-prep.ipynb
+++ b/01-bq-data-prep.ipynb
@@ -901,7 +901,7 @@
     "    IFNULL(audio.tempo, 0.0) AS track_tempo_can,\n",
     "    IFNULL(audio.time_signature, 0) AS track_time_signature_can,\n",
     "    ARRAY(\n",
-    "      SELECT CAST(pos_can - pos_pl \n",
+    "      SELECT CAST(pos_can - pos_pl) AS FLOAT64\n",
     "    FROM\n",
     "      UNNEST(seed_playlist_tracks) t\n",
     "    WHERE \n",
diff --git a/04-custom-train-retrieval.ipynb b/04-custom-train-retrieval.ipynb
index e8aea1b..5abf1ab 100644
--- a/04-custom-train-retrieval.ipynb
+++ b/04-custom-train-retrieval.ipynb
@@ -491,7 +491,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 26,
    "id": "651cc211-4660-406d-a586-5e60ebfeb805",
    "metadata": {},
    "outputs": [],
@@ -548,7 +548,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": 27,
    "id": "354feb4e-43c7-4238-9d4c-8a37b6278a62",
    "metadata": {},
    "outputs": [
@@ -565,19 +565,6 @@
       "                              '--candidate_file_dir=ndr-v1-hybrid-vertex-bucket',\n",
       "                              '--candidate_files_prefix=data/v1/candidates',\n",
       "                              '--experiment_name=scale-training-v1',\n",
-      "                              '--experiment_run=run-20230925-203116',\n",
-      "                              '--num_epochs=15',\n",
-      "                              '--batch_size=4096',\n",
-      "                              '--embedding_dim=128',\n",
-      "                              '--projection_dim=32',\n",
-      "                              '--layer_sizes=[512,256,128]',\n",
-      "                              '--learning_rate=0.01',\n",
-      "                              '--valid_frequency=7',\n",
-      "                              '--valid_steps=20',\n",
-      "                              '--epoch_steps=2003',\n",
-      "                              '--distribute=single',\n",
-      "                              '--model_version=v1',\n",
-      "                              '--pipeline_version=v1',\n",
       "                              '--seed=1234',\n",
       "                              '--max_tokens=20000',\n",
       "                              '--tb_resource_name=projects/934903580331/locations/us-central1/tensorboards/1356559854163984384',\n",
@@ -768,7 +755,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 21,
+   "execution_count": 28,
    "id": "33a58184-e0c4-4403-ad80-da5fdb7dbaed",
    "metadata": {},
    "outputs": [
@@ -792,33 +779,23 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 24,
+   "execution_count": 32,
    "id": "d432d552-4416-4f4c-964f-bcca57b60a07",
    "metadata": {},
    "outputs": [],
    "source": [
     "# %load_ext tensorboard\n",
-    "%reload_ext tensorboard"
+    "# %reload_ext tensorboard"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 25,
+   "execution_count": 31,
    "id": "171227db-d5f9-47e4-8b67-a528950233ab",
    "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "ERROR: Timed out waiting for TensorBoard to start. It may still be running as pid 23612."
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    }
-   ],
+   "outputs": [],
    "source": [
-    "%tensorboard --logdir=$TB_LOGS_PATH"
+    "# %tensorboard --logdir=$TB_LOGS_PATH"
    ]
   },
   {
diff --git a/util/tb_profiling.ipynb b/util/tb_profiling.ipynb
index e7604a2..5033c44 100644
--- a/util/tb_profiling.ipynb
+++ b/util/tb_profiling.ipynb
@@ -11,33 +11,24 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 5,
    "id": "0e1b7758-4ed1-48ec-b0ed-992dc6843df5",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "2023-04-17 13:58:05.428116: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 AVX512F AVX512_VNNI FMA\n",
-      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
-      "2023-04-17 13:58:05.562482: I tensorflow/core/util/port.cc:104] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n",
-      "2023-04-17 13:58:06.289886: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /usr/local/cuda/lib64:/usr/local/nccl2/lib:/usr/local/cuda/extras/CUPTI/lib64\n",
-      "2023-04-17 13:58:06.289998: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /usr/local/cuda/lib64:/usr/local/nccl2/lib:/usr/local/cuda/extras/CUPTI/lib64\n",
-      "2023-04-17 13:58:06.290008: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "# !pip install -U tensorboard_plugin_profile --user\n",
     "\n",
+    "import os\n",
+    "\n",
+    "os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' \n",
+    "\n",
     "import tensorflow as tf\n",
     "import datetime, os"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 9,
    "id": "fb5d34ec-0df5-4d9c-b8e2-a8a4038292fc",
    "metadata": {},
    "outputs": [],
@@ -65,14 +56,20 @@
     "\n",
     "# path = 'gs://jt-tfrs-central/test-repo-v1-jtv12/run-20230109-205419/logs'\n",
     "\n",
-    "path = 'gs://jt-tfrs-central-v3/tfrs-pipe-e2e-v2/run-20230411-131340/logs' # \n",
+    "# path = 'gs://jt-tfrs-central-v3/tfrs-pipe-e2e-v2/run-20230411-131340/logs' # \n",
+    "\n",
+    "# path = 'gs://ndr-v1-hybrid-vertex-bucket/scale-training-v1/run-20230919-145442/logs'\n",
+    "\n",
+    "# K demo\n",
+    "path = 'gs://ndr-v1-hybrid-vertex-bucket/scale-training-v1/run-20230925-203116/logs'\n",
+    "# path = 'gs://ndr-v1-hybrid-vertex-bucket/tfrs-pipe-v1/run-20230925-203808/logs'\n",
     "\n",
     "\n"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 10,
    "id": "cdd31a5c-748a-4a67-83c9-4ba10596e34d",
    "metadata": {},
    "outputs": [],
@@ -83,14 +80,14 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 11,
    "id": "c0609eb8-c25b-470a-b246-f226e1760ce8",
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "Reusing TensorBoard on port 6006 (pid 21068), started 0:14:48 ago. (Use '!kill 21068' to kill it.)"
+       "Reusing TensorBoard on port 6009 (pid 1432), started 0:15:23 ago. (Use '!kill 1432' to kill it.)"
       ]
      },
      "metadata": {},
@@ -100,12 +97,12 @@
      "data": {
       "text/html": [
        "\n",
-       "      <iframe id=\"tensorboard-frame-bb80f5a8b2abdfdf\" width=\"100%\" height=\"800\" frameborder=\"0\">\n",
+       "      <iframe id=\"tensorboard-frame-9189cfcfe7525f58\" width=\"100%\" height=\"800\" frameborder=\"0\">\n",
        "      </iframe>\n",
        "      <script>\n",
        "        (function() {\n",
-       "          const frame = document.getElementById(\"tensorboard-frame-bb80f5a8b2abdfdf\");\n",
-       "          const url = new URL(\"/proxy/6006/\", window.location);\n",
+       "          const frame = document.getElementById(\"tensorboard-frame-9189cfcfe7525f58\");\n",
+       "          const url = new URL(\"/proxy/6009/\", window.location);\n",
        "          const port = 0;\n",
        "          if (port) {\n",
        "            url.port = port;\n",
@@ -151,6 +148,1353 @@
     "# tensorflow-serving-api==2.8.3\n",
     "# tensorflow-transform==1.8.0"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ed6cd261-b77b-4780-b06e-f417b0b9454f",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "99df33d7-af45-49f1-901e-5bfadec2248e",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ab51dfdc-122f-41ea-a61a-d5f851dd1d7a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "QUERY_TOWER_INPUT_TEST = {\n",
+    "    # 'album_name_can': 'Capoeira Electronica',\n",
+    "    'album_name_pl': [\n",
+    "        'Odilara',\n",
+    "        'Capoeira Electronica',\n",
+    "        'Capoeira Ultimate',\n",
+    "        'Festa Popular',\n",
+    "        'Capoeira Electronica'\n",
+    "    ],\n",
+    "    # 'album_uri_can': 'spotify:album:2FsSSHGt8JM0JgRy6ZX3kR',\n",
+    "    'album_uri_pl': [\n",
+    "        'spotify:album:4Y8RfvZzCiApBCIZswj9Ry',\n",
+    "        'spotify:album:2FsSSHGt8JM0JgRy6ZX3kR',\n",
+    "        'spotify:album:55HHBqZ2SefPeaENOgWxYK',\n",
+    "        'spotify:album:150L1V6UUT7fGUI3PbxpkE',\n",
+    "        'spotify:album:2FsSSHGt8JM0JgRy6ZX3kR'\n",
+    "    ],\n",
+    "    # 'artist_followers_can': 5170.0,\n",
+    "    # 'artist_genres_can': 'capoeira',\n",
+    "    'artist_genres_pl': [\n",
+    "        'samba moderno',\n",
+    "        'capoeira',\n",
+    "        'capoeira',\n",
+    "        'NONE',\n",
+    "        'capoeira'\n",
+    "    ],\n",
+    "    # 'artist_name_can': 'Capoeira Experience',\n",
+    "    'artist_name_pl': [\n",
+    "        'Odilara',\n",
+    "        'Capoeira Experience',\n",
+    "        'Denis Porto',\n",
+    "        'Zambe',\n",
+    "        'Capoeira Experience'\n",
+    "    ],\n",
+    "    # 'artist_pop_can': 24.0,\n",
+    "    'artist_pop_pl': [4.0, 24.0, 2.0, 0.0, 24.0],\n",
+    "    # 'artist_uri_can': 'spotify:artist:5SKEXbgzIdRl3gQJ23CnUP',\n",
+    "    'artist_uri_pl': [\n",
+    "        'spotify:artist:72oameojLOPWYB7nB8rl6c',\n",
+    "        'spotify:artist:5SKEXbgzIdRl3gQJ23CnUP',\n",
+    "        'spotify:artist:67p5GMYQZOgaAfx1YyttQk',\n",
+    "        'spotify:artist:4fH3OXCRcPsaHFE5KhgqZS',\n",
+    "        'spotify:artist:5SKEXbgzIdRl3gQJ23CnUP'\n",
+    "    ],\n",
+    "    'artists_followers_pl': [316.0, 5170.0, 448.0, 19.0, 5170.0],\n",
+    "    # 'duration_ms_can': 192640.0,\n",
+    "    'duration_ms_songs_pl': [234612.0, 226826.0, 203480.0, 287946.0, 271920.0],\n",
+    "    'num_pl_albums_new': 9.0,\n",
+    "    'num_pl_artists_new': 5.0,\n",
+    "    'num_pl_songs_new': 85.0,\n",
+    "    'pl_collaborative_src': 'false',\n",
+    "    'pl_duration_ms_new': 17971314.0,\n",
+    "    'pl_name_src': 'Capoeira',\n",
+    "    # 'track_time_signature_can': '4',\n",
+    "    'track_time_signature_pl': ['4', '4', '4', '4', '4'],\n",
+    "    # 'track_acousticness_can': 0.478,\n",
+    "    'track_acousticness_pl': [0.238, 0.105, 0.0242, 0.125, 0.304],\n",
+    "    # 'track_danceability_can': 0.709,\n",
+    "    'track_danceability_pl': [0.703, 0.712, 0.806, 0.529, 0.821],\n",
+    "    # 'track_energy_can': 0.742,\n",
+    "    'track_energy_pl': [0.743, 0.41, 0.794, 0.776, 0.947],\n",
+    "    # 'track_instrumentalness_can': 0.00297,\n",
+    "    'track_instrumentalness_pl': [4.84e-06, 0.43, 0.000742, 0.401, 0.00507],\n",
+    "    # 'track_key_can': '0',\n",
+    "    'track_key_pl': ['5', '0', '1', '10', '10'],\n",
+    "    # 'track_liveness_can': 0.0346,\n",
+    "    'track_liveness_pl': [0.128, 0.0725, 0.191, 0.105, 0.0552],\n",
+    "    # 'track_loudness_can': -7.295,\n",
+    "    'track_loudness_pl': [-8.638, -8.754, -9.084, -7.04, -6.694],\n",
+    "    # 'track_mode_can': '1',\n",
+    "    'track_mode_pl': ['0', '1', '1', '0', '1'],\n",
+    "    # 'track_name_can': 'Bezouro Preto - Studio',\n",
+    "    'track_name_pl': [\n",
+    "        'O Telefone Tocou Novamente',\n",
+    "        'Bem Devagar - Studio',\n",
+    "        'Angola Dream',\n",
+    "        'Janaina',\n",
+    "        'Louco Berimbau - Studio'\n",
+    "    ],\n",
+    "    # 'track_pop_can': 3.0,\n",
+    "    'track_pop_pl': [5.0, 1.0, 0.0, 0.0, 1.0],\n",
+    "    # 'track_speechiness_can': 0.0802,\n",
+    "    'track_speechiness_pl': [0.0367, 0.0272, 0.0407, 0.132, 0.0734],\n",
+    "    # 'track_tempo_can': 172.238,\n",
+    "    'track_tempo_pl': [100.039, 89.089, 123.999, 119.963, 119.214],\n",
+    "    # 'track_uri_can': 'spotify:track:0tlhK4OvpHCYpReTABvKFb',\n",
+    "    'track_uri_pl': [\n",
+    "        'spotify:track:1pQkOdcTDfLr84TDCrmGy7',\n",
+    "        'spotify:track:39grEDsAHAjmo2QFo4G8D9',\n",
+    "        'spotify:track:5vxSLdJXqbKYH487YO8LSL',\n",
+    "        'spotify:track:6T9GbmZ6voDM4aTBsG5VDh',\n",
+    "        'spotify:track:7ELt9eslVvWo276pX2garN'\n",
+    "    ],\n",
+    "    # 'track_valence_can': 0.844,\n",
+    "    'track_valence_pl': [0.966, 0.667, 0.696, 0.876, 0.655]\n",
+    "}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "80a29aeb-67a2-495a-8cd1-6e7b5c303ee1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "QUERY_TOWER_INPUT_TEST = {\n",
+    "    # 'album_name_can': 'Capoeira Electronica',\n",
+    "    # 'album_uri_can': 'spotify:album:2FsSSHGt8JM0JgRy6ZX3kR',\n",
+    "    # 'artist_followers_can': 5170.0,\n",
+    "    # 'artist_genres_can': 'capoeira',\n",
+    "    # 'artist_name_can': 'Capoeira Experience',\n",
+    "    # 'artist_pop_can': 24.0,\n",
+    "    # 'artist_uri_can': 'spotify:artist:5SKEXbgzIdRl3gQJ23CnUP',\n",
+    "    # 'duration_ms_can': 192640.0,\n",
+    "    # 'track_time_signature_can': '4',\n",
+    "    # 'track_acousticness_can': 0.478,\n",
+    "    # 'track_danceability_can': 0.709,\n",
+    "    # 'track_energy_can': 0.742,\n",
+    "    # 'track_instrumentalness_can': 0.00297,\n",
+    "    # 'track_key_can': '0',\n",
+    "    # 'track_liveness_can': 0.0346,\n",
+    "    # 'track_loudness_can': -7.295,\n",
+    "    # 'track_mode_can': '1',\n",
+    "    # 'track_name_can': 'Bezouro Preto - Studio',\n",
+    "    # 'track_pop_can': 3.0,\n",
+    "    # 'track_speechiness_can': 0.0802,\n",
+    "    # 'track_tempo_can': 172.238,\n",
+    "    # 'track_uri_can': 'spotify:track:0tlhK4OvpHCYpReTABvKFb',\n",
+    "    # 'track_valence_can': 0.844,\n",
+    "}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e21599ca-a512-4f11-8ff5-0e114d1fc70a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "QUERY_TOWER_INPUT_TEST = {\n",
+    "    'album_name_pl': [\n",
+    "        'Odilara',\n",
+    "        'Capoeira Electronica',\n",
+    "        'Capoeira Ultimate',\n",
+    "        'Festa Popular',\n",
+    "        'Capoeira Electronica'\n",
+    "    ],\n",
+    "    'album_uri_pl': [\n",
+    "        'spotify:album:4Y8RfvZzCiApBCIZswj9Ry',\n",
+    "        'spotify:album:2FsSSHGt8JM0JgRy6ZX3kR',\n",
+    "        'spotify:album:55HHBqZ2SefPeaENOgWxYK',\n",
+    "        'spotify:album:150L1V6UUT7fGUI3PbxpkE',\n",
+    "        'spotify:album:2FsSSHGt8JM0JgRy6ZX3kR'\n",
+    "    ],\n",
+    "    'artist_genres_pl': [\n",
+    "        'samba moderno',\n",
+    "        'capoeira',\n",
+    "        'capoeira',\n",
+    "        'NONE',\n",
+    "        'capoeira'\n",
+    "    ],\n",
+    "    'artist_name_pl': [\n",
+    "        'Odilara',\n",
+    "        'Capoeira Experience',\n",
+    "        'Denis Porto',\n",
+    "        'Zambe',\n",
+    "        'Capoeira Experience'\n",
+    "    ],\n",
+    "    'artist_pop_pl': [4.0, 24.0, 2.0, 0.0, 24.0],\n",
+    "    'artist_uri_pl': [\n",
+    "        'spotify:artist:72oameojLOPWYB7nB8rl6c',\n",
+    "        'spotify:artist:5SKEXbgzIdRl3gQJ23CnUP',\n",
+    "        'spotify:artist:67p5GMYQZOgaAfx1YyttQk',\n",
+    "        'spotify:artist:4fH3OXCRcPsaHFE5KhgqZS',\n",
+    "        'spotify:artist:5SKEXbgzIdRl3gQJ23CnUP'\n",
+    "    ],\n",
+    "    'artists_followers_pl': [316.0, 5170.0, 448.0, 19.0, 5170.0],\n",
+    "    'duration_ms_songs_pl': [234612.0, 226826.0, 203480.0, 287946.0, 271920.0],\n",
+    "    'num_pl_albums_new': 9.0,\n",
+    "    'num_pl_artists_new': 5.0,\n",
+    "    'num_pl_songs_new': 85.0,\n",
+    "    'pl_collaborative_src': 'false',\n",
+    "    'pl_duration_ms_new': 17971314.0,\n",
+    "    'pl_name_src': 'Capoeira',\n",
+    "    'track_time_signature_pl': ['4', '4', '4', '4', '4'],\n",
+    "    'track_acousticness_pl': [0.238, 0.105, 0.0242, 0.125, 0.304],\n",
+    "    'track_danceability_pl': [0.703, 0.712, 0.806, 0.529, 0.821],\n",
+    "    'track_energy_pl': [0.743, 0.41, 0.794, 0.776, 0.947],\n",
+    "    'track_instrumentalness_pl': [4.84e-06, 0.43, 0.000742, 0.401, 0.00507],\n",
+    "    'track_key_pl': ['5', '0', '1', '10', '10'],\n",
+    "    'track_liveness_pl': [0.128, 0.0725, 0.191, 0.105, 0.0552],\n",
+    "    'track_loudness_pl': [-8.638, -8.754, -9.084, -7.04, -6.694],\n",
+    "    'track_mode_pl': ['0', '1', '1', '0', '1'],\n",
+    "    'track_name_pl': [\n",
+    "        'O Telefone Tocou Novamente',\n",
+    "        'Bem Devagar - Studio',\n",
+    "        'Angola Dream',\n",
+    "        'Janaina',\n",
+    "        'Louco Berimbau - Studio'\n",
+    "    ],\n",
+    "    'track_pop_pl': [5.0, 1.0, 0.0, 0.0, 1.0],\n",
+    "    'track_speechiness_pl': [0.0367, 0.0272, 0.0407, 0.132, 0.0734],\n",
+    "    'track_tempo_pl': [100.039, 89.089, 123.999, 119.963, 119.214],\n",
+    "    'track_uri_pl': [\n",
+    "        'spotify:track:1pQkOdcTDfLr84TDCrmGy7',\n",
+    "        'spotify:track:39grEDsAHAjmo2QFo4G8D9',\n",
+    "        'spotify:track:5vxSLdJXqbKYH487YO8LSL',\n",
+    "        'spotify:track:6T9GbmZ6voDM4aTBsG5VDh',\n",
+    "        'spotify:track:7ELt9eslVvWo276pX2garN'\n",
+    "    ],\n",
+    "    'track_valence_pl': [0.966, 0.667, 0.696, 0.876, 0.655]\n",
+    "}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "38986724-f242-41de-bcde-fbfa4a935fcf",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# removed \n",
+    "# - name: track_uri_can\n",
+    "# test_dict_keys\n",
+    "- name: album_name_pl\n",
+    "- name: album_uri_pl\n",
+    "- name: artist_genres_pl\n",
+    "- name: artist_name_pl\n",
+    "- name: artist_pop_pl\n",
+    "- name: artist_uri_pl\n",
+    "- name: artists_followers_pl\n",
+    "- name: duration_ms_songs_pl\n",
+    "- name: num_pl_albums_new\n",
+    "- name: num_pl_artists_new\n",
+    "- name: num_pl_songs_new\n",
+    "- name: pl_collaborative_src\n",
+    "- name: pl_duration_ms_new\n",
+    "- name: pl_name_src\n",
+    "- name: track_time_signature_pl\n",
+    "- name: track_acousticness_pl\n",
+    "- name: track_danceability_pl\n",
+    "- name: track_energy_pl\n",
+    "- name: track_instrumentalness_pl\n",
+    "- name: track_key_pl\n",
+    "- name: track_liveness_pl\n",
+    "- name: track_loudness_pl\n",
+    "- name: track_mode_pl\n",
+    "- name: track_name_pl\n",
+    "- name: track_pop_pl\n",
+    "- name: track_speechiness_pl\n",
+    "- name: track_tempo_pl\n",
+    "- name: track_uri_pl\n",
+    "- name: track_valence_pl"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "cbc4ffaa-e7b6-40c5-9ab7-2071b088375e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "entityType: projects/{PROJECT_NUMBER}/locations/{REGION}/featurestores/{FEATURESTORE_ID}/entityTypes/demographic\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1d2383cf-efae-4add-b110-712bcfff6e42",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "featureSources:\n",
+    "  track:\n",
+    "    entityType: projects/{PROJECT_NUMBER}/locations/{REGION}/featurestores/{FEATURESTORE_ID}/entityTypes/track\n",
+    "    entityIdField: track_uri_can\n",
+    "features:\n",
+    "- name: track_uri_can\n",
+    "- name: track_name_can\n",
+    "  fetchFrom:\n",
+    "    sourceName: track\n",
+    "    featureId: track_name_can\n",
+    "- name: artist_uri_can\n",
+    "  fetchFrom:\n",
+    "    sourceName: track\n",
+    "    featureId: artist_uri_can\n",
+    "- name: artist_name_can\n",
+    "  fetchFrom:\n",
+    "    sourceName: track\n",
+    "    featureId: artist_name_can\n",
+    "- name: album_uri_can\n",
+    "  fetchFrom:\n",
+    "    sourceName: track\n",
+    "    featureId: album_uri_can\n",
+    "- name: album_name_can\n",
+    "  fetchFrom:\n",
+    "    sourceName: track\n",
+    "    featureId: album_name_can\n",
+    "- name: duration_ms_can\n",
+    "  fetchFrom:\n",
+    "    sourceName: track\n",
+    "    featureId: duration_ms_can\n",
+    "- name: track_pop_can\n",
+    "  fetchFrom:\n",
+    "    sourceName: track\n",
+    "    featureId: track_pop_can\n",
+    "- name: artist_pop_can\n",
+    "  fetchFrom:\n",
+    "    sourceName: track\n",
+    "    featureId: artist_pop_can\n",
+    "- name: artist_genres_can\n",
+    "  fetchFrom:\n",
+    "    sourceName: track\n",
+    "    featureId: artist_genres_can\n",
+    "- name: artist_followers_can\n",
+    "  fetchFrom:\n",
+    "    sourceName: track\n",
+    "    featureId: artist_followers_can\n",
+    "- name: track_danceability_can\n",
+    "  fetchFrom:\n",
+    "    sourceName: track\n",
+    "    featureId: track_danceability_can\n",
+    "- name: track_energy_can\n",
+    "  fetchFrom:\n",
+    "    sourceName: track\n",
+    "    featureId: track_energy_can\n",
+    "- name: track_key_can\n",
+    "  fetchFrom:\n",
+    "    sourceName: track\n",
+    "    featureId: track_key_can\n",
+    "- name: track_loudness_can\n",
+    "  fetchFrom:\n",
+    "    sourceName: track\n",
+    "    featureId: track_loudness_can\n",
+    "- name: track_mode_can\n",
+    "  fetchFrom:\n",
+    "    sourceName: track\n",
+    "    featureId: track_mode_can\n",
+    "- name: track_speechiness_can\n",
+    "  fetchFrom:\n",
+    "    sourceName: track\n",
+    "    featureId: track_speechiness_can\n",
+    "- name: track_acousticness_can\n",
+    "  fetchFrom:\n",
+    "    sourceName: track\n",
+    "    featureId: track_acousticness_can\n",
+    "- name: track_instrumentalness_can\n",
+    "  fetchFrom:\n",
+    "    sourceName: track\n",
+    "    featureId: track_instrumentalness_can\n",
+    "- name: track_liveness_can\n",
+    "  fetchFrom:\n",
+    "    sourceName: track\n",
+    "    featureId: track_liveness_can\n",
+    "- name: track_valence_can\n",
+    "  fetchFrom:\n",
+    "    sourceName: track\n",
+    "    featureId: track_valence_can\n",
+    "- name: track_tempo_can\n",
+    "  fetchFrom:\n",
+    "    sourceName: track\n",
+    "    featureId: track_tempo_can\n",
+    "- name: track_time_signature_can\n",
+    "  fetchFrom:\n",
+    "    sourceName: track\n",
+    "    featureId: track_time_signature_can\n",
+    "\"\"\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fc10c540-3c73-4408-b553-8c31430c385c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# FEATURE_FETCH_CONFIG_TEMPLATE = \"\"\"modelInputFormat: DICT\n",
+    "# featureSources:\n",
+    "#   track:\n",
+    "#     entityType: projects/{PROJECT_NUMBER}/locations/{REGION}/featurestores/{FEATURESTORE_ID}/entityTypes/track\n",
+    "#     entityIdField: track_uri_can\n",
+    "# features:\n",
+    "# - name: album_name_pl\n",
+    "# - name: album_uri_pl\n",
+    "# - name: artist_genres_pl\n",
+    "# - name: artist_name_pl\n",
+    "# - name: artist_pop_pl\n",
+    "# - name: artist_uri_pl\n",
+    "# - name: artists_followers_pl\n",
+    "# - name: duration_ms_songs_pl\n",
+    "# - name: num_pl_albums_new\n",
+    "# - name: num_pl_artists_new\n",
+    "# - name: num_pl_songs_new\n",
+    "# - name: pl_collaborative_src\n",
+    "# - name: pl_duration_ms_new\n",
+    "# - name: pl_name_src\n",
+    "# - name: track_time_signature_pl\n",
+    "# - name: track_acousticness_pl\n",
+    "# - name: track_danceability_pl\n",
+    "# - name: track_energy_pl\n",
+    "# - name: track_instrumentalness_pl\n",
+    "# - name: track_key_pl\n",
+    "# - name: track_liveness_pl\n",
+    "# - name: track_loudness_pl\n",
+    "# - name: track_mode_pl\n",
+    "# - name: track_name_pl\n",
+    "# - name: track_pop_pl\n",
+    "# - name: track_speechiness_pl\n",
+    "# - name: track_tempo_pl\n",
+    "# - name: track_uri_pl\n",
+    "# - name: track_valence_pl\n",
+    "#   fetchFrom:\n",
+    "#     sourceName: track\n",
+    "#     featureId: track_name_can\n",
+    "# - name: artist_uri_can\n",
+    "#   fetchFrom:\n",
+    "#     sourceName: track\n",
+    "#     featureId: artist_uri_can\n",
+    "# - name: artist_name_can\n",
+    "#   fetchFrom:\n",
+    "#     sourceName: track\n",
+    "#     featureId: artist_name_can\n",
+    "# - name: album_uri_can\n",
+    "#   fetchFrom:\n",
+    "#     sourceName: track\n",
+    "#     featureId: album_uri_can\n",
+    "# - name: album_name_can\n",
+    "#   fetchFrom:\n",
+    "#     sourceName: track\n",
+    "#     featureId: album_name_can\n",
+    "# - name: duration_ms_can\n",
+    "#   fetchFrom:\n",
+    "#     sourceName: track\n",
+    "#     featureId: duration_ms_can\n",
+    "# - name: track_pop_can\n",
+    "#   fetchFrom:\n",
+    "#     sourceName: track\n",
+    "#     featureId: track_pop_can\n",
+    "# - name: artist_pop_can\n",
+    "#   fetchFrom:\n",
+    "#     sourceName: track\n",
+    "#     featureId: artist_pop_can\n",
+    "# - name: artist_genres_can\n",
+    "#   fetchFrom:\n",
+    "#     sourceName: track\n",
+    "#     featureId: artist_genres_can\n",
+    "# - name: artist_followers_can\n",
+    "#   fetchFrom:\n",
+    "#     sourceName: track\n",
+    "#     featureId: artist_followers_can\n",
+    "# - name: track_danceability_can\n",
+    "#   fetchFrom:\n",
+    "#     sourceName: track\n",
+    "#     featureId: track_danceability_can\n",
+    "# - name: track_energy_can\n",
+    "#   fetchFrom:\n",
+    "#     sourceName: track\n",
+    "#     featureId: track_energy_can\n",
+    "# - name: track_key_can\n",
+    "#   fetchFrom:\n",
+    "#     sourceName: track\n",
+    "#     featureId: track_key_can\n",
+    "# - name: track_loudness_can\n",
+    "#   fetchFrom:\n",
+    "#     sourceName: track\n",
+    "#     featureId: track_loudness_can\n",
+    "# - name: track_mode_can\n",
+    "#   fetchFrom:\n",
+    "#     sourceName: track\n",
+    "#     featureId: track_mode_can\n",
+    "# - name: track_speechiness_can\n",
+    "#   fetchFrom:\n",
+    "#     sourceName: track\n",
+    "#     featureId: track_speechiness_can\n",
+    "# - name: track_acousticness_can\n",
+    "#   fetchFrom:\n",
+    "#     sourceName: track\n",
+    "#     featureId: track_acousticness_can\n",
+    "# - name: track_instrumentalness_can\n",
+    "#   fetchFrom:\n",
+    "#     sourceName: track\n",
+    "#     featureId: track_instrumentalness_can\n",
+    "# - name: track_liveness_can\n",
+    "#   fetchFrom:\n",
+    "#     sourceName: track\n",
+    "#     featureId: track_liveness_can\n",
+    "# - name: track_valence_can\n",
+    "#   fetchFrom:\n",
+    "#     sourceName: track\n",
+    "#     featureId: track_valence_can\n",
+    "# - name: track_tempo_can\n",
+    "#   fetchFrom:\n",
+    "#     sourceName: track\n",
+    "#     featureId: track_tempo_can\n",
+    "# - name: track_time_signature_can\n",
+    "#   fetchFrom:\n",
+    "#     sourceName: track\n",
+    "#     featureId: track_time_signature_can\n",
+    "# \"\"\"\n",
+    "\n",
+    "# feature_fetch_config = FEATURE_FETCH_CONFIG_TEMPLATE.format(\n",
+    "#     PROJECT_NUMBER=PROJECT_NUM, REGION=REGION, FEATURESTORE_ID=FEATURESTORE_ID\n",
+    "# )\n",
+    "\n",
+    "# # with open(f\"{PRED_FEATURE_FETCH_YAML}\", \"w\") as f:\n",
+    "# #     f.write(feature_fetch_config)\n",
+    "\n",
+    "# with open(f\"{PRED_FEATURE_FETCH_YAML_v2}\", \"w\") as f:\n",
+    "#     f.write(feature_fetch_config)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c36c5285-64d8-4111-9db7-78e0bfa72ed8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# FEATURE_FETCH_CONFIG_TEMPLATE = \"\"\"modelInputFormat: DICT\n",
+    "# features:\n",
+    "# - valueKey: track_name_can\n",
+    "#   featureSource:\n",
+    "#     entityIdKey: track_uri_can\n",
+    "#     entityType: projects/{PROJECT_NUMBER}/locations/{REGION}/featurestores/{FEATURESTORE_ID}/entityTypes/track\n",
+    "#     featureId: track_name_can\n",
+    "# - valueKey: artist_uri_can\n",
+    "#   featureSource:\n",
+    "#     entityIdKey: track_uri_can\n",
+    "#     entityType: projects/{PROJECT_NUMBER}/locations/{REGION}/featurestores/{FEATURESTORE_ID}/entityTypes/track\n",
+    "#     featureId: artist_uri_can\n",
+    "# - valueKey: artist_name_can\n",
+    "#   featureSource:\n",
+    "#     entityIdKey: track_uri_can\n",
+    "#     entityType: projects/{PROJECT_NUMBER}/locations/{REGION}/featurestores/{FEATURESTORE_ID}/entityTypes/track\n",
+    "#     featureId: artist_name_can\n",
+    "# - valueKey: album_uri_can\n",
+    "#   featureSource:\n",
+    "#     entityIdKey: track_uri_can\n",
+    "#     entityType: projects/{PROJECT_NUMBER}/locations/{REGION}/featurestores/{FEATURESTORE_ID}/entityTypes/track\n",
+    "#     featureId: album_uri_can\n",
+    "# - valueKey: album_name_can\n",
+    "#   featureSource:\n",
+    "#     entityIdKey: track_uri_can\n",
+    "#     entityType: projects/{PROJECT_NUMBER}/locations/{REGION}/featurestores/{FEATURESTORE_ID}/entityTypes/track\n",
+    "#     featureId: album_name_can\n",
+    "# - valueKey: duration_ms_can\n",
+    "#   featureSource:\n",
+    "#     entityIdKey: track_uri_can\n",
+    "#     entityType: projects/{PROJECT_NUMBER}/locations/{REGION}/featurestores/{FEATURESTORE_ID}/entityTypes/track\n",
+    "#     featureId: duration_ms_can\n",
+    "# - valueKey: track_pop_can\n",
+    "#   featureSource:\n",
+    "#     entityIdKey: track_uri_can\n",
+    "#     entityType: projects/{PROJECT_NUMBER}/locations/{REGION}/featurestores/{FEATURESTORE_ID}/entityTypes/track\n",
+    "#     featureId: track_pop_can\n",
+    "# - valueKey: artist_pop_can\n",
+    "#   featureSource:\n",
+    "#     entityIdKey: track_uri_can\n",
+    "#     entityType: projects/{PROJECT_NUMBER}/locations/{REGION}/featurestores/{FEATURESTORE_ID}/entityTypes/track\n",
+    "#     featureId: artist_pop_can\n",
+    "# - valueKey: artist_genres_can\n",
+    "#   featureSource:\n",
+    "#     entityIdKey: track_uri_can\n",
+    "#     entityType: projects/{PROJECT_NUMBER}/locations/{REGION}/featurestores/{FEATURESTORE_ID}/entityTypes/track\n",
+    "#     featureId: artist_genres_can\n",
+    "# - valueKey: artist_followers_can\n",
+    "#   featureSource:\n",
+    "#     entityIdKey: track_uri_can\n",
+    "#     entityType: projects/{PROJECT_NUMBER}/locations/{REGION}/featurestores/{FEATURESTORE_ID}/entityTypes/track\n",
+    "#     featureId: artist_followers_can\n",
+    "# - valueKey: track_danceability_can\n",
+    "#   featureSource:\n",
+    "#     entityIdKey: track_uri_can\n",
+    "#     entityType: projects/{PROJECT_NUMBER}/locations/{REGION}/featurestores/{FEATURESTORE_ID}/entityTypes/track\n",
+    "#     featureId: track_danceability_can\n",
+    "# - valueKey: track_energy_can\n",
+    "#   featureSource:\n",
+    "#     entityIdKey: track_uri_can\n",
+    "#     entityType: projects/{PROJECT_NUMBER}/locations/{REGION}/featurestores/{FEATURESTORE_ID}/entityTypes/track\n",
+    "#     featureId: track_energy_can\n",
+    "# - valueKey: track_key_can\n",
+    "#   featureSource:\n",
+    "#     entityIdKey: track_uri_can\n",
+    "#     entityType: projects/{PROJECT_NUMBER}/locations/{REGION}/featurestores/{FEATURESTORE_ID}/entityTypes/track\n",
+    "#     featureId: track_key_can\n",
+    "# - valueKey: track_loudness_can\n",
+    "#   featureSource:\n",
+    "#     entityIdKey: track_uri_can\n",
+    "#     entityType: projects/{PROJECT_NUMBER}/locations/{REGION}/featurestores/{FEATURESTORE_ID}/entityTypes/track\n",
+    "#     featureId: track_loudness_can\n",
+    "# - valueKey: track_speechiness_can\n",
+    "#   featureSource:\n",
+    "#     entityIdKey: track_uri_can\n",
+    "#     entityType: projects/{PROJECT_NUMBER}/locations/{REGION}/featurestores/{FEATURESTORE_ID}/entityTypes/track\n",
+    "#     featureId: track_speechiness_can\n",
+    "# - valueKey: track_acousticness_can\n",
+    "#   featureSource:\n",
+    "#     entityIdKey: track_uri_can\n",
+    "#     entityType: projects/{PROJECT_NUMBER}/locations/{REGION}/featurestores/{FEATURESTORE_ID}/entityTypes/track\n",
+    "#     featureId: track_acousticness_can\n",
+    "# - valueKey: track_instrumentalness_can\n",
+    "#   featureSource:\n",
+    "#     entityIdKey: track_uri_can\n",
+    "#     entityType: projects/{PROJECT_NUMBER}/locations/{REGION}/featurestores/{FEATURESTORE_ID}/entityTypes/track\n",
+    "#     featureId: track_instrumentalness_can\n",
+    "# - valueKey: track_liveness_can\n",
+    "#   featureSource:\n",
+    "#     entityIdKey: track_uri_can\n",
+    "#     entityType: projects/{PROJECT_NUMBER}/locations/{REGION}/featurestores/{FEATURESTORE_ID}/entityTypes/track\n",
+    "#     featureId: track_liveness_can\n",
+    "# - valueKey: track_valence_can\n",
+    "#   featureSource:\n",
+    "#     entityIdKey: track_uri_can\n",
+    "#     entityType: projects/{PROJECT_NUMBER}/locations/{REGION}/featurestores/{FEATURESTORE_ID}/entityTypes/track\n",
+    "#     featureId: track_valence_can\n",
+    "# - valueKey: track_tempo_can\n",
+    "#   featureSource:\n",
+    "#     entityIdKey: track_uri_can\n",
+    "#     entityType: projects/{PROJECT_NUMBER}/locations/{REGION}/featurestores/{FEATURESTORE_ID}/entityTypes/track\n",
+    "#     featureId: track_tempo_can\n",
+    "# - valueKey: track_time_signature_can\n",
+    "#   featureSource:\n",
+    "#     entityIdKey: track_uri_can\n",
+    "#     entityType: projects/{PROJECT_NUMBER}/locations/{REGION}/featurestores/{FEATURESTORE_ID}/entityTypes/track\n",
+    "#     featureId: track_time_signature_can\n",
+    "# \"\"\"\n",
+    "\n",
+    "# feature_fetch_config = FEATURE_FETCH_CONFIG_TEMPLATE.format(\n",
+    "#     PROJECT_NUMBER=PROJECT_NUM, REGION=REGION, FEATURESTORE_ID=FEATURESTORE_ID\n",
+    "# )\n",
+    "\n",
+    "# # with open(f\"{PRED_FEATURE_FETCH_YAML}\", \"w\") as f:\n",
+    "# #     f.write(feature_fetch_config)\n",
+    "\n",
+    "# with open(f\"{PRED_FEATURE_FETCH_YAML_v3}\", \"w\") as f:\n",
+    "#     f.write(feature_fetch_config)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1266a33f-7f20-4f5d-bb73-800c25bcabf5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "QUERY_TOWER_FS_PRED_REQ = {'album_name_pl': ['Odilara',\n",
+    "                   'Capoeira Electronica',\n",
+    "                   'Capoeira Ultimate',\n",
+    "                   'Festa Popular',\n",
+    "                   'Capoeira Electronica'],\n",
+    " 'album_uri_pl': ['spotify:album:4Y8RfvZzCiApBCIZswj9Ry',\n",
+    "                  'spotify:album:2FsSSHGt8JM0JgRy6ZX3kR',\n",
+    "                  'spotify:album:55HHBqZ2SefPeaENOgWxYK',\n",
+    "                  'spotify:album:150L1V6UUT7fGUI3PbxpkE',\n",
+    "                  'spotify:album:2FsSSHGt8JM0JgRy6ZX3kR'],\n",
+    " 'artist_genres_pl': ['samba moderno',\n",
+    "                      'capoeira',\n",
+    "                      'capoeira',\n",
+    "                      'NONE',\n",
+    "                      'capoeira'],\n",
+    " 'artist_name_pl': ['Odilara',\n",
+    "                    'Capoeira Experience',\n",
+    "                    'Denis Porto',\n",
+    "                    'Zambe',\n",
+    "                    'Capoeira Experience'],\n",
+    " 'artist_pop_pl': [4.0, 24.0, 2.0, 0.0, 24.0],\n",
+    " 'artist_uri_pl': ['spotify:artist:72oameojLOPWYB7nB8rl6c',\n",
+    "                   'spotify:artist:5SKEXbgzIdRl3gQJ23CnUP',\n",
+    "                   'spotify:artist:67p5GMYQZOgaAfx1YyttQk',\n",
+    "                   'spotify:artist:4fH3OXCRcPsaHFE5KhgqZS',\n",
+    "                   'spotify:artist:5SKEXbgzIdRl3gQJ23CnUP'],\n",
+    " 'artists_followers_pl': [316.0, 5170.0, 448.0, 19.0, 5170.0],\n",
+    " 'duration_ms_songs_pl': [234612.0, 226826.0, 203480.0, 287946.0, 271920.0],\n",
+    " 'num_pl_albums_new': 9.0,\n",
+    " 'num_pl_artists_new': 5.0,\n",
+    " 'num_pl_songs_new': 85.0,\n",
+    " 'pl_collaborative_src': 'false',\n",
+    " 'pl_duration_ms_new': 17971314.0,\n",
+    " 'pl_name_src': 'Capoeira',\n",
+    " 'track_acousticness_pl': [0.238, 0.105, 0.0242, 0.125, 0.304],\n",
+    " 'track_danceability_pl': [0.703, 0.712, 0.806, 0.529, 0.821],\n",
+    " 'track_energy_pl': [0.743, 0.41, 0.794, 0.776, 0.947],\n",
+    " 'track_instrumentalness_pl': [4.84e-06, 0.43, 0.000742, 0.401, 0.00507],\n",
+    " 'track_key_pl': ['5', '0', '1', '10', '10'],\n",
+    " 'track_liveness_pl': [0.128, 0.0725, 0.191, 0.105, 0.0552],\n",
+    " 'track_loudness_pl': [-8.638, -8.754, -9.084, -7.04, -6.694],\n",
+    " 'track_mode_pl': ['0', '1', '1', '0', '1'],\n",
+    " 'track_name_pl': ['O Telefone Tocou Novamente',\n",
+    "                   'Bem Devagar - Studio',\n",
+    "                   'Angola Dream',\n",
+    "                   'Janaina',\n",
+    "                   'Louco Berimbau - Studio'],\n",
+    " 'track_pop_pl': [5.0, 1.0, 0.0, 0.0, 1.0],\n",
+    " 'track_speechiness_pl': [0.0367, 0.0272, 0.0407, 0.132, 0.0734],\n",
+    " 'track_tempo_pl': [100.039, 89.089, 123.999, 119.963, 119.214],\n",
+    " 'track_time_signature_pl': ['4', '4', '4', '4', '4'],\n",
+    " 'track_uri_can': 'spotify:track:0tlhK4OvpHCYpReTABvKFb',\n",
+    " 'track_uri_pl': ['spotify:track:1pQkOdcTDfLr84TDCrmGy7',\n",
+    "                  'spotify:track:39grEDsAHAjmo2QFo4G8D9',\n",
+    "                  'spotify:track:5vxSLdJXqbKYH487YO8LSL',\n",
+    "                  'spotify:track:6T9GbmZ6voDM4aTBsG5VDh',\n",
+    "                  'spotify:track:7ELt9eslVvWo276pX2garN'],\n",
+    " 'track_valence_pl': [0.966, 0.667, 0.696, 0.876, 0.655]}\n",
+    "\n",
+    "print(QUERY_TOWER_FS_PRED_REQ)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "90a12347-7746-40c1-aeea-85ebd8752b68",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "CREATE OR REPLACE TABLE\n",
+    "  `hybrid-vertex.spotify_e2e_test.v2_falt_valid_last_5_fs` AS (\n",
+    "  SELECT\n",
+    "    * EXCEPT(track_time_signature_can,\n",
+    "      track_time_signature_pl,\n",
+    "      track_mode_can,\n",
+    "      track_mode_pl,\n",
+    "      track_key_can,\n",
+    "      track_key_pl),\n",
+    "    CAST(track_mode_can AS STRING) AS track_mode_can,\n",
+    "    CAST(track_time_signature_can AS STRING) AS track_time_signature_can,\n",
+    "    CAST(track_key_can AS STRING) AS track_key_can,\n",
+    "    CURRENT_TIMESTAMP() AS feature_ts\n",
+    "  FROM\n",
+    "    `hybrid-vertex.spotify_e2e_test.v2_train_flatten_valid_last_5` \n",
+    "  )"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "12cd1982-12a1-4c43-88fd-aac97757f8b6",
+   "metadata": {},
+   "source": [
+    "## Simulate endpoint traffic\n",
+    "\n",
+    "[src](https://colab.sandbox.google.com/drive/1x4EoCoEgzqVPY-MGfD9KzkBMN1gxl8ET#scrollTo=Hw42_E7qcYG-)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b141f5d1-781b-46c6-8eb8-b31918be0977",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Setup clients\n",
+    "aiplatform.init(project=PROJECT_ID,location=REGION)\n",
+    "client = bigquery.Client(project=PROJECT_ID, location=REGION) # TODO: paramaterize\n",
+    "\n",
+    "# Find Vertex Endpoint\n",
+    "# logging.info(\"vertex_endpoint_uri: %s\", vertex_endpoint_uri)\n",
+    "\n",
+    "endpoint = aiplatform.Endpoint(_vertex_endpoint_uri)\n",
+    "endpoint_uri = endpoint.resource_name\n",
+    "# logging.info(\"endpoint: %s\", endpoint)\n",
+    "# logging.info(\"endpoint_uri: %s\", endpoint_uri) # should match vertex_endpoint_uri\n",
+    "\n",
+    "# ==============================================================================\n",
+    "# Create skewed traffic to trigger alerting\n",
+    "# ==============================================================================\n",
+    "print(\"Creating Skewed Data...\")\n",
+    "# BQ_TRAINING_DATA = f'bq://{PROJECT_ID}.{BQ_DATASET_NAME}.{BQ_TABLE_NAME_TRAIN_SPLIT}'\n",
+    "sql_script = f'''\n",
+    "SELECT * \n",
+    "FROM {PROJECT_ID}.{BQ_DATASET_NAME}.{BQ_TABLE_NAME_TRAIN_SPLIT}\n",
+    "'''\n",
+    "\n",
+    "data = client.query(sql_script).to_dataframe()\n",
+    "\n",
+    "NUMERIC_FEATURES = {\n",
+    "    'crim': (data['crim'].mean(), data['crim'].std()),\n",
+    "    'zn': (data['zn'].mean(), data['zn'].std()),\n",
+    "    'indus': (data['indus'].mean(), data['indus'].std()),\n",
+    "    'chas': (data['chas'].mean(), data['chas'].std()),\n",
+    "    'nox': (data['nox'].mean(), data['nox'].std()),\n",
+    "    'rm': (data['rm'].mean(), data['rm'].std()),\n",
+    "    'age': (data['age'].mean(), data['age'].std()),\n",
+    "    'dis': (data['dis'].mean(), data['dis'].std()),\n",
+    "    'rad': (data['rad'].mean(), data['rad'].std()),\n",
+    "    'tax': (data['tax'].mean(), data['tax'].std()),\n",
+    "    'ptratio': (data['ptratio'].mean(), data['ptratio'].std()),\n",
+    "    'b': (data['b'].mean(), data['b'].std()),\n",
+    "    'lstat': (data['lstat'].mean(), data['lstat'].std()),\n",
+    "}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a8d92f7c-02d0-4a3a-9892-9e60cf1b0a36",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def monitoring_test(endpoint, count=1, seconds=3, perturb_num={}):\n",
+    "    \n",
+    "    for i in range(0, count):\n",
+    "        instance = {}\n",
+    "    \n",
+    "        for key, stats in NUMERIC_FEATURES.items():\n",
+    "            mean = stats[0]\n",
+    "            std = stats[1]\n",
+    "            if key in perturb_num.keys():\n",
+    "                mean = perturb_num[key][0](mean)\n",
+    "                std = perturb_num[key][1](std)\n",
+    "            \n",
+    "            value = float(np.random.normal(mean, std, 1))\n",
+    "            value = max(value, 0)\n",
+    "            instance[key] = value\n",
+    "        \n",
+    "        print(\"Prediction request:\", [[instance[k] for k in instance]])\n",
+    "        predictions = endpoint.predict(instances= [[instance[k] for k in instance]])\n",
+    "        print(predictions)\n",
+    "        # explains = endpoint.explain(instances=list(instance.values()))\n",
+    "          # my_explanations = response.explanations\n",
+    "        # explanations = endpoint.explain(instances=list(instance.values()))\n",
+    "        # print(explanations)\n",
+    "        time.sleep(seconds)\n",
+    "\n",
+    "perturb_num = {\n",
+    "    'crim': (lambda x: x * 3, lambda x: x / 3),\n",
+    "    'zn': (lambda x: x * 3, lambda x: x / 3),\n",
+    "    'indus': (lambda x: x * 3, lambda x: x / 3),\n",
+    "    'chas': (lambda x: x * 3, lambda x: x / 3),\n",
+    "    'nox': (lambda x: x * 3, lambda x: x / 3),\n",
+    "    'rm': (lambda x: x * 3, lambda x: x / 3),\n",
+    "    'age': (lambda x: x * 3, lambda x: x / 3),\n",
+    "    'dis': (lambda x: x * 3, lambda x: x / 3),\n",
+    "    'rad': (lambda x: x * 3, lambda x: x / 3),\n",
+    "    'tax': (lambda x: x * 3, lambda x: x / 3),\n",
+    "    'ptratio': (lambda x: x * 3, lambda x: x / 3),\n",
+    "    'b': (lambda x: x * 3, lambda x: x / 3),\n",
+    "    'lstat':(lambda x: x * 3, lambda x: x / 3),\n",
+    "}\n",
+    "\n",
+    "print(\"Simulating Skewed Traffic...\\n\")\n",
+    "start=2\n",
+    "end=6\n",
+    "for multiplier in range(start, end+1):\n",
+    "  test_time=1200\n",
+    "  tests_per_second = 1\n",
+    "  sleep_time = 1 / tests_per_second\n",
+    "  iterations = test_time * tests_per_second\n",
+    "  perturb_num = {\n",
+    "      'crim': (lambda x: x * multiplier, lambda x: x / multiplier),\n",
+    "      'zn': (lambda x: x * multiplier, lambda x: x / multiplier),\n",
+    "      'indus': (lambda x: x * multiplier, lambda x: x / multiplier),\n",
+    "      'chas': (lambda x: x * multiplier, lambda x: x / multiplier),\n",
+    "      'nox': (lambda x: x * multiplier, lambda x: x / multiplier),\n",
+    "      'rm': (lambda x: x * multiplier, lambda x: x / multiplier),\n",
+    "      'age': (lambda x: x * multiplier, lambda x: x / multiplier),\n",
+    "      'dis': (lambda x: x * multiplier, lambda x: x / multiplier),\n",
+    "      'rad': (lambda x: x * multiplier, lambda x: x / multiplier),\n",
+    "      'tax': (lambda x: x * multiplier, lambda x: x / multiplier),\n",
+    "      'ptratio': (lambda x: x * multiplier, lambda x: x / multiplier),\n",
+    "      'b': (lambda x: x * multiplier, lambda x: x / multiplier),\n",
+    "      'lstat':(lambda x: x * multiplier, lambda x: x / multiplier),\n",
+    "  }\n",
+    "  print(f\"Test {multiplier} out of {end+1}: {iterations} prediction requests in {test_time} seconds; traffic perturbed by a factor of {multiplier}\")\n",
+    "  monitoring_test(endpoint, count=iterations, seconds=sleep_time, perturb_num=perturb_num)\n",
+    "\n",
+    "# print(f\"First test: {skewed_samples_count} requests, {sim_traffic_interval} second intervals\")\n",
+    "# monitoring_test(endpoint, count=iterations, perturb_num=perturb_num, seconds=sim_traffic_interval)\n",
+    "\n",
+    "print(\"Simulated Traffic Complete \\n\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c113e92f-4662-4b1a-954a-7bffdc282fe3",
+   "metadata": {},
+   "source": [
+    "# MAtching Engine CRUD"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a8a3f638-6a29-4d5c-91cb-a8345b17b3c8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "mengine_ann = MatchingEngineCRUD(\n",
+    "    project_id=PROJECT_ID \n",
+    "    , project_num=PROJECT_NUM\n",
+    "    , region=LOCATION \n",
+    "    , index_type = \"ann\"       # \"ann\" | \"bf\"\n",
+    "    , index_name=DISPLAY_NAME\n",
+    "    , vpc_network_name=VPC_NETWORK_FULL\n",
+    ")\n",
+    "\n",
+    "mengine_ann.index_name\n",
+    "\n",
+    "\n",
+    "mengine_bf = MatchingEngineCRUD(\n",
+    "    project_id=PROJECT_ID \n",
+    "    , project_num=PROJECT_NUM\n",
+    "    , region=LOCATION \n",
+    "    , index_type = \"bf\"         # \"ann\" | \"bf\"\n",
+    "    , index_name=BF_DISPLAY_NAME\n",
+    "    , vpc_network_name=VPC_NETWORK_FULL\n",
+    ")\n",
+    "\n",
+    "mengine_bf.index_name"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4854ef17-1d93-46ff-b117-eb41195543a2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "start = time.time()\n",
+    "# create ME index\n",
+    "tree_ah_index = mengine_ann.create_index(\n",
+    "    f\"{EMBEDDINGS_INITIAL_URI}\"\n",
+    "    , int(DIMENSIONS)\n",
+    ")\n",
+    "\n",
+    "end = time.time()\n",
+    "print(f\"elapsed time: {end - start}\")\n",
+    "\n",
+    "if tree_ah_index:\n",
+    "    print(tree_ah_index.name)\n",
+    "    \n",
+    "start = time.time()\n",
+    "# create ME index\n",
+    "brute_force_index = mengine_bf.create_index(\n",
+    "    f\"{EMBEDDINGS_INITIAL_URI}\"\n",
+    "    , int(DIMENSIONS)\n",
+    ")\n",
+    "\n",
+    "end = time.time()\n",
+    "print(f\"elapsed time: {end - start}\")\n",
+    "\n",
+    "if brute_force_index:\n",
+    "    print(brute_force_index.name)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "29848e50-9dd6-4f6d-9c66-5296f82e99c3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "start = time.time()\n",
+    "\n",
+    "ann_index_endpoint=mengine_ann.create_index_endpoint(\n",
+    "    endpoint_name=ANN_ENDPOINT_DISPLAY_NAME\n",
+    "    , network=VPC_NETWORK_FULL\n",
+    ")\n",
+    "\n",
+    "EXISTING_INDEX_ENDPOINT = f'projects/{PROJECT_NUM}/locations/{REGION}/indexEndpoints/5661297410488401920'\n",
+    "ann_index_endpoint = aiplatform.MatchingEngineIndexEndpoint(EXISTING_INDEX_ENDPOINT)\n",
+    "\n",
+    "end = time.time()\n",
+    "print(f\"elapsed time: {end - start}\")\n",
+    "\n",
+    "\n",
+    "start = time.time()\n",
+    "\n",
+    "bf_index_endpoint=mengine_bf.create_index_endpoint(\n",
+    "    endpoint_name=BF_ENDPOINT_DISPLAY_NAME\n",
+    "    , network=VPC_NETWORK_FULL\n",
+    ")\n",
+    "\n",
+    "# EXISTING = '1049611392061014016'\n",
+    "\n",
+    "end = time.time()\n",
+    "print(f\"elapsed time: {end - start}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1cea3032-68f7-4840-b5db-12fcfb2b5671",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# if CREATE_NEW_ASSETS:\n",
+    "    \n",
+    "ann_index_endpoint = mengine_ann.deploy_index(\n",
+    "    index_name = DISPLAY_NAME\n",
+    "    , endpoint_name = ANN_ENDPOINT_DISPLAY_NAME\n",
+    "    , machine_type = \"n1-standard-16\"\n",
+    "    , min_replica_count = 2\n",
+    "    , max_replica_count = 2\n",
+    ")\n",
+    "\n",
+    "    deployed_ann_index = my_ann_index_endpoint.deploy_index(\n",
+    "        index=tree_ah_index, \n",
+    "        deployed_index_id=DEPLOYED_ANN_INDEX_ID\n",
+    "    )\n",
+    "\n",
+    "# ann_index_endpoint\n",
+    "\n",
+    "\n",
+    "# if CREATE_NEW_ASSETS:\n",
+    "    \n",
+    "bf_index_endpoint = mengine_bf.deploy_index(\n",
+    "    index_name = BF_DISPLAY_NAME\n",
+    "    , endpoint_name = BF_ENDPOINT_DISPLAY_NAME\n",
+    "    , machine_type = \"n1-standard-16\"\n",
+    "    , min_replica_count = 2\n",
+    "    , max_replica_count = 2\n",
+    ")\n",
+    "\n",
+    "if bf_index_endpoint:\n",
+    "    print(f\"Index endpoint resource name: {bf_index_endpoint.name}\")\n",
+    "    print(f\"Index endpoint VPC network name: {bf_index_endpoint.network}\")\n",
+    "    print(f\"Deployed indexes on the index endpoint:\")\n",
+    "    for d in bf_index_endpoint.deployed_indexes:\n",
+    "        print(f\"    {d.id}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6c2b62a1-ac4c-4d6f-ab90-10b1a316b904",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ME_ANN_INDEX_RESOURCE_NAME, ME_ANN_INDEX_ENDPOINT_ID = mengine_ann.get_index_and_endpoint()\n",
+    "ME_ANN_INDEX_ID=ME_ANN_INDEX_RESOURCE_NAME.split(\"/\")[5]\n",
+    "\n",
+    "print(f\"ME_ANN_INDEX_RESOURCE_NAME  = {ME_ANN_INDEX_RESOURCE_NAME}\")\n",
+    "print(f\"ME_ANN_INDEX_ENDPOINT_ID    = {ME_ANN_INDEX_ENDPOINT_ID}\")\n",
+    "print(f\"ME_ANN_INDEX_ID             = {ME_ANN_INDEX_ID}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d84b9a1b-2a34-4344-b788-3195d74f0116",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ME_BF_INDEX_RESOURCE_NAME, ME_BF_INDEX_ENDPOINT_ID = mengine_bf.get_index_and_endpoint()\n",
+    "ME_BF_INDEX_ID=ME_BF_INDEX_RESOURCE_NAME.split(\"/\")[5]\n",
+    "\n",
+    "print(f\"ME_BF_INDEX_RESOURCE_NAME  = {ME_BF_INDEX_RESOURCE_NAME}\")\n",
+    "print(f\"ME_BF_INDEX_ENDPOINT_ID    = {ME_BF_INDEX_ENDPOINT_ID}\")\n",
+    "print(f\"ME_BF_INDEX_ID             = {ME_BF_INDEX_ID}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "bb6ee8af-1efa-490a-b6de-ebdbf0e37dec",
+   "metadata": {},
+   "source": [
+    "# Candidate Generation"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "08319e3e-3ba7-4067-8e33-da083ed7a91d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# ====================================================== #\n",
+    "\n",
+    "# PATH_TO_INDEX_DIR = 'customtest5-trainerv3/run-20230321-005425' # 50e-8m-128d\n",
+    "# PATH_TO_INDEX_DIR = 'test-e2e-v1-trainerv4/run-20230410-232135' #/model-dir/candidate_model'\n",
+    "# PATH_TO_INDEX_DIR = '8m-tfrs-v100-jtv15/run-20230125-205451' # 50e-8m-128d\n",
+    "# PATH_TO_INDEX_DIR = 'customtest5-trainerv3/run-20230321-005425' # 50e-8m-128d"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7b780ed8-a9b2-4e14-b07e-186a89a3546d",
+   "metadata": {},
+   "source": [
+    "# Requirements txt"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3ea0e970-e1ef-402d-bef1-0c7eb6f2d8fd",
+   "metadata": {},
+   "source": [
+    "## Retreival"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "84ba27ff-e272-4274-b460-18ad441e7d60",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# tensorboard==2.13.0\n",
+    "# tensorboard-data-server==0.7.1\n",
+    "# tensorboard-plugin-profile==2.13.1\n",
+    "# tensorboard-plugin-wit==1.8.1\n",
+    "# tensorboardX==2.6\n",
+    "# tensorflow-io==0.27.0"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "aa3e28c5-1831-4aa0-b876-b4a73e2f4f4b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# tensorboard==2.10.1\n",
+    "# # tensorboard==2.11.2 \n",
+    "# tensorboard-data-server==0.6.1\n",
+    "# tensorboard-plugin-profile==2.11.1\n",
+    "# tensorflow-io==0.27.0"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4ea4587b-05c5-4d0a-8793-1c2139e20c83",
+   "metadata": {},
+   "source": [
+    "# Docker debugging"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6d2d7514-1a31-4471-81f0-ce7f9ebfa716",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# sudo ln -s /usr/bin/stat /usr/bin/lstat"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4796cf87-75d0-4cad-aa0a-1cde4cfc5a73",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# *__pycache__\n",
+    "# *cpython-37.pyc\n",
+    "# spotipy_secret_creds.py\n",
+    "# custom_pipeline_spec.json\n",
+    "# /WIP/*\n",
+    "# beam_candidates/*\n",
+    "# beam_training/*\n",
+    "# learning/*\n",
+    "# src/vocab_pipes/*\n",
+    "# src/train_pipes/*\n",
+    "# src/feature_pipes/*\n",
+    "# custom_track_meta_pipeline_spec.json\n",
+    "# candidate_embeddings.json\n",
+    "# pip_freeze.txt\n",
+    "# README.md\n",
+    "# .gitignore\n",
+    "# .DS_Store\n",
+    "# ARCHIVED/*\n",
+    "# *.yaml\n",
+    "# **/__pycache__\n",
+    "# **/*.pyc\n",
+    "# __pycache__/train_pipes/*"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b7eb19fc-bda8-496a-b7d0-6e7763a76172",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# import modules\n",
+    "\n",
+    "# import train_config as cfg\n",
+    "# import train_utils\n",
+    "\n",
+    "# from . import train_config as cfg\n",
+    "# from . import train_utils\n",
+    "\n",
+    "# from src.two_tower_jt import train_config as cfg\n",
+    "# from src.two_tower_jt import train_utils\n",
+    "\n",
+    "# from .train_utils import (\n",
+    "#     full_parse, \n",
+    "#     get_train_strategy, \n",
+    "#     _is_chief, \n",
+    "#     get_arch_from_string, \n",
+    "#     tf_if_null_return_zero, \n",
+    "#     get_buckets_20, \n",
+    "#     upload_blob\n",
+    "# )"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a919abd6-5434-443e-9954-ee4bec4af8fb",
+   "metadata": {},
+   "source": [
+    "# Ranker"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6492f12d-c2c7-4fd9-8cf8-1c346eafb455",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# orirignal - works\n",
+    "\n",
+    "# ========================================\n",
+    "# Ranking - compute predictions\n",
+    "# ========================================\n",
+    "\n",
+    "class TheRankingModel(tfrs.models.Model):\n",
+    "\n",
+    "    def __init__(\n",
+    "        self\n",
+    "        , vocab_dict\n",
+    "        , layer_sizes: List = [512, 256, 128]\n",
+    "        , embedding_dim: int = 128\n",
+    "        , projection_dim: int = 50\n",
+    "        , seed: int = 1234\n",
+    "        , use_dropout: bool = True\n",
+    "        , dropout_rate: float = 0.33\n",
+    "        , max_tokens: int = 20000\n",
+    "    ):\n",
+    "        super().__init__()\n",
+    "        self.ranking_model: tf.keras.Model = RankingAudio(\n",
+    "            seed = seed\n",
+    "            , vocab_dict = vocab_dict\n",
+    "            , max_tokens = max_tokens\n",
+    "            , use_dropout = use_dropout\n",
+    "            , layer_sizes = layer_sizes\n",
+    "            , dropout_rate = dropout_rate\n",
+    "            , embedding_dim = embedding_dim\n",
+    "            , projection_dim = projection_dim\n",
+    "        )\n",
+    "\n",
+    "        self.task: tf.keras.layers.Layer = tfrs.tasks.Ranking(\n",
+    "            loss = tf.keras.losses.MeanSquaredError(\n",
+    "                reduction=tf.keras.losses.Reduction.SUM\n",
+    "            )\n",
+    "            , metrics = [\n",
+    "                tf.keras.metrics.RootMeanSquaredError(name=\"rmse_metric\")\n",
+    "                # , tfr.keras.metrics.NDCGMetric(\n",
+    "                #     name=\"ndcg_metric\"\n",
+    "                #     , ragged = True\n",
+    "                # )\n",
+    "            ]\n",
+    "        )\n",
+    "\n",
+    "    def call(self, features: Dict[str, tf.Tensor]) -> tf.Tensor:\n",
+    "        return self.ranking_model(features)\n",
+    "\n",
+    "    def compute_loss(self, features: Dict[Text, tf.Tensor], training=False) -> tf.Tensor:\n",
+    "        \n",
+    "        labels = features.pop(\"candidate_rank\")\n",
+    "\n",
+    "        rating_predictions = self(features)\n",
+    "\n",
+    "        # The task computes the loss and the metrics.\n",
+    "        return self.task(\n",
+    "            labels=labels\n",
+    "            # , predictions=rating_predictions\n",
+    "            , predictions=tf.squeeze(rating_predictions, axis=-1),\n",
+    "        )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "bfd5404d-64be-43bd-9d07-430c5b093bbc",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# # # ===================================================\n",
+    "# # # baseline ranker feats\n",
+    "# # # ===================================================\n",
+    "# # train_dataset = train_dataset.interleave(\n",
+    "# #     train_utils.full_parse,\n",
+    "# #     cycle_length=tf.data.AUTOTUNE, \n",
+    "# #     num_parallel_calls=tf.data.AUTOTUNE,\n",
+    "# #     deterministic=False,\n",
+    "# # ).map(\n",
+    "# #     feature_sets.parse_tfrecord, \n",
+    "# #     num_parallel_calls=tf.data.AUTOTUNE\n",
+    "# # ).batch(\n",
+    "# #     batch_size \n",
+    "# # ).prefetch(\n",
+    "# #     tf.data.AUTOTUNE,\n",
+    "# # ).with_options(\n",
+    "# #     options\n",
+    "# # )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e1d9f33c-fd95-429f-9758-cecb99d313d0",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "586b48d2-c6cb-4461-95bd-e4d580144896",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9c62076e-8acb-4361-95c2-549a29103a3e",
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {