refactor: improved comments and markdown

notebooks/03_fine_tuning.ipynb

@@ -8,29 +8,32 @@
     "\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",
+    "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",
     "**🔧 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",
+    "**⚡ 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 ~10min to reach acceptable performances for a 32kb functional tracks prediction task. 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",
+    "\n",
+    "**⚠️ Important Note on Hardware Requirements**: If the pipeline is designed to run on limited resources (e.g., Google Colab with a T4 GPU and 2CPUs), the timing mentioned was obtained on an **H100 GPU with 16 CPUs**. If you want to reach similar performance levels, you should be aware that you'll need **significant hardware resources** (high-end GPUs with substantial memory and multiple data loading workers). Training times will vary significantly based on your hardware configuration.\n",
+    "\n",
     "\n",
-    "- Loading genomic sequences from FASTA files and their corresponding BigWig signal tracks\n",
+    "The pipeline walks through the complete fine-tuning workflow:\n",
+    "- Loading genomic FASTA files sequences and their corresponding BigWig signal tracks from Hugging Face dataset\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",
+    "- Evaluation of the fine-tuned model on the test set\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",
+    "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 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"
+    "📝 Note for Google Colab users: This notebook is compatible with Colab and designed to work with limited resources! For faster training, make sure to enable GPU: Runtime → Change runtime type → GPU (T4 or better recommended). However, keep in mind that the timing benchmarks mentioned above were obtained on much more powerful hardware (H100 GPU), so your training times on Colab may be significantly longer."
    ]
   },
   {
@@ -56,28 +59,29 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import random\n",
+    "# Standard library imports\n",
     "import functools\n",
-    "from typing import List, Dict, Callable\n",
-    "import os\n",
     "import fnmatch\n",
+    "import os\n",
+    "import random\n",
     "from pathlib import Path\n",
-    "from huggingface_hub import HfApi, snapshot_download\n",
+    "from typing import Callable, Dict, List\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 pandas as pd\n",
+    "# Third-party imports\n",
+    "from huggingface_hub import HfApi, snapshot_download\n",
+    "from IPython.display import display\n",
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
+    "import pandas as pd\n",
     "import pyBigWig\n",
     "from pyfaidx import Fasta\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.nn.functional as F\n",
+    "from torch.optim import AdamW\n",
+    "from torch.utils.data import DataLoader, Dataset\n",
     "from torchmetrics import PearsonCorrCoef\n",
-    "import plotly.graph_objects as go\n",
-    "from IPython.display import display\n",
+    "from transformers import AutoConfig, AutoModelForMaskedLM, AutoTokenizer\n",
     "from tqdm import tqdm"
    ]
   },
@@ -131,29 +135,29 @@
     "    \n",
     "    # Data\n",
     "    \"hf_repo_id\": \"InstaDeepAI/NTv3_benchmark_dataset\",\n",
-    "    \"species\": \"arabidopsis\",\n",
+    "    \"species\": \"human\",\n",
     "    \"data_cache_dir\": \"./data\",\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",
+    "    \"batch_size\": 12,\n",
+    "    \"num_steps_training\": 5315,  # reproduce 10% of benchmark training length\n",
+    "    \"log_every_n_steps\": 20,\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",
+    "    \"validate_every_n_steps\": 100,\n",
+    "    \"num_validation_samples\": 1_000,\n",
     "\n",
     "    # Test\n",
-    "    \"num_test_samples\": 10000,\n",
+    "    \"num_test_samples\": 10_000,\n",
     "    \n",
     "    # General\n",
-    "    \"seed\": 0,\n",
+    "    \"seed\": 0,  # for reproducibility\n",
     "    \"device\": \"cuda\" if torch.cuda.is_available() else \"cpu\",\n",
-    "    \"num_workers\": 16,\n",
+    "    \"num_workers\": 2,  # If resources allows it, consider increasing for faster data loading\n",
     "}\n",
     "\n",
     "# Set random seed\n",
@@ -171,7 +175,15 @@
    "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."
+    "Download all needed data for fine-tuning from our Hugging Face NTv3 benchmark dataset for the selected species.\n",
+    "\n",
+    "These files contain:\n",
+    "- genomic sequences \n",
+    "- experimental signal data (e.g., ChIP-seq, ATAC-seq)\n",
+    "- ready-to-use splits (as bed files)\n",
+    "- metadata for data normalization\n",
+    "\n",
+    "If you want to fine-tune a model a species that is not available on our Hugging Face dataset or on other bigwig tracks, you should consider downloading the genome FASTA file and/or the BigWig files from URLs, using `wget`."
    ]
   },
   {
@@ -204,13 +216,7 @@
     "    # Find all files to download: species directory + metadata at root\n",
     "    species_pattern = f\"{species}/**\"\n",
     "    metadata_file = \"benchmark_metadata.tsv\"\n",
-    "    \n",
     "    species_files = [p for p in files if fnmatch.fnmatch(p, species_pattern)]\n",
-    "    if not species_files:\n",
-    "        raise ValueError(f\"No files found matching '{species_pattern}' in '{hf_repo_id}'\")\n",
-    "    \n",
-    "    if metadata_file not in files:\n",
-    "        raise ValueError(f\"No metadata file found at '{metadata_file}' in '{hf_repo_id}'\")\n",
     "    \n",
     "    # Download all needed files\n",
     "    download_patterns = [species_pattern, metadata_file]\n",
@@ -255,13 +261,8 @@
     "    \n",
     "    # Metadata file\n",
     "    metadata_path = local_dir / metadata_file\n",
-    "    if not metadata_path.is_file():\n",
-    "        raise ValueError(f\"Metadata file not found at '{metadata_path}'\")\n",
     "    metadata_df = pd.read_csv(metadata_path, sep=\"\\t\")\n",
     "\n",
-    "    if \"species\" not in metadata_df.columns:\n",
-    "        raise ValueError(\"benchmark_metadata.tsv has no 'species' column\")\n",
-    "\n",
     "    # Filter metadata according to species\n",
     "    metadata_df = metadata_df[metadata_df[\"species\"] == species].reset_index(drop=True)\n",
     "\n",
@@ -414,6 +415,13 @@
     "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": "markdown",
+   "metadata": {},
+   "source": [
+    "### Dataset Setup"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -573,7 +581,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "### Data preprocessing utilities"
+    "### Scaling function Setup"
    ]
   },
   {
@@ -621,6 +629,13 @@
     "    return transform_fn"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Instantiate datasets and dataloaders"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -736,19 +751,12 @@
     "    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.pearson_metric = PearsonCorrCoef(num_outputs=self.num_tracks).to(device)\n",
+    "        self.pearson_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.pearson_metric.reset()\n",
     "        self.losses = []\n",
     "    \n",
     "    def update(\n",
@@ -771,10 +779,7 @@
     "        # 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",
+    "        self.pearson_metric.update(pred_flat, target_flat)\n",
     "        \n",
     "        self.losses.append(loss)\n",
     "    \n",
@@ -782,13 +787,12 @@
     "        \"\"\"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",
+    "        # Compute Pearson correlation per track\n",
+    "        correlations = self.pearson_metric.compute().numpy()\n",
+    "        for i, track_name in enumerate(self.track_names):\n",
+    "            metrics_dict[f\"{track_name}/pearson\"] = correlations[i]\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",
@@ -887,7 +891,7 @@
     "    # Combine losses\n",
     "    loss = shape_loss + scale_loss / shape_loss_coefficient\n",
     "\n",
-    "    return loss, scale_loss, shape_loss\n"
+    "    return loss\n"
    ]
   },
   {
@@ -918,7 +922,7 @@
     "    bigwig_logits = outputs[\"bigwig_tracks_logits\"]\n",
     "    \n",
     "    # Compute loss\n",
-    "    loss, _, _ = poisson_multinomial_loss(\n",
+    "    loss = poisson_multinomial_loss(\n",
     "        logits=bigwig_logits,\n",
     "        targets=bigwig_targets,\n",
     "    )\n",
@@ -945,7 +949,7 @@
     "        bigwig_logits = outputs[\"bigwig_tracks_logits\"]\n",
     "        \n",
     "        # Compute loss\n",
-    "        loss, _, _ = poisson_multinomial_loss(\n",
+    "        loss = poisson_multinomial_loss(\n",
     "            logits=bigwig_logits,\n",
     "            targets=bigwig_targets,\n",
     "        )\n",
@@ -960,13 +964,6 @@
     "    return loss.item()"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### Interactive plotting is temporary for debug"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -1038,7 +1035,7 @@
     "            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",
+    "            f\"Pearson per track: {\", \".join([f\"{track_name}/pearson: {train_metrics_dict[f'{track_name}/pearson']:.4f}\" for track_name in bigwig_ids])}\"\n",
     "        )\n",
     "        train_metrics.reset()\n",
     "    \n",
@@ -1063,7 +1060,7 @@
     "        \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",
+    "        for track_name in bigwig_ids:\n",
     "            print(f\"    {track_name}/pearson: {val_metrics_dict[f'{track_name}/pearson']:.4f}\")\n",
     "        \n",
     "        model.train()  # Back to training mode\n",
@@ -1154,6 +1151,8 @@
     "\n",
     "Performances reached at ~1.5B tokens (~1500 steps in current 32kb sequences setup with batch_size=32)\n",
     "\n",
+    "**Hardware configuration**: These results were obtained on an **H100 GPU with 16 workers** for data loading in approximately **~10 minutes** of training.\n",
+    "\n",
     "Mean Pearson: 0.5835\n",
     "- ENCSR325NFE/pearson: 0.6081\n",
     "- ENCSR962OTG/pearson: 0.7286\n",