{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e57e6fc9",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/lucascamillomd/pyaging/blob/main/tutorials/tutorial_cpgptgrimage3.ipynb) [![Open In nbviewer](https://img.shields.io/badge/View%20in-nbviewer-orange)](https://nbviewer.jupyter.org/github/lucascamillomd/pyaging/blob/main/tutorials/tutorial_cpgptgrimage3.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0fbbe22",
   "metadata": {
    "papermill": {
     "duration": 0.005095,
     "end_time": "2025-03-31T09:46:56.905323",
     "exception": false,
     "start_time": "2025-03-31T09:46:56.900228",
     "status": "completed"
    },
    "tags": [],
    "vscode": {
     "languageId": "plaintext"
    }
   },
   "source": [
    "# The best DNAm mortality predictor: CpGPTGrimAge3\n",
    "\n",
    "## Table of Contents\n",
    "\n",
    "0. <a href=\"#0-read-quick-setup-tutorial\">Read Quick Setup Tutorial</a>\n",
    "1. <a href=\"#1-setup-environment\">Setup Environment</a>\n",
    "2. <a href=\"#2-load-data\">Load Data</a>\n",
    "3. <a href=\"#3-load-model-and-dependencies\">Load Model and Dependencies</a>\n",
    "4. <a href=\"#4-prepare-data-objects\">Prepare Data Objects</a>\n",
    "5. <a href=\"#5-compute-protein-proxies\">Compute Protein Proxies</a>\n",
    "6. <a href=\"#6-calculate-cpgptgrimage3\">Calculate CpGPTGrimAge3</a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c68a9cb",
   "metadata": {},
   "source": [
    "## 0. Read Quick Setup Tutorial"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "658a745d",
   "metadata": {},
   "source": [
    "Before, going through this tutorial, please familiarize yourself with the [quick setup tutorial](https://github.com/lcamillo/CpGPT/blob/main/tutorials/quick_setup.ipynb)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4966e7b",
   "metadata": {
    "papermill": {
     "duration": 0.004234,
     "end_time": "2025-03-31T09:46:56.931413",
     "exception": false,
     "start_time": "2025-03-31T09:46:56.927179",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## 1. Setup Environment\n",
    "\n",
    "CpGPT needs to be installed. The easiest is to use the following:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "249f058a",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install CpGPT==0.0.10 --quiet"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b937a4a",
   "metadata": {},
   "source": [
    "Please check out more instructions in the [offical CpGPT repo](https://github.com/lucascamillomd/CpGPT).\n",
    "\n",
    "We'll import the necessary Python packages and set up our environment for CpGPT. We'll be using a mix of standard data science libraries and CpGPT-specific modules. We'll also set some important variables that will be used throughout the notebook. Pay attention to these as you may need to adjust them based on your specific setup and requirements.\n",
    "\n",
    "CpGPT model files and DNA-sequence dependencies are hosted on Hugging Face.\n",
    "The next Python cell downloads any missing files into the standard Hugging Face cache and reuses them on later runs; no command-line download is needed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "hf-download",
   "metadata": {},
   "outputs": [],
   "source": [
    "from cpgpt import download_cpgpt\n",
    "\n",
    "resources = download_cpgpt(model=\"proteins\", species=\"human\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d73f9bda",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-31T09:46:56.941492Z",
     "iopub.status.busy": "2025-03-31T09:46:56.941081Z",
     "iopub.status.idle": "2025-03-31T09:46:56.945145Z",
     "shell.execute_reply": "2025-03-31T09:46:56.944678Z"
    },
    "papermill": {
     "duration": 0.009909,
     "end_time": "2025-03-31T09:46:56.945906",
     "exception": false,
     "start_time": "2025-03-31T09:46:56.935997",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Random seed for reproducibility\n",
    "RANDOM_SEED = 42\n",
    "\n",
    "# Directory paths\n",
    "DEPENDENCIES_DIR = \"../dependencies\"\n",
    "DATA_DIR = \"../data\"\n",
    "PROCESSED_DIR = \"../data/tutorials/processed/predict_mortality\"\n",
    "\n",
    "MODEL_NAME = \"proteins\" # this is the name of the model checkpoint required for CpGPTGrimAge3\n",
    "\n",
    "BETAS_PATH = \"../data/cpgcorpus/raw/GSE237561/GPL13534/betas/QCDPB.arrow\"\n",
    "FILTERED_BETAS_PATH = \"../data/cpgcorpus/raw/GSE237561/GPL13534/betas/QCDPB_filtered.arrow\"\n",
    "METADATA_PATH = \"../data/cpgcorpus/raw/GSE237561/GPL13534/metadata/metadata.arrow\"\n",
    "\n",
    "# The maximum context length to give to the model\n",
    "MAX_INPUT_LENGTH = 10_000 # you might wanna go higher hardware permitting\n",
    "\n",
    "LLM_DEPENDENCIES_DIR = str(resources.dependencies_path)\n",
    "MODEL_CHECKPOINT_PATH = str(resources.checkpoint_path)\n",
    "MODEL_CONFIG_PATH = str(resources.config_path)\n",
    "MODEL_VOCAB_PATH = str(resources.vocab_path) if resources.vocab_path is not None else None"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "107e02d2",
   "metadata": {
    "papermill": {
     "duration": 0.004296,
     "end_time": "2025-03-31T09:46:56.954551",
     "exception": false,
     "start_time": "2025-03-31T09:46:56.950255",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "> **⚠️ Warning**\n",
    "> \n",
    "> It is recommended to have a GPU for inference as CPU might be slow.\n",
    "> \n",
    "> Reconstructing the methylome for a few hundred samples might take up to one hour on a CPU. ⌛\n",
    ">\n",
    "> This might be a great exercise in testing your patience."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "974e26e7",
   "metadata": {
    "papermill": {
     "duration": 0.004292,
     "end_time": "2025-03-31T09:46:56.963100",
     "exception": false,
     "start_time": "2025-03-31T09:46:56.958808",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "### 1.2 Import packages\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c0ab299",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-31T09:46:56.972288Z",
     "iopub.status.busy": "2025-03-31T09:46:56.972092Z",
     "iopub.status.idle": "2025-03-31T09:47:03.556135Z",
     "shell.execute_reply": "2025-03-31T09:47:03.555626Z"
    },
    "papermill": {
     "duration": 6.589622,
     "end_time": "2025-03-31T09:47:03.556954",
     "exception": false,
     "start_time": "2025-03-31T09:46:56.967332",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:torchao.kernel.intmm:Warning: Detected no triton, on systems without Triton certain kernels will not work\n",
      "Seed set to 42\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "42"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Standard library imports\n",
    "import warnings\n",
    "import os\n",
    "import json\n",
    "\n",
    "warnings.simplefilter(action=\"ignore\", category=FutureWarning)\n",
    "\n",
    "# Plotting imports\n",
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pyaging as pya\n",
    "import seaborn as sns\n",
    "from tqdm.rich import tqdm\n",
    "\n",
    "# Lightning imports\n",
    "from lightning.pytorch import seed_everything\n",
    "\n",
    "# cpgpt-specific imports\n",
    "from cpgpt.data.components.cpgpt_datasaver import CpGPTDataSaver\n",
    "from cpgpt.data.cpgpt_datamodule import CpGPTDataModule\n",
    "from cpgpt.trainer.cpgpt_trainer import CpGPTTrainer\n",
    "from cpgpt.data.components.dna_llm_embedder import DNALLMEmbedder\n",
    "from cpgpt.data.components.illumina_methylation_prober import IlluminaMethylationProber\n",
    "from cpgpt.infer.cpgpt_inferencer import CpGPTInferencer\n",
    "from cpgpt.model.cpgpt_module import m_to_beta\n",
    "\n",
    "# Set random seed for reproducibility\n",
    "seed_everything(RANDOM_SEED, workers=True)\n",
    "try:\n",
    "    torch.backends.cudnn.deterministic = True\n",
    "    torch.backends.cudnn.benchmark = False\n",
    "except:\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca00d8db",
   "metadata": {
    "papermill": {
     "duration": 0.004402,
     "end_time": "2025-03-31T09:47:03.566152",
     "exception": false,
     "start_time": "2025-03-31T09:47:03.561750",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## 2. Load Data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68ddcd70",
   "metadata": {
    "papermill": {
     "duration": 0.00437,
     "end_time": "2025-03-31T09:47:03.574896",
     "exception": false,
     "start_time": "2025-03-31T09:47:03.570526",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "If you have your own data, please feel free to skip the following step but make sure it is saved in a .arrow format. Here, as an example target dataset, we'll use GSE237561, which contains methylation profiling data from 126 peripheral whole blood samples collected from 26 individuals across two independent cohorts. These samples were collected at three timepoints: prior to clozapine initiation, 4-12 weeks after initiation, and 6 months after initiation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "608f10d0",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-31T09:47:03.584602Z",
     "iopub.status.busy": "2025-03-31T09:47:03.584214Z",
     "iopub.status.idle": "2025-03-31T09:47:04.121930Z",
     "shell.execute_reply": "2025-03-31T09:47:04.121356Z"
    },
    "papermill": {
     "duration": 0.543485,
     "end_time": "2025-03-31T09:47:04.122718",
     "exception": false,
     "start_time": "2025-03-31T09:47:03.579233",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[34mcpgpt\u001b[0m\u001b[1m\u001b[0m -\u001b[36mCpGPTInferencer\u001b[0m: \u001b[1mInitializing class CpGPTInferencer.\u001b[0m\n",
      "\u001b[1m\u001b[34mcpgpt\u001b[0m\u001b[1m\u001b[0m -\u001b[36mCpGPTInferencer\u001b[0m: \u001b[1mUsing device: cpu.\u001b[0m\n",
      "\u001b[1m\u001b[34mcpgpt\u001b[0m\u001b[1m\u001b[0m -\u001b[36mCpGPTInferencer\u001b[0m: \u001b[1mUsing dependencies directory: ../dependencies\u001b[0m\n",
      "\u001b[1m\u001b[34mcpgpt\u001b[0m\u001b[1m\u001b[0m -\u001b[36mCpGPTInferencer\u001b[0m: \u001b[1mUsing data directory: ../data\u001b[0m\n",
      "\u001b[1m\u001b[34mcpgpt\u001b[0m\u001b[1m\u001b[0m -\u001b[36mCpGPTInferencer\u001b[0m: \u001b[1mThere are 19 CpGPT models available such as age, age_cot, average_adultweight, etc.\u001b[0m\n",
      "\u001b[1m\u001b[34mcpgpt\u001b[0m\u001b[1m\u001b[0m -\u001b[36mCpGPTInferencer\u001b[0m: \u001b[1mThere are 2088 GSE datasets available such as GSE100184, GSE100208, GSE100209, etc.\u001b[0m\n",
      "\u001b[1m\u001b[34mcpgpt\u001b[0m\u001b[1m\u001b[0m -\u001b[36mCpGPTInferencer\u001b[0m: \u001b[1mDataset GSE237561 already exists at ../data/cpgcorpus/raw/GSE237561 (skipping download).\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "# First let's declare the inferencer\n",
    "inferencer = CpGPTInferencer(dependencies_dir=DEPENDENCIES_DIR, data_dir=DATA_DIR)\n",
    "\n",
    "inferencer.download_cpgcorpus_dataset(\"GSE237561\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "99b48b7e",
   "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>cg00000029</th>\n",
       "      <th>cg00000108</th>\n",
       "      <th>cg00000109</th>\n",
       "      <th>cg00000165</th>\n",
       "      <th>cg00000236</th>\n",
       "      <th>cg00000289</th>\n",
       "      <th>cg00000292</th>\n",
       "      <th>cg00000321</th>\n",
       "      <th>cg00000363</th>\n",
       "      <th>cg00000622</th>\n",
       "      <th>...</th>\n",
       "      <th>rs7746156</th>\n",
       "      <th>rs798149</th>\n",
       "      <th>rs845016</th>\n",
       "      <th>rs877309</th>\n",
       "      <th>rs9292570</th>\n",
       "      <th>rs9363764</th>\n",
       "      <th>rs939290</th>\n",
       "      <th>rs951295</th>\n",
       "      <th>rs966367</th>\n",
       "      <th>rs9839873</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>GSM_ID</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></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>GSM7625568</th>\n",
       "      <td>0.592157</td>\n",
       "      <td>0.964411</td>\n",
       "      <td>0.899373</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.861353</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.894893</td>\n",
       "      <td>0.280911</td>\n",
       "      <td>0.386535</td>\n",
       "      <td>0.017273</td>\n",
       "      <td>...</td>\n",
       "      <td>0.974810</td>\n",
       "      <td>0.022425</td>\n",
       "      <td>0.086804</td>\n",
       "      <td>0.031910</td>\n",
       "      <td>0.535025</td>\n",
       "      <td>0.548239</td>\n",
       "      <td>0.064924</td>\n",
       "      <td>0.536217</td>\n",
       "      <td>0.091576</td>\n",
       "      <td>0.787829</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>GSM7625569</th>\n",
       "      <td>0.657346</td>\n",
       "      <td>0.962779</td>\n",
       "      <td>0.920897</td>\n",
       "      <td>0.170290</td>\n",
       "      <td>0.868804</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.945775</td>\n",
       "      <td>0.303094</td>\n",
       "      <td>0.409573</td>\n",
       "      <td>0.015473</td>\n",
       "      <td>...</td>\n",
       "      <td>0.977200</td>\n",
       "      <td>0.023908</td>\n",
       "      <td>0.093790</td>\n",
       "      <td>0.024449</td>\n",
       "      <td>0.524899</td>\n",
       "      <td>0.554774</td>\n",
       "      <td>0.048559</td>\n",
       "      <td>0.521857</td>\n",
       "      <td>0.081250</td>\n",
       "      <td>0.772327</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>GSM7625570</th>\n",
       "      <td>0.662022</td>\n",
       "      <td>0.964065</td>\n",
       "      <td>0.903984</td>\n",
       "      <td>0.180436</td>\n",
       "      <td>0.867933</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.915102</td>\n",
       "      <td>0.241706</td>\n",
       "      <td>0.392485</td>\n",
       "      <td>0.015086</td>\n",
       "      <td>...</td>\n",
       "      <td>0.979132</td>\n",
       "      <td>0.022763</td>\n",
       "      <td>0.091286</td>\n",
       "      <td>0.028332</td>\n",
       "      <td>0.518550</td>\n",
       "      <td>0.584879</td>\n",
       "      <td>0.054734</td>\n",
       "      <td>0.529968</td>\n",
       "      <td>0.101047</td>\n",
       "      <td>0.765933</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>GSM7625571</th>\n",
       "      <td>0.599778</td>\n",
       "      <td>0.961087</td>\n",
       "      <td>0.903260</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.845338</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.910445</td>\n",
       "      <td>0.277753</td>\n",
       "      <td>0.405914</td>\n",
       "      <td>0.016514</td>\n",
       "      <td>...</td>\n",
       "      <td>0.978393</td>\n",
       "      <td>0.019485</td>\n",
       "      <td>0.069463</td>\n",
       "      <td>0.024282</td>\n",
       "      <td>0.508195</td>\n",
       "      <td>0.569030</td>\n",
       "      <td>0.052701</td>\n",
       "      <td>0.508199</td>\n",
       "      <td>0.083252</td>\n",
       "      <td>0.787774</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>GSM7625572</th>\n",
       "      <td>0.556610</td>\n",
       "      <td>0.960655</td>\n",
       "      <td>0.893885</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.846172</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.916346</td>\n",
       "      <td>0.285945</td>\n",
       "      <td>0.404618</td>\n",
       "      <td>0.014193</td>\n",
       "      <td>...</td>\n",
       "      <td>0.978121</td>\n",
       "      <td>0.019678</td>\n",
       "      <td>0.541755</td>\n",
       "      <td>0.982736</td>\n",
       "      <td>0.528156</td>\n",
       "      <td>0.524965</td>\n",
       "      <td>0.965231</td>\n",
       "      <td>0.041688</td>\n",
       "      <td>0.672611</td>\n",
       "      <td>0.924847</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5 rows × 485578 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "            cg00000029  cg00000108  cg00000109  cg00000165  cg00000236  \\\n",
       "GSM_ID                                                                   \n",
       "GSM7625568    0.592157    0.964411    0.899373         NaN    0.861353   \n",
       "GSM7625569    0.657346    0.962779    0.920897    0.170290    0.868804   \n",
       "GSM7625570    0.662022    0.964065    0.903984    0.180436    0.867933   \n",
       "GSM7625571    0.599778    0.961087    0.903260         NaN    0.845338   \n",
       "GSM7625572    0.556610    0.960655    0.893885         NaN    0.846172   \n",
       "\n",
       "            cg00000289  cg00000292  cg00000321  cg00000363  cg00000622  ...  \\\n",
       "GSM_ID                                                                  ...   \n",
       "GSM7625568         NaN    0.894893    0.280911    0.386535    0.017273  ...   \n",
       "GSM7625569         NaN    0.945775    0.303094    0.409573    0.015473  ...   \n",
       "GSM7625570         NaN    0.915102    0.241706    0.392485    0.015086  ...   \n",
       "GSM7625571         NaN    0.910445    0.277753    0.405914    0.016514  ...   \n",
       "GSM7625572         NaN    0.916346    0.285945    0.404618    0.014193  ...   \n",
       "\n",
       "            rs7746156  rs798149  rs845016  rs877309  rs9292570  rs9363764  \\\n",
       "GSM_ID                                                                      \n",
       "GSM7625568   0.974810  0.022425  0.086804  0.031910   0.535025   0.548239   \n",
       "GSM7625569   0.977200  0.023908  0.093790  0.024449   0.524899   0.554774   \n",
       "GSM7625570   0.979132  0.022763  0.091286  0.028332   0.518550   0.584879   \n",
       "GSM7625571   0.978393  0.019485  0.069463  0.024282   0.508195   0.569030   \n",
       "GSM7625572   0.978121  0.019678  0.541755  0.982736   0.528156   0.524965   \n",
       "\n",
       "            rs939290  rs951295  rs966367  rs9839873  \n",
       "GSM_ID                                               \n",
       "GSM7625568  0.064924  0.536217  0.091576   0.787829  \n",
       "GSM7625569  0.048559  0.521857  0.081250   0.772327  \n",
       "GSM7625570  0.054734  0.529968  0.101047   0.765933  \n",
       "GSM7625571  0.052701  0.508199  0.083252   0.787774  \n",
       "GSM7625572  0.965231  0.041688  0.672611   0.924847  \n",
       "\n",
       "[5 rows x 485578 columns]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Load betas matrix\n",
    "df = pd.read_feather(BETAS_PATH)\n",
    "df.set_index(\"GSM_ID\", inplace=True)\n",
    "\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "75c65ae2",
   "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>title</th>\n",
       "      <th>geo_accession</th>\n",
       "      <th>status</th>\n",
       "      <th>submission_date</th>\n",
       "      <th>last_update_date</th>\n",
       "      <th>type</th>\n",
       "      <th>channel_count</th>\n",
       "      <th>source_name_ch1</th>\n",
       "      <th>organism_ch1</th>\n",
       "      <th>characteristics_ch1</th>\n",
       "      <th>...</th>\n",
       "      <th>cd8t:ch1</th>\n",
       "      <th>days.on.clozapine:ch1</th>\n",
       "      <th>gran:ch1</th>\n",
       "      <th>institute:ch1</th>\n",
       "      <th>mono:ch1</th>\n",
       "      <th>nk:ch1</th>\n",
       "      <th>participant_id:ch1</th>\n",
       "      <th>Sex:ch1</th>\n",
       "      <th>smokingscore:ch1</th>\n",
       "      <th>visit:ch1</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>GSM_ID</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></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>GSM7625568</th>\n",
       "      <td>genomic DNA from ID 0003 for 'a' visit</td>\n",
       "      <td>GSM7625568</td>\n",
       "      <td>Public on Jul 17 2024</td>\n",
       "      <td>Jul 17 2023</td>\n",
       "      <td>Jul 17 2024</td>\n",
       "      <td>genomic</td>\n",
       "      <td>1</td>\n",
       "      <td>peripheral whole blood</td>\n",
       "      <td>Homo sapiens</td>\n",
       "      <td>participant_id: 0003</td>\n",
       "      <td>...</td>\n",
       "      <td>0.124088032132122</td>\n",
       "      <td>0</td>\n",
       "      <td>0.623747063903508</td>\n",
       "      <td>KCL</td>\n",
       "      <td>0.0577366406423333</td>\n",
       "      <td>0.0181886611163226</td>\n",
       "      <td>0003</td>\n",
       "      <td>M</td>\n",
       "      <td>0.744140048408035</td>\n",
       "      <td>a</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>GSM7625569</th>\n",
       "      <td>genomic DNA from ID 0003 for 'b' visit</td>\n",
       "      <td>GSM7625569</td>\n",
       "      <td>Public on Jul 17 2024</td>\n",
       "      <td>Jul 17 2023</td>\n",
       "      <td>Jul 17 2024</td>\n",
       "      <td>genomic</td>\n",
       "      <td>1</td>\n",
       "      <td>peripheral whole blood</td>\n",
       "      <td>Homo sapiens</td>\n",
       "      <td>participant_id: 0003</td>\n",
       "      <td>...</td>\n",
       "      <td>0.140642508939063</td>\n",
       "      <td>42</td>\n",
       "      <td>0.532222309707589</td>\n",
       "      <td>KCL</td>\n",
       "      <td>0.0794055065754302</td>\n",
       "      <td>0.0165444940880789</td>\n",
       "      <td>0003</td>\n",
       "      <td>M</td>\n",
       "      <td>0.778521295727892</td>\n",
       "      <td>b</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>GSM7625570</th>\n",
       "      <td>genomic DNA from ID 0003 for 'd' visit</td>\n",
       "      <td>GSM7625570</td>\n",
       "      <td>Public on Jul 17 2024</td>\n",
       "      <td>Jul 17 2023</td>\n",
       "      <td>Jul 17 2024</td>\n",
       "      <td>genomic</td>\n",
       "      <td>1</td>\n",
       "      <td>peripheral whole blood</td>\n",
       "      <td>Homo sapiens</td>\n",
       "      <td>participant_id: 0003</td>\n",
       "      <td>...</td>\n",
       "      <td>0.112389247455345</td>\n",
       "      <td>84</td>\n",
       "      <td>0.610861991010505</td>\n",
       "      <td>KCL</td>\n",
       "      <td>0.0694956639909667</td>\n",
       "      <td>0.00363827604122652</td>\n",
       "      <td>0003</td>\n",
       "      <td>M</td>\n",
       "      <td>0.591346224352136</td>\n",
       "      <td>d</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>GSM7625571</th>\n",
       "      <td>genomic DNA from ID 0003 for 'e' visit</td>\n",
       "      <td>GSM7625571</td>\n",
       "      <td>Public on Jul 17 2024</td>\n",
       "      <td>Jul 17 2023</td>\n",
       "      <td>Jul 17 2024</td>\n",
       "      <td>genomic</td>\n",
       "      <td>1</td>\n",
       "      <td>peripheral whole blood</td>\n",
       "      <td>Homo sapiens</td>\n",
       "      <td>participant_id: 0003</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0927185883637476</td>\n",
       "      <td>168</td>\n",
       "      <td>0.647578326303527</td>\n",
       "      <td>KCL</td>\n",
       "      <td>0.0886719923597006</td>\n",
       "      <td>0.0163329200193279</td>\n",
       "      <td>0003</td>\n",
       "      <td>M</td>\n",
       "      <td>-0.771644717383253</td>\n",
       "      <td>e</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>GSM7625572</th>\n",
       "      <td>genomic DNA from ID 0005 for 'a' visit</td>\n",
       "      <td>GSM7625572</td>\n",
       "      <td>Public on Jul 17 2024</td>\n",
       "      <td>Jul 17 2023</td>\n",
       "      <td>Jul 17 2024</td>\n",
       "      <td>genomic</td>\n",
       "      <td>1</td>\n",
       "      <td>peripheral whole blood</td>\n",
       "      <td>Homo sapiens</td>\n",
       "      <td>participant_id: 0005</td>\n",
       "      <td>...</td>\n",
       "      <td>0.109718862397854</td>\n",
       "      <td>0</td>\n",
       "      <td>0.505235489278915</td>\n",
       "      <td>KCL</td>\n",
       "      <td>0.0514219148451151</td>\n",
       "      <td>0.0610566697720137</td>\n",
       "      <td>0005</td>\n",
       "      <td>M</td>\n",
       "      <td>1.32031205622569</td>\n",
       "      <td>a</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5 rows × 57 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                                             title geo_accession  \\\n",
       "GSM_ID                                                             \n",
       "GSM7625568  genomic DNA from ID 0003 for 'a' visit    GSM7625568   \n",
       "GSM7625569  genomic DNA from ID 0003 for 'b' visit    GSM7625569   \n",
       "GSM7625570  genomic DNA from ID 0003 for 'd' visit    GSM7625570   \n",
       "GSM7625571  genomic DNA from ID 0003 for 'e' visit    GSM7625571   \n",
       "GSM7625572  genomic DNA from ID 0005 for 'a' visit    GSM7625572   \n",
       "\n",
       "                           status submission_date last_update_date     type  \\\n",
       "GSM_ID                                                                        \n",
       "GSM7625568  Public on Jul 17 2024     Jul 17 2023      Jul 17 2024  genomic   \n",
       "GSM7625569  Public on Jul 17 2024     Jul 17 2023      Jul 17 2024  genomic   \n",
       "GSM7625570  Public on Jul 17 2024     Jul 17 2023      Jul 17 2024  genomic   \n",
       "GSM7625571  Public on Jul 17 2024     Jul 17 2023      Jul 17 2024  genomic   \n",
       "GSM7625572  Public on Jul 17 2024     Jul 17 2023      Jul 17 2024  genomic   \n",
       "\n",
       "           channel_count         source_name_ch1  organism_ch1  \\\n",
       "GSM_ID                                                           \n",
       "GSM7625568             1  peripheral whole blood  Homo sapiens   \n",
       "GSM7625569             1  peripheral whole blood  Homo sapiens   \n",
       "GSM7625570             1  peripheral whole blood  Homo sapiens   \n",
       "GSM7625571             1  peripheral whole blood  Homo sapiens   \n",
       "GSM7625572             1  peripheral whole blood  Homo sapiens   \n",
       "\n",
       "             characteristics_ch1  ...            cd8t:ch1  \\\n",
       "GSM_ID                            ...                       \n",
       "GSM7625568  participant_id: 0003  ...   0.124088032132122   \n",
       "GSM7625569  participant_id: 0003  ...   0.140642508939063   \n",
       "GSM7625570  participant_id: 0003  ...   0.112389247455345   \n",
       "GSM7625571  participant_id: 0003  ...  0.0927185883637476   \n",
       "GSM7625572  participant_id: 0005  ...   0.109718862397854   \n",
       "\n",
       "           days.on.clozapine:ch1           gran:ch1 institute:ch1  \\\n",
       "GSM_ID                                                              \n",
       "GSM7625568                     0  0.623747063903508           KCL   \n",
       "GSM7625569                    42  0.532222309707589           KCL   \n",
       "GSM7625570                    84  0.610861991010505           KCL   \n",
       "GSM7625571                   168  0.647578326303527           KCL   \n",
       "GSM7625572                     0  0.505235489278915           KCL   \n",
       "\n",
       "                      mono:ch1               nk:ch1 participant_id:ch1  \\\n",
       "GSM_ID                                                                   \n",
       "GSM7625568  0.0577366406423333   0.0181886611163226               0003   \n",
       "GSM7625569  0.0794055065754302   0.0165444940880789               0003   \n",
       "GSM7625570  0.0694956639909667  0.00363827604122652               0003   \n",
       "GSM7625571  0.0886719923597006   0.0163329200193279               0003   \n",
       "GSM7625572  0.0514219148451151   0.0610566697720137               0005   \n",
       "\n",
       "           Sex:ch1    smokingscore:ch1 visit:ch1  \n",
       "GSM_ID                                            \n",
       "GSM7625568       M   0.744140048408035         a  \n",
       "GSM7625569       M   0.778521295727892         b  \n",
       "GSM7625570       M   0.591346224352136         d  \n",
       "GSM7625571       M  -0.771644717383253         e  \n",
       "GSM7625572       M    1.32031205622569         a  \n",
       "\n",
       "[5 rows x 57 columns]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Load metadata\n",
    "metadata = pd.read_feather(METADATA_PATH)\n",
    "metadata.set_index(\"GSM_ID\", inplace=True)\n",
    "\n",
    "metadata.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "748ed7c4",
   "metadata": {
    "papermill": {
     "duration": 0.004714,
     "end_time": "2025-03-31T09:47:04.180021",
     "exception": false,
     "start_time": "2025-03-31T09:47:04.175307",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## 3. Load Model and Dependencies"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2e2696e",
   "metadata": {
    "papermill": {
     "duration": 0.004647,
     "end_time": "2025-03-31T09:47:04.189334",
     "exception": false,
     "start_time": "2025-03-31T09:47:04.184687",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "In order to calculate CpGPTGrimAge3, we need to calculate several DNA methylation plasma protein proxies with a finetuned model. The checkpoint is called `proteins` and it predicts 322 plasma protein levels which are normalized with mean 0 and variance 1 (μ = 0, σ² = 1)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ddc6aae",
   "metadata": {
    "papermill": {
     "duration": 0.004652,
     "end_time": "2025-03-31T09:47:04.198681",
     "exception": false,
     "start_time": "2025-03-31T09:47:04.194029",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "### 3.1 Use Downloaded Checkpoint and Configuration Files"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "163f6731",
   "metadata": {
    "papermill": {
     "duration": 0.004843,
     "end_time": "2025-03-31T09:47:04.249630",
     "exception": false,
     "start_time": "2025-03-31T09:47:04.244787",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "### 3.2 Load Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "7f96b1d9",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-31T09:47:04.260102Z",
     "iopub.status.busy": "2025-03-31T09:47:04.259782Z",
     "iopub.status.idle": "2025-03-31T09:47:04.487548Z",
     "shell.execute_reply": "2025-03-31T09:47:04.486991Z"
    },
    "papermill": {
     "duration": 0.233919,
     "end_time": "2025-03-31T09:47:04.488360",
     "exception": false,
     "start_time": "2025-03-31T09:47:04.254441",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[34mcpgpt\u001b[0m\u001b[1m\u001b[0m -\u001b[36mCpGPTInferencer\u001b[0m: \u001b[1mLoaded CpGPT model config.\u001b[0m\n",
      "\u001b[1m\u001b[34mcpgpt\u001b[0m\u001b[1m\u001b[0m -\u001b[36mCpGPTInferencer\u001b[0m: \u001b[1mInstantiated CpGPT model from config.\u001b[0m\n",
      "\u001b[1m\u001b[34mcpgpt\u001b[0m\u001b[1m\u001b[0m -\u001b[36mCpGPTInferencer\u001b[0m: \u001b[1mUsing device: cpu.\u001b[0m\n",
      "\u001b[1m\u001b[34mcpgpt\u001b[0m\u001b[1m\u001b[0m -\u001b[36mCpGPTInferencer\u001b[0m: \u001b[1mLoading checkpoint from: ../dependencies/model/weights/proteins.ckpt\u001b[0m\n",
      "\u001b[1m\u001b[34mcpgpt\u001b[0m\u001b[1m\u001b[0m -\u001b[36mCpGPTInferencer\u001b[0m: \u001b[1mCheckpoint loaded into the model.\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "# Load the model configuration\n",
    "config = inferencer.load_cpgpt_config(MODEL_CONFIG_PATH)\n",
    "\n",
    "# Load the model weights\n",
    "model = inferencer.load_cpgpt_model(\n",
    "    config,\n",
    "    model_ckpt_path=MODEL_CHECKPOINT_PATH,\n",
    "    strict_load=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5cb9d72f",
   "metadata": {},
   "source": [
    "### 3.3 Load Vocab"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4400ee5",
   "metadata": {},
   "source": [
    "The `proteins` model was trained with a vocabulary of about 4,689 CpG sites. Ideally, the data would be filtered to only include those features (or a subset thereof)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "c50bdb94",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the vocabulary\n",
    "with open(MODEL_VOCAB_PATH, \"r\") as f:\n",
    "    vocab = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "3a96ab81",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['cg21830050', 'cg10381813', 'cg08067365', 'cg09864227', 'cg07213830']"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_input_features = vocab['input']\n",
    "\n",
    "model_input_features[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "af4cd5fb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['cpgpt_tnfsf13', 'cpgpt_il33', 'cpgpt_calca', 'cpgpt_npy', 'cpgpt_hla-dra']"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_output_features = vocab['output']\n",
    "\n",
    "model_output_features[:5]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9ef1bab",
   "metadata": {},
   "source": [
    "### 3.4 Use Downloaded Dependencies"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96f32ad6",
   "metadata": {
    "papermill": {
     "duration": 0.005003,
     "end_time": "2025-03-31T09:47:04.498618",
     "exception": false,
     "start_time": "2025-03-31T09:47:04.493615",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## 4. Prepare Data Objects"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4dbb17d8",
   "metadata": {
    "papermill": {
     "duration": 0.004932,
     "end_time": "2025-03-31T09:47:04.508645",
     "exception": false,
     "start_time": "2025-03-31T09:47:04.503713",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "### 4.1 Declare Embedder and Prober"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4d29b20",
   "metadata": {
    "papermill": {
     "duration": 0.005004,
     "end_time": "2025-03-31T09:47:04.518600",
     "exception": false,
     "start_time": "2025-03-31T09:47:04.513596",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "In order to retrieve the sample embeddings, we need to memory-map the data. This is done by using the `CpGPTDataSaver` class. We first need to define the `DNALLMEmbedder` and `IlluminaMethylationProber` classes, which contain the information about the DNA LLM Embeddings and the conversion between Illumina array probes to genomic locations, respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "25958258",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-31T09:47:04.529316Z",
     "iopub.status.busy": "2025-03-31T09:47:04.529031Z",
     "iopub.status.idle": "2025-03-31T09:47:15.919094Z",
     "shell.execute_reply": "2025-03-31T09:47:15.918549Z"
    },
    "papermill": {
     "duration": 11.396311,
     "end_time": "2025-03-31T09:47:15.919844",
     "exception": false,
     "start_time": "2025-03-31T09:47:04.523533",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[34mcpgpt\u001b[0m\u001b[1m\u001b[0m -\u001b[36mDNALLMEmbedder\u001b[0m: \u001b[1mInitializing class DNALLMEmbedder.\u001b[0m\n",
      "\u001b[1m\u001b[34mcpgpt\u001b[0m\u001b[1m\u001b[0m -\u001b[36mDNALLMEmbedder\u001b[0m: \u001b[1mGenome files will be stored under ../dependencies/human/genomes.\u001b[0m\n",
      "\u001b[1m\u001b[34mcpgpt\u001b[0m\u001b[1m\u001b[0m -\u001b[36mDNALLMEmbedder\u001b[0m: \u001b[1mDNA embeddings will be stored under ../dependencies/human/dna_embeddings and subdirectories.\u001b[0m\n",
      "\u001b[1m\u001b[34mcpgpt\u001b[0m\u001b[1m\u001b[0m -\u001b[36mDNALLMEmbedder\u001b[0m: \u001b[1mEnsembl metadata dictionary loaded successfully\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "embedder = DNALLMEmbedder(dependencies_dir=LLM_DEPENDENCIES_DIR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "4d655c8f",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-31T09:47:15.931727Z",
     "iopub.status.busy": "2025-03-31T09:47:15.931336Z",
     "iopub.status.idle": "2025-03-31T09:47:16.406092Z",
     "shell.execute_reply": "2025-03-31T09:47:16.405524Z"
    },
    "papermill": {
     "duration": 0.481207,
     "end_time": "2025-03-31T09:47:16.406869",
     "exception": false,
     "start_time": "2025-03-31T09:47:15.925662",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[34mcpgpt\u001b[0m\u001b[1m\u001b[0m -\u001b[36mIlluminaMethylationProber\u001b[0m: \u001b[1mInitializing class IlluminaMethylationProber.\u001b[0m\n",
      "\u001b[1m\u001b[34mcpgpt\u001b[0m\u001b[1m\u001b[0m -\u001b[36mIlluminaMethylationProber\u001b[0m: \u001b[1mIllumina methylation manifest files will be stored under ../dependencies/human/manifests.\u001b[0m\n",
      "\u001b[1m\u001b[34mcpgpt\u001b[0m\u001b[1m\u001b[0m -\u001b[36mIlluminaMethylationProber\u001b[0m: \u001b[1mIllumina metadata dictionary loaded successfully.\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "prober = IlluminaMethylationProber(dependencies_dir=LLM_DEPENDENCIES_DIR, embedder=embedder)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ccbe37e9",
   "metadata": {},
   "source": [
    "### 4.2 Filter Vocab"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "40e82ed8",
   "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>cg26259818</th>\n",
       "      <th>cg26866325</th>\n",
       "      <th>cg23263937</th>\n",
       "      <th>cg15779600</th>\n",
       "      <th>cg16075139</th>\n",
       "      <th>cg04820362</th>\n",
       "      <th>cg11782409</th>\n",
       "      <th>cg07870237</th>\n",
       "      <th>cg01802397</th>\n",
       "      <th>cg04634427</th>\n",
       "      <th>...</th>\n",
       "      <th>cg01431830</th>\n",
       "      <th>cg18241647</th>\n",
       "      <th>cg14009688</th>\n",
       "      <th>cg26866482</th>\n",
       "      <th>cg23191950</th>\n",
       "      <th>cg03998338</th>\n",
       "      <th>cg06532212</th>\n",
       "      <th>cg23268677</th>\n",
       "      <th>cg20405584</th>\n",
       "      <th>cg14984434</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>GSM_ID</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></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>GSM7625568</th>\n",
       "      <td>0.900749</td>\n",
       "      <td>0.199913</td>\n",
       "      <td>0.107259</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.819918</td>\n",
       "      <td>0.903952</td>\n",
       "      <td>0.252276</td>\n",
       "      <td>0.049561</td>\n",
       "      <td>0.781003</td>\n",
       "      <td>...</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.940283</td>\n",
       "      <td>0.644005</td>\n",
       "      <td>0.047353</td>\n",
       "      <td>0.251154</td>\n",
       "      <td>0.916748</td>\n",
       "      <td>0.488639</td>\n",
       "      <td>0.063751</td>\n",
       "      <td>0.027925</td>\n",
       "      <td>0.946631</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>GSM7625569</th>\n",
       "      <td>0.884256</td>\n",
       "      <td>0.148837</td>\n",
       "      <td>0.118698</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.877955</td>\n",
       "      <td>0.885812</td>\n",
       "      <td>0.286435</td>\n",
       "      <td>0.055951</td>\n",
       "      <td>0.736555</td>\n",
       "      <td>...</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.932127</td>\n",
       "      <td>0.624015</td>\n",
       "      <td>0.044781</td>\n",
       "      <td>0.184450</td>\n",
       "      <td>0.921210</td>\n",
       "      <td>0.510761</td>\n",
       "      <td>0.057729</td>\n",
       "      <td>0.028602</td>\n",
       "      <td>0.928797</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>GSM7625570</th>\n",
       "      <td>0.869737</td>\n",
       "      <td>0.174457</td>\n",
       "      <td>0.083522</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.885510</td>\n",
       "      <td>0.885832</td>\n",
       "      <td>0.298032</td>\n",
       "      <td>0.058532</td>\n",
       "      <td>0.709248</td>\n",
       "      <td>...</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.932670</td>\n",
       "      <td>0.684649</td>\n",
       "      <td>0.044844</td>\n",
       "      <td>0.177397</td>\n",
       "      <td>0.916927</td>\n",
       "      <td>0.503245</td>\n",
       "      <td>0.073991</td>\n",
       "      <td>0.031689</td>\n",
       "      <td>0.933811</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>GSM7625571</th>\n",
       "      <td>0.902602</td>\n",
       "      <td>0.196606</td>\n",
       "      <td>0.081571</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.879783</td>\n",
       "      <td>0.891532</td>\n",
       "      <td>0.255523</td>\n",
       "      <td>0.056742</td>\n",
       "      <td>0.780125</td>\n",
       "      <td>...</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.950026</td>\n",
       "      <td>0.616295</td>\n",
       "      <td>0.047359</td>\n",
       "      <td>0.171478</td>\n",
       "      <td>0.940402</td>\n",
       "      <td>0.496729</td>\n",
       "      <td>0.073860</td>\n",
       "      <td>0.035205</td>\n",
       "      <td>0.942940</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>GSM7625572</th>\n",
       "      <td>0.899257</td>\n",
       "      <td>0.153354</td>\n",
       "      <td>0.107341</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.829583</td>\n",
       "      <td>0.864145</td>\n",
       "      <td>0.282356</td>\n",
       "      <td>0.060830</td>\n",
       "      <td>0.691230</td>\n",
       "      <td>...</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.958417</td>\n",
       "      <td>0.649126</td>\n",
       "      <td>0.043928</td>\n",
       "      <td>0.232675</td>\n",
       "      <td>0.752656</td>\n",
       "      <td>0.506659</td>\n",
       "      <td>0.065188</td>\n",
       "      <td>0.047482</td>\n",
       "      <td>0.937227</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5 rows × 4689 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "            cg26259818  cg26866325  cg23263937  cg15779600  cg16075139  \\\n",
       "GSM_ID                                                                   \n",
       "GSM7625568    0.900749    0.199913    0.107259         NaN         NaN   \n",
       "GSM7625569    0.884256    0.148837    0.118698         NaN         NaN   \n",
       "GSM7625570    0.869737    0.174457    0.083522         NaN         NaN   \n",
       "GSM7625571    0.902602    0.196606    0.081571         NaN         NaN   \n",
       "GSM7625572    0.899257    0.153354    0.107341         NaN         NaN   \n",
       "\n",
       "            cg04820362  cg11782409  cg07870237  cg01802397  cg04634427  ...  \\\n",
       "GSM_ID                                                                  ...   \n",
       "GSM7625568    0.819918    0.903952    0.252276    0.049561    0.781003  ...   \n",
       "GSM7625569    0.877955    0.885812    0.286435    0.055951    0.736555  ...   \n",
       "GSM7625570    0.885510    0.885832    0.298032    0.058532    0.709248  ...   \n",
       "GSM7625571    0.879783    0.891532    0.255523    0.056742    0.780125  ...   \n",
       "GSM7625572    0.829583    0.864145    0.282356    0.060830    0.691230  ...   \n",
       "\n",
       "            cg01431830  cg18241647  cg14009688  cg26866482  cg23191950  \\\n",
       "GSM_ID                                                                   \n",
       "GSM7625568         NaN    0.940283    0.644005    0.047353    0.251154   \n",
       "GSM7625569         NaN    0.932127    0.624015    0.044781    0.184450   \n",
       "GSM7625570         NaN    0.932670    0.684649    0.044844    0.177397   \n",
       "GSM7625571         NaN    0.950026    0.616295    0.047359    0.171478   \n",
       "GSM7625572         NaN    0.958417    0.649126    0.043928    0.232675   \n",
       "\n",
       "            cg03998338  cg06532212  cg23268677  cg20405584  cg14984434  \n",
       "GSM_ID                                                                  \n",
       "GSM7625568    0.916748    0.488639    0.063751    0.027925    0.946631  \n",
       "GSM7625569    0.921210    0.510761    0.057729    0.028602    0.928797  \n",
       "GSM7625570    0.916927    0.503245    0.073991    0.031689    0.933811  \n",
       "GSM7625571    0.940402    0.496729    0.073860    0.035205    0.942940  \n",
       "GSM7625572    0.752656    0.506659    0.065188    0.047482    0.937227  \n",
       "\n",
       "[5 rows x 4689 columns]"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "common_features = list(set(model_input_features) & set(df.columns))\n",
    "df_filtered = df.loc[:, common_features]\n",
    "df_filtered.to_feather(FILTERED_BETAS_PATH)\n",
    "\n",
    "df_filtered.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06bae259",
   "metadata": {
    "papermill": {
     "duration": 0.033169,
     "end_time": "2025-03-31T09:47:16.445679",
     "exception": false,
     "start_time": "2025-03-31T09:47:16.412510",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "### 4.3 Memory-Map Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "7b6f1516",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-31T09:47:16.457082Z",
     "iopub.status.busy": "2025-03-31T09:47:16.456921Z",
     "iopub.status.idle": "2025-03-31T09:47:16.478389Z",
     "shell.execute_reply": "2025-03-31T09:47:16.477969Z"
    },
    "papermill": {
     "duration": 0.028039,
     "end_time": "2025-03-31T09:47:16.479126",
     "exception": false,
     "start_time": "2025-03-31T09:47:16.451087",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[34mcpgpt\u001b[0m\u001b[1m\u001b[0m -\u001b[36mCpGPTDataSaver\u001b[0m: \u001b[1mInitializing class CpGPTDataSaver.\u001b[0m\n",
      "\u001b[1m\u001b[34mcpgpt\u001b[0m\u001b[1m\u001b[0m -\u001b[36mCpGPTDataSaver\u001b[0m: \u001b[1mDataset folders will be stored under ../data/tutorials/processed/predict_mortality.\u001b[0m\n",
      "\u001b[1m\u001b[34mcpgpt\u001b[0m\u001b[1m\u001b[0m -\u001b[36mCpGPTDataSaver\u001b[0m: \u001b[1mLoaded existing dataset metrics.\u001b[0m\n",
      "\u001b[1m\u001b[34mcpgpt\u001b[0m\u001b[1m\u001b[0m -\u001b[36mCpGPTDataSaver\u001b[0m: \u001b[1mLoaded existing genomic locations.\u001b[0m\n",
      "\u001b[1m\u001b[34mcpgpt\u001b[0m\u001b[1m\u001b[0m -\u001b[36mCpGPTDataSaver\u001b[0m: \u001b[1mStarting file processing.\u001b[0m\n",
      "\u001b[1m\u001b[34mcpgpt\u001b[0m\u001b[1m\u001b[0m -\u001b[36mCpGPTDataSaver\u001b[0m: \u001b[1m1 files already processed. Skipping those.\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "# Define datasaver\n",
    "datasaver = CpGPTDataSaver(data_paths=FILTERED_BETAS_PATH, processed_dir=PROCESSED_DIR)\n",
    "\n",
    "# Process the file\n",
    "datasaver.process_files(prober, embedder)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7c6a8fd",
   "metadata": {
    "papermill": {
     "duration": 0.005685,
     "end_time": "2025-03-31T09:47:16.522951",
     "exception": false,
     "start_time": "2025-03-31T09:47:16.517266",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "### 4.4 Declare data module"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47981f22",
   "metadata": {
    "papermill": {
     "duration": 0.005645,
     "end_time": "2025-03-31T09:47:16.534299",
     "exception": false,
     "start_time": "2025-03-31T09:47:16.528654",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Let's define one data module to use with our model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "61e64ea6",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-31T09:47:16.546311Z",
     "iopub.status.busy": "2025-03-31T09:47:16.546147Z",
     "iopub.status.idle": "2025-03-31T09:47:27.450125Z",
     "shell.execute_reply": "2025-03-31T09:47:27.449670Z"
    },
    "papermill": {
     "duration": 10.910884,
     "end_time": "2025-03-31T09:47:27.450876",
     "exception": false,
     "start_time": "2025-03-31T09:47:16.539992",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[34mcpgpt\u001b[0m\u001b[1m\u001b[0m -\u001b[36mDNALLMEmbedder\u001b[0m: \u001b[1mInitializing class DNALLMEmbedder.\u001b[0m\n",
      "\u001b[1m\u001b[34mcpgpt\u001b[0m\u001b[1m\u001b[0m -\u001b[36mDNALLMEmbedder\u001b[0m: \u001b[1mGenome files will be stored under ../dependencies/human/genomes.\u001b[0m\n",
      "\u001b[1m\u001b[34mcpgpt\u001b[0m\u001b[1m\u001b[0m -\u001b[36mDNALLMEmbedder\u001b[0m: \u001b[1mDNA embeddings will be stored under ../dependencies/human/dna_embeddings and subdirectories.\u001b[0m\n",
      "\u001b[1m\u001b[34mcpgpt\u001b[0m\u001b[1m\u001b[0m -\u001b[36mDNALLMEmbedder\u001b[0m: \u001b[1mEnsembl metadata dictionary loaded successfully\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "# Define datamodule\n",
    "datamodule = CpGPTDataModule(\n",
    "    predict_dir=PROCESSED_DIR,\n",
    "    dependencies_dir=LLM_DEPENDENCIES_DIR,\n",
    "    batch_size=1,\n",
    "    num_workers=0,\n",
    "    max_length=MAX_INPUT_LENGTH,\n",
    "    dna_llm=config.data.dna_llm,\n",
    "    dna_context_len=config.data.dna_context_len,\n",
    "    sorting_strategy=config.data.sorting_strategy,\n",
    "    pin_memory=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0909773a",
   "metadata": {
    "papermill": {
     "duration": 0.006123,
     "end_time": "2025-03-31T09:47:38.472302",
     "exception": false,
     "start_time": "2025-03-31T09:47:38.466179",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## 5. Compute Protein Proxies"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c9c3b8f",
   "metadata": {
    "papermill": {
     "duration": 0.006032,
     "end_time": "2025-03-31T09:47:38.484506",
     "exception": false,
     "start_time": "2025-03-31T09:47:38.478474",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "### 5.1 Declare Trainer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84e46211",
   "metadata": {
    "papermill": {
     "duration": 0.006033,
     "end_time": "2025-03-31T09:47:38.496656",
     "exception": false,
     "start_time": "2025-03-31T09:47:38.490623",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Given all models were trained under mixed precision, we'll use the `precision=\"16-mixed\"` argument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "85e45335",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-31T09:47:38.509682Z",
     "iopub.status.busy": "2025-03-31T09:47:38.509296Z",
     "iopub.status.idle": "2025-03-31T09:47:38.543734Z",
     "shell.execute_reply": "2025-03-31T09:47:38.543291Z"
    },
    "papermill": {
     "duration": 0.041764,
     "end_time": "2025-03-31T09:47:38.544478",
     "exception": false,
     "start_time": "2025-03-31T09:47:38.502714",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using 16bit Automatic Mixed Precision (AMP)\n",
      "/Users/lucascamillo/mambaforge/envs/cpgpt/lib/python3.10/site-packages/torch/amp/grad_scaler.py:132: UserWarning: torch.cuda.amp.GradScaler is enabled, but CUDA is not available.  Disabling.\n",
      "  warnings.warn(\n",
      "GPU available: True (mps), used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "HPU available: False, using: 0 HPUs\n"
     ]
    }
   ],
   "source": [
    "trainer = CpGPTTrainer(precision=\"16-mixed\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c8f7f7e",
   "metadata": {
    "papermill": {
     "duration": 0.006249,
     "end_time": "2025-03-31T09:47:38.556896",
     "exception": false,
     "start_time": "2025-03-31T09:47:38.550647",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "### 5.2 Get Predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "7111b26b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-03-31T09:47:38.570391Z",
     "iopub.status.busy": "2025-03-31T09:47:38.569946Z",
     "iopub.status.idle": "2025-03-31T09:47:41.671896Z",
     "shell.execute_reply": "2025-03-31T09:47:41.671452Z"
    },
    "papermill": {
     "duration": 3.109451,
     "end_time": "2025-03-31T09:47:41.672639",
     "exception": false,
     "start_time": "2025-03-31T09:47:38.563188",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[34mcpgpt\u001b[0m\u001b[1m\u001b[0m -\u001b[36mCpGPTDataset\u001b[0m: \u001b[1mInitializing class CpGPTDataset.\u001b[0m\n",
      "\u001b[1m\u001b[34mcpgpt\u001b[0m\u001b[1m\u001b[0m -\u001b[36mCpGPTDataset\u001b[0m: \u001b[1mLoaded existing dataset metrics.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ca986c0cec2b4a42b859d76fcdbd4af3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/lucascamillo/mambaforge/envs/cpgpt/lib/python3.10/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:424: The 'predict_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=7` in the `DataLoader` to improve performance.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">/Users/lucascamillo/mambaforge/envs/cpgpt/lib/python3.10/site-packages/torch/amp/autocast_mode.py:265: UserWarning:\n",
       "User provided device_type of 'cuda', but CUDA is not available. Disabling\n",
       "  warnings.warn(\n",
       "</pre>\n"
      ],
      "text/plain": [
       "/Users/lucascamillo/mambaforge/envs/cpgpt/lib/python3.10/site-packages/torch/amp/autocast_mode.py:265: UserWarning:\n",
       "User provided device_type of 'cuda', but CUDA is not available. Disabling\n",
       "  warnings.warn(\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"></pre>\n"
      ],
      "text/plain": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "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>cpgpt_tnfsf13</th>\n",
       "      <th>cpgpt_il33</th>\n",
       "      <th>cpgpt_calca</th>\n",
       "      <th>cpgpt_npy</th>\n",
       "      <th>cpgpt_hla-dra</th>\n",
       "      <th>cpgpt_c1qa</th>\n",
       "      <th>cpgpt_fth1</th>\n",
       "      <th>cpgpt_s100b</th>\n",
       "      <th>cpgpt_ceacam5</th>\n",
       "      <th>cpgpt_mme</th>\n",
       "      <th>...</th>\n",
       "      <th>cpgpt_ccl15</th>\n",
       "      <th>cpgpt_ccl14</th>\n",
       "      <th>cpgpt_ccl13</th>\n",
       "      <th>cpgpt_ccl11</th>\n",
       "      <th>cpgpt_ccl1</th>\n",
       "      <th>cpgpt_saa1</th>\n",
       "      <th>cpgpt_s100a9</th>\n",
       "      <th>cpgpt_s100a12</th>\n",
       "      <th>cpgpt_bdnf</th>\n",
       "      <th>cpgpt_vgf</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>GSM_ID</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></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>GSM7625568</th>\n",
       "      <td>0.159939</td>\n",
       "      <td>-0.109247</td>\n",
       "      <td>-0.123497</td>\n",
       "      <td>0.268346</td>\n",
       "      <td>0.311079</td>\n",
       "      <td>-0.104481</td>\n",
       "      <td>0.307234</td>\n",
       "      <td>0.013596</td>\n",
       "      <td>-0.029327</td>\n",
       "      <td>0.266559</td>\n",
       "      <td>...</td>\n",
       "      <td>-0.100315</td>\n",
       "      <td>0.267266</td>\n",
       "      <td>0.010796</td>\n",
       "      <td>0.257249</td>\n",
       "      <td>0.012845</td>\n",
       "      <td>-0.144005</td>\n",
       "      <td>-0.199204</td>\n",
       "      <td>-0.288781</td>\n",
       "      <td>0.283817</td>\n",
       "      <td>-0.094640</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>GSM7625569</th>\n",
       "      <td>0.047051</td>\n",
       "      <td>-0.255467</td>\n",
       "      <td>-0.281614</td>\n",
       "      <td>0.243430</td>\n",
       "      <td>0.189521</td>\n",
       "      <td>-0.217191</td>\n",
       "      <td>0.198763</td>\n",
       "      <td>0.019737</td>\n",
       "      <td>-0.110035</td>\n",
       "      <td>0.173720</td>\n",
       "      <td>...</td>\n",
       "      <td>-0.218816</td>\n",
       "      <td>0.086749</td>\n",
       "      <td>-0.205998</td>\n",
       "      <td>0.042463</td>\n",
       "      <td>-0.109191</td>\n",
       "      <td>-0.249524</td>\n",
       "      <td>-0.164367</td>\n",
       "      <td>-0.338399</td>\n",
       "      <td>0.262606</td>\n",
       "      <td>-0.106812</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>GSM7625570</th>\n",
       "      <td>0.050562</td>\n",
       "      <td>-0.225324</td>\n",
       "      <td>-0.236831</td>\n",
       "      <td>0.222035</td>\n",
       "      <td>0.186863</td>\n",
       "      <td>-0.197240</td>\n",
       "      <td>0.215832</td>\n",
       "      <td>-0.016688</td>\n",
       "      <td>-0.087977</td>\n",
       "      <td>0.180057</td>\n",
       "      <td>...</td>\n",
       "      <td>-0.197247</td>\n",
       "      <td>0.127577</td>\n",
       "      <td>-0.130641</td>\n",
       "      <td>0.094274</td>\n",
       "      <td>-0.094252</td>\n",
       "      <td>-0.226423</td>\n",
       "      <td>-0.189340</td>\n",
       "      <td>-0.316972</td>\n",
       "      <td>0.234579</td>\n",
       "      <td>-0.125441</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>GSM7625571</th>\n",
       "      <td>0.194388</td>\n",
       "      <td>-0.103120</td>\n",
       "      <td>-0.060972</td>\n",
       "      <td>0.246503</td>\n",
       "      <td>0.228984</td>\n",
       "      <td>-0.092588</td>\n",
       "      <td>0.311946</td>\n",
       "      <td>-0.041283</td>\n",
       "      <td>-0.011208</td>\n",
       "      <td>0.265450</td>\n",
       "      <td>...</td>\n",
       "      <td>-0.103527</td>\n",
       "      <td>0.264676</td>\n",
       "      <td>0.018594</td>\n",
       "      <td>0.178806</td>\n",
       "      <td>0.027473</td>\n",
       "      <td>-0.150791</td>\n",
       "      <td>-0.140418</td>\n",
       "      <td>-0.208958</td>\n",
       "      <td>0.266402</td>\n",
       "      <td>-0.067963</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>GSM7625572</th>\n",
       "      <td>0.017105</td>\n",
       "      <td>-0.295679</td>\n",
       "      <td>-0.310758</td>\n",
       "      <td>0.232463</td>\n",
       "      <td>0.167696</td>\n",
       "      <td>-0.245472</td>\n",
       "      <td>0.177904</td>\n",
       "      <td>0.006786</td>\n",
       "      <td>-0.177094</td>\n",
       "      <td>0.163067</td>\n",
       "      <td>...</td>\n",
       "      <td>-0.246631</td>\n",
       "      <td>0.001140</td>\n",
       "      <td>-0.265735</td>\n",
       "      <td>-0.007235</td>\n",
       "      <td>-0.135532</td>\n",
       "      <td>-0.271881</td>\n",
       "      <td>-0.144465</td>\n",
       "      <td>-0.389837</td>\n",
       "      <td>0.226135</td>\n",
       "      <td>-0.125433</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5 rows × 322 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "            cpgpt_tnfsf13  cpgpt_il33  cpgpt_calca  cpgpt_npy  cpgpt_hla-dra  \\\n",
       "GSM_ID                                                                         \n",
       "GSM7625568       0.159939   -0.109247    -0.123497   0.268346       0.311079   \n",
       "GSM7625569       0.047051   -0.255467    -0.281614   0.243430       0.189521   \n",
       "GSM7625570       0.050562   -0.225324    -0.236831   0.222035       0.186863   \n",
       "GSM7625571       0.194388   -0.103120    -0.060972   0.246503       0.228984   \n",
       "GSM7625572       0.017105   -0.295679    -0.310758   0.232463       0.167696   \n",
       "\n",
       "            cpgpt_c1qa  cpgpt_fth1  cpgpt_s100b  cpgpt_ceacam5  cpgpt_mme  \\\n",
       "GSM_ID                                                                      \n",
       "GSM7625568   -0.104481    0.307234     0.013596      -0.029327   0.266559   \n",
       "GSM7625569   -0.217191    0.198763     0.019737      -0.110035   0.173720   \n",
       "GSM7625570   -0.197240    0.215832    -0.016688      -0.087977   0.180057   \n",
       "GSM7625571   -0.092588    0.311946    -0.041283      -0.011208   0.265450   \n",
       "GSM7625572   -0.245472    0.177904     0.006786      -0.177094   0.163067   \n",
       "\n",
       "            ...  cpgpt_ccl15  cpgpt_ccl14  cpgpt_ccl13  cpgpt_ccl11  \\\n",
       "GSM_ID      ...                                                       \n",
       "GSM7625568  ...    -0.100315     0.267266     0.010796     0.257249   \n",
       "GSM7625569  ...    -0.218816     0.086749    -0.205998     0.042463   \n",
       "GSM7625570  ...    -0.197247     0.127577    -0.130641     0.094274   \n",
       "GSM7625571  ...    -0.103527     0.264676     0.018594     0.178806   \n",
       "GSM7625572  ...    -0.246631     0.001140    -0.265735    -0.007235   \n",
       "\n",
       "            cpgpt_ccl1  cpgpt_saa1  cpgpt_s100a9  cpgpt_s100a12  cpgpt_bdnf  \\\n",
       "GSM_ID                                                                        \n",
       "GSM7625568    0.012845   -0.144005     -0.199204      -0.288781    0.283817   \n",
       "GSM7625569   -0.109191   -0.249524     -0.164367      -0.338399    0.262606   \n",
       "GSM7625570   -0.094252   -0.226423     -0.189340      -0.316972    0.234579   \n",
       "GSM7625571    0.027473   -0.150791     -0.140418      -0.208958    0.266402   \n",
       "GSM7625572   -0.135532   -0.271881     -0.144465      -0.389837    0.226135   \n",
       "\n",
       "            cpgpt_vgf  \n",
       "GSM_ID                 \n",
       "GSM7625568  -0.094640  \n",
       "GSM7625569  -0.106812  \n",
       "GSM7625570  -0.125441  \n",
       "GSM7625571  -0.067963  \n",
       "GSM7625572  -0.125433  \n",
       "\n",
       "[5 rows x 322 columns]"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Get the target sample embeddings\n",
    "forward_pass = trainer.predict(\n",
    "    model=model,\n",
    "    datamodule=datamodule,\n",
    "    predict_mode=\"forward\",\n",
    "    return_keys=[\"pred_conditions\"]\n",
    ")\n",
    "\n",
    "pred_conditions_df = pd.DataFrame(forward_pass['pred_conditions'], index=df.index, columns=model_output_features)\n",
    "pred_conditions_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "910361b8",
   "metadata": {},
   "source": [
    "## 6. Calculate CpGPTGrimAge3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a0e0fde",
   "metadata": {},
   "source": [
    "In the last step, we need to join together all features that are necessary to calculate CpGPTGrimAge3, namely:\n",
    "- age: chronological age of the sample;\n",
    "- GrimAge2 proxies: protein and lifestyle proxies from GrimAge version 2; \n",
    "- CpGPT protein proxies: protein levels predicted with CpGPT."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dfe7cfa9",
   "metadata": {},
   "source": [
    "### Join All Features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9746cb86",
   "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>grimage2timp1</th>\n",
       "      <th>grimage2packyrs</th>\n",
       "      <th>grimage2logcrp</th>\n",
       "      <th>grimage2b2m</th>\n",
       "      <th>grimage2adm</th>\n",
       "      <th>grimage2leptin</th>\n",
       "      <th>grimage2gdf15</th>\n",
       "      <th>cpgpt_tnfsf13</th>\n",
       "      <th>cpgpt_il33</th>\n",
       "      <th>...</th>\n",
       "      <th>cpgpt_ccl15</th>\n",
       "      <th>cpgpt_ccl14</th>\n",
       "      <th>cpgpt_ccl13</th>\n",
       "      <th>cpgpt_ccl11</th>\n",
       "      <th>cpgpt_ccl1</th>\n",
       "      <th>cpgpt_saa1</th>\n",
       "      <th>cpgpt_s100a9</th>\n",
       "      <th>cpgpt_s100a12</th>\n",
       "      <th>cpgpt_bdnf</th>\n",
       "      <th>cpgpt_vgf</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>GSM_ID</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></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>GSM7625568</th>\n",
       "      <td>27.713889</td>\n",
       "      <td>33443.488165</td>\n",
       "      <td>4.889224</td>\n",
       "      <td>-0.207662</td>\n",
       "      <td>1.657720e+06</td>\n",
       "      <td>309.088007</td>\n",
       "      <td>6452.711603</td>\n",
       "      <td>525.537382</td>\n",
       "      <td>0.159939</td>\n",
       "      <td>-0.109247</td>\n",
       "      <td>...</td>\n",
       "      <td>-0.100315</td>\n",
       "      <td>0.267266</td>\n",
       "      <td>0.010796</td>\n",
       "      <td>0.257249</td>\n",
       "      <td>0.012845</td>\n",
       "      <td>-0.144005</td>\n",
       "      <td>-0.199204</td>\n",
       "      <td>-0.288781</td>\n",
       "      <td>0.283817</td>\n",
       "      <td>-0.094640</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>GSM7625569</th>\n",
       "      <td>27.713889</td>\n",
       "      <td>33603.182758</td>\n",
       "      <td>-0.675848</td>\n",
       "      <td>0.151504</td>\n",
       "      <td>1.679469e+06</td>\n",
       "      <td>319.938681</td>\n",
       "      <td>5657.270706</td>\n",
       "      <td>550.893628</td>\n",
       "      <td>0.047051</td>\n",
       "      <td>-0.255467</td>\n",
       "      <td>...</td>\n",
       "      <td>-0.218816</td>\n",
       "      <td>0.086749</td>\n",
       "      <td>-0.205998</td>\n",
       "      <td>0.042463</td>\n",
       "      <td>-0.109191</td>\n",
       "      <td>-0.249524</td>\n",
       "      <td>-0.164367</td>\n",
       "      <td>-0.338399</td>\n",
       "      <td>0.262606</td>\n",
       "      <td>-0.106812</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>GSM7625570</th>\n",
       "      <td>27.713889</td>\n",
       "      <td>33807.609589</td>\n",
       "      <td>3.680232</td>\n",
       "      <td>-0.113398</td>\n",
       "      <td>1.682385e+06</td>\n",
       "      <td>312.595083</td>\n",
       "      <td>6654.595126</td>\n",
       "      <td>545.927825</td>\n",
       "      <td>0.050562</td>\n",
       "      <td>-0.225324</td>\n",
       "      <td>...</td>\n",
       "      <td>-0.197247</td>\n",
       "      <td>0.127577</td>\n",
       "      <td>-0.130641</td>\n",
       "      <td>0.094274</td>\n",
       "      <td>-0.094252</td>\n",
       "      <td>-0.226423</td>\n",
       "      <td>-0.189340</td>\n",
       "      <td>-0.316972</td>\n",
       "      <td>0.234579</td>\n",
       "      <td>-0.125441</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>GSM7625571</th>\n",
       "      <td>27.713889</td>\n",
       "      <td>33880.497852</td>\n",
       "      <td>3.483722</td>\n",
       "      <td>0.168095</td>\n",
       "      <td>1.680143e+06</td>\n",
       "      <td>317.126936</td>\n",
       "      <td>6792.263369</td>\n",
       "      <td>531.686439</td>\n",
       "      <td>0.194388</td>\n",
       "      <td>-0.103120</td>\n",
       "      <td>...</td>\n",
       "      <td>-0.103527</td>\n",
       "      <td>0.264676</td>\n",
       "      <td>0.018594</td>\n",
       "      <td>0.178806</td>\n",
       "      <td>0.027473</td>\n",
       "      <td>-0.150791</td>\n",
       "      <td>-0.140418</td>\n",
       "      <td>-0.208958</td>\n",
       "      <td>0.266402</td>\n",
       "      <td>-0.067963</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>GSM7625572</th>\n",
       "      <td>23.183333</td>\n",
       "      <td>33192.500811</td>\n",
       "      <td>4.700933</td>\n",
       "      <td>-0.575709</td>\n",
       "      <td>1.590738e+06</td>\n",
       "      <td>317.324722</td>\n",
       "      <td>4798.638056</td>\n",
       "      <td>522.853005</td>\n",
       "      <td>0.017105</td>\n",
       "      <td>-0.295679</td>\n",
       "      <td>...</td>\n",
       "      <td>-0.246631</td>\n",
       "      <td>0.001140</td>\n",
       "      <td>-0.265735</td>\n",
       "      <td>-0.007235</td>\n",
       "      <td>-0.135532</td>\n",
       "      <td>-0.271881</td>\n",
       "      <td>-0.144465</td>\n",
       "      <td>-0.389837</td>\n",
       "      <td>0.226135</td>\n",
       "      <td>-0.125433</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5 rows × 330 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                  age  grimage2timp1  grimage2packyrs  grimage2logcrp  \\\n",
       "GSM_ID                                                                  \n",
       "GSM7625568  27.713889   33443.488165         4.889224       -0.207662   \n",
       "GSM7625569  27.713889   33603.182758        -0.675848        0.151504   \n",
       "GSM7625570  27.713889   33807.609589         3.680232       -0.113398   \n",
       "GSM7625571  27.713889   33880.497852         3.483722        0.168095   \n",
       "GSM7625572  23.183333   33192.500811         4.700933       -0.575709   \n",
       "\n",
       "             grimage2b2m  grimage2adm  grimage2leptin  grimage2gdf15  \\\n",
       "GSM_ID                                                                 \n",
       "GSM7625568  1.657720e+06   309.088007     6452.711603     525.537382   \n",
       "GSM7625569  1.679469e+06   319.938681     5657.270706     550.893628   \n",
       "GSM7625570  1.682385e+06   312.595083     6654.595126     545.927825   \n",
       "GSM7625571  1.680143e+06   317.126936     6792.263369     531.686439   \n",
       "GSM7625572  1.590738e+06   317.324722     4798.638056     522.853005   \n",
       "\n",
       "            cpgpt_tnfsf13  cpgpt_il33  ...  cpgpt_ccl15  cpgpt_ccl14  \\\n",
       "GSM_ID                                 ...                             \n",
       "GSM7625568       0.159939   -0.109247  ...    -0.100315     0.267266   \n",
       "GSM7625569       0.047051   -0.255467  ...    -0.218816     0.086749   \n",
       "GSM7625570       0.050562   -0.225324  ...    -0.197247     0.127577   \n",
       "GSM7625571       0.194388   -0.103120  ...    -0.103527     0.264676   \n",
       "GSM7625572       0.017105   -0.295679  ...    -0.246631     0.001140   \n",
       "\n",
       "            cpgpt_ccl13  cpgpt_ccl11  cpgpt_ccl1  cpgpt_saa1  cpgpt_s100a9  \\\n",
       "GSM_ID                                                                       \n",
       "GSM7625568     0.010796     0.257249    0.012845   -0.144005     -0.199204   \n",
       "GSM7625569    -0.205998     0.042463   -0.109191   -0.249524     -0.164367   \n",
       "GSM7625570    -0.130641     0.094274   -0.094252   -0.226423     -0.189340   \n",
       "GSM7625571     0.018594     0.178806    0.027473   -0.150791     -0.140418   \n",
       "GSM7625572    -0.265735    -0.007235   -0.135532   -0.271881     -0.144465   \n",
       "\n",
       "            cpgpt_s100a12  cpgpt_bdnf  cpgpt_vgf  \n",
       "GSM_ID                                            \n",
       "GSM7625568      -0.288781    0.283817  -0.094640  \n",
       "GSM7625569      -0.338399    0.262606  -0.106812  \n",
       "GSM7625570      -0.316972    0.234579  -0.125441  \n",
       "GSM7625571      -0.208958    0.266402  -0.067963  \n",
       "GSM7625572      -0.389837    0.226135  -0.125433  \n",
       "\n",
       "[5 rows x 330 columns]"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Get age from the metadata\n",
    "age = metadata.loc[:, ['age:ch1']].astype(float)\n",
    "age.columns = ['age']\n",
    "\n",
    "# Add age to the filtered betas\n",
    "df_filtered['age'] = metadata.loc[:, ['age:ch1']].astype(float)\n",
    "\n",
    "# Get GrimAge2 proxies\n",
    "grimage2_proxies = [\n",
    "    \"grimage2timp1\",\n",
    "    \"grimage2packyrs\",\n",
    "    \"grimage2logcrp\",\n",
    "    \"grimage2b2m\",\n",
    "    \"grimage2adm\",\n",
    "    \"grimage2leptin\",\n",
    "    \"grimage2gdf15\",\n",
    "    \"grimage2pai1\",\n",
    "]\n",
    "adata_grimage2 = pya.pp.df_to_adata(df_filtered, verbose=False)\n",
    "pya.pred.predict_age(adata_grimage2, clock_names=grimage2_proxies, verbose=False)\n",
    "\n",
    "# Combine all features\n",
    "combined_df = pd.concat([age, adata_grimage2.obs, pred_conditions_df], axis=1)\n",
    "\n",
    "combined_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "807a3c25",
   "metadata": {},
   "source": [
    "### 6.2 Calculate CpGPTGrimAge3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "e67be17b",
   "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>cpgptgrimage3</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>GSM_ID</th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>GSM7625568</th>\n",
       "      <td>35.866626</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>GSM7625569</th>\n",
       "      <td>34.805627</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>GSM7625570</th>\n",
       "      <td>34.625747</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>GSM7625571</th>\n",
       "      <td>37.634215</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>GSM7625572</th>\n",
       "      <td>30.980501</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "            cpgptgrimage3\n",
       "GSM_ID                   \n",
       "GSM7625568      35.866626\n",
       "GSM7625569      34.805627\n",
       "GSM7625570      34.625747\n",
       "GSM7625571      37.634215\n",
       "GSM7625572      30.980501"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "adata = pya.pp.df_to_adata(combined_df, verbose=False)\n",
    "pya.pred.predict_age(adata, clock_names=[\"cpgptgrimage3\"], verbose=False)\n",
    "\n",
    "adata.obs.head()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cpgpt",
   "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.15"
  },
  "papermill": {
   "default_parameters": {},
   "duration": 50.820106,
   "end_time": "2025-03-31T09:47:46.741485",
   "environment_variables": {},
   "exception": true,
   "input_path": "reference_map.ipynb",
   "output_path": "reference_map.ipynb",
   "parameters": {},
   "start_time": "2025-03-31T09:46:55.921379",
   "version": "2.6.0"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {
     "03b98b9c01444efebca082d782981bc4": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "4e2cd906a01843f482f99558ed02d638": {
      "model_module": "@jupyter-widgets/output",
      "model_module_version": "1.0.0",
      "model_name": "OutputModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/output",
       "_model_module_version": "1.0.0",
       "_model_name": "OutputModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/output",
       "_view_module_version": "1.0.0",
       "_view_name": "OutputView",
       "layout": "IPY_MODEL_03b98b9c01444efebca082d782981bc4",
       "msg_id": "",
       "outputs": [
        {
         "data": {
          "text/html": "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Predicting <span style=\"color: #6206e0; text-decoration-color: #6206e0\">━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━</span> 13/13 <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">0:00:02 • 0:00:00</span> <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f; text-decoration: underline\">5.07it/s</span>  \n</pre>\n",
          "text/plain": "Predicting \u001b[38;2;98;6;224m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m 13/13 \u001b[2m0:00:02 • 0:00:00\u001b[0m \u001b[2;4m5.07it/s\u001b[0m  \n"
         },
         "metadata": {},
         "output_type": "display_data"
        }
       ],
       "tabbable": null,
       "tooltip": null
      }
     }
    },
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}