Upload folder using huggingface_hub (#1)

- Upload folder using huggingface_hub (ae1d3d64cb38197827fff3b93c5036f46960246e)

README.md

@@ -1,3 +1,71 @@
----
-license: mit
----
+---
+license: mit
+library_name: pytorch-lightning
+pipeline_tag: tabular-regression
+tags:
+- biology
+- genomics
+datasets:
+- Genentech/borzoi-data
+---
+
+# Borzoi Model (Linder et al. 2025)
+
+## Model Description
+This repository contains the weights for the Borzoi model, a model designed to predict functional genomic tracks from genomic DNA sequences.
+
+- **Architecture:** Convolutions followed by Transformer and U-net layers.
+- **Input:** 524,288 bp of genomic DNA sequence.
+- **Output Resolution:** 32 bp bins.
+- **Source:** [Linder, J. et al. Nature Genetics (2025)](https://www.nature.com/articles/s41588-024-02053-6)
+
+
+## Repository Content
+The repository includes both full PyTorch Lightning checkpoints and raw state dictionaries for the human and mouse versions of the model. Each model has 4 independent replicates. 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_rep0.ckpt` | PyTorch Lightning | Full checkpoint including base model and human head, replicate 0. |
+| `human_rep1.ckpt` | PyTorch Lightning | Full checkpoint including base model and human head, replicate 1. |
+| `human_rep2.ckpt` | PyTorch Lightning | Full checkpoint including base model and human head, replicate 2. |
+| `human_rep3.ckpt` | PyTorch Lightning | Full checkpoint including base model and human head, replicate 3. |
+| `mouse_rep0.ckpt` | PyTorch Lightning | Full checkpoint including base model and mouse head, replicate 0. |
+| `mouse_rep1.ckpt` | PyTorch Lightning | Full checkpoint including base model and mouse head, replicate 1. |
+| `mouse_rep2.ckpt` | PyTorch Lightning | Full checkpoint including base model and mouse head, replicate 2. |
+| `mouse_rep3.ckpt` | PyTorch Lightning | Full checkpoint including base model and mouse head, replicate 3. |
+| `human_state_dict_rep0.h5` | HDF5 | Weights-only state dictionary for the human model, replicate 0. |
+| `human_state_dict_rep1.h5` | HDF5 | Weights-only state dictionary for the human model, replicate 1. |
+| `human_state_dict_rep2.h5` | HDF5 | Weights-only state dictionary for the human model, replicate 2. |
+| `human_state_dict_rep3.h5` | HDF5 | Weights-only state dictionary for the human model, replicate 3. |
+| `mouse_state_dict_rep0.h5` | HDF5 | Weights-only state dictionary for the mouse model, replicate 0. |
+| `mouse_state_dict_rep1.h5` | HDF5 | Weights-only state dictionary for the mouse model, replicate 1. |
+| `mouse_state_dict_rep2.h5` | HDF5 | Weights-only state dictionary for the mouse model, replicate 2. |
+| `mouse_state_dict_rep3.h5` | HDF5 | Weights-only state dictionary for the mouse model, replicate 3. |
+| `save_wandb_borzoi_human.ipynb` | Jupyter Notebook | Code used to create human checkpoints |
+| `save_wandb_borzoi_mouse.ipynb` | Jupyter Notebook | Code used to create mouse checkpoints |
+
+## Model Heads & Output Tracks
+The uman and mouse `.ckpt` files for the same replicate utilize the same core transformer trunk but differ in their species-specific output heads.
+
+### Outputs
+
+Human Head: 7,611 total tracks
+Mouse Head: 2,608 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/borzoi-model", 
+    filename="human_rep0.ckpt"
+)
+
+# Load the model
+model = LightningModel.load_from_checkpoint(ckpt_path)
+model.eval()
+```

human_rep0.ckpt

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

human_rep1.ckpt

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

human_rep2.ckpt

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

human_rep3.ckpt

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

human_state_dict_rep0.h5

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

human_state_dict_rep1.h5

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

human_state_dict_rep2.h5

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

human_state_dict_rep3.h5

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

mouse_rep0.ckpt

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

mouse_rep1.ckpt

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

mouse_rep2.ckpt

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

mouse_rep3.ckpt

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

mouse_state_dict_rep0.h5

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

mouse_state_dict_rep1.h5

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

mouse_state_dict_rep2.h5

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

mouse_state_dict_rep3.h5

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

save_wandb_borzoi_ckpt_human.ipynb

@@ -0,0 +1,684 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "c31446f3-a976-42a0-9dfb-bc873049464b",
+   "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 os\n",
+    "import pandas as pd\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": "45e5daae-fe31-4cc7-b876-683768acd3cf",
+   "metadata": {},
+   "source": [
+    "## set up wandb"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "18c561ab-ba82-4718-abe9-96c44133b314",
+   "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": "48cf7ec6-1b32-4f96-a0f1-c7ffb1cb5fcf",
+   "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/borzoi/wandb/run-20250306_055353-qtz95621</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/borzoi/runs/qtz95621' target=\"_blank\">copy-human</a></strong> to <a href='https://wandb.ai/grelu/borzoi' 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/borzoi' target=\"_blank\">https://wandb.ai/grelu/borzoi</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/borzoi/runs/qtz95621' target=\"_blank\">https://wandb.ai/grelu/borzoi/runs/qtz95621</a>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "run = wandb.init(\n",
+    "    entity='grelu', project='borzoi', job_type='copy', name='copy-human',\n",
+    "    settings=wandb.Settings(\n",
+    "        program_relpath='save_wandb_borzoi_ckpt_human.ipynb',\n",
+    "        program_abspath='/code/github/gReLU-applications/borzoi/save_wandb_borzoi_ckpt_human.ipynb'\n",
+    "    ))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "de1c87dc-cacf-48f0-a081-5260e325cb7b",
+   "metadata": {},
+   "source": [
+    "## Load tasks and intervals"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "556b9768-f3e4-466a-b097-0b3226e788dd",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tasks_path = '/code/borzoi/data/targets_human.txt.gz'"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "f8616dde-2567-4a3c-8fb9-de412d342078",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "7611\n"
+     ]
+    }
+   ],
+   "source": [
+    "tasks = pd.read_table(tasks_path, index_col=0)\n",
+    "print(len(tasks))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "0bf621ee-bd42-4299-98b1-d4f0e5ebb6ec",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\u001b[34m\u001b[1mwandb\u001b[0m:   1 of 1 files downloaded.  \n"
+     ]
+    }
+   ],
+   "source": [
+    "artifact = run.use_artifact('human_intervals:latest')\n",
+    "dir = artifact.download()\n",
+    "intervals = pd.read_table(os.path.join(dir, \"data.tsv\"))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "882d78ab-832c-4166-b130-91e37cda7327",
+   "metadata": {},
+   "source": [
+    "## Process tasks"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "66f24ff2-74bc-41ee-9cc9-5cf9e75c2ebf",
+   "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>clip_soft</th>\n",
+       "      <th>scale</th>\n",
+       "      <th>sum_stat</th>\n",
+       "      <th>strand_pair</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>CNhs10608+</td>\n",
+       "      <td>/home/drk/tillage/datasets/human/cage/fantom/C...</td>\n",
+       "      <td>768</td>\n",
+       "      <td>384</td>\n",
+       "      <td>1.0</td>\n",
+       "      <td>sum</td>\n",
+       "      <td>1</td>\n",
+       "      <td>CAGE:Clontech Human Universal Reference Total ...</td>\n",
+       "      <td>CAGE</td>\n",
+       "      <td>Clontech Human Universal Reference Total RNA, ...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>CNhs10608-</td>\n",
+       "      <td>/home/drk/tillage/datasets/human/cage/fantom/C...</td>\n",
+       "      <td>768</td>\n",
+       "      <td>384</td>\n",
+       "      <td>1.0</td>\n",
+       "      <td>sum</td>\n",
+       "      <td>0</td>\n",
+       "      <td>CAGE:Clontech Human Universal Reference Total ...</td>\n",
+       "      <td>CAGE</td>\n",
+       "      <td>Clontech Human Universal Reference Total RNA, ...</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "         name                                               file  clip  \\\n",
+       "0  CNhs10608+  /home/drk/tillage/datasets/human/cage/fantom/C...   768   \n",
+       "1  CNhs10608-  /home/drk/tillage/datasets/human/cage/fantom/C...   768   \n",
+       "\n",
+       "   clip_soft  scale sum_stat  strand_pair  \\\n",
+       "0        384    1.0      sum            1   \n",
+       "1        384    1.0      sum            0   \n",
+       "\n",
+       "                                         description assay  \\\n",
+       "0  CAGE:Clontech Human Universal Reference Total ...  CAGE   \n",
+       "1  CAGE:Clontech Human Universal Reference Total ...  CAGE   \n",
+       "\n",
+       "                                              sample  \n",
+       "0  Clontech Human Universal Reference Total RNA, ...  \n",
+       "1  Clontech Human Universal Reference Total RNA, ...  "
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "tasks=tasks.reset_index(drop=True)\n",
+    "tasks[\"assay\"] = tasks[\"description\"].apply(lambda x: x.split(\":\")[0])\n",
+    "tasks[\"assay\"].value_counts()\n",
+    "tasks[\"sample\"] = tasks[\"description\"].apply(lambda x: \":\".join(x.split(\":\")[1:]))\n",
+    "tasks = tasks.rename(columns={\"identifier\":\"name\"})\n",
+    "tasks.head(2)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "4f96ecc7-238e-4707-a3af-64b0d1828bd4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tasks = tasks.to_dict(orient=\"list\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3056ecfa-1b75-46b7-9e36-4e72c4cb3618",
+   "metadata": {},
+   "source": [
+    "## Process intervals"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "86aae057-b106-4884-81b3-d1aa563ad4f7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_intervals = intervals[intervals.split=='train'].iloc[:, :3]\n",
+    "val_intervals = intervals[intervals.split=='val'].iloc[:, :3]\n",
+    "test_intervals = intervals[intervals.split=='test'].iloc[:, :3]\n",
+    "del intervals"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e4bd10a9-40c2-4299-aa74-6474ea5a71ef",
+   "metadata": {},
+   "source": [
+    "## Initialize model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "06dba49d-11ba-471b-83d7-05fe45a9bd5a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model_params = {\n",
+    "    \"model_type\":\"BorzoiModel\",\n",
+    "    \"n_tasks\":7611,\n",
+    "    \"final_act_func\": 'softplus',\n",
+    "    \"final_pool_func\":None,\n",
+    "    \"crop_len\":5120,\n",
+    "}\n",
+    "train_params={\n",
+    "    \"task\":\"regression\",\n",
+    "    \"loss\":\"mse\",\n",
+    "}\n",
+    "lm = LightningModel(model_params, train_params)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ab6a3055-d338-40f0-9fd3-f8f3134df643",
+   "metadata": {},
+   "source": [
+    "## Save checkpoints"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "bbccda9a-7a62-4472-9c0b-34abf98a7e1c",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "0\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/tmp/ipykernel_3881069/2231986157.py:16: 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(f\"/data/borzoi/torch_weights/fold{rep}.h5\")\n",
+      "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"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "1\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/tmp/ipykernel_3881069/2231986157.py:16: 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(f\"/data/borzoi/torch_weights/fold{rep}.h5\")\n",
+      "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",
+      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3,4,5,6,7]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/tmp/ipykernel_3881069/2231986157.py:16: 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(f\"/data/borzoi/torch_weights/fold{rep}.h5\")\n",
+      "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",
+      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3,4,5,6,7]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "3\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/tmp/ipykernel_3881069/2231986157.py:16: 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(f\"/data/borzoi/torch_weights/fold{rep}.h5\")\n",
+      "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",
+      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3,4,5,6,7]\n"
+     ]
+    }
+   ],
+   "source": [
+    "for rep in range(4):\n",
+    "    print(rep)\n",
+    "    \n",
+    "    lm.data_params[\"tasks\"] = tasks\n",
+    "    lm.data_params[\"train\"] = dict()\n",
+    "    lm.data_params[\"val\"] = dict()\n",
+    "    lm.data_params[\"test\"] = dict()\n",
+    "    lm.data_params[\"train\"][\"seq_len\"] = 524288\n",
+    "    lm.data_params[\"train\"][\"label_len\"] = 6144*32\n",
+    "    lm.data_params[\"train\"][\"genome\"] = \"hg38\"\n",
+    "    lm.data_params[\"train\"][\"bin_size\"] = 32\n",
+    "    lm.data_params[\"train\"][\"intervals\"] = train_intervals.to_dict(orient='list')\n",
+    "    lm.data_params[\"val\"][\"intervals\"] = val_intervals.to_dict(orient='list')\n",
+    "    lm.data_params[\"test\"][\"intervals\"] = test_intervals.to_dict(orient='list')\n",
+    "\n",
+    "    state_dict = torch.load(f\"/data/borzoi/torch_weights/fold{rep}.h5\")\n",
+    "    lm.model.load_state_dict(state_dict)\n",
+    "    \n",
+    "    trainer = pl.Trainer()\n",
+    "    try:\n",
+    "        trainer.predict(lm)\n",
+    "    except:\n",
+    "        pass\n",
+    "\n",
+    "    trainer.save_checkpoint(f'/data/borzoi/torch_weights/human_rep{rep}.ckpt')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4b5fbc09-abca-4cb7-a8b7-7eeba7b39cf8",
+   "metadata": {},
+   "source": [
+    "## Upload to wandb"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "c53a69b4-9092-44b9-b8a9-20b72aa8b511",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "metadata={\n",
+    "            'model_params':lm.model_params, \n",
+    "            'train_params':lm.train_params, \n",
+    "            'data_params':lm.data_params\n",
+    "        }"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "2a25b878-ca1f-4872-9549-ab502975ca2e",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "0\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 351064 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 351064 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 351064 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 351064 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 351064 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 351064 bytes\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "1\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 351064 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 351064 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 351064 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 351064 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 351064 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 351064 bytes\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 351064 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 351064 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 351064 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 351064 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 351064 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 351064 bytes\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "3\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 351064 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 351064 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 351064 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 351064 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 351064 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 351064 bytes\n"
+     ]
+    }
+   ],
+   "source": [
+    "for rep in range(4):\n",
+    "    print(rep)\n",
+    "    artifact = wandb.Artifact(\n",
+    "        f'human_rep{rep}', \n",
+    "        type='model',\n",
+    "        metadata=metadata,\n",
+    "    )\n",
+    "    artifact.add_file(local_path=f'/data/borzoi/torch_weights/human_rep{rep}.ckpt', name='model.ckpt')\n",
+    "    run.log_artifact(artifact)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "98dd40ec-f276-4099-a4a7-39ccf1600a58",
+   "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/borzoi/runs/qtz95621' target=\"_blank\">https://wandb.ai/grelu/borzoi/runs/qtz95621</a><br> View project at: <a href='https://wandb.ai/grelu/borzoi' target=\"_blank\">https://wandb.ai/grelu/borzoi</a><br>Synced 6 W&B file(s), 0 media file(s), 8 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-20250306_055353-qtz95621/logs</code>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "run.finish()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1c5d6a57-a2a1-495d-90ff-16508e4e236c",
+   "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_borzoi_ckpt_mouse.ipynb

@@ -0,0 +1,717 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "c31446f3-a976-42a0-9dfb-bc873049464b",
+   "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",
+    "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": "45e5daae-fe31-4cc7-b876-683768acd3cf",
+   "metadata": {},
+   "source": [
+    "## set up wandb"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "18c561ab-ba82-4718-abe9-96c44133b314",
+   "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": "48cf7ec6-1b32-4f96-a0f1-c7ffb1cb5fcf",
+   "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/borzoi/wandb/run-20250304_221737-jxwlrbzp</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/borzoi/runs/jxwlrbzp' target=\"_blank\">copy-mouse</a></strong> to <a href='https://wandb.ai/grelu/borzoi' 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/borzoi' target=\"_blank\">https://wandb.ai/grelu/borzoi</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/borzoi/runs/jxwlrbzp' target=\"_blank\">https://wandb.ai/grelu/borzoi/runs/jxwlrbzp</a>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "run = wandb.init(entity='grelu', project='borzoi', job_type='copy', name='copy-mouse') # Initialize a W&B Run"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7fdb23c5-371e-4473-835c-cf21f82e610f",
+   "metadata": {},
+   "source": [
+    "## Paths"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "556b9768-f3e4-466a-b097-0b3226e788dd",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tasks_path = '/code/borzoi/data/targets_mouse.txt.gz'\n",
+    "intervals_path = '/gstore/data/resbioai/grelu/borzoi-data/mm10/sequences.bed'"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "de1c87dc-cacf-48f0-a081-5260e325cb7b",
+   "metadata": {},
+   "source": [
+    "## Process tasks"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "f8616dde-2567-4a3c-8fb9-de412d342078",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2608\n"
+     ]
+    }
+   ],
+   "source": [
+    "tasks = pd.read_table(tasks_path, index_col=0)\n",
+    "print(len(tasks))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "66f24ff2-74bc-41ee-9cc9-5cf9e75c2ebf",
+   "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>clip_soft</th>\n",
+       "      <th>scale</th>\n",
+       "      <th>sum_stat</th>\n",
+       "      <th>strand_pair</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>CNhs10464+</td>\n",
+       "      <td>/home/drk/tillage/datasets/mouse/cage/fantom/C...</td>\n",
+       "      <td>768</td>\n",
+       "      <td>384</td>\n",
+       "      <td>1.0</td>\n",
+       "      <td>sum</td>\n",
+       "      <td>1</td>\n",
+       "      <td>CAGE:placenta, adult pregnant day17</td>\n",
+       "      <td>CAGE</td>\n",
+       "      <td>placenta, adult pregnant day17</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>CNhs10464-</td>\n",
+       "      <td>/home/drk/tillage/datasets/mouse/cage/fantom/C...</td>\n",
+       "      <td>768</td>\n",
+       "      <td>384</td>\n",
+       "      <td>1.0</td>\n",
+       "      <td>sum</td>\n",
+       "      <td>0</td>\n",
+       "      <td>CAGE:placenta, adult pregnant day17</td>\n",
+       "      <td>CAGE</td>\n",
+       "      <td>placenta, adult pregnant day17</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "         name                                               file  clip  \\\n",
+       "0  CNhs10464+  /home/drk/tillage/datasets/mouse/cage/fantom/C...   768   \n",
+       "1  CNhs10464-  /home/drk/tillage/datasets/mouse/cage/fantom/C...   768   \n",
+       "\n",
+       "   clip_soft  scale sum_stat  strand_pair  \\\n",
+       "0        384    1.0      sum            1   \n",
+       "1        384    1.0      sum            0   \n",
+       "\n",
+       "                           description assay                          sample  \n",
+       "0  CAGE:placenta, adult pregnant day17  CAGE  placenta, adult pregnant day17  \n",
+       "1  CAGE:placenta, adult pregnant day17  CAGE  placenta, adult pregnant day17  "
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "tasks=tasks.reset_index(drop=True)\n",
+    "tasks[\"assay\"] = tasks[\"description\"].apply(lambda x: x.split(\":\")[0])\n",
+    "tasks[\"assay\"].value_counts()\n",
+    "tasks[\"sample\"] = tasks[\"description\"].apply(lambda x: \":\".join(x.split(\":\")[1:]))\n",
+    "tasks = tasks.rename(columns={\"identifier\":\"name\"})\n",
+    "tasks.head(2)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "4f96ecc7-238e-4707-a3af-64b0d1828bd4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tasks = tasks.to_dict(orient=\"list\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3056ecfa-1b75-46b7-9e36-4e72c4cb3618",
+   "metadata": {},
+   "source": [
+    "## Process intervals"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "e768d909-c699-47b6-86fd-c3ff54357c74",
+   "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>fold</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>chr1</td>\n",
+       "      <td>46257174</td>\n",
+       "      <td>46453782</td>\n",
+       "      <td>fold0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>chr2</td>\n",
+       "      <td>83512641</td>\n",
+       "      <td>83709249</td>\n",
+       "      <td>fold0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>chr7</td>\n",
+       "      <td>16218353</td>\n",
+       "      <td>16414961</td>\n",
+       "      <td>fold0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>chr3</td>\n",
+       "      <td>113724419</td>\n",
+       "      <td>113921027</td>\n",
+       "      <td>fold0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>chr3</td>\n",
+       "      <td>107470140</td>\n",
+       "      <td>107666748</td>\n",
+       "      <td>fold0</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "  chrom      start        end   fold\n",
+       "0  chr1   46257174   46453782  fold0\n",
+       "1  chr2   83512641   83709249  fold0\n",
+       "2  chr7   16218353   16414961  fold0\n",
+       "3  chr3  113724419  113921027  fold0\n",
+       "4  chr3  107470140  107666748  fold0"
+      ]
+     },
+     "execution_count": 8,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "intervals = pd.read_table(intervals_path, header=None)\n",
+    "intervals.columns = ['chrom', 'start', 'end', 'fold']\n",
+    "intervals.head()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "fc4f093c-9d9d-4b0d-9883-050099492944",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "split\n",
+       "train    36950\n",
+       "val       6318\n",
+       "test      6101\n",
+       "Name: count, dtype: int64"
+      ]
+     },
+     "execution_count": 9,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "intervals['split'] = 'train'\n",
+    "intervals.loc[intervals.fold=='fold3', 'split'] = 'test'\n",
+    "intervals.loc[intervals.fold=='fold4', 'split'] = 'val'\n",
+    "intervals.split.value_counts()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "7aafdb2e-93b3-406a-97fb-aab39d33aa33",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "196608"
+      ]
+     },
+     "execution_count": 10,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "get_unique_length(intervals)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "c45008d0-9dda-4660-8e1c-574570eef373",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "intervals = resize(intervals, 524288)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "86aae057-b106-4884-81b3-d1aa563ad4f7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_intervals = intervals[intervals.split=='train'].iloc[:, :3]\n",
+    "val_intervals = intervals[intervals.split=='val'].iloc[:, :3]\n",
+    "test_intervals = intervals[intervals.split=='test'].iloc[:, :3]\n",
+    "del intervals"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e4bd10a9-40c2-4299-aa74-6474ea5a71ef",
+   "metadata": {},
+   "source": [
+    "## Initialize model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "06dba49d-11ba-471b-83d7-05fe45a9bd5a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model_params = {\n",
+    "    \"model_type\":\"BorzoiModel\",\n",
+    "    \"n_tasks\":2608,\n",
+    "    \"final_act_func\": 'softplus',\n",
+    "    \"final_pool_func\":None,\n",
+    "    \"crop_len\":5120,\n",
+    "}\n",
+    "train_params={\n",
+    "    \"task\":\"regression\",\n",
+    "    \"loss\":\"mse\",\n",
+    "}\n",
+    "lm = LightningModel(model_params, train_params)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ab6a3055-d338-40f0-9fd3-f8f3134df643",
+   "metadata": {},
+   "source": [
+    "## Save checkpoints"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "bbccda9a-7a62-4472-9c0b-34abf98a7e1c",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/tmp/ipykernel_3294550/1480313343.py:18: 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(f\"/data/borzoi/torch_weights/mouse_fold{rep}.h5\")\n",
+      "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",
+      "/tmp/ipykernel_3294550/1480313343.py:18: 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(f\"/data/borzoi/torch_weights/mouse_fold{rep}.h5\")\n",
+      "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",
+      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3,4,5,6,7]\n",
+      "/tmp/ipykernel_3294550/1480313343.py:18: 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(f\"/data/borzoi/torch_weights/mouse_fold{rep}.h5\")\n",
+      "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",
+      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3,4,5,6,7]\n",
+      "/tmp/ipykernel_3294550/1480313343.py:18: 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(f\"/data/borzoi/torch_weights/mouse_fold{rep}.h5\")\n",
+      "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",
+      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3,4,5,6,7]\n"
+     ]
+    }
+   ],
+   "source": [
+    "for rep in range(4):\n",
+    "    \n",
+    "    lm.data_params[\"tasks\"] = tasks\n",
+    "\n",
+    "    lm.data_params[\"train\"] = dict()\n",
+    "    lm.data_params[\"val\"] = dict()\n",
+    "    lm.data_params[\"test\"] = dict()\n",
+    "    \n",
+    "    lm.data_params[\"train\"][\"seq_len\"] = 524288\n",
+    "    lm.data_params[\"train\"][\"label_len\"] = 6144*32\n",
+    "    lm.data_params[\"train\"][\"genome\"] = \"hg38\"\n",
+    "    lm.data_params[\"train\"][\"bin_size\"] = 32\n",
+    "    \n",
+    "    lm.data_params[\"train\"][\"intervals\"] = train_intervals.to_dict(orient='list')\n",
+    "    lm.data_params[\"val\"][\"intervals\"] = val_intervals.to_dict(orient='list')\n",
+    "    lm.data_params[\"test\"][\"intervals\"] = test_intervals.to_dict(orient='list')\n",
+    "\n",
+    "    state_dict = torch.load(f\"/data/borzoi/torch_weights/mouse_fold{rep}.h5\")\n",
+    "    lm.model.load_state_dict(state_dict)\n",
+    "    \n",
+    "    trainer = pl.Trainer()\n",
+    "    try:\n",
+    "        trainer.predict(lm)\n",
+    "    except:\n",
+    "        pass\n",
+    "    \n",
+    "    trainer.save_checkpoint(f'/data/borzoi/torch_weights/mouse_rep{rep}.ckpt')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4b5fbc09-abca-4cb7-a8b7-7eeba7b39cf8",
+   "metadata": {},
+   "source": [
+    "## Upload to wandb"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "id": "c53a69b4-9092-44b9-b8a9-20b72aa8b511",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "metadata={\n",
+    "            'model_params':lm.model_params, \n",
+    "            'train_params':lm.train_params, \n",
+    "            'data_params':lm.data_params\n",
+    "        }"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "id": "2a25b878-ca1f-4872-9549-ab502975ca2e",
+   "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 312024 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 312024 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 312024 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 312024 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 312024 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 312024 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 312024 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 312024 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 312024 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 312024 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 312024 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 312024 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 312024 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 312024 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 312024 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 312024 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 312024 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 312024 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 312024 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 312024 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 312024 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 312024 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 312024 bytes\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Serializing object of type list that is 312024 bytes\n"
+     ]
+    }
+   ],
+   "source": [
+    "for rep in range(4):\n",
+    "    ckpt = f'/data/borzoi/torch_weights/mouse_rep{rep}.ckpt'\n",
+    "    artifact = wandb.Artifact(\n",
+    "        f'mouse_rep{rep}', \n",
+    "        type='model',\n",
+    "        metadata=metadata,\n",
+    "    )\n",
+    "    artifact.add_file(local_path=ckpt, name='model.ckpt')\n",
+    "    run.log_artifact(artifact)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "id": "98dd40ec-f276-4099-a4a7-39ccf1600a58",
+   "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/borzoi/runs/jxwlrbzp' target=\"_blank\">https://wandb.ai/grelu/borzoi/runs/jxwlrbzp</a><br> View project at: <a href='https://wandb.ai/grelu/borzoi' target=\"_blank\">https://wandb.ai/grelu/borzoi</a><br>Synced 5 W&B file(s), 0 media file(s), 8 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_221737-jxwlrbzp/logs</code>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "run.finish()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1c5d6a57-a2a1-495d-90ff-16508e4e236c",
+   "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
+}