workflow_two_datasets.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example with TWO data sources\n",
    "The second data source is `Health` data, which contains rows with missing values"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. Define the `Population`\n",
    "The `Population` object/class in `src\\dataloaders\\populations\\base.py` is used to store information about the cohort we want to use for our dataset. Based on the populations, we further create a dataset that contains specific cohorts of people. You can define multiple populations and create various datasets based on various specifications of users.\n",
    "\n",
    "#### How to use the `Population`?\n",
    "Preferably, you would define a new populations object for each new cohort.\n",
    "\n",
    "For my dummy cohort, I defined a `UserPopulation` object in `src\\dataloaders\\populations\\base.py` that specifies what information (and which users) to keep in the dataset. Here is what happens inside of the `UserPopulation`:\n",
    "1. it takes the **raw** input file containing some information about the users (for example, some kind of demographic information as such as birthday, sex),\n",
    "2. preprocess and filter users based on some criteria (for example, you want to exclude people below certain age),\n",
    "3. create data splits (train, val, test).\n",
    "\n",
    "The `UserPopulation` object runs all these processes and saves outputs in `data\\processed\\populations` folder. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the Population object\n",
    "from src.dataloaders.populations.users import UserPopulation\n",
    "users = UserPopulation()\n",
    "## run the preprocessing part\n",
    "users.population()\n",
    "## create datasplits \n",
    "users.data_split()\n",
    "## You can also just run the prepare function to do both of the above steps\n",
    "users.prepare() "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After you run these commands for the first time, you will see that the results are saved in `data\\procesed\\populations`. \n",
    "Next time you the same functions, instead of calculating everything, the object would read the data from the `data\\processed\\populations` folder. **This is important for very big datasets**.\n",
    "\n",
    "If you want to redo the calculations, you need to empty the `data\\processed\\populations` folder. It also applies to cases, when you change the code of the `src\\dataloaders\\populations\\users.py`, as the Population object saves `arguments` so it can validate that you call a specific version of the Population object. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. Define the `TokenSource`\n",
    "The `TokenSource` class/object in `src\\data\\sources\\base.py` specifies how to process a specific *source* of data (for example, the *Synthetic Labor* dataset). You have to define a `TokenSource` object for each new data source. For example, for the *Synthetic Labor Dataset*, I have defined a `SyntheticLaborSource` class that is specifically tailored for the *Synthetic Labor Dataset*. \n",
    "\n",
    "If any of the variables is continious and you need to bin it, you can use the `Binned` class (see example in `src\\dataloaders\\sources\\synth_labor.py`).\n",
    "\n",
    "All in all, it makes it easier to process data from different data streams (or datasets) and specify how to convert each variable to tokens.\n",
    "\n",
    "#### How to use `TokenSource`?\n",
    "For example, `SyntheticLaborSource` (in `src\\dataloaders\\sources\\synth_labor.py`) specifies how to tokenize the `data\\rawdata\\synth_labor.csv`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from src.dataloaders.sources.synth_labor import SyntheticLaborSource"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "## First we initialize the TokenSource instance.\n",
    "synth_labor = SyntheticLaborSource()\n",
    "## process the raw file (and maybe do some preprocessing)\n",
    "#synth_labor.parsed()\n",
    "## index the files\n",
    "#synth_labor.indexed()\n",
    "## tokenize the files\n",
    "#synth_labor.tokenized()\n",
    "### Or use the prepare function to do all of the above\n",
    "synth_labor.prepare()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After you run these commands for the first time, you will see that the results are saved in `data\\procesed\\sources`. \n",
    "Next time you the same functions, instead of calculating everything, the object would read the data from the `data\\processed\\sources` folder. **This is important for very big datasets**.\n",
    "\n",
    "If you want to redo the calculations, you need to empty the *corresponding* file in the `data\\processed\\sources` folder. It also applies to cases, when you change the code of the `src\\dataloaders\\sources\\synth_labor.py`, as the `SyntheticLaborSource` object saves `arguments` so it can validate that you call it on future runs. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Add `Health` data\n",
    "`SyntheticHealthSource` (in `src\\dataloaders\\sources\\synth_heakth.py`) specifies how to tokenize the `data\\rawdata\\synth_health.csv`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from src.dataloaders.sources.synth_health import SyntheticHealthSource"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/carlomarx/.local/share/virtualenvs/life2vec-light-Ez8u7ZRp-python/lib/python3.11/site-packages/dask/dataframe/io/parquet/core.py:1215: FutureWarning: engine='pyarrow-dataset' has been deprecated. Please use engine='pyarrow' instead.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "## First we initialize the TokenSource instance.\n",
    "synth_health = SyntheticHealthSource()\n",
    "## process the raw file (and maybe do some preprocessing)\n",
    "#synth_health.parsed()\n",
    "## index the files\n",
    "#synth_health.indexed()\n",
    "## tokenize the files\n",
    "#synth_health.tokenized()\n",
    "### Or use the prepare function to do all of the above\n",
    "synth_health.prepare()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3. Define Corpus, Vocabulary and Task\n",
    "### 3.1. Let's assemble a corpus\n",
    "Now we can reuse both `Populations` and `Source` objects to actually create a dataset (based on the specification in both `UserPopulation` and `SyntheticLaborSource`). \n",
    "\n",
    "The `Corpus` object in `src\\dataloaders\\datamodule.py` takes all the specifications and creates a dataset (i.e. creates sentences out of tabular records). It also saves data in the corresponding data splits."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from src.dataloaders.datamodule import Corpus"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/carlomarx/.local/share/virtualenvs/life2vec-light-Ez8u7ZRp-python/lib/python3.11/site-packages/dask/dataframe/io/parquet/core.py:1215: FutureWarning: engine='pyarrow-dataset' has been deprecated. Please use engine='pyarrow' instead.\n",
      "  warnings.warn(\n",
      "/Users/carlomarx/.local/share/virtualenvs/life2vec-light-Ez8u7ZRp-python/lib/python3.11/site-packages/dask/dataframe/io/parquet/core.py:1215: FutureWarning: engine='pyarrow-dataset' has been deprecated. Please use engine='pyarrow' instead.\n",
      "  warnings.warn(\n",
      "/Users/carlomarx/.local/share/virtualenvs/life2vec-light-Ez8u7ZRp-python/lib/python3.11/site-packages/dask/dataframe/io/parquet/core.py:1215: FutureWarning: engine='pyarrow-dataset' has been deprecated. Please use engine='pyarrow' instead.\n",
      "  warnings.warn(\n",
      "/Users/carlomarx/.local/share/virtualenvs/life2vec-light-Ez8u7ZRp-python/lib/python3.11/site-packages/dask/dataframe/io/parquet/core.py:1215: FutureWarning: engine='pyarrow-dataset' has been deprecated. Please use engine='pyarrow' instead.\n",
      "  warnings.warn(\n",
      "/Users/carlomarx/.local/share/virtualenvs/life2vec-light-Ez8u7ZRp-python/lib/python3.11/site-packages/dask/dataframe/io/parquet/core.py:1215: FutureWarning: engine='pyarrow-dataset' has been deprecated. Please use engine='pyarrow' instead.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "corpus = Corpus(population=users, sources=[synth_labor, synth_health], name=\"synthetic_two\")\n",
    "corpus.prepare()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "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>RECORD_DATE</th>\n",
       "      <th>SENTENCE</th>\n",
       "      <th>BIRTHDAY</th>\n",
       "      <th>SEX</th>\n",
       "      <th>AGE</th>\n",
       "      <th>AFTER_THRESHOLD</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>USER_ID</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>122</td>\n",
       "      <td>INCOME_45 CITY_6 OCC_29 IND_32</td>\n",
       "      <td>1957-06-25</td>\n",
       "      <td>Male</td>\n",
       "      <td>63</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>130</td>\n",
       "      <td>TYPE_Outpatient</td>\n",
       "      <td>1957-06-25</td>\n",
       "      <td>Male</td>\n",
       "      <td>63</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>271</td>\n",
       "      <td>INCOME_20 CITY_18 OCC_91 IND_20</td>\n",
       "      <td>1957-06-25</td>\n",
       "      <td>Male</td>\n",
       "      <td>63</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>284</td>\n",
       "      <td>INCOME_2 CITY_21 OCC_29 IND_18</td>\n",
       "      <td>1957-06-25</td>\n",
       "      <td>Male</td>\n",
       "      <td>63</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>292</td>\n",
       "      <td>INCOME_31 CITY_17 OCC_93 IND_15</td>\n",
       "      <td>1957-06-25</td>\n",
       "      <td>Male</td>\n",
       "      <td>63</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "         RECORD_DATE                         SENTENCE   BIRTHDAY   SEX  AGE  \\\n",
       "USER_ID                                                                       \n",
       "1                122   INCOME_45 CITY_6 OCC_29 IND_32 1957-06-25  Male   63   \n",
       "1                130                  TYPE_Outpatient 1957-06-25  Male   63   \n",
       "1                271  INCOME_20 CITY_18 OCC_91 IND_20 1957-06-25  Male   63   \n",
       "1                284   INCOME_2 CITY_21 OCC_29 IND_18 1957-06-25  Male   63   \n",
       "1                292  INCOME_31 CITY_17 OCC_93 IND_15 1957-06-25  Male   63   \n",
       "\n",
       "         AFTER_THRESHOLD  \n",
       "USER_ID                   \n",
       "1                  False  \n",
       "1                  False  \n",
       "1                  False  \n",
       "1                  False  \n",
       "1                  False  "
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "corpus.combined_sentences(split=\"train\").head()\n",
    "#corpus.combined_sentences(split=\"val\").head()\n",
    "#corpus.combined_sentences(split=\"test\").head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.2. Let's create the vocabulary\n",
    "The `CorpusVocabulary` object takes the information about the `Corpus` (i.e. sentences that exist in your `train` dataset) and creates a vocabulary! Here you can choose to remove words that appear with low frequency."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from src.dataloaders.vocabulary import CorpusVocabulary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "vocab = CorpusVocabulary(corpus, name=\"synthetic_two\")\n",
    "vocab.prepare()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### NOTE:\n",
    "After running all the cells above, you should see new folders under the `data` folder!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.3. Let's speficy the task\n",
    "The `Task` object in `src\\dataloaders\\tasks\\base.py` specifies how to further process data to feed it into the model. For example, in case of the `MLM` task (specified in `src\\dataloaders\\tasks\\pretrain.py`), we specify the data augmentation procedures, as well as how to mask tokens (and create targets for the prediction task)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "from src.dataloaders.tasks.pretrain import MLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Specify the task we are going to use with the data\n",
    "task = MLM(name=\"mlm\", \n",
    "           max_length=200, # the maximum length of the sequence\n",
    "           no_sep = False, # you can decide to create data with or without the [SEP] token\n",
    "           # Augmentation\n",
    "            p_sequence_timecut = 0.01,\n",
    "            p_sequence_resample = 0.01,\n",
    "            p_sequence_abspos_noise = 0.1,\n",
    "            p_sequence_hide_background = 0.01,\n",
    "            p_sentence_drop_tokens = 0.01,\n",
    "            shuffle_within_sentences = True,\n",
    "            # MLM specific options\n",
    "            mask_ratio = 0.1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4. DataModule\n",
    "Datamodule object takes information about the `Corpus`, `Task` and `CorpusVocabulary` and assembles inputs that we further provide to the model. To learn more about the datamodules see the [Pytorch Lightning: Dataloader](https://lightning.ai/docs/pytorch/stable/data/datamodule.html) documentation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "from src.dataloaders.datamodule import L2VDataModule"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "## On the first initialization of the datamodule, we do not have a vocabulary object\n",
    "datamodule = L2VDataModule(corpus, batch_size=2, task=task, vocabulary=vocab, num_workers=2)\n",
    "datamodule.prepare()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Structure of the batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = next(iter(datamodule.train_dataloader()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys(['sequence_id', 'input_ids', 'padding_mask', 'target_tokens', 'target_pos', 'target_sop', 'original_sequence'])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([411, 165])"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X[\"sequence_id\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 5. Add model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "### use code from the experiments\n",
    "from src.models.pretrain import TransformerEncoder\n",
    "from pytorch_lightning import Trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "hparams = {\n",
    "    \"hidden_size\": 96,  # size of the hidden layers and embeddings\n",
    "    \"hidden_ff\": 128,  # size of the position-wise feed-forward layer\n",
    "    \"n_encoders\": 4,  # number of encoder blocks\n",
    "    \"n_heads\": 8,  # number of attention heads in the multiheadattention module\n",
    "    \"n_local\": 2,  # number of local attention heads\n",
    "    \"local_window_size\": 4,  # size of the window for local attention\n",
    "    \"max_length\": task.max_length,  # maximum length of the input sequence\n",
    "    \"vocab_size\": vocab.size(),  # size of the vocabulary\n",
    "    \"num_classes\": -1,  \n",
    "    \"cls_num_targs\": 3, # number of classes for the SOP class (we have 3: original, reversed, shuffled)\n",
    "    \"learning_rate\": 0.001,\n",
    "    \"batch_size\": datamodule.batch_size,\n",
    "    \"num_epochs\": 30,\n",
    "    #\"device\": 'cpu', # 'cuda' if you have a GPU or 'mps' if you have an Apple M1\n",
    "    \"attention_type\": \"performer\",\n",
    "    \"norm_type\": \"rezero\",\n",
    "    \"num_random_features\": 32,  # number of random features for the Attention module (Performer uses this)\n",
    "    \"parametrize_emb\": True,  # whether to center the token embedding matrix\n",
    "    \"emb_dropout\": 0.1,  # dropout for the embedding block\n",
    "    \"fw_dropout\": 0.1,  # dropout for the position-wise feed-forward layer\n",
    "    \"att_dropout\": 0.1,  # dropout for the multiheadattention module\n",
    "    \"dc_dropout\": 0.1,  # dropout for the decoder block\n",
    "    \"hidden_act\": \"swish\",  # activation function for the hidden layers (attention layers use ReLU)\n",
    "    \"training_task\": task.name,\n",
    "    \"weight_tying\": True,\n",
    "    \"norm_output_emb\": True,\n",
    "    \"epsilon\": 1e-8,\n",
    "    \"weight_decay\": 0.01,\n",
    "    \"beta1\": 0.9,\n",
    "    \"beta2\": 0.999,\n",
    "}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "l2v = TransformerEncoder(hparams=hparams).float()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True (mps), used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "HPU available: False, using: 0 HPUs\n"
     ]
    }
   ],
   "source": [
    "trainer = Trainer(max_epochs=hparams[\"num_epochs\"],\n",
    "                  limit_train_batches=100, accelerator=\"cpu\", )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "  | Name        | Type                | Params | Mode \n",
      "------------------------------------------------------------\n",
      "0 | transformer | Transformer         | 277 K  | train\n",
      "1 | mlm_decoder | MaskedLanguageModel | 39.2 K | train\n",
      "2 | sop_decoder | SOP_Decoder         | 9.6 K  | train\n",
      "3 | sop_loss    | CrossEntropyLoss    | 0      | train\n",
      "4 | mlm_loss    | CrossEntropyLoss    | 0      | train\n",
      "------------------------------------------------------------\n",
      "326 K     Trainable params\n",
      "0         Non-trainable params\n",
      "326 K     Total params\n",
      "1.306     Total estimated model params size (MB)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "data/processed/datasets/synthetic_two/mlm/_arguments\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 3:  83%|████████▎ | 83/100 [00:03<00:00, 23.24it/s, v_num=6]         "
     ]
    }
   ],
   "source": [
    "trainer.fit(model=l2v, datamodule=datamodule)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## More annotations to come."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "torch",
   "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}