demo/extended_inference.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c7bc575c",
   "metadata": {},
   "source": [
    "# Extended inference\n",
    "\n",
    "This demo shows the following inference options:\n",
    "\n",
    "- big data inference\n",
    "- iterable data load and inference\n",
    "- labels passing\n",
    "\n",
    "We will generage wyth `pyspark` a big dataset with transactions and other fields.\n",
    "\n",
    "Since the dataset is large, we will use an iterable dataset to save memory and enable data loading and GPU output in parallel.\n",
    "\n",
    "We don't need only embeddings. We need them with user_id, target, labels and other fields. We should keep it together for large iterable dataset."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70d465e5",
   "metadata": {},
   "source": [
    "# Data generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6974773f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark import SparkConf\n",
    "from pyspark.sql import SparkSession\n",
    "from pyspark.sql import Row\n",
    "from pyspark.sql import Window\n",
    "import pyspark.sql.functions as F\n",
    "import pyspark.sql.types as T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "50384ed8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "22/07/22 14:23:23 WARN Utils: Your hostname, vm2 resolves to a loopback address: 127.0.1.1; using 192.168.0.6 instead (on interface ens192)\n",
      "22/07/22 14:23:23 WARN Utils: Set SPARK_LOCAL_IP if you need to bind to another address\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting default log level to \"WARN\".\n",
      "To adjust logging level use sc.setLogLevel(newLevel). For SparkR, use setLogLevel(newLevel).\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "22/07/22 14:23:24 WARN NativeCodeLoader: Unable to load native-hadoop library for your platform... using builtin-java classes where applicable\n",
      "22/07/22 14:23:24 WARN SparkConf: Note that spark.local.dir will be overridden by the value set by the cluster manager (via SPARK_LOCAL_DIRS in mesos/standalone/kubernetes and LOCAL_DIRS in YARN).\n",
      "22/07/22 14:23:25 WARN Utils: Service 'SparkUI' could not bind on port 4040. Attempting port 4041.\n",
      "22/07/22 14:23:25 WARN Utils: Service 'SparkUI' could not bind on port 4041. Attempting port 4042.\n",
      "22/07/22 14:23:25 WARN Utils: Service 'SparkUI' could not bind on port 4042. Attempting port 4043.\n"
     ]
    }
   ],
   "source": [
    "spark_conf = SparkConf()\n",
    "spark_conf.setMaster(\"local[8]\").setAppName(\"pyspark_data_generation\")\n",
    "spark_conf.set(\"spark.executor.memory\", \"32g\")\n",
    "spark_conf.set(\"spark.executor.memoryOverhead\", \"4g\")\n",
    "spark_conf.set(\"spark.driver.memory\", \"32g\")\n",
    "spark_conf.set(\"spark.driver.memoryOverhead\", \"4g\")\n",
    "spark_conf.set(\"spark.local.dir\", \"./spark_local_dir\")\n",
    "\n",
    "spark = SparkSession.builder.config(conf=spark_conf).getOrCreate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "5f214099",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "            <div>\n",
       "                <p><b>SparkSession - in-memory</b></p>\n",
       "                \n",
       "        <div>\n",
       "            <p><b>SparkContext</b></p>\n",
       "\n",
       "            <p><a href=\"http://192.168.0.6:4043\">Spark UI</a></p>\n",
       "\n",
       "            <dl>\n",
       "              <dt>Version</dt>\n",
       "                <dd><code>v3.3.0</code></dd>\n",
       "              <dt>Master</dt>\n",
       "                <dd><code>local[8]</code></dd>\n",
       "              <dt>AppName</dt>\n",
       "                <dd><code>pyspark_data_generation</code></dd>\n",
       "            </dl>\n",
       "        </div>\n",
       "        \n",
       "            </div>\n",
       "        "
      ],
      "text/plain": [
       "<pyspark.sql.session.SparkSession at 0x7f373f331280>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "spark"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "2ec99680",
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "04451d05",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a9d08382",
   "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>trans_date</th>\n",
       "      <th>small_group</th>\n",
       "      <th>amount_rur</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>6.0</td>\n",
       "      <td>4</td>\n",
       "      <td>71.462997</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>6.0</td>\n",
       "      <td>35</td>\n",
       "      <td>45.016998</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>8.0</td>\n",
       "      <td>11</td>\n",
       "      <td>13.887000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>9.0</td>\n",
       "      <td>11</td>\n",
       "      <td>15.983000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>10.0</td>\n",
       "      <td>11</td>\n",
       "      <td>21.341000</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   trans_date small_group  amount_rur\n",
       "0         6.0           4   71.462997\n",
       "1         6.0          35   45.016998\n",
       "2         8.0          11   13.887000\n",
       "3         9.0          11   15.983000\n",
       "4        10.0          11   21.341000"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from urllib.request import urlretrieve\n",
    "data, _ = urlretrieve(\n",
    "    'https://huggingface.co/datasets/dllllb/age-group-prediction/resolve/main/transactions_train.csv.gz?download=true',\n",
    "    'transactions_train.csv.gz')\n",
    "\n",
    "df_real_trx = spark.read.csv(data, header=True) \\\n",
    ".select(\n",
    "    F.col('trans_date').cast('float'), \n",
    "    'small_group',\n",
    "    F.col('amount_rur').cast('float'),\n",
    ")\n",
    "df_real_trx.limit(5).toPandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d873d75a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    },
    {
     "data": {
      "text/plain": [
       "26450577"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "real_trx_count = df_real_trx.count()\n",
    "real_trx_count"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "848095a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "USER_CNT = 1e6\n",
    "TOTAL_TRX_FOR_GENERATE = 800 * USER_CNT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "1f28aeb2",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_generated = df_real_trx.sample(fraction=TOTAL_TRX_FOR_GENERATE / real_trx_count, withReplacement=True)\n",
    "df_generated = df_generated.repartition(200)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "e50cb85d",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_generated = df_generated.withColumn('user_id', F.concat(F.lit('u_'), F.floor(F.rand() * USER_CNT)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "e464dcfb",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_generated = df_generated.withColumn('event_time', F.rand())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "a0c58c7c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    },
    {
     "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>trans_date</th>\n",
       "      <th>small_group</th>\n",
       "      <th>amount_rur</th>\n",
       "      <th>user_id</th>\n",
       "      <th>event_time</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>649.0</td>\n",
       "      <td>18</td>\n",
       "      <td>10.991000</td>\n",
       "      <td>u_599206</td>\n",
       "      <td>0.246618</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>593.0</td>\n",
       "      <td>15</td>\n",
       "      <td>5.908000</td>\n",
       "      <td>u_195427</td>\n",
       "      <td>0.439432</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>354.0</td>\n",
       "      <td>1</td>\n",
       "      <td>3.704000</td>\n",
       "      <td>u_527251</td>\n",
       "      <td>0.167604</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>419.0</td>\n",
       "      <td>34</td>\n",
       "      <td>63.700001</td>\n",
       "      <td>u_44570</td>\n",
       "      <td>0.924908</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>216.0</td>\n",
       "      <td>4</td>\n",
       "      <td>10.221000</td>\n",
       "      <td>u_403353</td>\n",
       "      <td>0.706942</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   trans_date small_group  amount_rur   user_id  event_time\n",
       "0       649.0          18   10.991000  u_599206    0.246618\n",
       "1       593.0          15    5.908000  u_195427    0.439432\n",
       "2       354.0           1    3.704000  u_527251    0.167604\n",
       "3       419.0          34   63.700001   u_44570    0.924908\n",
       "4       216.0           4   10.221000  u_403353    0.706942"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_generated.limit(5).toPandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "efdeff30",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 195 ms, sys: 8.79 ms, total: 203 ms\n",
      "Wall time: 4min 2s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "df_generated.write.parquet('df_generated.parquet', mode='overwrite')\n",
    "df_generated = spark.read.parquet('df_generated.parquet')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a934eaf5",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "31a1f7c1",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAsUAAACcCAYAAACX1dYOAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8qNh9FAAAACXBIWXMAAAsTAAALEwEAmpwYAAAQlklEQVR4nO3de7BdZXnH8e9PUC7JlKuewRAMFkZLzYiSKi1T5yDOlIsa26EYiwoWm3a8VDQq0PqH/cMWO1WKQ8tMKiow1kipHahFqxM9bfkDLBRGBKpGIJIYgheIBqo1+vSPvUIPx5yTvcm+nb2+n5kzZ613rb33uyZP1n7Ou5613lQVkiRJUps9bdQdkCRJkkbNpFiSJEmtZ1IsSZKk1jMpliRJUuuZFEuSJKn1TIolSZLUeibFkiRJaj2TYklqkSQPJHnFqPshSePGpFiSJEmtZ1IsSSOUZHmSzyT5bpLvJ7kiydOSvC/J5iQPJ7kmySHN/tNJtsx5jydGf5O8P8l1zWt+lOTuJKuabdcCxwD/nGRnkvcO+3glaVyZFEvSiCTZD/gssBlYASwDNgDnNz+nAs8FlgJX9PDWr27e51Dgxt2vrao3AN8GXlVVS6vqL/f9KCRpMpgUS9LovAR4NvCeqnqsqn5cVTcD5wIfrqr7qmoncAmwJsn+Xb7vzVV1U1X9DLgWeOFAei9JE8SkWJJGZzmwuap2zWl/Np3R4902A/sDU12+70Ozlh8HDuwhoZakVjIplqTReRA4Zg8J63eA58xaPwbYBWwHHgMO3r2hKcF4Zg+fWU+tq5I02UyKJWl0vgJsAy5NsiTJgUlOAT4FvDPJsUmWAn8OfLoZUf4GnZHfs5I8HXgfcEAPn7mdTp2yJGkWk2JJGpGm5vdVwHF0boDbArwW+BidWuB/B+4Hfgy8vXnNDuAtwEeBrXRGjrfMfe8F/AXwviSPJnl3f45Ekha/VHklTZIkSe3mSLEkSZJaz6RYkiRJrWdSLEmSpNYzKZYkSVLrmRRLkiSp9cZihqMjjzyyVqxYMepujMRjjz3GkiVLRt0NjSFjQwsxPjQfY0PzMTY6br/99u9V1S9MejQWSfGKFSu47bbbRt2NkZiZmWF6enrU3dAYMja0EOND8zE2NB9joyPJ5j21Wz4hSZKk1jMpliRJUuuZFEuSJKn1xqKmWJIEKy7+l673XbdyF9OD64oktY4jxZIkSWo9R4olqSV6GYkGeODSswbUE0kaP44US5IkqfVMiiVJktR6JsWSJElqPZNiSZIktZ432klSF3q9SQ28UU2SFpOuRoqTHJrk+iT/neTeJL+e5PAkX0zyzeb3Yc2+SfKRJJuSfDXJiwd7CJIkSdK+6bZ84nLg81X1fOCFwL3AxcDGqjoe2NisA5wBHN/8rAWu7GuPJUmSpD7ba1Kc5BDgZcBVAFX1v1X1KLAauLrZ7WrgNc3yauCa6rgFODTJUX3utyRJktQ33YwUHwt8F/h4kjuSfDTJEmCqqrY1+zwETDXLy4AHZ71+S9MmSZIkjaVU1cI7JKuAW4BTqurWJJcDPwTeXlWHztrvkao6LMlngUur6uamfSNwUVXdNud919Ipr2BqauqkDRs29PGwFo+dO3eydOnSUXdDY8jYGC93bd3R82tWLjtkYJ8xdRA86/DBvT/03n+NB88dmo+x0XHqqafeXlWr5rZ38/SJLcCWqrq1Wb+eTv3w9iRHVdW2pjzi4Wb7VmD5rNcf3bQ9SVWtB9YDrFq1qqanp7s9lokyMzNDW49dCzM2xsv5T+XpE+dOD+wz1q3cxTk9xkevx9Br/zUePHdoPsbGwvZaPlFVDwEPJnle03QacA9wI3Be03YecEOzfCPwxuYpFCcDO2aVWUiSJEljp9vnFL8d+GSSZwD3AW+ik1Bfl+QCYDNwTrPvTcCZwCbg8WZfSZIkaWx1lRRX1Z3AL9Re0Bk1nrtvAW/dt25JkiRJw+OMdpKkvnDWP0mLmUmxJC1STyUJlSTtWbcz2kmSJEkTy6RYkiRJrWdSLEmSpNazpliSBsSaX0laPEyKJU2EXhNQn3ogSZrNpFiStEeOdEtqE2uKJUmS1HomxZIkSWo9k2JJkiS1nkmxJEmSWs+kWJIkSa1nUixJkqTW85FsklrJx41JkmZzpFiSJEmt50ixJGnRcOZCSYPiSLEkSZJaz6RYkiRJrWdSLEmSpNYzKZYkSVLreaOdpKHwBilJ0jjreqQ4yX5J7kjy2Wb92CS3JtmU5NNJntG0H9Csb2q2rxhQ3yVJkqS+6KV84h3AvbPWPwhcVlXHAY8AFzTtFwCPNO2XNftJkiRJY6ur8okkRwNnAR8A3pUkwMuB32t2uRp4P3AlsLpZBrgeuCJJqqr6121Jk84Z5yRJw9TtSPFfA+8Fft6sHwE8WlW7mvUtwLJmeRnwIECzfUezvyRJkjSWsrcB3CSvBM6sqrckmQbeDZwP3NKUSJBkOfC5qnpBkq8Bp1fVlmbbt4CXVtX35rzvWmAtwNTU1EkbNmzo53EtGjt37mTp0qWj7obG0KTFxl1bd4y6CxNl6iDY/j+j7sX4W7nskFF3Yegm7dyh/jE2Ok499dTbq2rV3PZuyidOAV6d5EzgQOCXgMuBQ5Ps34wGHw1sbfbfCiwHtiTZHzgE+P7cN62q9cB6gFWrVtX09HTPBzUJZmZmaOuxa2GTFhvnWw7RV+tW7uJDd/kAob154NzpUXdh6Cbt3KH+MTYWttfyiaq6pKqOrqoVwBrgS1V1LvBl4Oxmt/OAG5rlG5t1mu1fsp5YkiRJ42xfJu+4iM5Nd5vo1Axf1bRfBRzRtL8LuHjfuihJkiQNVk/X3qpqBphplu8DXrKHfX4M/G4f+iZJkiQNhdM8S5IkqfVMiiVJktR6JsWSJElqPZNiSZIktZ4PuZQkTaxepwt/4NKzBtQTSePOpFiSpIZJtNRelk9IkiSp9UyKJUmS1HomxZIkSWo9k2JJkiS1nkmxJEmSWs+kWJIkSa1nUixJkqTWMymWJElS6zl5hyQnLJAktZ5JsaSe9ZpES5I07iyfkCRJUuuZFEuSJKn1TIolSZLUeibFkiRJaj2TYkmSJLXeXpPiJMuTfDnJPUnuTvKOpv3wJF9M8s3m92FNe5J8JMmmJF9N8uJBH4QkSZK0L7oZKd4FrKuqE4CTgbcmOQG4GNhYVccDG5t1gDOA45uftcCVfe+1JEmS1Ed7fU5xVW0DtjXLP0pyL7AMWA1MN7tdDcwAFzXt11RVAbckOTTJUc37SJI0MZz4RpocPdUUJ1kBvAi4FZialeg+BEw1y8uAB2e9bEvTJkmSJI2lrme0S7IU+Efgwqr6YZIntlVVJalePjjJWjrlFUxNTTEzM9PLyyfGzp07W3vsWtgwY2Pdyl1D+Rz1z9RB/rstRsP4P+33iuZjbCysq6Q4ydPpJMSfrKrPNM3bd5dFJDkKeLhp3wosn/Xyo5u2J6mq9cB6gFWrVtX09PRTO4JFbmZmhrYeuxY2zNg432mbF511K3fxobu6HtfQuLjrsZ5f0mvJhd8rmo+xsbBunj4R4Crg3qr68KxNNwLnNcvnATfMan9j8xSKk4Ed1hNLkiRpnHUzzHAK8AbgriR3Nm1/AlwKXJfkAmAzcE6z7SbgTGAT8Djwpn52WNLe9XrzjyRJbdfN0yduBjLP5tP2sH8Bb93HfkmSJElDY0GatAg48itJ0mA5zbMkSZJaz6RYkiRJrWdSLEmSpNYzKZYkSVLrmRRLkiSp9Xz6hCRJY6zXp8984vQlA+qJNNlMiiVJmiB3bd3R09TtvU4jLU0qyyckSZLUeo4USyPQzeXQdSt39TTaI0mSnjpHiiVJktR6jhRLfeA0zJIkLW4mxZIktVivf9R7Y54mleUTkiRJaj2TYkmSJLWe5ROSJKlrlltoUjlSLEmSpNYzKZYkSVLrWT6hieelPkmStDcmxdIcPnNYkqT2sXxCkiRJredIsSRJGphhXH2z7E39MJCkOMnpwOXAfsBHq+rSQXyOJoPlCpKkfeG9I+qHvpdPJNkP+BvgDOAE4HVJTuj350iSJEn9MoiR4pcAm6rqPoAkG4DVwD0D+CwNgSO5kqRJ4siy9mQQSfEy4MFZ61uAlw7gcyaSCagkSeNlUr6b163cxfldHkuvfwhMwh8aqar+vmFyNnB6Vb25WX8D8NKqetuc/dYCa5vV5wFf72tHFo8jge+NuhMaS8aGFmJ8aD7GhuZjbHQ8p6qeObdxECPFW4Hls9aPbtqepKrWA+sH8PmLSpLbqmrVqPuh8WNsaCHGh+ZjbGg+xsbCBvGc4v8Ejk9ybJJnAGuAGwfwOZIkSVJf9H2kuKp2JXkb8K90Hsn2saq6u9+fI0mSJPXLQJ5TXFU3ATcN4r0nUOtLSDQvY0MLMT40H2ND8zE2FtD3G+0kSZKkxWYQNcWSJEnSomJSPEBJnpfkzlk/P0xyYZLDk3wxyTeb34c1+yfJR5JsSvLVJC8e9TFocBaIj/cn2Tqr/cxZr7mkiY+vJ/mtUfZfg5XknUnuTvK1JJ9KcmBzA/OtTQx8urmZmSQHNOubmu0rRtx9DdA8sfGJJPfPOm+c2Ozr90qLJHlHExd3J7mwaTPn6JJJ8QBV1der6sSqOhE4CXgc+CfgYmBjVR0PbGzWoTM19vHNz1rgyqF3WkOzQHwAXLZ7W1OjTzNd+hrgV4HTgb9tplXXhEmyDPhjYFVVvYDOTctrgA/SiY3jgEeAC5qXXAA80rRf1uynCbRAbAC8Z9Z5486mze+VlkjyAuAP6Mws/ELglUmOw5yjaybFw3Ma8K2q2kxn2uurm/argdc0y6uBa6rjFuDQJEcNvacahdnxMZ/VwIaq+klV3Q9sonPy02TaHzgoyf7AwcA24OXA9c32ueeO3eeU64HTkmR4XdWQzY2N7yywr98r7fErwK1V9XhV7QL+DfgdzDm6ZlI8PGuATzXLU1W1rVl+CJhqlvc0Rfay4XRPIzY7PgDe1lzO+tjuS10YH61RVVuBvwK+TScZ3gHcDjzafNnBk//9n4iNZvsO4Ihh9lnDsafYqKovNJs/0Jw3LktyQNPmeaM9vgb8ZpIjkhwMnElnMjVzji6ZFA9BU/f3auAf5m6rzuM/fARIi+0hPq4Efhk4kc6X3odG0zONSvOH0GrgWODZwBI6JTNquT3FRpLXA5cAzwd+DTgcuGhkndRIVNW9dEqnvgB8HrgT+Nmcfcw5FmBSPBxnAP9VVdub9e27L1E0vx9u2ruaIlsT50nxUVXbq+pnVfVz4O/4/xIJ46M9XgHcX1XfraqfAp8BTqFzeXP38+Vn//s/ERvN9kOA7w+3yxqSPcXGb1TVtuYy+E+Aj+N5o5Wq6qqqOqmqXkbnvoNvYM7RNZPi4XgdT740fiNwXrN8HnDDrPY3NneEnkznstg2NOmeFB9zarp+m84lMejEx5rmSQPH0rk54itD66WG6dvAyUkObmqDTwPuAb4MnN3sM/fcsfuccjbwpfIh9JNqT7Fx76ykJ3RqRmefN/xeaYkkz2p+H0OnnvjvMefompN3DFiSJXROYs+tqh1N2xHAdcAxwGbgnKr6QXMyu4LOZdLHgTdV1W2j6bmGYZ74uJZO6UQBDwB/uPtEleRPgd8HdgEXVtXnRtBtDUGSPwNeS+ff+g7gzXTq/TbQuTx+B/D6qvpJkgOBa4EXAT8A1lTVfSPpuAZuntj4HPBMIHQum/9RVe30e6VdkvwHnfsJfgq8q6o2mnN0z6RYkiRJrWf5hCRJklrPpFiSJEmtZ1IsSZKk1jMpliRJUuuZFEuSJKn1TIolSZLUeibFkiRJaj2TYkmSJLXe/wEi/TXTmfrsVAAAAABJRU5ErkJggg==",
      "text/plain": [
       "<Figure size 864x144 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 606 ms, sys: 785 ms, total: 1.39 s\n",
      "Wall time: 1min 20s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "df_generated.groupby('user_id').count().select('count') \\\n",
    ".sample(True, fraction=1e4 / USER_CNT).toPandas().hist(bins=50, figsize=(12, 2))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "79e59529",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ptls.preprocessing import PysparkDataPreprocessor\n",
    "\n",
    "preprocessor = PysparkDataPreprocessor(\n",
    "    col_id='user_id',\n",
    "    col_event_time='event_time',\n",
    "    event_time_transformation='none',\n",
    "    cols_category=['small_group'],\n",
    "    cols_numerical=['amount_rur'],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "6ddba40c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "22/07/22 14:31:17 WARN WindowExec: No Partition Defined for Window operation! Moving all data to a single partition, this can cause serious performance degradation.\n",
      "22/07/22 14:31:17 WARN WindowExec: No Partition Defined for Window operation! Moving all data to a single partition, this can cause serious performance degradation.\n",
      "22/07/22 14:31:17 WARN WindowExec: No Partition Defined for Window operation! Moving all data to a single partition, this can cause serious performance degradation.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 15:=====================================================>(198 + 2) / 200]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "22/07/22 14:31:26 WARN WindowExec: No Partition Defined for Window operation! Moving all data to a single partition, this can cause serious performance degradation.\n",
      "22/07/22 14:31:26 WARN WindowExec: No Partition Defined for Window operation! Moving all data to a single partition, this can cause serious performance degradation.\n",
      "22/07/22 14:31:26 WARN WindowExec: No Partition Defined for Window operation! Moving all data to a single partition, this can cause serious performance degradation.\n",
      "22/07/22 14:31:26 WARN WindowExec: No Partition Defined for Window operation! Moving all data to a single partition, this can cause serious performance degradation.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r\n",
      "                                                                                \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 70.1 ms, sys: 65.8 ms, total: 136 ms\n",
      "Wall time: 9.78 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "df_gen_features = preprocessor.fit_transform(df_generated)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "58d7595e",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_gen_features = df_gen_features.withColumn('target', F.round(F.rand()).cast('int'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "58b7d569",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 23:=====================================================>(199 + 1) / 200]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 350 ms, sys: 32.1 ms, total: 382 ms\n",
      "Wall time: 8min 26s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r\n",
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "df_gen_features.write.parquet('df_gen_features.parquet', mode='overwrite')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "d8eb8287",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "9.5G\tdf_gen_features.parquet\r\n"
     ]
    }
   ],
   "source": [
    "!du -sh df_gen_features.parquet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "c8668b16",
   "metadata": {},
   "outputs": [],
   "source": [
    "spark.stop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "843e1c2a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'small_group': 204}"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "preprocessor.get_category_dictionary_sizes()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83ab7d74",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "d704f4b7",
   "metadata": {},
   "source": [
    "# Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "59fd372b",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e2052c16",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import torch\n",
    "import pytorch_lightning as pl\n",
    "\n",
    "from tqdm.auto import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f248df9",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "b254e6c3",
   "metadata": {},
   "source": [
    "## Data for inferenfce"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "893d9b0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ptls.data_load.datasets import ParquetDataset, ParquetFiles\n",
    "from ptls.data_load.iterable_processing import ISeqLenLimit\n",
    "from ptls.data_load.utils import collate_feature_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "1fe1dec2",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'user_id': 'u_100002',\n",
       " 'trans_date': tensor([425., 297., 227., 199., 165., 236., 387., 608.,  18., 160., 569., 586.,\n",
       "         641., 650., 605., 579., 551., 568., 702., 383.,  34., 422., 221., 342.,\n",
       "         719., 255.,  57., 492.,  58.,  61., 509., 181., 396., 601., 253., 303.,\n",
       "         663., 532., 521., 519.,  55., 400., 399., 322., 539., 638., 531., 317.,\n",
       "         133.,   5.,  32., 456., 239., 572., 266.,  33., 627., 306., 653., 100.,\n",
       "         650., 383.,  10., 175., 484., 694., 610., 226., 125., 577., 134., 115.,\n",
       "         577., 631., 102., 245.,  15., 568., 607., 164., 601., 476., 208., 657.,\n",
       "         524.,  47., 418., 379., 713.,  34.,  68., 177., 619., 466., 158., 425.,\n",
       "         250., 583., 341., 319., 382., 141., 652., 147.,  11., 607., 531., 507.,\n",
       "         692., 603.,   0., 374., 496., 138., 724., 414., 728.,  16., 360., 214.,\n",
       "         130.,  41., 348., 339., 567.,  56., 618., 113., 250., 189., 553., 445.,\n",
       "         707., 669., 512.,  16., 138., 687., 543., 364., 438., 668., 197., 414.,\n",
       "         303., 725., 483., 656., 372., 667., 504., 623., 131., 367.,  74., 213.,\n",
       "         536., 658.,  40., 275., 426., 385.,   9., 506., 130., 209., 224., 238.,\n",
       "         692., 221., 490.,  63., 456., 711.,  98., 683., 715., 650., 294., 686.,\n",
       "         698., 707., 652., 260., 524., 117., 529., 288., 581., 269., 396.,  16.,\n",
       "         363., 570.,   7., 372., 696., 522.,  40., 671., 151., 508., 520., 305.,\n",
       "         500., 122., 407., 601., 639., 720., 509., 626., 567., 728., 466., 470.,\n",
       "         128., 678., 721., 658., 626., 459., 442., 275., 233., 351., 438., 586.,\n",
       "          26., 114., 697., 164.,  59., 707.,  82., 672., 626., 106., 693., 353.,\n",
       "         551., 442., 555., 195., 350., 312., 718., 357., 325., 586., 699., 278.,\n",
       "         518., 245.,  99., 450., 116., 522., 247., 182., 697., 108., 600., 495.,\n",
       "         607., 594., 567., 291., 244., 700., 717., 175., 197., 107., 611., 384.,\n",
       "         227., 464., 680., 428., 238., 440., 692., 251., 192., 478., 551., 240.,\n",
       "         183., 662.,  35., 433., 266., 536., 462., 594., 220.,  19., 181., 162.,\n",
       "          72., 541., 185., 313., 357.,  25., 337., 598., 426., 628., 478., 587.,\n",
       "          39., 241., 345.,  81., 264., 112., 690., 413., 603., 342., 465., 339.,\n",
       "          55., 332., 333., 350.,  94., 209., 420.,  40., 540., 295., 561., 118.,\n",
       "         689., 617., 548., 546., 453., 668., 475.,  85., 169., 533., 683., 656.,\n",
       "         509.,  92., 570., 491.,  59., 210., 671.,  89.,  31., 387., 656., 447.,\n",
       "         134., 705.,  44., 273., 416., 645., 312., 111., 125., 715.,  82., 481.,\n",
       "          31.,  61., 372., 686., 589., 241., 211., 339., 559., 626., 399., 492.,\n",
       "         627., 221., 186.,  17., 526., 542., 194., 239., 193., 663.,  21., 180.,\n",
       "         383., 503., 279., 663., 325., 155., 141., 261.,  83., 661., 540., 316.,\n",
       "         725., 181., 392., 114.,  20., 460., 444., 576., 243., 506., 345.,  61.,\n",
       "         725., 120., 472.,   3., 317., 661., 715., 588., 634., 332., 725., 150.,\n",
       "         251.,  62., 529., 417., 596., 723., 122., 702., 272., 583., 363., 522.,\n",
       "         329., 537., 326., 542., 361., 586., 104., 605., 403., 612., 518., 307.,\n",
       "         401., 508., 611., 554., 483.,  64., 405., 258., 470., 646., 516., 678.,\n",
       "         628., 190.,  23., 696., 608., 474., 413., 308., 703., 580., 430., 565.,\n",
       "         486., 273., 492., 301., 490., 411., 638., 463., 650., 401., 412., 285.,\n",
       "         512., 320., 590., 412.,  29., 529., 515., 506., 650., 183., 336., 714.,\n",
       "         446., 591., 537., 672., 667., 188., 177.,  46., 631.,  30., 628., 714.,\n",
       "         503., 693., 316., 387., 695.,  88., 661., 598., 523., 289., 379., 653.,\n",
       "         695., 261., 159., 168.,   9., 121., 581., 263., 498.,  39.,   9., 249.,\n",
       "         604., 652., 334., 682., 372., 218., 377., 334., 246.,  94.,  29., 416.,\n",
       "         569., 437., 591.,  86., 336., 229., 323., 558., 407.,  50., 636., 430.,\n",
       "          52., 298., 181., 494.,  18., 130., 706., 403., 284., 612., 376., 456.,\n",
       "         305., 426., 687.,  11., 342., 198., 238., 578.,  17., 180., 113., 691.,\n",
       "         204., 573., 482., 636.,  23., 670.,  69., 633., 515., 293., 691., 453.,\n",
       "         273., 390., 256., 724., 728., 364., 521., 432., 410., 704., 278., 629.,\n",
       "         578., 605., 430., 611.,  85., 223., 132., 303., 376., 434., 299., 559.,\n",
       "         477., 708., 183., 467., 354., 120., 681., 499., 652., 550., 487., 696.,\n",
       "         533., 330., 218., 678., 424., 539., 408.,  10.,  49., 485., 588., 151.,\n",
       "         122., 546., 713., 648., 223., 185., 229., 583., 720., 467., 679., 295.,\n",
       "         499., 719., 724., 702., 547., 264., 254., 621., 447., 688., 274.,  97.,\n",
       "         181., 245., 131., 410., 536., 542., 387., 603., 578., 506., 713., 123.,\n",
       "         431., 334., 408., 158.,  36., 662., 364., 209., 153., 625., 175., 701.,\n",
       "         530., 575., 483., 624., 397., 323., 604., 676., 372., 254., 682., 128.,\n",
       "         354., 647., 446.,  19., 458.,   6., 431., 687., 679., 163., 257., 719.,\n",
       "         644., 201., 198., 290.,  58., 133., 390., 309., 339.,  96.,  22., 246.,\n",
       "         109.,  15., 538., 446., 638., 521., 481., 126., 717., 285., 254., 535.,\n",
       "         723., 317., 322., 300.,  80., 485., 546., 123., 370., 129., 328., 192.,\n",
       "         251., 451., 261., 589.,  34., 620., 354., 652., 501., 577., 658., 108.,\n",
       "         317., 637., 541., 412., 631., 261., 623., 703., 397., 711., 565., 589.,\n",
       "         613., 593., 284., 714., 663., 138., 331., 406.,  77., 501., 719., 363.,\n",
       "         531., 591., 110., 535.]),\n",
       " 'small_group': tensor([45,  1,  3, 10,  9,  2,  3, 13,  1,  8, 26, 79,  1,  1,  1,  5,  3,  2,\n",
       "          8,  1,  7,  3,  3, 24, 47,  1, 39,  5,  1, 39,  2,  9,  2,  1,  1,  1,\n",
       "          4,  4,  1,  1,  2, 34,  1,  6,  1, 10,  6,  1,  5,  2,  4,  5, 47,  3,\n",
       "          3, 10,  1,  1,  1,  1,  8,  9,  1,  1, 13,  2,  1, 14,  8,  2,  1,  1,\n",
       "          3,  8, 16, 13,  6,  1, 66,  1, 16,  1, 18,  6,  4,  1,  1, 23,  1, 12,\n",
       "          7,  1,  1,  3,  1,  1,  5,  1, 33, 36,  4, 27,  2,  1,  7, 39, 61, 28,\n",
       "          1,  4,  3,  1,  9,  1, 23, 22,  1,  7, 44,  6, 37,  2,  1, 16,  2,  8,\n",
       "          2, 56,  4,  2,  1,  3,  2,  9, 48,  1,  1,  6,  8, 10,  1,  4,  7, 20,\n",
       "          1, 13,  4, 38,  7,  8,  5, 26,  1,  1,  3,  3,  1,  1, 16, 23, 61,  8,\n",
       "          4,  1,  3,  1,  8, 17,  1,  1,  1,  1, 11,  8,  3,  1, 11,  1,  4, 15,\n",
       "         12,  1,  4,  1,  1,  5, 22,  2, 13,  2, 17,  1,  8,  1,  8,  8,  8,  1,\n",
       "          1, 11,  9, 46,  2,  2,  1,  4,  1,  3,  1,  1, 30,  9,  1,  4,  2,  5,\n",
       "          1,  1, 15,  1,  1, 39,  1, 83,  1,  1,  1, 43,  1, 45,  7,  7,  3,  1,\n",
       "          2,  8, 42, 90,  5, 13,  1,  6, 81,  1,  5, 22,  2, 13,  1, 26, 15,  7,\n",
       "          1,  1,  3,  5, 68,  1,  2,  3,  1,  1,  5, 32,  1, 20,  6,  1,  4,  5,\n",
       "          6,  2,  5, 21,  6,  2, 50, 16, 12,  2,  5,  1,  1,  1, 18,  4,  2, 46,\n",
       "          1,  6,  1,  4,  4,  9,  1,  6,  1,  1,  1, 10,  4,  8,  1,  3,  2,  2,\n",
       "          2, 32,  2,  1,  3, 50,  2,  1,  2,  1,  2,  8,  4, 25,  1,  2,  3,  9,\n",
       "          4, 61,  2, 52,  5,  3,  3,  4,  5,  1,  1,  3,  2,  1,  5,  5,  7,  6,\n",
       "          1,  4, 10,  2,  6,  2,  4,  8,  7,  3,  2, 49,  3,  1,  4,  3,  3,  2,\n",
       "          9, 44,  3,  3,  1,  9,  1,  1, 49,  2,  3,  1,  2, 12,  3,  2,  2,  3,\n",
       "         18,  9,  1,  4, 12,  1, 94,  3,  1, 16,  7,  2,  2,  6, 36,  2,  4,  5,\n",
       "          4, 11,  1,  6, 10,  2,  9, 30,  2, 55,  2,  6,  2,  1,  4,  4,  1,  1,\n",
       "          1, 46,  1, 11,  8,  2,  5,  7, 29,  2,  5, 12,  2, 10,  2,  2,  2, 66,\n",
       "          1,  1,  9,  6,  6,  2, 18,  2, 11,  5,  1,  1,  1,  3,  4,  7,  3,  1,\n",
       "          7,  5,  9,  2,  2,  2,  7,  1,  6,  1,  2, 10, 21,  1, 39,  1,  3,  7,\n",
       "          2,  2,  1, 10,  1,  1,  1,  1,  3,  4, 13, 62, 15,  6,  1,  1, 39, 38,\n",
       "          2,  4,  9,  3,  1,  2,  2, 12,  2, 12,  6,  3,  2, 34,  1,  3,  1, 19,\n",
       "          1,  3,  2,  6,  8,  3,  7,  3,  9,  1, 53, 10,  1,  1,  1,  1, 40,  1,\n",
       "         48,  1,  1,  1,  5, 68,  2,  3, 47,  5,  4,  4,  4, 18, 37,  1,  2,  1,\n",
       "          1,  2, 10,  8,  1,  4,  7,  1,  1,  4, 23, 24, 12,  6, 13, 12,  1,  9,\n",
       "          3,  3,  4,  3,  8,  3, 28,  1,  9, 25,  1, 14, 13,  9, 34,  6,  2, 78,\n",
       "          2,  1, 11, 29,  7,  8, 24,  6,  1,  7,  1, 10,  1,  5,  5,  4,  3,  4,\n",
       "          1,  4, 37,  6,  5, 60,  3,  5, 19,  1,  1,  8,  3,  4,  1,  4,  1,  4,\n",
       "         52, 39, 33,  7,  1,  4, 22, 23, 36,  1,  1,  8,  4,  2,  6,  1,  1, 11,\n",
       "          6,  4,  2,  2, 92, 65,  1,  2, 14,  1,  4,  4, 22,  3,  1,  3,  6,  1,\n",
       "          1,  4,  4,  4,  9,  1,  6,  3,  1,  3,  1,  3,  4,  8,  1,  1,  2, 10,\n",
       "          1,  1, 10, 56, 31,  2,  1,  3,  2,  1,  3, 11, 34,  1,  4, 10,  1,  3,\n",
       "          3,  1,  1,  1,  4, 10,  7,  1, 17,  3,  4,  2,  1,  1,  1,  9,  1,  2,\n",
       "          1,  5,  3,  2, 15,  1,  1,  5,  1, 30,  1,  3,  1,  2,  2, 10,  7,  5,\n",
       "         13,  5,  3, 66,  4, 14,  1,  2,  5,  3,  1,  2, 33, 28,  5,  9,  5, 10,\n",
       "          1,  3,  3, 13, 29, 14,  1,  1,  1,  7,  3,  1,  9,  7,  1,  2, 12,  2,\n",
       "          1,  1,  5,  6,  1,  3, 74,  1,  1,  3,  1,  5,  6,  2,  1,  1,  3,  6,\n",
       "          6,  2,  2,  3,  1,  2,  1,  1,  2, 11,  9,  1, 40,  5, 28,  1,  3, 32,\n",
       "         12, 13,  6,  8], dtype=torch.int32),\n",
       " 'amount_rur': tensor([6.6244e+02, 7.8070e+00, 4.5790e+00, 4.3300e+01, 3.5370e+00, 6.4690e+00,\n",
       "         5.7702e+01, 1.8933e+01, 1.5569e+01, 5.3580e+00, 9.9366e+01, 5.4038e+01,\n",
       "         7.3180e+00, 5.3910e+00, 6.9820e+00, 3.1000e-01, 2.2898e+01, 2.4272e+01,\n",
       "         1.6432e+01, 1.2135e+01, 2.1564e+01, 2.2898e+01, 4.5795e+01, 1.1709e+02,\n",
       "         5.3910e+01, 5.2284e+01, 4.5790e+00, 9.0130e+00, 1.0543e+01, 9.1590e+00,\n",
       "         4.5750e+00, 1.5095e+01, 5.3580e+00, 4.4380e+00, 2.4180e+01, 6.1250e+00,\n",
       "         2.1980e+00, 1.7768e+01, 8.2710e+00, 1.7876e+01, 1.7173e+01, 4.3081e+02,\n",
       "         1.7941e+01, 2.3677e+01, 3.5390e+01, 7.0520e+00, 1.7768e+02, 1.0772e+02,\n",
       "         2.0787e+01, 2.7480e+00, 1.3948e+01, 7.4180e+00, 1.1907e+02, 1.8550e+01,\n",
       "         2.1560e+00, 2.1598e+02, 1.0534e+01, 1.3153e+02, 4.5790e+00, 2.3030e+01,\n",
       "         5.0847e+01, 1.2507e+01, 2.1799e+01, 6.4907e+01, 2.5828e+01, 3.3431e+01,\n",
       "         2.9258e+01, 3.2303e+01, 1.7173e+01, 1.0991e+01, 3.3746e+01, 4.8039e+01,\n",
       "         4.3127e+01, 2.7213e+01, 3.0770e+01, 1.6074e+01, 1.4472e+01, 4.1900e+01,\n",
       "         3.3994e+01, 3.1940e+01, 1.6820e+01, 4.4331e+01, 2.0104e+02, 1.8502e+01,\n",
       "         1.2076e+01, 1.1023e+01, 3.9668e+01, 6.4690e+00, 1.2868e+02, 1.5571e+01,\n",
       "         9.1591e+01, 2.7661e+01, 2.4655e+01, 4.5795e+01, 2.6137e+01, 1.3318e+02,\n",
       "         6.8570e+00, 3.6647e+01, 1.6173e+01, 1.8167e+02, 1.7036e+01, 2.1132e+02,\n",
       "         1.0375e+01, 1.4617e+01, 1.0688e+02, 2.2900e+00, 1.3281e+01, 6.7290e+01,\n",
       "         4.2559e+01, 1.0524e+01, 3.8815e+01, 8.1790e+00, 2.3720e+00, 1.5224e+01,\n",
       "         4.3127e+01, 8.2432e+01, 1.2964e+01, 1.2934e+02, 1.8776e+02, 1.1311e+02,\n",
       "         5.1753e+01, 2.9721e+01, 3.0910e+00, 4.1187e+01, 2.4561e+01, 1.8933e+01,\n",
       "         1.6289e+01, 1.0808e+02, 7.8035e+01, 5.1320e+00, 1.7906e+01, 1.2938e+01,\n",
       "         3.8560e+01, 2.9320e+00, 2.0264e+02, 3.8010e+01, 1.1607e+01, 2.1132e+01,\n",
       "         1.6474e+01, 3.4071e+02, 1.7026e+01, 3.9810e+00, 2.2898e+01, 1.0991e+02,\n",
       "         1.9463e+01, 1.3509e+01, 4.1346e+01, 9.1600e-01, 3.2041e+01, 1.4620e+01,\n",
       "         8.0783e+01, 1.2731e+01, 7.5387e+01, 9.6822e+01, 4.5790e+00, 9.1591e+01,\n",
       "         5.1538e+01, 5.3470e+01, 2.2050e+02, 1.8114e+01, 1.2135e+01, 7.8490e+00,\n",
       "         6.9010e+00, 2.6528e+01, 8.6250e+00, 5.5092e+01, 1.1403e+01, 1.0538e+02,\n",
       "         2.1906e+01, 1.0991e+02, 3.2197e+01, 2.2234e+01, 1.1989e+01, 5.5081e+01,\n",
       "         4.5790e+00, 6.3752e+01, 1.5177e+02, 5.0483e+01, 5.2660e+00, 2.3230e+00,\n",
       "         5.9039e+01, 4.3649e+01, 3.4350e+00, 1.2960e+01, 4.4465e+01, 1.8012e+01,\n",
       "         7.2356e+01, 8.1510e+00, 2.3246e+01, 3.6170e+00, 8.2495e+01, 4.3276e+01,\n",
       "         1.0990e+00, 3.6640e+00, 6.5130e+00, 9.6170e+00, 4.8730e+00, 7.3101e+01,\n",
       "         1.4105e+02, 8.9209e+01, 1.7402e+01, 3.1829e+01, 9.9190e+00, 3.3980e+01,\n",
       "         1.5283e+01, 1.4534e+01, 1.5698e+01, 1.3738e+01, 4.9738e+01, 2.9137e+01,\n",
       "         4.6949e+01, 4.0300e+00, 1.1975e+01, 5.2980e+00, 2.7386e+01, 8.0600e+00,\n",
       "         3.7305e+01, 6.7581e+01, 3.3640e+00, 1.5090e+00, 1.5232e+01, 1.0351e+01,\n",
       "         1.9234e+01, 7.2810e+00, 3.1315e+01, 2.7340e+01, 9.7690e+00, 4.3560e+00,\n",
       "         2.4809e+01, 8.0650e+00, 6.8701e+01, 8.8467e+01, 4.3130e+00, 8.7584e+01,\n",
       "         7.4610e+00, 2.9629e+01, 3.1156e+02, 1.7467e+02, 4.0540e+00, 7.2815e+01,\n",
       "         7.0785e+01, 5.4772e+01, 9.1000e+00, 1.3221e+01, 3.9634e+01, 1.2938e+01,\n",
       "         1.1861e+01, 3.0362e+01, 3.7960e+00, 1.7410e+00, 2.3585e+01, 5.3910e+00,\n",
       "         5.9779e+01, 2.4107e+01, 3.0190e+00, 4.0066e+01, 1.2895e+01, 2.2601e+01,\n",
       "         1.0135e+01, 1.2938e+01, 2.1101e+02, 1.3024e+01, 1.0625e+02, 9.1000e+00,\n",
       "         2.7981e+01, 9.1132e+01, 1.2076e+01, 1.5224e+01, 6.1250e+00, 6.6420e+00,\n",
       "         3.6660e+00, 2.4195e+01, 6.4131e+01, 1.3097e+01, 2.1520e+01, 9.8330e+00,\n",
       "         7.2360e+00, 7.2360e+00, 2.3030e+01, 3.6316e+01, 2.3000e+01, 6.5490e+00,\n",
       "         2.9647e+01, 1.0235e+01, 2.5553e+01, 1.1407e+01, 1.0872e+02, 2.2446e+01,\n",
       "         1.7106e+02, 8.0060e+00, 9.9879e+01, 2.3148e+01, 5.1926e+01, 7.6120e+01,\n",
       "         4.6508e+01, 3.2515e+01, 2.8982e+01, 1.4103e+01, 1.3569e+02, 4.5365e+01,\n",
       "         1.7220e+01, 1.5353e+01, 2.5650e+00, 2.0097e+01, 3.9840e+00, 2.1647e+01,\n",
       "         6.8870e+00, 2.0602e+02, 2.1560e+00, 9.2210e+00, 4.3130e+00, 1.0991e+01,\n",
       "         2.8348e+01, 4.7633e+01, 1.8319e+01, 4.7731e+01, 1.4577e+01, 1.6000e+01,\n",
       "         1.2938e+01, 4.3130e+00, 7.1970e+00, 1.4965e+01, 1.3740e+00, 2.2898e+01,\n",
       "         4.4200e+00, 1.2895e+01, 9.5950e+00, 1.6350e+02, 4.2101e+01, 1.7251e+01,\n",
       "         1.7251e+01, 9.7470e+00, 7.3300e-01, 4.3986e+01, 5.4400e+00, 1.2938e+01,\n",
       "         1.6518e+01, 7.2424e+01, 4.1220e+00, 8.2890e+00, 8.4630e+01, 1.1540e+01,\n",
       "         1.1450e+00, 2.0380e+00, 7.7390e+00, 5.9950e+00, 6.4120e+00, 1.0954e+01,\n",
       "         1.8027e+01, 5.0460e+00, 6.8693e+01, 3.2057e+01, 2.4729e+01, 2.1530e+00,\n",
       "         1.8319e+01, 4.1630e+00, 9.8460e+00, 4.3130e+00, 1.3738e+01, 2.9309e+01,\n",
       "         7.3316e+01, 1.3326e+02, 8.6250e+00, 1.2938e+01, 4.9788e+01, 1.0075e+01,\n",
       "         2.3720e+00, 3.1461e+01, 8.8584e+01, 2.6607e+01, 2.2900e+00, 4.9002e+01,\n",
       "         9.8000e+00, 6.1640e+01, 1.0782e+01, 2.5186e+01, 9.7040e+00, 3.8815e+01,\n",
       "         1.1562e+01, 8.0140e+00, 1.4931e+01, 3.9250e+00, 2.9996e+01, 1.0566e+01,\n",
       "         1.8673e+01, 7.8600e-01, 6.0450e+00, 2.3426e+01, 4.3123e+01, 9.7040e+00,\n",
       "         6.4690e+00, 1.3281e+01, 7.3227e+01, 1.9796e+01, 2.1560e+00, 2.7844e+01,\n",
       "         1.8001e+01, 1.8318e+02, 1.3143e+01, 8.7550e+00, 1.5896e+02, 1.1860e+01,\n",
       "         7.3316e+01, 1.3826e+02, 7.5110e+00, 6.1820e+00, 1.0075e+01, 5.9530e+00,\n",
       "         2.7386e+01, 4.5284e+01, 6.5550e+00, 7.9226e+01, 1.9666e+01, 2.6829e+01,\n",
       "         1.5966e+01, 2.2898e+01, 6.5507e+01, 5.1710e+01, 4.3130e+00, 2.1132e+01,\n",
       "         1.4879e+01, 6.2856e+01, 4.6262e+01, 3.3855e+01, 4.5280e+00, 7.5040e+00,\n",
       "         1.5818e+01, 2.2943e+01, 3.1782e+01, 1.2076e+01, 1.5112e+01, 6.8573e+01,\n",
       "         4.0588e+01, 9.1060e+00, 7.6120e+01, 5.4950e+00, 1.0163e+01, 1.2822e+01,\n",
       "         4.7440e+00, 3.6640e+00, 6.8918e+01, 7.1215e+01, 2.3291e+01, 1.5581e+01,\n",
       "         8.0398e+01, 4.3130e+00, 1.0653e+01, 2.7515e+01, 1.0782e+01, 7.3621e+01,\n",
       "         1.2181e+02, 2.4042e+01, 4.7440e+00, 1.1601e+01, 9.4340e+00, 4.2700e+00,\n",
       "         1.0479e+02, 7.4164e+01, 5.9991e+01, 1.7807e+01, 6.4690e+00, 1.9408e+01,\n",
       "         3.9936e+01, 8.2249e+01, 5.4950e+00, 3.9699e+01, 7.3270e+00, 1.0016e+02,\n",
       "         1.0533e+01, 1.0264e+02, 2.4583e+01, 1.3601e+01, 2.5010e+01, 3.2234e+01,\n",
       "         1.5871e+01, 2.2303e+01, 8.6250e+00, 1.2723e+01, 1.7768e+01, 1.9579e+01,\n",
       "         1.7579e+02, 1.9495e+02, 2.2943e+01, 1.1429e+01, 3.5790e+00, 2.0013e+01,\n",
       "         5.4950e+00, 2.5180e+00, 1.8970e+00, 2.5877e+01, 4.1261e+01, 1.2999e+01,\n",
       "         8.6040e+00, 5.4490e+00, 8.8850e+00, 3.3934e+01, 2.6696e+01, 1.3738e+01,\n",
       "         3.0190e+00, 3.0189e+01, 1.1577e+02, 1.7251e+01, 7.2684e+01, 8.8412e+01,\n",
       "         5.6650e+00, 3.0189e+01, 1.6518e+01, 1.3197e+01, 1.5846e+01, 4.5790e+00,\n",
       "         4.5789e+01, 2.1560e+00, 1.5095e+01, 1.3133e+02, 4.6138e+01, 1.7682e+01,\n",
       "         1.8135e+01, 1.7377e+01, 1.3038e+01, 7.3180e+00, 1.3738e+01, 4.1101e+01,\n",
       "         5.2070e+00, 1.7238e+01, 1.9504e+02, 1.9581e+01, 2.3720e+01, 1.8272e+01,\n",
       "         2.3063e+01, 9.1590e+00, 1.9234e+01, 6.0770e+01, 8.9710e+00, 5.3910e+00,\n",
       "         5.0200e+01, 1.2895e+01, 5.6109e+01, 2.9801e+01, 1.5397e+01, 7.4610e+00,\n",
       "         3.8920e+00, 1.6030e+00, 2.2254e+01, 1.7632e+01, 5.2660e+00, 1.3370e+01,\n",
       "         5.6175e+01, 5.1150e+00, 2.1219e+01, 8.5830e+00, 1.3648e+01, 7.1898e+01,\n",
       "         9.9410e+00, 1.5224e+01, 8.7980e+00, 4.6200e+00, 7.2260e+01, 7.0520e+00,\n",
       "         6.7260e+00, 2.2898e+01, 4.5629e+01, 4.3130e+00, 1.2119e+01, 2.1564e+01,\n",
       "         6.8901e+01, 5.9950e+00, 5.0380e+00, 4.0575e+01, 1.2464e+01, 7.7390e+00,\n",
       "         7.8265e+01, 1.7402e+01, 3.3967e+02, 6.8690e+00, 3.3019e+01, 1.2507e+02,\n",
       "         9.4020e+00, 2.1797e+01, 3.1310e+02, 9.9798e+01, 5.1750e+00, 5.1360e+00,\n",
       "         8.2432e+01, 9.9370e+00, 8.0463e+01, 2.1690e+02, 1.0653e+01, 5.9450e+00,\n",
       "         5.9530e+00, 6.4754e+01, 3.8681e+01, 2.5180e+00, 4.3127e+01, 4.3130e+00,\n",
       "         2.9103e+01, 7.7850e+00, 6.7754e+01, 5.3809e+01, 4.5060e+00, 3.4255e+01,\n",
       "         9.1591e+01, 2.9320e+00, 2.5531e+01, 1.8011e+02, 1.0207e+02, 2.0425e+01,\n",
       "         2.6103e+01, 3.1802e+01, 4.4261e+01, 1.9967e+01, 1.4333e+01, 5.7140e+00,\n",
       "         1.3601e+02, 6.8690e+00, 9.1132e+01, 5.8626e+01, 1.1219e+01, 1.6940e+00,\n",
       "         1.3739e+02, 3.0189e+01, 5.1667e+01, 8.8110e+00, 3.7965e+01, 2.6079e+01,\n",
       "         9.6170e+00, 2.2440e+01, 1.5112e+01, 7.4084e+01, 4.1463e+02, 3.2011e+01,\n",
       "         6.9263e+01, 3.4350e+00, 1.5483e+01, 2.5440e+00, 7.3270e+00, 3.0410e+02,\n",
       "         4.7831e+01, 1.4655e+01, 1.3693e+01, 6.7320e+00, 1.2594e+01, 1.4975e+01,\n",
       "         3.7745e+02, 5.9530e+00, 1.7585e+01, 1.5095e+01, 2.1981e+01, 7.9614e+01,\n",
       "         2.5660e+00, 8.6100e+00, 7.7850e+00, 7.8920e+00, 9.4890e+00, 4.3510e+00,\n",
       "         3.6636e+01, 4.3130e+00, 4.7010e+00, 4.5790e+00, 1.0505e+01, 4.5790e+00,\n",
       "         5.5504e+01, 1.8089e+01, 3.2194e+01, 1.0070e+00, 8.5830e+00, 5.1750e+00,\n",
       "         1.2468e+01, 9.4880e+00, 4.1704e+01, 4.2700e+00, 1.7356e+02, 1.0991e+01,\n",
       "         2.1697e+02, 2.1560e+00, 1.5741e+01, 2.3120e+00, 2.1564e+01, 1.2550e+02,\n",
       "         5.7020e+02, 7.3270e+00, 7.3270e+00, 2.5600e+01, 2.1737e+01, 3.6640e+00,\n",
       "         2.2900e+00, 1.0912e+01, 2.3360e+01, 2.6959e+01, 3.4346e+01, 6.8270e+02,\n",
       "         1.3722e+01, 1.7338e+01, 4.5212e+01, 8.6250e+00, 4.7010e+00, 8.0220e+00,\n",
       "         3.1193e+01, 2.5140e+00, 3.3811e+01, 2.9320e+00, 6.2650e+00, 1.0991e+01,\n",
       "         9.8029e+01, 6.5260e+00, 2.2900e+00, 1.1213e+01, 5.1657e+01, 6.3182e+01,\n",
       "         5.4047e+01, 2.2715e+01, 5.9446e+01, 3.0386e+01, 7.3510e+00, 2.1560e+00,\n",
       "         1.2895e+01, 1.0991e+01, 3.0190e+00, 9.7938e+01, 4.5795e+01, 1.9321e+01,\n",
       "         2.2956e+02, 8.0140e+00, 1.7251e+01, 1.3740e+00, 5.7440e+00, 2.0745e+01,\n",
       "         6.7529e+01, 4.3130e+00, 5.8007e+01, 4.1215e+01, 1.8135e+01, 1.3281e+01,\n",
       "         3.8815e+01, 1.2074e+02, 2.2211e+01, 4.5790e+00, 1.4151e+01, 1.5051e+01,\n",
       "         1.9721e+01, 4.5800e-01, 1.7251e+01, 2.3264e+01, 1.1964e+01, 2.7020e+00,\n",
       "         7.6983e+01, 1.7639e+01, 2.6741e+02, 6.9131e+01, 4.8310e+00, 2.8181e+01,\n",
       "         1.6388e+01, 3.1224e+01, 4.3510e+00, 1.8349e+01, 8.2420e+00, 1.1386e+01,\n",
       "         9.7840e+00, 3.2731e+01, 1.6850e+02, 5.3480e+00, 5.2016e+01, 9.1590e+00,\n",
       "         4.2773e+01, 6.8690e+01, 3.8916e+01, 1.4448e+02, 3.7152e+01, 2.1560e+00,\n",
       "         1.6807e+01, 5.0380e+00, 3.1600e+00, 1.4869e+02, 1.3739e+02, 1.7860e+01,\n",
       "         1.1907e+01, 3.3888e+01, 1.7208e+01, 4.5795e+01, 1.0898e+02, 1.2464e+01,\n",
       "         8.3182e+01, 2.6278e+01, 1.8502e+01, 1.5051e+02, 2.2898e+01, 2.5272e+01,\n",
       "         1.1609e+01, 7.1160e+00, 2.0147e+02, 4.1384e+01, 4.3130e+00, 1.6944e+02,\n",
       "         2.0120e+01, 2.7429e+01, 1.8319e+01, 1.1540e+01]),\n",
       " 'event_time': tensor([0.0016, 0.0019, 0.0021, 0.0035, 0.0057, 0.0075, 0.0089, 0.0095, 0.0108,\n",
       "         0.0114, 0.0122, 0.0134, 0.0141, 0.0148, 0.0155, 0.0166, 0.0171, 0.0172,\n",
       "         0.0178, 0.0204, 0.0204, 0.0213, 0.0239, 0.0273, 0.0275, 0.0276, 0.0310,\n",
       "         0.0313, 0.0314, 0.0321, 0.0345, 0.0346, 0.0347, 0.0347, 0.0392, 0.0407,\n",
       "         0.0409, 0.0428, 0.0454, 0.0460, 0.0465, 0.0499, 0.0505, 0.0507, 0.0511,\n",
       "         0.0515, 0.0518, 0.0522, 0.0523, 0.0531, 0.0566, 0.0566, 0.0579, 0.0601,\n",
       "         0.0616, 0.0626, 0.0637, 0.0640, 0.0651, 0.0675, 0.0681, 0.0685, 0.0694,\n",
       "         0.0694, 0.0702, 0.0720, 0.0726, 0.0726, 0.0734, 0.0745, 0.0750, 0.0771,\n",
       "         0.0804, 0.0822, 0.0826, 0.0828, 0.0832, 0.0833, 0.0841, 0.0853, 0.0860,\n",
       "         0.0867, 0.0867, 0.0872, 0.0888, 0.0916, 0.0926, 0.0928, 0.0929, 0.0938,\n",
       "         0.0940, 0.0944, 0.0959, 0.0959, 0.0960, 0.0971, 0.0982, 0.0986, 0.0992,\n",
       "         0.1005, 0.1014, 0.1015, 0.1031, 0.1052, 0.1061, 0.1095, 0.1139, 0.1149,\n",
       "         0.1149, 0.1156, 0.1162, 0.1186, 0.1187, 0.1195, 0.1203, 0.1218, 0.1223,\n",
       "         0.1226, 0.1232, 0.1237, 0.1247, 0.1271, 0.1285, 0.1286, 0.1291, 0.1298,\n",
       "         0.1300, 0.1300, 0.1301, 0.1302, 0.1306, 0.1344, 0.1428, 0.1465, 0.1479,\n",
       "         0.1480, 0.1491, 0.1504, 0.1514, 0.1519, 0.1519, 0.1521, 0.1532, 0.1549,\n",
       "         0.1560, 0.1560, 0.1588, 0.1598, 0.1600, 0.1658, 0.1685, 0.1686, 0.1693,\n",
       "         0.1696, 0.1699, 0.1699, 0.1701, 0.1706, 0.1723, 0.1735, 0.1754, 0.1758,\n",
       "         0.1762, 0.1767, 0.1779, 0.1815, 0.1818, 0.1820, 0.1820, 0.1824, 0.1826,\n",
       "         0.1858, 0.1871, 0.1887, 0.1905, 0.1928, 0.1950, 0.1962, 0.1991, 0.1995,\n",
       "         0.2000, 0.2007, 0.2034, 0.2044, 0.2054, 0.2063, 0.2067, 0.2097, 0.2102,\n",
       "         0.2116, 0.2117, 0.2121, 0.2123, 0.2138, 0.2139, 0.2152, 0.2167, 0.2177,\n",
       "         0.2181, 0.2182, 0.2189, 0.2193, 0.2216, 0.2226, 0.2233, 0.2259, 0.2262,\n",
       "         0.2269, 0.2282, 0.2288, 0.2311, 0.2314, 0.2327, 0.2330, 0.2347, 0.2378,\n",
       "         0.2379, 0.2414, 0.2419, 0.2422, 0.2427, 0.2449, 0.2450, 0.2455, 0.2471,\n",
       "         0.2536, 0.2576, 0.2585, 0.2597, 0.2605, 0.2639, 0.2641, 0.2659, 0.2682,\n",
       "         0.2683, 0.2695, 0.2710, 0.2717, 0.2733, 0.2747, 0.2760, 0.2772, 0.2774,\n",
       "         0.2774, 0.2776, 0.2791, 0.2791, 0.2805, 0.2807, 0.2837, 0.2841, 0.2848,\n",
       "         0.2873, 0.2876, 0.2892, 0.2896, 0.2899, 0.2918, 0.2939, 0.3010, 0.3040,\n",
       "         0.3044, 0.3044, 0.3046, 0.3075, 0.3091, 0.3116, 0.3154, 0.3156, 0.3158,\n",
       "         0.3227, 0.3241, 0.3243, 0.3244, 0.3245, 0.3294, 0.3303, 0.3312, 0.3334,\n",
       "         0.3335, 0.3343, 0.3346, 0.3354, 0.3354, 0.3379, 0.3388, 0.3411, 0.3412,\n",
       "         0.3441, 0.3464, 0.3504, 0.3517, 0.3527, 0.3533, 0.3537, 0.3550, 0.3605,\n",
       "         0.3610, 0.3616, 0.3621, 0.3626, 0.3629, 0.3631, 0.3639, 0.3651, 0.3660,\n",
       "         0.3676, 0.3694, 0.3695, 0.3704, 0.3713, 0.3714, 0.3733, 0.3740, 0.3785,\n",
       "         0.3788, 0.3795, 0.3799, 0.3803, 0.3809, 0.3823, 0.3872, 0.3873, 0.3877,\n",
       "         0.3884, 0.3933, 0.3961, 0.3969, 0.3969, 0.3982, 0.4013, 0.4031, 0.4042,\n",
       "         0.4139, 0.4140, 0.4158, 0.4182, 0.4188, 0.4192, 0.4219, 0.4227, 0.4245,\n",
       "         0.4248, 0.4251, 0.4255, 0.4292, 0.4292, 0.4293, 0.4307, 0.4334, 0.4377,\n",
       "         0.4401, 0.4416, 0.4458, 0.4461, 0.4469, 0.4484, 0.4513, 0.4545, 0.4559,\n",
       "         0.4573, 0.4578, 0.4612, 0.4672, 0.4689, 0.4692, 0.4704, 0.4705, 0.4714,\n",
       "         0.4729, 0.4735, 0.4742, 0.4759, 0.4771, 0.4773, 0.4784, 0.4786, 0.4817,\n",
       "         0.4817, 0.4820, 0.4832, 0.4869, 0.4884, 0.4889, 0.4893, 0.4899, 0.4906,\n",
       "         0.4917, 0.4919, 0.4922, 0.4966, 0.4969, 0.5006, 0.5012, 0.5021, 0.5023,\n",
       "         0.5040, 0.5043, 0.5044, 0.5057, 0.5069, 0.5091, 0.5095, 0.5108, 0.5110,\n",
       "         0.5119, 0.5126, 0.5182, 0.5190, 0.5191, 0.5193, 0.5196, 0.5210, 0.5218,\n",
       "         0.5275, 0.5278, 0.5288, 0.5300, 0.5300, 0.5304, 0.5316, 0.5329, 0.5346,\n",
       "         0.5353, 0.5413, 0.5416, 0.5426, 0.5427, 0.5433, 0.5438, 0.5446, 0.5473,\n",
       "         0.5482, 0.5488, 0.5489, 0.5522, 0.5528, 0.5533, 0.5546, 0.5556, 0.5562,\n",
       "         0.5564, 0.5572, 0.5582, 0.5587, 0.5599, 0.5608, 0.5628, 0.5630, 0.5631,\n",
       "         0.5642, 0.5663, 0.5670, 0.5676, 0.5688, 0.5696, 0.5697, 0.5721, 0.5723,\n",
       "         0.5730, 0.5749, 0.5761, 0.5766, 0.5768, 0.5798, 0.5836, 0.5846, 0.5867,\n",
       "         0.5871, 0.5892, 0.5907, 0.5912, 0.5930, 0.5945, 0.5974, 0.5986, 0.5990,\n",
       "         0.5996, 0.6000, 0.6007, 0.6018, 0.6040, 0.6042, 0.6064, 0.6069, 0.6075,\n",
       "         0.6075, 0.6082, 0.6084, 0.6086, 0.6098, 0.6117, 0.6143, 0.6151, 0.6152,\n",
       "         0.6178, 0.6181, 0.6216, 0.6222, 0.6239, 0.6245, 0.6247, 0.6272, 0.6273,\n",
       "         0.6286, 0.6294, 0.6315, 0.6321, 0.6322, 0.6327, 0.6331, 0.6377, 0.6393,\n",
       "         0.6405, 0.6422, 0.6425, 0.6428, 0.6435, 0.6470, 0.6480, 0.6484, 0.6495,\n",
       "         0.6500, 0.6516, 0.6540, 0.6551, 0.6561, 0.6582, 0.6636, 0.6649, 0.6655,\n",
       "         0.6686, 0.6687, 0.6699, 0.6706, 0.6736, 0.6737, 0.6743, 0.6763, 0.6768,\n",
       "         0.6769, 0.6781, 0.6820, 0.6852, 0.6859, 0.6896, 0.6907, 0.6907, 0.6918,\n",
       "         0.6920, 0.6945, 0.6948, 0.6953, 0.6956, 0.6957, 0.6990, 0.7000, 0.7004,\n",
       "         0.7010, 0.7011, 0.7013, 0.7024, 0.7050, 0.7067, 0.7074, 0.7075, 0.7088,\n",
       "         0.7107, 0.7112, 0.7116, 0.7122, 0.7169, 0.7170, 0.7171, 0.7171, 0.7178,\n",
       "         0.7185, 0.7187, 0.7189, 0.7194, 0.7230, 0.7248, 0.7256, 0.7270, 0.7289,\n",
       "         0.7293, 0.7297, 0.7369, 0.7408, 0.7410, 0.7415, 0.7416, 0.7418, 0.7442,\n",
       "         0.7464, 0.7475, 0.7506, 0.7524, 0.7568, 0.7593, 0.7597, 0.7603, 0.7617,\n",
       "         0.7620, 0.7656, 0.7671, 0.7673, 0.7709, 0.7713, 0.7716, 0.7735, 0.7737,\n",
       "         0.7740, 0.7746, 0.7760, 0.7790, 0.7802, 0.7812, 0.7818, 0.7818, 0.7837,\n",
       "         0.7848, 0.7854, 0.7858, 0.7862, 0.7874, 0.7879, 0.7889, 0.7904, 0.7912,\n",
       "         0.7914, 0.7928, 0.7934, 0.7940, 0.7953, 0.7964, 0.7980, 0.7984, 0.7984,\n",
       "         0.7986, 0.7987, 0.7990, 0.8001, 0.8028, 0.8040, 0.8077, 0.8094, 0.8100,\n",
       "         0.8106, 0.8119, 0.8120, 0.8127, 0.8128, 0.8175, 0.8212, 0.8219, 0.8222,\n",
       "         0.8234, 0.8234, 0.8236, 0.8240, 0.8260, 0.8267, 0.8328, 0.8330, 0.8331,\n",
       "         0.8359, 0.8382, 0.8391, 0.8415, 0.8458, 0.8461, 0.8472, 0.8484, 0.8487,\n",
       "         0.8493, 0.8497, 0.8498, 0.8508, 0.8510, 0.8516, 0.8522, 0.8530, 0.8553,\n",
       "         0.8563, 0.8570, 0.8577, 0.8583, 0.8641, 0.8660, 0.8666, 0.8683, 0.8700,\n",
       "         0.8703, 0.8704, 0.8726, 0.8730, 0.8736, 0.8756, 0.8759, 0.8772, 0.8794,\n",
       "         0.8803, 0.8805, 0.8839, 0.8875, 0.8876, 0.8877, 0.8888, 0.8900, 0.8904,\n",
       "         0.8921, 0.8925, 0.8930, 0.8931, 0.8934, 0.8935, 0.8937, 0.8954, 0.8955,\n",
       "         0.8969, 0.9015, 0.9016, 0.9017, 0.9028, 0.9030, 0.9054, 0.9056, 0.9068,\n",
       "         0.9069, 0.9082, 0.9102, 0.9110, 0.9120, 0.9122, 0.9122, 0.9124, 0.9204,\n",
       "         0.9204, 0.9220, 0.9241, 0.9253, 0.9258, 0.9258, 0.9264, 0.9283, 0.9293,\n",
       "         0.9293, 0.9298, 0.9303, 0.9337, 0.9393, 0.9417, 0.9435, 0.9440, 0.9441,\n",
       "         0.9478, 0.9480, 0.9496, 0.9500, 0.9502, 0.9515, 0.9522, 0.9525, 0.9535,\n",
       "         0.9547, 0.9549, 0.9557, 0.9588, 0.9605, 0.9612, 0.9614, 0.9622, 0.9643,\n",
       "         0.9677, 0.9711, 0.9719, 0.9756, 0.9782, 0.9794, 0.9799, 0.9812, 0.9827,\n",
       "         0.9848, 0.9857, 0.9877, 0.9880, 0.9894, 0.9904, 0.9928, 0.9954, 0.9960,\n",
       "         0.9971, 0.9977, 0.9982, 0.9992], dtype=torch.float64),\n",
       " 'target': 0}"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "iterable_inference_dataset = ParquetDataset(\n",
    "    data_files=ParquetFiles('df_gen_features.parquet').data_files,\n",
    "    i_filters=[\n",
    "        ISeqLenLimit(max_seq_len=2000), \n",
    "    ],\n",
    ")\n",
    "next(iter(iterable_inference_dataset))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "2aa2e33e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "685dba284be5427eaca48d27e1a69f3a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1min 1s, sys: 12 s, total: 1min 13s\n",
      "Wall time: 1min 10s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "for rec in tqdm(iterable_inference_dataset):\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "68346c45",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'user_id': array(['u_100002', 'u_100093', 'u_100292', 'u_100302', 'u_100467',\n",
       "        'u_100543', 'u_1006', 'u_100643', 'u_100697', 'u_100875',\n",
       "        'u_100906', 'u_10141', 'u_102102', 'u_102346', 'u_102662',\n",
       "        'u_102862', 'u_102942', 'u_102967', 'u_103141', 'u_103333',\n",
       "        'u_103364', 'u_103406', 'u_103412', 'u_103587', 'u_103616',\n",
       "        'u_103761', 'u_104104', 'u_104166', 'u_104405', 'u_104454',\n",
       "        'u_104702', 'u_104965', 'u_105212', 'u_10554', 'u_105673',\n",
       "        'u_106243', 'u_106396', 'u_106406', 'u_106428', 'u_106478',\n",
       "        'u_106919', 'u_1070', 'u_107084', 'u_107220', 'u_107749',\n",
       "        'u_10788', 'u_107905', 'u_108305', 'u_108622', 'u_108786',\n",
       "        'u_108924', 'u_108929', 'u_109173', 'u_10929', 'u_109433',\n",
       "        'u_109435', 'u_109465', 'u_110009', 'u_110042', 'u_110060',\n",
       "        'u_110168', 'u_110258', 'u_110574', 'u_110697', 'u_110745',\n",
       "        'u_110892', 'u_110909', 'u_111003', 'u_111037', 'u_111066',\n",
       "        'u_111100', 'u_111664', 'u_112134', 'u_112139', 'u_112228',\n",
       "        'u_112325', 'u_112507', 'u_112636', 'u_112728', 'u_112841',\n",
       "        'u_112972', 'u_113384', 'u_113651', 'u_113686', 'u_113741',\n",
       "        'u_113748', 'u_113869', 'u_113890', 'u_114293', 'u_114396',\n",
       "        'u_114432', 'u_114891', 'u_115532', 'u_115583', 'u_115609',\n",
       "        'u_115622', 'u_115973', 'u_116257', 'u_116603', 'u_116719',\n",
       "        'u_116933', 'u_117180', 'u_117231', 'u_117370', 'u_117412',\n",
       "        'u_117438', 'u_117465', 'u_117581', 'u_117605', 'u_117666',\n",
       "        'u_118432', 'u_11852', 'u_118706', 'u_119097', 'u_119203',\n",
       "        'u_119294', 'u_119314', 'u_119324', 'u_119505', 'u_119628',\n",
       "        'u_119659', 'u_119859', 'u_119962', 'u_12005', 'u_120088',\n",
       "        'u_120260', 'u_120286', 'u_120327', 'u_120789', 'u_120962',\n",
       "        'u_120984', 'u_121070', 'u_121130', 'u_121409', 'u_121432',\n",
       "        'u_12155', 'u_121694', 'u_121759', 'u_121885', 'u_121912',\n",
       "        'u_122120', 'u_122275', 'u_12243', 'u_12279', 'u_122813',\n",
       "        'u_122904', 'u_122907', 'u_122994', 'u_123289', 'u_123457',\n",
       "        'u_123488', 'u_123542', 'u_12402', 'u_124186', 'u_124255',\n",
       "        'u_124646', 'u_124745', 'u_124823', 'u_1249', 'u_125147',\n",
       "        'u_125393', 'u_126076', 'u_126754', 'u_126770', 'u_127011',\n",
       "        'u_127274', 'u_127361', 'u_127551', 'u_128112', 'u_128126',\n",
       "        'u_128196', 'u_128284', 'u_128714', 'u_128818', 'u_128849',\n",
       "        'u_129104', 'u_129154', 'u_1295', 'u_129669', 'u_129846',\n",
       "        'u_129930', 'u_129951', 'u_130040', 'u_130275', 'u_130368',\n",
       "        'u_130919', 'u_130947', 'u_131068', 'u_131139', 'u_131207',\n",
       "        'u_131326', 'u_131453', 'u_131551', 'u_131952', 'u_132058',\n",
       "        'u_132536', 'u_132575', 'u_132798', 'u_132956', 'u_133383',\n",
       "        'u_133420', 'u_133449', 'u_133451', 'u_133494', 'u_133543',\n",
       "        'u_133589', 'u_133804', 'u_133968', 'u_13449', 'u_135708',\n",
       "        'u_136227', 'u_136369', 'u_136568', 'u_137334', 'u_137415',\n",
       "        'u_137891', 'u_137947', 'u_137965', 'u_138015', 'u_138497',\n",
       "        'u_138507', 'u_138715', 'u_138718', 'u_138752', 'u_138876',\n",
       "        'u_138962', 'u_139133', 'u_139224', 'u_139650', 'u_139755',\n",
       "        'u_139891', 'u_139975', 'u_140026', 'u_140111', 'u_140135',\n",
       "        'u_140229', 'u_140342', 'u_140372', 'u_140749', 'u_140765',\n",
       "        'u_140875', 'u_140951', 'u_141076', 'u_141163', 'u_141318',\n",
       "        'u_142458', 'u_142464', 'u_142581', 'u_142618', 'u_14279',\n",
       "        'u_143065', 'u_143678', 'u_14390', 'u_14397', 'u_144049',\n",
       "        'u_144058', 'u_144192', 'u_144371', 'u_144482', 'u_145006',\n",
       "        'u_145072', 'u_145272', 'u_145274', 'u_145308', 'u_145320',\n",
       "        'u_145574', 'u_145625', 'u_145631', 'u_145719', 'u_145781',\n",
       "        'u_145833', 'u_146138', 'u_146170', 'u_146207', 'u_146270',\n",
       "        'u_146346', 'u_146552', 'u_146825', 'u_146843', 'u_147007',\n",
       "        'u_147589', 'u_14780', 'u_147807', 'u_147939', 'u_148451',\n",
       "        'u_148484', 'u_148752', 'u_148755', 'u_149708', 'u_149759',\n",
       "        'u_14980', 'u_149924', 'u_150063', 'u_150113', 'u_150240',\n",
       "        'u_150350', 'u_150543', 'u_150567', 'u_150750', 'u_150882',\n",
       "        'u_15102', 'u_151200', 'u_151847', 'u_151858', 'u_152122',\n",
       "        'u_152248', 'u_152334', 'u_152457', 'u_152546', 'u_152659',\n",
       "        'u_153086', 'u_153264', 'u_153510', 'u_15369', 'u_153850',\n",
       "        'u_154140', 'u_154142', 'u_154436', 'u_15445', 'u_154540',\n",
       "        'u_154899', 'u_155063', 'u_155106', 'u_155508', 'u_156606',\n",
       "        'u_157111', 'u_157168', 'u_157218', 'u_157420', 'u_157508',\n",
       "        'u_157618', 'u_157995', 'u_158014', 'u_158051', 'u_158151',\n",
       "        'u_158382', 'u_158559', 'u_159626', 'u_159810', 'u_159812',\n",
       "        'u_160423', 'u_160472', 'u_16083', 'u_161136', 'u_161194',\n",
       "        'u_161227', 'u_161685', 'u_161950', 'u_162089', 'u_162365',\n",
       "        'u_162814', 'u_163034', 'u_163384', 'u_163809', 'u_163823',\n",
       "        'u_163979', 'u_164157', 'u_164283', 'u_164783', 'u_165039',\n",
       "        'u_165164', 'u_165218', 'u_165242', 'u_165392', 'u_165399',\n",
       "        'u_165496', 'u_165557', 'u_165568', 'u_165689', 'u_165810',\n",
       "        'u_165987', 'u_166257', 'u_166433', 'u_166522', 'u_166636',\n",
       "        'u_16683', 'u_166857', 'u_167020', 'u_167616', 'u_167777',\n",
       "        'u_168086', 'u_168179', 'u_168237', 'u_168564', 'u_168575',\n",
       "        'u_168724', 'u_169378', 'u_169512', 'u_170303', 'u_17054',\n",
       "        'u_170865', 'u_171071', 'u_171406', 'u_171446', 'u_171662',\n",
       "        'u_171671', 'u_172107', 'u_172807', 'u_173088', 'u_173506',\n",
       "        'u_173507', 'u_173523', 'u_17360', 'u_173902', 'u_174303',\n",
       "        'u_174368', 'u_174385', 'u_174542', 'u_174682', 'u_174702',\n",
       "        'u_17491', 'u_175178', 'u_175462', 'u_175990', 'u_176236',\n",
       "        'u_177235', 'u_177436', 'u_177846', 'u_177901', 'u_178145',\n",
       "        'u_179110', 'u_179133', 'u_179521', 'u_179599', 'u_179724',\n",
       "        'u_179757', 'u_179771', 'u_179916', 'u_180538', 'u_18065',\n",
       "        'u_180841', 'u_181013', 'u_181297', 'u_181484', 'u_181579',\n",
       "        'u_18166', 'u_182260', 'u_182290', 'u_182310', 'u_182396',\n",
       "        'u_18251', 'u_182614', 'u_18270', 'u_183446', 'u_18359', 'u_18364',\n",
       "        'u_183899', 'u_184436', 'u_184512', 'u_184621', 'u_1847',\n",
       "        'u_184765', 'u_184832', 'u_185046', 'u_185091', 'u_185153',\n",
       "        'u_18531', 'u_185549', 'u_185689', 'u_185814', 'u_185829',\n",
       "        'u_18587', 'u_186108', 'u_186198', 'u_186241', 'u_187077',\n",
       "        'u_187145', 'u_187187', 'u_18732', 'u_187330', 'u_187331',\n",
       "        'u_187363', 'u_187421', 'u_187507', 'u_187529', 'u_187556',\n",
       "        'u_187636', 'u_188008', 'u_18815', 'u_188196', 'u_188372',\n",
       "        'u_188387', 'u_188630', 'u_188717', 'u_189020', 'u_189033',\n",
       "        'u_189284', 'u_189526', 'u_189531', 'u_189667', 'u_189822',\n",
       "        'u_190069', 'u_19018', 'u_190225', 'u_1903', 'u_190447',\n",
       "        'u_190565', 'u_190589', 'u_190921', 'u_19105', 'u_191143',\n",
       "        'u_191864', 'u_191952', 'u_192020', 'u_192289', 'u_192340',\n",
       "        'u_192483', 'u_192503', 'u_192626', 'u_192876', 'u_192984',\n",
       "        'u_193183', 'u_19350', 'u_193505', 'u_193525', 'u_193644',\n",
       "        'u_193842', 'u_19389', 'u_19438', 'u_194405', 'u_194730',\n",
       "        'u_194856', 'u_194914', 'u_194965', 'u_194967', 'u_195086',\n",
       "        'u_195256', 'u_195371', 'u_195379', 'u_195415', 'u_19546',\n",
       "        'u_195532', 'u_1959', 'u_195935', 'u_196240', 'u_196295',\n",
       "        'u_196358', 'u_196755', 'u_1970', 'u_197027', 'u_197405',\n",
       "        'u_197407', 'u_197490', 'u_197506', 'u_197661', 'u_197704',\n",
       "        'u_19786', 'u_197890', 'u_198003', 'u_198748', 'u_19881',\n",
       "        'u_19904', 'u_1992', 'u_199239', 'u_199420', 'u_199627', 'u_200',\n",
       "        'u_200248', 'u_200586', 'u_200619', 'u_200720', 'u_200849',\n",
       "        'u_200866', 'u_200872', 'u_201660', 'u_201707', 'u_201819',\n",
       "        'u_202395', 'u_20276', 'u_203038', 'u_20389', 'u_20418',\n",
       "        'u_204408', 'u_204532', 'u_204780', 'u_204861', 'u_204864',\n",
       "        'u_205053', 'u_20530', 'u_205423', 'u_205571', 'u_205670',\n",
       "        'u_206275', 'u_206621', 'u_207849', 'u_208521', 'u_208835',\n",
       "        'u_209305', 'u_209444', 'u_209628', 'u_209889', 'u_209891',\n",
       "        'u_209899', 'u_209900', 'u_209962', 'u_210441', 'u_210703',\n",
       "        'u_210812', 'u_210845', 'u_210918', 'u_211187', 'u_211468',\n",
       "        'u_2115', 'u_211949', 'u_212013', 'u_212209', 'u_212284',\n",
       "        'u_21285', 'u_212886', 'u_212927', 'u_213527', 'u_213943',\n",
       "        'u_213966', 'u_214448', 'u_214538', 'u_214621', 'u_214919',\n",
       "        'u_215429', 'u_215432', 'u_215894', 'u_216028', 'u_216425',\n",
       "        'u_216606', 'u_216957', 'u_217037', 'u_217119', 'u_21714',\n",
       "        'u_217196', 'u_217256', 'u_217465', 'u_217604', 'u_217794',\n",
       "        'u_217895', 'u_21825', 'u_218301', 'u_218415', 'u_218670',\n",
       "        'u_218883', 'u_21952', 'u_21961', 'u_219758', 'u_21977',\n",
       "        'u_219825', 'u_220055', 'u_220182', 'u_220411', 'u_220446',\n",
       "        'u_220545', 'u_220558', 'u_220630', 'u_220735', 'u_220749',\n",
       "        'u_22078', 'u_220925', 'u_220949', 'u_221013', 'u_221067',\n",
       "        'u_221155', 'u_221158', 'u_221210', 'u_221244', 'u_221348',\n",
       "        'u_221509', 'u_221712', 'u_221782', 'u_221894', 'u_222007',\n",
       "        'u_222066', 'u_222123', 'u_222179', 'u_222264', 'u_2223',\n",
       "        'u_222351', 'u_222522', 'u_222568', 'u_222767', 'u_223492',\n",
       "        'u_223651', 'u_223815', 'u_223907', 'u_224408', 'u_22484',\n",
       "        'u_225044', 'u_225133', 'u_225208', 'u_22529', 'u_225303',\n",
       "        'u_22572', 'u_225849', 'u_225981', 'u_226173', 'u_226276',\n",
       "        'u_226561', 'u_22672', 'u_226844', 'u_226938', 'u_227111',\n",
       "        'u_227155', 'u_227226', 'u_227333', 'u_227657', 'u_227832',\n",
       "        'u_227964', 'u_228326', 'u_228399', 'u_228401', 'u_228578',\n",
       "        'u_228606', 'u_228692', 'u_228883', 'u_229008', 'u_229106',\n",
       "        'u_229359', 'u_229521', 'u_230102', 'u_230234', 'u_230558',\n",
       "        'u_230569', 'u_230659', 'u_230931', 'u_231289', 'u_231345',\n",
       "        'u_231564', 'u_23186', 'u_231889', 'u_231988', 'u_232399',\n",
       "        'u_232750', 'u_232825', 'u_232830', 'u_232891', 'u_23295',\n",
       "        'u_233349', 'u_233441', 'u_233463', 'u_233510', 'u_233619',\n",
       "        'u_233688', 'u_233998', 'u_234036', 'u_234742', 'u_235050',\n",
       "        'u_235096', 'u_23513', 'u_235188', 'u_235256', 'u_235683',\n",
       "        'u_235750', 'u_236006', 'u_23607', 'u_236422', 'u_23686',\n",
       "        'u_237036', 'u_237128', 'u_237184', 'u_237337', 'u_23785',\n",
       "        'u_237863', 'u_237989', 'u_238229', 'u_238318', 'u_238437',\n",
       "        'u_238459', 'u_238651', 'u_238923', 'u_239033', 'u_239270',\n",
       "        'u_239448', 'u_239527', 'u_239634', 'u_239659', 'u_239712',\n",
       "        'u_239924', 'u_24020', 'u_240536', 'u_240707', 'u_240820',\n",
       "        'u_241032', 'u_241104', 'u_241165', 'u_241364', 'u_241772',\n",
       "        'u_241867', 'u_242690', 'u_242733', 'u_242744', 'u_242908',\n",
       "        'u_243191', 'u_243434', 'u_243495', 'u_24351', 'u_243700',\n",
       "        'u_243884', 'u_244175', 'u_244319', 'u_244410', 'u_244495',\n",
       "        'u_244660', 'u_244763', 'u_244815', 'u_244918', 'u_245146',\n",
       "        'u_245357', 'u_245534', 'u_245646', 'u_245741', 'u_245764',\n",
       "        'u_246035', 'u_246128', 'u_246136', 'u_2468', 'u_246883',\n",
       "        'u_246997', 'u_247016', 'u_24724', 'u_2474', 'u_247509', 'u_24783',\n",
       "        'u_247926', 'u_248003', 'u_248264', 'u_248333', 'u_248709',\n",
       "        'u_248915', 'u_248960', 'u_249000', 'u_249013', 'u_249274',\n",
       "        'u_249360', 'u_249395', 'u_249628', 'u_249657', 'u_24966',\n",
       "        'u_249668', 'u_249684', 'u_249788', 'u_249816', 'u_249903',\n",
       "        'u_249952', 'u_249993', 'u_25004', 'u_250309', 'u_250373',\n",
       "        'u_250386', 'u_250531', 'u_250644', 'u_250804', 'u_251207',\n",
       "        'u_251425', 'u_251562', 'u_251705', 'u_251869', 'u_251886',\n",
       "        'u_251896', 'u_252079', 'u_25228', 'u_252351', 'u_252566',\n",
       "        'u_253685', 'u_253869', 'u_254349', 'u_254500', 'u_254720',\n",
       "        'u_254758', 'u_254924', 'u_255353', 'u_255805', 'u_255895',\n",
       "        'u_255951', 'u_256008', 'u_256125', 'u_256173', 'u_256308',\n",
       "        'u_256478', 'u_25670', 'u_256728', 'u_256920', 'u_257098',\n",
       "        'u_257125', 'u_25721', 'u_257346', 'u_257381', 'u_257606',\n",
       "        'u_257762', 'u_257929', 'u_25798', 'u_25799', 'u_258040',\n",
       "        'u_258078', 'u_258408', 'u_258463', 'u_258772', 'u_258807',\n",
       "        'u_2590', 'u_259366', 'u_259421', 'u_259431', 'u_259433',\n",
       "        'u_25974', 'u_259891', 'u_260261', 'u_2604', 'u_260528',\n",
       "        'u_260609', 'u_260678', 'u_260729', 'u_260756', 'u_260912',\n",
       "        'u_261068', 'u_2611', 'u_261359', 'u_261539', 'u_261942',\n",
       "        'u_262184', 'u_262185', 'u_262247', 'u_262938', 'u_263059',\n",
       "        'u_263107', 'u_263151', 'u_263442', 'u_26349', 'u_26376',\n",
       "        'u_263792', 'u_264026', 'u_264590', 'u_265037', 'u_26543',\n",
       "        'u_265514', 'u_265608', 'u_265635', 'u_265705', 'u_265808',\n",
       "        'u_265846', 'u_265859', 'u_266399', 'u_266526', 'u_267101',\n",
       "        'u_267120', 'u_267303', 'u_268173', 'u_268379', 'u_268478',\n",
       "        'u_268770', 'u_268928', 'u_26916', 'u_269340', 'u_269416',\n",
       "        'u_269468', 'u_270083', 'u_270145', 'u_270290', 'u_270551',\n",
       "        'u_270639', 'u_270654', 'u_270764', 'u_271336', 'u_271507',\n",
       "        'u_271665', 'u_272155', 'u_272167', 'u_272376', 'u_272691',\n",
       "        'u_272753', 'u_273231', 'u_273443', 'u_273474', 'u_273516',\n",
       "        'u_273813', 'u_273956', 'u_273968', 'u_273973', 'u_274024',\n",
       "        'u_274140', 'u_27435', 'u_274399', 'u_274649', 'u_274726',\n",
       "        'u_274773', 'u_274840', 'u_275051', 'u_275227', 'u_275287',\n",
       "        'u_275441', 'u_275798', 'u_275853', 'u_27586', 'u_275879',\n",
       "        'u_275972', 'u_276026', 'u_276141', 'u_276243', 'u_276270',\n",
       "        'u_276365', 'u_276503', 'u_27656', 'u_276935', 'u_277096',\n",
       "        'u_277131', 'u_277166'], dtype='<U8'),\n",
       " 'trans_date': tensor([[425., 297., 227.,  ...,   0.,   0.,   0.],\n",
       "         [193., 465., 706.,  ...,   0.,   0.,   0.],\n",
       "         [712., 102., 433.,  ...,   0.,   0.,   0.],\n",
       "         ...,\n",
       "         [163., 640., 353.,  ...,   0.,   0.,   0.],\n",
       "         [582.,   7., 308.,  ...,   0.,   0.,   0.],\n",
       "         [178., 727., 515.,  ...,   0.,   0.,   0.]]),\n",
       " 'small_group': tensor([[45,  1,  3,  ...,  0,  0,  0],\n",
       "         [ 3,  2, 23,  ...,  0,  0,  0],\n",
       "         [ 6, 11,  1,  ...,  0,  0,  0],\n",
       "         ...,\n",
       "         [ 1, 23,  2,  ...,  0,  0,  0],\n",
       "         [ 6,  1,  1,  ...,  0,  0,  0],\n",
       "         [ 1,  3, 88,  ...,  0,  0,  0]], dtype=torch.int32),\n",
       " 'amount_rur': tensor([[662.4440,   7.8070,   4.5790,  ...,   0.0000,   0.0000,   0.0000],\n",
       "         [  2.1560,  13.7580,  22.8980,  ...,   0.0000,   0.0000,   0.0000],\n",
       "         [ 14.0160,  55.9800,  46.6160,  ...,   0.0000,   0.0000,   0.0000],\n",
       "         ...,\n",
       "         [  3.0190,  25.8770,  11.2870,  ...,   0.0000,   0.0000,   0.0000],\n",
       "         [ 11.3470,  21.9360,  24.6180,  ...,   0.0000,   0.0000,   0.0000],\n",
       "         [ 36.7450,  15.0950,  92.4610,  ...,   0.0000,   0.0000,   0.0000]]),\n",
       " 'event_time': tensor([[1.5926e-03, 1.9414e-03, 2.0840e-03,  ..., 0.0000e+00, 0.0000e+00,\n",
       "          0.0000e+00],\n",
       "         [1.0965e-03, 1.5581e-03, 5.0477e-03,  ..., 0.0000e+00, 0.0000e+00,\n",
       "          0.0000e+00],\n",
       "         [8.8511e-05, 2.9591e-03, 5.8103e-03,  ..., 0.0000e+00, 0.0000e+00,\n",
       "          0.0000e+00],\n",
       "         ...,\n",
       "         [1.0713e-03, 1.6172e-03, 2.6454e-03,  ..., 0.0000e+00, 0.0000e+00,\n",
       "          0.0000e+00],\n",
       "         [5.0058e-05, 2.5230e-03, 4.6929e-03,  ..., 0.0000e+00, 0.0000e+00,\n",
       "          0.0000e+00],\n",
       "         [3.0563e-03, 8.6677e-03, 1.0614e-02,  ..., 0.0000e+00, 0.0000e+00,\n",
       "          0.0000e+00]], dtype=torch.float64),\n",
       " 'target': tensor([0, 0, 1, 1, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 0, 0, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 1, 1, 0, 1, 1, 0, 1, 0, 0, 0,\n",
       "         0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0, 0, 1, 0, 0, 1, 1, 1,\n",
       "         0, 1, 1, 1, 1, 0, 0, 1, 1, 0, 1, 0, 1, 1, 1, 1, 0, 1, 1, 0, 1, 0, 1, 0,\n",
       "         0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 1, 0, 1, 1, 1, 1, 1, 1,\n",
       "         0, 1, 1, 0, 0, 0, 1, 0, 0, 1, 1, 0, 1, 1, 0, 0, 0, 1, 0, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 0, 1, 0, 0, 0, 0, 0, 1, 0, 1, 0, 1, 0, 1, 0, 0, 0, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 0, 1, 1, 0, 0, 1, 0, 0, 1, 1, 1, 0, 1, 0, 0, 1, 1, 1, 0, 0,\n",
       "         1, 1, 0, 1, 1, 1, 0, 0, 1, 1, 0, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 0, 1, 0,\n",
       "         0, 0, 0, 1, 0, 1, 1, 0, 1, 0, 1, 1, 0, 1, 1, 1, 1, 0, 1, 0, 1, 0, 1, 0,\n",
       "         1, 0, 0, 1, 1, 1, 0, 1, 0, 0, 0, 0, 1, 1, 1, 0, 1, 0, 0, 1, 0, 1, 0, 0,\n",
       "         1, 1, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 1, 1,\n",
       "         0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 1, 0, 0, 1, 1, 1, 0, 0, 0, 0, 1, 1,\n",
       "         1, 1, 1, 1, 0, 1, 1, 0, 1, 0, 0, 0, 1, 0, 1, 0, 1, 0, 0, 0, 1, 1, 0, 0,\n",
       "         1, 1, 0, 1, 0, 0, 0, 1, 0, 0, 1, 0, 1, 0, 0, 1, 0, 0, 0, 0, 1, 0, 1, 0,\n",
       "         0, 1, 0, 1, 0, 1, 1, 1, 1, 1, 0, 1, 1, 0, 0, 1, 0, 1, 1, 1, 0, 0, 0, 1,\n",
       "         1, 1, 1, 1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1,\n",
       "         1, 0, 1, 1, 0, 0, 1, 1, 1, 1, 0, 0, 0, 1, 1, 0, 1, 1, 0, 1, 1, 1, 1, 1,\n",
       "         0, 1, 0, 1, 1, 0, 1, 0, 1, 1, 0, 1, 0, 1, 0, 0, 1, 0, 0, 1, 1, 0, 1, 0,\n",
       "         0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 1, 0, 1, 0, 1, 0, 1, 1, 0, 1, 1, 0, 0, 1,\n",
       "         1, 1, 1, 0, 0, 0, 0, 1, 1, 0, 0, 0, 1, 1, 1, 1, 1, 0, 1, 0, 1, 1, 0, 1,\n",
       "         1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0,\n",
       "         0, 1, 0, 1, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 1, 1, 1, 1,\n",
       "         0, 1, 1, 1, 0, 0, 1, 1, 0, 1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1, 1, 1, 1, 1,\n",
       "         0, 1, 0, 1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0,\n",
       "         1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 1, 1, 1, 1, 0, 0, 1, 1, 0, 0, 0,\n",
       "         1, 0, 0, 0, 0, 0, 1, 1, 1, 1, 0, 1, 0, 1, 1, 0, 1, 0, 1, 0, 0, 1, 0, 1,\n",
       "         1, 0, 1, 0, 1, 0, 1, 0, 1, 1, 0, 1, 0, 1, 1, 0, 0, 0, 0, 1, 1, 1, 1, 1,\n",
       "         0, 0, 0, 1, 0, 0, 0, 1, 0, 1, 1, 1, 1, 0, 0, 1, 1, 0, 1, 0, 0, 1, 1, 1,\n",
       "         0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 1, 0, 0, 1, 1, 0,\n",
       "         0, 0, 0, 1, 0, 1, 1, 1, 0, 1, 1, 1, 0, 0, 1, 1, 1, 0, 0, 0, 0, 1, 0, 0,\n",
       "         1, 1, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 0, 0, 1, 1, 0, 1, 1, 1,\n",
       "         0, 0, 0, 1, 1, 0, 1, 1, 0, 1, 0, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 0, 0,\n",
       "         0, 1, 0, 1, 1, 0, 0, 1, 0, 0, 1, 0, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1,\n",
       "         0, 1, 1, 0, 1, 1, 0, 0, 1, 0, 0, 1, 1, 1, 1, 0, 1, 0, 0, 1, 1, 0, 0, 1,\n",
       "         1, 1, 1, 0, 1, 0, 0, 1, 0, 0, 0, 1, 1, 0, 1, 1, 0, 0, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 1, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1,\n",
       "         1, 1, 1, 1, 0, 0, 1, 0, 1, 0, 1, 1, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 0,\n",
       "         1, 1, 0, 1, 0, 1, 0, 1, 0, 0, 1, 1, 1, 1, 1, 1, 0, 1, 1, 0, 1, 0, 1, 1,\n",
       "         1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 0, 1, 0, 0, 0, 1, 0, 0, 1, 0,\n",
       "         0, 1, 0, 0, 0, 0, 0, 1, 1, 1, 0, 1, 1, 1, 0, 1])}"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "inference_dl = torch.utils.data.DataLoader(\n",
    "    dataset=iterable_inference_dataset,\n",
    "    collate_fn=collate_feature_dict,\n",
    "    shuffle=False,\n",
    "    batch_size=1000,\n",
    "    num_workers=12,\n",
    ")\n",
    "next(iter(inference_dl)).payload"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "dda97025",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7d0b7d0a87ae40e397755dcf85d5ff41",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1.73 s, sys: 2.52 s, total: 4.25 s\n",
      "Wall time: 13.3 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "for rec in tqdm(inference_dl):\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62263aca",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "819d70d1",
   "metadata": {},
   "source": [
    "## Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "85cfc731",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ptls.nn import TrxEncoder, RnnSeqEncoder\n",
    "from ptls.frames.inference_module import InferenceModule"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "6bde712c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# random model in demo\n",
    "# use trained model for real usage\n",
    "seq_encoder = RnnSeqEncoder(\n",
    "    trx_encoder=TrxEncoder(\n",
    "        embeddings={\n",
    "            'trans_date': {'in': 800, 'out': 32}, \n",
    "            'small_group': {'in': 200, 'out': 32},\n",
    "        },\n",
    "        numeric_values={'amount_rur': 'log'},\n",
    "    ),\n",
    "    hidden_size=256,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "18ccbe2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = InferenceModule(model=seq_encoder, pandas_output=True, model_out_name='emb')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "791e6562",
   "metadata": {},
   "source": [
    "## Iterable inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "7c52b52f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/kireev/pipenv_envs/pytorch-lifestream-KxQJF1XF/lib/python3.8/site-packages/pytorch_lightning/loops/utilities.py:91: PossibleUserWarning: `max_epochs` was not set. Setting it to 1000 epochs. To train without an epoch limit, set `max_epochs=-1`.\n",
      "  rank_zero_warn(\n",
      "GPU available: True, used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1]\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "301fdff9331c40a2807b8c5158232070",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Predicting: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/kireev/pipenv_envs/pytorch-lifestream-KxQJF1XF/lib/python3.8/site-packages/pytorch_lightning/loops/epoch/prediction_epoch_loop.py:175: UserWarning: Lightning couldn't infer the indices fetched for your dataloader.\n",
      "  warning_cache.warn(\"Lightning couldn't infer the indices fetched for your dataloader.\")\n"
     ]
    }
   ],
   "source": [
    "predict = pl.Trainer(gpus=1).predict(model, inference_dl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "231ca698",
   "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>user_id</th>\n",
       "      <th>target</th>\n",
       "      <th>emb_0000</th>\n",
       "      <th>emb_0001</th>\n",
       "      <th>emb_0002</th>\n",
       "      <th>emb_0003</th>\n",
       "      <th>emb_0004</th>\n",
       "      <th>emb_0005</th>\n",
       "      <th>emb_0006</th>\n",
       "      <th>emb_0007</th>\n",
       "      <th>...</th>\n",
       "      <th>emb_0246</th>\n",
       "      <th>emb_0247</th>\n",
       "      <th>emb_0248</th>\n",
       "      <th>emb_0249</th>\n",
       "      <th>emb_0250</th>\n",
       "      <th>emb_0251</th>\n",
       "      <th>emb_0252</th>\n",
       "      <th>emb_0253</th>\n",
       "      <th>emb_0254</th>\n",
       "      <th>emb_0255</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>u_100002</td>\n",
       "      <td>0</td>\n",
       "      <td>0.028287</td>\n",
       "      <td>-0.165995</td>\n",
       "      <td>0.126620</td>\n",
       "      <td>-0.032016</td>\n",
       "      <td>0.016137</td>\n",
       "      <td>0.039828</td>\n",
       "      <td>0.159993</td>\n",
       "      <td>-0.224086</td>\n",
       "      <td>...</td>\n",
       "      <td>0.077100</td>\n",
       "      <td>0.044393</td>\n",
       "      <td>0.242713</td>\n",
       "      <td>-0.043127</td>\n",
       "      <td>0.023614</td>\n",
       "      <td>-0.002254</td>\n",
       "      <td>-0.142776</td>\n",
       "      <td>-0.221038</td>\n",
       "      <td>0.057583</td>\n",
       "      <td>-0.007864</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>u_100093</td>\n",
       "      <td>0</td>\n",
       "      <td>0.236783</td>\n",
       "      <td>0.364345</td>\n",
       "      <td>-0.108432</td>\n",
       "      <td>0.070790</td>\n",
       "      <td>0.099409</td>\n",
       "      <td>-0.119790</td>\n",
       "      <td>0.510778</td>\n",
       "      <td>-0.231560</td>\n",
       "      <td>...</td>\n",
       "      <td>-0.031448</td>\n",
       "      <td>0.040263</td>\n",
       "      <td>0.045545</td>\n",
       "      <td>-0.006918</td>\n",
       "      <td>0.129049</td>\n",
       "      <td>0.071746</td>\n",
       "      <td>-0.210788</td>\n",
       "      <td>0.288435</td>\n",
       "      <td>-0.107366</td>\n",
       "      <td>-0.176545</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>u_100292</td>\n",
       "      <td>1</td>\n",
       "      <td>0.060974</td>\n",
       "      <td>0.015578</td>\n",
       "      <td>0.074314</td>\n",
       "      <td>-0.015297</td>\n",
       "      <td>-0.072065</td>\n",
       "      <td>-0.228197</td>\n",
       "      <td>0.084782</td>\n",
       "      <td>-0.316223</td>\n",
       "      <td>...</td>\n",
       "      <td>0.065913</td>\n",
       "      <td>-0.209312</td>\n",
       "      <td>0.186426</td>\n",
       "      <td>-0.014417</td>\n",
       "      <td>0.184326</td>\n",
       "      <td>-0.065463</td>\n",
       "      <td>-0.388418</td>\n",
       "      <td>0.221885</td>\n",
       "      <td>0.136650</td>\n",
       "      <td>-0.183059</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>u_100302</td>\n",
       "      <td>1</td>\n",
       "      <td>0.166716</td>\n",
       "      <td>-0.381586</td>\n",
       "      <td>-0.064041</td>\n",
       "      <td>-0.077300</td>\n",
       "      <td>0.254902</td>\n",
       "      <td>0.121023</td>\n",
       "      <td>0.516134</td>\n",
       "      <td>0.047728</td>\n",
       "      <td>...</td>\n",
       "      <td>-0.099801</td>\n",
       "      <td>-0.444753</td>\n",
       "      <td>-0.323648</td>\n",
       "      <td>0.066342</td>\n",
       "      <td>0.204678</td>\n",
       "      <td>-0.088744</td>\n",
       "      <td>-0.335036</td>\n",
       "      <td>0.024014</td>\n",
       "      <td>0.137864</td>\n",
       "      <td>-0.232802</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>u_100467</td>\n",
       "      <td>0</td>\n",
       "      <td>0.102175</td>\n",
       "      <td>0.026828</td>\n",
       "      <td>-0.128542</td>\n",
       "      <td>-0.134966</td>\n",
       "      <td>0.112983</td>\n",
       "      <td>0.093400</td>\n",
       "      <td>0.189535</td>\n",
       "      <td>0.082395</td>\n",
       "      <td>...</td>\n",
       "      <td>-0.100227</td>\n",
       "      <td>-0.016402</td>\n",
       "      <td>-0.132673</td>\n",
       "      <td>0.066391</td>\n",
       "      <td>-0.018519</td>\n",
       "      <td>-0.252219</td>\n",
       "      <td>-0.096244</td>\n",
       "      <td>-0.078941</td>\n",
       "      <td>0.241028</td>\n",
       "      <td>0.102450</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>995</th>\n",
       "      <td>u_27656</td>\n",
       "      <td>1</td>\n",
       "      <td>0.013051</td>\n",
       "      <td>0.091127</td>\n",
       "      <td>0.090885</td>\n",
       "      <td>-0.243966</td>\n",
       "      <td>0.262145</td>\n",
       "      <td>0.183707</td>\n",
       "      <td>0.482889</td>\n",
       "      <td>-0.085619</td>\n",
       "      <td>...</td>\n",
       "      <td>0.017391</td>\n",
       "      <td>-0.012522</td>\n",
       "      <td>-0.188911</td>\n",
       "      <td>0.075924</td>\n",
       "      <td>0.288637</td>\n",
       "      <td>0.184379</td>\n",
       "      <td>0.000888</td>\n",
       "      <td>0.193466</td>\n",
       "      <td>0.131888</td>\n",
       "      <td>-0.288630</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>996</th>\n",
       "      <td>u_276935</td>\n",
       "      <td>1</td>\n",
       "      <td>0.167983</td>\n",
       "      <td>0.154561</td>\n",
       "      <td>-0.059162</td>\n",
       "      <td>-0.200115</td>\n",
       "      <td>-0.012947</td>\n",
       "      <td>-0.068747</td>\n",
       "      <td>0.398722</td>\n",
       "      <td>-0.073746</td>\n",
       "      <td>...</td>\n",
       "      <td>0.054446</td>\n",
       "      <td>-0.009374</td>\n",
       "      <td>0.225024</td>\n",
       "      <td>-0.086913</td>\n",
       "      <td>0.225870</td>\n",
       "      <td>0.094627</td>\n",
       "      <td>0.005300</td>\n",
       "      <td>0.057958</td>\n",
       "      <td>0.311267</td>\n",
       "      <td>-0.215860</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>997</th>\n",
       "      <td>u_277096</td>\n",
       "      <td>1</td>\n",
       "      <td>0.102818</td>\n",
       "      <td>-0.004892</td>\n",
       "      <td>0.256763</td>\n",
       "      <td>-0.242788</td>\n",
       "      <td>0.213398</td>\n",
       "      <td>0.252951</td>\n",
       "      <td>0.247898</td>\n",
       "      <td>-0.035680</td>\n",
       "      <td>...</td>\n",
       "      <td>0.056761</td>\n",
       "      <td>-0.405258</td>\n",
       "      <td>-0.005567</td>\n",
       "      <td>0.203767</td>\n",
       "      <td>0.260439</td>\n",
       "      <td>0.091790</td>\n",
       "      <td>0.115536</td>\n",
       "      <td>0.036998</td>\n",
       "      <td>0.060348</td>\n",
       "      <td>-0.020778</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>998</th>\n",
       "      <td>u_277131</td>\n",
       "      <td>0</td>\n",
       "      <td>0.234514</td>\n",
       "      <td>0.465364</td>\n",
       "      <td>0.182645</td>\n",
       "      <td>-0.267335</td>\n",
       "      <td>0.109032</td>\n",
       "      <td>0.169204</td>\n",
       "      <td>0.133215</td>\n",
       "      <td>-0.113944</td>\n",
       "      <td>...</td>\n",
       "      <td>0.107012</td>\n",
       "      <td>0.138886</td>\n",
       "      <td>0.146164</td>\n",
       "      <td>0.094512</td>\n",
       "      <td>0.062334</td>\n",
       "      <td>-0.091048</td>\n",
       "      <td>-0.295387</td>\n",
       "      <td>0.091000</td>\n",
       "      <td>0.437284</td>\n",
       "      <td>-0.400930</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>999</th>\n",
       "      <td>u_277166</td>\n",
       "      <td>1</td>\n",
       "      <td>0.263952</td>\n",
       "      <td>0.048121</td>\n",
       "      <td>-0.087273</td>\n",
       "      <td>-0.159209</td>\n",
       "      <td>0.262990</td>\n",
       "      <td>0.132703</td>\n",
       "      <td>0.436050</td>\n",
       "      <td>0.360127</td>\n",
       "      <td>...</td>\n",
       "      <td>-0.200030</td>\n",
       "      <td>-0.096696</td>\n",
       "      <td>-0.114389</td>\n",
       "      <td>-0.039275</td>\n",
       "      <td>-0.089717</td>\n",
       "      <td>-0.051737</td>\n",
       "      <td>-0.109379</td>\n",
       "      <td>0.234441</td>\n",
       "      <td>0.060527</td>\n",
       "      <td>-0.122428</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>1000 rows × 258 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "      user_id  target  emb_0000  emb_0001  emb_0002  emb_0003  emb_0004  \\\n",
       "0    u_100002       0  0.028287 -0.165995  0.126620 -0.032016  0.016137   \n",
       "1    u_100093       0  0.236783  0.364345 -0.108432  0.070790  0.099409   \n",
       "2    u_100292       1  0.060974  0.015578  0.074314 -0.015297 -0.072065   \n",
       "3    u_100302       1  0.166716 -0.381586 -0.064041 -0.077300  0.254902   \n",
       "4    u_100467       0  0.102175  0.026828 -0.128542 -0.134966  0.112983   \n",
       "..        ...     ...       ...       ...       ...       ...       ...   \n",
       "995   u_27656       1  0.013051  0.091127  0.090885 -0.243966  0.262145   \n",
       "996  u_276935       1  0.167983  0.154561 -0.059162 -0.200115 -0.012947   \n",
       "997  u_277096       1  0.102818 -0.004892  0.256763 -0.242788  0.213398   \n",
       "998  u_277131       0  0.234514  0.465364  0.182645 -0.267335  0.109032   \n",
       "999  u_277166       1  0.263952  0.048121 -0.087273 -0.159209  0.262990   \n",
       "\n",
       "     emb_0005  emb_0006  emb_0007  ...  emb_0246  emb_0247  emb_0248  \\\n",
       "0    0.039828  0.159993 -0.224086  ...  0.077100  0.044393  0.242713   \n",
       "1   -0.119790  0.510778 -0.231560  ... -0.031448  0.040263  0.045545   \n",
       "2   -0.228197  0.084782 -0.316223  ...  0.065913 -0.209312  0.186426   \n",
       "3    0.121023  0.516134  0.047728  ... -0.099801 -0.444753 -0.323648   \n",
       "4    0.093400  0.189535  0.082395  ... -0.100227 -0.016402 -0.132673   \n",
       "..        ...       ...       ...  ...       ...       ...       ...   \n",
       "995  0.183707  0.482889 -0.085619  ...  0.017391 -0.012522 -0.188911   \n",
       "996 -0.068747  0.398722 -0.073746  ...  0.054446 -0.009374  0.225024   \n",
       "997  0.252951  0.247898 -0.035680  ...  0.056761 -0.405258 -0.005567   \n",
       "998  0.169204  0.133215 -0.113944  ...  0.107012  0.138886  0.146164   \n",
       "999  0.132703  0.436050  0.360127  ... -0.200030 -0.096696 -0.114389   \n",
       "\n",
       "     emb_0249  emb_0250  emb_0251  emb_0252  emb_0253  emb_0254  emb_0255  \n",
       "0   -0.043127  0.023614 -0.002254 -0.142776 -0.221038  0.057583 -0.007864  \n",
       "1   -0.006918  0.129049  0.071746 -0.210788  0.288435 -0.107366 -0.176545  \n",
       "2   -0.014417  0.184326 -0.065463 -0.388418  0.221885  0.136650 -0.183059  \n",
       "3    0.066342  0.204678 -0.088744 -0.335036  0.024014  0.137864 -0.232802  \n",
       "4    0.066391 -0.018519 -0.252219 -0.096244 -0.078941  0.241028  0.102450  \n",
       "..        ...       ...       ...       ...       ...       ...       ...  \n",
       "995  0.075924  0.288637  0.184379  0.000888  0.193466  0.131888 -0.288630  \n",
       "996 -0.086913  0.225870  0.094627  0.005300  0.057958  0.311267 -0.215860  \n",
       "997  0.203767  0.260439  0.091790  0.115536  0.036998  0.060348 -0.020778  \n",
       "998  0.094512  0.062334 -0.091048 -0.295387  0.091000  0.437284 -0.400930  \n",
       "999 -0.039275 -0.089717 -0.051737 -0.109379  0.234441  0.060527 -0.122428  \n",
       "\n",
       "[1000 rows x 258 columns]"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "predict[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "3e1457a9",
   "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>user_id</th>\n",
       "      <th>target</th>\n",
       "      <th>emb_0000</th>\n",
       "      <th>emb_0001</th>\n",
       "      <th>emb_0002</th>\n",
       "      <th>emb_0003</th>\n",
       "      <th>emb_0004</th>\n",
       "      <th>emb_0005</th>\n",
       "      <th>emb_0006</th>\n",
       "      <th>emb_0007</th>\n",
       "      <th>...</th>\n",
       "      <th>emb_0246</th>\n",
       "      <th>emb_0247</th>\n",
       "      <th>emb_0248</th>\n",
       "      <th>emb_0249</th>\n",
       "      <th>emb_0250</th>\n",
       "      <th>emb_0251</th>\n",
       "      <th>emb_0252</th>\n",
       "      <th>emb_0253</th>\n",
       "      <th>emb_0254</th>\n",
       "      <th>emb_0255</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>u_100002</td>\n",
       "      <td>0</td>\n",
       "      <td>0.028287</td>\n",
       "      <td>-0.165995</td>\n",
       "      <td>0.126620</td>\n",
       "      <td>-0.032016</td>\n",
       "      <td>0.016137</td>\n",
       "      <td>0.039828</td>\n",
       "      <td>0.159993</td>\n",
       "      <td>-0.224086</td>\n",
       "      <td>...</td>\n",
       "      <td>0.077100</td>\n",
       "      <td>0.044393</td>\n",
       "      <td>0.242713</td>\n",
       "      <td>-0.043127</td>\n",
       "      <td>0.023614</td>\n",
       "      <td>-0.002254</td>\n",
       "      <td>-0.142776</td>\n",
       "      <td>-0.221038</td>\n",
       "      <td>0.057583</td>\n",
       "      <td>-0.007864</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>u_100093</td>\n",
       "      <td>0</td>\n",
       "      <td>0.236783</td>\n",
       "      <td>0.364345</td>\n",
       "      <td>-0.108432</td>\n",
       "      <td>0.070790</td>\n",
       "      <td>0.099409</td>\n",
       "      <td>-0.119790</td>\n",
       "      <td>0.510778</td>\n",
       "      <td>-0.231560</td>\n",
       "      <td>...</td>\n",
       "      <td>-0.031448</td>\n",
       "      <td>0.040263</td>\n",
       "      <td>0.045545</td>\n",
       "      <td>-0.006918</td>\n",
       "      <td>0.129049</td>\n",
       "      <td>0.071746</td>\n",
       "      <td>-0.210788</td>\n",
       "      <td>0.288435</td>\n",
       "      <td>-0.107366</td>\n",
       "      <td>-0.176545</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>u_100292</td>\n",
       "      <td>1</td>\n",
       "      <td>0.060974</td>\n",
       "      <td>0.015578</td>\n",
       "      <td>0.074314</td>\n",
       "      <td>-0.015297</td>\n",
       "      <td>-0.072065</td>\n",
       "      <td>-0.228197</td>\n",
       "      <td>0.084782</td>\n",
       "      <td>-0.316223</td>\n",
       "      <td>...</td>\n",
       "      <td>0.065913</td>\n",
       "      <td>-0.209312</td>\n",
       "      <td>0.186426</td>\n",
       "      <td>-0.014417</td>\n",
       "      <td>0.184326</td>\n",
       "      <td>-0.065463</td>\n",
       "      <td>-0.388418</td>\n",
       "      <td>0.221885</td>\n",
       "      <td>0.136650</td>\n",
       "      <td>-0.183059</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>u_100302</td>\n",
       "      <td>1</td>\n",
       "      <td>0.166716</td>\n",
       "      <td>-0.381586</td>\n",
       "      <td>-0.064041</td>\n",
       "      <td>-0.077300</td>\n",
       "      <td>0.254902</td>\n",
       "      <td>0.121023</td>\n",
       "      <td>0.516134</td>\n",
       "      <td>0.047728</td>\n",
       "      <td>...</td>\n",
       "      <td>-0.099801</td>\n",
       "      <td>-0.444753</td>\n",
       "      <td>-0.323648</td>\n",
       "      <td>0.066342</td>\n",
       "      <td>0.204678</td>\n",
       "      <td>-0.088744</td>\n",
       "      <td>-0.335036</td>\n",
       "      <td>0.024014</td>\n",
       "      <td>0.137864</td>\n",
       "      <td>-0.232802</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>u_100467</td>\n",
       "      <td>0</td>\n",
       "      <td>0.102175</td>\n",
       "      <td>0.026828</td>\n",
       "      <td>-0.128542</td>\n",
       "      <td>-0.134966</td>\n",
       "      <td>0.112983</td>\n",
       "      <td>0.093400</td>\n",
       "      <td>0.189535</td>\n",
       "      <td>0.082395</td>\n",
       "      <td>...</td>\n",
       "      <td>-0.100227</td>\n",
       "      <td>-0.016402</td>\n",
       "      <td>-0.132673</td>\n",
       "      <td>0.066391</td>\n",
       "      <td>-0.018519</td>\n",
       "      <td>-0.252219</td>\n",
       "      <td>-0.096244</td>\n",
       "      <td>-0.078941</td>\n",
       "      <td>0.241028</td>\n",
       "      <td>0.102450</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5 rows × 258 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "    user_id  target  emb_0000  emb_0001  emb_0002  emb_0003  emb_0004  \\\n",
       "0  u_100002       0  0.028287 -0.165995  0.126620 -0.032016  0.016137   \n",
       "1  u_100093       0  0.236783  0.364345 -0.108432  0.070790  0.099409   \n",
       "2  u_100292       1  0.060974  0.015578  0.074314 -0.015297 -0.072065   \n",
       "3  u_100302       1  0.166716 -0.381586 -0.064041 -0.077300  0.254902   \n",
       "4  u_100467       0  0.102175  0.026828 -0.128542 -0.134966  0.112983   \n",
       "\n",
       "   emb_0005  emb_0006  emb_0007  ...  emb_0246  emb_0247  emb_0248  emb_0249  \\\n",
       "0  0.039828  0.159993 -0.224086  ...  0.077100  0.044393  0.242713 -0.043127   \n",
       "1 -0.119790  0.510778 -0.231560  ... -0.031448  0.040263  0.045545 -0.006918   \n",
       "2 -0.228197  0.084782 -0.316223  ...  0.065913 -0.209312  0.186426 -0.014417   \n",
       "3  0.121023  0.516134  0.047728  ... -0.099801 -0.444753 -0.323648  0.066342   \n",
       "4  0.093400  0.189535  0.082395  ... -0.100227 -0.016402 -0.132673  0.066391   \n",
       "\n",
       "   emb_0250  emb_0251  emb_0252  emb_0253  emb_0254  emb_0255  \n",
       "0  0.023614 -0.002254 -0.142776 -0.221038  0.057583 -0.007864  \n",
       "1  0.129049  0.071746 -0.210788  0.288435 -0.107366 -0.176545  \n",
       "2  0.184326 -0.065463 -0.388418  0.221885  0.136650 -0.183059  \n",
       "3  0.204678 -0.088744 -0.335036  0.024014  0.137864 -0.232802  \n",
       "4 -0.018519 -0.252219 -0.096244 -0.078941  0.241028  0.102450  \n",
       "\n",
       "[5 rows x 258 columns]"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "full_predict = pd.concat(predict, axis=0)\n",
    "full_predict.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "dc5dd87d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1000000, 258)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "full_predict.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "35e0eea3",
   "metadata": {},
   "outputs": [],
   "source": [
    "full_predict.to_parquet('full_predict.parquet')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "57979aed",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.2G\tfull_predict.parquet\r\n"
     ]
    }
   ],
   "source": [
    "!du -sh full_predict.parquet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f0f878e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ptls",
   "language": "python",
   "name": "ptls"
  },
  "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}