feat: update notebooks

notebooks_tutorials/{02_fine_tuning.ipynb → 02_fine_tuning_pretrained_model.ipynb}

@@ -4,9 +4,9 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# 🧬 Fine-Tuning a Model on BigWig Tracks Prediction\n",
+    "# 🧬 Fine-Tuning a Pre-Trained 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",
+    "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",
     "📊 We provide access to the NTv3-benchmark data that we released on our Hugging Face dataset: `InstaDeepAI/NTv3_benchmark_dataset`. In this repository, you will find ready-to-use genome FASTA files, Bigwig tracks, metadata, but also the splits that were used for the benchmark.\n",
     "\n",
@@ -15,7 +15,7 @@
     "- **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",
-    "**⚡ Key Advantage**: This simplified pipeline achieves close performance to more complex training approaches while enabling fast fine-tuning: on a H100 GPU and using 16 workers for data loading, it takes ~15min to reach acceptable performances for a 32kb functional tracks prediction task on **NTv3_8M_pre** model. The training speed benefits from the efficient NTv3 model architecture, but of course depends on your hardware capabilities (GPU acceleration and multi-worker data loading significantly reduce training time).\n"
+    "**⚡ Key Advantage**: This simplified pipeline achieves close performance to more complex training approaches while enabling fast fine-tuning: on a H100 GPU and using 16 workers for data loading, it takes ~15min to reach acceptable performances for a 32kb functional tracks prediction task on **NTv3_8M_pre** model. The training speed benefits from the efficient NTv3 model architecture, but of course depends on your hardware capabilities (GPU acceleration and multi-worker data loading significantly reduce training time)."
    ]
   },
   {

notebooks_tutorials/03_fine_tuning_posttrained_model.ipynb

@@ -0,0 +1,1141 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 🧬 Fine-Tuning a Post-trained Model on BigWig Tracks Prediction\n",
+    "\n",
+    "This notebook demonstrates a **simplified fine-tuning setup** that enables training of a **post-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",
+    "**🎯 Notebook purpose:**\n",
+    "This notebook is configured to train the `NTv3_650M_post` model on the `human` species from the NTv3 benchmark dataset. To run this training, you will need a large GPU (either A100 or H100).\n",
+    "For a simplified version of this notebook that uses the `NTv3_8M_pre` model and runs on a CPU, please see the [02_fine_tuning.ipynb](https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_tutorials/02_fine_tuning_pretrained_model.ipynb) notebook.\n",
+    "The notebook uses the same \"simplified setup\" as described there. \n",
+    "\n",
+    "📝 Note for Google Colab users: This notebook is compatible with Colab! This notebook is designed to be run on a high-performance GPU. The default parameters can be used with a H100 with 80GB of HBM."
+   ]
+  },
+  {
+   "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"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import random\n",
+    "import functools\n",
+    "from typing import List, Dict, Callable\n",
+    "import os\n",
+    "import fnmatch\n",
+    "from pathlib import Path\n",
+    "from huggingface_hub import HfApi, snapshot_download\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, AutoModel, AutoTokenizer\n",
+    "import pandas as pd\n",
+    "import matplotlib.pyplot as plt\n",
+    "import numpy as np\n",
+    "import pyBigWig\n",
+    "from pyfaidx import Fasta\n",
+    "from torchmetrics import PearsonCorrCoef\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",
+    "- **`hf_repo_id`**: HuggingFace dataset repository ID containing the benchmark data\n",
+    "- **`species`**: Species name (e.g., \"human\") to select data from the benchmark dataset\n",
+    "- **`data_cache_dir`**: Directory where downloaded data files (FASTA, bigWig) will be stored\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",
+    "### Test\n",
+    "- **`num_test_samples`**: Number of samples to use for test set evaluation\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": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "config = {\n",
+    "    # Model\n",
+    "    \"model_name\": \"InstaDeepAI/NTv3_650M_post\",\n",
+    "    \n",
+    "    # Data\n",
+    "    \"hf_repo_id\": \"InstaDeepAI/NTv3_benchmark_dataset\",\n",
+    "    \"species_name\": \"human\",\n",
+    "    \"data_cache_dir\": \"./data\",\n",
+    "    \"sequence_length\": 32_768,\n",
+    "    \"keep_target_center_fraction\": 0.375,\n",
+    "    \n",
+    "    # Training\n",
+    "    \"batch_size\": 4,\n",
+    "    \"num_steps_training\": 15000, #~2B tokens\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\": 0,\n",
+    "    \"device\": \"cuda\" if torch.cuda.is_available() else \"cpu\",\n",
+    "    \"num_workers\": 16,\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 prepare_genomics_inputs(\n",
+    "    species: str,\n",
+    "    data_cache_dir: str | Path = \"data\",\n",
+    "    hf_repo_id: str = \"InstaDeepAI/NTv3_benchmark_dataset\",\n",
+    "    bigwig_file_ids: list[str] | None = None,\n",
+    ") -> tuple[str, list[str], list[str], pd.DataFrame, pd.DataFrame]:\n",
+    "    \"\"\"\n",
+    "    Downloads:\n",
+    "      1) FASTA from HF dataset under: <species>/genome.fasta\n",
+    "      2) BigWigs from HF dataset under: <species>/functional_tracks/**\n",
+    "         (filtered by bigwig_file_ids if provided)\n",
+    "      3) Splits from HF dataset under: <species>/splits.bed\n",
+    "      4) Metadata from HF dataset under: benchmark_metadata.tsv\n",
+    "    \n",
+    "    Args:\n",
+    "        species: Species name (e.g., \"human\", \"arabidopsis\")\n",
+    "        data_cache_dir: Directory where downloaded data files will be stored\n",
+    "        hf_repo_id: HuggingFace dataset repository ID\n",
+    "        bigwig_file_ids: Optional list of BigWig file IDs to download. If None,\n",
+    "            downloads all available BigWig files for the species.\n",
+    "    \n",
+    "    Returns:\n",
+    "      (fasta_path, bigwig_path_list, bigwig_file_ids)\n",
+    "    \"\"\"\n",
+    "    cache = Path(data_cache_dir).expanduser().resolve()\n",
+    "    cache.mkdir(parents=True, exist_ok=True)\n",
+    "    \n",
+    "    # --- Download metadata + <species> files (FASTA, BigWigs, Splits) ---\n",
+    "    metadata_file = \"benchmark_metadata.tsv\"\n",
+    "    download_patterns = [metadata_file, f\"{species}/genome.fasta\", f\"{species}/splits.bed\"]\n",
+    "    \n",
+    "    if bigwig_file_ids is not None:\n",
+    "        # List files to validate requested BigWig files exist\n",
+    "        api = HfApi()\n",
+    "        files = api.list_repo_files(repo_id=hf_repo_id, repo_type=\"dataset\")\n",
+    "        species_pattern = f\"{species}/**\"\n",
+    "        species_files = [p for p in files if fnmatch.fnmatch(p, species_pattern)]\n",
+    "        \n",
+    "        # Get all available BigWig file IDs and their paths\n",
+    "        available_bigwig_files = {\n",
+    "            Path(p).stem: p for p in species_files \n",
+    "            if Path(p).suffix == \".bigwig\"\n",
+    "        }\n",
+    "        \n",
+    "        # Check that all requested files exist\n",
+    "        missing_files = set(bigwig_file_ids) - set(available_bigwig_files.keys())\n",
+    "        if missing_files:\n",
+    "            raise ValueError(\n",
+    "                f\"Requested BigWig files not found: {missing_files}. \"\n",
+    "                f\"Available files: {list(available_bigwig_files.keys())}\"\n",
+    "            )\n",
+    "        \n",
+    "        # Add specific patterns for requested BigWig files only\n",
+    "        for file_id in bigwig_file_ids:\n",
+    "            download_patterns.append(available_bigwig_files[file_id])\n",
+    "    else:\n",
+    "        # Download all BigWig files\n",
+    "        download_patterns.append(f\"{species}/functional_tracks/*.bigwig\")\n",
+    "    local_dir = Path(\n",
+    "        snapshot_download(\n",
+    "            repo_id=hf_repo_id,\n",
+    "            repo_type=\"dataset\",\n",
+    "            allow_patterns=download_patterns,\n",
+    "            local_dir=str(cache),\n",
+    "        )\n",
+    "    )\n",
+    "    \n",
+    "    # --- Organize outputs ---\n",
+    "    # FASTA file\n",
+    "    fasta_path_repo = f\"{species}/genome.fasta\"\n",
+    "    fasta_path = str(local_dir / fasta_path_repo)\n",
+    "    \n",
+    "    # BigWig files - use downloaded files directly\n",
+    "    bigwig_dir = local_dir / species / \"functional_tracks\"\n",
+    "    \n",
+    "    if bigwig_file_ids is not None:\n",
+    "        bigwig_paths = [str(bigwig_dir / f\"{file_id}.bigwig\") for file_id in bigwig_file_ids]\n",
+    "        bigwig_ids = bigwig_file_ids\n",
+    "    else:\n",
+    "        # Find all downloaded BigWig files\n",
+    "        bigwig_paths = [str(bigwig_file) for bigwig_file in bigwig_dir.glob(\"*.bigwig\")]\n",
+    "        bigwig_ids = [bigwig_file.stem for bigwig_file in bigwig_dir.glob(\"*.bigwig\")]         \n",
+    "    \n",
+    "    # Splits file\n",
+    "    splits_path_repo = f\"{species}/splits.bed\"\n",
+    "    splits_path = local_dir / splits_path_repo\n",
+    "\n",
+    "    splits_df = pd.read_csv(\n",
+    "        splits_path, \n",
+    "        sep=\"\\t\", \n",
+    "        header=None, \n",
+    "        names=[\"chr_name\", \"start\", \"end\", \"split\"],\n",
+    "        dtype={\"chr_name\": str, \"start\": int, \"end\": int, \"split\": str},\n",
+    "    )\n",
+    "    \n",
+    "    # Metadata file\n",
+    "    metadata_path = local_dir / metadata_file\n",
+    "    metadata_df = pd.read_csv(metadata_path, sep=\"\\t\")\n",
+    "\n",
+    "    # Filter metadata according to species\n",
+    "    metadata_df = metadata_df[metadata_df[\"species_name\"] == species].reset_index(drop=True)\n",
+    "\n",
+    "    # Order metadata according to bigwig file ids\n",
+    "    metadata_df = (\n",
+    "      metadata_df.set_index(\"file_id\")\n",
+    "        .loc[bigwig_ids]\n",
+    "        .reset_index()\n",
+    "    )\n",
+    "\n",
+    "    return fasta_path, bigwig_paths, bigwig_ids, splits_df, metadata_df"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "os.makedirs(config[\"data_cache_dir\"], exist_ok=True)\n",
+    "\n",
+    "# Download all species files + load the splits, and metadata\n",
+    "(\n",
+    "    fasta_path, \n",
+    "    bigwig_paths, \n",
+    "    bigwig_ids, \n",
+    "    species_splits_df,\n",
+    "    metadata_df \n",
+    ") = prepare_genomics_inputs(\n",
+    "    config[\"species_name\"], \n",
+    "    config[\"data_cache_dir\"], \n",
+    "    config[\"hf_repo_id\"]\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": [
+    "def crop_center(x: np.ndarray, keep_target_center_fraction: float = 0.375) -> np.ndarray:\n",
+    "    \"\"\"Crop the central sequence-length fraction for arrays of size (..., seq_len, num_tracks)\"\"\"\n",
+    "    seq_len = x.shape[-2]\n",
+    "    target_offset = int(seq_len * (1 - keep_target_center_fraction) // 2)\n",
+    "    target_length = seq_len - 2 * target_offset\n",
+    "    return x[..., target_offset:target_offset + target_length, :]\n",
+    "\n",
+    "\n",
+    "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",
+    "        species_str: 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 = AutoModel.from_pretrained(\n",
+    "            model_name, \n",
+    "            trust_remote_code=True,\n",
+    "            config=self.config,\n",
+    "        )\n",
+    "        if species_str in self.backbone.supported_species:\n",
+    "            self.species_ids = self.backbone.encode_species(species_str)\n",
+    "            print(f\"Using species: {species_str} with ids: {self.species_ids}\")\n",
+    "        else:\n",
+    "            # Mask token id\n",
+    "            print(f\"{species_str} not in supported species, using mask token id\")\n",
+    "            self.species_ids = torch.LongTensor([2])\n",
+    "\n",
+    "        self.keep_target_center_fraction = keep_target_center_fraction\n",
+    "\n",
+    "        # Bigwig head (NTv3 outputs at single-nucleotide resolution)\n",
+    "        self.bigwig_head = LinearHead(self.config.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",
+    "        species_tokens = torch.repeat_interleave(self.species_ids, tokens.shape[0])\n",
+    "        species_tokens = species_tokens.to(tokens.device)\n",
+    "        outputs = self.backbone(input_ids=tokens, species_ids=species_tokens, output_hidden_states=True)\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",
+    "            embedding = crop_center(embedding, self.keep_target_center_fraction)\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=bigwig_ids,\n",
+    "    species_str=config[\"full_species_name\"],\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(bigwig_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 file handles (one per worker process)\n",
+    "# This allows safe multi-worker DataLoader usage\n",
+    "_fasta_cache = {}  # Maps (process_id, file_path) -> Fasta handle\n",
+    "_bigwig_cache = {}  # Maps (process_id, file_path) -> pyBigWig handle\n",
+    "\n",
+    "\n",
+    "def _get_fasta_handle(fasta_path: str) -> Fasta:\n",
+    "    \"\"\"Get or create a FASTA file handle for the current process.\"\"\"\n",
+    "    process_id = os.getpid()\n",
+    "    abs_path = str(Path(fasta_path).resolve())\n",
+    "    cache_key = (process_id, abs_path)\n",
+    "    \n",
+    "    if cache_key not in _fasta_cache:\n",
+    "        _fasta_cache[cache_key] = Fasta(abs_path, as_raw=True, sequence_always_upper=True)\n",
+    "    \n",
+    "    return _fasta_cache[cache_key]\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",
+    "    abs_path = str(Path(bigwig_path).resolve())\n",
+    "    cache_key = (process_id, abs_path)\n",
+    "    \n",
+    "    if cache_key not in _bigwig_cache:\n",
+    "        # Check if file exists before trying to open\n",
+    "        if not Path(abs_path).exists():\n",
+    "            raise FileNotFoundError(\n",
+    "                f\"BigWig file not found: {abs_path}\\n\"\n",
+    "                f\"Original path: {bigwig_path}\\n\"\n",
+    "                f\"Current working directory: {os.getcwd()}\"\n",
+    "            )\n",
+    "        \n",
+    "        try:\n",
+    "            _bigwig_cache[cache_key] = pyBigWig.open(abs_path)\n",
+    "        except Exception as e:\n",
+    "            raise RuntimeError(\n",
+    "                f\"Failed to open BigWig file: {abs_path} with error: {str(e)}\\n\"\n",
+    "                f\"File exists: {Path(abs_path).exists()}\\n\"\n",
+    "                f\"File size: {Path(abs_path).stat().st_size if Path(abs_path).exists() else 'N/A'} bytes\"\n",
+    "            ) from e\n",
+    "    \n",
+    "    return _bigwig_cache[cache_key]\n",
+    "\n",
+    "\n",
+    "class GenomeBigWigDataset(Dataset):\n",
+    "    \"\"\"\n",
+    "    A PyTorch dataset to access a reference genome and bigwig tracks. The dataset is \n",
+    "    compatible with multi-worker DataLoaders (using process-local file handles and lazy \n",
+    "    loading). For each sample, a random genomic region is picked from the specified split,\n",
+    "    and a random window of length `sequence_length` within that region is returned.\n",
+    "    \"\"\"\n",
+    "\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        fasta_path: str,\n",
+    "        bigwig_path_list: list[str],\n",
+    "        chrom_regions: pd.DataFrame,\n",
+    "        split: 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",
+    "        # Store paths instead of opening files immediately (for multi-worker compatibility)\n",
+    "        self.fasta_path = fasta_path\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\n",
+    "        self.keep_target_center_fraction = keep_target_center_fraction\n",
+    "        self.chrom_regions = chrom_regions\n",
+    "\n",
+    "        # Filter regions by split\n",
+    "        split_regions = self.chrom_regions[self.chrom_regions[\"split\"] == split].copy()\n",
+    "\n",
+    "        # Filter valid regions (must be large enough for sequence_length)\n",
+    "        self.valid_regions = []\n",
+    "        for _, row in split_regions.iterrows():\n",
+    "\n",
+    "            region_length = row.end - row.start\n",
+    "            if region_length < self.sequence_length:\n",
+    "                continue\n",
+    "            \n",
+    "            # Store valid region\n",
+    "            self.valid_regions.append((row.chr_name, row.start, row.end))\n",
+    "\n",
+    "    def __len__(self):\n",
+    "        return self.num_samples\n",
+    "\n",
+    "    def __getitem__(self, idx):\n",
+    "        # Sample a random region from the valid regions\n",
+    "        chrom, region_start, region_end = random.choice(self.valid_regions)\n",
+    "        \n",
+    "        # Sample a random window within this region\n",
+    "        max_start = region_end - self.sequence_length\n",
+    "        start = random.randint(region_start, max_start)\n",
+    "        end = start + self.sequence_length\n",
+    "\n",
+    "        # Sequence - get FASTA handle lazily (cached per worker process)\n",
+    "        fasta = _get_fasta_handle(self.fasta_path)\n",
+    "        seq = 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",
+    "            bigwig_targets = crop_center(bigwig_targets, self.keep_target_center_fraction)\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": [
+    "def create_targets_scaling_fn(\n",
+    "    metadata_df: pd.DataFrame\n",
+    ") -> Callable[[torch.Tensor], torch.Tensor]:\n",
+    "    \"\"\"\n",
+    "    Build a scaling function that uses the track means to normalise and softclip the targets.\n",
+    "    \"\"\"\n",
+    "    # Open bigwig files and compute track statistics\n",
+    "    track_means = metadata_df[\"mean\"].to_numpy()\n",
+    "    print(f\"Track means: {track_means}\")\n",
+    "    print(f\"Number of tracks: {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",
+    "        # 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 datasets & dataloaders\n",
+    "create_dataset_fn = functools.partial(\n",
+    "    GenomeBigWigDataset,\n",
+    "    fasta_path=fasta_path,\n",
+    "    bigwig_path_list=bigwig_paths,\n",
+    "    chrom_regions=species_splits_df,\n",
+    "    sequence_length=config[\"sequence_length\"],\n",
+    "    tokenizer=tokenizer,\n",
+    "    transform_fn=create_targets_scaling_fn(metadata_df),\n",
+    "    keep_target_center_fraction=config[\"keep_target_center_fraction\"],\n",
+    ")\n",
+    "\n",
+    "train_dataset = create_dataset_fn(\n",
+    "    split=\"train\",\n",
+    "    num_samples=config[\"num_steps_training\"] * config[\"batch_size\"],\n",
+    ")\n",
+    "\n",
+    "val_dataset = create_dataset_fn(\n",
+    "    split=\"val\",\n",
+    "    num_samples=config[\"num_validation_samples\"],\n",
+    ")\n",
+    "\n",
+    "test_dataset = create_dataset_fn(\n",
+    "    split=\"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\"\\nTrain 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",
+    "    \"\"\"Metrics to handle multi-track pearson correlations and losses\"\"\"\n",
+    "    \n",
+    "    def __init__(self, track_names: List[str], split: str):\n",
+    "        self.track_names = track_names\n",
+    "        self.num_tracks = len(track_names)\n",
+    "        self.split = split\n",
+    "\n",
+    "        # Initialise metrics \n",
+    "        self.pearson = PearsonCorrCoef(num_outputs=self.num_tracks).to(device)\n",
+    "        self.pearson.set_dtype(torch.float64) # Use float64 for improved numerical stability\n",
+    "        self.losses = []\n",
+    "\n",
+    "        # Record mean metrics per logging interval\n",
+    "        self.step_idxs = []\n",
+    "        self.mean_pearsons = []\n",
+    "        self.mean_losses = []\n",
+    "    \n",
+    "    def reset(self):\n",
+    "        self.pearson.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 the metrics with predictions and targets of shape (..., num_tracks) and a scalar loss.\n",
+    "        \"\"\"\n",
+    "        # Flatten batch and sequence dimensions\n",
+    "        pred_flat = predictions.detach().reshape(-1, self.num_tracks).to(torch.float64)  # (N, num_tracks)\n",
+    "        target_flat = targets.detach().reshape(-1, self.num_tracks).to(torch.float64)  # (N, num_tracks)\n",
+    "        \n",
+    "        # Update metrics\n",
+    "        self.pearson.update(pred_flat, target_flat)\n",
+    "        self.losses.append(loss)\n",
+    "    \n",
+    "    def compute(self) -> Dict[str, float]:\n",
+    "        \"\"\"Compute the pearson correlations and loss and return a dictionary of metrics.\"\"\"\n",
+    "        # Per-track Pearson correlations\n",
+    "        correlations = self.pearson.compute().cpu().numpy()\n",
+    "        metrics_dict = {\n",
+    "            f\"{track_name}/pearson\": correlations[i] for i, track_name in enumerate(self.track_names)\n",
+    "        }\n",
+    "        metrics_dict[\"mean/pearson\"] = correlations.mean()\n",
+    "        \n",
+    "        # Mean loss\n",
+    "        metrics_dict[\"loss\"] = np.mean(self.losses)\n",
+    "        \n",
+    "        return metrics_dict\n",
+    "\n",
+    "    def update_mean_metrics(self, step_idx: int):\n",
+    "        \"\"\"Update the mean metrics over the logging interval and save to a csv file.\"\"\"\n",
+    "        # Update mean metrics with the mean pearson & average loss\n",
+    "        metrics_dict = self.compute()\n",
+    "        self.step_idxs.append(step_idx)\n",
+    "        self.mean_pearsons.append(metrics_dict[\"mean/pearson\"])\n",
+    "        self.mean_losses.append(metrics_dict[\"loss\"])\n",
+    "\n",
+    "        # Save metrics to a csv for plotting\n",
+    "        data = {\n",
+    "            \"step\": self.step_idxs,\n",
+    "            \"mean_loss\": self.mean_losses,\n",
+    "            \"mean_pearson\": self.mean_pearsons,\n",
+    "        }\n",
+    "        df = pd.DataFrame(data)\n",
+    "        df.to_csv(f\"metrics_{self.split}.csv\", index=False)\n",
+    "        \n",
+    "    def print_metrics(self, print_per_track: bool = False):\n",
+    "        \"\"\"Print a summary of the metrics.\"\"\"\n",
+    "        print(\n",
+    "            f\"Step {self.step_idxs[-1]}/{config['num_steps_training']} | \"\n",
+    "            f\"Loss: {self.mean_losses[-1]:.4f} | \"\n",
+    "            f\"Mean Pearson: {self.mean_pearsons[-1]:.4f}\"\n",
+    "        )\n",
+    "        metrics_dict = self.compute()\n",
+    "        if print_per_track:\n",
+    "            for metric_key, metric_value in metrics_dict.items():\n",
+    "                print(f\"    {metric_key}: {metric_value:.4f}\")\n",
+    "    "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_metrics = TracksMetrics(bigwig_ids, \"train\")\n",
+    "val_metrics = TracksMetrics(bigwig_ids, \"val\")\n",
+    "test_metrics = TracksMetrics(bigwig_ids, \"test\")"
+   ]
+  },
+  {
+   "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",
+    ") -> torch.Tensor: \n",
+    "    \"\"\"\n",
+    "    Regression loss for bigwig tracks (Poisson-Multinomial). The logits and targets are\n",
+    "    expected to be of shape (batch, seq_length, num_tracks).\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\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",
+    "    optimizer: torch.optim.Optimizer,\n",
+    "    batch: Dict[str, torch.Tensor],\n",
+    "    train_metrics: TracksMetrics,\n",
+    ") -> None:\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",
+    "    optimizer.zero_grad()\n",
+    "    loss.backward()\n",
+    "    optimizer.step()\n",
+    "\n",
+    "    # Update metrics\n",
+    "    train_metrics.update(\n",
+    "        predictions=bigwig_logits,\n",
+    "        targets=bigwig_targets,\n",
+    "        loss=loss.item()\n",
+    "    )\n",
+    "    \n",
+    "\n",
+    "\n",
+    "def validation_step(\n",
+    "    model: nn.Module,\n",
+    "    batch: Dict[str, torch.Tensor],\n",
+    "    metrics: TracksMetrics,\n",
+    ") -> None:\n",
+    "    \"\"\"Single validation step.\"\"\"\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",
+    "        )"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Run Training Loop"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Training loop\n",
+    "print(f\"Starting training for {config['num_steps_training']} steps\\n\")\n",
+    "\n",
+    "# Create iterator for training data (will cycle if needed)\n",
+    "train_iter = iter(train_loader)\n",
+    "model.train()\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",
+    "    # Take a training step\n",
+    "    train_step(model, optimizer, batch, train_metrics)\n",
+    "\n",
+    "    # Logging\n",
+    "    if (step_idx + 1) % config[\"log_every_n_steps\"] == 0:\n",
+    "        train_metrics.update_mean_metrics(step_idx + 1)\n",
+    "        train_metrics.print_metrics()\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",
+    "        model.eval()\n",
+    "        \n",
+    "        for val_batch in val_loader:\n",
+    "            validation_step(model, val_batch, val_metrics)\n",
+    "        \n",
+    "        val_metrics.update_mean_metrics(step_idx + 1)\n",
+    "        val_metrics.print_metrics(print_per_track=True)\n",
+    "        val_metrics.reset()\n",
+    "\n",
+    "        # Back to training mode\n",
+    "        print(\"\\n\" + \"-\"*100 + \"\\nTraining metrics:\")\n",
+    "        model.train()  \n",
+    "\n",
+    "print(f\"\\nTraining completed after {config['num_steps_training']} steps.\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Plot training results\n",
+    "fig, axes = plt.subplots(1, 2, figsize=(12, 5))\n",
+    "\n",
+    "df_train = pd.read_csv(\"metrics_train.csv\")\n",
+    "df_val = pd.read_csv(\"metrics_val.csv\")\n",
+    "\n",
+    "# Plot Loss\n",
+    "axes[0].plot(df_train[\"step\"], df_train[\"mean_loss\"], 'b-o', label='Train Loss', markersize=4, linewidth=1.5)\n",
+    "axes[0].plot(df_val[\"step\"], df_val[\"mean_loss\"], 'r-s', label='Val Loss', markersize=4, linewidth=1.5)\n",
+    "axes[0].set_xlabel('Step')\n",
+    "axes[0].set_ylabel('Loss')\n",
+    "axes[0].set_title('Loss')\n",
+    "axes[0].legend()\n",
+    "axes[0].grid(True, alpha=0.3)\n",
+    "\n",
+    "# Plot Pearson Correlation\n",
+    "axes[1].plot(df_train[\"step\"], df_train[\"mean_pearson\"], 'g-o', label='Train Pearson', markersize=4, linewidth=1.5)\n",
+    "axes[1].plot(df_val[\"step\"], df_val[\"mean_pearson\"], 'orange', marker='s', label='Val Pearson', markersize=4, linewidth=1.5)\n",
+    "axes[1].set_xlabel('Step')\n",
+    "axes[1].set_ylabel('Pearson Correlation')\n",
+    "axes[1].set_title('Mean Pearson Correlation')\n",
+    "axes[1].legend()\n",
+    "axes[1].grid(True, alpha=0.3)"
+   ]
+  },
+  {
+   "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 bigwig_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",
+    "#TODO: Add test set results after run!"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "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.14"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

tabs/home.html

@@ -84,10 +84,11 @@
       <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><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><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_tutorials/03_model_interpretation.ipynb" target="_blank" rel="noopener noreferrer">🔍 03 — Model interpretation</a></li>
-        <li>🧪 04 — Training NTv3-generative <em>(coming soon)</em></li>
-        <li>🪰 05 — Generating enhancer sequences <em>(coming soon)</em></li>
+        <li><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_tutorials/02_fine_tuning_pretrained_model.ipynb" target="_blank" rel="noopener noreferrer">🎯 02 — Fine-tune a pre-trained model on bigwig tracks</a></li>
+        <li><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_tutorials/03_fine_tuning_posttrained_model.ipynb" target="_blank" rel="noopener noreferrer">🎯 03 — Fine-tune a post-trained model on bigwig tracks</a></li>
+        <li><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_tutorials/04_model_interpretation.ipynb" target="_blank" rel="noopener noreferrer">🔍 04 — Model interpretation</a></li>
+        <li>🧪 05 — Training NTv3-generative <em>(coming soon)</em></li>
+        <li>🪰 06 — Generating enhancer sequences <em>(coming soon)</em></li>
       </ul>
     </div>
     <div class="card">