T4C Schema.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "c015c926",
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "Descriptors cannot not be created directly.\nIf this call came from a _pb2.py file, your generated code is out of date and must be regenerated with protoc >= 3.19.0.\nIf you cannot immediately regenerate your protos, some other possible workarounds are:\n 1. Downgrade the protobuf package to 3.20.x or lower.\n 2. Set PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python (but this will use pure-Python parsing and will be much slower).\n\nMore information: https://developers.google.com/protocol-buffers/docs/news/2022-05-06#python-updates",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[13], line 5\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mshutil\u001b[39;00m\n\u001b[1;32m      4\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mnumpy\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mnp\u001b[39;00m\n\u001b[0;32m----> 5\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mnvtabular\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mops\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;241m*\u001b[39m\n\u001b[1;32m      6\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mmerlin\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mschema\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mtags\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Tags\n\u001b[1;32m      8\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mnvtabular\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mnvt\u001b[39;00m\n",
      "File \u001b[0;32m~/miniconda3/envs/rl/lib/python3.8/site-packages/nvtabular/__init__.py:19\u001b[0m\n\u001b[1;32m     16\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mwarnings\u001b[39;00m\n\u001b[1;32m     18\u001b[0m \u001b[38;5;66;03m# this is to mimic the previous API, but all of these should probably be removed\u001b[39;00m\n\u001b[0;32m---> 19\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mmerlin\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdag\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mgraph\u001b[39;00m  \u001b[38;5;66;03m# noqa\u001b[39;00m\n\u001b[1;32m     20\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mmerlin\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m io\n\u001b[1;32m     21\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mmerlin\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m dispatch, utils  \u001b[38;5;66;03m# noqa\u001b[39;00m\n",
      "File \u001b[0;32m~/miniconda3/envs/rl/lib/python3.8/site-packages/merlin/dag/__init__.py:18\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;66;03m#\u001b[39;00m\n\u001b[1;32m      2\u001b[0m \u001b[38;5;66;03m# Copyright (c) 2022, NVIDIA CORPORATION.\u001b[39;00m\n\u001b[1;32m      3\u001b[0m \u001b[38;5;66;03m#\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     16\u001b[0m \n\u001b[1;32m     17\u001b[0m \u001b[38;5;66;03m# flake8: noqa\u001b[39;00m\n\u001b[0;32m---> 18\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mmerlin\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdag\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mbase_operator\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m BaseOperator, DataFormats, Supports\n\u001b[1;32m     19\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mmerlin\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdag\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mgraph\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Graph\n\u001b[1;32m     20\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mmerlin\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdag\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mnode\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Node, iter_nodes, postorder_iter_nodes, preorder_iter_nodes\n",
      "File \u001b[0;32m~/miniconda3/envs/rl/lib/python3.8/site-packages/merlin/dag/base_operator.py:23\u001b[0m\n\u001b[1;32m     21\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mmerlin\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdag\u001b[39;00m\n\u001b[1;32m     22\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mmerlin\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mprotocols\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Transformable\n\u001b[0;32m---> 23\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mmerlin\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdag\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mselector\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m ColumnSelector\n\u001b[1;32m     24\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mmerlin\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mschema\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m ColumnSchema, Schema\n\u001b[1;32m     27\u001b[0m \u001b[38;5;28;01mclass\u001b[39;00m \u001b[38;5;21;01mSupports\u001b[39;00m(Flag):\n",
      "File \u001b[0;32m~/miniconda3/envs/rl/lib/python3.8/site-packages/merlin/dag/selector.py:19\u001b[0m\n\u001b[1;32m     16\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtyping\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m List, Union\n\u001b[1;32m     18\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mmerlin\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdag\u001b[39;00m\n\u001b[0;32m---> 19\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mmerlin\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mschema\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Tags\n\u001b[1;32m     22\u001b[0m \u001b[38;5;28;01mclass\u001b[39;00m \u001b[38;5;21;01mColumnSelector\u001b[39;00m:\n\u001b[1;32m     23\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"A ColumnSelector describes a group of columns to be transformed by Operators in a\u001b[39;00m\n\u001b[1;32m     24\u001b[0m \u001b[38;5;124;03m    Graph. Operators can be applied to the selected columns by shifting (>>) operators\u001b[39;00m\n\u001b[1;32m     25\u001b[0m \u001b[38;5;124;03m    on to the ColumnSelector, which returns a new Node with the transformations\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     39\u001b[0m \u001b[38;5;124;03m        at-least-one of the tags provided will be considered.\u001b[39;00m\n\u001b[1;32m     40\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n",
      "File \u001b[0;32m~/miniconda3/envs/rl/lib/python3.8/site-packages/merlin/schema/__init__.py:18\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;66;03m#\u001b[39;00m\n\u001b[1;32m      2\u001b[0m \u001b[38;5;66;03m# Copyright (c) 2022, NVIDIA CORPORATION.\u001b[39;00m\n\u001b[1;32m      3\u001b[0m \u001b[38;5;66;03m#\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     16\u001b[0m \n\u001b[1;32m     17\u001b[0m \u001b[38;5;66;03m# flake8: noqa\u001b[39;00m\n\u001b[0;32m---> 18\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mmerlin\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mschema\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mschema\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m ColumnSchema, Schema\n\u001b[1;32m     19\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mmerlin\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mschema\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mtags\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Tags, TagSet, TagsType\n",
      "File \u001b[0;32m~/miniconda3/envs/rl/lib/python3.8/site-packages/merlin/schema/schema.py:22\u001b[0m\n\u001b[1;32m     18\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtyping\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Dict, List, Optional, Text, Tuple, Union\n\u001b[1;32m     20\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mpandas\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mpd\u001b[39;00m\n\u001b[0;32m---> 22\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mmerlin\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdtypes\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mmd\u001b[39;00m\n\u001b[1;32m     23\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mmerlin\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdtypes\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m DType\n\u001b[1;32m     24\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mmerlin\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdtypes\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mshape\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Shape\n",
      "File \u001b[0;32m~/miniconda3/envs/rl/lib/python3.8/site-packages/merlin/dtypes/__init__.py:18\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;66;03m#\u001b[39;00m\n\u001b[1;32m      2\u001b[0m \u001b[38;5;66;03m# Copyright (c) 2022, NVIDIA CORPORATION.\u001b[39;00m\n\u001b[1;32m      3\u001b[0m \u001b[38;5;66;03m#\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     16\u001b[0m \n\u001b[1;32m     17\u001b[0m \u001b[38;5;66;03m# flake8: noqa\u001b[39;00m\n\u001b[0;32m---> 18\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mmerlin\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdtypes\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m aliases, mappings\n\u001b[1;32m     19\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mmerlin\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdtypes\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01maliases\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;241m*\u001b[39m\n\u001b[1;32m     20\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mmerlin\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdtypes\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mbase\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m DType\n",
      "File \u001b[0;32m~/miniconda3/envs/rl/lib/python3.8/site-packages/merlin/dtypes/mappings/__init__.py:18\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;66;03m#\u001b[39;00m\n\u001b[1;32m      2\u001b[0m \u001b[38;5;66;03m# Copyright (c) 2022, NVIDIA CORPORATION.\u001b[39;00m\n\u001b[1;32m      3\u001b[0m \u001b[38;5;66;03m#\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     16\u001b[0m \n\u001b[1;32m     17\u001b[0m \u001b[38;5;66;03m# flake8: noqa\u001b[39;00m\n\u001b[0;32m---> 18\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mmerlin\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdtypes\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mmappings\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m cudf, merlin, numpy, pandas, python, tf, torch, triton\n",
      "File \u001b[0;32m~/miniconda3/envs/rl/lib/python3.8/site-packages/merlin/dtypes/mappings/tf.py:21\u001b[0m\n\u001b[1;32m     18\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mmerlin\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdtypes\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mregistry\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m _dtype_registry\n\u001b[1;32m     20\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m---> 21\u001b[0m     \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m dtypes \u001b[38;5;28;01mas\u001b[39;00m tf_dtypes\n\u001b[1;32m     23\u001b[0m     tf_dtypes \u001b[38;5;241m=\u001b[39m DTypeMapping(\n\u001b[1;32m     24\u001b[0m         {\n\u001b[1;32m     25\u001b[0m             \u001b[38;5;66;03m# Unsigned Integer\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     45\u001b[0m         ),\n\u001b[1;32m     46\u001b[0m     )\n\u001b[1;32m     47\u001b[0m     _dtype_registry\u001b[38;5;241m.\u001b[39mregister(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtf\u001b[39m\u001b[38;5;124m\"\u001b[39m, tf_dtypes)\n",
      "File \u001b[0;32m~/miniconda3/envs/rl/lib/python3.8/site-packages/tensorflow/__init__.py:37\u001b[0m\n\u001b[1;32m     34\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01msys\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01m_sys\u001b[39;00m\n\u001b[1;32m     35\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mtyping\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01m_typing\u001b[39;00m\n\u001b[0;32m---> 37\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mpython\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mtools\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m module_util \u001b[38;5;28;01mas\u001b[39;00m _module_util\n\u001b[1;32m     38\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mpython\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mutil\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mlazy_loader\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m LazyLoader \u001b[38;5;28;01mas\u001b[39;00m _LazyLoader\n\u001b[1;32m     40\u001b[0m \u001b[38;5;66;03m# Make sure code inside the TensorFlow codebase can use tf2.enabled() at import.\u001b[39;00m\n",
      "File \u001b[0;32m~/miniconda3/envs/rl/lib/python3.8/site-packages/tensorflow/python/__init__.py:37\u001b[0m\n\u001b[1;32m     29\u001b[0m \u001b[38;5;66;03m# We aim to keep this file minimal and ideally remove completely.\u001b[39;00m\n\u001b[1;32m     30\u001b[0m \u001b[38;5;66;03m# If you are adding a new file with @tf_export decorators,\u001b[39;00m\n\u001b[1;32m     31\u001b[0m \u001b[38;5;66;03m# import it in modules_with_exports.py instead.\u001b[39;00m\n\u001b[1;32m     32\u001b[0m \n\u001b[1;32m     33\u001b[0m \u001b[38;5;66;03m# go/tf-wildcard-import\u001b[39;00m\n\u001b[1;32m     34\u001b[0m \u001b[38;5;66;03m# pylint: disable=wildcard-import,g-bad-import-order,g-import-not-at-top\u001b[39;00m\n\u001b[1;32m     36\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mpython\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m pywrap_tensorflow \u001b[38;5;28;01mas\u001b[39;00m _pywrap_tensorflow\n\u001b[0;32m---> 37\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mpython\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01meager\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m context\n\u001b[1;32m     39\u001b[0m \u001b[38;5;66;03m# pylint: enable=wildcard-import\u001b[39;00m\n\u001b[1;32m     40\u001b[0m \n\u001b[1;32m     41\u001b[0m \u001b[38;5;66;03m# Bring in subpackages.\u001b[39;00m\n\u001b[1;32m     42\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mpython\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m data\n",
      "File \u001b[0;32m~/miniconda3/envs/rl/lib/python3.8/site-packages/tensorflow/python/eager/context.py:28\u001b[0m\n\u001b[1;32m     25\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mabsl\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m logging\n\u001b[1;32m     26\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mnumpy\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mnp\u001b[39;00m\n\u001b[0;32m---> 28\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mframework\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m function_pb2\n\u001b[1;32m     29\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mprotobuf\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m config_pb2\n\u001b[1;32m     30\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mprotobuf\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m coordination_config_pb2\n",
      "File \u001b[0;32m~/miniconda3/envs/rl/lib/python3.8/site-packages/tensorflow/core/framework/function_pb2.py:16\u001b[0m\n\u001b[1;32m     11\u001b[0m \u001b[38;5;66;03m# @@protoc_insertion_point(imports)\u001b[39;00m\n\u001b[1;32m     13\u001b[0m _sym_db \u001b[38;5;241m=\u001b[39m _symbol_database\u001b[38;5;241m.\u001b[39mDefault()\n\u001b[0;32m---> 16\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mframework\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m attr_value_pb2 \u001b[38;5;28;01mas\u001b[39;00m tensorflow_dot_core_dot_framework_dot_attr__value__pb2\n\u001b[1;32m     17\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mframework\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m node_def_pb2 \u001b[38;5;28;01mas\u001b[39;00m tensorflow_dot_core_dot_framework_dot_node__def__pb2\n\u001b[1;32m     18\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mframework\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m op_def_pb2 \u001b[38;5;28;01mas\u001b[39;00m tensorflow_dot_core_dot_framework_dot_op__def__pb2\n",
      "File \u001b[0;32m~/miniconda3/envs/rl/lib/python3.8/site-packages/tensorflow/core/framework/attr_value_pb2.py:16\u001b[0m\n\u001b[1;32m     11\u001b[0m \u001b[38;5;66;03m# @@protoc_insertion_point(imports)\u001b[39;00m\n\u001b[1;32m     13\u001b[0m _sym_db \u001b[38;5;241m=\u001b[39m _symbol_database\u001b[38;5;241m.\u001b[39mDefault()\n\u001b[0;32m---> 16\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mframework\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m tensor_pb2 \u001b[38;5;28;01mas\u001b[39;00m tensorflow_dot_core_dot_framework_dot_tensor__pb2\n\u001b[1;32m     17\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mframework\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m tensor_shape_pb2 \u001b[38;5;28;01mas\u001b[39;00m tensorflow_dot_core_dot_framework_dot_tensor__shape__pb2\n\u001b[1;32m     18\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mframework\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m types_pb2 \u001b[38;5;28;01mas\u001b[39;00m tensorflow_dot_core_dot_framework_dot_types__pb2\n",
      "File \u001b[0;32m~/miniconda3/envs/rl/lib/python3.8/site-packages/tensorflow/core/framework/tensor_pb2.py:16\u001b[0m\n\u001b[1;32m     11\u001b[0m \u001b[38;5;66;03m# @@protoc_insertion_point(imports)\u001b[39;00m\n\u001b[1;32m     13\u001b[0m _sym_db \u001b[38;5;241m=\u001b[39m _symbol_database\u001b[38;5;241m.\u001b[39mDefault()\n\u001b[0;32m---> 16\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mframework\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m resource_handle_pb2 \u001b[38;5;28;01mas\u001b[39;00m tensorflow_dot_core_dot_framework_dot_resource__handle__pb2\n\u001b[1;32m     17\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mframework\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m tensor_shape_pb2 \u001b[38;5;28;01mas\u001b[39;00m tensorflow_dot_core_dot_framework_dot_tensor__shape__pb2\n\u001b[1;32m     18\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mframework\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m types_pb2 \u001b[38;5;28;01mas\u001b[39;00m tensorflow_dot_core_dot_framework_dot_types__pb2\n",
      "File \u001b[0;32m~/miniconda3/envs/rl/lib/python3.8/site-packages/tensorflow/core/framework/resource_handle_pb2.py:16\u001b[0m\n\u001b[1;32m     11\u001b[0m \u001b[38;5;66;03m# @@protoc_insertion_point(imports)\u001b[39;00m\n\u001b[1;32m     13\u001b[0m _sym_db \u001b[38;5;241m=\u001b[39m _symbol_database\u001b[38;5;241m.\u001b[39mDefault()\n\u001b[0;32m---> 16\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mframework\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m tensor_shape_pb2 \u001b[38;5;28;01mas\u001b[39;00m tensorflow_dot_core_dot_framework_dot_tensor__shape__pb2\n\u001b[1;32m     17\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mframework\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m types_pb2 \u001b[38;5;28;01mas\u001b[39;00m tensorflow_dot_core_dot_framework_dot_types__pb2\n\u001b[1;32m     20\u001b[0m DESCRIPTOR \u001b[38;5;241m=\u001b[39m _descriptor\u001b[38;5;241m.\u001b[39mFileDescriptor(\n\u001b[1;32m     21\u001b[0m   name\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mtensorflow/core/framework/resource_handle.proto\u001b[39m\u001b[38;5;124m'\u001b[39m,\n\u001b[1;32m     22\u001b[0m   package\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mtensorflow\u001b[39m\u001b[38;5;124m'\u001b[39m,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     26\u001b[0m   ,\n\u001b[1;32m     27\u001b[0m   dependencies\u001b[38;5;241m=\u001b[39m[tensorflow_dot_core_dot_framework_dot_tensor__shape__pb2\u001b[38;5;241m.\u001b[39mDESCRIPTOR,tensorflow_dot_core_dot_framework_dot_types__pb2\u001b[38;5;241m.\u001b[39mDESCRIPTOR,])\n",
      "File \u001b[0;32m~/miniconda3/envs/rl/lib/python3.8/site-packages/tensorflow/core/framework/tensor_shape_pb2.py:36\u001b[0m\n\u001b[1;32m     13\u001b[0m _sym_db \u001b[38;5;241m=\u001b[39m _symbol_database\u001b[38;5;241m.\u001b[39mDefault()\n\u001b[1;32m     18\u001b[0m DESCRIPTOR \u001b[38;5;241m=\u001b[39m _descriptor\u001b[38;5;241m.\u001b[39mFileDescriptor(\n\u001b[1;32m     19\u001b[0m   name\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mtensorflow/core/framework/tensor_shape.proto\u001b[39m\u001b[38;5;124m'\u001b[39m,\n\u001b[1;32m     20\u001b[0m   package\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mtensorflow\u001b[39m\u001b[38;5;124m'\u001b[39m,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     23\u001b[0m   serialized_pb\u001b[38;5;241m=\u001b[39m_b(\u001b[38;5;124m'\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;124m,tensorflow/core/framework/tensor_shape.proto\u001b[39m\u001b[38;5;130;01m\\x12\u001b[39;00m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;124mtensorflow\u001b[39m\u001b[38;5;130;01m\\\"\u001b[39;00m\u001b[38;5;124mz\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;130;01m\\x10\u001b[39;00m\u001b[38;5;124mTensorShapeProto\u001b[39m\u001b[38;5;130;01m\\x12\u001b[39;00m\u001b[38;5;124m-\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;130;01m\\x03\u001b[39;00m\u001b[38;5;130;01m\\x64\u001b[39;00m\u001b[38;5;124mim\u001b[39m\u001b[38;5;130;01m\\x18\u001b[39;00m\u001b[38;5;130;01m\\x02\u001b[39;00m\u001b[38;5;124m \u001b[39m\u001b[38;5;130;01m\\x03\u001b[39;00m\u001b[38;5;124m(\u001b[39m\u001b[38;5;130;01m\\x0b\u001b[39;00m\u001b[38;5;130;01m\\x32\u001b[39;00m\u001b[38;5;124m .tensorflow.TensorShapeProto.Dim\u001b[39m\u001b[38;5;130;01m\\x12\u001b[39;00m\u001b[38;5;130;01m\\x14\u001b[39;00m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;130;01m\\x0c\u001b[39;00m\u001b[38;5;124munknown_rank\u001b[39m\u001b[38;5;130;01m\\x18\u001b[39;00m\u001b[38;5;130;01m\\x03\u001b[39;00m\u001b[38;5;124m \u001b[39m\u001b[38;5;130;01m\\x01\u001b[39;00m\u001b[38;5;124m(\u001b[39m\u001b[38;5;130;01m\\x08\u001b[39;00m\u001b[38;5;130;01m\\x1a\u001b[39;00m\u001b[38;5;124m!\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;130;01m\\x03\u001b[39;00m\u001b[38;5;130;01m\\x44\u001b[39;00m\u001b[38;5;124mim\u001b[39m\u001b[38;5;130;01m\\x12\u001b[39;00m\u001b[38;5;130;01m\\x0c\u001b[39;00m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;130;01m\\x04\u001b[39;00m\u001b[38;5;124msize\u001b[39m\u001b[38;5;130;01m\\x18\u001b[39;00m\u001b[38;5;130;01m\\x01\u001b[39;00m\u001b[38;5;124m \u001b[39m\u001b[38;5;130;01m\\x01\u001b[39;00m\u001b[38;5;124m(\u001b[39m\u001b[38;5;130;01m\\x03\u001b[39;00m\u001b[38;5;130;01m\\x12\u001b[39;00m\u001b[38;5;130;01m\\x0c\u001b[39;00m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;130;01m\\x04\u001b[39;00m\u001b[38;5;124mname\u001b[39m\u001b[38;5;130;01m\\x18\u001b[39;00m\u001b[38;5;130;01m\\x02\u001b[39;00m\u001b[38;5;124m \u001b[39m\u001b[38;5;130;01m\\x01\u001b[39;00m\u001b[38;5;124m(\u001b[39m\u001b[38;5;130;01m\\t\u001b[39;00m\u001b[38;5;124mB\u001b[39m\u001b[38;5;130;01m\\x87\u001b[39;00m\u001b[38;5;130;01m\\x01\u001b[39;00m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;130;01m\\x18\u001b[39;00m\u001b[38;5;124morg.tensorflow.frameworkB\u001b[39m\u001b[38;5;130;01m\\x11\u001b[39;00m\u001b[38;5;124mTensorShapeProtosP\u001b[39m\u001b[38;5;130;01m\\x01\u001b[39;00m\u001b[38;5;124mZSgithub.com/tensorflow/tensorflow/tensorflow/go/core/framework/tensor_shape_go_proto\u001b[39m\u001b[38;5;130;01m\\xf8\u001b[39;00m\u001b[38;5;130;01m\\x01\u001b[39;00m\u001b[38;5;130;01m\\x01\u001b[39;00m\u001b[38;5;130;01m\\x62\u001b[39;00m\u001b[38;5;130;01m\\x06\u001b[39;00m\u001b[38;5;124mproto3\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[1;32m     24\u001b[0m )\n\u001b[1;32m     29\u001b[0m _TENSORSHAPEPROTO_DIM \u001b[38;5;241m=\u001b[39m _descriptor\u001b[38;5;241m.\u001b[39mDescriptor(\n\u001b[1;32m     30\u001b[0m   name\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mDim\u001b[39m\u001b[38;5;124m'\u001b[39m,\n\u001b[1;32m     31\u001b[0m   full_name\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mtensorflow.TensorShapeProto.Dim\u001b[39m\u001b[38;5;124m'\u001b[39m,\n\u001b[1;32m     32\u001b[0m   filename\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m     33\u001b[0m   file\u001b[38;5;241m=\u001b[39mDESCRIPTOR,\n\u001b[1;32m     34\u001b[0m   containing_type\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m     35\u001b[0m   fields\u001b[38;5;241m=\u001b[39m[\n\u001b[0;32m---> 36\u001b[0m     \u001b[43m_descriptor\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mFieldDescriptor\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m     37\u001b[0m \u001b[43m      \u001b[49m\u001b[43mname\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43msize\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfull_name\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mtensorflow.TensorShapeProto.Dim.size\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mindex\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     38\u001b[0m \u001b[43m      \u001b[49m\u001b[43mnumber\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mtype\u001b[39;49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m3\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcpp_type\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m2\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mlabel\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     39\u001b[0m \u001b[43m      \u001b[49m\u001b[43mhas_default_value\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mFalse\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdefault_value\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     40\u001b[0m \u001b[43m      \u001b[49m\u001b[43mmessage_type\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43menum_type\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcontaining_type\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m     41\u001b[0m \u001b[43m      \u001b[49m\u001b[43mis_extension\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mFalse\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mextension_scope\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m     42\u001b[0m \u001b[43m      \u001b[49m\u001b[43mserialized_options\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfile\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mDESCRIPTOR\u001b[49m\u001b[43m)\u001b[49m,\n\u001b[1;32m     43\u001b[0m     _descriptor\u001b[38;5;241m.\u001b[39mFieldDescriptor(\n\u001b[1;32m     44\u001b[0m       name\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mname\u001b[39m\u001b[38;5;124m'\u001b[39m, full_name\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mtensorflow.TensorShapeProto.Dim.name\u001b[39m\u001b[38;5;124m'\u001b[39m, index\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1\u001b[39m,\n\u001b[1;32m     45\u001b[0m       number\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m2\u001b[39m, \u001b[38;5;28mtype\u001b[39m\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m9\u001b[39m, cpp_type\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m9\u001b[39m, label\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1\u001b[39m,\n\u001b[1;32m     46\u001b[0m       has_default_value\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mFalse\u001b[39;00m, default_value\u001b[38;5;241m=\u001b[39m_b(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m\"\u001b[39m)\u001b[38;5;241m.\u001b[39mdecode(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mutf-8\u001b[39m\u001b[38;5;124m'\u001b[39m),\n\u001b[1;32m     47\u001b[0m       message_type\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m, enum_type\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m, containing_type\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m     48\u001b[0m       is_extension\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mFalse\u001b[39;00m, extension_scope\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m     49\u001b[0m       serialized_options\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m, file\u001b[38;5;241m=\u001b[39mDESCRIPTOR),\n\u001b[1;32m     50\u001b[0m   ],\n\u001b[1;32m     51\u001b[0m   extensions\u001b[38;5;241m=\u001b[39m[\n\u001b[1;32m     52\u001b[0m   ],\n\u001b[1;32m     53\u001b[0m   nested_types\u001b[38;5;241m=\u001b[39m[],\n\u001b[1;32m     54\u001b[0m   enum_types\u001b[38;5;241m=\u001b[39m[\n\u001b[1;32m     55\u001b[0m   ],\n\u001b[1;32m     56\u001b[0m   serialized_options\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m     57\u001b[0m   is_extendable\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mFalse\u001b[39;00m,\n\u001b[1;32m     58\u001b[0m   syntax\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mproto3\u001b[39m\u001b[38;5;124m'\u001b[39m,\n\u001b[1;32m     59\u001b[0m   extension_ranges\u001b[38;5;241m=\u001b[39m[],\n\u001b[1;32m     60\u001b[0m   oneofs\u001b[38;5;241m=\u001b[39m[\n\u001b[1;32m     61\u001b[0m   ],\n\u001b[1;32m     62\u001b[0m   serialized_start\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m149\u001b[39m,\n\u001b[1;32m     63\u001b[0m   serialized_end\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m182\u001b[39m,\n\u001b[1;32m     64\u001b[0m )\n\u001b[1;32m     66\u001b[0m _TENSORSHAPEPROTO \u001b[38;5;241m=\u001b[39m _descriptor\u001b[38;5;241m.\u001b[39mDescriptor(\n\u001b[1;32m     67\u001b[0m   name\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mTensorShapeProto\u001b[39m\u001b[38;5;124m'\u001b[39m,\n\u001b[1;32m     68\u001b[0m   full_name\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mtensorflow.TensorShapeProto\u001b[39m\u001b[38;5;124m'\u001b[39m,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    100\u001b[0m   serialized_end\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m182\u001b[39m,\n\u001b[1;32m    101\u001b[0m )\n\u001b[1;32m    103\u001b[0m _TENSORSHAPEPROTO_DIM\u001b[38;5;241m.\u001b[39mcontaining_type \u001b[38;5;241m=\u001b[39m _TENSORSHAPEPROTO\n",
      "File \u001b[0;32m~/miniconda3/envs/rl/lib/python3.8/site-packages/google/protobuf/descriptor.py:561\u001b[0m, in \u001b[0;36mFieldDescriptor.__new__\u001b[0;34m(cls, name, full_name, index, number, type, cpp_type, label, default_value, message_type, enum_type, containing_type, is_extension, extension_scope, options, serialized_options, has_default_value, containing_oneof, json_name, file, create_key)\u001b[0m\n\u001b[1;32m    555\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__new__\u001b[39m(\u001b[38;5;28mcls\u001b[39m, name, full_name, index, number, \u001b[38;5;28mtype\u001b[39m, cpp_type, label,\n\u001b[1;32m    556\u001b[0m             default_value, message_type, enum_type, containing_type,\n\u001b[1;32m    557\u001b[0m             is_extension, extension_scope, options\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m    558\u001b[0m             serialized_options\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m    559\u001b[0m             has_default_value\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m, containing_oneof\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m, json_name\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m    560\u001b[0m             file\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m, create_key\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m):  \u001b[38;5;66;03m# pylint: disable=redefined-builtin\u001b[39;00m\n\u001b[0;32m--> 561\u001b[0m   \u001b[43m_message\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mMessage\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_CheckCalledFromGeneratedFile\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    562\u001b[0m   \u001b[38;5;28;01mif\u001b[39;00m is_extension:\n\u001b[1;32m    563\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m _message\u001b[38;5;241m.\u001b[39mdefault_pool\u001b[38;5;241m.\u001b[39mFindExtensionByName(full_name)\n",
      "\u001b[0;31mTypeError\u001b[0m: Descriptors cannot not be created directly.\nIf this call came from a _pb2.py file, your generated code is out of date and must be regenerated with protoc >= 3.19.0.\nIf you cannot immediately regenerate your protos, some other possible workarounds are:\n 1. Downgrade the protobuf package to 3.20.x or lower.\n 2. Set PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python (but this will use pure-Python parsing and will be much slower).\n\nMore information: https://developers.google.com/protocol-buffers/docs/news/2022-05-06#python-updates"
     ]
    }
   ],
   "source": [
    "# External dependencies\n",
    "import os\n",
    "import shutil\n",
    "import numpy as np\n",
    "from nvtabular.ops import *\n",
    "from merlin.schema.tags import Tags\n",
    "\n",
    "import nvtabular as nvt\n",
    "\n",
    "from os import path\n",
    "\n",
    "# Get dataframe library - cudf or pandas\n",
    "from merlin.core.dispatch import get_lib\n",
    "df_lib = get_lib()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "87bfc236",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting merlin-models\n",
      "  Downloading merlin-models-23.6.0.tar.gz (418 kB)\n",
      "\u001b[2K     \u001b[38;2;114;156;31m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m418.9/418.9 kB\u001b[0m \u001b[31m5.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m MB/s\u001b[0m eta \u001b[36m0:00:01\u001b[0m\n",
      "\u001b[?25h  Installing build dependencies ... \u001b[?25ldone\n",
      "\u001b[?25h  Getting requirements to build wheel ... \u001b[?25ldone\n",
      "\u001b[?25h  Preparing metadata (pyproject.toml) ... \u001b[?25ldone\n",
      "\u001b[?25hRequirement already satisfied: merlin-dataloader<23.07,>=23.06 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from merlin-models) (23.6.0)\n",
      "Requirement already satisfied: merlin-core<23.07,>=23.06 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from merlin-models) (23.6.0)\n",
      "Requirement already satisfied: packaging in ./miniconda3/envs/rl/lib/python3.8/site-packages (from merlin-core<23.07,>=23.06->merlin-models) (22.0)\n",
      "Requirement already satisfied: pandas<1.6.0dev0,>=1.2.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from merlin-core<23.07,>=23.06->merlin-models) (1.5.2)\n",
      "Requirement already satisfied: pyarrow>=5.0.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from merlin-core<23.07,>=23.06->merlin-models) (11.0.0)\n",
      "Requirement already satisfied: numba>=0.54 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from merlin-core<23.07,>=23.06->merlin-models) (0.57.1)\n",
      "Requirement already satisfied: dask>=2022.11.1 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from merlin-core<23.07,>=23.06->merlin-models) (2023.3.2)\n",
      "Requirement already satisfied: fsspec>=2022.7.1 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from merlin-core<23.07,>=23.06->merlin-models) (2023.6.0)\n",
      "Requirement already satisfied: distributed>=2022.11.1 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from merlin-core<23.07,>=23.06->merlin-models) (2023.3.2.1)\n",
      "Requirement already satisfied: protobuf>=3.0.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from merlin-core<23.07,>=23.06->merlin-models) (4.23.3)\n",
      "Requirement already satisfied: pynvml<11.5,>=11.0.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from merlin-core<23.07,>=23.06->merlin-models) (11.4.1)\n",
      "Requirement already satisfied: dask-cuda>=22.12.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from merlin-core<23.07,>=23.06->merlin-models) (23.4.0)\n",
      "Requirement already satisfied: tensorflow-metadata>=1.2.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from merlin-core<23.07,>=23.06->merlin-models) (1.13.1)\n",
      "Requirement already satisfied: tqdm>=4.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from merlin-core<23.07,>=23.06->merlin-models) (4.65.0)\n",
      "Requirement already satisfied: npy-append-array in ./miniconda3/envs/rl/lib/python3.8/site-packages (from merlin-core<23.07,>=23.06->merlin-models) (0.9.16)\n",
      "Requirement already satisfied: numpy>=1.22.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from merlin-core<23.07,>=23.06->merlin-models) (1.24.1)\n",
      "Requirement already satisfied: betterproto<2.0.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from merlin-core<23.07,>=23.06->merlin-models) (1.2.5)\n",
      "Requirement already satisfied: grpclib in ./miniconda3/envs/rl/lib/python3.8/site-packages (from betterproto<2.0.0->merlin-core<23.07,>=23.06->merlin-models) (0.4.5)\n",
      "Requirement already satisfied: stringcase in ./miniconda3/envs/rl/lib/python3.8/site-packages (from betterproto<2.0.0->merlin-core<23.07,>=23.06->merlin-models) (1.2.0)\n",
      "Requirement already satisfied: toolz>=0.8.2 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from dask>=2022.11.1->merlin-core<23.07,>=23.06->merlin-models) (0.12.0)\n",
      "Requirement already satisfied: pyyaml>=5.3.1 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from dask>=2022.11.1->merlin-core<23.07,>=23.06->merlin-models) (6.0)\n",
      "Requirement already satisfied: partd>=1.2.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from dask>=2022.11.1->merlin-core<23.07,>=23.06->merlin-models) (1.4.0)\n",
      "Requirement already satisfied: click>=7.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from dask>=2022.11.1->merlin-core<23.07,>=23.06->merlin-models) (8.1.3)\n",
      "Requirement already satisfied: importlib-metadata>=4.13.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from dask>=2022.11.1->merlin-core<23.07,>=23.06->merlin-models) (4.13.0)\n",
      "Requirement already satisfied: cloudpickle>=1.1.1 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from dask>=2022.11.1->merlin-core<23.07,>=23.06->merlin-models) (2.2.0)\n",
      "Requirement already satisfied: zict>=0.1.3 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from dask-cuda>=22.12.0->merlin-core<23.07,>=23.06->merlin-models) (3.0.0)\n",
      "Requirement already satisfied: msgpack>=1.0.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from distributed>=2022.11.1->merlin-core<23.07,>=23.06->merlin-models) (1.0.5)\n",
      "Requirement already satisfied: tblib>=1.6.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from distributed>=2022.11.1->merlin-core<23.07,>=23.06->merlin-models) (2.0.0)\n",
      "Requirement already satisfied: urllib3>=1.24.3 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from distributed>=2022.11.1->merlin-core<23.07,>=23.06->merlin-models) (1.26.13)\n",
      "Requirement already satisfied: jinja2>=2.10.3 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from distributed>=2022.11.1->merlin-core<23.07,>=23.06->merlin-models) (3.1.2)\n",
      "Requirement already satisfied: locket>=1.0.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from distributed>=2022.11.1->merlin-core<23.07,>=23.06->merlin-models) (1.0.0)\n",
      "Requirement already satisfied: sortedcontainers>=2.0.5 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from distributed>=2022.11.1->merlin-core<23.07,>=23.06->merlin-models) (2.4.0)\n",
      "Requirement already satisfied: tornado>=6.0.3 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from distributed>=2022.11.1->merlin-core<23.07,>=23.06->merlin-models) (6.2)\n",
      "Requirement already satisfied: psutil>=5.7.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from distributed>=2022.11.1->merlin-core<23.07,>=23.06->merlin-models) (5.9.4)\n",
      "Requirement already satisfied: llvmlite<0.41,>=0.40.0dev0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from numba>=0.54->merlin-core<23.07,>=23.06->merlin-models) (0.40.1)\n",
      "Requirement already satisfied: pytz>=2020.1 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from pandas<1.6.0dev0,>=1.2.0->merlin-core<23.07,>=23.06->merlin-models) (2022.7)\n",
      "Requirement already satisfied: python-dateutil>=2.8.1 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from pandas<1.6.0dev0,>=1.2.0->merlin-core<23.07,>=23.06->merlin-models) (2.8.2)\n",
      "Requirement already satisfied: googleapis-common-protos<2,>=1.52.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from tensorflow-metadata>=1.2.0->merlin-core<23.07,>=23.06->merlin-models) (1.59.1)\n",
      "Requirement already satisfied: absl-py<2.0.0,>=0.9 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from tensorflow-metadata>=1.2.0->merlin-core<23.07,>=23.06->merlin-models) (1.3.0)\n",
      "Requirement already satisfied: zipp>=0.5 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from importlib-metadata>=4.13.0->dask>=2022.11.1->merlin-core<23.07,>=23.06->merlin-models) (3.11.0)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from jinja2>=2.10.3->distributed>=2022.11.1->merlin-core<23.07,>=23.06->merlin-models) (2.1.1)\n",
      "Requirement already satisfied: six>=1.5 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from python-dateutil>=2.8.1->pandas<1.6.0dev0,>=1.2.0->merlin-core<23.07,>=23.06->merlin-models) (1.16.0)\n",
      "Requirement already satisfied: h2<5,>=3.1.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from grpclib->betterproto<2.0.0->merlin-core<23.07,>=23.06->merlin-models) (4.1.0)\n",
      "Requirement already satisfied: multidict in ./miniconda3/envs/rl/lib/python3.8/site-packages (from grpclib->betterproto<2.0.0->merlin-core<23.07,>=23.06->merlin-models) (6.0.4)\n",
      "Requirement already satisfied: hpack<5,>=4.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from h2<5,>=3.1.0->grpclib->betterproto<2.0.0->merlin-core<23.07,>=23.06->merlin-models) (4.0.0)\n",
      "Requirement already satisfied: hyperframe<7,>=6.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from h2<5,>=3.1.0->grpclib->betterproto<2.0.0->merlin-core<23.07,>=23.06->merlin-models) (6.0.1)\n",
      "Building wheels for collected packages: merlin-models\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  Building wheel for merlin-models (pyproject.toml) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for merlin-models: filename=merlin_models-23.6.0-py3-none-any.whl size=386073 sha256=9131847b4e2ca5ed9225246f91730aef3688e922fadb57c3dfd95bf73525e51c\n",
      "  Stored in directory: /Users/rajathdb/Library/Caches/pip/wheels/9b/7f/83/44b6f42cde2fff7339924dfe733df5e63b3f5c7b49b17be6c9\n",
      "Successfully built merlin-models\n",
      "Installing collected packages: merlin-models\n",
      "Successfully installed merlin-models-23.6.0\n"
     ]
    }
   ],
   "source": [
    "!pip install merlin-models\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "196d0f19",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting protobuf==3.19\n",
      "  Downloading protobuf-3.19.0-py2.py3-none-any.whl (162 kB)\n",
      "\u001b[2K     \u001b[38;2;114;156;31m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m162.6/162.6 kB\u001b[0m \u001b[31m4.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hInstalling collected packages: protobuf\n",
      "  Attempting uninstall: protobuf\n",
      "    Found existing installation: protobuf 3.20.0\n",
      "    Uninstalling protobuf-3.20.0:\n",
      "      Successfully uninstalled protobuf-3.20.0\n",
      "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
      "tensorflow-metadata 1.13.1 requires protobuf<5,>=3.20.3, but you have protobuf 3.19.0 which is incompatible.\n",
      "googleapis-common-protos 1.59.1 requires protobuf!=3.20.0,!=3.20.1,!=4.21.1,!=4.21.2,!=4.21.3,!=4.21.4,!=4.21.5,<5.0.0.dev0,>=3.19.5, but you have protobuf 3.19.0 which is incompatible.\u001b[0m\u001b[31m\n",
      "\u001b[0mSuccessfully installed protobuf-3.19.0\n"
     ]
    }
   ],
   "source": [
    "!pip install protobuf==3.19\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "d2eb81fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "!export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e906ac5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# External dependencies\n",
    "import os\n",
    "\n",
    "from merlin.core.utils import download_file\n",
    "\n",
    "# Get dataframe library - cudf or pandas\n",
    "from merlin.core.dispatch import get_lib\n",
    "df_lib = get_lib()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "dbee1417",
   "metadata": {},
   "outputs": [],
   "source": [
    "INPUT_DATA_DIR = os.environ.get(\n",
    "    \"INPUT_DATA_DIR\", os.path.expanduser(\"nvtabular/movielens/data/\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "8dc1a562",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "downloading ml-25m.zip: 262MB [02:01, 2.16MB/s]                                         \n",
      "unzipping files: 100%|█████████████████████████████████| 8/8 [00:03<00:00,  2.11files/s]\n"
     ]
    }
   ],
   "source": [
    "download_file(\n",
    "    \"http://files.grouplens.org/datasets/movielens/ml-25m.zip\",\n",
    "    os.path.join(INPUT_DATA_DIR, \"ml-25m.zip\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "472fce7b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>movieId</th>\n",
       "      <th>title</th>\n",
       "      <th>genres</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>Toy Story (1995)</td>\n",
       "      <td>Adventure|Animation|Children|Comedy|Fantasy</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2</td>\n",
       "      <td>Jumanji (1995)</td>\n",
       "      <td>Adventure|Children|Fantasy</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3</td>\n",
       "      <td>Grumpier Old Men (1995)</td>\n",
       "      <td>Comedy|Romance</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>4</td>\n",
       "      <td>Waiting to Exhale (1995)</td>\n",
       "      <td>Comedy|Drama|Romance</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>5</td>\n",
       "      <td>Father of the Bride Part II (1995)</td>\n",
       "      <td>Comedy</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   movieId                               title  \\\n",
       "0        1                    Toy Story (1995)   \n",
       "1        2                      Jumanji (1995)   \n",
       "2        3             Grumpier Old Men (1995)   \n",
       "3        4            Waiting to Exhale (1995)   \n",
       "4        5  Father of the Bride Part II (1995)   \n",
       "\n",
       "                                        genres  \n",
       "0  Adventure|Animation|Children|Comedy|Fantasy  \n",
       "1                   Adventure|Children|Fantasy  \n",
       "2                               Comedy|Romance  \n",
       "3                         Comedy|Drama|Romance  \n",
       "4                                       Comedy  "
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "movies = df_lib.read_csv(os.path.join(INPUT_DATA_DIR, \"ml-25m/movies.csv\"))\n",
    "movies.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "19167933",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>movieId</th>\n",
       "      <th>genres</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>[Adventure, Animation, Children, Comedy, Fantasy]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2</td>\n",
       "      <td>[Adventure, Children, Fantasy]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3</td>\n",
       "      <td>[Comedy, Romance]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>4</td>\n",
       "      <td>[Comedy, Drama, Romance]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>5</td>\n",
       "      <td>[Comedy]</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   movieId                                             genres\n",
       "0        1  [Adventure, Animation, Children, Comedy, Fantasy]\n",
       "1        2                     [Adventure, Children, Fantasy]\n",
       "2        3                                  [Comedy, Romance]\n",
       "3        4                           [Comedy, Drama, Romance]\n",
       "4        5                                           [Comedy]"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "movies[\"genres\"] = movies[\"genres\"].str.split(\"|\")\n",
    "movies = movies.drop(\"title\", axis=1)\n",
    "movies.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "84089208",
   "metadata": {},
   "outputs": [],
   "source": [
    "movies.to_parquet(os.path.join(INPUT_DATA_DIR, \"movies_converted.parquet\"))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "5cc64ab8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>userId</th>\n",
       "      <th>movieId</th>\n",
       "      <th>rating</th>\n",
       "      <th>timestamp</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>296</td>\n",
       "      <td>5.0</td>\n",
       "      <td>1147880044</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>306</td>\n",
       "      <td>3.5</td>\n",
       "      <td>1147868817</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1</td>\n",
       "      <td>307</td>\n",
       "      <td>5.0</td>\n",
       "      <td>1147868828</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1</td>\n",
       "      <td>665</td>\n",
       "      <td>5.0</td>\n",
       "      <td>1147878820</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1</td>\n",
       "      <td>899</td>\n",
       "      <td>3.5</td>\n",
       "      <td>1147868510</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   userId  movieId  rating   timestamp\n",
       "0       1      296     5.0  1147880044\n",
       "1       1      306     3.5  1147868817\n",
       "2       1      307     5.0  1147868828\n",
       "3       1      665     5.0  1147878820\n",
       "4       1      899     3.5  1147868510"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ratings = df_lib.read_csv(os.path.join(INPUT_DATA_DIR, \"ml-25m\", \"ratings.csv\"))\n",
    "ratings.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "51acf93b",
   "metadata": {},
   "outputs": [],
   "source": [
    "ratings = ratings.drop(\"timestamp\", axis=1)\n",
    "\n",
    "# shuffle the dataset\n",
    "ratings = ratings.sample(len(ratings), replace=False)\n",
    "\n",
    "# split the train_df as training and validation data sets.\n",
    "num_valid = int(len(ratings) * 0.2)\n",
    "\n",
    "train = ratings[:-num_valid]\n",
    "valid = ratings[-num_valid:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "b1821627",
   "metadata": {},
   "outputs": [],
   "source": [
    "train.to_parquet(os.path.join(INPUT_DATA_DIR, \"train.parquet\"))\n",
    "valid.to_parquet(os.path.join(INPUT_DATA_DIR, \"valid.parquet\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "80d44421",
   "metadata": {},
   "outputs": [],
   "source": [
    "# External dependencies\n",
    "import os\n",
    "import shutil\n",
    "import numpy as np\n",
    "from nvtabular.ops import *\n",
    "from merlin.schema.tags import Tags\n",
    "\n",
    "import nvtabular as nvt\n",
    "\n",
    "from os import path\n",
    "\n",
    "# Get dataframe library - cudf or pandas\n",
    "from merlin.core.dispatch import get_lib\n",
    "df_lib = get_lib()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "62cc245f",
   "metadata": {},
   "outputs": [],
   "source": [
    "CATEGORICAL_COLUMNS = [\"userId\", \"movieId\"]\n",
    "LABEL_COLUMNS = [\"rating\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "bc4bcb5a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/rajathdb/miniconda3/envs/rl/lib/python3.8/site-packages/merlin/io/dataset.py:267: UserWarning: Initializing an NVTabular Dataset in CPU mode.This is an experimental feature with extremely limited support!\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "userId = [\"userId\"] >> TagAsUserID()\n",
    "movieId = [\"movieId\"] >> TagAsItemID()\n",
    "\n",
    "joined = userId + movieId >> JoinExternal(movies, on=[\"movieId\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "65150712",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/svg+xml": [
       "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
       "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
       " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
       "<!-- Generated by graphviz version 8.0.5 (20230430.1635)\n",
       " -->\n",
       "<!-- Pages: 1 -->\n",
       "<svg width=\"276pt\" height=\"404pt\"\n",
       " viewBox=\"0.00 0.00 276.22 404.00\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 400)\">\n",
       "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-400 272.22,-400 272.22,4 -4,4\"/>\n",
       "<!-- 0 -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>0</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"133.61\" cy=\"-90\" rx=\"57.49\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"133.61\" y=\"-84.95\" font-family=\"Times,serif\" font-size=\"14.00\">JoinExternal</text>\n",
       "</g>\n",
       "<!-- 6 -->\n",
       "<g id=\"node9\" class=\"node\">\n",
       "<title>6</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"133.61\" cy=\"-18\" rx=\"51.86\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"133.61\" y=\"-12.95\" font-family=\"Times,serif\" font-size=\"14.00\">output cols</text>\n",
       "</g>\n",
       "<!-- 0&#45;&gt;6 -->\n",
       "<g id=\"edge8\" class=\"edge\">\n",
       "<title>0&#45;&gt;6</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M133.61,-71.7C133.61,-64.24 133.61,-55.32 133.61,-46.97\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"137.11,-47.1 133.61,-37.1 130.11,-47.1 137.11,-47.1\"/>\n",
       "</g>\n",
       "<!-- 4 -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>4</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"133.61\" cy=\"-162\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"133.61\" y=\"-156.95\" font-family=\"Times,serif\" font-size=\"14.00\">+</text>\n",
       "</g>\n",
       "<!-- 4&#45;&gt;0 -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>4&#45;&gt;0</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M133.61,-143.7C133.61,-136.24 133.61,-127.32 133.61,-118.97\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"137.11,-119.1 133.61,-109.1 130.11,-119.1 137.11,-119.1\"/>\n",
       "</g>\n",
       "<!-- 1 -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>1</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"62.61\" cy=\"-234\" rx=\"62.61\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"62.61\" y=\"-228.95\" font-family=\"Times,serif\" font-size=\"14.00\">TagAsUserID</text>\n",
       "</g>\n",
       "<!-- 1&#45;&gt;4 -->\n",
       "<g id=\"edge5\" class=\"edge\">\n",
       "<title>1&#45;&gt;4</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M79.44,-216.41C88.89,-207.1 100.8,-195.35 111.05,-185.24\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"113.17,-188.08 117.83,-178.57 108.25,-183.1 113.17,-188.08\"/>\n",
       "</g>\n",
       "<!-- 2 -->\n",
       "<g id=\"node4\" class=\"node\">\n",
       "<title>2</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"62.61\" cy=\"-306\" rx=\"56.47\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"62.61\" y=\"-300.95\" font-family=\"Times,serif\" font-size=\"14.00\">SelectionOp</text>\n",
       "</g>\n",
       "<!-- 2&#45;&gt;1 -->\n",
       "<g id=\"edge2\" class=\"edge\">\n",
       "<title>2&#45;&gt;1</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M62.61,-287.7C62.61,-280.24 62.61,-271.32 62.61,-262.97\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"66.11,-263.1 62.61,-253.1 59.11,-263.1 66.11,-263.1\"/>\n",
       "</g>\n",
       "<!-- 2_selector -->\n",
       "<g id=\"node5\" class=\"node\">\n",
       "<title>2_selector</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"62.61\" cy=\"-378\" rx=\"43.16\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"62.61\" y=\"-372.95\" font-family=\"Times,serif\" font-size=\"14.00\">[&#39;userId&#39;]</text>\n",
       "</g>\n",
       "<!-- 2_selector&#45;&gt;2 -->\n",
       "<g id=\"edge3\" class=\"edge\">\n",
       "<title>2_selector&#45;&gt;2</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M62.61,-359.7C62.61,-352.24 62.61,-343.32 62.61,-334.97\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"66.11,-335.1 62.61,-325.1 59.11,-335.1 66.11,-335.1\"/>\n",
       "</g>\n",
       "<!-- 3 -->\n",
       "<g id=\"node6\" class=\"node\">\n",
       "<title>3</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"205.61\" cy=\"-306\" rx=\"56.47\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"205.61\" y=\"-300.95\" font-family=\"Times,serif\" font-size=\"14.00\">SelectionOp</text>\n",
       "</g>\n",
       "<!-- 5 -->\n",
       "<g id=\"node8\" class=\"node\">\n",
       "<title>5</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"205.61\" cy=\"-234\" rx=\"62.61\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"205.61\" y=\"-228.95\" font-family=\"Times,serif\" font-size=\"14.00\">TagAsItemID</text>\n",
       "</g>\n",
       "<!-- 3&#45;&gt;5 -->\n",
       "<g id=\"edge7\" class=\"edge\">\n",
       "<title>3&#45;&gt;5</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M205.61,-287.7C205.61,-280.24 205.61,-271.32 205.61,-262.97\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"209.11,-263.1 205.61,-253.1 202.11,-263.1 209.11,-263.1\"/>\n",
       "</g>\n",
       "<!-- 3_selector -->\n",
       "<g id=\"node7\" class=\"node\">\n",
       "<title>3_selector</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"205.61\" cy=\"-378\" rx=\"51.35\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"205.61\" y=\"-372.95\" font-family=\"Times,serif\" font-size=\"14.00\">[&#39;movieId&#39;]</text>\n",
       "</g>\n",
       "<!-- 3_selector&#45;&gt;3 -->\n",
       "<g id=\"edge4\" class=\"edge\">\n",
       "<title>3_selector&#45;&gt;3</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M205.61,-359.7C205.61,-352.24 205.61,-343.32 205.61,-334.97\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"209.11,-335.1 205.61,-325.1 202.11,-335.1 209.11,-335.1\"/>\n",
       "</g>\n",
       "<!-- 5&#45;&gt;4 -->\n",
       "<g id=\"edge6\" class=\"edge\">\n",
       "<title>5&#45;&gt;4</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M188.55,-216.41C178.74,-206.87 166.31,-194.79 155.75,-184.53\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"158.73,-182.57 149.12,-178.11 153.85,-187.59 158.73,-182.57\"/>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "<graphviz.graphs.Digraph at 0x15dce2ee0>"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "joined.graph\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "580559c1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting graphviz\n",
      "  Downloading graphviz-0.20.1-py3-none-any.whl (47 kB)\n",
      "\u001b[2K     \u001b[38;2;114;156;31m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m47.0/47.0 kB\u001b[0m \u001b[31m2.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hInstalling collected packages: graphviz\n",
      "Successfully installed graphviz-0.20.1\n"
     ]
    }
   ],
   "source": [
    "!pip install graphviz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "8583f80c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running `brew update --auto-update`...\n",
      "\u001b[34m==>\u001b[0m \u001b[1mAuto-updated Homebrew!\u001b[0m\n",
      "Updated 2 taps (homebrew/core and homebrew/cask).\n",
      "\u001b[34m==>\u001b[0m \u001b[1mNew Formulae\u001b[0m\n",
      "cppinsights  pypy3.10     pypy3.9      python-toml  rio          scrapy\n",
      "\u001b[34m==>\u001b[0m \u001b[1mNew Casks\u001b[0m\n",
      "drata-agent                picoscope                  score\n",
      "\n",
      "You have \u001b[1m19\u001b[0m outdated formulae and \u001b[1m1\u001b[0m outdated cask installed.\n",
      "\n",
      "\u001b[32m==>\u001b[0m \u001b[1mFetching dependencies for graphviz: \u001b[32mfreetype\u001b[39m, \u001b[32mimath\u001b[39m, \u001b[32mopenexr\u001b[39m, \u001b[32mjpeg-xl\u001b[39m, \u001b[32mlibavif\u001b[39m, \u001b[32mgd\u001b[39m, \u001b[32mjasper\u001b[39m, \u001b[32mnetpbm\u001b[39m, \u001b[32mgts\u001b[39m, \u001b[32mlibx11\u001b[39m, \u001b[32mgdk-pixbuf\u001b[39m, \u001b[32mlibrsvg\u001b[39m, \u001b[32mm4\u001b[39m and \u001b[32mlibtool\u001b[39m\u001b[0m\n",
      "\u001b[32m==>\u001b[0m \u001b[1mFetching \u001b[32mfreetype\u001b[39m\u001b[0m\n",
      "\u001b[34m==>\u001b[0m \u001b[1mDownloading https://ghcr.io/v2/homebrew/core/freetype/manifests/2.13.1\u001b[0m\n",
      "######################################################################### 100.0%\n",
      "\u001b[34m==>\u001b[0m \u001b[1mDownloading https://ghcr.io/v2/homebrew/core/freetype/blobs/sha256:87a44e9a3\u001b[0m\n",
      "######################################################################### 100.0%\n",
      "\u001b[32m==>\u001b[0m \u001b[1mFetching \u001b[32mimath\u001b[39m\u001b[0m\n",
      "\u001b[34m==>\u001b[0m \u001b[1mDownloading https://ghcr.io/v2/homebrew/core/imath/manifests/3.1.9\u001b[0m\n",
      "######################################################################### 100.0%\n",
      "\u001b[34m==>\u001b[0m \u001b[1mDownloading https://ghcr.io/v2/homebrew/core/imath/blobs/sha256:089d0a19694c\u001b[0m\n",
      "######################################################################### 100.0%\n",
      "\u001b[32m==>\u001b[0m \u001b[1mFetching \u001b[32mopenexr\u001b[39m\u001b[0m\n",
      "\u001b[34m==>\u001b[0m \u001b[1mDownloading https://ghcr.io/v2/homebrew/core/openexr/manifests/3.1.9\u001b[0m\n",
      "######################################################################### 100.0%\n",
      "\u001b[34m==>\u001b[0m \u001b[1mDownloading https://ghcr.io/v2/homebrew/core/openexr/blobs/sha256:522d4717b7\u001b[0m\n",
      "######################################################################### 100.0%\n",
      "\u001b[32m==>\u001b[0m \u001b[1mFetching \u001b[32mjpeg-xl\u001b[39m\u001b[0m\n",
      "\u001b[34m==>\u001b[0m \u001b[1mDownloading https://ghcr.io/v2/homebrew/core/jpeg-xl/manifests/0.8.2\u001b[0m\n",
      "######################################################################### 100.0%\n",
      "\u001b[34m==>\u001b[0m \u001b[1mDownloading https://ghcr.io/v2/homebrew/core/jpeg-xl/blobs/sha256:6186d463dd\u001b[0m\n",
      "######################################################################### 100.0%\n",
      "\u001b[32m==>\u001b[0m \u001b[1mFetching \u001b[32mlibavif\u001b[39m\u001b[0m\n",
      "\u001b[34m==>\u001b[0m \u001b[1mDownloading https://ghcr.io/v2/homebrew/core/libavif/manifests/0.11.1\u001b[0m\n",
      "######################################################################### 100.0%\n",
      "\u001b[34m==>\u001b[0m \u001b[1mDownloading https://ghcr.io/v2/homebrew/core/libavif/blobs/sha256:9b4735a706\u001b[0m\n",
      "######################################################################### 100.0%\n",
      "\u001b[32m==>\u001b[0m \u001b[1mFetching \u001b[32mgd\u001b[39m\u001b[0m\n",
      "\u001b[34m==>\u001b[0m \u001b[1mDownloading https://ghcr.io/v2/homebrew/core/gd/manifests/2.3.3_5\u001b[0m\n",
      "######################################################################### 100.0%\n",
      "\u001b[34m==>\u001b[0m \u001b[1mDownloading https://ghcr.io/v2/homebrew/core/gd/blobs/sha256:3e23180f273c110\u001b[0m\n",
      "######################################################################### 100.0%\n",
      "\u001b[32m==>\u001b[0m \u001b[1mFetching \u001b[32mjasper\u001b[39m\u001b[0m\n",
      "\u001b[34m==>\u001b[0m \u001b[1mDownloading https://ghcr.io/v2/homebrew/core/jasper/manifests/4.0.0\u001b[0m\n",
      "######################################################################### 100.0%\n",
      "\u001b[34m==>\u001b[0m \u001b[1mDownloading https://ghcr.io/v2/homebrew/core/jasper/blobs/sha256:a24af85ebb8\u001b[0m\n",
      "######################################################################### 100.0%\n",
      "\u001b[32m==>\u001b[0m \u001b[1mFetching \u001b[32mnetpbm\u001b[39m\u001b[0m\n",
      "\u001b[34m==>\u001b[0m \u001b[1mDownloading https://ghcr.io/v2/homebrew/core/netpbm/manifests/10.86.38_1-1\u001b[0m\n",
      "######################################################################### 100.0%\n",
      "\u001b[34m==>\u001b[0m \u001b[1mDownloading https://ghcr.io/v2/homebrew/core/netpbm/blobs/sha256:9ff16b71c63\u001b[0m\n",
      "######################################################################### 100.0%\n",
      "\u001b[32m==>\u001b[0m \u001b[1mFetching \u001b[32mgts\u001b[39m\u001b[0m\n",
      "\u001b[34m==>\u001b[0m \u001b[1mDownloading https://ghcr.io/v2/homebrew/core/gts/manifests/0.7.6_2\u001b[0m\n",
      "######################################################################### 100.0%\n",
      "\u001b[34m==>\u001b[0m \u001b[1mDownloading https://ghcr.io/v2/homebrew/core/gts/blobs/sha256:c35739e0ed8143\u001b[0m\n",
      "######################################################################### 100.0%\n",
      "\u001b[32m==>\u001b[0m \u001b[1mFetching \u001b[32mlibx11\u001b[39m\u001b[0m\n",
      "\u001b[34m==>\u001b[0m \u001b[1mDownloading https://ghcr.io/v2/homebrew/core/libx11/manifests/1.8.6\u001b[0m\n",
      "######################################################################### 100.0%\n",
      "\u001b[34m==>\u001b[0m \u001b[1mDownloading https://ghcr.io/v2/homebrew/core/libx11/blobs/sha256:12634191390\u001b[0m\n",
      "######################################################################### 100.0%\n",
      "\u001b[32m==>\u001b[0m \u001b[1mFetching \u001b[32mgdk-pixbuf\u001b[39m\u001b[0m\n",
      "\u001b[34m==>\u001b[0m \u001b[1mDownloading https://ghcr.io/v2/homebrew/core/gdk-pixbuf/manifests/2.42.10_1\u001b[0m\n",
      "######################################################################### 100.0%\n",
      "\u001b[34m==>\u001b[0m \u001b[1mDownloading https://ghcr.io/v2/homebrew/core/gdk-pixbuf/blobs/sha256:fe74fd0\u001b[0m\n",
      "######################################################################### 100.0%\n",
      "\u001b[32m==>\u001b[0m \u001b[1mFetching \u001b[32mlibrsvg\u001b[39m\u001b[0m\n",
      "\u001b[34m==>\u001b[0m \u001b[1mDownloading https://ghcr.io/v2/homebrew/core/librsvg/manifests/2.56.1\u001b[0m\n",
      "######################################################################### 100.0%\n",
      "\u001b[34m==>\u001b[0m \u001b[1mDownloading https://ghcr.io/v2/homebrew/core/librsvg/blobs/sha256:c09f8982a8\u001b[0m\n",
      "######################################################################### 100.0%\n",
      "\u001b[32m==>\u001b[0m \u001b[1mFetching \u001b[32mm4\u001b[39m\u001b[0m\n",
      "\u001b[34m==>\u001b[0m \u001b[1mDownloading https://ghcr.io/v2/homebrew/core/m4/manifests/1.4.19\u001b[0m\n",
      "######################################################################### 100.0%\n",
      "\u001b[34m==>\u001b[0m \u001b[1mDownloading https://ghcr.io/v2/homebrew/core/m4/blobs/sha256:11308abe8d607be\u001b[0m\n",
      "######################################################################### 100.0%\n",
      "\u001b[32m==>\u001b[0m \u001b[1mFetching \u001b[32mlibtool\u001b[39m\u001b[0m\n",
      "\u001b[34m==>\u001b[0m \u001b[1mDownloading https://ghcr.io/v2/homebrew/core/libtool/manifests/2.4.7-1\u001b[0m\n",
      "######################################################################### 100.0%\n",
      "\u001b[34m==>\u001b[0m \u001b[1mDownloading https://ghcr.io/v2/homebrew/core/libtool/blobs/sha256:a7196b340a\u001b[0m\n",
      "######################################################################### 100.0%\n",
      "\u001b[32m==>\u001b[0m \u001b[1mFetching \u001b[32mgraphviz\u001b[39m\u001b[0m\n",
      "\u001b[34m==>\u001b[0m \u001b[1mDownloading https://ghcr.io/v2/homebrew/core/graphviz/manifests/8.0.5\u001b[0m\n",
      "######################################################################### 100.0%\n",
      "\u001b[34m==>\u001b[0m \u001b[1mDownloading https://ghcr.io/v2/homebrew/core/graphviz/blobs/sha256:d9351344f\u001b[0m\n",
      "######################################################################### 100.0%\n",
      "\u001b[32m==>\u001b[0m \u001b[1mInstalling dependencies for graphviz: \u001b[32mfreetype\u001b[39m, \u001b[32mimath\u001b[39m, \u001b[32mopenexr\u001b[39m, \u001b[32mjpeg-xl\u001b[39m, \u001b[32mlibavif\u001b[39m, \u001b[32mgd\u001b[39m, \u001b[32mjasper\u001b[39m, \u001b[32mnetpbm\u001b[39m, \u001b[32mgts\u001b[39m, \u001b[32mlibx11\u001b[39m, \u001b[32mgdk-pixbuf\u001b[39m, \u001b[32mlibrsvg\u001b[39m, \u001b[32mm4\u001b[39m and \u001b[32mlibtool\u001b[39m\u001b[0m\n",
      "\u001b[32m==>\u001b[0m \u001b[1mInstalling graphviz dependency: \u001b[32mfreetype\u001b[39m\u001b[0m\n",
      "\u001b[34m==>\u001b[0m \u001b[1mPouring freetype--2.13.1.arm64_ventura.bottle.tar.gz\u001b[0m\n",
      "🍺  /opt/homebrew/Cellar/freetype/2.13.1: 67 files, 2.4MB\n",
      "\u001b[32m==>\u001b[0m \u001b[1mInstalling graphviz dependency: \u001b[32mimath\u001b[39m\u001b[0m\n",
      "\u001b[34m==>\u001b[0m \u001b[1mPouring imath--3.1.9.arm64_ventura.bottle.tar.gz\u001b[0m\n",
      "🍺  /opt/homebrew/Cellar/imath/3.1.9: 49 files, 935.7KB\n",
      "\u001b[32m==>\u001b[0m \u001b[1mInstalling graphviz dependency: \u001b[32mopenexr\u001b[39m\u001b[0m\n",
      "\u001b[34m==>\u001b[0m \u001b[1mPouring openexr--3.1.9.arm64_ventura.bottle.tar.gz\u001b[0m\n",
      "🍺  /opt/homebrew/Cellar/openexr/3.1.9: 194 files, 7.7MB\n",
      "\u001b[32m==>\u001b[0m \u001b[1mInstalling graphviz dependency: \u001b[32mjpeg-xl\u001b[39m\u001b[0m\n",
      "\u001b[34m==>\u001b[0m \u001b[1mPouring jpeg-xl--0.8.2.arm64_ventura.bottle.tar.gz\u001b[0m\n",
      "🍺  /opt/homebrew/Cellar/jpeg-xl/0.8.2: 43 files, 10.8MB\n",
      "\u001b[32m==>\u001b[0m \u001b[1mInstalling graphviz dependency: \u001b[32mlibavif\u001b[39m\u001b[0m\n",
      "\u001b[34m==>\u001b[0m \u001b[1mPouring libavif--0.11.1.arm64_ventura.bottle.tar.gz\u001b[0m\n",
      "🍺  /opt/homebrew/Cellar/libavif/0.11.1: 19 files, 496.6KB\n",
      "\u001b[32m==>\u001b[0m \u001b[1mInstalling graphviz dependency: \u001b[32mgd\u001b[39m\u001b[0m\n",
      "\u001b[34m==>\u001b[0m \u001b[1mPouring gd--2.3.3_5.arm64_ventura.bottle.tar.gz\u001b[0m\n",
      "🍺  /opt/homebrew/Cellar/gd/2.3.3_5: 33 files, 1.5MB\n",
      "\u001b[32m==>\u001b[0m \u001b[1mInstalling graphviz dependency: \u001b[32mjasper\u001b[39m\u001b[0m\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[34m==>\u001b[0m \u001b[1mPouring jasper--4.0.0.arm64_ventura.bottle.tar.gz\u001b[0m\n",
      "🍺  /opt/homebrew/Cellar/jasper/4.0.0: 44 files, 1.6MB\n",
      "\u001b[32m==>\u001b[0m \u001b[1mInstalling graphviz dependency: \u001b[32mnetpbm\u001b[39m\u001b[0m\n",
      "\u001b[34m==>\u001b[0m \u001b[1mPouring netpbm--10.86.38_1.arm64_ventura.bottle.1.tar.gz\u001b[0m\n",
      "🍺  /opt/homebrew/Cellar/netpbm/10.86.38_1: 409 files, 18.5MB\n",
      "\u001b[32m==>\u001b[0m \u001b[1mInstalling graphviz dependency: \u001b[32mgts\u001b[39m\u001b[0m\n",
      "\u001b[34m==>\u001b[0m \u001b[1mPouring gts--0.7.6_2.arm64_ventura.bottle.tar.gz\u001b[0m\n",
      "🍺  /opt/homebrew/Cellar/gts/0.7.6_2: 25 files, 1.5MB\n",
      "\u001b[32m==>\u001b[0m \u001b[1mInstalling graphviz dependency: \u001b[32mlibx11\u001b[39m\u001b[0m\n",
      "\u001b[34m==>\u001b[0m \u001b[1mPouring libx11--1.8.6.arm64_ventura.bottle.tar.gz\u001b[0m\n",
      "🍺  /opt/homebrew/Cellar/libx11/1.8.6: 1,054 files, 7MB\n",
      "\u001b[32m==>\u001b[0m \u001b[1mInstalling graphviz dependency: \u001b[32mgdk-pixbuf\u001b[39m\u001b[0m\n",
      "\u001b[34m==>\u001b[0m \u001b[1mPouring gdk-pixbuf--2.42.10_1.arm64_ventura.bottle.tar.gz\u001b[0m\n",
      "\u001b[34m==>\u001b[0m \u001b[1m/opt/homebrew/Cellar/gdk-pixbuf/2.42.10_1/bin/gdk-pixbuf-query-loaders --upd\u001b[0m\n",
      "🍺  /opt/homebrew/Cellar/gdk-pixbuf/2.42.10_1: 150 files, 3.9MB\n",
      "\u001b[32m==>\u001b[0m \u001b[1mInstalling graphviz dependency: \u001b[32mlibrsvg\u001b[39m\u001b[0m\n",
      "\u001b[34m==>\u001b[0m \u001b[1mPouring librsvg--2.56.1.arm64_ventura.bottle.tar.gz\u001b[0m\n",
      "\u001b[34m==>\u001b[0m \u001b[1m/opt/homebrew/opt/gdk-pixbuf/bin/gdk-pixbuf-query-loaders --update-cache\u001b[0m\n",
      "🍺  /opt/homebrew/Cellar/librsvg/2.56.1: 22 files, 43.4MB\n",
      "\u001b[32m==>\u001b[0m \u001b[1mInstalling graphviz dependency: \u001b[32mm4\u001b[39m\u001b[0m\n",
      "\u001b[34m==>\u001b[0m \u001b[1mPouring m4--1.4.19.arm64_ventura.bottle.tar.gz\u001b[0m\n",
      "🍺  /opt/homebrew/Cellar/m4/1.4.19: 13 files, 742.4KB\n",
      "\u001b[32m==>\u001b[0m \u001b[1mInstalling graphviz dependency: \u001b[32mlibtool\u001b[39m\u001b[0m\n",
      "\u001b[34m==>\u001b[0m \u001b[1mPouring libtool--2.4.7.arm64_ventura.bottle.1.tar.gz\u001b[0m\n",
      "🍺  /opt/homebrew/Cellar/libtool/2.4.7: 75 files, 3.8MB\n",
      "\u001b[32m==>\u001b[0m \u001b[1mInstalling \u001b[32mgraphviz\u001b[39m\u001b[0m\n",
      "\u001b[34m==>\u001b[0m \u001b[1mPouring graphviz--8.0.5.arm64_ventura.bottle.tar.gz\u001b[0m\n",
      "🍺  /opt/homebrew/Cellar/graphviz/8.0.5: 283 files, 7MB\n",
      "\u001b[34m==>\u001b[0m \u001b[1mRunning `brew cleanup graphviz`...\u001b[0m\n",
      "Disable this behaviour by setting HOMEBREW_NO_INSTALL_CLEANUP.\n",
      "Hide these hints with HOMEBREW_NO_ENV_HINTS (see `man brew`).\n",
      "\u001b[34m==>\u001b[0m \u001b[1mUpgrading 1 dependent of upgraded formulae:\u001b[0m\n",
      "Disable this behaviour by setting HOMEBREW_NO_INSTALLED_DEPENDENTS_CHECK.\n",
      "Hide these hints with HOMEBREW_NO_ENV_HINTS (see `man brew`).\n",
      "qt@5 5.15.8_3 -> 5.15.10\n",
      "\u001b[32m==>\u001b[0m \u001b[1mFetching \u001b[32mqt@5\u001b[39m\u001b[0m\n",
      "\u001b[34m==>\u001b[0m \u001b[1mDownloading https://ghcr.io/v2/homebrew/core/qt/5/manifests/5.15.10-1\u001b[0m\n",
      "######################################################################### 100.0%\n",
      "\u001b[34m==>\u001b[0m \u001b[1mDownloading https://ghcr.io/v2/homebrew/core/qt/5/blobs/sha256:8caeb580012da\u001b[0m\n",
      "######################################################################### 100.0%\n",
      "\u001b[32m==>\u001b[0m \u001b[1mUpgrading \u001b[32mqt@5\u001b[39m\n",
      "  5.15.8_3 -> 5.15.10 \n",
      "\u001b[0m\n",
      "\u001b[34m==>\u001b[0m \u001b[1mPouring qt@5--5.15.10.arm64_ventura.bottle.1.tar.gz\u001b[0m\n",
      "🍺  /opt/homebrew/Cellar/qt@5/5.15.10: 10,845 files, 342.8MB\n",
      "\u001b[34m==>\u001b[0m \u001b[1mRunning `brew cleanup qt@5`...\u001b[0m\n",
      "Removing: /opt/homebrew/Cellar/qt@5/5.15.8_3... (10,840 files, 342.5MB)\n",
      "\u001b[32m==>\u001b[0m \u001b[1mChecking for dependents of upgraded formulae...\u001b[0m\n",
      "\u001b[34m==>\u001b[0m \u001b[1mNo broken dependents found!\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "!brew install graphviz\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "6943518f",
   "metadata": {},
   "outputs": [],
   "source": [
    "cat_features = joined >> Categorify()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "f9fcdd4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "ratings = nvt.ColumnGroup([\"rating\"]) >> LambdaOp(lambda col: (col > 3).astype(\"int8\")) >> AddTags(Tags.TARGET)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "586e5873",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/svg+xml": [
       "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
       "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
       " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
       "<!-- Generated by graphviz version 8.0.5 (20230430.1635)\n",
       " -->\n",
       "<!-- Pages: 1 -->\n",
       "<svg width=\"605pt\" height=\"548pt\"\n",
       " viewBox=\"0.00 0.00 604.73 548.00\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 544)\">\n",
       "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-544 600.73,-544 600.73,4 -4,4\"/>\n",
       "<!-- 0 -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>0</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"462.12\" cy=\"-234\" rx=\"57.49\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"462.12\" y=\"-228.95\" font-family=\"Times,serif\" font-size=\"14.00\">JoinExternal</text>\n",
       "</g>\n",
       "<!-- 2 -->\n",
       "<g id=\"node5\" class=\"node\">\n",
       "<title>2</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"423.12\" cy=\"-162\" rx=\"50.33\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"423.12\" y=\"-156.95\" font-family=\"Times,serif\" font-size=\"14.00\">Categorify</text>\n",
       "</g>\n",
       "<!-- 0&#45;&gt;2 -->\n",
       "<g id=\"edge3\" class=\"edge\">\n",
       "<title>0&#45;&gt;2</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M452.68,-216.05C448.18,-207.97 442.69,-198.12 437.67,-189.11\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"440.4,-187.82 432.48,-180.79 434.29,-191.23 440.4,-187.82\"/>\n",
       "</g>\n",
       "<!-- 7 -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>7</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"462.12\" cy=\"-306\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"462.12\" y=\"-300.95\" font-family=\"Times,serif\" font-size=\"14.00\">+</text>\n",
       "</g>\n",
       "<!-- 7&#45;&gt;0 -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>7&#45;&gt;0</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M462.12,-287.7C462.12,-280.24 462.12,-271.32 462.12,-262.97\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"465.62,-263.1 462.12,-253.1 458.62,-263.1 465.62,-263.1\"/>\n",
       "</g>\n",
       "<!-- 1 -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>1</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"391.12\" cy=\"-378\" rx=\"62.61\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"391.12\" y=\"-372.95\" font-family=\"Times,serif\" font-size=\"14.00\">TagAsUserID</text>\n",
       "</g>\n",
       "<!-- 1&#45;&gt;7 -->\n",
       "<g id=\"edge9\" class=\"edge\">\n",
       "<title>1&#45;&gt;7</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M407.95,-360.41C417.4,-351.1 429.31,-339.35 439.56,-329.24\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"441.68,-332.08 446.34,-322.57 436.77,-327.1 441.68,-332.08\"/>\n",
       "</g>\n",
       "<!-- 3 -->\n",
       "<g id=\"node4\" class=\"node\">\n",
       "<title>3</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"391.12\" cy=\"-450\" rx=\"56.47\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"391.12\" y=\"-444.95\" font-family=\"Times,serif\" font-size=\"14.00\">SelectionOp</text>\n",
       "</g>\n",
       "<!-- 3&#45;&gt;1 -->\n",
       "<g id=\"edge2\" class=\"edge\">\n",
       "<title>3&#45;&gt;1</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M391.12,-431.7C391.12,-424.24 391.12,-415.32 391.12,-406.97\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"394.62,-407.1 391.12,-397.1 387.62,-407.1 394.62,-407.1\"/>\n",
       "</g>\n",
       "<!-- 4 -->\n",
       "<g id=\"node7\" class=\"node\">\n",
       "<title>4</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"327.12\" cy=\"-90\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"327.12\" y=\"-84.95\" font-family=\"Times,serif\" font-size=\"14.00\">+</text>\n",
       "</g>\n",
       "<!-- 2&#45;&gt;4 -->\n",
       "<g id=\"edge5\" class=\"edge\">\n",
       "<title>2&#45;&gt;4</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M401.82,-145.46C387.5,-135.03 368.54,-121.2 353.31,-110.09\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"355.97,-106.98 345.83,-103.91 351.85,-112.63 355.97,-106.98\"/>\n",
       "</g>\n",
       "<!-- 3_selector -->\n",
       "<g id=\"node6\" class=\"node\">\n",
       "<title>3_selector</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"391.12\" cy=\"-522\" rx=\"43.16\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"391.12\" y=\"-516.95\" font-family=\"Times,serif\" font-size=\"14.00\">[&#39;userId&#39;]</text>\n",
       "</g>\n",
       "<!-- 3_selector&#45;&gt;3 -->\n",
       "<g id=\"edge4\" class=\"edge\">\n",
       "<title>3_selector&#45;&gt;3</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M391.12,-503.7C391.12,-496.24 391.12,-487.32 391.12,-478.97\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"394.62,-479.1 391.12,-469.1 387.62,-479.1 394.62,-479.1\"/>\n",
       "</g>\n",
       "<!-- 11 -->\n",
       "<g id=\"node15\" class=\"node\">\n",
       "<title>11</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"327.12\" cy=\"-18\" rx=\"51.86\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"327.12\" y=\"-12.95\" font-family=\"Times,serif\" font-size=\"14.00\">output cols</text>\n",
       "</g>\n",
       "<!-- 4&#45;&gt;11 -->\n",
       "<g id=\"edge14\" class=\"edge\">\n",
       "<title>4&#45;&gt;11</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M327.12,-71.7C327.12,-64.24 327.12,-55.32 327.12,-46.97\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"330.62,-47.1 327.12,-37.1 323.62,-47.1 330.62,-47.1\"/>\n",
       "</g>\n",
       "<!-- 5 -->\n",
       "<g id=\"node8\" class=\"node\">\n",
       "<title>5</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"271.12\" cy=\"-162\" rx=\"44.7\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"271.12\" y=\"-156.95\" font-family=\"Times,serif\" font-size=\"14.00\">AddTags</text>\n",
       "</g>\n",
       "<!-- 5&#45;&gt;4 -->\n",
       "<g id=\"edge6\" class=\"edge\">\n",
       "<title>5&#45;&gt;4</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M284.39,-144.41C291.53,-135.5 300.44,-124.35 308.28,-114.56\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"311.58,-117.03 315.09,-107.04 306.11,-112.66 311.58,-117.03\"/>\n",
       "</g>\n",
       "<!-- 9 -->\n",
       "<g id=\"node9\" class=\"node\">\n",
       "<title>9</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"193.12\" cy=\"-234\" rx=\"193.12\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"193.12\" y=\"-228.95\" font-family=\"Times,serif\" font-size=\"14.00\">LambdaOp(lambda col: (col &gt; 3).astype(&quot;int8&quot;))</text>\n",
       "</g>\n",
       "<!-- 9&#45;&gt;5 -->\n",
       "<g id=\"edge7\" class=\"edge\">\n",
       "<title>9&#45;&gt;5</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M212.4,-215.7C222.4,-206.73 234.74,-195.65 245.52,-185.98\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"247.29,-189.19 252.4,-179.91 242.62,-183.98 247.29,-189.19\"/>\n",
       "</g>\n",
       "<!-- 6 -->\n",
       "<g id=\"node10\" class=\"node\">\n",
       "<title>6</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"534.12\" cy=\"-450\" rx=\"56.47\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"534.12\" y=\"-444.95\" font-family=\"Times,serif\" font-size=\"14.00\">SelectionOp</text>\n",
       "</g>\n",
       "<!-- 10 -->\n",
       "<g id=\"node12\" class=\"node\">\n",
       "<title>10</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"534.12\" cy=\"-378\" rx=\"62.61\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"534.12\" y=\"-372.95\" font-family=\"Times,serif\" font-size=\"14.00\">TagAsItemID</text>\n",
       "</g>\n",
       "<!-- 6&#45;&gt;10 -->\n",
       "<g id=\"edge13\" class=\"edge\">\n",
       "<title>6&#45;&gt;10</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M534.12,-431.7C534.12,-424.24 534.12,-415.32 534.12,-406.97\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"537.62,-407.1 534.12,-397.1 530.62,-407.1 537.62,-407.1\"/>\n",
       "</g>\n",
       "<!-- 6_selector -->\n",
       "<g id=\"node11\" class=\"node\">\n",
       "<title>6_selector</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"534.12\" cy=\"-522\" rx=\"51.35\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"534.12\" y=\"-516.95\" font-family=\"Times,serif\" font-size=\"14.00\">[&#39;movieId&#39;]</text>\n",
       "</g>\n",
       "<!-- 6_selector&#45;&gt;6 -->\n",
       "<g id=\"edge8\" class=\"edge\">\n",
       "<title>6_selector&#45;&gt;6</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M534.12,-503.7C534.12,-496.24 534.12,-487.32 534.12,-478.97\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"537.62,-479.1 534.12,-469.1 530.62,-479.1 537.62,-479.1\"/>\n",
       "</g>\n",
       "<!-- 10&#45;&gt;7 -->\n",
       "<g id=\"edge10\" class=\"edge\">\n",
       "<title>10&#45;&gt;7</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M517.06,-360.41C507.25,-350.87 494.82,-338.79 484.26,-328.53\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"487.24,-326.57 477.63,-322.11 482.36,-331.59 487.24,-326.57\"/>\n",
       "</g>\n",
       "<!-- 8 -->\n",
       "<g id=\"node13\" class=\"node\">\n",
       "<title>8</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"193.12\" cy=\"-306\" rx=\"56.47\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"193.12\" y=\"-300.95\" font-family=\"Times,serif\" font-size=\"14.00\">SelectionOp</text>\n",
       "</g>\n",
       "<!-- 8&#45;&gt;9 -->\n",
       "<g id=\"edge12\" class=\"edge\">\n",
       "<title>8&#45;&gt;9</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M193.12,-287.7C193.12,-280.24 193.12,-271.32 193.12,-262.97\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"196.62,-263.1 193.12,-253.1 189.62,-263.1 196.62,-263.1\"/>\n",
       "</g>\n",
       "<!-- 8_selector -->\n",
       "<g id=\"node14\" class=\"node\">\n",
       "<title>8_selector</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"193.12\" cy=\"-378\" rx=\"41.63\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"193.12\" y=\"-372.95\" font-family=\"Times,serif\" font-size=\"14.00\">[&#39;rating&#39;]</text>\n",
       "</g>\n",
       "<!-- 8_selector&#45;&gt;8 -->\n",
       "<g id=\"edge11\" class=\"edge\">\n",
       "<title>8_selector&#45;&gt;8</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M193.12,-359.7C193.12,-352.24 193.12,-343.32 193.12,-334.97\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"196.62,-335.1 193.12,-325.1 189.62,-335.1 196.62,-335.1\"/>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "<graphviz.graphs.Digraph at 0x15dce29a0>"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output = cat_features + ratings\n",
    "(output).graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "42096c96",
   "metadata": {},
   "outputs": [],
   "source": [
    "workflow = nvt.Workflow(output)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "391c05a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "dict_dtypes = {}\n",
    "\n",
    "for col in CATEGORICAL_COLUMNS:\n",
    "    dict_dtypes[col] = np.int64\n",
    "\n",
    "for col in LABEL_COLUMNS:\n",
    "    dict_dtypes[col] = np.float32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "140e3300",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/rajathdb/miniconda3/envs/rl/lib/python3.8/site-packages/merlin/io/dataset.py:267: UserWarning: Initializing an NVTabular Dataset in CPU mode.This is an experimental feature with extremely limited support!\n",
      "  warnings.warn(\n",
      "/Users/rajathdb/miniconda3/envs/rl/lib/python3.8/site-packages/merlin/io/dataset.py:267: UserWarning: Initializing an NVTabular Dataset in CPU mode.This is an experimental feature with extremely limited support!\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "train_dataset = nvt.Dataset([os.path.join(INPUT_DATA_DIR, \"train.parquet\")])\n",
    "valid_dataset = nvt.Dataset([os.path.join(INPUT_DATA_DIR, \"valid.parquet\")])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "983b6a62",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/rajathdb/miniconda3/envs/rl/lib/python3.8/site-packages/merlin/io/dataset.py:457: UserWarning: Changing an NVTabular Dataset to CPU mode.This is an experimental feature with extremely limited support!\n",
      "  warnings.warn(\n",
      "/Users/rajathdb/miniconda3/envs/rl/lib/python3.8/site-packages/merlin/io/dataset.py:267: UserWarning: Initializing an NVTabular Dataset in CPU mode.This is an experimental feature with extremely limited support!\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 34 s, sys: 10.9 s, total: 44.8 s\n",
      "Wall time: 47.9 s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<nvtabular.workflow.workflow.Workflow at 0x15df86dc0>"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "workflow.fit(train_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "2f947891",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make sure we have a clean output path\n",
    "if path.exists(os.path.join(INPUT_DATA_DIR, \"train\")):\n",
    "    shutil.rmtree(os.path.join(INPUT_DATA_DIR, \"train\"))\n",
    "if path.exists(os.path.join(INPUT_DATA_DIR, \"valid\")):\n",
    "    shutil.rmtree(os.path.join(INPUT_DATA_DIR, \"valid\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "5d7348a7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 4 µs, sys: 1 µs, total: 5 µs\n",
      "Wall time: 14.1 µs\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/rajathdb/miniconda3/envs/rl/lib/python3.8/site-packages/merlin/io/dataset.py:267: UserWarning: Initializing an NVTabular Dataset in CPU mode.This is an experimental feature with extremely limited support!\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "%time\n",
    "workflow.transform(train_dataset).to_parquet(\n",
    "    output_path=os.path.join(INPUT_DATA_DIR, \"train\"),\n",
    "    shuffle=nvt.io.Shuffle.PER_PARTITION,\n",
    "    cats=[\"userId\", \"movieId\", \"genres\"],\n",
    "    labels=[\"rating\"],\n",
    "    dtypes=dict_dtypes,\n",
    "    write_hugectr_keyset=False  # only needed if using this ETL Notebook for training with HugeCTR\n",
    "                               # should be removed otherwise to speed up computation\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "82641a5c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 3 µs, sys: 2 µs, total: 5 µs\n",
      "Wall time: 17.9 µs\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/rajathdb/miniconda3/envs/rl/lib/python3.8/site-packages/merlin/io/dataset.py:267: UserWarning: Initializing an NVTabular Dataset in CPU mode.This is an experimental feature with extremely limited support!\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "%time\n",
    "workflow.transform(valid_dataset).to_parquet(\n",
    "    output_path=os.path.join(INPUT_DATA_DIR, \"valid\"),\n",
    "    shuffle=False,\n",
    "    cats=[\"userId\", \"movieId\", \"genres\"],\n",
    "    labels=[\"rating\"],\n",
    "    dtypes=dict_dtypes,\n",
    "    write_hugectr_keyset=True  # only needed if using this ETL Notebook for training with HugeCTR\n",
    "                               # should be removed otherwise to speed up computation\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "ddc4b00b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>name</th>\n",
       "      <th>tags</th>\n",
       "      <th>dtype</th>\n",
       "      <th>is_list</th>\n",
       "      <th>is_ragged</th>\n",
       "      <th>properties.num_buckets</th>\n",
       "      <th>properties.freq_threshold</th>\n",
       "      <th>properties.max_size</th>\n",
       "      <th>properties.cat_path</th>\n",
       "      <th>properties.domain.min</th>\n",
       "      <th>properties.domain.max</th>\n",
       "      <th>properties.domain.name</th>\n",
       "      <th>properties.embedding_sizes.cardinality</th>\n",
       "      <th>properties.embedding_sizes.dimension</th>\n",
       "      <th>properties.value_count.min</th>\n",
       "      <th>properties.value_count.max</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>userId</td>\n",
       "      <td>(Tags.CATEGORICAL)</td>\n",
       "      <td>DType(name='int64', element_type=&lt;ElementType....</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>.//categories/unique.userId.parquet</td>\n",
       "      <td>0.0</td>\n",
       "      <td>162543.0</td>\n",
       "      <td>userId</td>\n",
       "      <td>162544.0</td>\n",
       "      <td>512.0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>movieId</td>\n",
       "      <td>(Tags.CATEGORICAL)</td>\n",
       "      <td>DType(name='int64', element_type=&lt;ElementType....</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>.//categories/unique.movieId.parquet</td>\n",
       "      <td>0.0</td>\n",
       "      <td>56639.0</td>\n",
       "      <td>movieId</td>\n",
       "      <td>56640.0</td>\n",
       "      <td>512.0</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>genres</td>\n",
       "      <td>(Tags.CATEGORICAL)</td>\n",
       "      <td>DType(name='int64', element_type=&lt;ElementType....</td>\n",
       "      <td>True</td>\n",
       "      <td>True</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>.//categories/unique.genres.parquet</td>\n",
       "      <td>0.0</td>\n",
       "      <td>22.0</td>\n",
       "      <td>genres</td>\n",
       "      <td>23.0</td>\n",
       "      <td>16.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>rating</td>\n",
       "      <td>(Tags.TARGET)</td>\n",
       "      <td>DType(name='int8', element_type=&lt;ElementType.I...</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "[{'name': 'userId', 'tags': {<Tags.CATEGORICAL: 'categorical'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0, 'max_size': 0, 'cat_path': './/categories/unique.userId.parquet', 'domain': {'min': 0, 'max': 162543, 'name': 'userId'}, 'embedding_sizes': {'cardinality': 162544, 'dimension': 512}}, 'dtype': DType(name='int64', element_type=<ElementType.Int: 'int'>, element_size=64, element_unit=None, signed=True, shape=Shape(dims=(Dimension(min=0, max=None),))), 'is_list': False, 'is_ragged': False}, {'name': 'movieId', 'tags': {<Tags.CATEGORICAL: 'categorical'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0, 'max_size': 0, 'cat_path': './/categories/unique.movieId.parquet', 'domain': {'min': 0, 'max': 56639, 'name': 'movieId'}, 'embedding_sizes': {'cardinality': 56640, 'dimension': 512}}, 'dtype': DType(name='int64', element_type=<ElementType.Int: 'int'>, element_size=64, element_unit=None, signed=True, shape=Shape(dims=(Dimension(min=0, max=None),))), 'is_list': False, 'is_ragged': False}, {'name': 'genres', 'tags': {<Tags.CATEGORICAL: 'categorical'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0, 'max_size': 0, 'cat_path': './/categories/unique.genres.parquet', 'domain': {'min': 0, 'max': 22, 'name': 'genres'}, 'embedding_sizes': {'cardinality': 23, 'dimension': 16}, 'value_count': {'min': 0, 'max': None}}, 'dtype': DType(name='int64', element_type=<ElementType.Int: 'int'>, element_size=64, element_unit=None, signed=True, shape=Shape(dims=(Dimension(min=0, max=None), Dimension(min=0, max=None)))), 'is_list': True, 'is_ragged': True}, {'name': 'rating', 'tags': {<Tags.TARGET: 'target'>}, 'properties': {}, 'dtype': DType(name='int8', element_type=<ElementType.Int: 'int'>, element_size=8, element_unit=None, signed=True, shape=Shape(dims=(Dimension(min=0, max=None),))), 'is_list': False, 'is_ragged': False}]"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "workflow.save(os.path.join(INPUT_DATA_DIR, \"workflow\"))\n",
    "workflow.output_schema"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "0b95c15c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(['nvtabular/movielens/data/train/part_0.parquet'],\n",
       " ['nvtabular/movielens/data/valid/part_0.parquet'])"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import glob\n",
    "\n",
    "TRAIN_PATHS = sorted(glob.glob(os.path.join(INPUT_DATA_DIR, \"train\", \"*.parquet\")))\n",
    "VALID_PATHS = sorted(glob.glob(os.path.join(INPUT_DATA_DIR, \"valid\", \"*.parquet\")))\n",
    "TRAIN_PATHS, VALID_PATHS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "2a00ccc6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>userId</th>\n",
       "      <th>movieId</th>\n",
       "      <th>genres</th>\n",
       "      <th>rating</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>434</td>\n",
       "      <td>5492</td>\n",
       "      <td>[3, 6]</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>48824</td>\n",
       "      <td>1548</td>\n",
       "      <td>[5, 4, 10]</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>36806</td>\n",
       "      <td>1729</td>\n",
       "      <td>[4, 3]</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>40682</td>\n",
       "      <td>4625</td>\n",
       "      <td>[4, 3, 8]</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>7569</td>\n",
       "      <td>1049</td>\n",
       "      <td>[5, 7, 4, 8]</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   userId  movieId        genres  rating\n",
       "0     434     5492        [3, 6]     0.0\n",
       "1   48824     1548    [5, 4, 10]     1.0\n",
       "2   36806     1729        [4, 3]     0.0\n",
       "3   40682     4625     [4, 3, 8]     0.0\n",
       "4    7569     1049  [5, 7, 4, 8]     0.0"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = df_lib.read_parquet(TRAIN_PATHS[0])\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "f043ae36",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['nvtabular/movielens/data/train/part_0.parquet']"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "TRAIN_PATHS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8356315b",
   "metadata": {},
   "outputs": [],
   "source": [
    "!export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4234a6e2",
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "Descriptors cannot not be created directly.\nIf this call came from a _pb2.py file, your generated code is out of date and must be regenerated with protoc >= 3.19.0.\nIf you cannot immediately regenerate your protos, some other possible workarounds are:\n 1. Downgrade the protobuf package to 3.20.x or lower.\n 2. Set PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python (but this will use pure-Python parsing and will be much slower).\n\nMore information: https://developers.google.com/protocol-buffers/docs/news/2022-05-06#python-updates",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[4], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mmerlin\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mschema\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Schema\n\u001b[1;32m      2\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mmerlin\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mio\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Dataset\n\u001b[1;32m      4\u001b[0m train \u001b[38;5;241m=\u001b[39m Dataset(os\u001b[38;5;241m.\u001b[39mpath\u001b[38;5;241m.\u001b[39mjoin(INPUT_DATA_DIR, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtrain/part_0.parquet\u001b[39m\u001b[38;5;124m\"\u001b[39m))\n",
      "File \u001b[0;32m~/miniconda3/envs/rl/lib/python3.8/site-packages/merlin/schema/__init__.py:18\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;66;03m#\u001b[39;00m\n\u001b[1;32m      2\u001b[0m \u001b[38;5;66;03m# Copyright (c) 2022, NVIDIA CORPORATION.\u001b[39;00m\n\u001b[1;32m      3\u001b[0m \u001b[38;5;66;03m#\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     16\u001b[0m \n\u001b[1;32m     17\u001b[0m \u001b[38;5;66;03m# flake8: noqa\u001b[39;00m\n\u001b[0;32m---> 18\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mmerlin\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mschema\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mschema\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m ColumnSchema, Schema\n\u001b[1;32m     19\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mmerlin\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mschema\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mtags\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Tags, TagSet, TagsType\n",
      "File \u001b[0;32m~/miniconda3/envs/rl/lib/python3.8/site-packages/merlin/schema/schema.py:22\u001b[0m\n\u001b[1;32m     18\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtyping\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Dict, List, Optional, Text, Tuple, Union\n\u001b[1;32m     20\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mpandas\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mpd\u001b[39;00m\n\u001b[0;32m---> 22\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mmerlin\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdtypes\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mmd\u001b[39;00m\n\u001b[1;32m     23\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mmerlin\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdtypes\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m DType\n\u001b[1;32m     24\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mmerlin\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdtypes\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mshape\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Shape\n",
      "File \u001b[0;32m~/miniconda3/envs/rl/lib/python3.8/site-packages/merlin/dtypes/__init__.py:18\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;66;03m#\u001b[39;00m\n\u001b[1;32m      2\u001b[0m \u001b[38;5;66;03m# Copyright (c) 2022, NVIDIA CORPORATION.\u001b[39;00m\n\u001b[1;32m      3\u001b[0m \u001b[38;5;66;03m#\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     16\u001b[0m \n\u001b[1;32m     17\u001b[0m \u001b[38;5;66;03m# flake8: noqa\u001b[39;00m\n\u001b[0;32m---> 18\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mmerlin\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdtypes\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m aliases, mappings\n\u001b[1;32m     19\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mmerlin\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdtypes\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01maliases\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;241m*\u001b[39m\n\u001b[1;32m     20\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mmerlin\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdtypes\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mbase\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m DType\n",
      "File \u001b[0;32m~/miniconda3/envs/rl/lib/python3.8/site-packages/merlin/dtypes/mappings/__init__.py:18\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;66;03m#\u001b[39;00m\n\u001b[1;32m      2\u001b[0m \u001b[38;5;66;03m# Copyright (c) 2022, NVIDIA CORPORATION.\u001b[39;00m\n\u001b[1;32m      3\u001b[0m \u001b[38;5;66;03m#\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     16\u001b[0m \n\u001b[1;32m     17\u001b[0m \u001b[38;5;66;03m# flake8: noqa\u001b[39;00m\n\u001b[0;32m---> 18\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mmerlin\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdtypes\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mmappings\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m cudf, merlin, numpy, pandas, python, tf, torch, triton\n",
      "File \u001b[0;32m~/miniconda3/envs/rl/lib/python3.8/site-packages/merlin/dtypes/mappings/tf.py:21\u001b[0m\n\u001b[1;32m     18\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mmerlin\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdtypes\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mregistry\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m _dtype_registry\n\u001b[1;32m     20\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m---> 21\u001b[0m     \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m dtypes \u001b[38;5;28;01mas\u001b[39;00m tf_dtypes\n\u001b[1;32m     23\u001b[0m     tf_dtypes \u001b[38;5;241m=\u001b[39m DTypeMapping(\n\u001b[1;32m     24\u001b[0m         {\n\u001b[1;32m     25\u001b[0m             \u001b[38;5;66;03m# Unsigned Integer\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     45\u001b[0m         ),\n\u001b[1;32m     46\u001b[0m     )\n\u001b[1;32m     47\u001b[0m     _dtype_registry\u001b[38;5;241m.\u001b[39mregister(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtf\u001b[39m\u001b[38;5;124m\"\u001b[39m, tf_dtypes)\n",
      "File \u001b[0;32m~/miniconda3/envs/rl/lib/python3.8/site-packages/tensorflow/__init__.py:37\u001b[0m\n\u001b[1;32m     34\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01msys\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01m_sys\u001b[39;00m\n\u001b[1;32m     35\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mtyping\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01m_typing\u001b[39;00m\n\u001b[0;32m---> 37\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mpython\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mtools\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m module_util \u001b[38;5;28;01mas\u001b[39;00m _module_util\n\u001b[1;32m     38\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mpython\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mutil\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mlazy_loader\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m LazyLoader \u001b[38;5;28;01mas\u001b[39;00m _LazyLoader\n\u001b[1;32m     40\u001b[0m \u001b[38;5;66;03m# Make sure code inside the TensorFlow codebase can use tf2.enabled() at import.\u001b[39;00m\n",
      "File \u001b[0;32m~/miniconda3/envs/rl/lib/python3.8/site-packages/tensorflow/python/__init__.py:37\u001b[0m\n\u001b[1;32m     29\u001b[0m \u001b[38;5;66;03m# We aim to keep this file minimal and ideally remove completely.\u001b[39;00m\n\u001b[1;32m     30\u001b[0m \u001b[38;5;66;03m# If you are adding a new file with @tf_export decorators,\u001b[39;00m\n\u001b[1;32m     31\u001b[0m \u001b[38;5;66;03m# import it in modules_with_exports.py instead.\u001b[39;00m\n\u001b[1;32m     32\u001b[0m \n\u001b[1;32m     33\u001b[0m \u001b[38;5;66;03m# go/tf-wildcard-import\u001b[39;00m\n\u001b[1;32m     34\u001b[0m \u001b[38;5;66;03m# pylint: disable=wildcard-import,g-bad-import-order,g-import-not-at-top\u001b[39;00m\n\u001b[1;32m     36\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mpython\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m pywrap_tensorflow \u001b[38;5;28;01mas\u001b[39;00m _pywrap_tensorflow\n\u001b[0;32m---> 37\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mpython\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01meager\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m context\n\u001b[1;32m     39\u001b[0m \u001b[38;5;66;03m# pylint: enable=wildcard-import\u001b[39;00m\n\u001b[1;32m     40\u001b[0m \n\u001b[1;32m     41\u001b[0m \u001b[38;5;66;03m# Bring in subpackages.\u001b[39;00m\n\u001b[1;32m     42\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mpython\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m data\n",
      "File \u001b[0;32m~/miniconda3/envs/rl/lib/python3.8/site-packages/tensorflow/python/eager/context.py:28\u001b[0m\n\u001b[1;32m     25\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mabsl\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m logging\n\u001b[1;32m     26\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mnumpy\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mnp\u001b[39;00m\n\u001b[0;32m---> 28\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mframework\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m function_pb2\n\u001b[1;32m     29\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mprotobuf\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m config_pb2\n\u001b[1;32m     30\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mprotobuf\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m coordination_config_pb2\n",
      "File \u001b[0;32m~/miniconda3/envs/rl/lib/python3.8/site-packages/tensorflow/core/framework/function_pb2.py:16\u001b[0m\n\u001b[1;32m     11\u001b[0m \u001b[38;5;66;03m# @@protoc_insertion_point(imports)\u001b[39;00m\n\u001b[1;32m     13\u001b[0m _sym_db \u001b[38;5;241m=\u001b[39m _symbol_database\u001b[38;5;241m.\u001b[39mDefault()\n\u001b[0;32m---> 16\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mframework\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m attr_value_pb2 \u001b[38;5;28;01mas\u001b[39;00m tensorflow_dot_core_dot_framework_dot_attr__value__pb2\n\u001b[1;32m     17\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mframework\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m node_def_pb2 \u001b[38;5;28;01mas\u001b[39;00m tensorflow_dot_core_dot_framework_dot_node__def__pb2\n\u001b[1;32m     18\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mframework\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m op_def_pb2 \u001b[38;5;28;01mas\u001b[39;00m tensorflow_dot_core_dot_framework_dot_op__def__pb2\n",
      "File \u001b[0;32m~/miniconda3/envs/rl/lib/python3.8/site-packages/tensorflow/core/framework/attr_value_pb2.py:16\u001b[0m\n\u001b[1;32m     11\u001b[0m \u001b[38;5;66;03m# @@protoc_insertion_point(imports)\u001b[39;00m\n\u001b[1;32m     13\u001b[0m _sym_db \u001b[38;5;241m=\u001b[39m _symbol_database\u001b[38;5;241m.\u001b[39mDefault()\n\u001b[0;32m---> 16\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mframework\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m tensor_pb2 \u001b[38;5;28;01mas\u001b[39;00m tensorflow_dot_core_dot_framework_dot_tensor__pb2\n\u001b[1;32m     17\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mframework\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m tensor_shape_pb2 \u001b[38;5;28;01mas\u001b[39;00m tensorflow_dot_core_dot_framework_dot_tensor__shape__pb2\n\u001b[1;32m     18\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mframework\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m types_pb2 \u001b[38;5;28;01mas\u001b[39;00m tensorflow_dot_core_dot_framework_dot_types__pb2\n",
      "File \u001b[0;32m~/miniconda3/envs/rl/lib/python3.8/site-packages/tensorflow/core/framework/tensor_pb2.py:16\u001b[0m\n\u001b[1;32m     11\u001b[0m \u001b[38;5;66;03m# @@protoc_insertion_point(imports)\u001b[39;00m\n\u001b[1;32m     13\u001b[0m _sym_db \u001b[38;5;241m=\u001b[39m _symbol_database\u001b[38;5;241m.\u001b[39mDefault()\n\u001b[0;32m---> 16\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mframework\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m resource_handle_pb2 \u001b[38;5;28;01mas\u001b[39;00m tensorflow_dot_core_dot_framework_dot_resource__handle__pb2\n\u001b[1;32m     17\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mframework\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m tensor_shape_pb2 \u001b[38;5;28;01mas\u001b[39;00m tensorflow_dot_core_dot_framework_dot_tensor__shape__pb2\n\u001b[1;32m     18\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mframework\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m types_pb2 \u001b[38;5;28;01mas\u001b[39;00m tensorflow_dot_core_dot_framework_dot_types__pb2\n",
      "File \u001b[0;32m~/miniconda3/envs/rl/lib/python3.8/site-packages/tensorflow/core/framework/resource_handle_pb2.py:16\u001b[0m\n\u001b[1;32m     11\u001b[0m \u001b[38;5;66;03m# @@protoc_insertion_point(imports)\u001b[39;00m\n\u001b[1;32m     13\u001b[0m _sym_db \u001b[38;5;241m=\u001b[39m _symbol_database\u001b[38;5;241m.\u001b[39mDefault()\n\u001b[0;32m---> 16\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mframework\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m tensor_shape_pb2 \u001b[38;5;28;01mas\u001b[39;00m tensorflow_dot_core_dot_framework_dot_tensor__shape__pb2\n\u001b[1;32m     17\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtensorflow\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mframework\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m types_pb2 \u001b[38;5;28;01mas\u001b[39;00m tensorflow_dot_core_dot_framework_dot_types__pb2\n\u001b[1;32m     20\u001b[0m DESCRIPTOR \u001b[38;5;241m=\u001b[39m _descriptor\u001b[38;5;241m.\u001b[39mFileDescriptor(\n\u001b[1;32m     21\u001b[0m   name\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mtensorflow/core/framework/resource_handle.proto\u001b[39m\u001b[38;5;124m'\u001b[39m,\n\u001b[1;32m     22\u001b[0m   package\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mtensorflow\u001b[39m\u001b[38;5;124m'\u001b[39m,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     26\u001b[0m   ,\n\u001b[1;32m     27\u001b[0m   dependencies\u001b[38;5;241m=\u001b[39m[tensorflow_dot_core_dot_framework_dot_tensor__shape__pb2\u001b[38;5;241m.\u001b[39mDESCRIPTOR,tensorflow_dot_core_dot_framework_dot_types__pb2\u001b[38;5;241m.\u001b[39mDESCRIPTOR,])\n",
      "File \u001b[0;32m~/miniconda3/envs/rl/lib/python3.8/site-packages/tensorflow/core/framework/tensor_shape_pb2.py:36\u001b[0m\n\u001b[1;32m     13\u001b[0m _sym_db \u001b[38;5;241m=\u001b[39m _symbol_database\u001b[38;5;241m.\u001b[39mDefault()\n\u001b[1;32m     18\u001b[0m DESCRIPTOR \u001b[38;5;241m=\u001b[39m _descriptor\u001b[38;5;241m.\u001b[39mFileDescriptor(\n\u001b[1;32m     19\u001b[0m   name\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mtensorflow/core/framework/tensor_shape.proto\u001b[39m\u001b[38;5;124m'\u001b[39m,\n\u001b[1;32m     20\u001b[0m   package\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mtensorflow\u001b[39m\u001b[38;5;124m'\u001b[39m,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     23\u001b[0m   serialized_pb\u001b[38;5;241m=\u001b[39m_b(\u001b[38;5;124m'\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;124m,tensorflow/core/framework/tensor_shape.proto\u001b[39m\u001b[38;5;130;01m\\x12\u001b[39;00m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;124mtensorflow\u001b[39m\u001b[38;5;130;01m\\\"\u001b[39;00m\u001b[38;5;124mz\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;130;01m\\x10\u001b[39;00m\u001b[38;5;124mTensorShapeProto\u001b[39m\u001b[38;5;130;01m\\x12\u001b[39;00m\u001b[38;5;124m-\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;130;01m\\x03\u001b[39;00m\u001b[38;5;130;01m\\x64\u001b[39;00m\u001b[38;5;124mim\u001b[39m\u001b[38;5;130;01m\\x18\u001b[39;00m\u001b[38;5;130;01m\\x02\u001b[39;00m\u001b[38;5;124m \u001b[39m\u001b[38;5;130;01m\\x03\u001b[39;00m\u001b[38;5;124m(\u001b[39m\u001b[38;5;130;01m\\x0b\u001b[39;00m\u001b[38;5;130;01m\\x32\u001b[39;00m\u001b[38;5;124m .tensorflow.TensorShapeProto.Dim\u001b[39m\u001b[38;5;130;01m\\x12\u001b[39;00m\u001b[38;5;130;01m\\x14\u001b[39;00m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;130;01m\\x0c\u001b[39;00m\u001b[38;5;124munknown_rank\u001b[39m\u001b[38;5;130;01m\\x18\u001b[39;00m\u001b[38;5;130;01m\\x03\u001b[39;00m\u001b[38;5;124m \u001b[39m\u001b[38;5;130;01m\\x01\u001b[39;00m\u001b[38;5;124m(\u001b[39m\u001b[38;5;130;01m\\x08\u001b[39;00m\u001b[38;5;130;01m\\x1a\u001b[39;00m\u001b[38;5;124m!\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;130;01m\\x03\u001b[39;00m\u001b[38;5;130;01m\\x44\u001b[39;00m\u001b[38;5;124mim\u001b[39m\u001b[38;5;130;01m\\x12\u001b[39;00m\u001b[38;5;130;01m\\x0c\u001b[39;00m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;130;01m\\x04\u001b[39;00m\u001b[38;5;124msize\u001b[39m\u001b[38;5;130;01m\\x18\u001b[39;00m\u001b[38;5;130;01m\\x01\u001b[39;00m\u001b[38;5;124m \u001b[39m\u001b[38;5;130;01m\\x01\u001b[39;00m\u001b[38;5;124m(\u001b[39m\u001b[38;5;130;01m\\x03\u001b[39;00m\u001b[38;5;130;01m\\x12\u001b[39;00m\u001b[38;5;130;01m\\x0c\u001b[39;00m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;130;01m\\x04\u001b[39;00m\u001b[38;5;124mname\u001b[39m\u001b[38;5;130;01m\\x18\u001b[39;00m\u001b[38;5;130;01m\\x02\u001b[39;00m\u001b[38;5;124m \u001b[39m\u001b[38;5;130;01m\\x01\u001b[39;00m\u001b[38;5;124m(\u001b[39m\u001b[38;5;130;01m\\t\u001b[39;00m\u001b[38;5;124mB\u001b[39m\u001b[38;5;130;01m\\x87\u001b[39;00m\u001b[38;5;130;01m\\x01\u001b[39;00m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;130;01m\\x18\u001b[39;00m\u001b[38;5;124morg.tensorflow.frameworkB\u001b[39m\u001b[38;5;130;01m\\x11\u001b[39;00m\u001b[38;5;124mTensorShapeProtosP\u001b[39m\u001b[38;5;130;01m\\x01\u001b[39;00m\u001b[38;5;124mZSgithub.com/tensorflow/tensorflow/tensorflow/go/core/framework/tensor_shape_go_proto\u001b[39m\u001b[38;5;130;01m\\xf8\u001b[39;00m\u001b[38;5;130;01m\\x01\u001b[39;00m\u001b[38;5;130;01m\\x01\u001b[39;00m\u001b[38;5;130;01m\\x62\u001b[39;00m\u001b[38;5;130;01m\\x06\u001b[39;00m\u001b[38;5;124mproto3\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[1;32m     24\u001b[0m )\n\u001b[1;32m     29\u001b[0m _TENSORSHAPEPROTO_DIM \u001b[38;5;241m=\u001b[39m _descriptor\u001b[38;5;241m.\u001b[39mDescriptor(\n\u001b[1;32m     30\u001b[0m   name\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mDim\u001b[39m\u001b[38;5;124m'\u001b[39m,\n\u001b[1;32m     31\u001b[0m   full_name\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mtensorflow.TensorShapeProto.Dim\u001b[39m\u001b[38;5;124m'\u001b[39m,\n\u001b[1;32m     32\u001b[0m   filename\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m     33\u001b[0m   file\u001b[38;5;241m=\u001b[39mDESCRIPTOR,\n\u001b[1;32m     34\u001b[0m   containing_type\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m     35\u001b[0m   fields\u001b[38;5;241m=\u001b[39m[\n\u001b[0;32m---> 36\u001b[0m     \u001b[43m_descriptor\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mFieldDescriptor\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m     37\u001b[0m \u001b[43m      \u001b[49m\u001b[43mname\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43msize\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfull_name\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mtensorflow.TensorShapeProto.Dim.size\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mindex\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     38\u001b[0m \u001b[43m      \u001b[49m\u001b[43mnumber\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mtype\u001b[39;49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m3\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcpp_type\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m2\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mlabel\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     39\u001b[0m \u001b[43m      \u001b[49m\u001b[43mhas_default_value\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mFalse\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdefault_value\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     40\u001b[0m \u001b[43m      \u001b[49m\u001b[43mmessage_type\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43menum_type\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcontaining_type\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m     41\u001b[0m \u001b[43m      \u001b[49m\u001b[43mis_extension\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mFalse\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mextension_scope\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m     42\u001b[0m \u001b[43m      \u001b[49m\u001b[43mserialized_options\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfile\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mDESCRIPTOR\u001b[49m\u001b[43m)\u001b[49m,\n\u001b[1;32m     43\u001b[0m     _descriptor\u001b[38;5;241m.\u001b[39mFieldDescriptor(\n\u001b[1;32m     44\u001b[0m       name\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mname\u001b[39m\u001b[38;5;124m'\u001b[39m, full_name\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mtensorflow.TensorShapeProto.Dim.name\u001b[39m\u001b[38;5;124m'\u001b[39m, index\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1\u001b[39m,\n\u001b[1;32m     45\u001b[0m       number\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m2\u001b[39m, \u001b[38;5;28mtype\u001b[39m\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m9\u001b[39m, cpp_type\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m9\u001b[39m, label\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1\u001b[39m,\n\u001b[1;32m     46\u001b[0m       has_default_value\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mFalse\u001b[39;00m, default_value\u001b[38;5;241m=\u001b[39m_b(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m\"\u001b[39m)\u001b[38;5;241m.\u001b[39mdecode(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mutf-8\u001b[39m\u001b[38;5;124m'\u001b[39m),\n\u001b[1;32m     47\u001b[0m       message_type\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m, enum_type\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m, containing_type\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m     48\u001b[0m       is_extension\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mFalse\u001b[39;00m, extension_scope\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m     49\u001b[0m       serialized_options\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m, file\u001b[38;5;241m=\u001b[39mDESCRIPTOR),\n\u001b[1;32m     50\u001b[0m   ],\n\u001b[1;32m     51\u001b[0m   extensions\u001b[38;5;241m=\u001b[39m[\n\u001b[1;32m     52\u001b[0m   ],\n\u001b[1;32m     53\u001b[0m   nested_types\u001b[38;5;241m=\u001b[39m[],\n\u001b[1;32m     54\u001b[0m   enum_types\u001b[38;5;241m=\u001b[39m[\n\u001b[1;32m     55\u001b[0m   ],\n\u001b[1;32m     56\u001b[0m   serialized_options\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m     57\u001b[0m   is_extendable\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mFalse\u001b[39;00m,\n\u001b[1;32m     58\u001b[0m   syntax\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mproto3\u001b[39m\u001b[38;5;124m'\u001b[39m,\n\u001b[1;32m     59\u001b[0m   extension_ranges\u001b[38;5;241m=\u001b[39m[],\n\u001b[1;32m     60\u001b[0m   oneofs\u001b[38;5;241m=\u001b[39m[\n\u001b[1;32m     61\u001b[0m   ],\n\u001b[1;32m     62\u001b[0m   serialized_start\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m149\u001b[39m,\n\u001b[1;32m     63\u001b[0m   serialized_end\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m182\u001b[39m,\n\u001b[1;32m     64\u001b[0m )\n\u001b[1;32m     66\u001b[0m _TENSORSHAPEPROTO \u001b[38;5;241m=\u001b[39m _descriptor\u001b[38;5;241m.\u001b[39mDescriptor(\n\u001b[1;32m     67\u001b[0m   name\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mTensorShapeProto\u001b[39m\u001b[38;5;124m'\u001b[39m,\n\u001b[1;32m     68\u001b[0m   full_name\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mtensorflow.TensorShapeProto\u001b[39m\u001b[38;5;124m'\u001b[39m,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    100\u001b[0m   serialized_end\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m182\u001b[39m,\n\u001b[1;32m    101\u001b[0m )\n\u001b[1;32m    103\u001b[0m _TENSORSHAPEPROTO_DIM\u001b[38;5;241m.\u001b[39mcontaining_type \u001b[38;5;241m=\u001b[39m _TENSORSHAPEPROTO\n",
      "File \u001b[0;32m~/miniconda3/envs/rl/lib/python3.8/site-packages/google/protobuf/descriptor.py:561\u001b[0m, in \u001b[0;36mFieldDescriptor.__new__\u001b[0;34m(cls, name, full_name, index, number, type, cpp_type, label, default_value, message_type, enum_type, containing_type, is_extension, extension_scope, options, serialized_options, has_default_value, containing_oneof, json_name, file, create_key)\u001b[0m\n\u001b[1;32m    555\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__new__\u001b[39m(\u001b[38;5;28mcls\u001b[39m, name, full_name, index, number, \u001b[38;5;28mtype\u001b[39m, cpp_type, label,\n\u001b[1;32m    556\u001b[0m             default_value, message_type, enum_type, containing_type,\n\u001b[1;32m    557\u001b[0m             is_extension, extension_scope, options\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m    558\u001b[0m             serialized_options\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m    559\u001b[0m             has_default_value\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m, containing_oneof\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m, json_name\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m    560\u001b[0m             file\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m, create_key\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m):  \u001b[38;5;66;03m# pylint: disable=redefined-builtin\u001b[39;00m\n\u001b[0;32m--> 561\u001b[0m   \u001b[43m_message\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mMessage\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_CheckCalledFromGeneratedFile\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    562\u001b[0m   \u001b[38;5;28;01mif\u001b[39;00m is_extension:\n\u001b[1;32m    563\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m _message\u001b[38;5;241m.\u001b[39mdefault_pool\u001b[38;5;241m.\u001b[39mFindExtensionByName(full_name)\n",
      "\u001b[0;31mTypeError\u001b[0m: Descriptors cannot not be created directly.\nIf this call came from a _pb2.py file, your generated code is out of date and must be regenerated with protoc >= 3.19.0.\nIf you cannot immediately regenerate your protos, some other possible workarounds are:\n 1. Downgrade the protobuf package to 3.20.x or lower.\n 2. Set PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python (but this will use pure-Python parsing and will be much slower).\n\nMore information: https://developers.google.com/protocol-buffers/docs/news/2022-05-06#python-updates"
     ]
    }
   ],
   "source": [
    "from merlin.schema import Schema\n",
    "from merlin.io import Dataset\n",
    "\n",
    "train = Dataset(os.path.join(INPUT_DATA_DIR, \"train/part_0.parquet\"))\n",
    "schema = train.schema"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "id": "ce1113f4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'nvtabular/movielens/data/'"
      ]
     },
     "execution_count": 79,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "INPUT_DATA_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "id": "97e91811",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['userId', 'movieId', 'genres', 'rating']"
      ]
     },
     "execution_count": 80,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "schema.column_names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "id": "8975a1c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "schema = schema.select_by_name(\n",
    "   ['userId', 'movieId','genres']\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "id": "ae53ea94",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>name</th>\n",
       "      <th>tags</th>\n",
       "      <th>dtype</th>\n",
       "      <th>is_list</th>\n",
       "      <th>is_ragged</th>\n",
       "      <th>properties.num_buckets</th>\n",
       "      <th>properties.freq_threshold</th>\n",
       "      <th>properties.max_size</th>\n",
       "      <th>properties.cat_path</th>\n",
       "      <th>properties.embedding_sizes.cardinality</th>\n",
       "      <th>properties.embedding_sizes.dimension</th>\n",
       "      <th>properties.domain.min</th>\n",
       "      <th>properties.domain.max</th>\n",
       "      <th>properties.domain.name</th>\n",
       "      <th>properties.value_count.min</th>\n",
       "      <th>properties.value_count.max</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>userId</td>\n",
       "      <td>(Tags.CATEGORICAL)</td>\n",
       "      <td>DType(name='int64', element_type=&lt;ElementType....</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>None</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>.//categories/unique.userId.parquet</td>\n",
       "      <td>162544.0</td>\n",
       "      <td>512.0</td>\n",
       "      <td>0</td>\n",
       "      <td>162543</td>\n",
       "      <td>userId</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>movieId</td>\n",
       "      <td>(Tags.CATEGORICAL)</td>\n",
       "      <td>DType(name='int64', element_type=&lt;ElementType....</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>None</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>.//categories/unique.movieId.parquet</td>\n",
       "      <td>56640.0</td>\n",
       "      <td>512.0</td>\n",
       "      <td>0</td>\n",
       "      <td>56639</td>\n",
       "      <td>movieId</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>genres</td>\n",
       "      <td>(Tags.CATEGORICAL)</td>\n",
       "      <td>DType(name='int64', element_type=&lt;ElementType....</td>\n",
       "      <td>True</td>\n",
       "      <td>True</td>\n",
       "      <td>None</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>.//categories/unique.genres.parquet</td>\n",
       "      <td>23.0</td>\n",
       "      <td>16.0</td>\n",
       "      <td>0</td>\n",
       "      <td>22</td>\n",
       "      <td>genres</td>\n",
       "      <td>0.0</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "[{'name': 'userId', 'tags': {<Tags.CATEGORICAL: 'categorical'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'cat_path': './/categories/unique.userId.parquet', 'embedding_sizes': {'cardinality': 162544.0, 'dimension': 512.0}, 'domain': {'min': 0, 'max': 162543, 'name': 'userId'}}, 'dtype': DType(name='int64', element_type=<ElementType.Int: 'int'>, element_size=64, element_unit=None, signed=True, shape=Shape(dims=(Dimension(min=0, max=None),))), 'is_list': False, 'is_ragged': False}, {'name': 'movieId', 'tags': {<Tags.CATEGORICAL: 'categorical'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'cat_path': './/categories/unique.movieId.parquet', 'embedding_sizes': {'cardinality': 56640.0, 'dimension': 512.0}, 'domain': {'min': 0, 'max': 56639, 'name': 'movieId'}}, 'dtype': DType(name='int64', element_type=<ElementType.Int: 'int'>, element_size=64, element_unit=None, signed=True, shape=Shape(dims=(Dimension(min=0, max=None),))), 'is_list': False, 'is_ragged': False}, {'name': 'genres', 'tags': {<Tags.CATEGORICAL: 'categorical'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'cat_path': './/categories/unique.genres.parquet', 'embedding_sizes': {'cardinality': 23.0, 'dimension': 16.0}, 'domain': {'min': 0, 'max': 22, 'name': 'genres'}, 'value_count': {'min': 0, 'max': None}}, 'dtype': DType(name='int64', element_type=<ElementType.Int: 'int'>, element_size=64, element_unit=None, signed=True, shape=Shape(dims=(Dimension(min=0, max=None), Dimension(min=0, max=None)))), 'is_list': True, 'is_ragged': True}]"
      ]
     },
     "execution_count": 82,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "schema"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "bb7f315e",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: transformers4rec[nvtabular,pytorch] in ./miniconda3/envs/rl/lib/python3.8/site-packages (23.6.0)\n",
      "Requirement already satisfied: pyarrow>=1.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from transformers4rec[nvtabular,pytorch]) (11.0.0)\n",
      "Requirement already satisfied: merlin-models<23.07,>=23.06 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from transformers4rec[nvtabular,pytorch]) (23.6.0)\n",
      "Requirement already satisfied: tqdm>=4.27 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from transformers4rec[nvtabular,pytorch]) (4.65.0)\n",
      "Requirement already satisfied: torchmetrics>=0.10.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from transformers4rec[nvtabular,pytorch]) (0.11.4)\n",
      "Requirement already satisfied: transformers[torch]<5,>=4.12 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from transformers4rec[nvtabular,pytorch]) (4.27.4)\n",
      "Requirement already satisfied: nvtabular in ./miniconda3/envs/rl/lib/python3.8/site-packages (from transformers4rec[nvtabular,pytorch]) (23.6.0)\n",
      "Requirement already satisfied: torch>=1.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from transformers4rec[nvtabular,pytorch]) (2.0.0)\n",
      "Requirement already satisfied: merlin-dataloader<23.07,>=23.06 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from merlin-models<23.07,>=23.06->transformers4rec[nvtabular,pytorch]) (23.6.0)\n",
      "Requirement already satisfied: merlin-core<23.07,>=23.06 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from merlin-models<23.07,>=23.06->transformers4rec[nvtabular,pytorch]) (23.6.0)\n",
      "Requirement already satisfied: numpy>=1.16.6 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from pyarrow>=1.0->transformers4rec[nvtabular,pytorch]) (1.24.1)\n",
      "Requirement already satisfied: jinja2 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from torch>=1.0->transformers4rec[nvtabular,pytorch]) (3.1.2)\n",
      "Requirement already satisfied: filelock in ./miniconda3/envs/rl/lib/python3.8/site-packages (from torch>=1.0->transformers4rec[nvtabular,pytorch]) (3.10.7)\n",
      "Requirement already satisfied: networkx in ./miniconda3/envs/rl/lib/python3.8/site-packages (from torch>=1.0->transformers4rec[nvtabular,pytorch]) (3.0)\n",
      "Requirement already satisfied: sympy in ./miniconda3/envs/rl/lib/python3.8/site-packages (from torch>=1.0->transformers4rec[nvtabular,pytorch]) (1.11.1)\n",
      "Requirement already satisfied: typing-extensions in ./miniconda3/envs/rl/lib/python3.8/site-packages (from torch>=1.0->transformers4rec[nvtabular,pytorch]) (4.4.0)\n",
      "Requirement already satisfied: packaging in ./miniconda3/envs/rl/lib/python3.8/site-packages (from torchmetrics>=0.10.0->transformers4rec[nvtabular,pytorch]) (22.0)\n",
      "Requirement already satisfied: requests in ./miniconda3/envs/rl/lib/python3.8/site-packages (from transformers[torch]<5,>=4.12->transformers4rec[nvtabular,pytorch]) (2.28.1)\n",
      "Requirement already satisfied: huggingface-hub<1.0,>=0.11.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from transformers[torch]<5,>=4.12->transformers4rec[nvtabular,pytorch]) (0.13.3)\n",
      "Requirement already satisfied: pyyaml>=5.1 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from transformers[torch]<5,>=4.12->transformers4rec[nvtabular,pytorch]) (6.0)\n",
      "Requirement already satisfied: tokenizers!=0.11.3,<0.14,>=0.11.1 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from transformers[torch]<5,>=4.12->transformers4rec[nvtabular,pytorch]) (0.13.2)\n",
      "Requirement already satisfied: regex!=2019.12.17 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from transformers[torch]<5,>=4.12->transformers4rec[nvtabular,pytorch]) (2023.3.23)\n",
      "Requirement already satisfied: scipy in ./miniconda3/envs/rl/lib/python3.8/site-packages (from nvtabular->transformers4rec[nvtabular,pytorch]) (1.10.1)\n",
      "Requirement already satisfied: distributed>=2022.11.1 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from merlin-core<23.07,>=23.06->merlin-models<23.07,>=23.06->transformers4rec[nvtabular,pytorch]) (2023.3.2.1)\n",
      "Requirement already satisfied: pandas<1.6.0dev0,>=1.2.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from merlin-core<23.07,>=23.06->merlin-models<23.07,>=23.06->transformers4rec[nvtabular,pytorch]) (1.5.2)\n",
      "Requirement already satisfied: betterproto<2.0.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from merlin-core<23.07,>=23.06->merlin-models<23.07,>=23.06->transformers4rec[nvtabular,pytorch]) (1.2.5)\n",
      "Requirement already satisfied: npy-append-array in ./miniconda3/envs/rl/lib/python3.8/site-packages (from merlin-core<23.07,>=23.06->merlin-models<23.07,>=23.06->transformers4rec[nvtabular,pytorch]) (0.9.16)\n",
      "Requirement already satisfied: fsspec>=2022.7.1 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from merlin-core<23.07,>=23.06->merlin-models<23.07,>=23.06->transformers4rec[nvtabular,pytorch]) (2023.6.0)\n",
      "Requirement already satisfied: pynvml<11.5,>=11.0.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from merlin-core<23.07,>=23.06->merlin-models<23.07,>=23.06->transformers4rec[nvtabular,pytorch]) (11.4.1)\n",
      "Requirement already satisfied: dask>=2022.11.1 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from merlin-core<23.07,>=23.06->merlin-models<23.07,>=23.06->transformers4rec[nvtabular,pytorch]) (2023.3.2)\n",
      "Requirement already satisfied: protobuf>=3.0.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from merlin-core<23.07,>=23.06->merlin-models<23.07,>=23.06->transformers4rec[nvtabular,pytorch]) (4.23.3)\n",
      "Requirement already satisfied: dask-cuda>=22.12.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from merlin-core<23.07,>=23.06->merlin-models<23.07,>=23.06->transformers4rec[nvtabular,pytorch]) (23.4.0)\n",
      "Requirement already satisfied: numba>=0.54 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from merlin-core<23.07,>=23.06->merlin-models<23.07,>=23.06->transformers4rec[nvtabular,pytorch]) (0.57.1)\n",
      "Requirement already satisfied: tensorflow-metadata>=1.2.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from merlin-core<23.07,>=23.06->merlin-models<23.07,>=23.06->transformers4rec[nvtabular,pytorch]) (1.13.1)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from jinja2->torch>=1.0->transformers4rec[nvtabular,pytorch]) (2.1.1)\n",
      "Requirement already satisfied: charset-normalizer<3,>=2 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from requests->transformers[torch]<5,>=4.12->transformers4rec[nvtabular,pytorch]) (2.1.1)\n",
      "Requirement already satisfied: idna<4,>=2.5 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from requests->transformers[torch]<5,>=4.12->transformers4rec[nvtabular,pytorch]) (3.4)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from requests->transformers[torch]<5,>=4.12->transformers4rec[nvtabular,pytorch]) (2022.12.7)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from requests->transformers[torch]<5,>=4.12->transformers4rec[nvtabular,pytorch]) (1.26.13)\n",
      "Requirement already satisfied: mpmath>=0.19 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from sympy->torch>=1.0->transformers4rec[nvtabular,pytorch]) (1.3.0)\n",
      "Requirement already satisfied: stringcase in ./miniconda3/envs/rl/lib/python3.8/site-packages (from betterproto<2.0.0->merlin-core<23.07,>=23.06->merlin-models<23.07,>=23.06->transformers4rec[nvtabular,pytorch]) (1.2.0)\n",
      "Requirement already satisfied: grpclib in ./miniconda3/envs/rl/lib/python3.8/site-packages (from betterproto<2.0.0->merlin-core<23.07,>=23.06->merlin-models<23.07,>=23.06->transformers4rec[nvtabular,pytorch]) (0.4.5)\n",
      "Requirement already satisfied: importlib-metadata>=4.13.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from dask>=2022.11.1->merlin-core<23.07,>=23.06->merlin-models<23.07,>=23.06->transformers4rec[nvtabular,pytorch]) (4.13.0)\n",
      "Requirement already satisfied: cloudpickle>=1.1.1 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from dask>=2022.11.1->merlin-core<23.07,>=23.06->merlin-models<23.07,>=23.06->transformers4rec[nvtabular,pytorch]) (2.2.0)\n",
      "Requirement already satisfied: toolz>=0.8.2 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from dask>=2022.11.1->merlin-core<23.07,>=23.06->merlin-models<23.07,>=23.06->transformers4rec[nvtabular,pytorch]) (0.12.0)\n",
      "Requirement already satisfied: click>=7.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from dask>=2022.11.1->merlin-core<23.07,>=23.06->merlin-models<23.07,>=23.06->transformers4rec[nvtabular,pytorch]) (8.1.3)\n",
      "Requirement already satisfied: partd>=1.2.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from dask>=2022.11.1->merlin-core<23.07,>=23.06->merlin-models<23.07,>=23.06->transformers4rec[nvtabular,pytorch]) (1.4.0)\n",
      "Requirement already satisfied: zict>=0.1.3 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from dask-cuda>=22.12.0->merlin-core<23.07,>=23.06->merlin-models<23.07,>=23.06->transformers4rec[nvtabular,pytorch]) (3.0.0)\n",
      "Requirement already satisfied: locket>=1.0.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from distributed>=2022.11.1->merlin-core<23.07,>=23.06->merlin-models<23.07,>=23.06->transformers4rec[nvtabular,pytorch]) (1.0.0)\n",
      "Requirement already satisfied: msgpack>=1.0.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from distributed>=2022.11.1->merlin-core<23.07,>=23.06->merlin-models<23.07,>=23.06->transformers4rec[nvtabular,pytorch]) (1.0.5)\n",
      "Requirement already satisfied: psutil>=5.7.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from distributed>=2022.11.1->merlin-core<23.07,>=23.06->merlin-models<23.07,>=23.06->transformers4rec[nvtabular,pytorch]) (5.9.4)\n",
      "Requirement already satisfied: tornado>=6.0.3 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from distributed>=2022.11.1->merlin-core<23.07,>=23.06->merlin-models<23.07,>=23.06->transformers4rec[nvtabular,pytorch]) (6.2)\n",
      "Requirement already satisfied: tblib>=1.6.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from distributed>=2022.11.1->merlin-core<23.07,>=23.06->merlin-models<23.07,>=23.06->transformers4rec[nvtabular,pytorch]) (2.0.0)\n",
      "Requirement already satisfied: sortedcontainers>=2.0.5 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from distributed>=2022.11.1->merlin-core<23.07,>=23.06->merlin-models<23.07,>=23.06->transformers4rec[nvtabular,pytorch]) (2.4.0)\n",
      "Requirement already satisfied: llvmlite<0.41,>=0.40.0dev0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from numba>=0.54->merlin-core<23.07,>=23.06->merlin-models<23.07,>=23.06->transformers4rec[nvtabular,pytorch]) (0.40.1)\n",
      "Requirement already satisfied: python-dateutil>=2.8.1 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from pandas<1.6.0dev0,>=1.2.0->merlin-core<23.07,>=23.06->merlin-models<23.07,>=23.06->transformers4rec[nvtabular,pytorch]) (2.8.2)\n",
      "Requirement already satisfied: pytz>=2020.1 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from pandas<1.6.0dev0,>=1.2.0->merlin-core<23.07,>=23.06->merlin-models<23.07,>=23.06->transformers4rec[nvtabular,pytorch]) (2022.7)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: googleapis-common-protos<2,>=1.52.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from tensorflow-metadata>=1.2.0->merlin-core<23.07,>=23.06->merlin-models<23.07,>=23.06->transformers4rec[nvtabular,pytorch]) (1.59.1)\n",
      "Requirement already satisfied: absl-py<2.0.0,>=0.9 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from tensorflow-metadata>=1.2.0->merlin-core<23.07,>=23.06->merlin-models<23.07,>=23.06->transformers4rec[nvtabular,pytorch]) (1.3.0)\n",
      "Requirement already satisfied: zipp>=0.5 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from importlib-metadata>=4.13.0->dask>=2022.11.1->merlin-core<23.07,>=23.06->merlin-models<23.07,>=23.06->transformers4rec[nvtabular,pytorch]) (3.11.0)\n",
      "Requirement already satisfied: six>=1.5 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from python-dateutil>=2.8.1->pandas<1.6.0dev0,>=1.2.0->merlin-core<23.07,>=23.06->merlin-models<23.07,>=23.06->transformers4rec[nvtabular,pytorch]) (1.16.0)\n",
      "Requirement already satisfied: h2<5,>=3.1.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from grpclib->betterproto<2.0.0->merlin-core<23.07,>=23.06->merlin-models<23.07,>=23.06->transformers4rec[nvtabular,pytorch]) (4.1.0)\n",
      "Requirement already satisfied: multidict in ./miniconda3/envs/rl/lib/python3.8/site-packages (from grpclib->betterproto<2.0.0->merlin-core<23.07,>=23.06->merlin-models<23.07,>=23.06->transformers4rec[nvtabular,pytorch]) (6.0.4)\n",
      "Requirement already satisfied: hyperframe<7,>=6.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from h2<5,>=3.1.0->grpclib->betterproto<2.0.0->merlin-core<23.07,>=23.06->merlin-models<23.07,>=23.06->transformers4rec[nvtabular,pytorch]) (6.0.1)\n",
      "Requirement already satisfied: hpack<5,>=4.0 in ./miniconda3/envs/rl/lib/python3.8/site-packages (from h2<5,>=3.1.0->grpclib->betterproto<2.0.0->merlin-core<23.07,>=23.06->merlin-models<23.07,>=23.06->transformers4rec[nvtabular,pytorch]) (4.0.0)\n"
     ]
    }
   ],
   "source": [
    "!pip install 'transformers4rec[pytorch,nvtabular]'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "id": "d0e6cc95",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800000; text-decoration-color: #800000\">╭─────────────────────────────── </span><span style=\"color: #800000; text-decoration-color: #800000; font-weight: bold\">Traceback </span><span style=\"color: #bf7f7f; text-decoration-color: #bf7f7f; font-weight: bold\">(most recent call last)</span><span style=\"color: #800000; text-decoration-color: #800000\"> ────────────────────────────────╮</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">&lt;module&gt;</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">5</span>                                                                                    <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 2 </span>                                                                                            <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 3 </span>                                                                                            <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 4 # Define input module to process tabular input-features and to prepare masked inputs</span>        <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span> 5 input_module = tr.TabularSequenceFeatures.from_schema(                                      <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 6 │   </span>schema,                                                                                 <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 7 │   </span>aggregation=<span style=\"color: #808000; text-decoration-color: #808000\">\"concat\"</span>,                                                                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"> 8 │   </span>d_output=<span style=\"color: #0000ff; text-decoration-color: #0000ff\">256</span>,                                                                           <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">/Users/rajathdb/miniconda3/envs/rl/lib/python3.8/site-packages/transformers4rec/torch/features/</span><span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">s</span> <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">equence.py</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">194</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">from_schema</span>                                                                    <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">191 </span><span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">│   │   </span><span style=\"color: #808000; text-decoration-color: #808000\">TabularFeatures</span>                                                                    <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">192 </span><span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">│   │   │   </span><span style=\"color: #808000; text-decoration-color: #808000\">Returns ``TabularFeatures`` from a dataset schema</span>                              <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">193 </span><span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">│   │   </span><span style=\"color: #808000; text-decoration-color: #808000\">\"\"\"</span>                                                                                <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span>194 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   │   </span>output: TabularSequenceFeatures = <span style=\"color: #00ffff; text-decoration-color: #00ffff\">super</span>().from_schema(  <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\"># type: ignore</span>             <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">195 │   │   │   </span>schema=schema,                                                                 <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">196 │   │   │   </span>continuous_tags=continuous_tags,                                               <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">197 │   │   │   </span>categorical_tags=categorical_tags,                                             <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">/Users/rajathdb/miniconda3/envs/rl/lib/python3.8/site-packages/transformers4rec/torch/features/</span><span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">t</span> <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">abular.py</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">179</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">from_schema</span>                                                                     <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">176 │   │   │   │   │   </span>schema, tags=continuous_tags, **kwargs                                 <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">177 │   │   │   │   </span>)                                                                          <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">178 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">if</span> categorical_tags:                                                               <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span>179 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   │   │   </span>maybe_categorical_module = <span style=\"color: #00ffff; text-decoration-color: #00ffff\">cls</span>.EMBEDDING_MODULE_CLASS.from_schema(             <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">180 │   │   │   │   </span>schema, tags=categorical_tags, **kwargs                                    <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">181 │   │   │   </span>)                                                                              <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">182 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">if</span> pretrained_embeddings_tags:                                                     <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">/Users/rajathdb/miniconda3/envs/rl/lib/python3.8/site-packages/transformers4rec/torch/features/</span><span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">e</span> <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">mbedding.py</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">213</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">from_schema</span>                                                                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">210 │   │   </span>                                                                                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">211 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">if</span> automatic_build <span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">and</span> schema:                                                     <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">212 │   │   │   </span>output.build(                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span>213 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   │   │   │   </span>get_output_sizes_from_schema(                                              <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">214 │   │   │   │   │   </span>schema,                                                                <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">215 │   │   │   │   │   </span>kwargs.get(<span style=\"color: #808000; text-decoration-color: #808000\">\"batch_size\"</span>, -<span style=\"color: #0000ff; text-decoration-color: #0000ff\">1</span>),                                          <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">216 │   │   │   │   │   </span>max_sequence_length=max_sequence_length,                               <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #bfbf7f; text-decoration-color: #bfbf7f\">/Users/rajathdb/miniconda3/envs/rl/lib/python3.8/site-packages/transformers4rec/torch/utils/</span><span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">torc</span> <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #808000; text-decoration-color: #808000; font-weight: bold\">h_utils.py</span>:<span style=\"color: #0000ff; text-decoration-color: #0000ff\">186</span> in <span style=\"color: #00ff00; text-decoration-color: #00ff00\">get_output_sizes_from_schema</span>                                                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>                                                                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">183 │   │   </span>                                                                                   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">184 │   │   # Sequential or multi-hot feature</span>                                                  <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">185 │   │   </span><span style=\"color: #0000ff; text-decoration-color: #0000ff\">elif</span> _has_field(feature, <span style=\"color: #808000; text-decoration-color: #808000\">\"value_count\"</span>):                                           <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span> <span style=\"color: #800000; text-decoration-color: #800000\">❱ </span>186 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">│   │   │   </span>sizes[name] = torch.Size(                                                      <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">187 │   │   │   │   </span>[                                                                          <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">188 │   │   │   │   │   </span>batch_size,                                                            <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">│</span>   <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">189 │   │   │   │   │   </span>max_sequence_length <span style=\"color: #0000ff; text-decoration-color: #0000ff\">if</span> max_sequence_length <span style=\"color: #0000ff; text-decoration-color: #0000ff\">else</span> feature.value_count.   <span style=\"color: #800000; text-decoration-color: #800000\">│</span>\n",
       "<span style=\"color: #800000; text-decoration-color: #800000\">╰──────────────────────────────────────────────────────────────────────────────────────────────────╯</span>\n",
       "<span style=\"color: #ff0000; text-decoration-color: #ff0000; font-weight: bold\">TypeError: </span><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">torch.Size</span><span style=\"font-weight: bold\">()</span> takes an iterable of <span style=\"color: #008000; text-decoration-color: #008000\">'int'</span> <span style=\"font-weight: bold\">(</span>item <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span> is <span style=\"color: #008000; text-decoration-color: #008000\">'NoneType'</span><span style=\"font-weight: bold\">)</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[31m╭─\u001b[0m\u001b[31m──────────────────────────────\u001b[0m\u001b[31m \u001b[0m\u001b[1;31mTraceback \u001b[0m\u001b[1;2;31m(most recent call last)\u001b[0m\u001b[31m \u001b[0m\u001b[31m───────────────────────────────\u001b[0m\u001b[31m─╮\u001b[0m\n",
       "\u001b[31m│\u001b[0m in \u001b[92m<module>\u001b[0m:\u001b[94m5\u001b[0m                                                                                    \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 2 \u001b[0m                                                                                            \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 3 \u001b[0m                                                                                            \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 4 \u001b[0m\u001b[2m# Define input module to process tabular input-features and to prepare masked inputs\u001b[0m        \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m 5 input_module = tr.TabularSequenceFeatures.from_schema(                                      \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 6 \u001b[0m\u001b[2m│   \u001b[0mschema,                                                                                 \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 7 \u001b[0m\u001b[2m│   \u001b[0maggregation=\u001b[33m\"\u001b[0m\u001b[33mconcat\u001b[0m\u001b[33m\"\u001b[0m,                                                                   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m 8 \u001b[0m\u001b[2m│   \u001b[0md_output=\u001b[94m256\u001b[0m,                                                                           \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[2;33m/Users/rajathdb/miniconda3/envs/rl/lib/python3.8/site-packages/transformers4rec/torch/features/\u001b[0m\u001b[1;33ms\u001b[0m \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[1;33mequence.py\u001b[0m:\u001b[94m194\u001b[0m in \u001b[92mfrom_schema\u001b[0m                                                                    \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m191 \u001b[0m\u001b[2;33m│   │   \u001b[0m\u001b[33mTabularFeatures\u001b[0m                                                                    \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m192 \u001b[0m\u001b[2;33m│   │   │   \u001b[0m\u001b[33mReturns ``TabularFeatures`` from a dataset schema\u001b[0m                              \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m193 \u001b[0m\u001b[2;33m│   │   \u001b[0m\u001b[33m\"\"\"\u001b[0m                                                                                \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m194 \u001b[2m│   │   \u001b[0moutput: TabularSequenceFeatures = \u001b[96msuper\u001b[0m().from_schema(  \u001b[2m# type: ignore\u001b[0m             \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m195 \u001b[0m\u001b[2m│   │   │   \u001b[0mschema=schema,                                                                 \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m196 \u001b[0m\u001b[2m│   │   │   \u001b[0mcontinuous_tags=continuous_tags,                                               \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m197 \u001b[0m\u001b[2m│   │   │   \u001b[0mcategorical_tags=categorical_tags,                                             \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[2;33m/Users/rajathdb/miniconda3/envs/rl/lib/python3.8/site-packages/transformers4rec/torch/features/\u001b[0m\u001b[1;33mt\u001b[0m \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[1;33mabular.py\u001b[0m:\u001b[94m179\u001b[0m in \u001b[92mfrom_schema\u001b[0m                                                                     \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m176 \u001b[0m\u001b[2m│   │   │   │   │   \u001b[0mschema, tags=continuous_tags, **kwargs                                 \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m177 \u001b[0m\u001b[2m│   │   │   │   \u001b[0m)                                                                          \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m178 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mif\u001b[0m categorical_tags:                                                               \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m179 \u001b[2m│   │   │   \u001b[0mmaybe_categorical_module = \u001b[96mcls\u001b[0m.EMBEDDING_MODULE_CLASS.from_schema(             \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m180 \u001b[0m\u001b[2m│   │   │   │   \u001b[0mschema, tags=categorical_tags, **kwargs                                    \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m181 \u001b[0m\u001b[2m│   │   │   \u001b[0m)                                                                              \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m182 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mif\u001b[0m pretrained_embeddings_tags:                                                     \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[2;33m/Users/rajathdb/miniconda3/envs/rl/lib/python3.8/site-packages/transformers4rec/torch/features/\u001b[0m\u001b[1;33me\u001b[0m \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[1;33mmbedding.py\u001b[0m:\u001b[94m213\u001b[0m in \u001b[92mfrom_schema\u001b[0m                                                                   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m210 \u001b[0m\u001b[2m│   │   \u001b[0m                                                                                   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m211 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94mif\u001b[0m automatic_build \u001b[95mand\u001b[0m schema:                                                     \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m212 \u001b[0m\u001b[2m│   │   │   \u001b[0moutput.build(                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m213 \u001b[2m│   │   │   │   \u001b[0mget_output_sizes_from_schema(                                              \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m214 \u001b[0m\u001b[2m│   │   │   │   │   \u001b[0mschema,                                                                \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m215 \u001b[0m\u001b[2m│   │   │   │   │   \u001b[0mkwargs.get(\u001b[33m\"\u001b[0m\u001b[33mbatch_size\u001b[0m\u001b[33m\"\u001b[0m, -\u001b[94m1\u001b[0m),                                          \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m216 \u001b[0m\u001b[2m│   │   │   │   │   \u001b[0mmax_sequence_length=max_sequence_length,                               \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[2;33m/Users/rajathdb/miniconda3/envs/rl/lib/python3.8/site-packages/transformers4rec/torch/utils/\u001b[0m\u001b[1;33mtorc\u001b[0m \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[1;33mh_utils.py\u001b[0m:\u001b[94m186\u001b[0m in \u001b[92mget_output_sizes_from_schema\u001b[0m                                                   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m                                                                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m183 \u001b[0m\u001b[2m│   │   \u001b[0m                                                                                   \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m184 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[2m# Sequential or multi-hot feature\u001b[0m                                                  \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m185 \u001b[0m\u001b[2m│   │   \u001b[0m\u001b[94melif\u001b[0m _has_field(feature, \u001b[33m\"\u001b[0m\u001b[33mvalue_count\u001b[0m\u001b[33m\"\u001b[0m):                                           \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m \u001b[31m❱ \u001b[0m186 \u001b[2m│   │   │   \u001b[0msizes[name] = torch.Size(                                                      \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m187 \u001b[0m\u001b[2m│   │   │   │   \u001b[0m[                                                                          \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m188 \u001b[0m\u001b[2m│   │   │   │   │   \u001b[0mbatch_size,                                                            \u001b[31m│\u001b[0m\n",
       "\u001b[31m│\u001b[0m   \u001b[2m189 \u001b[0m\u001b[2m│   │   │   │   │   \u001b[0mmax_sequence_length \u001b[94mif\u001b[0m max_sequence_length \u001b[94melse\u001b[0m feature.value_count.   \u001b[31m│\u001b[0m\n",
       "\u001b[31m╰──────────────────────────────────────────────────────────────────────────────────────────────────╯\u001b[0m\n",
       "\u001b[1;91mTypeError: \u001b[0m\u001b[1;35mtorch.Size\u001b[0m\u001b[1m(\u001b[0m\u001b[1m)\u001b[0m takes an iterable of \u001b[32m'int'\u001b[0m \u001b[1m(\u001b[0mitem \u001b[1;36m1\u001b[0m is \u001b[32m'NoneType'\u001b[0m\u001b[1m)\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from transformers4rec import torch as tr\n",
    "\n",
    "\n",
    "# Define input module to process tabular input-features and to prepare masked inputs\n",
    "input_module = tr.TabularSequenceFeatures.from_schema(\n",
    "    schema,\n",
    "    aggregation=\"concat\",\n",
    "    d_output=256,\n",
    "    masking=\"mlm\",\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "id": "894a5772",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "[]"
      ]
     },
     "execution_count": 84,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "schema.select_by_tag(Tags.TARGET)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "id": "f1852362",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "id": "82c78233",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([3, 4, 5])"
      ]
     },
     "execution_count": 87,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.Size([3, 4, 5])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95f7a941",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}