04-data-transformation-tft.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8fd50e71-b400-4fc1-b2b1-b09334c2205e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TensorFlow version: 2.13.1\n",
      "TFX Transform version: 1.14.0\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import tensorflow as tf\n",
    "import apache_beam as beam\n",
    "import tensorflow_transform as tft\n",
    "import tensorflow_transform.beam as tft_beam\n",
    "from tensorflow_transform.tf_metadata import schema_utils\n",
    "from tensorflow_transform.tf_metadata import dataset_metadata\n",
    "from tfx_bsl.public import tfxio\n",
    "from tfx_bsl.coders.example_coder import RecordBatchToExamples\n",
    "\n",
    "import tempfile\n",
    "\n",
    "print(f'TensorFlow version: {tf.__version__}')\n",
    "print(f'TFX Transform version: {tft.__version__}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "aaff845a-436f-4aff-8960-8fe70f103111",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data_file = \"data/train_cleaned.csv\"\n",
    "train_df =  pd.read_csv(train_data_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ba4458bb-5ec8-4289-9d84-26889cd2a611",
   "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>age</th>\n",
       "      <th>workclass</th>\n",
       "      <th>fnlwgt</th>\n",
       "      <th>education</th>\n",
       "      <th>education_num</th>\n",
       "      <th>marital_status</th>\n",
       "      <th>occupation</th>\n",
       "      <th>relationship</th>\n",
       "      <th>race</th>\n",
       "      <th>sex</th>\n",
       "      <th>capital_gain</th>\n",
       "      <th>capital_loss</th>\n",
       "      <th>hours_per_week</th>\n",
       "      <th>native_country</th>\n",
       "      <th>income</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>39</td>\n",
       "      <td>State-gov</td>\n",
       "      <td>77516</td>\n",
       "      <td>Bachelors</td>\n",
       "      <td>13</td>\n",
       "      <td>Never-married</td>\n",
       "      <td>Adm-clerical</td>\n",
       "      <td>Not-in-family</td>\n",
       "      <td>White</td>\n",
       "      <td>Male</td>\n",
       "      <td>2174</td>\n",
       "      <td>0</td>\n",
       "      <td>40</td>\n",
       "      <td>United-States</td>\n",
       "      <td>&lt;=50K</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>50</td>\n",
       "      <td>Self-emp-not-inc</td>\n",
       "      <td>83311</td>\n",
       "      <td>Bachelors</td>\n",
       "      <td>13</td>\n",
       "      <td>Married-civ-spouse</td>\n",
       "      <td>Exec-managerial</td>\n",
       "      <td>Husband</td>\n",
       "      <td>White</td>\n",
       "      <td>Male</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>13</td>\n",
       "      <td>United-States</td>\n",
       "      <td>&lt;=50K</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>38</td>\n",
       "      <td>Private</td>\n",
       "      <td>215646</td>\n",
       "      <td>HS-grad</td>\n",
       "      <td>9</td>\n",
       "      <td>Divorced</td>\n",
       "      <td>Handlers-cleaners</td>\n",
       "      <td>Not-in-family</td>\n",
       "      <td>White</td>\n",
       "      <td>Male</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>40</td>\n",
       "      <td>United-States</td>\n",
       "      <td>&lt;=50K</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>53</td>\n",
       "      <td>Private</td>\n",
       "      <td>234721</td>\n",
       "      <td>11th</td>\n",
       "      <td>7</td>\n",
       "      <td>Married-civ-spouse</td>\n",
       "      <td>Handlers-cleaners</td>\n",
       "      <td>Husband</td>\n",
       "      <td>Black</td>\n",
       "      <td>Male</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>40</td>\n",
       "      <td>United-States</td>\n",
       "      <td>&lt;=50K</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>28</td>\n",
       "      <td>Private</td>\n",
       "      <td>338409</td>\n",
       "      <td>Bachelors</td>\n",
       "      <td>13</td>\n",
       "      <td>Married-civ-spouse</td>\n",
       "      <td>Prof-specialty</td>\n",
       "      <td>Wife</td>\n",
       "      <td>Black</td>\n",
       "      <td>Female</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>40</td>\n",
       "      <td>Cuba</td>\n",
       "      <td>&lt;=50K</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   age         workclass  fnlwgt  education  education_num  \\\n",
       "0   39         State-gov   77516  Bachelors             13   \n",
       "1   50  Self-emp-not-inc   83311  Bachelors             13   \n",
       "2   38           Private  215646    HS-grad              9   \n",
       "3   53           Private  234721       11th              7   \n",
       "4   28           Private  338409  Bachelors             13   \n",
       "\n",
       "       marital_status         occupation   relationship   race     sex  \\\n",
       "0       Never-married       Adm-clerical  Not-in-family  White    Male   \n",
       "1  Married-civ-spouse    Exec-managerial        Husband  White    Male   \n",
       "2            Divorced  Handlers-cleaners  Not-in-family  White    Male   \n",
       "3  Married-civ-spouse  Handlers-cleaners        Husband  Black    Male   \n",
       "4  Married-civ-spouse     Prof-specialty           Wife  Black  Female   \n",
       "\n",
       "   capital_gain  capital_loss  hours_per_week native_country  income  \n",
       "0          2174             0              40  United-States   <=50K  \n",
       "1             0             0              13  United-States   <=50K  \n",
       "2             0             0              40  United-States   <=50K  \n",
       "3             0             0              40  United-States   <=50K  \n",
       "4             0             0              40           Cuba   <=50K  "
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_df.head()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "57f51e3a-16c7-451a-ae0f-8de913bb29e2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'workclass': FixedLenFeature(shape=[], dtype=tf.string, default_value=None), 'education': FixedLenFeature(shape=[], dtype=tf.string, default_value=None), 'marital_status': FixedLenFeature(shape=[], dtype=tf.string, default_value=None), 'occupation': FixedLenFeature(shape=[], dtype=tf.string, default_value=None), 'relationship': FixedLenFeature(shape=[], dtype=tf.string, default_value=None), 'race': FixedLenFeature(shape=[], dtype=tf.string, default_value=None), 'sex': FixedLenFeature(shape=[], dtype=tf.string, default_value=None), 'native_country': FixedLenFeature(shape=[], dtype=tf.string, default_value=None), 'age': FixedLenFeature(shape=[], dtype=tf.float32, default_value=None), 'fnlwgt': FixedLenFeature(shape=[], dtype=tf.float32, default_value=None), 'capital_gain': FixedLenFeature(shape=[], dtype=tf.float32, default_value=None), 'capital_loss': FixedLenFeature(shape=[], dtype=tf.float32, default_value=None), 'hours_per_week': FixedLenFeature(shape=[], dtype=tf.float32, default_value=None), 'education_num': FixedLenFeature(shape=[], dtype=tf.float32, default_value=None), 'income': FixedLenFeature(shape=[], dtype=tf.string, default_value=None)}\n"
     ]
    }
   ],
   "source": [
    "CATEGORICAL_FEATURE_KEYS = [\n",
    "    'workclass',\n",
    "    'education',\n",
    "    'marital_status',\n",
    "    'occupation',\n",
    "    'relationship',\n",
    "    'race',\n",
    "    'sex',\n",
    "    'native_country',\n",
    "]\n",
    "\n",
    "NUMERIC_FEATURE_KEYS = [\n",
    "    'age',\n",
    "    'fnlwgt',\n",
    "    'capital_gain',\n",
    "    'capital_loss',\n",
    "    'hours_per_week',\n",
    "    'education_num',\n",
    "]\n",
    "\n",
    "LABEL_KEY = 'income'\n",
    "\n",
    "\n",
    "RAW_DATA_FEATURE_SPEC = dict(\n",
    "    [(name, tf.io.FixedLenFeature([], tf.string))\n",
    "     for name in CATEGORICAL_FEATURE_KEYS] +\n",
    "    [(name, tf.io.FixedLenFeature([], tf.float32))\n",
    "     for name in NUMERIC_FEATURE_KEYS] +\n",
    "    [(LABEL_KEY, tf.io.FixedLenFeature([], tf.string))]\n",
    ")\n",
    "print(RAW_DATA_FEATURE_SPEC)\n",
    "SCHEMA = dataset_metadata.DatasetMetadata(\n",
    "    schema_utils.schema_from_feature_spec(RAW_DATA_FEATURE_SPEC)\n",
    ").schema\n",
    "\n",
    "raw_data_metadata = dataset_metadata.DatasetMetadata(\n",
    "    schema_utils.schema_from_feature_spec(RAW_DATA_FEATURE_SPEC)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "57ef576e-149b-456e-a28a-83664482c9e2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "feature {\n",
      "  name: \"age\"\n",
      "  type: FLOAT\n",
      "  presence {\n",
      "    min_fraction: 1.0\n",
      "  }\n",
      "  shape {\n",
      "  }\n",
      "}\n",
      "feature {\n",
      "  name: \"capital_gain\"\n",
      "  type: FLOAT\n",
      "  presence {\n",
      "    min_fraction: 1.0\n",
      "  }\n",
      "  shape {\n",
      "  }\n",
      "}\n",
      "feature {\n",
      "  name: \"capital_loss\"\n",
      "  type: FLOAT\n",
      "  presence {\n",
      "    min_fraction: 1.0\n",
      "  }\n",
      "  shape {\n",
      "  }\n",
      "}\n",
      "feature {\n",
      "  name: \"education\"\n",
      "  type: BYTES\n",
      "  presence {\n",
      "    min_fraction: 1.0\n",
      "  }\n",
      "  shape {\n",
      "  }\n",
      "}\n",
      "feature {\n",
      "  name: \"education_num\"\n",
      "  type: FLOAT\n",
      "  presence {\n",
      "    min_fraction: 1.0\n",
      "  }\n",
      "  shape {\n",
      "  }\n",
      "}\n",
      "feature {\n",
      "  name: \"fnlwgt\"\n",
      "  type: FLOAT\n",
      "  presence {\n",
      "    min_fraction: 1.0\n",
      "  }\n",
      "  shape {\n",
      "  }\n",
      "}\n",
      "feature {\n",
      "  name: \"hours_per_week\"\n",
      "  type: FLOAT\n",
      "  presence {\n",
      "    min_fraction: 1.0\n",
      "  }\n",
      "  shape {\n",
      "  }\n",
      "}\n",
      "feature {\n",
      "  name: \"income\"\n",
      "  type: BYTES\n",
      "  presence {\n",
      "    min_fraction: 1.0\n",
      "  }\n",
      "  shape {\n",
      "  }\n",
      "}\n",
      "feature {\n",
      "  name: \"marital_status\"\n",
      "  type: BYTES\n",
      "  presence {\n",
      "    min_fraction: 1.0\n",
      "  }\n",
      "  shape {\n",
      "  }\n",
      "}\n",
      "feature {\n",
      "  name: \"native_country\"\n",
      "  type: BYTES\n",
      "  presence {\n",
      "    min_fraction: 1.0\n",
      "  }\n",
      "  shape {\n",
      "  }\n",
      "}\n",
      "feature {\n",
      "  name: \"occupation\"\n",
      "  type: BYTES\n",
      "  presence {\n",
      "    min_fraction: 1.0\n",
      "  }\n",
      "  shape {\n",
      "  }\n",
      "}\n",
      "feature {\n",
      "  name: \"race\"\n",
      "  type: BYTES\n",
      "  presence {\n",
      "    min_fraction: 1.0\n",
      "  }\n",
      "  shape {\n",
      "  }\n",
      "}\n",
      "feature {\n",
      "  name: \"relationship\"\n",
      "  type: BYTES\n",
      "  presence {\n",
      "    min_fraction: 1.0\n",
      "  }\n",
      "  shape {\n",
      "  }\n",
      "}\n",
      "feature {\n",
      "  name: \"sex\"\n",
      "  type: BYTES\n",
      "  presence {\n",
      "    min_fraction: 1.0\n",
      "  }\n",
      "  shape {\n",
      "  }\n",
      "}\n",
      "feature {\n",
      "  name: \"workclass\"\n",
      "  type: BYTES\n",
      "  presence {\n",
      "    min_fraction: 1.0\n",
      "  }\n",
      "  shape {\n",
      "  }\n",
      "}\n",
      "\n",
      "\n",
      "\n",
      "feature {\n",
      "  name: \"age\"\n",
      "  type: FLOAT\n",
      "  presence {\n",
      "    min_fraction: 1.0\n",
      "  }\n",
      "  shape {\n",
      "  }\n",
      "}\n",
      "feature {\n",
      "  name: \"capital_gain\"\n",
      "  type: FLOAT\n",
      "  presence {\n",
      "    min_fraction: 1.0\n",
      "  }\n",
      "  shape {\n",
      "  }\n",
      "}\n",
      "feature {\n",
      "  name: \"capital_loss\"\n",
      "  type: FLOAT\n",
      "  presence {\n",
      "    min_fraction: 1.0\n",
      "  }\n",
      "  shape {\n",
      "  }\n",
      "}\n",
      "feature {\n",
      "  name: \"education\"\n",
      "  type: BYTES\n",
      "  presence {\n",
      "    min_fraction: 1.0\n",
      "  }\n",
      "  shape {\n",
      "  }\n",
      "}\n",
      "feature {\n",
      "  name: \"education_num\"\n",
      "  type: FLOAT\n",
      "  presence {\n",
      "    min_fraction: 1.0\n",
      "  }\n",
      "  shape {\n",
      "  }\n",
      "}\n",
      "feature {\n",
      "  name: \"fnlwgt\"\n",
      "  type: FLOAT\n",
      "  presence {\n",
      "    min_fraction: 1.0\n",
      "  }\n",
      "  shape {\n",
      "  }\n",
      "}\n",
      "feature {\n",
      "  name: \"hours_per_week\"\n",
      "  type: FLOAT\n",
      "  presence {\n",
      "    min_fraction: 1.0\n",
      "  }\n",
      "  shape {\n",
      "  }\n",
      "}\n",
      "feature {\n",
      "  name: \"income\"\n",
      "  type: BYTES\n",
      "  presence {\n",
      "    min_fraction: 1.0\n",
      "  }\n",
      "  shape {\n",
      "  }\n",
      "}\n",
      "feature {\n",
      "  name: \"marital_status\"\n",
      "  type: BYTES\n",
      "  presence {\n",
      "    min_fraction: 1.0\n",
      "  }\n",
      "  shape {\n",
      "  }\n",
      "}\n",
      "feature {\n",
      "  name: \"native_country\"\n",
      "  type: BYTES\n",
      "  presence {\n",
      "    min_fraction: 1.0\n",
      "  }\n",
      "  shape {\n",
      "  }\n",
      "}\n",
      "feature {\n",
      "  name: \"occupation\"\n",
      "  type: BYTES\n",
      "  presence {\n",
      "    min_fraction: 1.0\n",
      "  }\n",
      "  shape {\n",
      "  }\n",
      "}\n",
      "feature {\n",
      "  name: \"race\"\n",
      "  type: BYTES\n",
      "  presence {\n",
      "    min_fraction: 1.0\n",
      "  }\n",
      "  shape {\n",
      "  }\n",
      "}\n",
      "feature {\n",
      "  name: \"relationship\"\n",
      "  type: BYTES\n",
      "  presence {\n",
      "    min_fraction: 1.0\n",
      "  }\n",
      "  shape {\n",
      "  }\n",
      "}\n",
      "feature {\n",
      "  name: \"sex\"\n",
      "  type: BYTES\n",
      "  presence {\n",
      "    min_fraction: 1.0\n",
      "  }\n",
      "  shape {\n",
      "  }\n",
      "}\n",
      "feature {\n",
      "  name: \"workclass\"\n",
      "  type: BYTES\n",
      "  presence {\n",
      "    min_fraction: 1.0\n",
      "  }\n",
      "  shape {\n",
      "  }\n",
      "}\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(raw_data_metadata._schema)\n",
    "print(\"\\n\")\n",
    "print(SCHEMA)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6e1769f7-d42f-4b84-bb8c-c9b01fb2155e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensorflow_metadata.proto.v0.schema_pb2.Schema"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(SCHEMA)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "4a9c1c7b-8d04-4f3b-a18c-a9c1f6f8b8b8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['age',\n",
       " 'workclass',\n",
       " 'fnlwgt',\n",
       " 'education',\n",
       " 'education_num',\n",
       " 'marital_status',\n",
       " 'occupation',\n",
       " 'relationship',\n",
       " 'race',\n",
       " 'sex',\n",
       " 'capital_gain',\n",
       " 'capital_loss',\n",
       " 'hours_per_week',\n",
       " 'native_country',\n",
       " 'income']"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ORDERED_CSV_COLUMNS = list(train_df.columns)\n",
    "ORDERED_CSV_COLUMNS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "2d0e3fa0-0276-4e9e-b28f-cb7a8140d3bf",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:apache_beam.runners.interactive.interactive_environment:Dependencies required for Interactive Beam PCollection visualization are not available, please use: `pip install apache-beam[interactive]` to install necessary dependencies to enable all data visualization features.\n"
     ]
    },
    {
     "data": {
      "application/javascript": [
       "\n",
       "        if (typeof window.interactive_beam_jquery == 'undefined') {\n",
       "          var jqueryScript = document.createElement('script');\n",
       "          jqueryScript.src = 'https://code.jquery.com/jquery-3.4.1.slim.min.js';\n",
       "          jqueryScript.type = 'text/javascript';\n",
       "          jqueryScript.onload = function() {\n",
       "            var datatableScript = document.createElement('script');\n",
       "            datatableScript.src = 'https://cdn.datatables.net/1.10.20/js/jquery.dataTables.min.js';\n",
       "            datatableScript.type = 'text/javascript';\n",
       "            datatableScript.onload = function() {\n",
       "              window.interactive_beam_jquery = jQuery.noConflict(true);\n",
       "              window.interactive_beam_jquery(document).ready(function($){\n",
       "                \n",
       "              });\n",
       "            }\n",
       "            document.head.appendChild(datatableScript);\n",
       "          };\n",
       "          document.head.appendChild(jqueryScript);\n",
       "        } else {\n",
       "          window.interactive_beam_jquery(document).ready(function($){\n",
       "            \n",
       "          });\n",
       "        }"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Read CSV files and translate to RecordBatch\n",
    "pipeline = beam.Pipeline()\n",
    "csv_tfxio = tfxio.CsvTFXIO(train_data_file,\n",
    "                           telemetry_descriptors=[],\n",
    "                           column_names=ORDERED_CSV_COLUMNS,\n",
    "                           schema=SCHEMA,\n",
    "                          skip_header_lines=1)\n",
    "\n",
    "#pipeline = beam.Pipeline()\n",
    "raw_data = pipeline | 'TFXIORead' >> csv_tfxio.BeamSource()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "57689c10-66dd-4ccf-bf69-f2f726da3e68",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read CSV and perform some clean up\n",
    "#pipeline = beam.Pipeline()\n",
    "\n",
    "#csv_tfxio = tfxio.BeamRecordCsvTFXIO(\n",
    "#    physical_format='text', column_names=ORDERED_CSV_COLUMNS, schema=SCHEMA)\n",
    "\n",
    "#raw_data = (\n",
    "#    pipeline\n",
    "#    | 'ReadTrainData' >> beam.io.ReadFromText(\n",
    "#       train_data_file, coder=beam.coders.BytesCoder())\n",
    "#    | 'FixCommasTrainData' >> beam.Map(\n",
    "#        lambda line: line.replace(b', ', b','))\n",
    "#    | 'DecodeTrainData' >> csv_tfxio.BeamSource())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "164b2842-bc29-4896-a05b-1350e38926dd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<PCollection[[8]: TFXIORead/RawRecordToRecordBatch/CollectRecordBatchTelemetry/ProfileRecordBatches.None] at 0x1df320af280>"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "raw_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "1a1f18a4-954b-4642-a1bb-c23ea812885a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Features with string data types that will be converted to indices\n",
    "CATEGORICAL_FEATURE_KEYS = [\n",
    "    'education', 'marital_status', 'occupation', 'race', 'relationship', 'workclass', 'sex', 'native_country'\n",
    "]\n",
    "\n",
    "# Feature to replace before converting to string data type\n",
    "REPLACE_KEY = 'marital_status'\n",
    "\n",
    "# Numerical features that are marked as continuous\n",
    "NUMERIC_FEATURE_KEYS = ['fnlwgt', \n",
    "                        'education_num',\n",
    "                        'capital_gain', 'capital_loss', 'hours_per_week'\n",
    "                       ]\n",
    "\n",
    "# Feature that can be grouped into buckets\n",
    "BUCKET_FEATURE_KEYS = ['age']\n",
    "\n",
    "# Number of buckets used by tf.transform for encoding each bucket feature.\n",
    "FEATURE_BUCKET_COUNT = {'age': 4}\n",
    "\n",
    "# Feature that the model will predict\n",
    "LABEL_KEY = 'income'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "487da838-0a5f-4c87-8e28-ba86cd6bdacd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Utility function for renaming the feature\n",
    "def transformed_name(key):\n",
    "    return key + '_xf'\n",
    "    \n",
    "def preprocessing_fn(inputs):\n",
    "    \"\"\"Preprocess input columns into transformed columns.\n",
    "    Args: \n",
    "        inputs (): \n",
    "\n",
    "    Return:\n",
    "        dictionary\n",
    "        \n",
    "    \"\"\"\n",
    "    outputs = {}\n",
    "    # Replace and map some features to reduce the number of classes\n",
    "    inputs[REPLACE_KEY] = tf.strings.regex_replace(inputs[REPLACE_KEY], pattern='Married-civ-spouse', rewrite='Married')\n",
    "    inputs[REPLACE_KEY] = tf.strings.regex_replace(inputs[REPLACE_KEY], pattern='Never-married', rewrite='Single')\n",
    "    inputs[REPLACE_KEY] = tf.strings.regex_replace(inputs[REPLACE_KEY], pattern='Widowed', rewrite='Widowed')\n",
    "    inputs[REPLACE_KEY] = tf.strings.regex_replace(inputs[REPLACE_KEY], pattern='Separated', rewrite='Divorced')\n",
    "    inputs[REPLACE_KEY] = tf.strings.regex_replace(inputs[REPLACE_KEY], pattern='Divorced', rewrite='Divorced')\n",
    "    inputs[REPLACE_KEY] = tf.strings.regex_replace(inputs[REPLACE_KEY], pattern='Married-spouse-absent', rewrite='Divorced')\n",
    "    inputs[REPLACE_KEY] = tf.strings.regex_replace(inputs[REPLACE_KEY], pattern='Married-AF-spouse', rewrite='Divorced')\n",
    "\n",
    "    # Map features To reduce get binary classes as the percentage of data is much\n",
    "    # Feature to change to two classes before applying vocab\n",
    "    inputs['workclass'] = tf.where(inputs['workclass'] == 'Private', 'Private', 'Others')\n",
    "    inputs['race'] = tf.where(inputs['race'] ==  'White','White', 'Others')\n",
    "    inputs['native_country'] = tf.where(inputs['native_country'] == 'United-States', 'United-States', 'Others')\n",
    "    \n",
    "    # Scale these fatures to be [0,1]\n",
    "    for key in NUMERIC_FEATURE_KEYS:\n",
    "        outputs[transformed_name(key)] = tft.scale_to_0_1(inputs[key])\n",
    "\n",
    "    for key in CATEGORICAL_FEATURE_KEYS:\n",
    "        outputs[transformed_name(key)] = tft.compute_and_apply_vocabulary(inputs[key])\n",
    "\n",
    "    for key in BUCKET_FEATURE_KEYS:\n",
    "        outputs[transformed_name(key)] = tft.bucketize(inputs[key], FEATURE_BUCKET_COUNT[key])\n",
    "\n",
    "\n",
    "    outputs[transformed_name(LABEL_KEY)] = tft.compute_and_apply_vocabulary(inputs[LABEL_KEY])\n",
    "    \n",
    "    \n",
    "    return outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "43ea77bc-bb46-44b8-a18a-8e30e6b9df01",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pprint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9856e87a-d686-4dca-80cf-708c36b04365",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#train_df.to_dict('records')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "afc3f42c-4247-49c5-bc3a-70fcc5f518d7",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:absl:You are passing instance dicts and DatasetMetadata to TFT which will not provide optimal performance. Consider following the TFT guide to upgrade to the TFXIO format (Apache Arrow RecordBatch).\n",
      "WARNING:absl:You are passing instance dicts and DatasetMetadata to TFT which will not provide optimal performance. Consider following the TFT guide to upgrade to the TFXIO format (Apache Arrow RecordBatch).\n",
      "WARNING:absl:You are outputting instance dicts from `TransformDataset` which will not provide optimal performance. Consider setting  `output_record_batches=True` to upgrade to the TFXIO format (Apache Arrow RecordBatch). Encoding functionality in this module works with both formats.\n",
      "WARNING:apache_beam.options.pipeline_options:Discarding unparseable args: ['C:\\\\Users\\\\Windows 10\\\\.conda\\\\envs\\\\tensorflow_env\\\\lib\\\\site-packages\\\\ipykernel_launcher.py', '-f', 'C:\\\\Users\\\\Windows 10\\\\AppData\\\\Roaming\\\\jupyter\\\\runtime\\\\kernel-8901e1cf-982e-4ce2-ba6e-7251f3daffa1.json']\n"
     ]
    }
   ],
   "source": [
    "\n",
    "# Ignore the warnings\n",
    "tf.get_logger().setLevel('ERROR')\n",
    "\n",
    "# a temporary directory is needed when analyzing the data\n",
    "with tft_beam.Context(temp_dir=tempfile.mkdtemp()):\n",
    "    \n",
    "    # define the pipeline using Apache Beam syntax\n",
    "    transformed_dataset1, transform_fn1 = (\n",
    "        \n",
    "        # analyze and transform the dataset using the preprocessing function\n",
    "        (train_df.to_dict('records'), raw_data_metadata) | tft_beam.AnalyzeAndTransformDataset(\n",
    "            preprocessing_fn)\n",
    "    )\n",
    "\n",
    "# unpack the transformed dataset\n",
    "transformed_data1, transformed_metadata1 = transformed_dataset1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "93c9bade-0109-4a22-bb99-6a6c03464cad",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Raw data:\n",
      "[{'age': 39,\n",
      "  'capital_gain': 2174,\n",
      "  'capital_loss': 0,\n",
      "  'education': 'Bachelors',\n",
      "  'education_num': 13,\n",
      "  'fnlwgt': 77516,\n",
      "  'hours_per_week': 40,\n",
      "  'income': ' <=50K',\n",
      "  'marital_status': 'Never-married',\n",
      "  'native_country': 'United-States',\n",
      "  'occupation': 'Adm-clerical',\n",
      "  'race': 'White',\n",
      "  'relationship': 'Not-in-family',\n",
      "  'sex': 'Male',\n",
      "  'workclass': 'State-gov'},\n",
      " {'age': 50,\n",
      "  'capital_gain': 0,\n",
      "  'capital_loss': 0,\n",
      "  'education': 'Bachelors',\n",
      "  'education_num': 13,\n",
      "  'fnlwgt': 83311,\n",
      "  'hours_per_week': 13,\n",
      "  'income': ' <=50K',\n",
      "  'marital_status': 'Married-civ-spouse',\n",
      "  'native_country': 'United-States',\n",
      "  'occupation': 'Exec-managerial',\n",
      "  'race': 'White',\n",
      "  'relationship': 'Husband',\n",
      "  'sex': 'Male',\n",
      "  'workclass': 'Self-emp-not-inc'},\n",
      " {'age': 38,\n",
      "  'capital_gain': 0,\n",
      "  'capital_loss': 0,\n",
      "  'education': 'HS-grad',\n",
      "  'education_num': 9,\n",
      "  'fnlwgt': 215646,\n",
      "  'hours_per_week': 40,\n",
      "  'income': ' <=50K',\n",
      "  'marital_status': 'Divorced',\n",
      "  'native_country': 'United-States',\n",
      "  'occupation': 'Handlers-cleaners',\n",
      "  'race': 'White',\n",
      "  'relationship': 'Not-in-family',\n",
      "  'sex': 'Male',\n",
      "  'workclass': 'Private'},\n",
      " {'age': 53,\n",
      "  'capital_gain': 0,\n",
      "  'capital_loss': 0,\n",
      "  'education': '11th',\n",
      "  'education_num': 7,\n",
      "  'fnlwgt': 234721,\n",
      "  'hours_per_week': 40,\n",
      "  'income': ' <=50K',\n",
      "  'marital_status': 'Married-civ-spouse',\n",
      "  'native_country': 'United-States',\n",
      "  'occupation': 'Handlers-cleaners',\n",
      "  'race': 'Black',\n",
      "  'relationship': 'Husband',\n",
      "  'sex': 'Male',\n",
      "  'workclass': 'Private'},\n",
      " {'age': 28,\n",
      "  'capital_gain': 0,\n",
      "  'capital_loss': 0,\n",
      "  'education': 'Bachelors',\n",
      "  'education_num': 13,\n",
      "  'fnlwgt': 338409,\n",
      "  'hours_per_week': 40,\n",
      "  'income': ' <=50K',\n",
      "  'marital_status': 'Married-civ-spouse',\n",
      "  'native_country': 'Cuba',\n",
      "  'occupation': 'Prof-specialty',\n",
      "  'race': 'Black',\n",
      "  'relationship': 'Wife',\n",
      "  'sex': 'Female',\n",
      "  'workclass': 'Private'},\n",
      " {'age': 37,\n",
      "  'capital_gain': 0,\n",
      "  'capital_loss': 0,\n",
      "  'education': 'Masters',\n",
      "  'education_num': 14,\n",
      "  'fnlwgt': 284582,\n",
      "  'hours_per_week': 40,\n",
      "  'income': ' <=50K',\n",
      "  'marital_status': 'Married-civ-spouse',\n",
      "  'native_country': 'United-States',\n",
      "  'occupation': 'Exec-managerial',\n",
      "  'race': 'White',\n",
      "  'relationship': 'Wife',\n",
      "  'sex': 'Female',\n",
      "  'workclass': 'Private'},\n",
      " {'age': 49,\n",
      "  'capital_gain': 0,\n",
      "  'capital_loss': 0,\n",
      "  'education': '9th',\n",
      "  'education_num': 5,\n",
      "  'fnlwgt': 160187,\n",
      "  'hours_per_week': 16,\n",
      "  'income': ' <=50K',\n",
      "  'marital_status': 'Married-spouse-absent',\n",
      "  'native_country': 'Jamaica',\n",
      "  'occupation': 'Other-service',\n",
      "  'race': 'Black',\n",
      "  'relationship': 'Not-in-family',\n",
      "  'sex': 'Female',\n",
      "  'workclass': 'Private'},\n",
      " {'age': 52,\n",
      "  'capital_gain': 0,\n",
      "  'capital_loss': 0,\n",
      "  'education': 'HS-grad',\n",
      "  'education_num': 9,\n",
      "  'fnlwgt': 209642,\n",
      "  'hours_per_week': 45,\n",
      "  'income': ' >50K',\n",
      "  'marital_status': 'Married-civ-spouse',\n",
      "  'native_country': 'United-States',\n",
      "  'occupation': 'Exec-managerial',\n",
      "  'race': 'White',\n",
      "  'relationship': 'Husband',\n",
      "  'sex': 'Male',\n",
      "  'workclass': 'Self-emp-not-inc'},\n",
      " {'age': 31,\n",
      "  'capital_gain': 14084,\n",
      "  'capital_loss': 0,\n",
      "  'education': 'Masters',\n",
      "  'education_num': 14,\n",
      "  'fnlwgt': 45781,\n",
      "  'hours_per_week': 50,\n",
      "  'income': ' >50K',\n",
      "  'marital_status': 'Never-married',\n",
      "  'native_country': 'United-States',\n",
      "  'occupation': 'Prof-specialty',\n",
      "  'race': 'White',\n",
      "  'relationship': 'Not-in-family',\n",
      "  'sex': 'Female',\n",
      "  'workclass': 'Private'},\n",
      " {'age': 42,\n",
      "  'capital_gain': 5178,\n",
      "  'capital_loss': 0,\n",
      "  'education': 'Bachelors',\n",
      "  'education_num': 13,\n",
      "  'fnlwgt': 159449,\n",
      "  'hours_per_week': 40,\n",
      "  'income': ' >50K',\n",
      "  'marital_status': 'Married-civ-spouse',\n",
      "  'native_country': 'United-States',\n",
      "  'occupation': 'Exec-managerial',\n",
      "  'race': 'White',\n",
      "  'relationship': 'Husband',\n",
      "  'sex': 'Male',\n",
      "  'workclass': 'Private'}]\n",
      "\n",
      "Transformed data:\n",
      "[{'age_xf': 2,\n",
      "  'capital_gain_xf': 0.02174021676182747,\n",
      "  'capital_loss_xf': 0.0,\n",
      "  'education_num_xf': 0.800000011920929,\n",
      "  'education_xf': 2,\n",
      "  'fnlwgt_xf': 0.044301897287368774,\n",
      "  'hours_per_week_xf': 0.3979591727256775,\n",
      "  'income_xf': 0,\n",
      "  'marital_status_xf': 1,\n",
      "  'native_country_xf': 0,\n",
      "  'occupation_xf': 3,\n",
      "  'race_xf': 0,\n",
      "  'relationship_xf': 1,\n",
      "  'sex_xf': 0,\n",
      "  'workclass_xf': 1},\n",
      " {'age_xf': 3,\n",
      "  'capital_gain_xf': 0.0,\n",
      "  'capital_loss_xf': 0.0,\n",
      "  'education_num_xf': 0.800000011920929,\n",
      "  'education_xf': 2,\n",
      "  'fnlwgt_xf': 0.048237595707178116,\n",
      "  'hours_per_week_xf': 0.12244898080825806,\n",
      "  'income_xf': 0,\n",
      "  'marital_status_xf': 0,\n",
      "  'native_country_xf': 0,\n",
      "  'occupation_xf': 2,\n",
      "  'race_xf': 0,\n",
      "  'relationship_xf': 0,\n",
      "  'sex_xf': 0,\n",
      "  'workclass_xf': 1},\n",
      " {'age_xf': 2,\n",
      "  'capital_gain_xf': 0.0,\n",
      "  'capital_loss_xf': 0.0,\n",
      "  'education_num_xf': 0.5333333611488342,\n",
      "  'education_xf': 0,\n",
      "  'fnlwgt_xf': 0.13811343908309937,\n",
      "  'hours_per_week_xf': 0.3979591727256775,\n",
      "  'income_xf': 0,\n",
      "  'marital_status_xf': 2,\n",
      "  'native_country_xf': 0,\n",
      "  'occupation_xf': 9,\n",
      "  'race_xf': 0,\n",
      "  'relationship_xf': 1,\n",
      "  'sex_xf': 0,\n",
      "  'workclass_xf': 0},\n",
      " {'age_xf': 3,\n",
      "  'capital_gain_xf': 0.0,\n",
      "  'capital_loss_xf': 0.0,\n",
      "  'education_num_xf': 0.4000000059604645,\n",
      "  'education_xf': 5,\n",
      "  'fnlwgt_xf': 0.151068314909935,\n",
      "  'hours_per_week_xf': 0.3979591727256775,\n",
      "  'income_xf': 0,\n",
      "  'marital_status_xf': 0,\n",
      "  'native_country_xf': 0,\n",
      "  'occupation_xf': 9,\n",
      "  'race_xf': 1,\n",
      "  'relationship_xf': 0,\n",
      "  'sex_xf': 0,\n",
      "  'workclass_xf': 0},\n",
      " {'age_xf': 1,\n",
      "  'capital_gain_xf': 0.0,\n",
      "  'capital_loss_xf': 0.0,\n",
      "  'education_num_xf': 0.800000011920929,\n",
      "  'education_xf': 2,\n",
      "  'fnlwgt_xf': 0.22148843109607697,\n",
      "  'hours_per_week_xf': 0.3979591727256775,\n",
      "  'income_xf': 0,\n",
      "  'marital_status_xf': 0,\n",
      "  'native_country_xf': 1,\n",
      "  'occupation_xf': 0,\n",
      "  'race_xf': 1,\n",
      "  'relationship_xf': 4,\n",
      "  'sex_xf': 1,\n",
      "  'workclass_xf': 0},\n",
      " {'age_xf': 2,\n",
      "  'capital_gain_xf': 0.0,\n",
      "  'capital_loss_xf': 0.0,\n",
      "  'education_num_xf': 0.8666666746139526,\n",
      "  'education_xf': 3,\n",
      "  'fnlwgt_xf': 0.18493160605430603,\n",
      "  'hours_per_week_xf': 0.3979591727256775,\n",
      "  'income_xf': 0,\n",
      "  'marital_status_xf': 0,\n",
      "  'native_country_xf': 0,\n",
      "  'occupation_xf': 2,\n",
      "  'race_xf': 0,\n",
      "  'relationship_xf': 4,\n",
      "  'sex_xf': 1,\n",
      "  'workclass_xf': 0},\n",
      " {'age_xf': 3,\n",
      "  'capital_gain_xf': 0.0,\n",
      "  'capital_loss_xf': 0.0,\n",
      "  'education_num_xf': 0.2666666805744171,\n",
      "  'education_xf': 10,\n",
      "  'fnlwgt_xf': 0.10044824331998825,\n",
      "  'hours_per_week_xf': 0.15306122601032257,\n",
      "  'income_xf': 0,\n",
      "  'marital_status_xf': 2,\n",
      "  'native_country_xf': 1,\n",
      "  'occupation_xf': 5,\n",
      "  'race_xf': 1,\n",
      "  'relationship_xf': 1,\n",
      "  'sex_xf': 1,\n",
      "  'workclass_xf': 0},\n",
      " {'age_xf': 3,\n",
      "  'capital_gain_xf': 0.0,\n",
      "  'capital_loss_xf': 0.0,\n",
      "  'education_num_xf': 0.5333333611488342,\n",
      "  'education_xf': 0,\n",
      "  'fnlwgt_xf': 0.13403581082820892,\n",
      "  'hours_per_week_xf': 0.44897958636283875,\n",
      "  'income_xf': 1,\n",
      "  'marital_status_xf': 0,\n",
      "  'native_country_xf': 0,\n",
      "  'occupation_xf': 2,\n",
      "  'race_xf': 0,\n",
      "  'relationship_xf': 0,\n",
      "  'sex_xf': 0,\n",
      "  'workclass_xf': 1},\n",
      " {'age_xf': 1,\n",
      "  'capital_gain_xf': 0.14084140956401825,\n",
      "  'capital_loss_xf': 0.0,\n",
      "  'education_num_xf': 0.8666666746139526,\n",
      "  'education_xf': 3,\n",
      "  'fnlwgt_xf': 0.02274894341826439,\n",
      "  'hours_per_week_xf': 0.5,\n",
      "  'income_xf': 1,\n",
      "  'marital_status_xf': 1,\n",
      "  'native_country_xf': 0,\n",
      "  'occupation_xf': 0,\n",
      "  'race_xf': 0,\n",
      "  'relationship_xf': 1,\n",
      "  'sex_xf': 1,\n",
      "  'workclass_xf': 0},\n",
      " {'age_xf': 2,\n",
      "  'capital_gain_xf': 0.051780518144369125,\n",
      "  'capital_loss_xf': 0.0,\n",
      "  'education_num_xf': 0.800000011920929,\n",
      "  'education_xf': 2,\n",
      "  'fnlwgt_xf': 0.09994702786207199,\n",
      "  'hours_per_week_xf': 0.3979591727256775,\n",
      "  'income_xf': 1,\n",
      "  'marital_status_xf': 0,\n",
      "  'native_country_xf': 0,\n",
      "  'occupation_xf': 2,\n",
      "  'race_xf': 0,\n",
      "  'relationship_xf': 0,\n",
      "  'sex_xf': 0,\n",
      "  'workclass_xf': 0}]\n"
     ]
    }
   ],
   "source": [
    "# print the results\n",
    "print('\\nRaw data:\\n{}\\n'.format(pprint.pformat(train_df.to_dict('records')[:10])))\n",
    "print('Transformed data:\\n{}'.format(pprint.pformat(transformed_data1[:10])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9d3069f-be93-4a9c-b141-51f3c19cddbb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create constant module\n",
    "census_constant_module = \"transformation/income_constant.py\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13b0430a-c467-421b-b323-28c6cc33dc65",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile {census_constant_module}\n",
    "# Features with string data types that will be converted to indices\n",
    "CATEGORICAL_FEATURE_KEYS = [\n",
    "    'education', 'marital_status', 'occupation', 'race', 'relationship', 'workclass', 'sex', 'native_country'\n",
    "]\n",
    "\n",
    "# Numerical features that are marked as continuous\n",
    "NUMERIC_FEATURE_KEYS = ['fnlwgt', \n",
    "                        #'education_num',\n",
    "                        'capital_gain', 'capital_loss', 'hours_per_week'\n",
    "                       ]\n",
    "\n",
    "# Feature that can be grouped into buckets\n",
    "BUCKET_FEATURE_KEYS = ['age']\n",
    "\n",
    "# Number of buckets used by tf.transform for encoding each bucket feature.\n",
    "FEATURE_BUCKET_COUNT = {'age': 4}\n",
    "\n",
    "# Feature that the model will predict\n",
    "LABEL_KEY = 'income'\n",
    "\n",
    "# Utility function for renaming the feature\n",
    "def transformed_name(key):\n",
    "    return key + '_xf'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f9415e6-3c3c-45ab-844a-38700c7d494a",
   "metadata": {},
   "outputs": [],
   "source": [
    "census_transform_module_file = 'transformation/income_transform.py'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88b03453-7703-432a-83ef-51c3e34ae2e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile {census_transform_module_file}\n",
    "import income_constant\n",
    "import tensorflow as tf\n",
    "import tensorflow_transform as tft\n",
    "\n",
    "# Unpack the contents of the constants module\n",
    "_NUMERIC_FEATURE_KEYS = income_constant.NUMERIC_FEATURE_KEYS\n",
    "_CATEGORICAL_FEATURE_KEYS = income_constant.CATEGORICAL_FEATURE_KEYS\n",
    "_BUCKET_FEATURE_KEYS = income_constant.BUCKET_FEATURE_KEYS\n",
    "_FEATURE_BUCKET_COUNT = income_constant.FEATURE_BUCKET_COUNT\n",
    "_LABEL_KEY = income_constant.LABEL_KEY\n",
    "_transformed_name = income_constant.transformed_name\n",
    "\n",
    "def preprocessing_fn(inputs):\n",
    "    \"\"\"Preprocess input columns into transformed columns.\n",
    "    Args: \n",
    "        inputs (): \n",
    "\n",
    "    Return:\n",
    "        dictionary\n",
    "        \n",
    "    \"\"\"\n",
    "    outputs = {}\n",
    "    # Scale these fatures to be [0,1]\n",
    "    for key in _NUMERIC_FEATURE_KEYS:\n",
    "        outputs[_transformed_name(key)] = tft.scale_to_0_1(inputs[key])\n",
    "\n",
    "    for key in _CATEGORICAL_FEATURE_KEYS:\n",
    "        outputs[_transformed_name(key)] = tft.compute_and_apply_vocabulary(inputs[key])\n",
    "\n",
    "    for key in _BUCKET_FEATURE_KEYS:\n",
    "        outputs[_transformed_name(key)] = tft.bucketize(inputs[key], _FEATURE_BUCKET_COUNT[key])\n",
    "\n",
    "\n",
    "    outputs[_transformed_name(_LABEL_KEY)] = tft.compute_and_apply_vocabulary(inputs[_LABEL_KEY])\n",
    "    \n",
    "    \n",
    "    return outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "6e0fc7d5-1834-411a-98cb-3cd213e2cb39",
   "metadata": {},
   "outputs": [],
   "source": [
    "raw_dataset = (raw_data, csv_tfxio.TensorAdapterConfig())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "fbb4fb30-bdec-4cbd-8a71-b664a3825547",
   "metadata": {},
   "outputs": [],
   "source": [
    "#from transformation.income_transform import preprocessing_fn\n",
    "# faster\n",
    "# ignor Warnings\n",
    "\n",
    "tf.get_logger().setLevel('ERROR')\n",
    "\n",
    "# create a temp directory to analyze the data\n",
    "with tft_beam.Context(temp_dir=tempfile.mkdtemp()):\n",
    "    transformed_dataset2, transform_function2 = (\n",
    "        raw_dataset | tft_beam.AnalyzeAndTransformDataset(preprocessing_fn=preprocessing_fn, \n",
    "                                                          output_record_batches=True\n",
    "                                                         )\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "110d202e-db0b-4b41-895e-aaca4f677323",
   "metadata": {},
   "outputs": [],
   "source": [
    "#output_dir = tempfile.mkdtemp()\n",
    "transformed_data2, transformed_metadata2 = transformed_dataset2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "8682b849-4d0d-460b-9d83-8df687446221",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<PCollection[[18]: AnalyzeAndTransformDataset/TransformDataset/ConvertToRecordBatch.None] at 0x1df43366980>"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "transformed_data2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "6edcaf0a-fb22-48c8-8c3e-875e818c7c87",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.makedirs(\"transformation/transformed_data\", exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "83122e41-3ec7-4286-8fc9-a48f3b57e2c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "output_dir = \"transformation/transformed_data\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea29f557-6b40-472a-af94-b5298de54bbb",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ef4ec7b-1483-4e34-be28-9352d104feb7",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "e68e8e09-273b-413d-b2a3-4b62dc9dec57",
   "metadata": {},
   "outputs": [],
   "source": [
    "#transformed_data, _ = transformed_data\n",
    "#import os\n",
    "_ = (\n",
    "    transformed_data2\n",
    "    | 'EncodeTrainData' >>\n",
    "    beam.FlatMapTuple(lambda batch, transformed_metadata2: RecordBatchToExamples(batch))\n",
    "    | 'WriteTrainData' >> beam.io.WriteToTFRecord(\n",
    "        os.path.join(output_dir , 'transformed.tfrecord')))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "90705b13-8e35-4509-9e91-179bbbcd2f3b",
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = (\n",
    "    transform_function2\n",
    "    | 'WriteTransformFn' >> tft_beam.WriteTransformFn(output_dir))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "dde9939c-ab4f-473e-a4ba-3303bdec65c8",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:apache_beam.io.tfrecordio:Couldn't find python-snappy so the implementation of _TFRecordUtil._masked_crc32c is not as fast as it could be.\n"
     ]
    }
   ],
   "source": [
    "result = pipeline.run().wait_until_finish()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96c4300b-296e-404d-b0c8-6fbda52eb3b6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5316a6e2-414d-4295-8a29-e8ba8cbe451d",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}