Hyper-Scale Training with MinIO and TensorFlow - v2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Hyper-Scale Machine Learning with MinIO and TensorFlow"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "We are living in a transformative era defined by information and AI. Massive amounts of data are generated and collected every day to feed these voracious, state-of-the-art, AI/ML algorithms. The more data, the better the outcomes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "One of the frameworks that has emerged as the lead industry standards is [Google's TensorFlow](https://www.tensorflow.org/).  Highly versatile, one can get started quickly and write simple models with their [Keras](https://www.tensorflow.org/guide/keras?hl=en) framework. If you seek a more advanced approach TensorFlow also allows you to construct your own machine learning models using low level APIs. No matter what strategy you choose, TensorFlow will make sure that your algorithm gets optimized for whatever infrastructure you select for your algorithms - whether it's [CPU's](https://en.wikipedia.org/wiki/Central_processing_unit), [GPU's](https://en.wikipedia.org/wiki/Graphics_processing_unit) or [TPU's](https://en.wikipedia.org/wiki/Tensor_processing_unit)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "As datasets become too large to fit into memory or local disk, AI/ML pipelines now have the requirement to load data from an external data source. Take for example the [ImageNet](https://en.wikipedia.org/wiki/ImageNet) dataset with its `14 Million` Images with an estimated storage size of `1.31TB`. This dataset cannot be fit into memory nor on any machine local storage drive. These challenges are further complicated if your pipelines are running inside a stateless environment such a Kubernetes (which is increasingly the norm). "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "The emerging standard for this problem is to employ high performance object storage in the design of your AI/ML pipelines. MinIO is the leader in this space and has published a number of benchmarks that speak to its throughput capabilities. In this post, we will cover how to leverage MinIO for your TensorFlow projects."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## A Four Stage Hyper-Scale Data Pipeline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "To build a hyper-scale pipeline we will have each stage of the pipeline read from MinIO. In this example we are going to build four stages of a machine learning pipeline. This architecture will load the desired data on-demand from MinIO. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "First, we are going to preprocess our dataset and encode it in a format that TensorFlow can quickly digest. This format is the [tf.TFRecord](https://www.tensorflow.org/tutorials/load_data/tfrecord), which is a type of binary encoding for our data. We are taking this step because we do not want to waste time processing the data during the training as we are planning on loading each batch of training directly from MinIO as it's needed. If the data is pre-processed before we feed it into the model training we save a significant amount of time. Ideally, we create pre-processed chunks of data that group a good chunk of records - at least `100-200MB` in size."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "To speed up the data-loading and training stages we are going to leverage the excellent [tf.data](https://www.tensorflow.org/api_docs/python/tf/data) api. This API is designed to efficiently load data during the training/validation of our model. It prepares the next batch of data as the current one is being processed by the model. The advantage of this approach is that it ensures efficient utilization of expensive GPUs or TPUs which cannot sit idle due to slow loading data. MinIO does not encounter this problem - [it can saturate 100Gbps network with a few NVMe drives](https://min.io/resources/docs/MinIO-Throughput-Benchmarks-on-NVMe-SSD-32-Node.pdf) or also with [Hard Disk Drives](https://min.io/resources/docs/MinIO-Throughput-Benchmarks-on-HDD-24-Node.pdf) ensuring the pipeline is crunching data as fast as the hardware allows"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "During training we want to make sure we store the training checkpoints of our model as well as TensorBoard histograms. The checkpoints are useful in case the training gets interrupted and we want to resume the training or if we get more data and want to keep training our model with the new data and the TensorBoard histograms let us see how the training is going as it happens. TensorFlow supports writing both of these directly to MinIO."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "A quick side note. When the model is complete we will save it to MinIO as well - allowing us to serve it using [TensorFlow Serving](https://www.tensorflow.org/tfx/guide/serving)  - but that's a post for some other time."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "![title](pic1.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Building the Pipeline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "For our hyper-scale pipeline we are going to use a dataset that can easily fit into your local computer so you can follow along. The [Large Movie Review Dataset](https://ai.stanford.edu/~amaas/data/sentiment/) from Stanford is great since it has a large number of samples (25,000 for training and 25,000 for testing) so we are going to build a [sentiment analysis](https://en.wikipedia.org/wiki/Sentiment_analysis) model that will tell us whether a movie review is `positive` or `negative`. Keep in mind that each step can be applied to any larger dataset. The advantage of this dataset is that you can try on your own computer. Let's get started!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Download the dataset and upload it to MinIO using [MinIO Client](https://docs.min.io/docs/minio-client-quickstart-guide.html)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "```bash\n",
    "curl -O http://ai.stanford.edu/~amaas/data/sentiment/aclImdb_v1.tar.gz\n",
    "mc mb myminio/datasets\n",
    "mc cp aclImdb_v1.tar.gz myminio/datasets/\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Let's start by declaring some configurations for our pipeline,  such as `batch size`, location of our dataset and a fixed `random seed` so we can run this pipeline again and again and get the same results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import tensorflow_io\n",
    "from minio import Minio\n",
    "from minio.error import S3Error\n",
    "import tarfile\n",
    "import os\n",
    "import random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.9.1\n"
     ]
    }
   ],
   "source": [
    "print(tf.__version__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "random_seed = 44\n",
    "batch_size = 128\n",
    "datasets_bucket = \"datasets\"\n",
    "preprocessed_data_folder = \"preprocessed-data\"\n",
    "tf_record_file_size = 500\n",
    "# How to access MinIO\n",
    "# minio_address = \"localhost:9000\"\n",
    "# TODO: REMOVE ME, THIS IS MY LOCAL ENV :-)\n",
    "minio_address = \"192.168.86.197:9000\"\n",
    "minio_access_key = \"minioadmin\"\n",
    "minio_secret_key = \"minioadmin\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "We are going to download our dataset from MinIO using [minio-py](https://github.com/minio/minio-py)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "minioClient = Minio(minio_address,\n",
    "                  access_key=minio_access_key,\n",
    "                  secret_key=minio_secret_key,\n",
    "                  secure=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "try:\n",
    "       minioClient.fget_object(\n",
    "           datasets_bucket,\n",
    "           \"aclImdb_v1.tar.gz\",\n",
    "           \"/tmp/dataset.tar.gz\")\n",
    "except S3Error as err:\n",
    "       print(err)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Now let's uncompress the dataset to a temporary folder (`/tmp/dataset`) to preprocess our data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "extract_folder = f\"/tmp/{datasets_bucket}/\"\n",
    "\n",
    "with tarfile.open(\"/tmp/dataset.tar.gz\", \"r:gz\") as tar:\n",
    "    tar.extractall(path=extract_folder)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Pre-Processing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Due to the structure of the dataset we are going to read from four folders, initially `test` and `train` which hold `25,000` examples each, then, in each of those folders we have `12,500` of each label pos for positive comments and neg for negative comments. From these four folders, we are going to store all samples into two variables, `train` and `test`. If we were preprocessing a dataset that couldn't fit in the local machine we could simply load segments of the object, one at a time and process them as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "train = []\n",
    "test = []\n",
    "\n",
    "dirs_to_read = [\n",
    "    \"aclImdb/train/pos\",\n",
    "    \"aclImdb/train/neg\",\n",
    "    \"aclImdb/test/pos\",\n",
    "    \"aclImdb/test/neg\",\n",
    "]\n",
    "\n",
    "for dir_name in dirs_to_read:\n",
    "    parts = dir_name.split(\"/\")\n",
    "    dataset = parts[1]\n",
    "    label = parts[2]\n",
    "    for filename in os.listdir(os.path.join(extract_folder,dir_name)):\n",
    "        with open(os.path.join(extract_folder,dir_name,filename),\"r\") as f:\n",
    "            content = f.read()\n",
    "            if dataset == \"train\":\n",
    "                train.append({\n",
    "                    \"text\":content,\n",
    "                    \"label\":label\n",
    "                })\n",
    "            elif dataset == \"test\":\n",
    "                test.append({\n",
    "                    \"text\":content,\n",
    "                    \"label\":label\n",
    "                })"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "We will then shuffle the dataset so we don't introduce bias into the training by providing 12,500 consecutive positive examples followed by 12,500 consecutive negative examples. Our model would have a hard time generalizing that. By shuffling the data the model will get to see and learn from both positive and negative examples at the same time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "random.Random(random_seed).shuffle(train)\n",
    "random.Random(random_seed).shuffle(test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Since we are dealing with text we need to turn the text to a vector representation that accurately depicts the meanings of the sentence. If we were dealing with images we would resize the images and turn them into vector representations having each pixel be a value of the resized image. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "For text, however, we have a bigger challenge since a word doesn't really have a numerical representation. This is where [embeddings](https://en.wikipedia.org/wiki/Embedding) are useful. An embedding is a vector representation of some text, in this case we are going to represent the whole review as a single vector of 512 dimensions. Instead of doing the pre-processing of text manually (tokenizing, building vocabulary and training an embeddings layer) we are going to leverage an existing model called [USE (Universal Sentence Encoder)](https://arxiv.org/abs/1803.11175) to encode sentences into vectors so we can continue with our example. This is one of the wonders of deep learning, the ability to re-use different models alongside yours. Here we use TensorFlow Hub and we are going to load the latest `USE` model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "outputs": [],
   "source": [
    "import requests\n",
    "\n",
    "# Download the model\n",
    "model_url = \"https://tfhub.dev/google/universal-sentence-encoder/4?tf-hub-format=compressed\"\n",
    "response = requests.get(model_url)\n",
    "open(\"/tmp/universal-sentence-encoder_4.tar.gz\", \"wb\").write(response.content)\n",
    "\n",
    "se_model_prefix = \"universal-sentence-encoder/4\"\n",
    "extract_folder = f\"/tmp/{se_model_prefix}/\"\n",
    "\n",
    "with tarfile.open(\"/tmp/universal-sentence-encoder_4.tar.gz\", \"r:gz\") as tar:\n",
    "    tar.extractall(path=extract_folder)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('/tmp/universal-sentence-encoder/4/', ['assets', 'variables'], ['saved_model.pb'])\n",
      " - saved_model.pb: saved_model.pb\n",
      "('/tmp/universal-sentence-encoder/4/assets', [], [])\n",
      "('/tmp/universal-sentence-encoder/4/variables', [], ['variables.index', 'variables.data-00000-of-00001'])\n",
      "variables - variables.index: variables/variables.index\n",
      "variables - variables.data-00000-of-00001: variables/variables.data-00000-of-00001\n"
     ]
    }
   ],
   "source": [
    "\n",
    "dirs = os.walk(extract_folder)\n",
    "for f in dirs:\n",
    "    print(f)\n",
    "    for file in f[2]:\n",
    "\n",
    "        target = f[0].replace(extract_folder,\"\")\n",
    "        target_prefix = os.path.join(target,file)\n",
    "        print(f\"{target} - {file}: {target_prefix}\")\n",
    "\n",
    "        try:\n",
    "            minioClient.fput_object(\n",
    "                    datasets_bucket,\n",
    "                    f\"{se_model_prefix}/{target_prefix}\",\n",
    "                    f\"{os.path.join(f[0],file)}\")\n",
    "        except S3Error as err:\n",
    "            print(err)\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "outputs": [],
   "source": [
    "os.environ[\"AWS_ACCESS_KEY_ID\"] = minio_access_key\n",
    "os.environ[\"AWS_SECRET_ACCESS_KEY\"] = minio_secret_key\n",
    "os.environ[\"AWS_REGION\"] = \"us-east-1\"\n",
    "os.environ[\"S3_ENDPOINT\"] = minio_address\n",
    "os.environ[\"S3_USE_HTTPS\"] = \"0\"\n",
    "os.environ[\"S3_VERIFY_SSL\"] = \"0\""
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [
    {
     "ename": "OSError",
     "evalue": "SavedModel file does not exist at: s3://datasets/universal-sentence-encoder/4//{saved_model.pbtxt|saved_model.pb}",
     "output_type": "error",
     "traceback": [
      "\u001B[0;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[0;31mOSError\u001B[0m                                   Traceback (most recent call last)",
      "Input \u001B[0;32mIn [10]\u001B[0m, in \u001B[0;36m<cell line: 2>\u001B[0;34m()\u001B[0m\n\u001B[1;32m      1\u001B[0m \u001B[38;5;66;03m# embed = tf.compat.v1.saved_model.load_v2(f\"s3://datasets/{se_model_prefix}\")\u001B[39;00m\n\u001B[0;32m----> 2\u001B[0m embed \u001B[38;5;241m=\u001B[39m \u001B[43mtf\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mcompat\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mv1\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43msaved_model\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mload_v2\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;124;43mf\u001B[39;49m\u001B[38;5;124;43m\"\u001B[39;49m\u001B[38;5;124;43ms3://datasets/universal-sentence-encoder/4/\u001B[39;49m\u001B[38;5;124;43m\"\u001B[39;49m\u001B[43m)\u001B[49m\n",
      "File \u001B[0;32m~/anaconda3/envs/mlpipeline/lib/python3.9/site-packages/tensorflow/python/saved_model/load.py:782\u001B[0m, in \u001B[0;36mload\u001B[0;34m(export_dir, tags, options)\u001B[0m\n\u001B[1;32m    780\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;28misinstance\u001B[39m(export_dir, os\u001B[38;5;241m.\u001B[39mPathLike):\n\u001B[1;32m    781\u001B[0m   export_dir \u001B[38;5;241m=\u001B[39m os\u001B[38;5;241m.\u001B[39mfspath(export_dir)\n\u001B[0;32m--> 782\u001B[0m result \u001B[38;5;241m=\u001B[39m \u001B[43mload_partial\u001B[49m\u001B[43m(\u001B[49m\u001B[43mexport_dir\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;28;43;01mNone\u001B[39;49;00m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mtags\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43moptions\u001B[49m\u001B[43m)\u001B[49m[\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mroot\u001B[39m\u001B[38;5;124m\"\u001B[39m]\n\u001B[1;32m    783\u001B[0m \u001B[38;5;28;01mreturn\u001B[39;00m result\n",
      "File \u001B[0;32m~/anaconda3/envs/mlpipeline/lib/python3.9/site-packages/tensorflow/python/saved_model/load.py:887\u001B[0m, in \u001B[0;36mload_partial\u001B[0;34m(export_dir, filters, tags, options)\u001B[0m\n\u001B[1;32m    882\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m tags \u001B[38;5;129;01mis\u001B[39;00m \u001B[38;5;129;01mnot\u001B[39;00m \u001B[38;5;28;01mNone\u001B[39;00m \u001B[38;5;129;01mand\u001B[39;00m \u001B[38;5;129;01mnot\u001B[39;00m \u001B[38;5;28misinstance\u001B[39m(tags, \u001B[38;5;28mset\u001B[39m):\n\u001B[1;32m    883\u001B[0m   \u001B[38;5;66;03m# Supports e.g. tags=SERVING and tags=[SERVING]. Sets aren't considered\u001B[39;00m\n\u001B[1;32m    884\u001B[0m   \u001B[38;5;66;03m# sequences for nest.flatten, so we put those through as-is.\u001B[39;00m\n\u001B[1;32m    885\u001B[0m   tags \u001B[38;5;241m=\u001B[39m nest\u001B[38;5;241m.\u001B[39mflatten(tags)\n\u001B[1;32m    886\u001B[0m saved_model_proto, debug_info \u001B[38;5;241m=\u001B[39m (\n\u001B[0;32m--> 887\u001B[0m     \u001B[43mloader_impl\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mparse_saved_model_with_debug_info\u001B[49m\u001B[43m(\u001B[49m\u001B[43mexport_dir\u001B[49m\u001B[43m)\u001B[49m)\n\u001B[1;32m    889\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m (\u001B[38;5;28mlen\u001B[39m(saved_model_proto\u001B[38;5;241m.\u001B[39mmeta_graphs) \u001B[38;5;241m==\u001B[39m \u001B[38;5;241m1\u001B[39m \u001B[38;5;129;01mand\u001B[39;00m\n\u001B[1;32m    890\u001B[0m     saved_model_proto\u001B[38;5;241m.\u001B[39mmeta_graphs[\u001B[38;5;241m0\u001B[39m]\u001B[38;5;241m.\u001B[39mHasField(\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mobject_graph_def\u001B[39m\u001B[38;5;124m\"\u001B[39m)):\n\u001B[1;32m    891\u001B[0m   metrics\u001B[38;5;241m.\u001B[39mIncrementReadApi(_LOAD_V2_LABEL)\n",
      "File \u001B[0;32m~/anaconda3/envs/mlpipeline/lib/python3.9/site-packages/tensorflow/python/saved_model/loader_impl.py:57\u001B[0m, in \u001B[0;36mparse_saved_model_with_debug_info\u001B[0;34m(export_dir)\u001B[0m\n\u001B[1;32m     44\u001B[0m \u001B[38;5;28;01mdef\u001B[39;00m \u001B[38;5;21mparse_saved_model_with_debug_info\u001B[39m(export_dir):\n\u001B[1;32m     45\u001B[0m   \u001B[38;5;124;03m\"\"\"Reads the savedmodel as well as the graph debug info.\u001B[39;00m\n\u001B[1;32m     46\u001B[0m \n\u001B[1;32m     47\u001B[0m \u001B[38;5;124;03m  Args:\u001B[39;00m\n\u001B[0;32m   (...)\u001B[0m\n\u001B[1;32m     55\u001B[0m \u001B[38;5;124;03m    parsed. Missing graph debug info file is fine.\u001B[39;00m\n\u001B[1;32m     56\u001B[0m \u001B[38;5;124;03m  \"\"\"\u001B[39;00m\n\u001B[0;32m---> 57\u001B[0m   saved_model \u001B[38;5;241m=\u001B[39m \u001B[43mparse_saved_model\u001B[49m\u001B[43m(\u001B[49m\u001B[43mexport_dir\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m     59\u001B[0m   debug_info_path \u001B[38;5;241m=\u001B[39m file_io\u001B[38;5;241m.\u001B[39mjoin(\n\u001B[1;32m     60\u001B[0m       saved_model_utils\u001B[38;5;241m.\u001B[39mget_debug_dir(export_dir),\n\u001B[1;32m     61\u001B[0m       constants\u001B[38;5;241m.\u001B[39mDEBUG_INFO_FILENAME_PB)\n\u001B[1;32m     62\u001B[0m   debug_info \u001B[38;5;241m=\u001B[39m graph_debug_info_pb2\u001B[38;5;241m.\u001B[39mGraphDebugInfo()\n",
      "File \u001B[0;32m~/anaconda3/envs/mlpipeline/lib/python3.9/site-packages/tensorflow/python/saved_model/loader_impl.py:115\u001B[0m, in \u001B[0;36mparse_saved_model\u001B[0;34m(export_dir)\u001B[0m\n\u001B[1;32m    113\u001B[0m     \u001B[38;5;28;01mraise\u001B[39;00m \u001B[38;5;167;01mIOError\u001B[39;00m(\u001B[38;5;124mf\u001B[39m\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mCannot parse file \u001B[39m\u001B[38;5;132;01m{\u001B[39;00mpath_to_pbtxt\u001B[38;5;132;01m}\u001B[39;00m\u001B[38;5;124m: \u001B[39m\u001B[38;5;132;01m{\u001B[39;00m\u001B[38;5;28mstr\u001B[39m(e)\u001B[38;5;132;01m}\u001B[39;00m\u001B[38;5;124m.\u001B[39m\u001B[38;5;124m\"\u001B[39m)\n\u001B[1;32m    114\u001B[0m \u001B[38;5;28;01melse\u001B[39;00m:\n\u001B[0;32m--> 115\u001B[0m   \u001B[38;5;28;01mraise\u001B[39;00m \u001B[38;5;167;01mIOError\u001B[39;00m(\n\u001B[1;32m    116\u001B[0m       \u001B[38;5;124mf\u001B[39m\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mSavedModel file does not exist at: \u001B[39m\u001B[38;5;132;01m{\u001B[39;00mexport_dir\u001B[38;5;132;01m}\u001B[39;00m\u001B[38;5;132;01m{\u001B[39;00mos\u001B[38;5;241m.\u001B[39mpath\u001B[38;5;241m.\u001B[39msep\u001B[38;5;132;01m}\u001B[39;00m\u001B[38;5;124m\"\u001B[39m\n\u001B[1;32m    117\u001B[0m       \u001B[38;5;124mf\u001B[39m\u001B[38;5;124m\"\u001B[39m\u001B[38;5;130;01m{{\u001B[39;00m\u001B[38;5;132;01m{\u001B[39;00mconstants\u001B[38;5;241m.\u001B[39mSAVED_MODEL_FILENAME_PBTXT\u001B[38;5;132;01m}\u001B[39;00m\u001B[38;5;124m|\u001B[39m\u001B[38;5;124m\"\u001B[39m\n\u001B[1;32m    118\u001B[0m       \u001B[38;5;124mf\u001B[39m\u001B[38;5;124m\"\u001B[39m\u001B[38;5;132;01m{\u001B[39;00mconstants\u001B[38;5;241m.\u001B[39mSAVED_MODEL_FILENAME_PB\u001B[38;5;132;01m}\u001B[39;00m\u001B[38;5;130;01m}}\u001B[39;00m\u001B[38;5;124m\"\u001B[39m)\n",
      "\u001B[0;31mOSError\u001B[0m: SavedModel file does not exist at: s3://datasets/universal-sentence-encoder/4//{saved_model.pbtxt|saved_model.pb}"
     ]
    }
   ],
   "source": [
    "# embed = tf.compat.v1.saved_model.load_v2(f\"s3://datasets/{se_model_prefix}\")\n",
    "embed = tf.compat.v1.saved_model.load_v2(f\"s3://datasets/universal-sentence-encoder/4/\")"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {
    "scrolled": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-08-03 16:01:00.416611: I tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.cc:113] Plugin optimizer for device_type GPU is enabled.\n"
     ]
    }
   ],
   "source": [
    "import tensorflow_hub as hub\n",
    "# embed = hub.load(\"https://tfhub.dev/google/universal-sentence-encoder-large/5\")\n",
    "embed = hub.load(\"https://tfhub.dev/google/universal-sentence-encoder/4\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "print(\"x\")"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Since it would be too much to create the embeddings of `25,000` sentences and keep that in memory, we are going to slice our datasets into chunks of `500`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "To store our data into a `TFRecord` we need to encode the features as `tf.train.Feature`.  We are going to store the label of our data as list of tf.int64 and our Movie Review as a list of floats since after we encode the sentence using `USE` we will end-up with a embedding of `512` dimensions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def _embedded_sentence_feature(value):\n",
    "    # convert tensor to list of float values\n",
    "    input = value.numpy().ravel().tolist()\n",
    "    return tf.train.Feature(float_list=tf.train.FloatList(value=input))\n",
    "def _label_feature(value):\n",
    "    # convert tensor to list of float values\n",
    "    input = value.numpy().ravel().tolist()\n",
    "    return tf.train.Feature(int64_list=tf.train.Int64List(value=input))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def encode_label(label):\n",
    "    if label == \"pos\":\n",
    "        return tf.constant([1,0])\n",
    "    elif label == \"neg\":\n",
    "        return tf.constant([0,1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def serialize_example(label, sentence_tensor):\n",
    "    feature = {\n",
    "      \"sentence\": _embedded_sentence_feature(sentence_tensor[0]),\n",
    "      \"label\": _label_feature(label),\n",
    "    }\n",
    "    example_proto = tf.train.Example(features=tf.train.Features(feature=feature))\n",
    "    return example_proto"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import timeit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total of 25000 elements\n",
      "Total of 50 files of 500 records - 0.0003850420002891042\n",
      "Done serializing example ( 0.06819304200007537 - 5.29579997419205e-05 - 0.001390417000038724 ). Buffer 500 - 32.3230268750003215454\n",
      "Records in buffer 500 - 32.32309037499999\n",
      "done write example file, size 947513 - 32.43944870900032\n",
      "Done with batch 1/50 - 32.473159625000335\n",
      "Done serializing example ( 0.051969999999982974 - 5.300000020724838e-05 - 0.0015033329996185785 ). Buffer 500 - 80.3795619590000665\n",
      "Records in buffer 500 - 80.37957916699997\n",
      "done write example file, size 947747 - 80.55104795900024\n",
      "Done with batch 2/50 - 80.58765270899994\n",
      "Done serializing example ( 0.05199454200010223 - 5.3125000249565346e-05 - 0.0013241659999039257 ). Buffer 500 - 115.422089708999925\n",
      "Records in buffer 500 - 115.42210462499997\n",
      "done write example file, size 947801 - 115.56911745900015\n",
      "Done with batch 3/50 - 115.62044733399989\n",
      "Done serializing example ( 0.052842874999896594 - 5.304199976308155e-05 - 0.001404458000251907 ). Buffer 500 - 154.1904891250001924\n",
      "Records in buffer 500 - 154.19050425000023\n",
      "done write example file, size 947343 - 154.30908737500022\n",
      "Done with batch 4/50 - 154.34423800000013\n",
      "Done serializing example ( 0.06957595899984881 - 5.795800007035723e-05 - 0.0024123329999383714 ). Buffer 500 - 189.9845063340003489\n",
      "Records in buffer 500 - 189.98452108400033\n",
      "done write example file, size 947566 - 190.10389829199994\n",
      "Done with batch 5/50 - 190.1378614170003\n",
      "Done serializing example ( 0.07539220800026669 - 6.41249998807325e-05 - 0.0016273339997496805 ). Buffer 500 - 228.33298287500023723\n",
      "Records in buffer 500 - 228.33299716700003\n",
      "done write example file, size 947684 - 228.4740027920002\n",
      "Done with batch 6/50 - 228.5345977500001\n",
      "Done serializing example ( 0.05367095900010099 - 6.245799977477873e-05 - 0.0015944579999995767 ). Buffer 500 - 266.4464146250002545\n",
      "Records in buffer 500 - 266.4464303750001\n",
      "done write example file, size 947650 - 266.5699428339999\n",
      "Done with batch 7/50 - 266.5965248340003\n",
      "Done serializing example ( 0.07249845900014407 - 5.279100014377036e-05 - 0.0014037920000191662 ). Buffer 500 - 308.7481906670026935\n",
      "Records in buffer 500 - 308.7482069590001\n",
      "done write example file, size 947609 - 308.8779299170001\n",
      "Done with batch 8/50 - 308.93753637500004\n",
      "Done serializing example ( 0.07160766700008026 - 5.287500016493141e-05 - 0.0014120000000730215 ). Buffer 500 - 355.7463965420002336\n",
      "Records in buffer 500 - 355.74641199999996\n",
      "done write example file, size 947470 - 355.87899779200006\n",
      "Done with batch 9/50 - 355.9307344590002\n",
      "Done serializing example ( 0.07983270899967465 - 6.233300018720911e-05 - 0.001546333000078448 ). Buffer 500 - 400.13246366700014017\n",
      "Records in buffer 500 - 400.132479375\n",
      "done write example file, size 947724 - 400.2578475\n",
      "Done with batch 10/50 - 400.28799500000014\n",
      "Done serializing example ( 0.07319079199987755 - 5.316700026014587e-05 - 0.0013666249997186242 ). Buffer 500 - 438.613678749999966\n",
      "Records in buffer 500 - 438.6136933749999\n",
      "done write example file, size 947491 - 438.73368804200027\n",
      "Done with batch 11/50 - 438.7547211670003\n",
      "Done serializing example ( 0.05205979099991964 - 4.891700018561096e-05 - 0.002346957999634469 ). Buffer 500 - 480.54872483400004376\n",
      "Records in buffer 500 - 480.5487804170002\n",
      "done write example file, size 947616 - 480.6665915000003\n",
      "Done with batch 12/50 - 480.70523708400015\n",
      "Done serializing example ( 0.07614375000002838 - 5.879199989067274e-05 - 0.0014050419999875885 ). Buffer 500 - 518.736912334002021\n",
      "Records in buffer 500 - 518.736927375\n",
      "done write example file, size 947466 - 518.8557559999999\n",
      "Done with batch 13/50 - 518.8812443340003\n",
      "Done serializing example ( 0.07225995799990415 - 5.2834000143775484e-05 - 0.0015566249999210413 ). Buffer 500 - 561.14455833400011\n",
      "Records in buffer 500 - 561.1445737090003\n",
      "done write example file, size 947490 - 561.2687886670001\n",
      "Done with batch 14/50 - 561.304222834\n",
      "Done serializing example ( 0.03252179099990826 - 5.254200004856102e-05 - 0.0015642500002286397 ). Buffer 500 - 601.973682584000313\n",
      "Records in buffer 500 - 601.9736982500003\n",
      "done write example file, size 947697 - 602.1390511670002\n",
      "Done with batch 15/50 - 602.177412\n",
      "Done serializing example ( 0.06356979100019089 - 5.945899965809076e-05 - 0.001773250000042026 ). Buffer 500 - 642.7623392929000392\n",
      "Records in buffer 500 - 642.762356917\n",
      "done write example file, size 947625 - 642.887461625\n",
      "Done with batch 16/50 - 642.923300084\n",
      "Done serializing example ( 0.07617162499991537 - 5.333299986887141e-05 - 0.0016950420003922773 ). Buffer 500 - 683.945117752500022\n",
      "Records in buffer 500 - 683.9451333340003\n",
      "done write example file, size 947288 - 684.063221459\n",
      "Done with batch 17/50 - 684.0886827500003\n",
      "Done serializing example ( 0.0827892500001326 - 5.291699972076458e-05 - 0.0016010829999686393 ). Buffer 500 - 724.7337587500001112\n",
      "Records in buffer 500 - 724.7337740420003\n",
      "done write example file, size 947919 - 724.8503746670003\n",
      "Done with batch 18/50 - 724.877708125\n",
      "Done serializing example ( 0.07485450000012861 - 5.4707999879610725e-05 - 0.001460500000121101 ). Buffer 500 - 765.401037250000203\n",
      "Records in buffer 500 - 765.4010547500002\n",
      "done write example file, size 947461 - 765.5194692089999\n",
      "Done with batch 19/50 - 765.5481408750002\n",
      "Done serializing example ( 0.07696254200027397 - 5.97079997533001e-05 - 0.0025608330001887225 ). Buffer 500 - 807.0901184170002921\n",
      "Records in buffer 500 - 807.0901335000003\n",
      "done write example file, size 947553 - 807.2137614170001\n",
      "Done with batch 20/50 - 807.2396373750003\n",
      "Done serializing example ( 0.04739312500032611 - 5.474999989019125e-05 - 0.0014773339999010204 ). Buffer 500 - 852.01539416700011\n",
      "Records in buffer 500 - 852.015408709\n",
      "done write example file, size 947708 - 852.148035834\n",
      "Done with batch 21/50 - 852.1723695840001\n",
      "Done serializing example ( 0.05806341700008488 - 5.687499970008503e-05 - 0.0017167500000141445 ). Buffer 500 - 894.028981417000113\n",
      "Records in buffer 500 - 894.0289979590002\n",
      "done write example file, size 947300 - 894.1458632920003\n",
      "Done with batch 22/50 - 894.171445167\n",
      "Done serializing example ( 0.0667071660000147 - 5.4458999784401385e-05 - 0.0013832080003339797 ). Buffer 500 - 934.289857000000222\n",
      "Records in buffer 500 - 934.2898713750001\n",
      "done write example file, size 947399 - 934.4029451250003\n",
      "Done with batch 23/50 - 934.4281697920001\n",
      "Done serializing example ( 0.08027120799988552 - 5.995799983793404e-05 - 0.0015148749998843414 ). Buffer 500 - 982.643952252001111\n",
      "Records in buffer 500 - 982.6439774590003\n",
      "done write example file, size 947436 - 982.7715171670002\n",
      "Done with batch 24/50 - 982.8017761670003\n",
      "Done serializing example ( 0.08473912499994185 - 7.24169999557489e-05 - 0.0017248330000256828 ). Buffer 500 - 1026.117718459000101\n",
      "Records in buffer 500 - 1026.1177885840002\n",
      "done write example file, size 947603 - 1026.2324266250002\n",
      "Done with batch 25/50 - 1026.261317292\n",
      "Done serializing example ( 0.08218270799989114 - 5.154200016477262e-05 - 0.0014302499998848361 ). Buffer 500 - 1068.17337120990011\n",
      "Records in buffer 500 - 1068.1733907920002\n",
      "done write example file, size 947575 - 1068.294720375\n",
      "Done with batch 26/50 - 1068.3256874170002\n",
      "Done serializing example ( 0.11148812500005079 - 5.437500021798769e-05 - 0.0015559999997094565 ). Buffer 500 - 1112.658102209900023\n",
      "Records in buffer 500 - 1112.6582062920002\n",
      "done write example file, size 947602 - 1112.8109065000003\n",
      "Done with batch 27/50 - 1112.845727334\n",
      "Done serializing example ( 0.04907312499926775 - 5.891700038773706e-05 - 0.0017684579997876426 ). Buffer 500 - 1154.164328916999742\n",
      "Records in buffer 500 - 1154.1643439589998\n",
      "done write example file, size 947588 - 1154.2806902499997\n",
      "Done with batch 28/50 - 1154.3019240419999\n",
      "Done serializing example ( 0.07918258299923764 - 5.983400023978902e-05 - 0.0015588750002279994 ). Buffer 500 - 1198.13299474999993\n",
      "Records in buffer 500 - 1198.133011625\n",
      "done write example file, size 947299 - 1198.2548190839998\n",
      "Done with batch 29/50 - 1198.2916361670004\n",
      "Done serializing example ( 0.07879287499963539 - 5.47500003449386e-05 - 0.0017222079995917738 ). Buffer 500 - 1240.1929655840004173\n",
      "Records in buffer 500 - 1240.1929819590005\n",
      "done write example file, size 947453 - 1240.3132262500003\n",
      "Done with batch 30/50 - 1240.3500568750005\n",
      "Done serializing example ( 0.09385587499946269 - 6.070900053600781e-05 - 0.0015996249994714162 ). Buffer 500 - 1281.959894000000532\n",
      "Records in buffer 500 - 1281.9600187920005\n",
      "done write example file, size 947169 - 1282.0746166670006\n",
      "Done with batch 31/50 - 1282.098663459\n",
      "Done serializing example ( 0.07878670799982501 - 5.4542000725632533e-05 - 0.001585582999723556 ). Buffer 500 - 1323.316020166999878\n",
      "Records in buffer 500 - 1323.3160803339997\n",
      "done write example file, size 947469 - 1323.430325292\n",
      "Done with batch 32/50 - 1323.4572905000005\n",
      "Done serializing example ( 0.08701566699983232 - 5.2708000112033915e-05 - 0.001465958999688155 ). Buffer 500 - 1364.298539334000573\n",
      "Records in buffer 500 - 1364.2986335000005\n",
      "done write example file, size 947288 - 1364.4132382919997\n",
      "Done with batch 33/50 - 1364.4379325\n",
      "Done serializing example ( 0.08286950000001525 - 5.170799977349816e-05 - 0.0014071250006963965 ). Buffer 500 - 1406.588431000000136\n",
      "Records in buffer 500 - 1406.5884452090004\n",
      "done write example file, size 947519 - 1406.706542\n",
      "Done with batch 34/50 - 1406.7364052499997\n",
      "Done serializing example ( 0.08489954100059549 - 0.00012595899988809833 - 0.0015568329999950947 ). Buffer 500 - 1448.8335257919998\n",
      "Records in buffer 500 - 1448.8335421250003\n",
      "done write example file, size 947697 - 1448.9580289590003\n",
      "Done with batch 35/50 - 1448.9900065419997\n",
      "Done serializing example ( 0.08356800000001385 - 5.400000009103678e-05 - 0.0013584579992311774 ). Buffer 500 - 1493.95482791700032\n",
      "Records in buffer 500 - 1493.9548918339997\n",
      "done write example file, size 947610 - 1494.0758848750002\n",
      "Done with batch 36/50 - 1494.1024076250005\n",
      "Done serializing example ( 0.09655133399974147 - 5.354099994292483e-05 - 0.0015381670000351733 ). Buffer 500 - 1536.394292584000147\n",
      "Records in buffer 500 - 1536.3943544169997\n",
      "done write example file, size 947423 - 1536.51041575\n",
      "Done with batch 37/50 - 1536.54223025\n",
      "Done serializing example ( 0.08055887499995151 - 5.1250000069558155e-05 - 0.0014386670000021695 ). Buffer 500 - 1579.78205566700033\n",
      "Records in buffer 500 - 1579.7820700419998\n",
      "done write example file, size 947504 - 1579.9061527920003\n",
      "Done with batch 38/50 - 1579.92726125\n",
      "Done serializing example ( 0.09029766699950414 - 5.825000062031904e-05 - 0.0014176670001688763 ). Buffer 500 - 1621.730556084000594\n",
      "Records in buffer 500 - 1621.730570959\n",
      "done write example file, size 947609 - 1621.846117084\n",
      "Done with batch 39/50 - 1621.8745385840002\n",
      "Done serializing example ( 0.08626062499934051 - 5.512500047188951e-05 - 0.0016973329993561492 ). Buffer 500 - 1665.118608667000444\n",
      "Records in buffer 500 - 1665.118690917\n",
      "done write example file, size 947736 - 1665.300328292\n",
      "Done with batch 40/50 - 1665.4539640000003\n",
      "Done serializing example ( 0.08355220900011773 - 5.454100028146058e-05 - 0.0015059999996083206 ). Buffer 500 - 1711.988235584400025\n",
      "Records in buffer 500 - 1711.9883317920003\n",
      "done write example file, size 947459 - 1712.1152767499998\n",
      "Done with batch 41/50 - 1712.1524317089998\n",
      "Done serializing example ( 0.07310874999984662 - 5.999999939376721e-05 - 0.0015339170004153857 ). Buffer 500 - 1754.81683804200024\n",
      "Records in buffer 500 - 1754.816854959\n",
      "done write example file, size 947704 - 1754.931158292\n",
      "Done with batch 42/50 - 1754.9535476250003\n",
      "Done serializing example ( 0.08967037499951402 - 5.958299971098313e-05 - 0.0031956669999999576 ). Buffer 500 - 1798.590929959000344\n",
      "Records in buffer 500 - 1798.5909456669997\n",
      "done write example file, size 947415 - 1798.7305677499999\n",
      "Done with batch 43/50 - 1798.765273375\n",
      "Done serializing example ( 0.08097766700029752 - 5.874999988009222e-05 - 0.0016185829999812995 ). Buffer 500 - 1841.396960500000217\n",
      "Records in buffer 500 - 1841.3970074170002\n",
      "done write example file, size 947542 - 1841.5161588339997\n",
      "Done with batch 44/50 - 1841.541820167\n",
      "Done serializing example ( 0.08116579200031993 - 5.18749993716483e-05 - 0.0015445000008185161 ). Buffer 500 - 1884.4080380840005017\n",
      "Records in buffer 500 - 1884.4080535419998\n",
      "done write example file, size 947154 - 1884.52180925\n",
      "Done with batch 45/50 - 1884.5452485840005\n",
      "Done serializing example ( 0.08937441699981719 - 5.995799983793404e-05 - 0.0019786670000030426 ). Buffer 500 - 1931.638364667000114\n",
      "Records in buffer 500 - 1931.6383794590001\n",
      "done write example file, size 947678 - 1931.7607162089998\n",
      "Done with batch 46/50 - 1931.7807290420005\n",
      "Done serializing example ( 0.08381433299928176 - 6.037500043021282e-05 - 0.0016189589996429277 ). Buffer 500 - 1976.203789166999747\n",
      "Records in buffer 500 - 1976.203852959\n",
      "done write example file, size 947962 - 1976.3268479589997\n",
      "Done with batch 47/50 - 1976.359565709\n",
      "Done serializing example ( 0.08806100000037986 - 5.254200004856102e-05 - 0.0014526249997288687 ). Buffer 500 - 2021.480726959997851\n",
      "Records in buffer 500 - 2021.480792292\n",
      "done write example file, size 947680 - 2021.5960282920005\n",
      "Done with batch 48/50 - 2021.6290334169998\n",
      "Done serializing example ( 0.08087525000064488 - 5.229199996392708e-05 - 0.0018137919996661367 ). Buffer 500 - 2065.437285084999742\n",
      "Records in buffer 500 - 2065.4373101250003\n",
      "done write example file, size 947551 - 2065.557782584\n",
      "Done with batch 49/50 - 2065.59102625\n",
      "Done serializing example ( 0.09072995799942873 - 5.487500038725557e-05 - 0.0014988329994594096 ). Buffer 500 - 2110.56934237500053\n",
      "Records in buffer 500 - 2110.569416542\n",
      "done write example file, size 947688 - 2110.6871291670004\n",
      "Done with batch 50/50 - 2110.722045875\n",
      "\n",
      "total time is : 2110.722068209\n",
      "Total of 25000 elements\n",
      "Total of 50 files of 500 records - 5.707999662263319e-06\n",
      "Done serializing example ( 0.08663687500029482 - 5.479199990077177e-05 - 0.001766374999533582 ). Buffer 500 - 46.117605957999339866\n",
      "Records in buffer 500 - 46.117620874999375\n",
      "done write example file, size 947445 - 46.23405908299992\n",
      "Done with batch 1/50 - 46.26431012499961\n",
      "Done serializing example ( 0.08658041699982277 - 5.175000023882603e-05 - 0.0015441249997820705 ). Buffer 500 - 90.26839799999925296\n",
      "Records in buffer 500 - 90.26846324999951\n",
      "done write example file, size 947881 - 90.38980249999986\n",
      "Done with batch 2/50 - 90.4153384579995\n",
      "Done serializing example ( 0.0719378339999821 - 5.07079994349624e-05 - 0.0014373330004673335 ). Buffer 500 - 138.57382995799986684\n",
      "Records in buffer 500 - 138.57389308299935\n",
      "done write example file, size 947610 - 138.7059161669995\n",
      "Done with batch 3/50 - 138.73427566699957\n",
      "Done serializing example ( 0.09196212500046386 - 5.241700000624405e-05 - 0.0015156669996940764 ). Buffer 500 - 185.168777457999577\n",
      "Records in buffer 500 - 185.16879304199938\n",
      "done write example file, size 947837 - 185.28772083299918\n",
      "Done with batch 4/50 - 185.31417591699937\n",
      "Done serializing example ( 0.094853833000343 - 6.13330003034207e-05 - 0.0016925419995459379 ). Buffer 500 - 231.565972041999889543\n",
      "Records in buffer 500 - 231.56598662499982\n",
      "done write example file, size 947647 - 231.68272895799964\n",
      "Done with batch 5/50 - 231.71197545799987\n",
      "Done serializing example ( 0.09935920800035092 - 5.762499949923949e-05 - 0.0015146670002650353 ). Buffer 500 - 276.7342232499995644\n",
      "Records in buffer 500 - 276.7342373329993\n",
      "done write example file, size 947685 - 276.84812104199955\n",
      "Done with batch 6/50 - 276.8710896249995\n",
      "Done serializing example ( 0.09009791700009373 - 5.095799951959634e-05 - 0.0015039580002849107 ). Buffer 500 - 322.178254124999597\n",
      "Records in buffer 500 - 322.1782689169995\n",
      "done write example file, size 947568 - 322.2938912919999\n",
      "Done with batch 7/50 - 322.3224571669998\n",
      "Done serializing example ( 0.1044711249996908 - 5.479199990077177e-05 - 0.0018984159996762173 ). Buffer 500 - 368.60407491699975775\n",
      "Records in buffer 500 - 368.6040900829994\n",
      "done write example file, size 947415 - 368.7250172079994\n",
      "Done with batch 8/50 - 368.76322187499954\n",
      "Done serializing example ( 0.08939350000036939 - 5.383400002756389e-05 - 0.0012520000000222353 ). Buffer 500 - 414.6758378749991676\n",
      "Records in buffer 500 - 414.6758995\n",
      "done write example file, size 947424 - 414.79698908299997\n",
      "Done with batch 9/50 - 414.8223159169993\n",
      "Done serializing example ( 0.08903583400024218 - 5.595799939328572e-05 - 0.0014373750000231666 ). Buffer 500 - 461.697843416999687\n",
      "Records in buffer 500 - 461.69785704199967\n",
      "done write example file, size 947735 - 461.8294665829999\n",
      "Done with batch 10/50 - 461.8677388749993\n",
      "Done serializing example ( 0.11401212500004476 - 5.9500000133994035e-05 - 0.001961250000022119 ). Buffer 500 - 509.393789207999364\n",
      "Records in buffer 500 - 509.39385395799945\n",
      "done write example file, size 947823 - 509.5391392499996\n",
      "Done with batch 11/50 - 509.58851208299984\n",
      "Done serializing example ( 0.09816599999976461 - 5.250000049272785e-05 - 0.0015779999994265381 ). Buffer 500 - 556.96589433299967\n",
      "Records in buffer 500 - 556.9659087499995\n",
      "done write example file, size 947811 - 557.1126960419997\n",
      "Done with batch 12/50 - 557.1672831669994\n",
      "Done serializing example ( 0.09213620799982891 - 5.3249999837134965e-05 - 0.0014188329996613902 ). Buffer 500 - 604.55867049999962\n",
      "Records in buffer 500 - 604.5587331669994\n",
      "done write example file, size 947835 - 604.6918852079998\n",
      "Done with batch 13/50 - 604.7403409999997\n",
      "Done serializing example ( 0.1081966250003461 - 5.1125000027241185e-05 - 0.001396999999997206 ). Buffer 500 - 652.5340234579999983\n",
      "Records in buffer 500 - 652.5340385\n",
      "done write example file, size 947552 - 652.6525449169994\n",
      "Done with batch 14/50 - 652.6814391669996\n",
      "Done serializing example ( 0.1011775420001868 - 5.408299966802588e-05 - 0.0015219580000120914 ). Buffer 500 - 699.754302624999696\n",
      "Records in buffer 500 - 699.754318583\n",
      "done write example file, size 947662 - 699.8712428749996\n",
      "Done with batch 15/50 - 699.8959765829995\n",
      "Done serializing example ( 0.08925041699967551 - 5.154100017534802e-05 - 0.0028722090000883327 ). Buffer 500 - 748.124247666999565\n",
      "Records in buffer 500 - 748.1242627499996\n",
      "done write example file, size 947806 - 748.247799625\n",
      "Done with batch 16/50 - 748.2858049999995\n",
      "Done serializing example ( 0.08676716599984502 - 5.0874999942607246e-05 - 0.0013476670001182356 ). Buffer 500 - 796.0870784999997\n",
      "Records in buffer 500 - 796.0870923749999\n",
      "done write example file, size 947578 - 796.2395564999997\n",
      "Done with batch 17/50 - 796.2636957079994\n",
      "Done serializing example ( 0.08667770800002472 - 5.2208999477443285e-05 - 0.0013753330003964948 ). Buffer 500 - 842.06461066699922\n",
      "Records in buffer 500 - 842.0646952079996\n",
      "done write example file, size 947649 - 842.1767653329998\n",
      "Done with batch 18/50 - 842.1975180419995\n",
      "Done serializing example ( 0.0860770830004185 - 5.120799960423028e-05 - 0.0014062090003790217 ). Buffer 500 - 887.9324311669998952\n",
      "Records in buffer 500 - 887.9324456249997\n",
      "done write example file, size 947851 - 888.0429802919998\n",
      "Done with batch 19/50 - 888.0635619579998\n",
      "Done serializing example ( 0.11070924999967247 - 5.3499999921768904e-05 - 0.0014238339999792515 ). Buffer 500 - 933.9825152499998\n",
      "Records in buffer 500 - 933.982576292\n",
      "done write example file, size 947736 - 934.0966731669996\n",
      "Done with batch 20/50 - 934.1150418329999\n",
      "Done serializing example ( 0.09627091700076562 - 5.3499999921768904e-05 - 0.0014582919993699761 ). Buffer 500 - 979.90023099999965\n",
      "Records in buffer 500 - 979.9003162919998\n",
      "done write example file, size 947345 - 980.0110119999999\n",
      "Done with batch 21/50 - 980.0302431249993\n",
      "Done serializing example ( 0.09310558300057892 - 4.933299987897044e-05 - 0.0014520419999826117 ). Buffer 500 - 1025.983113792999922\n",
      "Records in buffer 500 - 1025.9831973749997\n",
      "done write example file, size 947603 - 1026.1045718329997\n",
      "Done with batch 22/50 - 1026.1232467499995\n",
      "Done serializing example ( 0.09405887499997334 - 5.7499999456922524e-05 - 0.0014004160002514254 ). Buffer 500 - 1073.23950849999984\n",
      "Records in buffer 500 - 1073.2395220829994\n",
      "done write example file, size 947737 - 1073.3528998329994\n",
      "Done with batch 23/50 - 1073.3779205829996\n",
      "Done serializing example ( 0.10973508299957757 - 4.970800000592135e-05 - 0.0015630420002707979 ). Buffer 500 - 1120.366049041999866\n",
      "Records in buffer 500 - 1120.3661102919996\n",
      "done write example file, size 947727 - 1120.4776743329994\n",
      "Done with batch 24/50 - 1120.5016417499992\n",
      "Done serializing example ( 0.13109016599992174 - 5.679200057784328e-05 - 0.0017937080001502181 ). Buffer 500 - 1167.665649791999467\n",
      "Records in buffer 500 - 1167.6657175\n",
      "done write example file, size 948068 - 1167.7803157079998\n",
      "Done with batch 25/50 - 1167.803697458\n",
      "Done serializing example ( 0.10314608299995598 - 4.974999956175452e-05 - 0.0015199999998003477 ). Buffer 500 - 1215.30952341699967\n",
      "Records in buffer 500 - 1215.3095371669997\n",
      "done write example file, size 947773 - 1215.4225564999997\n",
      "Done with batch 26/50 - 1215.4465447079992\n",
      "Done serializing example ( 0.09157845899972017 - 5.1791000259981956e-05 - 0.0015847089998715091 ). Buffer 500 - 1263.31462016699977\n",
      "Records in buffer 500 - 1263.3146346249996\n",
      "done write example file, size 947645 - 1263.4267682079999\n",
      "Done with batch 27/50 - 1263.4522772499995\n",
      "Done serializing example ( 0.0946256249999351 - 4.904099932900863e-05 - 0.0013633750004373724 ). Buffer 500 - 1311.565826791999893\n",
      "Records in buffer 500 - 1311.5658407499996\n",
      "done write example file, size 947550 - 1311.6909045829998\n",
      "Done with batch 28/50 - 1311.7114448329994\n",
      "Done serializing example ( 0.12245558299946424 - 6.220900013431674e-05 - 0.00143441599993821 ). Buffer 500 - 1360.45486804199949214\n",
      "Records in buffer 500 - 1360.4549353329994\n",
      "done write example file, size 947775 - 1360.5772167919995\n",
      "Done with batch 29/50 - 1360.597066208\n",
      "Done serializing example ( 0.09287454199966305 - 4.9500000386615284e-05 - 0.0013284999995448743 ). Buffer 500 - 1409.1179433749994\n",
      "Records in buffer 500 - 1409.1179573749996\n",
      "done write example file, size 947609 - 1409.2303955419993\n",
      "Done with batch 30/50 - 1409.2564921249996\n",
      "Done serializing example ( 0.09228470900052343 - 4.974999956175452e-05 - 0.0014738750005562906 ). Buffer 500 - 1459.848406999999528\n",
      "Records in buffer 500 - 1459.8484224169997\n",
      "done write example file, size 947421 - 1459.968898792\n",
      "Done with batch 31/50 - 1459.9937315829993\n",
      "Done serializing example ( 0.09270808400015085 - 5.029099975217832e-05 - 0.0013575840002886252 ). Buffer 500 - 1510.649858999999269\n",
      "Records in buffer 500 - 1510.6498723329996\n",
      "done write example file, size 947180 - 1510.7626272079997\n",
      "Done with batch 32/50 - 1510.7857676249996\n",
      "Done serializing example ( 0.0905334999997649 - 5.183400025998708e-05 - 0.0014809160002187127 ). Buffer 500 - 1560.8171210829996468\n",
      "Records in buffer 500 - 1560.8171358749996\n",
      "done write example file, size 948029 - 1560.9317326669998\n",
      "Done with batch 33/50 - 1560.9549124579999\n",
      "Done serializing example ( 0.11768254199978401 - 5.229199996392708e-05 - 0.0013692080001419527 ). Buffer 500 - 1615.270245207999331\n",
      "Records in buffer 500 - 1615.2702596669997\n",
      "done write example file, size 947416 - 1615.3832831669997\n",
      "Done with batch 34/50 - 1615.4048712919994\n",
      "Done serializing example ( 0.10974787499981176 - 5.016599970986135e-05 - 0.0013555840005210484 ). Buffer 500 - 1671.01890812499912\n",
      "Records in buffer 500 - 1671.018922667\n",
      "done write example file, size 947645 - 1671.134116708\n",
      "Done with batch 35/50 - 1671.1554187079992\n",
      "Done serializing example ( 0.10708629100008693 - 5.7375000324100256e-05 - 0.0015435420000358135 ). Buffer 500 - 1726.83994891699978\n",
      "Records in buffer 500 - 1726.8399626669998\n",
      "done write example file, size 947694 - 1726.9547091249997\n",
      "Done with batch 36/50 - 1726.9916789169993\n",
      "Done serializing example ( 0.09668558400062466 - 4.783299937116681e-05 - 0.0013600420006696368 ). Buffer 500 - 1782.30911874999959\n",
      "Records in buffer 500 - 1782.3091772499993\n",
      "done write example file, size 947962 - 1782.4220218749997\n",
      "Done with batch 37/50 - 1782.444585792\n",
      "Done serializing example ( 0.07904412500010949 - 5.933299962634919e-05 - 0.00144216700027755 ). Buffer 500 - 1838.01644662499989969\n",
      "Records in buffer 500 - 1838.0164608749992\n",
      "done write example file, size 947428 - 1838.1290159579994\n",
      "Done with batch 38/50 - 1838.1518786249999\n",
      "Done serializing example ( 0.1279420000000755 - 5.1790999350487255e-05 - 0.0014501250007015187 ). Buffer 500 - 1895.682373041999224\n",
      "Records in buffer 500 - 1895.6824782079993\n",
      "done write example file, size 947468 - 1895.7984375419992\n",
      "Done with batch 39/50 - 1895.8202151669993\n",
      "Done serializing example ( 0.11391487499986397 - 4.891700064035831e-05 - 0.0014286659998106188 ). Buffer 500 - 1954.35115762499944\n",
      "Records in buffer 500 - 1954.3511721669993\n",
      "done write example file, size 947667 - 1954.4660776669998\n",
      "Done with batch 40/50 - 1954.4897196249995\n",
      "Done serializing example ( 0.0751572089993715 - 4.9791000492405146e-05 - 0.0013472499995259568 ). Buffer 500 - 2014.321589958997528\n",
      "Records in buffer 500 - 2014.3216043329994\n",
      "done write example file, size 947919 - 2014.4348205829992\n",
      "Done with batch 41/50 - 2014.458915708\n",
      "Done serializing example ( 0.13125012499949662 - 5.541599966818467e-05 - 0.001601209000000381 ). Buffer 500 - 2075.051760458599577\n",
      "Records in buffer 500 - 2075.051774291999\n",
      "done write example file, size 947674 - 2075.1654995829995\n",
      "Done with batch 42/50 - 2075.1866998329997\n",
      "Done serializing example ( 0.16558241600068868 - 5.616699945676373e-05 - 0.0014400000000023283 ). Buffer 500 - 2136.83846437499968\n",
      "Records in buffer 500 - 2136.8385267919994\n",
      "done write example file, size 947789 - 2136.9514080419995\n",
      "Done with batch 43/50 - 2136.9717267079996\n",
      "Done serializing example ( 0.12314795899965247 - 6.100000064179767e-05 - 0.0015309999998862622 ). Buffer 500 - 2199.93710966699935\n",
      "Records in buffer 500 - 2199.9371849579993\n",
      "done write example file, size 947576 - 2200.0673347499996\n",
      "Done with batch 44/50 - 2200.1018484999995\n",
      "Done serializing example ( 0.13232166599937045 - 5.120900004840223e-05 - 0.0014196660004017758 ). Buffer 500 - 2265.60511999999933\n",
      "Records in buffer 500 - 2265.6052055419996\n",
      "done write example file, size 947701 - 2265.7187194999997\n",
      "Done with batch 45/50 - 2265.7398069579995\n",
      "Done serializing example ( 0.11099216599995998 - 5.104200045025209e-05 - 0.001331124999524036 ). Buffer 500 - 2331.729639292999547\n",
      "Records in buffer 500 - 2331.7296995829993\n",
      "done write example file, size 947257 - 2331.8436557919995\n",
      "Done with batch 46/50 - 2331.864186125\n",
      "Done serializing example ( 0.13549812500059488 - 5.108299956191331e-05 - 0.0014302090003184276 ). Buffer 500 - 2396.36820837595975\n",
      "Records in buffer 500 - 2396.3682708749993\n",
      "done write example file, size 947753 - 2396.4826629169993\n",
      "Done with batch 47/50 - 2396.509070124999\n",
      "Done serializing example ( 0.1326366670000425 - 6.370800019794842e-05 - 0.001420499999767344 ). Buffer 500 - 2462.2447928759169996\n",
      "Records in buffer 500 - 2462.2448088329993\n",
      "done write example file, size 947937 - 2462.369689542\n",
      "Done with batch 48/50 - 2462.389671958\n",
      "Done serializing example ( 0.12713341699964076 - 4.933299987897044e-05 - 0.0014004169997861027 ). Buffer 500 - 2529.30225537499973\n",
      "Records in buffer 500 - 2529.302305333\n",
      "done write example file, size 947696 - 2529.415804542\n",
      "Done with batch 49/50 - 2529.4362684999996\n",
      "Done serializing example ( 0.19472954100001516 - 5.587499981629662e-05 - 0.0014076669995120028 ). Buffer 500 - 2599.76735849999933\n",
      "Records in buffer 500 - 2599.767421332999\n",
      "done write example file, size 947949 - 2599.894344167\n",
      "Done with batch 50/50 - 2599.9524559169995\n",
      "\n",
      "total time is : 2599.952476208\n",
      "Done!\n"
     ]
    }
   ],
   "source": [
    "def process_examples(records,prefix=\"\"):\n",
    "    starttime = timeit.default_timer()\n",
    "    total_training = len(records)\n",
    "    print(f\"Total of {total_training} elements\")\n",
    "    total_batches = math.floor(total_training / tf_record_file_size)\n",
    "    if total_training % tf_record_file_size != 0:\n",
    "        total_batches += 1 \n",
    "    print(f\"Total of {total_batches} files of {tf_record_file_size} records - {timeit.default_timer() - starttime}\")\n",
    "\n",
    "    counter = 0\n",
    "    file_counter = 0\n",
    "    buffer = []\n",
    "    file_list = []\n",
    "    for i in range(total_training):\n",
    "        counter += 1\n",
    "        prev_t = timeit.default_timer() - starttime\n",
    "        sentence_embedding = embed([records[i][\"text\"]])\n",
    "        emb_t = timeit.default_timer() - starttime - prev_t\n",
    "        label_encoded = encode_label(records[i][\"label\"])\n",
    "        lab_t = timeit.default_timer() - starttime - prev_t - emb_t\n",
    "        record = serialize_example(label_encoded, sentence_embedding)\n",
    "        rec_t = timeit.default_timer() - starttime - prev_t - emb_t - lab_t\n",
    "        buffer.append(record)\n",
    "        print(f\"\\rDone serializing example ( {emb_t} - {lab_t} - {rec_t} ). Buffer {len(buffer)} - {timeit.default_timer() - starttime}\",end=\"\")\n",
    "\n",
    "        if len(buffer) >= tf_record_file_size:\n",
    "            print(\"\")\n",
    "            print(f\"Records in buffer {len(buffer)} - {timeit.default_timer() - starttime}\")\n",
    "            # save this buffer of examples as a file to MinIO\n",
    "            counter = 0\n",
    "            file_counter+=1\n",
    "            file_name = f\"{prefix}_file{file_counter}.tfrecord\"\n",
    "            with open(file_name,\"w+\") as f:\n",
    "                with tf.io.TFRecordWriter(f.name,options=\"GZIP\") as writer:\n",
    "                    for example in buffer:\n",
    "                        writer.write(example.SerializeToString())\n",
    "            print(f\"done write example file, size {os.stat(file_name).st_size} - {timeit.default_timer() - starttime}\")\n",
    "            try:\n",
    "                minioClient.fput_object(datasets_bucket, f\"{preprocessed_data_folder}/{file_name}\", file_name)\n",
    "            except S3Error as err:\n",
    "                print(err)\n",
    "            file_list.append(file_name)\n",
    "            os.remove(file_name)\n",
    "            buffer=[]\n",
    "            print(f\"Done with batch {file_counter}/{total_batches} - {timeit.default_timer() - starttime}\")\n",
    "    print(\"\")\n",
    "    if len(buffer) > 0:\n",
    "        file_counter+=1\n",
    "        file_name = f\"file{file_counter}.tfrecord\"\n",
    "        with open(file_name,\"w+\") as f:\n",
    "            with tf.io.TFRecordWriter(f.name) as writer:\n",
    "                for example in buffer:\n",
    "                    writer.write(example.SerializeToString())\n",
    "        try:\n",
    "            minioClient.fput_object(datasets_bucket, f\"{preprocessed_data_folder}/{file_name}\", file_name)\n",
    "        except S3Error as err:\n",
    "            print(err)\n",
    "        file_list.append(file_name)\n",
    "        os.remove(file_name)\n",
    "        buffer=[]\n",
    "    print(\"total time is :\", timeit.default_timer() - starttime)\n",
    "    return file_list\n",
    "process_examples(train,prefix=\"train\")\n",
    "process_examples(test,prefix=\"test\")\n",
    "print(\"Done!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "At this point we are done preprocessing our data. We have a set of `.tfrecord` files stored on our bucket. We will now feed that to the model allowing it to consume and train concurrently. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "We are going to get a list of files (training data) from MinIO. Technically the pre-processing stage and the training stage could be completely decoupled so it's a good idea to list the file chunks we have in bucket."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# List all training tfrecord files\n",
    "objects = minioClient.list_objects(datasets_bucket, prefix=\"train\")\n",
    "training_files_list = []\n",
    "for obj in objects:\n",
    "    training_files_list.append(obj.object_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# List all testing tfrecord files\n",
    "objects = minioClient.list_objects(datasets_bucket, prefix=\"test\")\n",
    "testing_files_list = []\n",
    "for obj in objects:\n",
    "    testing_files_list.append(obj.object_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Now let us create a `tf.data.Dataset` that loads records from our files on MinIO as they become needed. To do that we are going to take the list of files we have and format them in a way that references the location of the actual objects. We will do this for the testing dataset as well. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "all_training_filenames = [f\"s3://datasets/{f}\" for f in training_files_list]\n",
    "testing_filenames = [f\"s3://datasets/{f}\" for f in testing_files_list]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "The following step is optional, but I recommend it. I am going to split my training dataset into two sets, `90%` of the data for training and `10%` of the data for validation, the model won't learn on the validation data but it will help the model train better."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "total_train_data_files = math.floor(len(all_training_filenames)*0.9)\n",
    "if total_train_data_files == len(all_training_filenames):\n",
    "    total_train_data_files -= 1\n",
    "training_files = all_training_filenames[0:total_train_data_files]\n",
    "validation_files = all_training_filenames[total_train_data_files:]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Now let's create the `tf.data` datasets:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "AUTO = tf.data.experimental.AUTOTUNE\n",
    "ignore_order = tf.data.Options()\n",
    "ignore_order.experimental_deterministic = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "dataset = tf.data.TFRecordDataset(training_files,num_parallel_reads=AUTO,compression_type=\"GZIP\")\n",
    "dataset = dataset.with_options(ignore_order)\n",
    "\n",
    "validation = tf.data.TFRecordDataset(validation_files,num_parallel_reads=AUTO,compression_type=\"GZIP\")\n",
    "validation = validation.with_options(ignore_order)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "testing_dataset = tf.data.TFRecordDataset(testing_filenames,num_parallel_reads=AUTO,compression_type=\"GZIP\")\n",
    "testing_dataset = testing_dataset.with_options(ignore_order)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "In order to decode our `TFRecord` encoded files we are going to need a decoding function that does the exact opposite of our `serialize_example` function. Since the data coming out of the `TFRecord` has shape `(512,)` and `(2,)` respectively, we are going to reshape it as well since that's the format our model will be expecting to receive."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def decode_fn(record_bytes): \n",
    "    schema = {\n",
    "        \"label\": tf.io.FixedLenFeature([2], dtype=tf.int64), \n",
    "        \"sentence\": tf.io.FixedLenFeature([512], dtype=tf.float32),\n",
    "        }\n",
    "    \n",
    "    tf_example = tf.io.parse_single_example(record_bytes,schema) \n",
    "    new_shape = tf.reshape(tf_example[\"sentence\"],[1,512])\n",
    "    label = tf.reshape(tf_example[\"label\"],[1,2])\n",
    "    return new_shape,label"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Let's build our model, nothing fancy, I'm just going to use a couple of Dense layers with a [softmax](https://en.wikipedia.org/wiki/Softmax_function) activation at the end.  We are trying to predict whether the input is `positive` or `negative` so we are going to get probabilities of the likelihood of each."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from tensorflow import keras"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "model = keras.Sequential()\n",
    "\n",
    "model.add(\n",
    "  keras.layers.Dense(\n",
    "    units=256,\n",
    "    input_shape=(1,512 ),\n",
    "    activation=\"relu\"\n",
    "  )\n",
    ")\n",
    "model.add(\n",
    "  keras.layers.Dropout(rate=0.5)\n",
    ")\n",
    "\n",
    "model.add(\n",
    "  keras.layers.Dense(\n",
    "    units=16,\n",
    "    activation=\"relu\"\n",
    "  )\n",
    ")\n",
    "model.add(\n",
    "  keras.layers.Dropout(rate=0.5)\n",
    ")\n",
    "\n",
    "model.add(keras.layers.Dense(2, activation=\"softmax\"))\n",
    "model.compile(\n",
    "    loss=\"categorical_crossentropy\",\n",
    "    optimizer=keras.optimizers.Adam(0.001),\n",
    "    metrics=[\"accuracy\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "![Structure of our Deep Learning model](pic2.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"sequential\"\n",
      "_________________________________________________________________\n",
      " Layer (type)                Output Shape              Param #   \n",
      "=================================================================\n",
      " dense (Dense)               (None, 1, 256)            131328    \n",
      "                                                                 \n",
      " dropout (Dropout)           (None, 1, 256)            0         \n",
      "                                                                 \n",
      " dense_1 (Dense)             (None, 1, 16)             4112      \n",
      "                                                                 \n",
      " dropout_1 (Dropout)         (None, 1, 16)             0         \n",
      "                                                                 \n",
      " dense_2 (Dense)             (None, 1, 2)              34        \n",
      "                                                                 \n",
      "=================================================================\n",
      "Total params: 135,474\n",
      "Trainable params: 135,474\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Let's prepare our datasets for the training stage by having them repeat themselves a little and batch `128` items at a time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "mapped_ds = dataset.map(decode_fn)\n",
    "mapped_ds = mapped_ds.repeat(5)\n",
    "mapped_ds = mapped_ds.batch(128)\n",
    "\n",
    "mapped_validation = validation.map(decode_fn)\n",
    "mapped_validation = mapped_validation.repeat(5)\n",
    "mapped_validation = mapped_validation.batch(128)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "testing_mapped_ds = testing_dataset.map(decode_fn)\n",
    "testing_mapped_ds = testing_mapped_ds.repeat(5)\n",
    "testing_mapped_ds = testing_mapped_ds.batch(128)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "As we train we wanna make sure to store checkpoints of our model in case the training gets interrupted and we wanna resume where we left off, so we are going to use the keras callback `tf.keras.callbacks.ModelCheckpoint` to have TensorFlow save the checkpoint to `MinIO` after every epoch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "checkpoint_path = f\"s3://{datasets_bucket}/checkpoints/cp.ckpt\"\n",
    "cp_callback = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_path,\n",
    "                                                 save_weights_only=True,\n",
    "                                                 verbose=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "We also want to save the `TensorBoard`  histograms so we are going to add a callback to store those in our bucket under the `logs/imdb/` prefix. We are identifying this run with a `model_note` and the current time, this is so we can tell apart different instances of training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from datetime import datetime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "model_note=\"256\"\n",
    "logdir = f\"s3://{datasets_bucket}/logs/imdb/{model_note}-\" + datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
    "tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=logdir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Finally we will train the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-08-04 14:11:10.479669: W tensorflow/core/framework/op_kernel.cc:1745] OP_REQUIRES failed at summary_kernels.cc:65 : UNIMPLEMENTED: File system scheme 's3' not implemented (file: 's3://datasets/logs/imdb/256-20220804-141016/train')\n"
     ]
    },
    {
     "ename": "UnimplementedError",
     "evalue": "File system scheme 's3' not implemented (file: 's3://datasets/logs/imdb/256-20220804-141016/train') [Op:CreateSummaryFileWriter]",
     "output_type": "error",
     "traceback": [
      "\u001B[0;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[0;31mUnimplementedError\u001B[0m                        Traceback (most recent call last)",
      "Input \u001B[0;32mIn [93]\u001B[0m, in \u001B[0;36m<cell line: 1>\u001B[0;34m()\u001B[0m\n\u001B[0;32m----> 1\u001B[0m history \u001B[38;5;241m=\u001B[39m \u001B[43mmodel\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mfit\u001B[49m\u001B[43m(\u001B[49m\n\u001B[1;32m      2\u001B[0m \u001B[43m    \u001B[49m\u001B[43mmapped_ds\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m      3\u001B[0m \u001B[43m    \u001B[49m\u001B[43mepochs\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[38;5;241;43m10\u001B[39;49m\u001B[43m,\u001B[49m\n\u001B[1;32m      4\u001B[0m \u001B[43m    \u001B[49m\u001B[43mcallbacks\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43m[\u001B[49m\u001B[43mcp_callback\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mtensorboard_callback\u001B[49m\u001B[43m]\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m      5\u001B[0m \u001B[43m    \u001B[49m\u001B[43mvalidation_data\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43mmapped_validation\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m      6\u001B[0m \u001B[43m)\u001B[49m\n",
      "File \u001B[0;32m~/miniforge3/envs/mlpipeline/lib/python3.9/site-packages/keras/utils/traceback_utils.py:67\u001B[0m, in \u001B[0;36mfilter_traceback.<locals>.error_handler\u001B[0;34m(*args, **kwargs)\u001B[0m\n\u001B[1;32m     65\u001B[0m \u001B[38;5;28;01mexcept\u001B[39;00m \u001B[38;5;167;01mException\u001B[39;00m \u001B[38;5;28;01mas\u001B[39;00m e:  \u001B[38;5;66;03m# pylint: disable=broad-except\u001B[39;00m\n\u001B[1;32m     66\u001B[0m   filtered_tb \u001B[38;5;241m=\u001B[39m _process_traceback_frames(e\u001B[38;5;241m.\u001B[39m__traceback__)\n\u001B[0;32m---> 67\u001B[0m   \u001B[38;5;28;01mraise\u001B[39;00m e\u001B[38;5;241m.\u001B[39mwith_traceback(filtered_tb) \u001B[38;5;28;01mfrom\u001B[39;00m \u001B[38;5;28mNone\u001B[39m\n\u001B[1;32m     68\u001B[0m \u001B[38;5;28;01mfinally\u001B[39;00m:\n\u001B[1;32m     69\u001B[0m   \u001B[38;5;28;01mdel\u001B[39;00m filtered_tb\n",
      "File \u001B[0;32m~/miniforge3/envs/mlpipeline/lib/python3.9/site-packages/tensorflow/python/framework/ops.py:7164\u001B[0m, in \u001B[0;36mraise_from_not_ok_status\u001B[0;34m(e, name)\u001B[0m\n\u001B[1;32m   7162\u001B[0m \u001B[38;5;28;01mdef\u001B[39;00m \u001B[38;5;21mraise_from_not_ok_status\u001B[39m(e, name):\n\u001B[1;32m   7163\u001B[0m   e\u001B[38;5;241m.\u001B[39mmessage \u001B[38;5;241m+\u001B[39m\u001B[38;5;241m=\u001B[39m (\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124m name: \u001B[39m\u001B[38;5;124m\"\u001B[39m \u001B[38;5;241m+\u001B[39m name \u001B[38;5;28;01mif\u001B[39;00m name \u001B[38;5;129;01mis\u001B[39;00m \u001B[38;5;129;01mnot\u001B[39;00m \u001B[38;5;28;01mNone\u001B[39;00m \u001B[38;5;28;01melse\u001B[39;00m \u001B[38;5;124m\"\u001B[39m\u001B[38;5;124m\"\u001B[39m)\n\u001B[0;32m-> 7164\u001B[0m   \u001B[38;5;28;01mraise\u001B[39;00m core\u001B[38;5;241m.\u001B[39m_status_to_exception(e) \u001B[38;5;28;01mfrom\u001B[39;00m \u001B[38;5;28mNone\u001B[39m\n",
      "\u001B[0;31mUnimplementedError\u001B[0m: File system scheme 's3' not implemented (file: 's3://datasets/logs/imdb/256-20220804-141016/train') [Op:CreateSummaryFileWriter]"
     ]
    }
   ],
   "source": [
    "history = model.fit(\n",
    "    mapped_ds,\n",
    "    epochs=10,\n",
    "    callbacks=[cp_callback, tensorboard_callback],\n",
    "    validation_data=mapped_validation,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Now that we have our model, we want to save it to MinIO:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "model.save(f\"s3://{datasets_bucket}/imdb_sentiment_analysis\") "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Let's take a look"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# list all data in history\n",
    "keys = history.history.keys()\n",
    "# summarize history for accuracy\n",
    "plt.plot(history.history[\"accuracy\"])\n",
    "if \"val_accuracy\" in keys:\n",
    "    plt.plot(history.history[\"val_accuracy\"])\n",
    "plt.title(\"model accuracy\")\n",
    "plt.ylabel(\"accuracy\")\n",
    "plt.xlabel(\"epoch\")\n",
    "plt.legend([\"train\", \"validation\"], loc=\"upper left\")\n",
    "plt.show()\n",
    "# summarize history for loss\n",
    "plt.plot(history.history[\"loss\"])\n",
    "if \"val_loss\" in keys:\n",
    "    plt.plot(history.history[\"val_loss\"])\n",
    "plt.title(\"model loss\")\n",
    "plt.ylabel(\"loss\")\n",
    "plt.xlabel(\"epoch\")\n",
    "plt.legend([\"train\", \"validation\"], loc=\"upper left\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Let's test our model and see how it performs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "testing = model.evaluate(testing_mapped_ds)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "This returns 85.63% accuracy, not state of the art, but also not bad for such a simple example."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Let's run `TensorBoard` to explore our models loading the data straight from MinIO"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "```bash\n",
    "AWS_ACCESS_KEY_ID=minioadmin AWS_SECRET_ACCESS_KEY=minioadmin AWS_REGION=us-east-1 S3_ENDPOINT=localhost:9000 S3_USE_HTTPS=0 S3_VERIFY_SSL=0 tensorboard --logdir s3://datasets/logs\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Then go to `http://localhost:6006` on your browser"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "![TensorBoard](pic3.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "We can play with our model and see if it works"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "samples = [\n",
    "    \"This movie sucks\",\n",
    "    \"This was extremely good, I loved it.\",\n",
    "    \"great acting\",\n",
    "    \"terrible acting\",\n",
    "    \"pure kahoot\",\n",
    "    \"I don't know what's the point of this movie, this movie sucks but the acting is great\",\n",
    "    \"This is not a good movie\",\n",
    "]\n",
    "sample_embedded = embed(samples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "res = model.predict(sample_embedded)\n",
    "for s in range(len(samples)):\n",
    "    if res[s][0] > res[s][1]:\n",
    "        print(f\"{samples[s]} - positive\")\n",
    "    else:\n",
    "        print(f\"{samples[s]} - negative\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Conclusion"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "As demonstrated, you can build large scale AI/ML pipelines that can rely entirely on MinIO. This is a function of both MinIO's performance charateristics but also its ability to seamlessly scale to Petabytes and Exabytes of data. By separating storage and compute, one can build a framework that is not dependant on local resources - allowing you to run them on a container inside Kubernetes. This adds considerable flexibility."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "You can see how TensorFlow was able to load the data as it was needed and no customization was needed at all, it simply worked. Moreover this approach could be quickly extended to training by running [TensorFlow in a distributed](https://www.tensorflow.org/guide/distributed_training) manner. This ensures there is very little data to shuffle over the network between training nodes as MinIO becomes the sole source of that data."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Copyright MinIO 2020"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}