fix: notebook simplification

notebooks/03_fine_tuning.ipynb

@@ -16,7 +16,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -28,14 +28,14 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 29,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
     "# 0. Imports\n",
     "import random\n",
     "import functools\n",
-    "from typing import List, Dict, Optional, Callable\n",
+    "from typing import List, Dict, Callable\n",
     "import os\n",
     "import subprocess\n",
     "\n",
@@ -48,19 +48,50 @@
     "import numpy as np\n",
     "import pyBigWig\n",
     "from pyfaidx import Fasta\n",
-    "from torchmetrics import PearsonCorrCoef"
+    "from torchmetrics import PearsonCorrCoef\n",
+    "import plotly.graph_objects as go\n",
+    "from plotly.subplots import make_subplots\n",
+    "from IPython.display import display"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# 1. Configuration setup"
+    "# 1. Configuration setup\n",
+    "\n",
+    "## Configuration Parameters\n",
+    "\n",
+    "### Model\n",
+    "- **`model_name`**: HuggingFace model name/identifier for the pretrained backbone model\n",
+    "\n",
+    "### Data\n",
+    "- **`data_cache_dir`**: Directory where downloaded data files (FASTA, bigWig) will be stored\n",
+    "- **`fasta_url`**: URL to download reference genome FASTA file\n",
+    "- **`bigwig_url_list`**: List of URLs for bigWig track files to download\n",
+    "- **`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",
+    "- **`validate_every_n_steps`**: Run validation every N steps\n",
+    "\n",
+    "### Validation\n",
+    "- **`num_validation_samples`**: Number of samples to use for validation set\n",
+    "\n",
+    "### General\n",
+    "- **`seed`**: Random seed for reproducibility\n",
+    "- **`device`**: Device to run training on (\"cuda\" or \"cpu\")\n",
+    "- **`num_workers`**: Number of worker processes for DataLoader (0 = single-threaded)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 30,
+   "execution_count": 15,
    "metadata": {},
    "outputs": [
     {
@@ -74,37 +105,32 @@
    "source": [
     "config = {\n",
     "    # Model\n",
-    "    \"model_name\": \"InstaDeepAI/ntv3_8M_7downsample_pretrained_le_1mb\",  # HuggingFace model name/identifier\n",
+    "    \"model_name\": \"InstaDeepAI/ntv3_8M_7downsample_pretrained_le_1mb\",\n",
     "    \n",
     "    # Data\n",
-    "    \"data_cache_dir\": \"./data\",  # Directory where downloaded data files (FASTA, bigWig) will be stored\n",
-    "    \"fasta_url\": \"https://hgdownload.gi.ucsc.edu/goldenPath/hg38/bigZips/hg38.fa.gz\",  # URL to download reference genome FASTA file\n",
-    "    \"bigwig_url_list\": [\"https://www.encodeproject.org/files/ENCFF884LDL/@@download/ENCFF884LDL.bigWig\"],  # List of URLs for bigWig track files to download\n",
-    "    \"sequence_length\": 1_024,  # Length of input sequences in base pairs (bp)\n",
-    "    \"keep_target_center_fraction\": 0.375,  # Fraction of center sequence to keep for target prediction (crops edges to focus on center)\n",
+    "    \"data_cache_dir\": \"./data\",\n",
+    "    \"fasta_url\": \"https://hgdownload.gi.ucsc.edu/goldenPath/hg38/bigZips/hg38.fa.gz\",\n",
+    "    \"bigwig_url_list\": [\n",
+    "        \"https://www.encodeproject.org/files/ENCFF884LDL/@@download/ENCFF884LDL.bigWig\"\n",
+    "    ],\n",
+    "    \"sequence_length\": 1_024,\n",
+    "    \"keep_target_center_fraction\": 0.375,\n",
     "    \n",
     "    # Training\n",
-    "    \"batch_size\": 2,  # Number of samples per batch\n",
-    "    \"learning_rate\": 1e-5,  # Constant learning rate for optimizer\n",
-    "    \"weight_decay\": 0.01,  # L2 regularization coefficient for optimizer\n",
-    "    \n",
-    "    \"num_tokens_training\": 131_072,  # Total training tokens budget (determines total training steps)\n",
-    "    \"num_tokens_per_update\": 4_096,  # Target tokens per optimizer update (batch_size * seq_len * grad_accum)\n",
-    "    \"num_tokens_per_log\": 8_192,  # Tokens between training logs (how often to print metrics)\n",
-    "    \"num_tokens_per_validation\": 16_384,  # Tokens between validation runs (how often to evaluate on validation set)\n",
+    "    \"batch_size\": 8,\n",
+    "    \"num_steps_training\": 1000,\n",
+    "    \"log_every_n_steps\": 10,\n",
+    "    \"learning_rate\": 1e-5,\n",
+    "    \"weight_decay\": 0.01,\n",
     "    \n",
     "    # Validation\n",
-    "    \"num_validation_samples\": 10,  # Number of samples to use for validation set\n",
-    "    \n",
-    "    # Loss\n",
-    "    \"bigwig_loss_weight\": 1.0,  # Weight multiplier for bigwig prediction loss\n",
-    "    \"bigwig_scalar_loss_function\": \"poisson-multinomial\",  # Loss function type for bigwig tracks\n",
-    "    \"bigwig_shape_loss_coefficient\": 5.0,  # Coefficient balancing shape loss vs scale loss in poisson-multinomial loss\n",
+    "    \"validate_every_n_steps\": 50,\n",
+    "    \"num_validation_samples\": 100,\n",
     "    \n",
     "    # General\n",
-    "    \"seed\": 42,  # Random seed for reproducibility\n",
-    "    \"device\": \"cuda\" if torch.cuda.is_available() else \"cpu\",  # Device to run training on (\"cuda\" or \"cpu\")\n",
-    "    \"num_workers\": 0,  # Number of worker processes for DataLoader (0 = single-threaded)\n",
+    "    \"seed\": 42,\n",
+    "    \"device\": \"cuda\" if torch.cuda.is_available() else \"cpu\",\n",
+    "    \"num_workers\": 0,\n",
     "}\n",
     "\n",
     "os.makedirs(config[\"data_cache_dir\"], exist_ok=True)\n",
@@ -216,7 +242,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 31,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -236,7 +262,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 32,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -304,7 +330,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 33,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [
     {
@@ -335,6 +361,49 @@
     "print(f\"Model parameters: {sum(p.numel() for p in model.parameters()):,}\")"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Scaling functions for targets\n",
+    "def get_track_means(bigwig_file_ids: List[str]) -> np.ndarray:\n",
+    "    \"\"\"\n",
+    "    Get track means for normalization.\n",
+    "    For now, return dummy values. In real pipeline, this loads from metadata.\n",
+    "    \"\"\"\n",
+    "    # Dummy values - in real pipeline, this would load from actual metadata\n",
+    "    return np.ones(len(bigwig_file_ids), dtype=np.float32) * 1.0\n",
+    "\n",
+    "\n",
+    "def create_targets_scaling_fn(bigwig_file_ids: List[str]) -> Callable[[torch.Tensor], torch.Tensor]:\n",
+    "    \"\"\"\n",
+    "    Build a scaling function based on track means.\n",
+    "    \"\"\"\n",
+    "    # Load track means\n",
+    "    track_means_np = get_track_means(bigwig_file_ids)\n",
+    "    track_means = torch.tensor(track_means_np, 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.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": {},
@@ -344,7 +413,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 34,
+   "execution_count": 7,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -388,6 +457,7 @@
     "        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",
     "        num_tracks: int = 1,\n",
     "    ):\n",
@@ -401,6 +471,7 @@
     "        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.num_tracks = num_tracks\n",
     "        self.chroms = chroms\n",
@@ -465,6 +536,9 @@
     "            target_length = seq_len - 2 * target_offset\n",
     "            bigwig_targets = bigwig_targets[target_offset:target_offset + target_length, :]\n",
     "\n",
+    "        # Apply scaling to targets\n",
+    "        bigwig_targets = self.transform_fn(bigwig_targets)\n",
+    "\n",
     "        sample = {\n",
     "            \"tokens\": tokens,\n",
     "            \"bigwig_targets\": bigwig_targets,\n",
@@ -477,7 +551,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 35,
+   "execution_count": 16,
    "metadata": {},
    "outputs": [
     {
@@ -485,18 +559,22 @@
      "output_type": "stream",
      "text": [
       "Train samples: 100\n",
-      "Val samples: 10\n",
-      "Test samples: 10\n"
+      "Val samples: 100\n",
+      "Test samples: 100\n"
      ]
     }
    ],
    "source": [
+    "# Create scaling function\n",
+    "transform_fn = create_targets_scaling_fn(config[\"bigwig_file_ids\"])\n",
+    "\n",
     "create_dataset_fn = functools.partial(\n",
     "    GenomeBigWigDataset,\n",
     "    fasta_path=fasta_path,\n",
     "    bigwig_path_list=bigwig_path_list,\n",
     "    sequence_length=config[\"sequence_length\"],\n",
     "    tokenizer=tokenizer,\n",
+    "    transform_fn=transform_fn,\n",
     "    keep_target_center_fraction=config[\"keep_target_center_fraction\"],\n",
     "    num_tracks=len(config[\"bigwig_file_ids\"]),\n",
     ")\n",
@@ -552,21 +630,18 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 36,
+   "execution_count": 17,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Gradient accumulation steps: 2\n",
-      "Effective batch size: 4\n",
-      "Effective tokens per update: 4096\n",
-      "\n",
-      "Training constants:\n",
-      "  Total training steps: 32\n",
-      "  Log training metrics every: 2 steps\n",
-      "  Run validation every: 4 steps\n",
+      "Training configuration:\n",
+      "  Batch size: 8\n",
+      "  Total training steps: 1000\n",
+      "  Log metrics every: 10 steps\n",
+      "  Validate every: 50 steps\n",
       "\n",
       "Optimizer setup:\n",
       "  Learning rate: 1e-05\n"
@@ -574,37 +649,12 @@
     }
    ],
    "source": [
-    "# Calculate gradient accumulation steps and effective batch size\n",
-    "num_devices = 1  # Single device for now\n",
-    "sequence_length = config[\"sequence_length\"]\n",
-    "batch_size = config[\"batch_size\"]\n",
-    "\n",
-    "# Calculate gradient accumulation steps\n",
-    "num_accumulation_gradient = max(1, int(config[\"num_tokens_per_update\"] // (batch_size * num_devices * sequence_length)))\n",
-    "\n",
-    "# Calculate effective batch size and tokens per update\n",
-    "effective_batch_size = batch_size * num_devices * num_accumulation_gradient\n",
-    "effective_num_tokens_per_update = effective_batch_size * sequence_length\n",
-    "\n",
-    "print(f\"Gradient accumulation steps: {num_accumulation_gradient}\")\n",
-    "print(f\"Effective batch size: {effective_batch_size}\")\n",
-    "print(f\"Effective tokens per update: {effective_num_tokens_per_update}\")\n",
-    "\n",
-    "# Compute logging constants (based on deepspeed pipeline: compute_logging_constants)\n",
-    "num_train_samples = len(train_dataset)\n",
-    "num_tokens_per_update = effective_num_tokens_per_update  # Same as effective_num_tokens_per_update\n",
-    "\n",
-    "# Total training steps based on token budget\n",
-    "num_steps_training = config[\"num_tokens_training\"] // num_tokens_per_update\n",
-    "\n",
-    "# Steps for logging and validation\n",
-    "log_train_step = int(np.ceil(config[\"num_tokens_per_log\"] / num_tokens_per_update))\n",
-    "log_validation_step = int(np.ceil(config[\"num_tokens_per_validation\"] / num_tokens_per_update))\n",
-    "\n",
-    "print(f\"\\nTraining constants:\")\n",
-    "print(f\"  Total training steps: {num_steps_training}\")\n",
-    "print(f\"  Log training metrics every: {log_train_step} steps\")\n",
-    "print(f\"  Run validation every: {log_validation_step} steps\")\n",
+    "# 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",
@@ -626,87 +676,62 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 37,
+   "execution_count": 18,
    "metadata": {},
    "outputs": [],
    "source": [
     "class TracksMetrics:\n",
-    "    \"\"\"Simple metrics tracker for tracks prediction with both scaled and raw metrics.\"\"\"\n",
+    "    \"\"\"Simple metrics tracker for tracks prediction.\"\"\"\n",
     "    \n",
     "    def __init__(self, track_names: List[str]):\n",
     "        self.track_names = track_names\n",
     "        self.num_tracks = len(track_names)\n",
-    "        # Scaled metrics: comparing scaled targets with scaled predictions\n",
-    "        self.pearson_metrics_scaled = [\n",
-    "            PearsonCorrCoef().to(device) for _ in range(self.num_tracks)\n",
-    "        ]\n",
-    "        # Raw metrics: comparing raw targets with unscaled predictions\n",
-    "        self.pearson_metrics_raw = [\n",
+    "        # Metrics: comparing scaled targets with scaled predictions\n",
+    "        self.pearson_metrics = [\n",
     "            PearsonCorrCoef().to(device) for _ in range(self.num_tracks)\n",
     "        ]\n",
     "        self.losses = []\n",
     "    \n",
     "    def reset(self):\n",
-    "        for metric in self.pearson_metrics_scaled:\n",
-    "            metric.reset()\n",
-    "        for metric in self.pearson_metrics_raw:\n",
+    "        for metric in self.pearson_metrics:\n",
     "            metric.reset()\n",
     "        self.losses = []\n",
     "    \n",
     "    def update(\n",
     "        self, \n",
-    "        predictions_scaled: torch.Tensor, \n",
-    "        targets_scaled: torch.Tensor,\n",
-    "        predictions_raw: torch.Tensor,\n",
-    "        targets_raw: torch.Tensor,\n",
+    "        predictions: torch.Tensor, \n",
+    "        targets: torch.Tensor,\n",
     "        loss: float\n",
     "    ):\n",
     "        \"\"\"\n",
-    "        Update both scaled and raw metrics.\n",
+    "        Update metrics.\n",
     "        Args:\n",
-    "            predictions_scaled: (batch, seq_len, num_tracks) - scaled predictions\n",
-    "            targets_scaled: (batch, seq_len, num_tracks) - scaled targets\n",
-    "            predictions_raw: (batch, seq_len, num_tracks) - raw/unscaled predictions\n",
-    "            targets_raw: (batch, seq_len, num_tracks) - raw targets\n",
+    "            predictions: (batch, seq_len, num_tracks)\n",
+    "            targets: (batch, seq_len, num_tracks)\n",
     "            loss: scalar loss value\n",
     "        \"\"\"\n",
     "        # Flatten batch and sequence dimensions\n",
-    "        pred_scaled_flat = predictions_scaled.detach().reshape(-1, self.num_tracks)  # (N, num_tracks)\n",
-    "        target_scaled_flat = targets_scaled.detach().reshape(-1, self.num_tracks)  # (N, num_tracks)\n",
-    "        pred_raw_flat = predictions_raw.detach().reshape(-1, self.num_tracks)  # (N, num_tracks)\n",
-    "        target_raw_flat = targets_raw.detach().reshape(-1, self.num_tracks)  # (N, num_tracks)\n",
-    "        \n",
-    "        # Update scaled metrics\n",
-    "        for i, metric in enumerate(self.pearson_metrics_scaled):\n",
-    "            metric.update(pred_scaled_flat[:, i], target_scaled_flat[:, i])\n",
+    "        pred_flat = predictions.detach().reshape(-1, self.num_tracks)  # (N, num_tracks)\n",
+    "        target_flat = targets.detach().reshape(-1, self.num_tracks)  # (N, num_tracks)\n",
     "        \n",
-    "        # Update raw metrics\n",
-    "        for i, metric in enumerate(self.pearson_metrics_raw):\n",
-    "            metric.update(pred_raw_flat[:, i], target_raw_flat[:, i])\n",
+    "        # Update metrics\n",
+    "        for i, metric in enumerate(self.pearson_metrics):\n",
+    "            metric.update(pred_flat[:, i], target_flat[:, i])\n",
     "        \n",
     "        self.losses.append(loss)\n",
     "    \n",
     "    def compute(self) -> Dict[str, float]:\n",
-    "        \"\"\"Compute and return all metrics (both scaled and raw).\"\"\"\n",
+    "        \"\"\"Compute and return all metrics.\"\"\"\n",
     "        metrics_dict = {}\n",
     "        \n",
-    "        # Scaled metrics: per-track Pearson correlations\n",
-    "        for i, (track_name, metric) in enumerate(zip(self.track_names, self.pearson_metrics_scaled)):\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\"metrics_scaled/{track_name}/pearson\"] = corr\n",
+    "            metrics_dict[f\"{track_name}/pearson\"] = corr\n",
     "        \n",
-    "        # Scaled metrics: mean Pearson correlation\n",
-    "        correlations_scaled = [metric.compute().item() for metric in self.pearson_metrics_scaled]\n",
-    "        metrics_dict[\"metrics_scaled/mean/pearson\"] = np.nanmean(correlations_scaled)\n",
-    "        \n",
-    "        # Raw metrics: per-track Pearson correlations\n",
-    "        for i, (track_name, metric) in enumerate(zip(self.track_names, self.pearson_metrics_raw)):\n",
-    "            corr = metric.compute().item()\n",
-    "            metrics_dict[f\"metrics_raw/{track_name}/pearson\"] = corr\n",
-    "        \n",
-    "        # Raw metrics: mean Pearson correlation\n",
-    "        correlations_raw = [metric.compute().item() for metric in self.pearson_metrics_raw]\n",
-    "        metrics_dict[\"metrics_raw/mean/pearson\"] = np.nanmean(correlations_raw)\n",
+    "        # Mean Pearson correlation\n",
+    "        correlations = [metric.compute().item() for metric in self.pearson_metrics]\n",
+    "        metrics_dict[\"mean/pearson\"] = np.nanmean(correlations)\n",
     "        \n",
     "        # Mean loss\n",
     "        metrics_dict[\"loss\"] = np.mean(self.losses) if self.losses else 0.0\n",
@@ -716,7 +741,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 38,
+   "execution_count": 19,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -729,148 +754,12 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# 7. Scaling functions setup (copied from pipeline)"
+    "# 7. Loss functions"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 39,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Scaling functions created\n"
-     ]
-    }
-   ],
-   "source": [
-    "def get_track_means(bigwig_file_ids: List[str]) -> np.ndarray:\n",
-    "    \"\"\"\n",
-    "    Get track means for normalization.\n",
-    "    For now, return dummy values. In real pipeline, this loads from metadata.\n",
-    "    \"\"\"\n",
-    "    # Dummy values - in real pipeline, this would load from actual metadata\n",
-    "    return np.ones(len(bigwig_file_ids), dtype=np.float32) * 1.0\n",
-    "\n",
-    "\n",
-    "def get_rna_seq_track_ids(bigwig_file_ids: List[str]) -> List[int]:\n",
-    "    \"\"\"\n",
-    "    Get RNA-seq track indices.\n",
-    "    For now, return empty list. In real pipeline, this identifies RNA-seq tracks.\n",
-    "    \"\"\"\n",
-    "    # Dummy - in real pipeline, this would identify RNA-seq tracks\n",
-    "    return []\n",
-    "\n",
-    "\n",
-    "def create_targets_scaling_fn(bigwig_file_ids: List[str]) -> Callable[[torch.Tensor], torch.Tensor]:\n",
-    "    \"\"\"\n",
-    "    Build a scaling function based on track means and RNA-seq squashing.\n",
-    "    Copied from the supervised tracks pipeline.\n",
-    "    \"\"\"\n",
-    "    # Load track means\n",
-    "    track_means_np = get_track_means(bigwig_file_ids)\n",
-    "    track_means = torch.tensor(track_means_np, dtype=torch.float32)\n",
-    "    \n",
-    "    # Get which tracks use squashing\n",
-    "    rna_ids = get_rna_seq_track_ids(bigwig_file_ids)\n",
-    "    apply_squashing = torch.zeros((len(bigwig_file_ids),), dtype=torch.bool)\n",
-    "    if len(rna_ids) > 0:\n",
-    "        apply_squashing[rna_ids] = True\n",
-    "    \n",
-    "    def transform_fn(x: torch.Tensor) -> torch.Tensor:\n",
-    "        \"\"\"\n",
-    "        x: torch.Tensor, shape (batch, seq_len, num_tracks)\n",
-    "        \"\"\"\n",
-    "        device = x.device\n",
-    "        \n",
-    "        # Move constants to correct device\n",
-    "        means = track_means.to(device)\n",
-    "        squash_mask = apply_squashing.to(device)\n",
-    "        \n",
-    "        # Normalize\n",
-    "        scaled = x / means\n",
-    "        \n",
-    "        # Power squashing where needed\n",
-    "        squashed = torch.where(\n",
-    "            squash_mask.view(1, 1, -1),\n",
-    "            scaled.pow(0.75),\n",
-    "            scaled,\n",
-    "        )\n",
-    "        \n",
-    "        # Smooth clipping: if > 10, apply formula\n",
-    "        clipped = torch.where(\n",
-    "            squashed > 10.0,\n",
-    "            2.0 * torch.sqrt(squashed * 10.0) - 10.0,\n",
-    "            squashed,\n",
-    "        )\n",
-    "        \n",
-    "        return clipped\n",
-    "    \n",
-    "    return transform_fn\n",
-    "\n",
-    "\n",
-    "def create_predictions_scaling_fn(bigwig_file_ids: List[str]) -> Callable[[torch.Tensor], torch.Tensor]:\n",
-    "    \"\"\"\n",
-    "    Inverse scaling function to apply on predictions before computing metrics.\n",
-    "    Copied from the supervised tracks pipeline.\n",
-    "    \"\"\"\n",
-    "    # Load means\n",
-    "    track_means_np = get_track_means(bigwig_file_ids)\n",
-    "    track_means = torch.tensor(track_means_np, dtype=torch.float32)\n",
-    "    \n",
-    "    # RNA-seq mask\n",
-    "    rna_ids = get_rna_seq_track_ids(bigwig_file_ids)\n",
-    "    apply_squashing = torch.zeros((len(bigwig_file_ids),), dtype=torch.bool)\n",
-    "    if len(rna_ids) > 0:\n",
-    "        apply_squashing[rna_ids] = True\n",
-    "    \n",
-    "    def inverse_transform_fn(x: torch.Tensor) -> torch.Tensor:\n",
-    "        \"\"\"\n",
-    "        x: torch.Tensor, shape (batch, seq_len, num_tracks)\n",
-    "        \"\"\"\n",
-    "        device = x.device\n",
-    "        means = track_means.to(device)\n",
-    "        squash_mask = apply_squashing.to(device)\n",
-    "        \n",
-    "        # Undo clipping\n",
-    "        unclipped = torch.where(\n",
-    "            x > 10.0,\n",
-    "            (x + 10.0).pow(2) / (4 * 10.0),\n",
-    "            x,\n",
-    "        )\n",
-    "        \n",
-    "        # Undo squashing\n",
-    "        unsquashed = torch.where(\n",
-    "            squash_mask.view(1, 1, -1),\n",
-    "            unclipped.pow(1.0 / 0.75),\n",
-    "            unclipped,\n",
-    "        )\n",
-    "        \n",
-    "        # Undo normalization\n",
-    "        return unsquashed * means\n",
-    "    \n",
-    "    return inverse_transform_fn\n",
-    "\n",
-    "\n",
-    "# Create scaling functions\n",
-    "scale_targets_fn = create_targets_scaling_fn(config[\"bigwig_file_ids\"])\n",
-    "scale_predictions_fn = create_predictions_scaling_fn(config[\"bigwig_file_ids\"])\n",
-    "\n",
-    "print(\"Scaling functions created\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# 8. Loss functions"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 40,
+   "execution_count": 20,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -887,49 +776,24 @@
     "def poisson_multinomial_loss(\n",
     "    logits: torch.Tensor,\n",
     "    targets: torch.Tensor,\n",
-    "    mask: torch.Tensor | None = None,\n",
     "    shape_loss_coefficient: float = 5.0,\n",
     "    epsilon: float = 1e-7,\n",
     ") -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]:\n",
     "    \"\"\"\n",
     "    Regression loss for bigwig tracks (MSE, Poisson, or Poisson-Multinomial).\n",
     "    \"\"\"\n",
-    "    scale_loss, shape_loss = None, None\n",
-    "    \n",
-    "    if mask is None:\n",
-    "        mask = torch.ones_like(targets, dtype=torch.float32, device=targets.device)\n",
-    "    else:\n",
-    "        mask = mask.float()\n",
-    "    \n",
-    "    mask_sum = mask.sum() + epsilon\n",
-    "    masked_logits = logits * mask\n",
-    "    masked_targets = targets * mask\n",
     "\n",
     "    # Scale loss\n",
-    "    mask_sum_per_track_per_seq = mask.sum(dim=1)  # (batch, num_tracks)\n",
-    "    mask_per_sequence = mask_sum_per_track_per_seq > 0.0  # (batch, num_tracks)\n",
-    "    \n",
-    "    sum_pred = masked_logits.sum(dim=1)  # (batch, num_tracks)\n",
-    "    sum_true = masked_targets.sum(dim=1)  # (batch, num_tracks)\n",
-    "    \n",
+    "    sum_pred = logits.sum(dim=1)  # (batch, num_tracks)\n",
+    "    sum_true = targets.sum(dim=1)  # (batch, num_tracks)\n",
     "    scale_loss = poisson_loss(sum_true, sum_pred, epsilon=epsilon)\n",
-    "    scale_loss = scale_loss / (mask_sum_per_track_per_seq + epsilon)\n",
-    "    \n",
-    "    if mask_per_sequence.any():\n",
-    "        scale_loss_filtered = scale_loss[mask_per_sequence]\n",
-    "        scale_loss = scale_loss_filtered.mean()\n",
-    "    else:\n",
-    "        scale_loss = torch.tensor(0.0, device=targets.device, dtype=targets.dtype)\n",
+    "    scale_loss = scale_loss.mean()\n",
     "    \n",
     "    # Shape loss\n",
-    "    predicted_counts = masked_logits + (epsilon * mask)\n",
-    "    masked_targets_with_epsilon = masked_targets + (epsilon * mask)\n",
-    "    \n",
-    "    denom = predicted_counts.sum(dim=1, keepdim=True) + epsilon\n",
-    "    p_pred = predicted_counts / denom\n",
-    "    \n",
+    "    denom = logits.sum(dim=1, keepdim=True) + epsilon\n",
+    "    p_pred = logits / denom\n",
     "    pl_pred = safe_for_grad_log_torch(p_pred)\n",
-    "    shape_loss = -(masked_targets_with_epsilon * pl_pred).sum() / mask_sum\n",
+    "    shape_loss = -(targets * pl_pred).mean()\n",
     "    \n",
     "    # Combine\n",
     "    loss = shape_loss + scale_loss / shape_loss_coefficient\n",
@@ -941,57 +805,42 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# 9. Training loop"
+    "# 8. Training loop"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 41,
+   "execution_count": 21,
    "metadata": {},
    "outputs": [],
    "source": [
     "def train_step(\n",
     "    model: nn.Module,\n",
     "    batch: Dict[str, torch.Tensor],\n",
-    "    optimizer: torch.optim.Optimizer,\n",
-    "    scale_targets_fn: Callable,\n",
-    "    config: Dict,\n",
-    "    num_accumulation_steps: int = 1,\n",
     ") -> float:\n",
-    "    \"\"\"Single training step with gradient accumulation support.\"\"\"\n",
+    "    \"\"\"Single training step.\"\"\"\n",
     "    tokens = batch[\"tokens\"].to(device)\n",
-    "    bigwig_targets = batch[\"bigwig_targets\"].to(device)  # Shape: (batch, seq_len_cropped, num_tracks)\n",
+    "    bigwig_targets = batch[\"bigwig_targets\"].to(device)\n",
     "    \n",
     "    # Forward pass\n",
     "    outputs = model(tokens=tokens)\n",
-    "    bigwig_logits = outputs[\"bigwig_tracks_logits\"]  # Shape: (batch, cropped_seq_len, num_tracks)\n",
-    "    \n",
-    "    # Scale targets\n",
-    "    scaled_targets = scale_targets_fn(bigwig_targets)\n",
+    "    bigwig_logits = outputs[\"bigwig_tracks_logits\"]\n",
     "    \n",
     "    # Compute loss\n",
     "    loss, _, _ = poisson_multinomial_loss(\n",
     "        logits=bigwig_logits,\n",
-    "        targets=scaled_targets,\n",
-    "        shape_loss_coefficient=config[\"bigwig_shape_loss_coefficient\"],\n",
+    "        targets=bigwig_targets,\n",
     "    )\n",
     "    \n",
-    "    # Scale loss by accumulation steps (for gradient accumulation)\n",
-    "    loss = loss / num_accumulation_steps\n",
-    "    \n",
-    "    # Backward pass (accumulate gradients)\n",
+    "    # Backward pass\n",
     "    loss.backward()\n",
-    "    \n",
-    "    return loss.item() * num_accumulation_steps  # Return unscaled loss for logging\n",
+    "    return loss.item()\n",
     "\n",
     "\n",
     "def validation_step(\n",
     "    model: nn.Module,\n",
     "    batch: Dict[str, torch.Tensor],\n",
-    "    scale_targets_fn: Callable,\n",
-    "    scale_predictions_fn: Callable,\n",
     "    metrics: TracksMetrics,\n",
-    "    config: Dict,\n",
     ") -> float:\n",
     "    \"\"\"Single validation step.\"\"\"\n",
     "    model.eval()\n",
@@ -1004,35 +853,32 @@
     "        outputs = model(tokens=tokens)\n",
     "        bigwig_logits = outputs[\"bigwig_tracks_logits\"]\n",
     "        \n",
-    "        # Scale targets for loss computation\n",
-    "        scaled_targets = scale_targets_fn(bigwig_targets)\n",
-    "        \n",
-    "        # Compute loss (using scaled targets)\n",
+    "        # Compute loss\n",
     "        loss, _, _ = poisson_multinomial_loss(\n",
     "            logits=bigwig_logits,\n",
-    "            targets=scaled_targets,\n",
-    "            shape_loss_coefficient=config[\"bigwig_shape_loss_coefficient\"],\n",
+    "            targets=bigwig_targets,\n",
     "        )\n",
     "        \n",
-    "        # Scale predictions back to original space for metrics\n",
-    "        # (predictions are in scaled space, need to inverse transform)\n",
-    "        unscaled_predictions = scale_predictions_fn(bigwig_logits)\n",
-    "        \n",
-    "        # Update metrics (using original space targets and predictions)\n",
+    "        # Update metrics\n",
     "        metrics.update(\n",
-    "            predictions_scaled=bigwig_logits,\n",
-    "            targets_scaled=scaled_targets,\n",
-    "            predictions_raw=unscaled_predictions,\n",
-    "            targets_raw=bigwig_targets,\n",
+    "            predictions=bigwig_logits,\n",
+    "            targets=bigwig_targets,\n",
     "            loss=loss.item()\n",
     "        )\n",
     "    \n",
     "    return loss.item()"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Interactive plotting is temporary for debug"
+   ]
+  },
   {
    "cell_type": "code",
-   "execution_count": 42,
+   "execution_count": 22,
    "metadata": {},
    "outputs": [
     {
@@ -1040,163 +886,455 @@
      "output_type": "stream",
      "text": [
       "Starting training...\n",
-      "Training for 32 steps with 2 gradient accumulation steps\n",
+      "Training for 1000 steps\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "5935c992adb7428bac8de1aa6873dd7e",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "FigureWidget({\n",
+       "    'data': [{'line': {'color': 'blue'},\n",
+       "              'mode': 'lines+markers',\n",
+       "              'name': 'Train Loss',\n",
+       "              'type': 'scatter',\n",
+       "              'uid': '5424e4af-13b6-48c8-a367-8aa145c3a9db',\n",
+       "              'x': [],\n",
+       "              'xaxis': 'x',\n",
+       "              'y': [],\n",
+       "              'yaxis': 'y'},\n",
+       "             {'line': {'color': 'red'},\n",
+       "              'mode': 'lines+markers',\n",
+       "              'name': 'Val Loss',\n",
+       "              'type': 'scatter',\n",
+       "              'uid': 'fe995660-5f01-4c12-9d7d-9ed19ddee785',\n",
+       "              'x': [],\n",
+       "              'xaxis': 'x',\n",
+       "              'y': [],\n",
+       "              'yaxis': 'y'},\n",
+       "             {'line': {'color': 'green'},\n",
+       "              'mode': 'lines+markers',\n",
+       "              'name': 'Train Pearson',\n",
+       "              'type': 'scatter',\n",
+       "              'uid': '8453b45b-4613-41bc-a46b-ac59ba9e6f97',\n",
+       "              'x': [],\n",
+       "              'xaxis': 'x2',\n",
+       "              'y': [],\n",
+       "              'yaxis': 'y2'},\n",
+       "             {'line': {'color': 'orange'},\n",
+       "              'mode': 'lines+markers',\n",
+       "              'name': 'Val Pearson',\n",
+       "              'type': 'scatter',\n",
+       "              'uid': '0887ea97-abf9-4fcf-8ea8-c638dc153a4d',\n",
+       "              'x': [],\n",
+       "              'xaxis': 'x2',\n",
+       "              'y': [],\n",
+       "              'yaxis': 'y2'}],\n",
+       "    'layout': {'annotations': [{'font': {'size': 16},\n",
+       "                                'showarrow': False,\n",
+       "                                'text': 'Loss',\n",
+       "                                'x': 0.2125,\n",
+       "                                'xanchor': 'center',\n",
+       "                                'xref': 'paper',\n",
+       "                                'y': 1.0,\n",
+       "                                'yanchor': 'bottom',\n",
+       "                                'yref': 'paper'},\n",
+       "                               {'font': {'size': 16},\n",
+       "                                'showarrow': False,\n",
+       "                                'text': 'Mean Pearson Correlation',\n",
+       "                                'x': 0.7875,\n",
+       "                                'xanchor': 'center',\n",
+       "                                'xref': 'paper',\n",
+       "                                'y': 1.0,\n",
+       "                                'yanchor': 'bottom',\n",
+       "                                'yref': 'paper'}],\n",
+       "               'height': 800,\n",
+       "               'showlegend': True,\n",
+       "               'template': '...',\n",
+       "               'title': {'text': 'Training'},\n",
+       "               'width': 1600,\n",
+       "               'xaxis': {'anchor': 'y', 'domain': [0.0, 0.425], 'title': {'text': 'Step'}},\n",
+       "               'xaxis2': {'anchor': 'y2', 'domain': [0.575, 1.0], 'title': {'text': 'Step'}},\n",
+       "               'yaxis': {'anchor': 'x', 'domain': [0.0, 1.0], 'title': {'text': 'Loss'}},\n",
+       "               'yaxis2': {'anchor': 'x2', 'domain': [0.0, 1.0], 'title': {'text': 'Pearson Correlation'}}}\n",
+       "})"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/y-bornachot/venvs/ntv3-env/lib/python3.12/site-packages/torch/amp/autocast_mode.py:287: UserWarning:\n",
+      "\n",
+      "In CPU autocast, but the target dtype is not supported. Disabling autocast.\n",
+      "CPU Autocast only supports dtype of torch.bfloat16, torch.float16 currently.\n",
+      "\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step 10/1000 | Loss: 0.2374 | Mean Pearson: 0.0382 | LR: 1.00e-05\n",
+      "Step 20/1000 | Loss: 2.2259 | Mean Pearson: -0.0884 | LR: 1.00e-05\n",
+      "Step 30/1000 | Loss: 20.0122 | Mean Pearson: 0.1379 | LR: 1.00e-05\n",
+      "Step 40/1000 | Loss: 9.6938 | Mean Pearson: -0.1497 | LR: 1.00e-05\n",
+      "Step 50/1000 | Loss: -1.8435 | Mean Pearson: -0.1875 | LR: 1.00e-05\n",
+      "\n",
+      "Running validation at step 50...\n",
+      "  Validation Loss: 11.5599\n",
+      "  Validation Mean Pearson: -0.1576\n",
+      "    ENCFF884LDL/pearson: -0.1576\n",
+      "Step 60/1000 | Loss: 1.4427 | Mean Pearson: 0.2841 | LR: 1.00e-05\n",
+      "Step 70/1000 | Loss: -3.4037 | Mean Pearson: -0.1362 | LR: 1.00e-05\n",
+      "Step 80/1000 | Loss: 9.0958 | Mean Pearson: -0.1319 | LR: 1.00e-05\n",
+      "Step 90/1000 | Loss: -7.8433 | Mean Pearson: -0.0576 | LR: 1.00e-05\n",
+      "Step 100/1000 | Loss: 7.3503 | Mean Pearson: -0.2150 | LR: 1.00e-05\n",
       "\n",
-      "Step 1/32 | Loss: 0.5661 | Mean Pearson: -0.0525 | Tokens: 4,096\n",
+      "Running validation at step 100...\n",
+      "  Validation Loss: 22.3383\n",
+      "  Validation Mean Pearson: -0.2867\n",
+      "    ENCFF884LDL/pearson: -0.2867\n",
+      "Step 110/1000 | Loss: -8.1600 | Mean Pearson: -0.1616 | LR: 1.00e-05\n",
+      "Step 120/1000 | Loss: -0.8743 | Mean Pearson: -0.1318 | LR: 1.00e-05\n",
+      "Step 130/1000 | Loss: -2.9825 | Mean Pearson: -0.0480 | LR: 1.00e-05\n",
+      "Step 140/1000 | Loss: -2.4524 | Mean Pearson: -0.0879 | LR: 1.00e-05\n",
+      "Step 150/1000 | Loss: 3.8818 | Mean Pearson: -0.0907 | LR: 1.00e-05\n",
       "\n",
-      "Running validation at step 0...\n",
-      "  Validation Loss: 0.3987\n",
-      "  Validation Mean Pearson: -0.0426\n",
-      "    ENCFF884LDL/pearson: -0.0426\n",
-      "Step 3/32 | Loss: 0.3825 | Mean Pearson: -0.0112 | Tokens: 12,288\n",
-      "Step 5/32 | Loss: 1.1384 | Mean Pearson: -0.0777 | Tokens: 20,480\n",
+      "Running validation at step 150...\n",
+      "  Validation Loss: 19.6866\n",
+      "  Validation Mean Pearson: -0.2207\n",
+      "    ENCFF884LDL/pearson: -0.2207\n",
+      "Step 160/1000 | Loss: -1.0933 | Mean Pearson: -0.1243 | LR: 1.00e-05\n",
+      "Step 170/1000 | Loss: -2.2577 | Mean Pearson: -0.0212 | LR: 1.00e-05\n",
+      "Step 180/1000 | Loss: 0.0738 | Mean Pearson: 0.5643 | LR: 1.00e-05\n",
+      "Step 190/1000 | Loss: -0.1097 | Mean Pearson: 0.0309 | LR: 1.00e-05\n",
+      "Step 200/1000 | Loss: -8.7972 | Mean Pearson: 0.4804 | LR: 1.00e-05\n",
       "\n",
-      "Running validation at step 4...\n",
-      "  Validation Loss: 0.4381\n",
-      "  Validation Mean Pearson: -0.0017\n",
-      "    ENCFF884LDL/pearson: -0.0017\n",
-      "Step 7/32 | Loss: 0.4961 | Mean Pearson: -0.0188 | Tokens: 28,672\n",
-      "Step 9/32 | Loss: 0.4903 | Mean Pearson: -0.1522 | Tokens: 36,864\n",
+      "Running validation at step 200...\n",
+      "  Validation Loss: -8.8160\n",
+      "  Validation Mean Pearson: 0.0912\n",
+      "    ENCFF884LDL/pearson: 0.0912\n",
+      "Step 210/1000 | Loss: -2.5429 | Mean Pearson: 0.3908 | LR: 1.00e-05\n",
+      "Step 220/1000 | Loss: -6.8421 | Mean Pearson: 0.4080 | LR: 1.00e-05\n",
+      "Step 230/1000 | Loss: -4.4312 | Mean Pearson: -0.0400 | LR: 1.00e-05\n",
+      "Step 240/1000 | Loss: -11.4732 | Mean Pearson: 0.6653 | LR: 1.00e-05\n",
+      "Step 250/1000 | Loss: -9.2648 | Mean Pearson: 0.0539 | LR: 1.00e-05\n",
       "\n",
-      "Running validation at step 8...\n",
-      "  Validation Loss: 0.3429\n",
-      "  Validation Mean Pearson: -0.0997\n",
-      "    ENCFF884LDL/pearson: -0.0997\n",
-      "Step 11/32 | Loss: 0.4597 | Mean Pearson: -0.0199 | Tokens: 45,056\n",
-      "Step 13/32 | Loss: 0.6507 | Mean Pearson: -0.0256 | Tokens: 53,248\n",
+      "Running validation at step 250...\n",
+      "  Validation Loss: -6.8987\n",
+      "  Validation Mean Pearson: 0.0654\n",
+      "    ENCFF884LDL/pearson: 0.0654\n",
+      "Step 260/1000 | Loss: -0.6699 | Mean Pearson: 0.0913 | LR: 1.00e-05\n",
+      "Step 270/1000 | Loss: -8.6625 | Mean Pearson: 0.3179 | LR: 1.00e-05\n",
+      "Step 280/1000 | Loss: -11.7691 | Mean Pearson: 0.0004 | LR: 1.00e-05\n",
+      "Step 290/1000 | Loss: -14.1622 | Mean Pearson: 0.0492 | LR: 1.00e-05\n",
+      "Step 300/1000 | Loss: 0.9208 | Mean Pearson: 0.0607 | LR: 1.00e-05\n",
       "\n",
-      "Running validation at step 12...\n",
-      "  Validation Loss: 0.3901\n",
-      "  Validation Mean Pearson: -0.0786\n",
-      "    ENCFF884LDL/pearson: -0.0786\n",
-      "Step 15/32 | Loss: 0.3911 | Mean Pearson: -0.0419 | Tokens: 61,440\n",
-      "Step 17/32 | Loss: 0.4202 | Mean Pearson: -0.0883 | Tokens: 69,632\n",
+      "Running validation at step 300...\n",
+      "  Validation Loss: -5.0427\n",
+      "  Validation Mean Pearson: 0.3464\n",
+      "    ENCFF884LDL/pearson: 0.3464\n",
+      "Step 310/1000 | Loss: -1.2881 | Mean Pearson: 0.1696 | LR: 1.00e-05\n",
+      "Step 320/1000 | Loss: -18.6637 | Mean Pearson: 0.0892 | LR: 1.00e-05\n",
+      "Step 330/1000 | Loss: -36.6038 | Mean Pearson: 0.3356 | LR: 1.00e-05\n",
+      "Step 340/1000 | Loss: -2.4984 | Mean Pearson: 0.2305 | LR: 1.00e-05\n",
+      "Step 350/1000 | Loss: -4.7985 | Mean Pearson: 0.0968 | LR: 1.00e-05\n",
       "\n",
-      "Running validation at step 16...\n",
-      "  Validation Loss: 0.3626\n",
-      "  Validation Mean Pearson: -0.0840\n",
-      "    ENCFF884LDL/pearson: -0.0840\n",
-      "Step 19/32 | Loss: 0.3608 | Mean Pearson: -0.1057 | Tokens: 77,824\n",
-      "Step 21/32 | Loss: 0.3942 | Mean Pearson: 0.1459 | Tokens: 86,016\n",
+      "Running validation at step 350...\n",
+      "  Validation Loss: -13.6500\n",
+      "  Validation Mean Pearson: 0.2737\n",
+      "    ENCFF884LDL/pearson: 0.2737\n",
+      "Step 360/1000 | Loss: -9.4795 | Mean Pearson: 0.0579 | LR: 1.00e-05\n",
+      "Step 370/1000 | Loss: 0.3531 | Mean Pearson: 0.0240 | LR: 1.00e-05\n",
+      "Step 380/1000 | Loss: -5.7921 | Mean Pearson: 0.4119 | LR: 1.00e-05\n",
+      "Step 390/1000 | Loss: -2.7049 | Mean Pearson: 0.1343 | LR: 1.00e-05\n",
+      "Step 400/1000 | Loss: -32.8422 | Mean Pearson: 0.1545 | LR: 1.00e-05\n",
       "\n",
-      "Running validation at step 20...\n",
-      "  Validation Loss: 0.3281\n",
-      "  Validation Mean Pearson: -0.0667\n",
-      "    ENCFF884LDL/pearson: -0.0667\n",
-      "Step 23/32 | Loss: 0.4090 | Mean Pearson: 0.0540 | Tokens: 94,208\n",
-      "Step 25/32 | Loss: 0.5151 | Mean Pearson: -0.0076 | Tokens: 102,400\n",
+      "Running validation at step 400...\n",
+      "  Validation Loss: -4.3502\n",
+      "  Validation Mean Pearson: 0.3124\n",
+      "    ENCFF884LDL/pearson: 0.3124\n",
+      "Step 410/1000 | Loss: -18.9574 | Mean Pearson: 0.0594 | LR: 1.00e-05\n",
+      "Step 420/1000 | Loss: -5.4032 | Mean Pearson: 0.2804 | LR: 1.00e-05\n",
+      "Step 430/1000 | Loss: -0.5171 | Mean Pearson: 0.1835 | LR: 1.00e-05\n",
+      "Step 440/1000 | Loss: -3.4071 | Mean Pearson: 0.0680 | LR: 1.00e-05\n",
+      "Step 450/1000 | Loss: -3.5580 | Mean Pearson: 0.0850 | LR: 1.00e-05\n",
       "\n",
-      "Running validation at step 24...\n",
-      "  Validation Loss: 0.2927\n",
-      "  Validation Mean Pearson: -0.0409\n",
-      "    ENCFF884LDL/pearson: -0.0409\n",
-      "Step 27/32 | Loss: 0.4339 | Mean Pearson: -0.0887 | Tokens: 110,592\n",
-      "Step 29/32 | Loss: 0.4516 | Mean Pearson: -0.0763 | Tokens: 118,784\n",
+      "Running validation at step 450...\n",
+      "  Validation Loss: -7.3308\n",
+      "  Validation Mean Pearson: 0.1128\n",
+      "    ENCFF884LDL/pearson: 0.1128\n",
+      "Step 460/1000 | Loss: -0.9750 | Mean Pearson: 0.1717 | LR: 1.00e-05\n",
+      "Step 470/1000 | Loss: -5.5775 | Mean Pearson: 0.1321 | LR: 1.00e-05\n",
+      "Step 480/1000 | Loss: -1.1170 | Mean Pearson: 0.1484 | LR: 1.00e-05\n",
+      "Step 490/1000 | Loss: -3.8053 | Mean Pearson: 0.1959 | LR: 1.00e-05\n",
+      "Step 500/1000 | Loss: -4.5933 | Mean Pearson: 0.1860 | LR: 1.00e-05\n",
       "\n",
-      "Running validation at step 28...\n",
-      "  Validation Loss: 0.3076\n",
-      "  Validation Mean Pearson: -0.0861\n",
-      "    ENCFF884LDL/pearson: -0.0861\n",
-      "Step 31/32 | Loss: 0.4121 | Mean Pearson: -0.0530 | Tokens: 126,976\n",
+      "Running validation at step 500...\n",
+      "  Validation Loss: -5.7617\n",
+      "  Validation Mean Pearson: 0.3155\n",
+      "    ENCFF884LDL/pearson: 0.3155\n",
+      "Step 510/1000 | Loss: -3.3306 | Mean Pearson: 0.2815 | LR: 1.00e-05\n",
+      "Step 520/1000 | Loss: -2.1962 | Mean Pearson: 0.1151 | LR: 1.00e-05\n",
+      "Step 530/1000 | Loss: -1.5388 | Mean Pearson: 0.3783 | LR: 1.00e-05\n",
+      "Step 540/1000 | Loss: -2.2349 | Mean Pearson: 0.0734 | LR: 1.00e-05\n",
+      "Step 550/1000 | Loss: -1.5502 | Mean Pearson: 0.2171 | LR: 1.00e-05\n",
       "\n",
-      "Training completed after 32 steps!\n"
+      "Running validation at step 550...\n",
+      "  Validation Loss: -3.0059\n",
+      "  Validation Mean Pearson: 0.2325\n",
+      "    ENCFF884LDL/pearson: 0.2325\n",
+      "Step 560/1000 | Loss: -2.0764 | Mean Pearson: -0.0049 | LR: 1.00e-05\n",
+      "Step 570/1000 | Loss: -1.7384 | Mean Pearson: 0.2989 | LR: 1.00e-05\n",
+      "Step 580/1000 | Loss: -6.7306 | Mean Pearson: 0.2522 | LR: 1.00e-05\n",
+      "Step 590/1000 | Loss: -3.2473 | Mean Pearson: 0.1042 | LR: 1.00e-05\n",
+      "Step 600/1000 | Loss: -4.2841 | Mean Pearson: 0.1936 | LR: 1.00e-05\n",
+      "\n",
+      "Running validation at step 600...\n",
+      "  Validation Loss: -4.5611\n",
+      "  Validation Mean Pearson: 0.2744\n",
+      "    ENCFF884LDL/pearson: 0.2744\n",
+      "Step 610/1000 | Loss: -3.5691 | Mean Pearson: 0.1803 | LR: 1.00e-05\n",
+      "Step 620/1000 | Loss: -7.2129 | Mean Pearson: 0.0901 | LR: 1.00e-05\n",
+      "Step 630/1000 | Loss: -6.0598 | Mean Pearson: 0.1795 | LR: 1.00e-05\n",
+      "Step 640/1000 | Loss: -2.8917 | Mean Pearson: 0.1111 | LR: 1.00e-05\n",
+      "Step 650/1000 | Loss: -2.7210 | Mean Pearson: 0.3566 | LR: 1.00e-05\n",
+      "\n",
+      "Running validation at step 650...\n",
+      "  Validation Loss: -4.3997\n",
+      "  Validation Mean Pearson: 0.3327\n",
+      "    ENCFF884LDL/pearson: 0.3327\n",
+      "Step 660/1000 | Loss: -3.4793 | Mean Pearson: 0.0441 | LR: 1.00e-05\n",
+      "Step 670/1000 | Loss: -1.9743 | Mean Pearson: 0.1364 | LR: 1.00e-05\n",
+      "Step 680/1000 | Loss: -5.7498 | Mean Pearson: 0.2330 | LR: 1.00e-05\n",
+      "Step 690/1000 | Loss: -12.8701 | Mean Pearson: 0.3182 | LR: 1.00e-05\n",
+      "Step 700/1000 | Loss: -1.5847 | Mean Pearson: 0.1971 | LR: 1.00e-05\n",
+      "\n",
+      "Running validation at step 700...\n",
+      "  Validation Loss: -2.0630\n",
+      "  Validation Mean Pearson: 0.1267\n",
+      "    ENCFF884LDL/pearson: 0.1267\n",
+      "Step 710/1000 | Loss: -6.0704 | Mean Pearson: 0.3715 | LR: 1.00e-05\n",
+      "Step 720/1000 | Loss: -2.6020 | Mean Pearson: 0.1244 | LR: 1.00e-05\n",
+      "Step 730/1000 | Loss: -58.8965 | Mean Pearson: 0.5625 | LR: 1.00e-05\n",
+      "Step 740/1000 | Loss: -1.2855 | Mean Pearson: 0.2658 | LR: 1.00e-05\n",
+      "Step 750/1000 | Loss: -4.4599 | Mean Pearson: 0.0137 | LR: 1.00e-05\n",
+      "\n",
+      "Running validation at step 750...\n",
+      "  Validation Loss: -11.1562\n",
+      "  Validation Mean Pearson: 0.0844\n",
+      "    ENCFF884LDL/pearson: 0.0844\n",
+      "Step 760/1000 | Loss: -11.6905 | Mean Pearson: 0.1914 | LR: 1.00e-05\n",
+      "Step 770/1000 | Loss: -4.0964 | Mean Pearson: 0.2022 | LR: 1.00e-05\n",
+      "Step 780/1000 | Loss: -1.5512 | Mean Pearson: 0.3568 | LR: 1.00e-05\n",
+      "Step 790/1000 | Loss: -5.5843 | Mean Pearson: 0.2058 | LR: 1.00e-05\n",
+      "Step 800/1000 | Loss: -3.9190 | Mean Pearson: 0.4362 | LR: 1.00e-05\n",
+      "\n",
+      "Running validation at step 800...\n",
+      "  Validation Loss: -4.7017\n",
+      "  Validation Mean Pearson: 0.3817\n",
+      "    ENCFF884LDL/pearson: 0.3817\n",
+      "Step 810/1000 | Loss: -7.6856 | Mean Pearson: 0.0672 | LR: 1.00e-05\n",
+      "Step 820/1000 | Loss: -5.3603 | Mean Pearson: 0.2325 | LR: 1.00e-05\n",
+      "Step 830/1000 | Loss: -3.8539 | Mean Pearson: 0.2808 | LR: 1.00e-05\n",
+      "Step 840/1000 | Loss: -8.1141 | Mean Pearson: 0.2529 | LR: 1.00e-05\n",
+      "Step 850/1000 | Loss: -10.5886 | Mean Pearson: 0.3454 | LR: 1.00e-05\n",
+      "\n",
+      "Running validation at step 850...\n",
+      "  Validation Loss: -4.9108\n",
+      "  Validation Mean Pearson: 0.2195\n",
+      "    ENCFF884LDL/pearson: 0.2195\n",
+      "Step 860/1000 | Loss: -4.1028 | Mean Pearson: 0.3304 | LR: 1.00e-05\n",
+      "Step 870/1000 | Loss: -7.1834 | Mean Pearson: 0.1206 | LR: 1.00e-05\n",
+      "Step 880/1000 | Loss: -8.9869 | Mean Pearson: 0.3584 | LR: 1.00e-05\n",
+      "Step 890/1000 | Loss: -2.2697 | Mean Pearson: 0.0943 | LR: 1.00e-05\n",
+      "Step 900/1000 | Loss: -14.0142 | Mean Pearson: 0.4761 | LR: 1.00e-05\n",
+      "\n",
+      "Running validation at step 900...\n",
+      "  Validation Loss: -3.2329\n",
+      "  Validation Mean Pearson: 0.3635\n",
+      "    ENCFF884LDL/pearson: 0.3635\n",
+      "Step 910/1000 | Loss: -9.0941 | Mean Pearson: 0.2754 | LR: 1.00e-05\n",
+      "Step 920/1000 | Loss: -4.6371 | Mean Pearson: 0.0167 | LR: 1.00e-05\n",
+      "Step 930/1000 | Loss: -7.9853 | Mean Pearson: 0.0941 | LR: 1.00e-05\n",
+      "Step 940/1000 | Loss: -22.9349 | Mean Pearson: 0.5140 | LR: 1.00e-05\n",
+      "Step 950/1000 | Loss: -2.0866 | Mean Pearson: 0.1746 | LR: 1.00e-05\n",
+      "\n",
+      "Running validation at step 950...\n",
+      "  Validation Loss: -8.8318\n",
+      "  Validation Mean Pearson: 0.1597\n",
+      "    ENCFF884LDL/pearson: 0.1597\n",
+      "Step 960/1000 | Loss: -4.8540 | Mean Pearson: 0.6318 | LR: 1.00e-05\n",
+      "Step 970/1000 | Loss: -4.1091 | Mean Pearson: 0.0985 | LR: 1.00e-05\n",
+      "Step 980/1000 | Loss: -5.1141 | Mean Pearson: 0.2031 | LR: 1.00e-05\n",
+      "Step 990/1000 | Loss: -4.1959 | Mean Pearson: 0.2404 | LR: 1.00e-05\n",
+      "Step 1000/1000 | Loss: -0.9942 | Mean Pearson: 0.2742 | LR: 1.00e-05\n",
+      "\n",
+      "Running validation at step 1000...\n",
+      "  Validation Loss: -4.2796\n",
+      "  Validation Mean Pearson: 0.1425\n",
+      "    ENCFF884LDL/pearson: 0.1425\n",
+      "\n",
+      "Training completed after 1000 steps.\n"
      ]
     }
    ],
    "source": [
-    "# Training loop (step-based with gradient accumulation)\n",
+    "# Training loop\n",
     "print(\"Starting training...\")\n",
-    "print(f\"Training for {num_steps_training} steps with {num_accumulation_gradient} gradient accumulation steps\\n\")\n",
+    "print(f\"Training for {config[\"num_steps_training\"]} steps\\n\")\n",
     "\n",
     "model.train()\n",
     "train_metrics.reset()\n",
     "optimizer.zero_grad()  # Initialize gradients\n",
     "\n",
+    "# Track metrics for plotting\n",
+    "train_steps = []\n",
+    "train_losses = []\n",
+    "train_pearson_scores = []\n",
+    "val_steps = []\n",
+    "val_losses = []\n",
+    "val_pearson_scores = []\n",
+    "\n",
+    "# Initialize interactive plots using FigureWidget for real-time updates\n",
+    "from plotly.graph_objects import FigureWidget\n",
+    "from plotly.subplots import make_subplots\n",
+    "\n",
+    "# Create base figure with subplots\n",
+    "fig_base = make_subplots(\n",
+    "    rows=1, cols=2,\n",
+    "    subplot_titles=('Loss', 'Mean Pearson Correlation'),\n",
+    "    horizontal_spacing=0.15,\n",
+    ")\n",
+    "\n",
+    "# Add empty traces for train and val metrics\n",
+    "fig_base.add_trace(\n",
+    "    go.Scatter(x=[], y=[], mode='lines+markers', name='Train Loss', line=dict(color='blue')),\n",
+    "    row=1, col=1\n",
+    ")\n",
+    "fig_base.add_trace(\n",
+    "    go.Scatter(x=[], y=[], mode='lines+markers', name='Val Loss', line=dict(color='red')),\n",
+    "    row=1, col=1\n",
+    ")\n",
+    "fig_base.add_trace(\n",
+    "    go.Scatter(x=[], y=[], mode='lines+markers', name='Train Pearson', line=dict(color='green')),\n",
+    "    row=1, col=2\n",
+    ")\n",
+    "fig_base.add_trace(\n",
+    "    go.Scatter(x=[], y=[], mode='lines+markers', name='Val Pearson', line=dict(color='orange')),\n",
+    "    row=1, col=2\n",
+    ")\n",
+    "\n",
+    "fig_base.update_xaxes(title_text=\"Step\", row=1, col=1)\n",
+    "fig_base.update_xaxes(title_text=\"Step\", row=1, col=2)\n",
+    "fig_base.update_yaxes(title_text=\"Loss\", row=1, col=1)\n",
+    "fig_base.update_yaxes(title_text=\"Pearson Correlation\", row=1, col=2)\n",
+    "fig_base.update_layout(height=800, width=1600, showlegend=True, title_text=\"Training\")\n",
+    "\n",
+    "# Convert to FigureWidget for interactive updates\n",
+    "fig = FigureWidget(fig_base)\n",
+    "\n",
+    "# Display initial plot (will update in place during training)\n",
+    "display(fig)\n",
+    "\n",
     "# Create iterator for training data (will cycle if needed)\n",
     "train_iter = iter(train_loader)\n",
-    "num_tokens_seen = 0\n",
-    "\n",
-    "# Main training loop: for loop over optimizer steps (like deepspeed pipeline)\n",
-    "for optimizer_step_idx in range(num_steps_training):\n",
-    "    # Gradient accumulation loop\n",
-    "    accumulated_loss = 0.0\n",
-    "    for acc_idx in range(num_accumulation_gradient):\n",
-    "        try:\n",
-    "            batch = next(train_iter)\n",
-    "        except StopIteration:\n",
-    "            # Restart iterator if we run out of data\n",
-    "            train_iter = iter(train_loader)\n",
-    "            batch = next(train_iter)\n",
-    "        \n",
-    "        # Forward pass and accumulate gradients\n",
-    "        loss = train_step(\n",
-    "            model, batch, optimizer, scale_targets_fn, config, \n",
-    "            num_accumulation_steps=num_accumulation_gradient\n",
-    "        )\n",
-    "        accumulated_loss += loss\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",
-    "    # Update optimizer (after accumulation)\n",
+    "    # Forward pass and backward pass\n",
+    "    loss = train_step(model, batch)\n",
+    "    \n",
+    "    # Update optimizer\n",
     "    optimizer.step()\n",
     "    optimizer.zero_grad()\n",
     "    \n",
-    "        # Update tokens seen\n",
-    "    num_tokens_seen += effective_num_tokens_per_update\n",
-    "    \n",
-    "    # Update metrics (on last batch of accumulation)\n",
+    "    # Update metrics\n",
     "    tokens = batch[\"tokens\"].to(device)\n",
     "    bigwig_targets = batch[\"bigwig_targets\"].to(device)\n",
     "    with torch.no_grad():\n",
     "        outputs = model(tokens=tokens)\n",
     "        bigwig_logits = outputs[\"bigwig_tracks_logits\"]\n",
     "        \n",
-    "        # Scale targets for scaled metrics\n",
-    "        scaled_targets = scale_targets_fn(bigwig_targets)\n",
-    "        \n",
-    "        # Unscale predictions for raw metrics\n",
-    "        unscaled_predictions = scale_predictions_fn(bigwig_logits)\n",
-    "        \n",
-    "        avg_loss = accumulated_loss / num_accumulation_gradient\n",
     "        train_metrics.update(\n",
-    "            predictions_scaled=bigwig_logits,\n",
-    "            targets_scaled=scaled_targets,\n",
-    "            predictions_raw=unscaled_predictions,\n",
-    "            targets_raw=bigwig_targets,\n",
-    "            loss=avg_loss\n",
+    "            predictions=bigwig_logits,\n",
+    "            targets=bigwig_targets,\n",
+    "            loss=loss\n",
     "        )\n",
     "    \n",
     "    # Logging\n",
-    "    if optimizer_step_idx % log_train_step == 0:\n",
+    "    if (step_idx + 1) % config[\"log_every_n_steps\"] == 0:\n",
     "        train_metrics_dict = train_metrics.compute()\n",
-    "        current_lr = config[\"learning_rate\"]\n",
-    "        print(f\"Step {optimizer_step_idx + 1}/{num_steps_training} | \"\n",
-    "              f\"Loss: {avg_loss:.4f} | \"\n",
-    "              f\"Mean Pearson: {train_metrics_dict['metrics_scaled/mean/pearson']:.4f} | \"\n",
-    "              f\"Tokens: {num_tokens_seen:,}\")\n",
+    "        current_lr = optimizer.param_groups[0]['lr']\n",
+    "        \n",
+    "        # Track metrics for plotting\n",
+    "        train_steps.append(step_idx + 1)\n",
+    "        train_losses.append(loss)\n",
+    "        train_pearson_scores.append(train_metrics_dict['mean/pearson'])\n",
+    "        \n",
+    "        # Update plots - direct assignment to FigureWidget data updates the plot automatically\n",
+    "        fig.data[0].x = train_steps\n",
+    "        fig.data[0].y = train_losses\n",
+    "        fig.data[2].x = train_steps\n",
+    "        fig.data[2].y = train_pearson_scores\n",
+    "        \n",
+    "        print(f\"Step {step_idx + 1}/{config[\"num_steps_training\"]} | \"\n",
+    "              f\"Loss: {loss:.4f} | \"\n",
+    "              f\"Mean Pearson: {train_metrics_dict['mean/pearson']:.4f} | \"\n",
+    "              f\"LR: {current_lr:.2e}\")\n",
     "        train_metrics.reset()\n",
     "    \n",
     "    # Validation\n",
-    "    if optimizer_step_idx % log_validation_step == 0:\n",
-    "        print(f\"\\nRunning validation at step {optimizer_step_idx}...\")\n",
+    "    if (step_idx + 1) % config[\"validate_every_n_steps\"] == 0:\n",
+    "        print(f\"\\nRunning validation at step {step_idx + 1}...\")\n",
     "        val_metrics.reset()\n",
     "        model.eval()\n",
     "        \n",
-    "        val_losses = []\n",
+    "        val_batch_losses = []\n",
     "        for val_batch in val_loader:\n",
-    "            val_loss = validation_step(\n",
-    "                model, val_batch, scale_targets_fn, scale_predictions_fn, val_metrics, config\n",
-    "            )\n",
-    "            val_losses.append(val_loss)\n",
+    "            val_loss = validation_step(model, val_batch, val_metrics)\n",
+    "            val_batch_losses.append(val_loss)\n",
     "        \n",
     "        # Print validation metrics\n",
     "        val_metrics_dict = val_metrics.compute()\n",
-    "        print(f\"  Validation Loss: {np.mean(val_losses):.4f}\")\n",
-    "        print(f\"  Validation Mean Pearson: {val_metrics_dict['metrics_scaled/mean/pearson']:.4f}\")\n",
+    "        val_loss_mean = np.mean(val_batch_losses)\n",
+    "        val_pearson_mean = val_metrics_dict['mean/pearson']\n",
+    "        \n",
+    "        # Track validation metrics\n",
+    "        val_steps.append(step_idx + 1)\n",
+    "        val_losses.append(val_loss_mean)\n",
+    "        val_pearson_scores.append(val_pearson_mean)\n",
+    "        \n",
+    "        # Update plots with validation data - direct assignment updates the plot automatically\n",
+    "        fig.data[1].x = val_steps\n",
+    "        fig.data[1].y = val_losses\n",
+    "        fig.data[3].x = val_steps\n",
+    "        fig.data[3].y = val_pearson_scores\n",
+    "        \n",
+    "        print(f\"  Validation Loss: {val_loss_mean:.4f}\")\n",
+    "        print(f\"  Validation Mean Pearson: {val_pearson_mean:.4f}\")\n",
     "        for track_name in config[\"bigwig_file_ids\"]:\n",
-    "            print(f\"    {track_name}/pearson: {val_metrics_dict[f'metrics_scaled/{track_name}/pearson']:.4f}\")\n",
+    "            print(f\"    {track_name}/pearson: {val_metrics_dict[f'{track_name}/pearson']:.4f}\")\n",
     "        \n",
     "        model.train()  # Back to training mode\n",
     "\n",
-    "print(f\"\\nTraining completed after {num_steps_training} steps!\")\n"
+    "print(f\"\\nTraining completed after {config[\"num_steps_training\"]} steps.\")"
    ]
   },
   {
@@ -1208,122 +1346,59 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 43,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def test_step(\n",
-    "    model: nn.Module,\n",
-    "    batch: Dict[str, torch.Tensor],\n",
-    "    scale_targets_fn: Callable,\n",
-    "    scale_predictions_fn: Callable,\n",
-    "    metrics: TracksMetrics,\n",
-    ") -> None:\n",
-    "    \"\"\"\n",
-    "    Pure evaluation step for test set (no loss computation).\n",
-    "    Based on tracks_evaluation_step_torch from deepspeed pipeline.\n",
-    "    \"\"\"\n",
-    "    tokens = batch[\"tokens\"].to(device)\n",
-    "    bigwig_targets = batch[\"bigwig_targets\"].to(device)  # Shape: (batch, seq_len_cropped, num_tracks)\n",
-    "    \n",
-    "    with torch.no_grad():\n",
-    "        # Forward pass\n",
-    "        outputs = model(tokens=tokens)\n",
-    "        bigwig_logits = outputs[\"bigwig_tracks_logits\"]  # Shape: (batch, cropped_seq_len, num_tracks)\n",
-    "        \n",
-    "        # Scale targets for scaled metrics\n",
-    "        scaled_targets = scale_targets_fn(bigwig_targets)\n",
-    "        \n",
-    "        # Unscale predictions for raw metrics\n",
-    "        unscaled_predictions = scale_predictions_fn(bigwig_logits)\n",
-    "        \n",
-    "        # Update metrics with both scaled and raw values\n",
-    "        # Pass 0.0 as loss since we don't compute loss in test evaluation\n",
-    "        metrics.update(\n",
-    "            predictions_scaled=bigwig_logits,\n",
-    "            targets_scaled=scaled_targets,\n",
-    "            predictions_raw=unscaled_predictions,\n",
-    "            targets_raw=bigwig_targets,\n",
-    "            loss=0.0\n",
-    "        )"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 28,
+   "execution_count": 24,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "\n",
-      "==================================================\n",
-      "Test Set Evaluation\n",
-      "==================================================\n",
-      "Running test evaluation with 5 steps (10 samples)\n",
+      "Running test evaluation with 12 steps (100 samples)\n",
       "\n",
       "==================================================\n",
       "Test Set Results\n",
       "==================================================\n",
       "\n",
-      "Scaled Metrics (scaled predictions vs scaled targets):\n",
-      "  Mean Pearson (scaled): -0.0020\n",
-      "    ENCFF884LDL/pearson: -0.0020\n",
-      "\n",
-      "Raw Metrics (raw predictions vs raw targets):\n",
-      "  Mean Pearson (raw): -0.0020\n",
-      "    ENCFF884LDL/pearson: -0.0020\n",
-      "==================================================\n"
+      "Metrics:\n",
+      "  Mean Pearson: 0.1787\n",
+      "    ENCFF884LDL/pearson: 0.1787\n"
      ]
     }
    ],
    "source": [
-    "print(\"\\n\" + \"=\"*50)\n",
-    "print(\"Test Set Evaluation\")\n",
-    "print(\"=\"*50)\n",
-    "\n",
     "# 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",
-    "\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",
-    "# Create iterator for test data\n",
-    "test_iter = iter(test_loader)\n",
+    "for test_batch in test_loader:        \n",
     "\n",
-    "# Run test evaluation (based on deepspeed pipeline: for loop over test steps)\n",
-    "for _ in range(num_test_steps):\n",
-    "    try:\n",
-    "        test_batch = next(test_iter)\n",
-    "    except StopIteration:\n",
-    "        break\n",
-    "    \n",
-    "    # Perform test evaluation (pure evaluation, no loss computation)\n",
-    "    test_step(\n",
-    "        model, test_batch, scale_targets_fn, scale_predictions_fn, test_metrics\n",
+    "    _ = validation_step(        \n",
+    "        model, \n",
+    "        test_batch, \n",
+    "        test_metrics,\n",
     "    )\n",
-    "\n",
+    "    \n",
     "# Compute final test metrics\n",
     "test_metrics_dict = test_metrics.compute()\n",
-    "\n",
     "print(\"\\n\" + \"=\"*50)\n",
     "print(\"Test Set Results\")\n",
     "print(\"=\"*50)\n",
-    "print(f\"\\nScaled Metrics (scaled predictions vs scaled targets):\")\n",
-    "print(f\"  Mean Pearson (scaled): {test_metrics_dict['metrics_scaled/mean/pearson']:.4f}\")\n",
-    "for track_name in config[\"bigwig_file_ids\"]:\n",
-    "    print(f\"    {track_name}/pearson: {test_metrics_dict[f'metrics_scaled/{track_name}/pearson']:.4f}\")\n",
-    "\n",
-    "print(f\"\\nRaw Metrics (raw predictions vs raw targets):\")\n",
-    "print(f\"  Mean Pearson (raw): {test_metrics_dict['metrics_raw/mean/pearson']:.4f}\")\n",
-    "for track_name in config[\"bigwig_file_ids\"]:\n",
-    "    print(f\"    {track_name}/pearson: {test_metrics_dict[f'metrics_raw/{track_name}/pearson']:.4f}\")\n",
-    "print(\"=\"*50)"
+    "print(f\"\\nMetrics:\")\n",
+    "print(f\"  Mean Pearson: {test_metrics_dict['mean/pearson']:.4f}\")\n",
+    "for track_name in config[\"bigwig_file_ids\"]:    \n",
+    "    print(f\"    {track_name}/pearson: {test_metrics_dict[f'{track_name}/pearson']:.4f}\")"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {