feat: multiproc data download + cleaned cells

notebooks/03_fine_tuning.ipynb

@@ -4,13 +4,13 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# \ud83e\uddec Fine-Tuning a Model on BigWig Tracks Prediction\n",
+    "# 🧬 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",
-    "**\u26a1 Key Advantage**: This simplified pipeline achieves **close performance to more complex training approaches** while enabling **relatively fast fine-tuning in approximately one hour**. The setup is designed for rapid experimentation and iteration, making it ideal for adapting pre-trained models to your specific genomic tracks or experimental conditions without the overhead of complex distributed training infrastructure.\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",
-    "**\ud83d\udd27 Main Simplifications**: Compared to the full supervised tracks pipeline, this notebook simplifies several aspects to enable faster iteration:\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",
@@ -30,7 +30,14 @@
     "\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",
-    "\ud83d\udcdd Note for Google Colab users: This notebook is compatible with Colab! For faster training, make sure to enable GPU: Runtime \u2192 Change runtime type \u2192 GPU (T4 or better recommended).\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"
    ]
   },
   {
@@ -43,11 +50,39 @@
     "!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. \ud83d\udce6 Imports + Configuration\n",
+    "# 1. ⚙️ Configuration\n",
     "\n",
     "## Configuration Parameters\n",
     "\n",
@@ -81,37 +116,17 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 21,
    "metadata": {},
-   "outputs": [],
-   "source": [
-    "# 0. Imports\n",
-    "import random\n",
-    "import functools\n",
-    "from typing import List, Dict, Callable\n",
-    "import os\n",
-    "import subprocess\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": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Using device: cpu\n"
+     ]
+    }
+   ],
    "source": [
     "config = {\n",
     "    # Model\n",
@@ -121,7 +136,6 @@
     "    \"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/ENCFF884LDL/@@download/ENCFF884LDL.bigWig\",\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",
@@ -167,6 +181,8 @@
     "    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",
@@ -190,7 +206,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# 2. \ud83d\udce5 Genome & Tracks Data Download\n",
+    "# 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."
    ]
@@ -201,36 +217,56 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Download fasta file\n",
+    "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",
-    "print(f\"Downloading {fasta_filename}...\")\n",
-    "subprocess.run([\"wget\", \"-c\", config[\"fasta_url\"], \"-O\", fasta_gz_path], check=True)\n",
-    "\n",
-    "print(f\"Extracting {fasta_filename}...\")\n",
-    "subprocess.run([\"gunzip\", \"-f\", fasta_gz_path], check=True)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Download bigwig files\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",
-    "    print(f\"Downloading {filename}...\")\n",
-    "    subprocess.run([\"wget\", \"-c\", bigwig_url, \"-O\", filepath], check=True)"
+    "    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"
+    "### Data Splits Definition"
    ]
   },
   {
@@ -250,7 +286,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# 3. \ud83e\udde0 Model and tokenizer setup\n",
+    "# 3. 🧠 Model and tokenizer setup\n",
     " \n",
     "In this section, we set up the model and tokenizer. \n",
     " \n",
@@ -354,163 +390,11 @@
     "print(f\"Model parameters: {sum(p.numel() for p in model.parameters()):,}\")"
    ]
   },
-  {
-   "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": "markdown",
    "metadata": {},
    "source": [
-    "# 4. \ud83d\udd04 Data loading\n",
+    "# 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."
    ]
@@ -678,6 +562,165 @@
     "        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,
@@ -751,7 +794,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# 5. \u2699\ufe0f Optimizer setup\n",
+    "# 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"
@@ -785,7 +828,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# 6. \ud83d\udcca Metrics setup\n",
+    "# 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."
    ]
@@ -878,7 +921,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# 7. \ud83d\udcc9 Loss functions\n",
+    "# 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."
    ]
@@ -960,7 +1003,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# 8. \ud83c\udfc3 Training loop\n",
+    "# 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."
    ]
@@ -1192,7 +1235,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# 9. \ud83e\uddea Test evaluation\n",
+    "# 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."
    ]
@@ -1236,6 +1279,8 @@
    "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",
@@ -1270,4 +1315,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 2
-}
+}