Upload folder using huggingface_hub

README.md

@@ -1,3 +1,59 @@
----
-license: mit
----
+---
+license: mit
+library_name: pytorch-lightning
+pipeline_tag: tabular-regression
+tags:
+- biology
+- genomics
+datasets:
+- Genentech/enformer-data
+---
+
+# Enformer Model (Avsec et al. 2021)
+
+## Model Description
+This repository contains the weights for the Enformer model, a long-range transformer architecture designed to predict functional genomic tracks from genomic DNA sequences.
+
+- **Architecture:** Convolutions followed by Transformer layers.
+- **Input:** 196,608 bp of genomic DNA sequence.
+- **Output Resolution:** 128 bp bins.
+- **Source:** [Avsec, Ž. et al. Nature Methods (2021)](https://www.nature.com/articles/s41592-021-01252-x)
+
+
+## Repository Content
+The repository includes both full PyTorch Lightning checkpoints and raw state dictionaries for the human and mouse versions of the model. Note that the weights are derived from the publication but the model has been converted into the PyTorch Lightning format used by gReLU (https://github.com/Genentech/gReLU).
+
+| File | Type | Description |
+| :--- | :--- | :--- |
+| `human.ckpt` | PyTorch Lightning | Full checkpoint including base model and human head. |
+| `mouse.ckpt` | PyTorch Lightning | Full checkpoint including base model and mouse head. |
+| `human_state_dict.h5` | HDF5 | Weights-only state dictionary for the human model. |
+| `mouse_state_dict.h5` | HDF5 | Weights-only state dictionary for the mouse model. |
+| `save_wandb_enformer_human.ipynb` | Jupyter Notebook | Code used to create `human.ckpt` |
+| `save_wandb_enformer_mouse.ipynb` | Jupyter Notebook | Code used to create `mouse.ckpt` |
+
+## Model Heads & Output Tracks
+Both `.ckpt` files utilize the same core transformer trunk but differ in their species-specific output heads.
+
+### Outputs
+
+Human Head: 5,313 total tracks
+Mouse Head: 1,643 total tracks
+
+## Usage
+The models are intended for use with the `grelu` library.
+
+```python
+from grelu.lightning import LightningModel
+from huggingface_hub import hf_hub_download
+
+# Download the desired checkpoint
+ckpt_path = hf_hub_download(
+    repo_id="Genentech/enformer-model", 
+    filename="human.ckpt"
+)
+
+# Load the model
+model = LightningModel.load_from_checkpoint(ckpt_path)
+model.eval()
+```

human.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:030e706b7b64c744be2c363b73e822694cf404164e8720cafb5940cf81da4db6
+size 986737142

human_state_dict.h5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b9b9080c6ddfe1285175bb909b1440c427527784db9134757bbbd986c826de4f
+size 984921590

mouse.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:19246b404028c4dd943265c519d51d796889c712091547539ecd7d8ad0d5e29b
+size 940724534

mouse_state_dict.h5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f97535122efb146d998ca7f4bbcee86ddf1d4d1c8cec5e5d9f8e5fcd42cbe2ff
+size 939809910

save_wandb_enformer_human.ipynb

@@ -0,0 +1,1002 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "0b100814-a834-4c18-abbe-9cab6ea1278c",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/opt/conda/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n"
+     ]
+    }
+   ],
+   "source": [
+    "import wandb\n",
+    "import torch\n",
+    "import pandas as pd\n",
+    "\n",
+    "from grelu.lightning import LightningModel\n",
+    "import pytorch_lightning as pl\n",
+    "from grelu.sequence.utils import get_unique_length, resize"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "cb22e3f0-8ef7-41f3-aefc-4a2d182af5ba",
+   "metadata": {},
+   "source": [
+    "## wandb login"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "137c5253-351e-4945-88a3-a4b7c555326c",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\u001b[34m\u001b[1mwandb\u001b[0m: Using wandb-core as the SDK backend.  Please refer to https://wandb.me/wandb-core for more information.\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mavantikalal\u001b[0m (\u001b[33mgrelu\u001b[0m) to \u001b[32mhttps://api.wandb.ai\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "True"
+      ]
+     },
+     "execution_count": 2,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "wandb.login(host=\"https://api.wandb.ai\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "d372dc19-d701-407f-9303-40adc76f475c",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "Tracking run with wandb version 0.19.7"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "Run data is saved locally in <code>/code/github/gReLU-applications/enformer/wandb/run-20250304_222603-rwin6f0o</code>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "Syncing run <strong><a href='https://wandb.ai/grelu/enformer/runs/rwin6f0o' target=\"_blank\">copy-human</a></strong> to <a href='https://wandb.ai/grelu/enformer' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/developer-guide' target=\"_blank\">docs</a>)<br>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       " View project at <a href='https://wandb.ai/grelu/enformer' target=\"_blank\">https://wandb.ai/grelu/enformer</a>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       " View run at <a href='https://wandb.ai/grelu/enformer/runs/rwin6f0o' target=\"_blank\">https://wandb.ai/grelu/enformer/runs/rwin6f0o</a>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "run = wandb.init(entity='grelu', project='enformer', job_type='copy', name='copy-human',\n",
+    "    settings=wandb.Settings(\n",
+    "        program_relpath='/code/github/gReLU-applications/enformer/save_wandb_enformer_human.ipynb',\n",
+    "        program_abspath='/code/github/gReLU-applications/enformer/save_wandb_enformer_human.ipynb'\n",
+    "    ))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "d116934e-657f-49af-a3de-1e0b4d6e895d",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m No relevant files were detected in the specified directory. No code will be logged to your run.\n"
+     ]
+    }
+   ],
+   "source": [
+    "wandb.run.log_code() "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4564271b-2a71-47e3-bfc4-9b0f4c616a73",
+   "metadata": {},
+   "source": [
+    "## Paths"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "86052145-d3a9-4f6a-9b23-8353b5dd38e1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "targets_path = 'https://raw.githubusercontent.com/calico/basenji/master/manuscripts/cross2020/targets_human.txt'"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "id": "97aed0b6-8b61-424a-8366-37c204b061d8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sequences_path = '/gstore/data/resbioai/grelu/enformer/sequences.bed'"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "915c74f7-15de-4ae6-a9c9-0142fa885f1e",
+   "metadata": {},
+   "source": [
+    "## Process tasks"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "id": "2d8d16a8-95a9-4902-bec9-a25b44a5876e",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "5313\n"
+     ]
+    },
+    {
+     "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>genome</th>\n",
+       "      <th>identifier</th>\n",
+       "      <th>file</th>\n",
+       "      <th>clip</th>\n",
+       "      <th>scale</th>\n",
+       "      <th>sum_stat</th>\n",
+       "      <th>description</th>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>index</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>0</th>\n",
+       "      <td>0</td>\n",
+       "      <td>ENCFF833POA</td>\n",
+       "      <td>/home/drk/tillage/datasets/human/dnase/encode/...</td>\n",
+       "      <td>32</td>\n",
+       "      <td>2</td>\n",
+       "      <td>mean</td>\n",
+       "      <td>DNASE:cerebellum male adult (27 years) and mal...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>0</td>\n",
+       "      <td>ENCFF110QGM</td>\n",
+       "      <td>/home/drk/tillage/datasets/human/dnase/encode/...</td>\n",
+       "      <td>32</td>\n",
+       "      <td>2</td>\n",
+       "      <td>mean</td>\n",
+       "      <td>DNASE:frontal cortex male adult (27 years) and...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>0</td>\n",
+       "      <td>ENCFF880MKD</td>\n",
+       "      <td>/home/drk/tillage/datasets/human/dnase/encode/...</td>\n",
+       "      <td>32</td>\n",
+       "      <td>2</td>\n",
+       "      <td>mean</td>\n",
+       "      <td>DNASE:chorion</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "       genome   identifier                                               file  \\\n",
+       "index                                                                           \n",
+       "0           0  ENCFF833POA  /home/drk/tillage/datasets/human/dnase/encode/...   \n",
+       "1           0  ENCFF110QGM  /home/drk/tillage/datasets/human/dnase/encode/...   \n",
+       "2           0  ENCFF880MKD  /home/drk/tillage/datasets/human/dnase/encode/...   \n",
+       "\n",
+       "       clip  scale sum_stat                                        description  \n",
+       "index                                                                           \n",
+       "0        32      2     mean  DNASE:cerebellum male adult (27 years) and mal...  \n",
+       "1        32      2     mean  DNASE:frontal cortex male adult (27 years) and...  \n",
+       "2        32      2     mean                                      DNASE:chorion  "
+      ]
+     },
+     "execution_count": 16,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "tasks = pd.read_csv(targets_path, sep='\\t', index_col=0)\n",
+    "print(len(tasks))\n",
+    "tasks.head(3)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "id": "68fd7e09-88be-4185-8983-090a4721dfbf",
+   "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>name</th>\n",
+       "      <th>file</th>\n",
+       "      <th>clip</th>\n",
+       "      <th>scale</th>\n",
+       "      <th>sum_stat</th>\n",
+       "      <th>description</th>\n",
+       "      <th>assay</th>\n",
+       "      <th>sample</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>ENCFF833POA</td>\n",
+       "      <td>/home/drk/tillage/datasets/human/dnase/encode/...</td>\n",
+       "      <td>32</td>\n",
+       "      <td>2</td>\n",
+       "      <td>mean</td>\n",
+       "      <td>DNASE:cerebellum male adult (27 years) and mal...</td>\n",
+       "      <td>DNASE</td>\n",
+       "      <td>cerebellum male adult (27 years) and male adul...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>ENCFF110QGM</td>\n",
+       "      <td>/home/drk/tillage/datasets/human/dnase/encode/...</td>\n",
+       "      <td>32</td>\n",
+       "      <td>2</td>\n",
+       "      <td>mean</td>\n",
+       "      <td>DNASE:frontal cortex male adult (27 years) and...</td>\n",
+       "      <td>DNASE</td>\n",
+       "      <td>frontal cortex male adult (27 years) and male ...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>ENCFF880MKD</td>\n",
+       "      <td>/home/drk/tillage/datasets/human/dnase/encode/...</td>\n",
+       "      <td>32</td>\n",
+       "      <td>2</td>\n",
+       "      <td>mean</td>\n",
+       "      <td>DNASE:chorion</td>\n",
+       "      <td>DNASE</td>\n",
+       "      <td>chorion</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>ENCFF463ZLQ</td>\n",
+       "      <td>/home/drk/tillage/datasets/human/dnase/encode/...</td>\n",
+       "      <td>32</td>\n",
+       "      <td>2</td>\n",
+       "      <td>mean</td>\n",
+       "      <td>DNASE:Ishikawa treated with 0.02% dimethyl sul...</td>\n",
+       "      <td>DNASE</td>\n",
+       "      <td>Ishikawa treated with 0.02% dimethyl sulfoxide...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>ENCFF890OGQ</td>\n",
+       "      <td>/home/drk/tillage/datasets/human/dnase/encode/...</td>\n",
+       "      <td>32</td>\n",
+       "      <td>2</td>\n",
+       "      <td>mean</td>\n",
+       "      <td>DNASE:GM03348</td>\n",
+       "      <td>DNASE</td>\n",
+       "      <td>GM03348</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "          name                                               file  clip  \\\n",
+       "0  ENCFF833POA  /home/drk/tillage/datasets/human/dnase/encode/...    32   \n",
+       "1  ENCFF110QGM  /home/drk/tillage/datasets/human/dnase/encode/...    32   \n",
+       "2  ENCFF880MKD  /home/drk/tillage/datasets/human/dnase/encode/...    32   \n",
+       "3  ENCFF463ZLQ  /home/drk/tillage/datasets/human/dnase/encode/...    32   \n",
+       "4  ENCFF890OGQ  /home/drk/tillage/datasets/human/dnase/encode/...    32   \n",
+       "\n",
+       "   scale sum_stat                                        description  assay  \\\n",
+       "0      2     mean  DNASE:cerebellum male adult (27 years) and mal...  DNASE   \n",
+       "1      2     mean  DNASE:frontal cortex male adult (27 years) and...  DNASE   \n",
+       "2      2     mean                                      DNASE:chorion  DNASE   \n",
+       "3      2     mean  DNASE:Ishikawa treated with 0.02% dimethyl sul...  DNASE   \n",
+       "4      2     mean                                      DNASE:GM03348  DNASE   \n",
+       "\n",
+       "                                              sample  \n",
+       "0  cerebellum male adult (27 years) and male adul...  \n",
+       "1  frontal cortex male adult (27 years) and male ...  \n",
+       "2                                            chorion  \n",
+       "3  Ishikawa treated with 0.02% dimethyl sulfoxide...  \n",
+       "4                                            GM03348  "
+      ]
+     },
+     "execution_count": 17,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "tasks = tasks.reset_index(drop=True)\n",
+    "tasks = tasks.drop(columns=[\"genome\"])\n",
+    "tasks[\"assay\"] = tasks[\"description\"].apply(lambda x: x.split(\":\")[0])\n",
+    "tasks[\"sample\"] = tasks[\"description\"].apply(lambda x: \":\".join(x.split(\":\")[1:]))\n",
+    "tasks = tasks.rename(columns={\"identifier\":\"name\"})\n",
+    "tasks.head()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "id": "cc958ef5-f033-43df-b091-22e03021dd48",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tasks = tasks.to_dict(orient=\"list\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0caaa50e-a1dd-4d2a-ab17-06e545cef485",
+   "metadata": {},
+   "source": [
+    "## Process intervals"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "id": "456bf351-55aa-426d-a10d-0edc9b3591f4",
+   "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>chrom</th>\n",
+       "      <th>start</th>\n",
+       "      <th>end</th>\n",
+       "      <th>split</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>chr18</td>\n",
+       "      <td>928386</td>\n",
+       "      <td>1059458</td>\n",
+       "      <td>train</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>chr4</td>\n",
+       "      <td>113630947</td>\n",
+       "      <td>113762019</td>\n",
+       "      <td>train</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>chr11</td>\n",
+       "      <td>18427720</td>\n",
+       "      <td>18558792</td>\n",
+       "      <td>train</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>chr16</td>\n",
+       "      <td>85805681</td>\n",
+       "      <td>85936753</td>\n",
+       "      <td>train</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>chr3</td>\n",
+       "      <td>158386188</td>\n",
+       "      <td>158517260</td>\n",
+       "      <td>train</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "   chrom      start        end  split\n",
+       "0  chr18     928386    1059458  train\n",
+       "1   chr4  113630947  113762019  train\n",
+       "2  chr11   18427720   18558792  train\n",
+       "3  chr16   85805681   85936753  train\n",
+       "4   chr3  158386188  158517260  train"
+      ]
+     },
+     "execution_count": 19,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "intervals = pd.read_table(sequences_path, header=None)\n",
+    "intervals.columns = ['chrom', 'start', 'end', 'split']\n",
+    "intervals.head()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "id": "1701f4c1-d7fe-42b6-80eb-323cb9755531",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "split\n",
+       "train    34021\n",
+       "valid     2213\n",
+       "test      1937\n",
+       "Name: count, dtype: int64"
+      ]
+     },
+     "execution_count": 20,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "intervals.split.value_counts()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "id": "94298aa9-04cf-44e5-82c5-e1786207524d",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "131072"
+      ]
+     },
+     "execution_count": 21,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "get_unique_length(intervals)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "id": "f4f97eb7-f087-46aa-a6c9-09ebbf5af775",
+   "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>chrom</th>\n",
+       "      <th>start</th>\n",
+       "      <th>end</th>\n",
+       "      <th>split</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>chr18</td>\n",
+       "      <td>895618</td>\n",
+       "      <td>1092226</td>\n",
+       "      <td>train</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>chr4</td>\n",
+       "      <td>113598179</td>\n",
+       "      <td>113794787</td>\n",
+       "      <td>train</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>chr11</td>\n",
+       "      <td>18394952</td>\n",
+       "      <td>18591560</td>\n",
+       "      <td>train</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>chr16</td>\n",
+       "      <td>85772913</td>\n",
+       "      <td>85969521</td>\n",
+       "      <td>train</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>chr3</td>\n",
+       "      <td>158353420</td>\n",
+       "      <td>158550028</td>\n",
+       "      <td>train</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "   chrom      start        end  split\n",
+       "0  chr18     895618    1092226  train\n",
+       "1   chr4  113598179  113794787  train\n",
+       "2  chr11   18394952   18591560  train\n",
+       "3  chr16   85772913   85969521  train\n",
+       "4   chr3  158353420  158550028  train"
+      ]
+     },
+     "execution_count": 22,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "intervals = resize(intervals, 196608)\n",
+    "intervals.head()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 23,
+   "id": "9337f97b-8463-4f6a-bc51-23e83af46c89",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_intervals = intervals[intervals.split=='train'].iloc[:, :3]\n",
+    "val_intervals = intervals[intervals.split=='valid'].iloc[:, :3]\n",
+    "test_intervals = intervals[intervals.split=='test'].iloc[:, :3]\n",
+    "del intervals"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "62ecac6c-f031-4269-9672-50f57b883368",
+   "metadata": {},
+   "source": [
+    "## Initialize model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 24,
+   "id": "a5f3af07-c670-4ceb-abc1-891f45b298d0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model_params={\n",
+    "    'model_type':'EnformerModel',\n",
+    "    'final_act_func': 'softplus',\n",
+    "    'final_pool_func':None,\n",
+    "    'n_tasks': 5313,\n",
+    "    'crop_len':320,\n",
+    "}\n",
+    "train_params={'task':'regression', 'loss':'mse'}\n",
+    "\n",
+    "model = LightningModel(model_params, train_params)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4e6981d5-fb40-486a-bb34-db296ac14778",
+   "metadata": {},
+   "source": [
+    "## Load weights"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 25,
+   "id": "08cef7a5-3a9c-4342-8d82-2eabf8eac4d6",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/tmp/ipykernel_3296127/1230005423.py:1: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
+      "  state_dict = torch.load(\"/data/enformer/torch_weights/human.h5\")\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "<All keys matched successfully>"
+      ]
+     },
+     "execution_count": 25,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "state_dict = torch.load(\"/data/enformer/torch_weights/human.h5\")\n",
+    "model.model.load_state_dict(state_dict)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9b11fb1f-2d53-4b08-ae6e-15c3dfa1c1b1",
+   "metadata": {},
+   "source": [
+    "## Add hparams"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 26,
+   "id": "9fd166f4-998b-4ff1-aabd-fe3b12fd3b82",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model.data_params[\"train\"] = dict()\n",
+    "model.data_params[\"val\"] = dict()\n",
+    "model.data_params[\"test\"] = dict()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 27,
+   "id": "09cbfeb5-9f73-44f8-81c5-7f5db0ea577e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model.data_params[\"train\"][\"seq_len\"] = 196608\n",
+    "model.data_params[\"train\"][\"label_len\"] = 896 * 128\n",
+    "model.data_params[\"train\"][\"genome\"] = \"hg38\"\n",
+    "model.data_params[\"train\"][\"bin_size\"] = 128\n",
+    "model.data_params[\"train\"][\"max_seq_shift\"] = 3\n",
+    "model.data_params[\"train\"][\"rc\"] = True"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "208f6727-0540-43dd-8554-4036b8e57180",
+   "metadata": {},
+   "source": [
+    "## Add tasks"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 28,
+   "id": "ba516087-edd9-46b8-9e76-44ff35d6e553",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model.data_params[\"tasks\"] = tasks"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ab5ef3dd-aaa4-4322-a25f-d6b818efb893",
+   "metadata": {},
+   "source": [
+    "## Add intervals"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 29,
+   "id": "be3564f3-2b4d-4618-9277-faceab43b7cd",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model.data_params[\"train\"][\"intervals\"] = train_intervals.to_dict(orient='list')\n",
+    "model.data_params[\"val\"][\"intervals\"] = val_intervals.to_dict(orient='list')\n",
+    "model.data_params[\"test\"][\"intervals\"] = test_intervals.to_dict(orient='list')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d8471c54-8716-41de-84b0-212ba7419adb",
+   "metadata": {},
+   "source": [
+    "## Save"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 30,
+   "id": "7c4e3973-40e4-46c7-9ac9-8985d82498e1",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Trainer will use only 1 of 8 GPUs because it is running inside an interactive / notebook environment. You may try to set `Trainer(devices=8)` but please note that multi-GPU inside interactive / notebook environments is considered experimental and unstable. Your mileage may vary.\n",
+      "GPU available: True (cuda), used: True\n",
+      "TPU available: False, using: 0 TPU cores\n",
+      "HPU available: False, using: 0 HPUs\n",
+      "You are using a CUDA device ('NVIDIA A100-SXM4-80GB') that has Tensor Cores. To properly utilize them, you should set `torch.set_float32_matmul_precision('medium' | 'high')` which will trade-off precision for performance. For more details, read https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html#torch.set_float32_matmul_precision\n",
+      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3,4,5,6,7]\n"
+     ]
+    }
+   ],
+   "source": [
+    "trainer = pl.Trainer()\n",
+    "try:\n",
+    "    trainer.predict(model) \n",
+    "except:\n",
+    "    trainer.save_checkpoint('/data/enformer/torch_weights/human.ckpt')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5dd274ed-ba38-41ab-8288-ca848c33b539",
+   "metadata": {},
+   "source": [
+    "## Upload"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 31,
+   "id": "4f460b22-9da7-40f8-bd77-0c330a2b1111",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 277336 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 277336 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 277336 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 277336 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 277336 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 277336 bytes\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "<Artifact human>"
+      ]
+     },
+     "execution_count": 31,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "artifact = wandb.Artifact(\n",
+    "    'human', \n",
+    "    type='model',\n",
+    "    metadata={\n",
+    "        'model_params':model.model_params, \n",
+    "        'train_params':model.train_params, \n",
+    "        'data_params':model.data_params\n",
+    "    }\n",
+    ")\n",
+    "artifact.add_file(local_path='/data/enformer/torch_weights/human.ckpt', name='model.ckpt')\n",
+    "run.log_artifact(artifact)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 32,
+   "id": "9269df1a-b6d4-410d-bbe6-39a4b70576cc",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       " View run <strong style=\"color:#cdcd00\">copy-human</strong> at: <a href='https://wandb.ai/grelu/enformer/runs/rwin6f0o' target=\"_blank\">https://wandb.ai/grelu/enformer/runs/rwin6f0o</a><br> View project at: <a href='https://wandb.ai/grelu/enformer' target=\"_blank\">https://wandb.ai/grelu/enformer</a><br>Synced 6 W&B file(s), 0 media file(s), 2 artifact file(s) and 0 other file(s)"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "Find logs at: <code>./wandb/run-20250304_222603-rwin6f0o/logs</code>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "run.finish() "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "55eef443-2d99-443f-a31a-d1aee4a31afe",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

save_wandb_enformer_mouse.ipynb

@@ -0,0 +1,979 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "0b100814-a834-4c18-abbe-9cab6ea1278c",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/opt/conda/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n"
+     ]
+    }
+   ],
+   "source": [
+    "import wandb\n",
+    "import torch\n",
+    "import pandas as pd\n",
+    "\n",
+    "from grelu.lightning import LightningModel\n",
+    "import pytorch_lightning as pl\n",
+    "from grelu.sequence.utils import get_unique_length, resize"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "cb22e3f0-8ef7-41f3-aefc-4a2d182af5ba",
+   "metadata": {},
+   "source": [
+    "## wandb login"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "137c5253-351e-4945-88a3-a4b7c555326c",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\u001b[34m\u001b[1mwandb\u001b[0m: Using wandb-core as the SDK backend.  Please refer to https://wandb.me/wandb-core for more information.\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mavantikalal\u001b[0m (\u001b[33mgrelu\u001b[0m) to \u001b[32mhttps://api.wandb.ai\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "True"
+      ]
+     },
+     "execution_count": 2,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "wandb.login(host=\"https://api.wandb.ai\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "b5be7b96-5f3c-4073-b152-22ae3e46db06",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "Tracking run with wandb version 0.19.7"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "Run data is saved locally in <code>/code/github/gReLU-applications/enformer/wandb/run-20250304_222920-jrgxfvad</code>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "Syncing run <strong><a href='https://wandb.ai/grelu/enformer/runs/jrgxfvad' target=\"_blank\">copy-mouse</a></strong> to <a href='https://wandb.ai/grelu/enformer' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/developer-guide' target=\"_blank\">docs</a>)<br>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       " View project at <a href='https://wandb.ai/grelu/enformer' target=\"_blank\">https://wandb.ai/grelu/enformer</a>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       " View run at <a href='https://wandb.ai/grelu/enformer/runs/jrgxfvad' target=\"_blank\">https://wandb.ai/grelu/enformer/runs/jrgxfvad</a>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "run = wandb.init(entity='grelu', project='enformer', job_type='copy', name='copy-mouse') # Initialize a W&B Run"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4564271b-2a71-47e3-bfc4-9b0f4c616a73",
+   "metadata": {},
+   "source": [
+    "## Paths"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "86052145-d3a9-4f6a-9b23-8353b5dd38e1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "targets_path = 'https://raw.githubusercontent.com/calico/basenji/master/manuscripts/cross2020/targets_mouse.txt'"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "97aed0b6-8b61-424a-8366-37c204b061d8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sequences_path = '/gstore/data/resbioai/grelu/enformer/sequences-mouse.bed'"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "915c74f7-15de-4ae6-a9c9-0142fa885f1e",
+   "metadata": {},
+   "source": [
+    "## Process tasks"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "6701b16c-e750-4eb6-b2e0-1223d49cc103",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "1643\n"
+     ]
+    },
+    {
+     "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>genome</th>\n",
+       "      <th>identifier</th>\n",
+       "      <th>file</th>\n",
+       "      <th>clip</th>\n",
+       "      <th>scale</th>\n",
+       "      <th>sum_stat</th>\n",
+       "      <th>description</th>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>index</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>5313</th>\n",
+       "      <td>1</td>\n",
+       "      <td>ENCFF866ZTV</td>\n",
+       "      <td>/home/drk/tillage/datasets/mouse/dnase/encode/...</td>\n",
+       "      <td>32</td>\n",
+       "      <td>2</td>\n",
+       "      <td>mean</td>\n",
+       "      <td>DNASE:B6D2F1/J 416B</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>5314</th>\n",
+       "      <td>1</td>\n",
+       "      <td>ENCFF695LHM</td>\n",
+       "      <td>/home/drk/tillage/datasets/mouse/dnase/encode/...</td>\n",
+       "      <td>32</td>\n",
+       "      <td>2</td>\n",
+       "      <td>mean</td>\n",
+       "      <td>DNASE:BALB/cAnN A20</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>5315</th>\n",
+       "      <td>1</td>\n",
+       "      <td>ENCFF079SPZ</td>\n",
+       "      <td>/home/drk/tillage/datasets/mouse/dnase/encode/...</td>\n",
+       "      <td>32</td>\n",
+       "      <td>2</td>\n",
+       "      <td>mean</td>\n",
+       "      <td>DNASE:C57BL/6 B cell male adult (8 weeks)</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "       genome   identifier                                               file  \\\n",
+       "index                                                                           \n",
+       "5313        1  ENCFF866ZTV  /home/drk/tillage/datasets/mouse/dnase/encode/...   \n",
+       "5314        1  ENCFF695LHM  /home/drk/tillage/datasets/mouse/dnase/encode/...   \n",
+       "5315        1  ENCFF079SPZ  /home/drk/tillage/datasets/mouse/dnase/encode/...   \n",
+       "\n",
+       "       clip  scale sum_stat                                description  \n",
+       "index                                                                   \n",
+       "5313     32      2     mean                        DNASE:B6D2F1/J 416B  \n",
+       "5314     32      2     mean                        DNASE:BALB/cAnN A20  \n",
+       "5315     32      2     mean  DNASE:C57BL/6 B cell male adult (8 weeks)  "
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "tasks = pd.read_csv(targets_path, sep='\\t', index_col=0)\n",
+    "print(len(tasks))\n",
+    "tasks.head(3)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "cc958ef5-f033-43df-b091-22e03021dd48",
+   "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>name</th>\n",
+       "      <th>file</th>\n",
+       "      <th>clip</th>\n",
+       "      <th>scale</th>\n",
+       "      <th>sum_stat</th>\n",
+       "      <th>description</th>\n",
+       "      <th>assay</th>\n",
+       "      <th>sample</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>ENCFF866ZTV</td>\n",
+       "      <td>/home/drk/tillage/datasets/mouse/dnase/encode/...</td>\n",
+       "      <td>32</td>\n",
+       "      <td>2</td>\n",
+       "      <td>mean</td>\n",
+       "      <td>DNASE:B6D2F1/J 416B</td>\n",
+       "      <td>DNASE</td>\n",
+       "      <td>B6D2F1/J 416B</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>ENCFF695LHM</td>\n",
+       "      <td>/home/drk/tillage/datasets/mouse/dnase/encode/...</td>\n",
+       "      <td>32</td>\n",
+       "      <td>2</td>\n",
+       "      <td>mean</td>\n",
+       "      <td>DNASE:BALB/cAnN A20</td>\n",
+       "      <td>DNASE</td>\n",
+       "      <td>BALB/cAnN A20</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>ENCFF079SPZ</td>\n",
+       "      <td>/home/drk/tillage/datasets/mouse/dnase/encode/...</td>\n",
+       "      <td>32</td>\n",
+       "      <td>2</td>\n",
+       "      <td>mean</td>\n",
+       "      <td>DNASE:C57BL/6 B cell male adult (8 weeks)</td>\n",
+       "      <td>DNASE</td>\n",
+       "      <td>C57BL/6 B cell male adult (8 weeks)</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>ENCFF798VSP</td>\n",
+       "      <td>/home/drk/tillage/datasets/mouse/dnase/encode/...</td>\n",
+       "      <td>32</td>\n",
+       "      <td>2</td>\n",
+       "      <td>mean</td>\n",
+       "      <td>DNASE:C57BL/6 splenic B cell male adult (8 weeks)</td>\n",
+       "      <td>DNASE</td>\n",
+       "      <td>C57BL/6 splenic B cell male adult (8 weeks)</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>ENCFF474GND</td>\n",
+       "      <td>/home/drk/tillage/datasets/mouse/dnase/encode/...</td>\n",
+       "      <td>32</td>\n",
+       "      <td>2</td>\n",
+       "      <td>mean</td>\n",
+       "      <td>DNASE:C57BL/6 cerebellum male adult (8 weeks)</td>\n",
+       "      <td>DNASE</td>\n",
+       "      <td>C57BL/6 cerebellum male adult (8 weeks)</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "          name                                               file  clip  \\\n",
+       "0  ENCFF866ZTV  /home/drk/tillage/datasets/mouse/dnase/encode/...    32   \n",
+       "1  ENCFF695LHM  /home/drk/tillage/datasets/mouse/dnase/encode/...    32   \n",
+       "2  ENCFF079SPZ  /home/drk/tillage/datasets/mouse/dnase/encode/...    32   \n",
+       "3  ENCFF798VSP  /home/drk/tillage/datasets/mouse/dnase/encode/...    32   \n",
+       "4  ENCFF474GND  /home/drk/tillage/datasets/mouse/dnase/encode/...    32   \n",
+       "\n",
+       "   scale sum_stat                                        description  assay  \\\n",
+       "0      2     mean                                DNASE:B6D2F1/J 416B  DNASE   \n",
+       "1      2     mean                                DNASE:BALB/cAnN A20  DNASE   \n",
+       "2      2     mean          DNASE:C57BL/6 B cell male adult (8 weeks)  DNASE   \n",
+       "3      2     mean  DNASE:C57BL/6 splenic B cell male adult (8 weeks)  DNASE   \n",
+       "4      2     mean      DNASE:C57BL/6 cerebellum male adult (8 weeks)  DNASE   \n",
+       "\n",
+       "                                        sample  \n",
+       "0                                B6D2F1/J 416B  \n",
+       "1                                BALB/cAnN A20  \n",
+       "2          C57BL/6 B cell male adult (8 weeks)  \n",
+       "3  C57BL/6 splenic B cell male adult (8 weeks)  \n",
+       "4      C57BL/6 cerebellum male adult (8 weeks)  "
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "tasks = tasks.reset_index(drop=True)\n",
+    "tasks = tasks.drop(columns=[\"genome\"])\n",
+    "tasks[\"assay\"] = tasks[\"description\"].apply(lambda x: x.split(\":\")[0])\n",
+    "tasks[\"sample\"] = tasks[\"description\"].apply(lambda x: \":\".join(x.split(\":\")[1:]))\n",
+    "tasks = tasks.rename(columns={\"identifier\":\"name\"})\n",
+    "tasks.head()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "d914129b-4dcd-4a11-9a8e-f99618200024",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tasks = tasks.to_dict(orient=\"list\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0caaa50e-a1dd-4d2a-ab17-06e545cef485",
+   "metadata": {},
+   "source": [
+    "## Process intervals"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "456bf351-55aa-426d-a10d-0edc9b3591f4",
+   "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>chrom</th>\n",
+       "      <th>start</th>\n",
+       "      <th>end</th>\n",
+       "      <th>split</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>chr4</td>\n",
+       "      <td>34106647</td>\n",
+       "      <td>34237719</td>\n",
+       "      <td>train</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>chr5</td>\n",
+       "      <td>52207747</td>\n",
+       "      <td>52338819</td>\n",
+       "      <td>train</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>chr19</td>\n",
+       "      <td>20136862</td>\n",
+       "      <td>20267934</td>\n",
+       "      <td>train</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>chr14</td>\n",
+       "      <td>61845439</td>\n",
+       "      <td>61976511</td>\n",
+       "      <td>train</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>chr15</td>\n",
+       "      <td>6592346</td>\n",
+       "      <td>6723418</td>\n",
+       "      <td>train</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "   chrom     start       end  split\n",
+       "0   chr4  34106647  34237719  train\n",
+       "1   chr5  52207747  52338819  train\n",
+       "2  chr19  20136862  20267934  train\n",
+       "3  chr14  61845439  61976511  train\n",
+       "4  chr15   6592346   6723418  train"
+      ]
+     },
+     "execution_count": 9,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "intervals = pd.read_table(sequences_path, header=None, names = ['chrom', 'start', 'end', 'split'])\n",
+    "intervals.head()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "237d58cd-4250-46de-af72-9fd5a481657e",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "split\n",
+       "train    29295\n",
+       "valid     2209\n",
+       "test      2017\n",
+       "Name: count, dtype: int64"
+      ]
+     },
+     "execution_count": 10,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "intervals.split.value_counts()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "8fb0c5d9-fc33-498f-b026-a9ac93b59418",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "131072"
+      ]
+     },
+     "execution_count": 11,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "get_unique_length(intervals)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "a1633ea5-82a6-45d7-b28e-9ed1cef59c24",
+   "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>chrom</th>\n",
+       "      <th>start</th>\n",
+       "      <th>end</th>\n",
+       "      <th>split</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>chr4</td>\n",
+       "      <td>34073879</td>\n",
+       "      <td>34270487</td>\n",
+       "      <td>train</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>chr5</td>\n",
+       "      <td>52174979</td>\n",
+       "      <td>52371587</td>\n",
+       "      <td>train</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>chr19</td>\n",
+       "      <td>20104094</td>\n",
+       "      <td>20300702</td>\n",
+       "      <td>train</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>chr14</td>\n",
+       "      <td>61812671</td>\n",
+       "      <td>62009279</td>\n",
+       "      <td>train</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>chr15</td>\n",
+       "      <td>6559578</td>\n",
+       "      <td>6756186</td>\n",
+       "      <td>train</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "   chrom     start       end  split\n",
+       "0   chr4  34073879  34270487  train\n",
+       "1   chr5  52174979  52371587  train\n",
+       "2  chr19  20104094  20300702  train\n",
+       "3  chr14  61812671  62009279  train\n",
+       "4  chr15   6559578   6756186  train"
+      ]
+     },
+     "execution_count": 12,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "intervals = resize(intervals, 196608)\n",
+    "intervals.head()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "75f1f860-dcd3-4e9f-8a2d-95ca95e63654",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_intervals = intervals[intervals.split=='train'].iloc[:, :3]\n",
+    "val_intervals = intervals[intervals.split=='valid'].iloc[:, :3]\n",
+    "test_intervals = intervals[intervals.split=='test'].iloc[:, :3]\n",
+    "del intervals"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "62ecac6c-f031-4269-9672-50f57b883368",
+   "metadata": {},
+   "source": [
+    "## Initialize model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "a5f3af07-c670-4ceb-abc1-891f45b298d0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model_params={\n",
+    "    'model_type':'EnformerModel',\n",
+    "    'final_act_func': 'softplus',\n",
+    "    'final_pool_func':None,\n",
+    "    'n_tasks': 1643,\n",
+    "    'crop_len':320,\n",
+    "}\n",
+    "train_params={'task':'regression', 'loss':'mse'}\n",
+    "\n",
+    "model = LightningModel(model_params, train_params)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4e6981d5-fb40-486a-bb34-db296ac14778",
+   "metadata": {},
+   "source": [
+    "## Load weights"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "id": "08cef7a5-3a9c-4342-8d82-2eabf8eac4d6",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/tmp/ipykernel_3297608/3379349373.py:1: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
+      "  state_dict = torch.load(\"/data/enformer/torch_weights/mouse.h5\")\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "<All keys matched successfully>"
+      ]
+     },
+     "execution_count": 15,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "state_dict = torch.load(\"/data/enformer/torch_weights/mouse.h5\")\n",
+    "model.model.load_state_dict(state_dict)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9b11fb1f-2d53-4b08-ae6e-15c3dfa1c1b1",
+   "metadata": {},
+   "source": [
+    "## Add hparams"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "id": "d9a32cf3-4c69-45d9-81b9-f7c5b83ae763",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model.data_params[\"train\"] = dict()\n",
+    "model.data_params[\"val\"] = dict()\n",
+    "model.data_params[\"test\"] = dict()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "id": "09cbfeb5-9f73-44f8-81c5-7f5db0ea577e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model.data_params[\"train\"][\"seq_len\"] = 196608\n",
+    "model.data_params[\"train\"][\"label_len\"] = 896 * 128\n",
+    "model.data_params[\"train\"][\"genome\"] = \"mm10\"\n",
+    "model.data_params[\"train\"][\"bin_size\"] = 128\n",
+    "model.data_params[\"train\"][\"max_seq_shift\"] = 3\n",
+    "model.data_params[\"train\"][\"rc\"] = True"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "208f6727-0540-43dd-8554-4036b8e57180",
+   "metadata": {},
+   "source": [
+    "## Add tasks"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "id": "ba516087-edd9-46b8-9e76-44ff35d6e553",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model.data_params[\"tasks\"] = tasks"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ab5ef3dd-aaa4-4322-a25f-d6b818efb893",
+   "metadata": {},
+   "source": [
+    "## Add intervals"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "id": "be3564f3-2b4d-4618-9277-faceab43b7cd",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model.data_params[\"train\"][\"intervals\"] = train_intervals.to_dict(orient='list')\n",
+    "model.data_params[\"val\"][\"intervals\"] = val_intervals.to_dict(orient='list')\n",
+    "model.data_params[\"test\"][\"intervals\"] = test_intervals.to_dict(orient='list')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d8471c54-8716-41de-84b0-212ba7419adb",
+   "metadata": {},
+   "source": [
+    "## Save"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "id": "7c4e3973-40e4-46c7-9ac9-8985d82498e1",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Trainer will use only 1 of 8 GPUs because it is running inside an interactive / notebook environment. You may try to set `Trainer(devices=8)` but please note that multi-GPU inside interactive / notebook environments is considered experimental and unstable. Your mileage may vary.\n",
+      "GPU available: True (cuda), used: True\n",
+      "TPU available: False, using: 0 TPU cores\n",
+      "HPU available: False, using: 0 HPUs\n",
+      "You are using a CUDA device ('NVIDIA A100-SXM4-80GB') that has Tensor Cores. To properly utilize them, you should set `torch.set_float32_matmul_precision('medium' | 'high')` which will trade-off precision for performance. For more details, read https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html#torch.set_float32_matmul_precision\n",
+      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3,4,5,6,7]\n"
+     ]
+    }
+   ],
+   "source": [
+    "trainer = pl.Trainer()\n",
+    "try:\n",
+    "    trainer.predict(model) \n",
+    "except:\n",
+    "    trainer.save_checkpoint('/data/enformer/torch_weights/mouse.ckpt')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5dd274ed-ba38-41ab-8288-ca848c33b539",
+   "metadata": {},
+   "source": [
+    "## Upload"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "id": "92ff666a-ac2a-41b4-a48f-dc2faeb1a3dd",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 246488 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 246488 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 246488 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 246488 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 246488 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 246488 bytes\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "<Artifact mouse>"
+      ]
+     },
+     "execution_count": 21,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "artifact = wandb.Artifact(\n",
+    "    'mouse', \n",
+    "    type='model',\n",
+    "    metadata={\n",
+    "        'model_params':model.model_params, \n",
+    "        'train_params':model.train_params, \n",
+    "        'data_params':model.data_params\n",
+    "    }\n",
+    ")\n",
+    "artifact.add_file(local_path='/data/enformer/torch_weights/mouse.ckpt', name='model.ckpt')\n",
+    "run.log_artifact(artifact)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "id": "9269df1a-b6d4-410d-bbe6-39a4b70576cc",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       " View run <strong style=\"color:#cdcd00\">copy-mouse</strong> at: <a href='https://wandb.ai/grelu/enformer/runs/jrgxfvad' target=\"_blank\">https://wandb.ai/grelu/enformer/runs/jrgxfvad</a><br> View project at: <a href='https://wandb.ai/grelu/enformer' target=\"_blank\">https://wandb.ai/grelu/enformer</a><br>Synced 5 W&B file(s), 0 media file(s), 2 artifact file(s) and 0 other file(s)"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "Find logs at: <code>./wandb/run-20250304_222920-jrgxfvad/logs</code>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "run.finish() "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b00d26d5-e7b6-47c1-a8b5-7d232d0c2592",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}