fine-tuning notebook

notebooks_tutorials/02_fine_tuning.ipynb

@@ -0,0 +1,1318 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 🧬 Fine-Tuning a Model on BigWig Tracks Prediction\n",
+    "\n",
+    "This notebook demonstrates a **simplified fine-tuning setup** that enables training of a pre-trained Nucleotide Transformer v3 (NTv3) model to predict BigWig signal tracks directly from DNA sequences. The streamlined approach leverages a pre-trained NTv3 backbone as a feature extractor and adds a custom prediction head that outputs single-nucleotide resolution signal values for various genomic tracks (e.g., ChIP-seq, ATAC-seq, RNA-seq).\n",
+    "\n",
+    "**⚡ Key Advantage**: This simplified pipeline achieves close performance to more complex training approaches while enabling fast fine-tuning. The training speed benefits from the efficient NTv3 model architecture and depends on your hardware capabilities (GPU acceleration and multi-worker data loading significantly reduce training time). With NTv3 models, meaningful Pearson correlations can typically be reached within ~10minutes of training on a 32kb functional tracks prediction task.\n",
+    "\n",
+    "**🔧 Main Simplifications**: Compared to the full supervised tracks pipeline, this notebook simplifies several aspects to enable faster iteration:\n",
+    "\n",
+    "- **Data splits**: Uses simple chromosome-based train/val/test splits (e.g., assigning entire chromosomes to each split) instead of more complex region-based splits\n",
+    "- **Random sequence sampling**: The dataset randomly samples sequences from chromosomes/regions on-the-fly, rather than using pre-computed sliding windows\n",
+    "- **Constant learning rate**: Uses a fixed learning rate throughout training without learning rate scheduling\n",
+    "- **No gradient accumulation**: Implements simple step-based training without gradient accumulation, making the training loop more straightforward\n",
+    "\n",
+    "The pipeline walks through the complete fine-tuning workflow:\n",
+    "\n",
+    "- Loading genomic sequences from FASTA files and their corresponding BigWig signal tracks\n",
+    "- Setting up a PyTorch dataset with proper train/validation/test splits\n",
+    "- Configuring the model architecture with a custom linear head\n",
+    "- Implementing a training loop with appropriate loss functions and evaluation metrics\n",
+    "\n",
+    "This provides a clean interface for training and evaluation.\n",
+    "\n",
+    "The model architecture consists of a pre-trained NTv3 backbone that processes DNA sequences and a custom linear head that predicts BigWig signal values at single-nucleotide resolution. Predictions are center-cropped to focus on the central portion of the input sequence (configurable via `keep_target_center_fraction`), which helps reduce edge effects from sequence context windows. The training uses a Poisson-Multinomial loss function that captures both the scale and shape of the signal distributions, and evaluation is performed using Pearson correlation metrics on both scaled and raw predictions.\n",
+    "\n",
+    "If you're interested in using pre-trained models for inference without fine-tuning, or exploring different model architectures, please refer to other notebooks in this collection. This notebook focuses specifically on the simplified fine-tuning process, which is useful when you want to quickly adapt a pre-trained model to your specific genomic tracks or improve performance on particular cell types or experimental conditions.\n",
+    "\n",
+    "📝 Note for Google Colab users: This notebook is compatible with Colab! For faster training, make sure to enable GPU: Runtime → Change runtime type → GPU (T4 or better recommended).\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 0. 📦 Imports dependencies"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Install dependencies\n",
+    "!pip install pyfaidx pyBigWig torchmetrics transformers plotly"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import random\n",
+    "import functools\n",
+    "from typing import List, Dict, Callable\n",
+    "import os\n",
+    "import subprocess\n",
+    "from concurrent.futures import ThreadPoolExecutor, as_completed\n",
+    "\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.nn.functional as F\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "from torch.optim import AdamW\n",
+    "from transformers import AutoConfig, AutoModelForMaskedLM, AutoTokenizer\n",
+    "import numpy as np\n",
+    "import pyBigWig\n",
+    "from pyfaidx import Fasta\n",
+    "from torchmetrics import PearsonCorrCoef\n",
+    "import plotly.graph_objects as go\n",
+    "from IPython.display import display\n",
+    "from tqdm import tqdm"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 1. ⚙️ Configuration\n",
+    "\n",
+    "## Configuration Parameters\n",
+    "\n",
+    "### Model\n",
+    "- **`model_name`**: HuggingFace model name/identifier for the pretrained backbone model\n",
+    "\n",
+    "### Data\n",
+    "- **`data_cache_dir`**: Directory where downloaded data files (FASTA, bigWig) will be stored\n",
+    "- **`fasta_url`**: URL to download reference genome FASTA file\n",
+    "- **`bigwig_url_list`**: List of URLs for bigWig track files to download\n",
+    "- **`bigwig_file_ids`**: List of identifiers/names for bigWig tracks (set after downloading, used for model head and metrics)\n",
+    "- **`sequence_length`**: Length of input sequences in base pairs (bp)\n",
+    "- **`keep_target_center_fraction`**: Fraction of center sequence to keep for target prediction (crops edges to focus on center)\n",
+    "\n",
+    "### Training\n",
+    "- **`batch_size`**: Number of samples per batch\n",
+    "- **`learning_rate`**: Constant learning rate for optimizer\n",
+    "- **`weight_decay`**: L2 regularization coefficient for optimizer\n",
+    "- **`num_steps_training`**: Total number of training steps\n",
+    "- **`log_every_n_steps`**: Log training metrics every N steps\n",
+    "\n",
+    "### Validation\n",
+    "- **`validate_every_n_steps`**: Run validation every N steps\n",
+    "- **`num_validation_samples`**: Number of samples to use for validation set\n",
+    "\n",
+    "### General\n",
+    "- **`seed`**: Random seed for reproducibility\n",
+    "- **`device`**: Device to run training on (\"cuda\" or \"cpu\")\n",
+    "- **`num_workers`**: Number of worker processes for DataLoader (0 = single-threaded)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Using device: cpu\n"
+     ]
+    }
+   ],
+   "source": [
+    "config = {\n",
+    "    # Model\n",
+    "    \"model_name\": \"InstaDeepAI/NTv3_8M_pre\",\n",
+    "    \n",
+    "    # Data\n",
+    "    \"data_cache_dir\": \"./data\",\n",
+    "    \"fasta_url\": \"https://hgdownload.gi.ucsc.edu/goldenPath/hg38/bigZips/hg38.fa.gz\",\n",
+    "    \"bigwig_url_list\": [\n",
+    "        \"https://www.encodeproject.org/files/ENCFF055QKS/@@download/ENCFF055QKS.bigWig\",\n",
+    "        \"https://www.encodeproject.org/files/ENCFF214GOQ/@@download/ENCFF214GOQ.bigWig\",\n",
+    "        \"https://www.encodeproject.org/files/ENCFF592NIB/@@download/ENCFF592NIB.bigWig\",\n",
+    "        \"https://www.encodeproject.org/files/ENCFF921PHQ/@@download/ENCFF921PHQ.bigWig\",\n",
+    "    ],\n",
+    "    \"sequence_length\": 32_768,\n",
+    "    \"keep_target_center_fraction\": 0.375,\n",
+    "    \n",
+    "    # Training\n",
+    "    \"batch_size\": 32,\n",
+    "    \"num_steps_training\": 19932,\n",
+    "    \"log_every_n_steps\": 40,\n",
+    "    \"learning_rate\": 1e-5,\n",
+    "    \"weight_decay\": 0.01,\n",
+    "    \n",
+    "    # Validation\n",
+    "    \"validate_every_n_steps\": 400,\n",
+    "    \"num_validation_samples\": 1000,\n",
+    "\n",
+    "    # Test\n",
+    "    \"num_test_samples\": 10000,\n",
+    "    \n",
+    "    # General\n",
+    "    \"seed\": 17,\n",
+    "    \"device\": \"cuda\" if torch.cuda.is_available() else \"cpu\",\n",
+    "    \"num_workers\": 16,\n",
+    "}\n",
+    "\n",
+    "os.makedirs(config[\"data_cache_dir\"], exist_ok=True)\n",
+    "\n",
+    "# Extract filenames from URLs\n",
+    "def extract_filename_from_url(url: str) -> str:\n",
+    "    \"\"\"Extract filename from URL, handling query parameters.\"\"\"\n",
+    "    # Remove query parameters if present\n",
+    "    url_clean = url.split('?')[0]\n",
+    "    # Get the last part of the URL path\n",
+    "    return url_clean.split('/')[-1]\n",
+    "\n",
+    "# Create paths for downloaded files\n",
+    "fasta_path = os.path.join(config[\"data_cache_dir\"], extract_filename_from_url(config[\"fasta_url\"]).replace('.gz', ''))\n",
+    "bigwig_path_list = [\n",
+    "    os.path.join(config[\"data_cache_dir\"], extract_filename_from_url(url))\n",
+    "    for url in config[\"bigwig_url_list\"]\n",
+    "]\n",
+    "\n",
+    "\n",
+    "# TODO: find a way to link the experiment accession to bigwig file ids\n",
+    "# Create bigwig_file_ids from filenames (without extension)\n",
+    "config[\"bigwig_file_ids\"] = [\n",
+    "    # os.path.splitext(extract_filename_from_url(url))[0]\n",
+    "    # for url in config[\"bigwig_url_list\"]\n",
+    "    \"ENCSR325NFE\",\n",
+    "    \"ENCSR962OTG\",\n",
+    "    \"ENCSR619DQO_P\",\n",
+    "    \"ENCSR619DQO_M\",\n",
+    "]\n",
+    "\n",
+    "# Set random seed\n",
+    "torch.manual_seed(config[\"seed\"])\n",
+    "np.random.seed(config[\"seed\"])\n",
+    "\n",
+    "# Set device\n",
+    "device = torch.device(config[\"device\"])\n",
+    "print(f\"Using device: {device}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 2. 📥 Genome & Tracks Data Download\n",
+    "\n",
+    "Download the reference genome FASTA file and BigWig signal tracks from public repositories. These files contain the genomic sequences and experimental signal data (e.g., ChIP-seq, ATAC-seq) that we'll use for training."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def _download_file(url: str, output_path: str) -> None:\n",
+    "    \"\"\"Download a file from URL to output_path using wget.\"\"\"\n",
+    "    subprocess.run([\"wget\", \"-c\", url, \"-O\", output_path], check=True)\n",
+    "\n",
+    "# Prepare download tasks: (url, output_path)\n",
+    "download_tasks = []\n",
+    "\n",
+    "# FASTA file\n",
+    "fasta_filename = extract_filename_from_url(config[\"fasta_url\"])\n",
+    "fasta_gz_path = os.path.join(config[\"data_cache_dir\"], fasta_filename)\n",
+    "download_tasks.append((config[\"fasta_url\"], fasta_gz_path))\n",
+    "\n",
+    "# BigWig files\n",
+    "for bigwig_url in config[\"bigwig_url_list\"]:\n",
+    "    filename = extract_filename_from_url(bigwig_url)\n",
+    "    filepath = os.path.join(config[\"data_cache_dir\"], filename)\n",
+    "    download_tasks.append((bigwig_url, filepath))\n",
+    "\n",
+    "# Download files in parallel\n",
+    "max_workers = min(len(download_tasks), 8)\n",
+    "\n",
+    "print(f\"Downloading {len(download_tasks)} files using {max_workers} workers...\")\n",
+    "with ThreadPoolExecutor(max_workers=max_workers) as executor:\n",
+    "    # Submit all download tasks\n",
+    "    future_to_path = {\n",
+    "        executor.submit(_download_file, url, path): path\n",
+    "        for url, path in download_tasks\n",
+    "    }\n",
+    "    \n",
+    "    # Wait for all downloads to complete\n",
+    "    for future in as_completed(future_to_path):\n",
+    "        try:\n",
+    "            future.result()  # Raises exception if download failed\n",
+    "            path = future_to_path[future]\n",
+    "            print(f\"✓ Downloaded: {os.path.basename(path)}\")\n",
+    "        except Exception as e:\n",
+    "            path = future_to_path[future]\n",
+    "            raise RuntimeError(f\"Failed to download {path}: {e}\") from e\n",
+    "\n",
+    "# Extract FASTA file after download\n",
+    "print(f\"\\nExtracting {fasta_filename}...\")\n",
+    "subprocess.run([\"gunzip\", \"-f\", fasta_gz_path], check=True)\n",
+    "print(\"✓ Extraction complete\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Data Splits Definition"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "chrom_splits = {\n",
+    "    \"train\": [f\"chr{i}\" for i in range(1, 21)] + ['chrX', 'chrY'],\n",
+    "    \"val\": ['chr22'],\n",
+    "    \"test\": ['chr21']\n",
+    "}"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 3. 🧠 Model and tokenizer setup\n",
+    " \n",
+    "In this section, we set up the model and tokenizer. \n",
+    " \n",
+    "Our approach uses any suitable pretrained backbone from HuggingFace Transformers (for example, `InstaDeepAI/ntv3_650M_pre`),\n",
+    "which is then extended with an additional linear head. \n",
+    " \n",
+    "This linear head is trained for regression on a set of genomic tracks, \n",
+    "allowing the model to make predictions for each track at single nucleotide resolution.\n",
+    " \n",
+    "The following code wraps the HuggingFace model together with this regression head for the end-to-end task.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class LinearHead(nn.Module):\n",
+    "    \"\"\"A linear head that predicts one scalar value per track.\"\"\"\n",
+    "    def __init__(self, embed_dim: int, num_labels: int):\n",
+    "        super().__init__()\n",
+    "        self.layer_norm = nn.LayerNorm(embed_dim)\n",
+    "        self.head = nn.Linear(embed_dim, num_labels)\n",
+    "    \n",
+    "    def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
+    "        x = self.layer_norm(x)\n",
+    "        x = self.head(x)\n",
+    "        x = F.softplus(x)  # Ensure positive values\n",
+    "        return x\n",
+    "\n",
+    "\n",
+    "class HFModelWithHead(nn.Module):\n",
+    "    \"\"\"Simple model wrapper: HF backbone + bigwig head.\"\"\"\n",
+    "    \n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        model_name: str,\n",
+    "        bigwig_track_names: List[str],\n",
+    "        keep_target_center_fraction: float = 0.375,\n",
+    "    ):\n",
+    "        super().__init__()\n",
+    "        \n",
+    "        # Load config and model\n",
+    "        self.config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)\n",
+    "        self.backbone = AutoModelForMaskedLM.from_pretrained(\n",
+    "            model_name, \n",
+    "            trust_remote_code=True,\n",
+    "            config=self.config\n",
+    "        )\n",
+    "        \n",
+    "        self.keep_target_center_fraction = keep_target_center_fraction\n",
+    "\n",
+    "        if hasattr(self.config, \"embed_dim\"):\n",
+    "            embed_dim = self.config.embed_dim\n",
+    "        else:\n",
+    "            raise ValueError(f\"Could not determine embed_dim for {model_name}\")\n",
+    "        \n",
+    "        # Bigwig head (NTv3 outputs at single-nucleotide resolution)\n",
+    "        self.bigwig_head = LinearHead(embed_dim, len(bigwig_track_names))\n",
+    "        self.model_name = model_name\n",
+    "    \n",
+    "    def forward(self, tokens: torch.Tensor, **kwargs) -> Dict[str, torch.Tensor]:\n",
+    "        # Forward through backbone\n",
+    "        outputs = self.backbone(input_ids=tokens)\n",
+    "        embedding = outputs.hidden_states[-1]  # Last hidden state\n",
+    "        \n",
+    "        # Crop to center fraction\n",
+    "        if self.keep_target_center_fraction < 1.0:\n",
+    "            seq_len = embedding.shape[1]\n",
+    "            target_offset = int(seq_len * (1 - self.keep_target_center_fraction) // 2)\n",
+    "            target_length = seq_len - 2 * target_offset\n",
+    "            embedding = embedding[:, target_offset:target_offset + target_length, :]\n",
+    "        \n",
+    "        # Predict bigwig tracks\n",
+    "        bigwig_logits = self.bigwig_head(embedding)\n",
+    "        \n",
+    "        return {\"bigwig_tracks_logits\": bigwig_logits}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Load tokenizer\n",
+    "tokenizer = AutoTokenizer.from_pretrained(config[\"model_name\"], trust_remote_code=True)\n",
+    "\n",
+    "# Create model\n",
+    "model = HFModelWithHead(\n",
+    "    model_name=config[\"model_name\"],\n",
+    "    bigwig_track_names=config[\"bigwig_file_ids\"],\n",
+    "    keep_target_center_fraction=config[\"keep_target_center_fraction\"],\n",
+    ")\n",
+    "model = model.to(device)\n",
+    "model.train()\n",
+    "\n",
+    "print(f\"Model loaded: {config['model_name']}\")\n",
+    "print(f\"Number of bigwig tracks: {len(config['bigwig_file_ids'])}\")\n",
+    "print(f\"Model parameters: {sum(p.numel() for p in model.parameters()):,}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 4. 🔄 Data loading\n",
+    "\n",
+    "Create PyTorch datasets and data loaders that efficiently sample random genomic windows from the reference genome and extract corresponding BigWig signal values. The dataset handles sequence tokenization, target scaling, and chromosome-based train/val/test splits."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Process-local cache for BigWig file handles (one per worker process)\n",
+    "# This allows safe multi-worker DataLoader usage\n",
+    "_bigwig_cache = {}  # Maps (process_id, file_path) -> pyBigWig handle\n",
+    "\n",
+    "\n",
+    "def _get_bigwig_handle(bigwig_path: str) -> pyBigWig.pyBigWig:\n",
+    "    \"\"\"Get or create a BigWig file handle for the current process.\"\"\"\n",
+    "    process_id = os.getpid()\n",
+    "    cache_key = (process_id, bigwig_path)\n",
+    "    \n",
+    "    if cache_key not in _bigwig_cache:\n",
+    "        _bigwig_cache[cache_key] = pyBigWig.open(bigwig_path)\n",
+    "    \n",
+    "    return _bigwig_cache[cache_key]\n",
+    "\n",
+    "\n",
+    "class GenomeBigWigDataset(Dataset):\n",
+    "    \"\"\"\n",
+    "    Random genomic windows from a reference genome + bigWig signal.\n",
+    "\n",
+    "    Each sample:\n",
+    "        - picks a chromosome/region (from `chroms` or `regions`),\n",
+    "        - picks a random window of length `sequence_length`,\n",
+    "        - returns (sequence, signal, chrom, start, end).\n",
+    "\n",
+    "    This dataset is compatible with multi-worker DataLoaders. BigWig files\n",
+    "    are opened lazily using a process-local cache, ensuring each worker process\n",
+    "    has its own file handles and avoiding concurrent access issues.\n",
+    "\n",
+    "    Args\n",
+    "    ----\n",
+    "    fasta_path : str\n",
+    "        Path to the reference genome FASTA (e.g. hg38.fna).\n",
+    "    bigwig_path_list : str\n",
+    "        Path to the bigWig file (e.g. ENCFF884LDL.bigWig).\n",
+    "    chroms : List[str]\n",
+    "        Chromosome names as they appear in the bigWig (e.g. [\"chr1\", \"chr2\", ...]).\n",
+    "        Used for backward compatibility or when regions=None.\n",
+    "    sequence_length : int\n",
+    "        Length of each random window (in bp).\n",
+    "    num_samples : int\n",
+    "        Number of samples the dataset will provide (len(dataset)).\n",
+    "    tokenizer : AutoTokenizer\n",
+    "        Tokenizer to use for tokenization.\n",
+    "    transform_fn : Callable\n",
+    "        Function to transform/scaling bigwig targets.\n",
+    "    keep_target_center_fraction : float\n",
+    "        Fraction of center sequence to keep for target prediction (crops edges to focus on center).\n",
+    "    regions : List[tuple[str, int, int]] | None\n",
+    "        Optional list of regions as (chromosome, start, end) tuples.\n",
+    "        If provided, samples are drawn randomly from within these regions only.\n",
+    "        This matches the JAX pipeline approach using BED file splits.\n",
+    "        If None, samples from entire chromosomes in `chroms`.\n",
+    "    \"\"\"\n",
+    "\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        fasta_path: str,\n",
+    "        bigwig_path_list: list[str],\n",
+    "        chroms: List[str],\n",
+    "        sequence_length: int,\n",
+    "        num_samples: int,\n",
+    "        tokenizer: AutoTokenizer,\n",
+    "        transform_fn: Callable[[torch.Tensor], torch.Tensor],\n",
+    "        keep_target_center_fraction: float = 1.0,\n",
+    "    ):\n",
+    "        super().__init__()\n",
+    "\n",
+    "        self.fasta = Fasta(fasta_path, as_raw=True, sequence_always_upper=True)\n",
+    "        # Store paths instead of opening files immediately (for multi-worker compatibility)\n",
+    "        self.bigwig_path_list = bigwig_path_list\n",
+    "        self.sequence_length = sequence_length\n",
+    "        self.num_samples = num_samples\n",
+    "        self.tokenizer = tokenizer\n",
+    "        self.transform_fn = transform_fn  # Use pre-computed transform function\n",
+    "        self.keep_target_center_fraction = keep_target_center_fraction\n",
+    "        self.chroms = chroms\n",
+    "\n",
+    "        # Get chromosome lengths from first BigWig file (lazy, cached per process)\n",
+    "        # We need this for validation, so open temporarily\n",
+    "        bw_handle = _get_bigwig_handle(bigwig_path_list[0])\n",
+    "        bw_chrom_lengths = bw_handle.chroms()  # dict: chrom -> length\n",
+    "\n",
+    "        self.valid_chroms = []\n",
+    "        self.chrom_lengths = {}\n",
+    "\n",
+    "        for c in chroms:\n",
+    "            if c not in bw_chrom_lengths or c not in self.fasta:\n",
+    "                continue\n",
+    "\n",
+    "            fa_len = len(self.fasta[c])\n",
+    "            bw_len = bw_chrom_lengths[c]\n",
+    "            L = min(fa_len, bw_len)\n",
+    "\n",
+    "            if L > self.sequence_length:\n",
+    "                self.valid_chroms.append(c)\n",
+    "                self.chrom_lengths[c] = L\n",
+    "\n",
+    "        if not self.valid_chroms:\n",
+    "            raise ValueError(\"No valid chromosomes after intersecting FASTA and bigWig.\")\n",
+    "\n",
+    "    def __len__(self):\n",
+    "        return self.num_samples\n",
+    "\n",
+    "    def __getitem__(self, idx):\n",
+    "\n",
+    "        # Sample from entire chromosomes\n",
+    "        chrom = random.choice(self.valid_chroms)\n",
+    "        chrom_len = self.chrom_lengths[chrom]\n",
+    "        max_start = chrom_len - self.sequence_length\n",
+    "        start = random.randint(0, max_start)\n",
+    "        end = start + self.sequence_length\n",
+    "\n",
+    "        # Sequence\n",
+    "        seq = self.fasta[chrom][start:end]  # string slice\n",
+    "        # Tokenize with padding and truncation to ensure consistent lengths for batching\n",
+    "        tokenized = self.tokenizer(\n",
+    "            seq,\n",
+    "            padding=\"max_length\",\n",
+    "            truncation=True,\n",
+    "            max_length=self.sequence_length,\n",
+    "            return_tensors=\"pt\",\n",
+    "        )\n",
+    "        tokens = tokenized[\"input_ids\"][0]  # Shape: (max_length,)\n",
+    "\n",
+    "        # Signal from bigWig tracks (numpy array) -> torch tensor\n",
+    "        # Get BigWig handles lazily (cached per worker process)\n",
+    "        bigwig_targets = np.array([\n",
+    "            _get_bigwig_handle(bw_path).values(chrom, start, end, numpy=True)\n",
+    "            for bw_path in self.bigwig_path_list\n",
+    "        ])  # shape (num_tracks, seq_len)\n",
+    "        # Transpose to (seq_len, num_tracks)\n",
+    "        bigwig_targets = bigwig_targets.T\n",
+    "        # pyBigWig returns NaN where no data; turn NaN into 0\n",
+    "        bigwig_targets = torch.tensor(bigwig_targets, dtype=torch.float32)\n",
+    "        bigwig_targets = torch.nan_to_num(bigwig_targets, nan=0.0)\n",
+    "        \n",
+    "        # Crop targets to center fraction\n",
+    "        if self.keep_target_center_fraction < 1.0:\n",
+    "            seq_len = bigwig_targets.shape[0]  # First dimension is sequence length\n",
+    "            target_offset = int(seq_len * (1 - self.keep_target_center_fraction) // 2)\n",
+    "            target_length = seq_len - 2 * target_offset\n",
+    "            bigwig_targets = bigwig_targets[target_offset:target_offset + target_length, :]\n",
+    "\n",
+    "        # Apply scaling to targets\n",
+    "        bigwig_targets = self.transform_fn(bigwig_targets)\n",
+    "\n",
+    "        sample = {\n",
+    "            \"tokens\": tokens,\n",
+    "            \"bigwig_targets\": bigwig_targets,\n",
+    "            \"chrom\": chrom,\n",
+    "            \"start\": start,\n",
+    "            \"end\": end,\n",
+    "        }\n",
+    "        return sample"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Data preprocessing utilities"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Scaling functions for targets\n",
+    "def compute_chromosome_stats(track_data: np.ndarray) -> dict:\n",
+    "    \"\"\"\n",
+    "    Compute minimal statistics needed for weighted mean computation.\n",
+    "    \n",
+    "    Args:\n",
+    "        track_data: numpy array of track values for a chromosome\n",
+    "        \n",
+    "    Returns:\n",
+    "        Dictionary with statistics: sum, mean, total_count\n",
+    "    \"\"\"\n",
+    "    track_data = track_data.astype(np.float32)\n",
+    "    \n",
+    "    # Compute statistics\n",
+    "    sum_all = np.sum(track_data)\n",
+    "    total_count = track_data.size\n",
+    "    mean_all = sum_all / total_count if total_count > 0 else 0.0\n",
+    "    \n",
+    "    return {\n",
+    "        \"sum\": sum_all,\n",
+    "        \"mean\": mean_all,\n",
+    "        \"total_count\": total_count,\n",
+    "    }\n",
+    "\n",
+    "\n",
+    "def aggregate_file_statistics(chr_stats_list: List[dict]) -> dict:\n",
+    "    \"\"\"\n",
+    "    Aggregate chromosome-level statistics into file-level statistics.\n",
+    "    \n",
+    "    Args:\n",
+    "        chr_stats_list: List of dictionaries, each containing chromosome-level statistics\n",
+    "        \n",
+    "    Returns:\n",
+    "        Dictionary with aggregated file-level statistics (only mean)\n",
+    "    \"\"\"\n",
+    "    # Convert to arrays for easier computation\n",
+    "    total_counts = np.array([s[\"total_count\"] for s in chr_stats_list], dtype=np.int64)\n",
+    "    means = np.array([s[\"mean\"] for s in chr_stats_list], dtype=np.float32)\n",
+    "    sums = np.array([s[\"sum\"] for s in chr_stats_list], dtype=np.float32)\n",
+    "    \n",
+    "    # Aggregate total count\n",
+    "    total_count = np.sum(total_counts)\n",
+    "    \n",
+    "    # Weighted mean: mean = sum(mean_chr * total_count_chr) / sum(total_count_chr)\n",
+    "    mean = np.sum(means * total_counts) / total_count if total_count > 0 else 0.0\n",
+    "    \n",
+    "    return {\n",
+    "        \"total_count\": total_count,\n",
+    "        \"sum\": np.sum(sums),\n",
+    "        \"mean\": mean,\n",
+    "    }\n",
+    "\n",
+    "\n",
+    "def get_track_means(bigwig_tracks_list: List[pyBigWig.pyBigWig]) -> np.ndarray:\n",
+    "    \"\"\"\n",
+    "    Get track means for normalization.\n",
+    "    Computes statistics per chromosome and aggregates using weighted averaging,\n",
+    "    \n",
+    "    Args:\n",
+    "        bigwig_tracks_list: List of pyBigWig file objects\n",
+    "        \n",
+    "    Returns:\n",
+    "        Array of track means, one per bigwig file\n",
+    "    \"\"\"\n",
+    "    track_means = []\n",
+    "    \n",
+    "    for bigwig_track in bigwig_tracks_list:\n",
+    "        chrom_lengths = bigwig_track.chroms()\n",
+    "        all_chr_stats = []\n",
+    "        \n",
+    "        # Compute statistics for each chromosome\n",
+    "        for chrom_name, chrom_length in chrom_lengths.items():\n",
+    "            try:\n",
+    "                # Get chromosome data as numpy array\n",
+    "                bw_array = np.array(\n",
+    "                    bigwig_track.values(chrom_name, 0, chrom_length, numpy=True),\n",
+    "                    dtype=np.float32\n",
+    "                )\n",
+    "                # Replace NaN with 0\n",
+    "                bw_array = np.nan_to_num(bw_array, nan=0.0)\n",
+    "                \n",
+    "                # Compute chromosome-level statistics\n",
+    "                chr_stats = compute_chromosome_stats(bw_array)\n",
+    "                all_chr_stats.append(chr_stats)\n",
+    "            except Exception as e:\n",
+    "                # Skip chromosomes that fail to load\n",
+    "                print(f\"Warning: Failed to load chromosome {chrom_name}: {e}\")\n",
+    "                continue\n",
+    "        \n",
+    "        if not all_chr_stats:\n",
+    "            raise ValueError(f\"No valid chromosomes found for bigwig track\")\n",
+    "        \n",
+    "        # Aggregate chromosome-level stats into file-level stats\n",
+    "        file_stats = aggregate_file_statistics(all_chr_stats)\n",
+    "        \n",
+    "        # Use the weighted mean for normalization\n",
+    "        track_means.append(file_stats[\"mean\"])\n",
+    "    \n",
+    "    return np.array(track_means, dtype=np.float32)\n",
+    "\n",
+    "\n",
+    "def create_targets_scaling_fn(bigwig_path_list: List[str]) -> Callable[[torch.Tensor], torch.Tensor]:\n",
+    "    \"\"\"\n",
+    "    Build a scaling function based on track means computed from bigwig files.\n",
+    "    \n",
+    "    Opens bigwig files, computes track statistics, and creates a transform function.\n",
+    "    The statistics are computed once and reused for all calls to the returned transform function.\n",
+    "    \n",
+    "    Args:\n",
+    "        bigwig_path_list: List of paths to bigwig files\n",
+    "        \n",
+    "    Returns:\n",
+    "        Transform function that scales input tensors\n",
+    "    \"\"\"\n",
+    "    # Open bigwig files and compute track statistics\n",
+    "    print(\"Computing track statistics (this may take a while)...\")\n",
+    "    bw_list = [\n",
+    "        pyBigWig.open(bigwig_path)\n",
+    "        for bigwig_path in bigwig_path_list\n",
+    "    ]\n",
+    "    track_means = get_track_means(bw_list)\n",
+    "    print(f\"Computed track means: {track_means}\")\n",
+    "    print(f\"Track means shape: {track_means.shape}\")\n",
+    "    \n",
+    "    # Create tensor from computed means\n",
+    "    track_means_tensor = torch.tensor(track_means, dtype=torch.float32)\n",
+    "    \n",
+    "    def transform_fn(x: torch.Tensor) -> torch.Tensor:\n",
+    "        \"\"\"\n",
+    "        x: torch.Tensor, shape (seq_len, num_tracks) or (batch, seq_len, num_tracks)\n",
+    "        \"\"\"\n",
+    "        # Move constants to correct device then normalize\n",
+    "        means = track_means_tensor.to(x.device)\n",
+    "        scaled = x / means\n",
+    "\n",
+    "        # Smooth clipping: if > 10, apply formula\n",
+    "        clipped = torch.where(\n",
+    "            scaled > 10.0,\n",
+    "            2.0 * torch.sqrt(scaled * 10.0) - 10.0,\n",
+    "            scaled,\n",
+    "        )\n",
+    "        return clipped\n",
+    "    \n",
+    "    return transform_fn"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Create scaling function\n",
+    "targets_transform_fn = create_targets_scaling_fn(bigwig_path_list)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Create datasets & dataloaders\n",
+    "create_dataset_fn = functools.partial(\n",
+    "    GenomeBigWigDataset,\n",
+    "    fasta_path=fasta_path,\n",
+    "    bigwig_path_list=bigwig_path_list,\n",
+    "    sequence_length=config[\"sequence_length\"],\n",
+    "    tokenizer=tokenizer,\n",
+    "    transform_fn=targets_transform_fn,\n",
+    "    keep_target_center_fraction=config[\"keep_target_center_fraction\"],\n",
+    ")\n",
+    "\n",
+    "train_dataset = create_dataset_fn(\n",
+    "    chroms=chrom_splits[\"train\"],\n",
+    "    num_samples=config[\"num_steps_training\"] * config[\"batch_size\"],\n",
+    ")\n",
+    "\n",
+    "val_dataset = create_dataset_fn(\n",
+    "    chroms=chrom_splits[\"val\"],\n",
+    "    num_samples=config[\"num_validation_samples\"],\n",
+    ")\n",
+    "\n",
+    "test_dataset = create_dataset_fn(\n",
+    "    chroms=chrom_splits[\"test\"],\n",
+    "    num_samples=config[\"num_test_samples\"],\n",
+    ")\n",
+    "\n",
+    "# Create dataloaders\n",
+    "train_loader = DataLoader(\n",
+    "    train_dataset,\n",
+    "    batch_size=config[\"batch_size\"],\n",
+    "    shuffle=True,\n",
+    "    num_workers=config[\"num_workers\"],\n",
+    ")\n",
+    "\n",
+    "val_loader = DataLoader(\n",
+    "    val_dataset,\n",
+    "    batch_size=config[\"batch_size\"],\n",
+    "    shuffle=False,\n",
+    "    num_workers=config[\"num_workers\"],\n",
+    ")\n",
+    "\n",
+    "test_loader = DataLoader(\n",
+    "    test_dataset,\n",
+    "    batch_size=config[\"batch_size\"],\n",
+    "    shuffle=False,\n",
+    "    num_workers=config[\"num_workers\"],\n",
+    ")\n",
+    "\n",
+    "print(f\"Train samples: {len(train_dataset)}\")\n",
+    "print(f\"Val samples: {len(val_dataset)}\")\n",
+    "print(f\"Test samples: {len(test_dataset)}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 5. ⚙️ Optimizer setup\n",
+    "\n",
+    "Configure the AdamW optimizer with learning rate and weight decay hyperparameters. This optimizer will update the model parameters during training to minimize the loss function.\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Training setup\n",
+    "print(f\"Training configuration:\")\n",
+    "print(f\"  Batch size: {config[\"batch_size\"]}\")\n",
+    "print(f\"  Total training steps: {config[\"num_steps_training\"]}\")\n",
+    "print(f\"  Log metrics every: {config[\"log_every_n_steps\"]} steps\")\n",
+    "print(f\"  Validate every: {config[\"validate_every_n_steps\"]} steps\")\n",
+    "\n",
+    "# Setup optimizer\n",
+    "optimizer = AdamW(\n",
+    "    model.parameters(),\n",
+    "    lr=config[\"learning_rate\"],\n",
+    "    weight_decay=config[\"weight_decay\"],\n",
+    ")\n",
+    "\n",
+    "print(f\"\\nOptimizer setup:\")\n",
+    "print(f\"  Learning rate: {config['learning_rate']}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 6. 📊 Metrics setup\n",
+    "\n",
+    "Set up evaluation metrics to track model performance during training and validation. We use Pearson correlation coefficients to measure how well the predicted BigWig signals match the ground truth signals."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class TracksMetrics:\n",
+    "    \"\"\"Simple metrics tracker for tracks prediction.\"\"\"\n",
+    "    \n",
+    "    def __init__(self, track_names: List[str]):\n",
+    "        self.track_names = track_names\n",
+    "        self.num_tracks = len(track_names)\n",
+    "        # Metrics: comparing scaled targets with scaled predictions\n",
+    "        # Configure to use float64 for improved numerical stability\n",
+    "        self.pearson_metrics = [\n",
+    "            PearsonCorrCoef().to(device) for _ in range(self.num_tracks)\n",
+    "        ]\n",
+    "        # Set dtype to float64 to prevent NaN warnings when variance is close to zero\n",
+    "        for metric in self.pearson_metrics:\n",
+    "            metric.set_dtype(torch.float64)\n",
+    "        self.losses = []\n",
+    "    \n",
+    "    def reset(self):\n",
+    "        for metric in self.pearson_metrics:\n",
+    "            metric.reset()\n",
+    "        self.losses = []\n",
+    "    \n",
+    "    def update(\n",
+    "        self, \n",
+    "        predictions: torch.Tensor, \n",
+    "        targets: torch.Tensor,\n",
+    "        loss: float\n",
+    "    ):\n",
+    "        \"\"\"\n",
+    "        Update metrics.\n",
+    "        Args:\n",
+    "            predictions: (batch, seq_len, num_tracks)\n",
+    "            targets: (batch, seq_len, num_tracks)\n",
+    "            loss: scalar loss value\n",
+    "        \"\"\"\n",
+    "        # Flatten batch and sequence dimensions\n",
+    "        pred_flat = predictions.detach().reshape(-1, self.num_tracks)  # (N, num_tracks)\n",
+    "        target_flat = targets.detach().reshape(-1, self.num_tracks)  # (N, num_tracks)\n",
+    "        \n",
+    "        # Convert to float64 for improved numerical stability in Pearson correlation\n",
+    "        pred_flat = pred_flat.to(torch.float64)\n",
+    "        target_flat = target_flat.to(torch.float64)\n",
+    "        \n",
+    "        # Update metrics\n",
+    "        for i, metric in enumerate(self.pearson_metrics):\n",
+    "            metric.update(pred_flat[:, i], target_flat[:, i])\n",
+    "        \n",
+    "        self.losses.append(loss)\n",
+    "    \n",
+    "    def compute(self) -> Dict[str, float]:\n",
+    "        \"\"\"Compute and return all metrics.\"\"\"\n",
+    "        metrics_dict = {}\n",
+    "        \n",
+    "        # Per-track Pearson correlations\n",
+    "        for i, (track_name, metric) in enumerate(zip(self.track_names, self.pearson_metrics)):\n",
+    "            corr = metric.compute().item()\n",
+    "            metrics_dict[f\"{track_name}/pearson\"] = corr\n",
+    "        \n",
+    "        # Mean Pearson correlation\n",
+    "        correlations = [metric.compute().item() for metric in self.pearson_metrics]\n",
+    "        metrics_dict[\"mean/pearson\"] = np.nanmean(correlations)\n",
+    "        \n",
+    "        # Mean loss\n",
+    "        metrics_dict[\"loss\"] = np.mean(self.losses) if self.losses else 0.0\n",
+    "        \n",
+    "        return metrics_dict"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_metrics = TracksMetrics(config[\"bigwig_file_ids\"])\n",
+    "val_metrics = TracksMetrics(config[\"bigwig_file_ids\"])\n",
+    "test_metrics = TracksMetrics(config[\"bigwig_file_ids\"])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 7. 📉 Loss functions\n",
+    "\n",
+    "Define the Poisson-Multinomial loss function that captures both the scale (total signal) and shape (distribution) of BigWig tracks. This loss is specifically designed for count-based genomic signal data."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def poisson_loss(ytrue: torch.Tensor, ypred: torch.Tensor, epsilon: float = 1e-7) -> torch.Tensor:\n",
+    "    \"\"\"Poisson loss per element: ypred - ytrue * log(ypred).\"\"\"\n",
+    "    return ypred - ytrue * torch.log(ypred + epsilon)\n",
+    "\n",
+    "\n",
+    "def safe_for_grad_log_torch(x: torch.Tensor) -> torch.Tensor:\n",
+    "    \"\"\"Guarantees that the log is defined for all x > 0 in a differentiable way.\"\"\"\n",
+    "    return torch.log(torch.where(x > 0.0, x, torch.ones_like(x)))\n",
+    "\n",
+    "\n",
+    "def poisson_multinomial_loss(\n",
+    "    logits: torch.Tensor,\n",
+    "    targets: torch.Tensor,\n",
+    "    shape_loss_coefficient: float = 5.0,\n",
+    "    epsilon: float = 1e-7,\n",
+    ") -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n",
+    "    \"\"\"\n",
+    "    Regression loss for bigwig tracks (Poisson-Multinomial).\n",
+    "    \n",
+    "    Args:\n",
+    "        logits: (batch, seq_length, num_tracks) - predicted counts\n",
+    "        targets: (batch, seq_length, num_tracks) - target counts\n",
+    "        shape_loss_coefficient: coefficient to weight scale loss\n",
+    "        epsilon: epsilon for numerical stability\n",
+    "    \n",
+    "    Returns:\n",
+    "        loss, scale_loss, shape_loss\n",
+    "    \"\"\"\n",
+    "    batch_size, seq_length, num_tracks = logits.shape\n",
+    "    \n",
+    "    # Scale loss: Poisson loss on total counts per sequence per track\n",
+    "    # Sum over sequence dimension (axis=1)\n",
+    "    sum_pred = logits.sum(dim=1)  # (batch, num_tracks)\n",
+    "    sum_true = targets.sum(dim=1)  # (batch, num_tracks)\n",
+    "    \n",
+    "    # Compute poisson loss per (batch, track)\n",
+    "    scale_loss = poisson_loss(sum_true, sum_pred, epsilon=epsilon)  # (batch, num_tracks)\n",
+    "    \n",
+    "    # Normalize by sequence length\n",
+    "    scale_loss = scale_loss / (seq_length + epsilon)\n",
+    "    \n",
+    "    # Average over batch and tracks\n",
+    "    scale_loss = scale_loss.mean()\n",
+    "    \n",
+    "    # Shape loss: Multinomial loss\n",
+    "    # Add epsilon to all positions\n",
+    "    predicted_counts = logits + epsilon\n",
+    "    targets_with_epsilon = targets + epsilon\n",
+    "    \n",
+    "    # Normalize predictions to get probabilities\n",
+    "    denom = predicted_counts.sum(dim=1, keepdim=True) + epsilon  # (batch, 1, num_tracks)\n",
+    "    p_pred = predicted_counts / denom\n",
+    "    \n",
+    "    # Compute shape loss: -sum(targets * log(p_pred))\n",
+    "    pl_pred = safe_for_grad_log_torch(p_pred)\n",
+    "    shape_loss = -(targets_with_epsilon * pl_pred)\n",
+    "    \n",
+    "    # Sum over all dimensions and normalize by total number of positions\n",
+    "    shape_denom = batch_size * seq_length * num_tracks + epsilon\n",
+    "    shape_loss = shape_loss.sum() / shape_denom\n",
+    "    \n",
+    "    # Combine losses\n",
+    "    loss = shape_loss + scale_loss / shape_loss_coefficient\n",
+    "\n",
+    "    return loss, scale_loss, shape_loss\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 8. 🏃 Training loop\n",
+    "\n",
+    "Run the main training loop that iterates through batches, computes gradients, and updates model parameters. The loop includes periodic validation checks and real-time metric visualization to monitor training progress."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def train_step(\n",
+    "    model: nn.Module,\n",
+    "    batch: Dict[str, torch.Tensor],\n",
+    ") -> float:\n",
+    "    \"\"\"Single training step.\"\"\"\n",
+    "    tokens = batch[\"tokens\"].to(device)\n",
+    "    bigwig_targets = batch[\"bigwig_targets\"].to(device)\n",
+    "    \n",
+    "    # Forward pass\n",
+    "    outputs = model(tokens=tokens)\n",
+    "    bigwig_logits = outputs[\"bigwig_tracks_logits\"]\n",
+    "    \n",
+    "    # Compute loss\n",
+    "    loss, _, _ = poisson_multinomial_loss(\n",
+    "        logits=bigwig_logits,\n",
+    "        targets=bigwig_targets,\n",
+    "    )\n",
+    "    \n",
+    "    # Backward pass\n",
+    "    loss.backward()\n",
+    "    return loss.item()\n",
+    "\n",
+    "\n",
+    "def validation_step(\n",
+    "    model: nn.Module,\n",
+    "    batch: Dict[str, torch.Tensor],\n",
+    "    metrics: TracksMetrics,\n",
+    ") -> float:\n",
+    "    \"\"\"Single validation step.\"\"\"\n",
+    "    model.eval()\n",
+    "    \n",
+    "    tokens = batch[\"tokens\"].to(device)\n",
+    "    bigwig_targets = batch[\"bigwig_targets\"].to(device)\n",
+    "    \n",
+    "    with torch.no_grad():\n",
+    "        # Forward pass\n",
+    "        outputs = model(tokens=tokens)\n",
+    "        bigwig_logits = outputs[\"bigwig_tracks_logits\"]\n",
+    "        \n",
+    "        # Compute loss\n",
+    "        loss, _, _ = poisson_multinomial_loss(\n",
+    "            logits=bigwig_logits,\n",
+    "            targets=bigwig_targets,\n",
+    "        )\n",
+    "        \n",
+    "        # Update metrics\n",
+    "        metrics.update(\n",
+    "            predictions=bigwig_logits,\n",
+    "            targets=bigwig_targets,\n",
+    "            loss=loss.item()\n",
+    "        )\n",
+    "    \n",
+    "    return loss.item()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Interactive plotting is temporary for debug"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Training loop\n",
+    "print(\"Starting training...\")\n",
+    "print(f\"Training for {config['num_steps_training']} steps\\n\")\n",
+    "\n",
+    "model.train()\n",
+    "train_metrics.reset()\n",
+    "optimizer.zero_grad()  # Initialize gradients\n",
+    "\n",
+    "# Track metrics for plotting\n",
+    "train_steps = []\n",
+    "train_losses = []\n",
+    "train_pearson_scores = []\n",
+    "val_steps = []\n",
+    "val_losses = []\n",
+    "val_pearson_scores = []\n",
+    "\n",
+    "# Initialize interactive plots using FigureWidget for real-time updates\n",
+    "from plotly.graph_objects import FigureWidget\n",
+    "from plotly.subplots import make_subplots\n",
+    "\n",
+    "# Create base figure with subplots\n",
+    "fig_base = make_subplots(\n",
+    "    rows=1, cols=2,\n",
+    "    subplot_titles=('Loss', 'Mean Pearson Correlation'),\n",
+    "    horizontal_spacing=0.15,\n",
+    ")\n",
+    "\n",
+    "# Add empty traces for train and val metrics\n",
+    "fig_base.add_trace(\n",
+    "    go.Scatter(x=[], y=[], mode='lines+markers', name='Train Loss', line=dict(color='blue')),\n",
+    "    row=1, col=1\n",
+    ")\n",
+    "fig_base.add_trace(\n",
+    "    go.Scatter(x=[], y=[], mode='lines+markers', name='Val Loss', line=dict(color='red')),\n",
+    "    row=1, col=1\n",
+    ")\n",
+    "fig_base.add_trace(\n",
+    "    go.Scatter(x=[], y=[], mode='lines+markers', name='Train Pearson', line=dict(color='green')),\n",
+    "    row=1, col=2\n",
+    ")\n",
+    "fig_base.add_trace(\n",
+    "    go.Scatter(x=[], y=[], mode='lines+markers', name='Val Pearson', line=dict(color='orange')),\n",
+    "    row=1, col=2\n",
+    ")\n",
+    "\n",
+    "fig_base.update_xaxes(title_text=\"Step\", row=1, col=1)\n",
+    "fig_base.update_xaxes(title_text=\"Step\", row=1, col=2)\n",
+    "fig_base.update_yaxes(title_text=\"Loss\", row=1, col=1)\n",
+    "fig_base.update_yaxes(title_text=\"Pearson Correlation\", row=1, col=2)\n",
+    "fig_base.update_layout(height=800, width=1600, showlegend=True)\n",
+    "\n",
+    "# Convert to FigureWidget for interactive updates\n",
+    "fig = FigureWidget(fig_base)\n",
+    "\n",
+    "# Display initial plot (will update in place during training)\n",
+    "display(fig)\n",
+    "\n",
+    "# Create iterator for training data (will cycle if needed)\n",
+    "train_iter = iter(train_loader)\n",
+    "\n",
+    "# Main training loop\n",
+    "for step_idx in range(config[\"num_steps_training\"]):\n",
+    "    try:\n",
+    "        batch = next(train_iter)\n",
+    "    except StopIteration:\n",
+    "        # Restart iterator if we run out of data\n",
+    "        train_iter = iter(train_loader)\n",
+    "        batch = next(train_iter)\n",
+    "    \n",
+    "    # Forward pass and backward pass\n",
+    "    loss = train_step(model, batch)\n",
+    "    \n",
+    "    # Update optimizer\n",
+    "    optimizer.step()\n",
+    "    optimizer.zero_grad()\n",
+    "    \n",
+    "    # Update metrics\n",
+    "    tokens = batch[\"tokens\"].to(device)\n",
+    "    bigwig_targets = batch[\"bigwig_targets\"].to(device)\n",
+    "    with torch.no_grad():\n",
+    "        outputs = model(tokens=tokens)\n",
+    "        bigwig_logits = outputs[\"bigwig_tracks_logits\"]\n",
+    "        \n",
+    "        train_metrics.update(\n",
+    "            predictions=bigwig_logits,\n",
+    "            targets=bigwig_targets,\n",
+    "            loss=loss\n",
+    "        )\n",
+    "    \n",
+    "    # Logging\n",
+    "    if (step_idx + 1) % config[\"log_every_n_steps\"] == 0:\n",
+    "        train_metrics_dict = train_metrics.compute()\n",
+    "        \n",
+    "        # Get accumulated mean loss across all batches since last reset\n",
+    "        mean_loss = train_metrics_dict['loss']\n",
+    "        \n",
+    "        # Track metrics for plotting\n",
+    "        train_steps.append(step_idx + 1)\n",
+    "        train_losses.append(mean_loss)\n",
+    "        train_pearson_scores.append(train_metrics_dict['mean/pearson'])\n",
+    "        \n",
+    "        # Update plots - direct assignment to FigureWidget data updates the plot automatically\n",
+    "        fig.data[0].x = train_steps\n",
+    "        fig.data[0].y = train_losses\n",
+    "        fig.data[2].x = train_steps\n",
+    "        fig.data[2].y = train_pearson_scores\n",
+    "        \n",
+    "        print(\n",
+    "            f\"Step {step_idx + 1}/{config['num_steps_training']} | \"\n",
+    "            f\"Loss: {mean_loss:.4f} | \"\n",
+    "            f\"Mean Pearson: {train_metrics_dict['mean/pearson']:.4f} | \"\n",
+    "            f\"Pearson per track: {train_metrics_dict[f'{track_name}/pearson']:.4f for track_name in config['bigwig_file_ids']}\"\n",
+    "        )\n",
+    "        train_metrics.reset()\n",
+    "    \n",
+    "    # Validation\n",
+    "    if (step_idx + 1) % config[\"validate_every_n_steps\"] == 0:\n",
+    "        print(f\"\\nRunning validation at step {step_idx + 1}...\")\n",
+    "        val_metrics.reset()\n",
+    "        model.eval()\n",
+    "        \n",
+    "        for val_batch in val_loader:\n",
+    "            val_loss = validation_step(model, val_batch, val_metrics)\n",
+    "        \n",
+    "        # Print validation metrics\n",
+    "        val_metrics_dict = val_metrics.compute()\n",
+    "        val_pearson_mean = val_metrics_dict['mean/pearson']\n",
+    "        \n",
+    "        # Track validation metrics\n",
+    "        val_steps.append(step_idx + 1)\n",
+    "        val_losses.append(val_metrics_dict['loss'])\n",
+    "        val_pearson_scores.append(val_pearson_mean)\n",
+    "        \n",
+    "        # Update plots with validation data - direct assignment updates the plot automatically\n",
+    "        fig.data[1].x = val_steps\n",
+    "        fig.data[1].y = val_losses\n",
+    "        fig.data[3].x = val_steps\n",
+    "        fig.data[3].y = val_pearson_scores\n",
+    "        \n",
+    "        print(f\"  Validation Loss: {val_metrics_dict['loss']:.4f}\")\n",
+    "        print(f\"  Validation Mean Pearson: {val_pearson_mean:.4f}\")\n",
+    "        for track_name in config[\"bigwig_file_ids\"]:\n",
+    "            print(f\"    {track_name}/pearson: {val_metrics_dict[f'{track_name}/pearson']:.4f}\")\n",
+    "        \n",
+    "        model.train()  # Back to training mode\n",
+    "\n",
+    "print(f\"\\nTraining completed after {config['num_steps_training']} steps.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 9. 🧪 Test evaluation\n",
+    "\n",
+    "Evaluate the fine-tuned model on the held-out test set to assess final performance. This provides an unbiased estimate of how well the model generalizes to unseen genomic regions."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Calculate number of test steps (based on deepspeed pipeline)\n",
+    "num_test_samples = len(test_dataset)\n",
+    "num_test_steps = num_test_samples // config[\"batch_size\"]\n",
+    "print(f\"Running test evaluation with {num_test_steps} steps ({num_test_samples} samples)\")\n",
+    "\n",
+    "# Set model to eval mode\n",
+    "model.eval()\n",
+    "\n",
+    "# Run test evaluation with progress bar\n",
+    "for test_batch in tqdm(test_loader, desc=\"Test evaluation\", total=num_test_steps):        \n",
+    "    _ = validation_step(        \n",
+    "        model, \n",
+    "        test_batch, \n",
+    "        test_metrics,\n",
+    "    )\n",
+    "    \n",
+    "# Compute final test metrics\n",
+    "test_metrics_dict = test_metrics.compute()\n",
+    "print(\"\\n\" + \"=\"*50)\n",
+    "print(\"Test Set Results\")\n",
+    "print(\"=\"*50)\n",
+    "print(f\"\\nMetrics:\")\n",
+    "print(f\"  Mean Pearson: {test_metrics_dict['mean/pearson']:.4f}\")\n",
+    "for track_name in config[\"bigwig_file_ids\"]:    \n",
+    "    print(f\"    {track_name}/pearson: {test_metrics_dict[f'{track_name}/pearson']:.4f}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    " ## Test set results\n",
+    "\n",
+    "Performances reached at ~1.5B tokens (~1500 steps in current 32kb sequences setup with batch_size=32)\n",
+    "\n",
+    "Mean Pearson: 0.5835\n",
+    "- ENCSR325NFE/pearson: 0.6081\n",
+    "- ENCSR962OTG/pearson: 0.7286\n",
+    "- ENCSR619DQO_P/pearson: 0.4976\n",
+    "- ENCSR619DQO_M/pearson: 0.4999"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3.12 (ntv3-env)",
+   "language": "python",
+   "name": "ntv3-env"
+  },
+  "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.12.3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

tabs/home.html

@@ -84,7 +84,7 @@
       <ul>
         <li><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_tutorials/00_quickstart_inference.ipynb" target="_blank" rel="noopener noreferrer">🚀 00 — Quickstart inference</a></li>
         <li><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_tutorials/01_tracks_prediction.ipynb" target="_blank" rel="noopener noreferrer">📊 01 — Tracks prediction</a></li>
-        <li>🎯 02 — Fine-tune on bigwig tracks</li>
+        <li><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_tutorials/02_fine_tuning.ipynb" target="_blank" rel="noopener noreferrer">🎯 02 — Fine-tune on bigwig tracks</a></li>
         <li>🔍 03 — Model interpretation</li>
         <li>🧪 04 — Training NTv3 generative </li>
       </ul>