{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "27b646f0",
   "metadata": {},
   "source": [
    "# Pretraining\n",
    "\n",
    "In this tutorial, we will pretrain a masked transformer model on our single-cell multiomics data using scDynOmics.\n",
    "The pretraining process involves masking a portion of the input data and training the model to predict the masked values,\n",
    "    which helps the model learn meaningful representations of the data.\n",
    "\n",
    "The associated script is available in the scDynOmics repository under ``scripts/pt.py``.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ceb37eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "from pytorch_lightning import Trainer, seed_everything\n",
    "import scdynomics.utils.json as JSON\n",
    "from scdynomics.utils.trainer import (\n",
    "    set_accelerator,\n",
    "    cal_accum_grad_batches,\n",
    "    make_callbacks,\n",
    ")\n",
    "from scdynomics.utils.data import Multimodal_Corpus\n",
    "from scdynomics.pt.reconstructor import load_or_init\n",
    "from scdynomics.pt.pipeline import bimodal as bimodal_pt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52d3b6c4",
   "metadata": {},
   "source": [
    "## 0. Configuration\n",
    "\n",
    "For the purpose of this tutorial,\n",
    "    the configuration parameters are explicitly defined below to make them easy to explore and modify.\n",
    "\n",
    "*Note: In a production environment,\n",
    "    especially when using SLURM or other high-performance computing (HPC) resource distribution systems,\n",
    "    these configurations are typically parsed from command-line arguments.\n",
    "This allows for greater flexibility and seamless integration with job scheduling pipelines.*\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "8761e144",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Seed set to 42\n"
     ]
    }
   ],
   "source": [
    "# File paths\n",
    "data_corpus_path = \"../../data/sample/mm10.100sample.h5ad\" # example data included in the scDynOmics repository relative to this notebook\n",
    "config_path = \"../../data/config/mm10.pt.sample.json\" # for production, also create a modified config file, especially with higher value than 1 for trainer_hparams->max_epochs \n",
    "token_dict_path = \"../../data/gene_dict/token_dict.json.gz\"\n",
    "default_root_dir = \"logs\"\n",
    "pretrainer_ckpt = None  # Set to a path if resuming from a checkpoint\n",
    "\n",
    "# Data and layers\n",
    "mod1_layer = \"X\"\n",
    "mod2_layer = \"pan_promoter\"\n",
    "\n",
    "# Training parameters\n",
    "valid_fraction = 0.1\n",
    "dataloader_worker = 10\n",
    "seed = 42\n",
    "\n",
    "# Apply seed\n",
    "seed_everything(seed)\n",
    "torch.set_float32_matmul_precision(\"medium\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb32d4f5",
   "metadata": {},
   "source": [
    "## 1. Summarized Pipeline Execution\n",
    "\n",
    "You can run the entire pretraining process using the high-level :func:`scdynomics.pt.pipeline.bimodal` function.\n",
    "This function abstracts away the individual data loading and training steps.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26d1c162",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Seed set to 42\n",
      "/usr/lib/python3.11/site-packages/lightning_fabric/plugins/environments/slurm.py:204: The `srun` command is available on your system but is not used. HINT: If your intention is to run Lightning on SLURM, prepend your python command with `srun` like so: srun python /usr/lib/python3 ...\n",
      "GPU available: True (cuda), used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "HPU available: False, using: 0 HPUs\n",
      "/usr/lib/python3.11/site-packages/pytorch_lightning/trainer/setup.py:177: GPU available but not used. You can set it by doing `Trainer(accelerator='gpu')`.\n",
      "\n",
      "  | Name           | Type             | Params | Mode \n",
      "------------------------------------------------------------\n",
      "0 | vbinning        | Value_Binning    | 0      | train\n",
      "1 | input_dropout  | Input_Dropout    | 0      | train\n",
      "2 | embedder       | Genome_Embedder  | 2.2 M  | train\n",
      "3 | encoder        | Encoder          | 75.8 M | train\n",
      "4 | pool           | Sequential       | 0      | train\n",
      "5 | mod0_rebuilder | Linear           | 1.5 K  | train\n",
      "6 | mod1_rebuilder | Linear           | 1.5 K  | train\n",
      "7 | criterion      | CrossEntropyLoss | 0      | train\n",
      "------------------------------------------------------------\n",
      "77.9 M    Trainable params\n",
      "0         Non-trainable params\n",
      "77.9 M    Total params\n",
      "311.676   Total estimated model params size (MB)\n",
      "185       Modules in train mode\n",
      "0         Modules in eval mode\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Random initialized pretrainer model\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`Trainer.fit` stopped: `max_epochs=1` reached.\n"
     ]
    }
   ],
   "source": [
    "# Run pretraining\n",
    "trainer, pretrainer_model = bimodal_pt(\n",
    "    config_path = config_path,\n",
    "    data_corpus_path = data_corpus_path,\n",
    "    mod1_layer = mod1_layer,\n",
    "    mod2_layer = mod2_layer,\n",
    "    valid_fraction = valid_fraction,\n",
    "    dataloader_worker = dataloader_worker,\n",
    "    pretrainer_ckpt = pretrainer_ckpt,\n",
    "    token_dict_path = token_dict_path,\n",
    "    default_root_dir = default_root_dir,\n",
    "    seed = seed,\n",
    "    ndevices = 1,\n",
    "    accelerator = 'cpu',\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f97ceb87",
   "metadata": {},
   "source": [
    "## 2. Detailed Step-by-Step Breakdown\n",
    "\n",
    "Below, we break down what happens under the hood of the :func:`scdynomics.pt.pipeline.bimodal` function to give you a deeper understanding of the framework's mechanics.\n",
    "\n",
    "---\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2525ecc",
   "metadata": {},
   "source": [
    "### 2.1 Data Loading\n",
    "\n",
    "We load our data into a :class:`scdynomics.utils.data.Multimodal_Corpus` and split it into training and validation sets.\n",
    "Then, we wrap them in PyTorch DataLoaders for efficient batching."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "bff5f4d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the config file\n",
    "config = JSON.decode(config_path)\n",
    "dl_hparams = config[\"dataloader_hparams\"]\n",
    "\n",
    "# Load training data corpus\n",
    "full_data = Multimodal_Corpus(\n",
    "    data_corpus_path,\n",
    "    layers = [mod1_layer, mod2_layer],\n",
    "    backed = False,\n",
    ")\n",
    "\n",
    "# Split train and validation data\n",
    "train_data_size = int(1-valid_fraction * len(full_data))\n",
    "train_data, val_data = torch.utils.data.random_split(\n",
    "    full_data,\n",
    "    [train_data_size, len(full_data) - train_data_size]\n",
    ")\n",
    "train_data = DataLoader(\n",
    "    train_data,\n",
    "    num_workers = dataloader_worker,\n",
    "    batch_size = dl_hparams[\"batch_per_device\"],\n",
    "    shuffle = True\n",
    ")\n",
    "val_data = DataLoader(\n",
    "    val_data,\n",
    "    num_workers = dataloader_worker,\n",
    "    batch_size = dl_hparams[\"batch_per_device\"],\n",
    "    shuffle = False\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf5a9cdd",
   "metadata": {},
   "source": [
    "### 2.2 Model Initialization\n",
    "\n",
    "We initialize our :class:`scdynomics.Masked_Pretrainer`.\n",
    "If a checkpoint path is provided, we load the model from the checkpoint to resume training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "db82346c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Random initialized pretrainer model\n"
     ]
    }
   ],
   "source": [
    "# Load or initialize the pretrainer model\n",
    "pretrainer_model = load_or_init(\n",
    "    ckpt_path = pretrainer_ckpt,\n",
    "    token_dict_path = token_dict_path,\n",
    "    config = config,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a6f6f4f",
   "metadata": {},
   "source": [
    "### 2.3 Trainer Setup & Execution\n",
    "\n",
    "Finally, we set up PyTorch Lightning's Trainer with our callbacks,\n",
    "    like Model Checkpointing,\n",
    "    and start the training process"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "d671b63e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/lib/python3.11/site-packages/lightning_fabric/plugins/environments/slurm.py:204: The `srun` command is available on your system but is not used. HINT: If your intention is to run Lightning on SLURM, prepend your python command with `srun` like so: srun python /usr/lib/python3 ...\n",
      "GPU available: True (cuda), used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "HPU available: False, using: 0 HPUs\n",
      "/usr/lib/python3.11/site-packages/pytorch_lightning/trainer/setup.py:177: GPU available but not used. You can set it by doing `Trainer(accelerator='gpu')`.\n",
      "\n",
      "  | Name           | Type             | Params | Mode \n",
      "------------------------------------------------------------\n",
      "0 | vbinning        | Value_Binning    | 0      | train\n",
      "1 | input_dropout  | Input_Dropout    | 0      | train\n",
      "2 | embedder       | Genome_Embedder  | 2.2 M  | train\n",
      "3 | encoder        | Encoder          | 75.8 M | train\n",
      "4 | pool           | Sequential       | 0      | train\n",
      "5 | mod0_rebuilder | Linear           | 1.5 K  | train\n",
      "6 | mod1_rebuilder | Linear           | 1.5 K  | train\n",
      "7 | criterion      | CrossEntropyLoss | 0      | train\n",
      "------------------------------------------------------------\n",
      "77.9 M    Trainable params\n",
      "0         Non-trainable params\n",
      "77.9 M    Total params\n",
      "311.676   Total estimated model params size (MB)\n",
      "185       Modules in train mode\n",
      "0         Modules in eval mode\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`Trainer.fit` stopped: `max_epochs=1` reached.\n"
     ]
    }
   ],
   "source": [
    "# Init trainer\n",
    "trainer = Trainer(\n",
    "    devices = 1,\n",
    "    accelerator = \"cpu\", # or \"cuda\",\n",
    "    default_root_dir = default_root_dir,\n",
    "    accumulate_grad_batches = cal_accum_grad_batches(config, 1),\n",
    "    callbacks = make_callbacks(config),\n",
    "    **config[\"trainer_hparams\"]\n",
    ")\n",
    "\n",
    "# Train the model\n",
    "trainer.fit(\n",
    "    model = pretrainer_model,\n",
    "    train_dataloaders = train_data,\n",
    "    val_dataloaders = val_data,\n",
    "    ckpt_path = pretrainer_ckpt,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "365ff300",
   "metadata": {},
   "source": [
    "## 3. TensorBoard Visualization\n",
    "\n",
    "The training progress and metrics can be monitored using TensorBoard by running the following command in the terminal:\n",
    "\n",
    "---\n",
    "```bash\n",
    "tensorboard --logdir logs\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "639303d5",
   "metadata": {},
   "source": [
    "## 4. Cleanup\n",
    "\n",
    "WARNING: Don't run this cell if you want to load the checkpoint for the fine-tuning tutorial.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73827122",
   "metadata": {},
   "outputs": [],
   "source": [
    "import shutil\n",
    "shutil.rmtree(default_root_dir)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}