Blame for notebooks_tutorials/05_diffusion_fine_tuning.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b55e2d88",
   "metadata": {},
   "source": [
    "# 🧬 Training a discrete diffusion model on biological sequences\n",
    "\n",
    "This notebook demonstrates how to **fine-tune NTv3** for controllable enhancer sequence generation using **Masked Diffusion Language Modeling (MDLM)** ([paper](https://arxiv.org/pdf/2406.07524)).\n",
    "\n",
    "---\n",
    "\n",
    "## 🎯 What You'll Learn\n",
    "\n",
    "This tutorial covers:\n",
    "1. How to extend the post-trained [`InstaDeepAI/NTv3_tracks_1mb`](https://huggingface.co/InstaDeepAI/NTv3_tracks_1mb) model with new conditioning signals\n",
    "2. The fundamentals of Masked Diffusion Language Modeling (MDLM) for DNA generation\n",
    "3. Best practices for training generative DNA models with EMA stabilization\n",
    "\n",
    "---\n",
    "\n",
    "## 🧪 Background: Generative DNA Modeling\n",
    "\n",
    "### What is MDLM?\n",
    "\n",
    "**Masked Diffusion Language Modeling (MDLM)** is a diffusion-based approach adapted for discrete sequences. Unlike continuous diffusion models that add Gaussian noise, MDLM progressively masks tokens:\n",
    "\n",
    "- **Forward process**: Gradually replace tokens with `<mask>` according to a noise schedule\n",
    "- **Reverse process**: Model learns to predict original tokens from masked sequences\n",
    "- **Sampling**: Start from fully masked sequence and iteratively unmask\n",
    "\n",
    "This approach combines the strengths of diffusion models (iterative refinement, stable training) with the discrete nature of DNA sequences.\n",
    "\n",
    "### Why Conditional Generation?\n",
    "\n",
    "We want to generate enhancer sequences with **specific properties**:\n",
    "- **Activity level**: Control how \"strong\" the enhancer should be (low → high activity)\n",
    "- **Sequence context**: Provided by the input sequences, allow the enhancer design to consider promoter-enhancer interaction.\n",
    "\n",
    "By conditioning on these attributes, we can generate enhancers that are both biologically realistic and functionally targeted.\n",
    "\n",
    "---\n",
    "\n",
    "## 📊 Overview\n",
    "\n",
    "We fine-tune **NTv3_tracks** to generate *Drosophila* enhancer sequences conditioned on:\n",
    "\n",
    "| Condition | Source | Training |\n",
    "|-----------|--------|----------|\n",
    "| **Species** | NTv3_tracks species embeddings | Frozen |\n",
    "| **Activity Level** | New embeddings (bins 0-4) | Trainable |\n",
    "\n",
    "The training uses MDLM-style diffusion with:\n",
    "- **Linear noise schedule** for sequence masking (uniform sampling of noise levels)\n",
    "- **Classifier-free guidance** through activity dropout (10% of samples)\n",
    "- **Exponential Moving Average (EMA)** for stable inference\n",
    "\n",
    "[Notebook 06](https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_tutorials/06_plant_promoter_generation.ipynb) is the other half of this workflow: it generates sequences with a released model trained by exactly this recipe, on a plant promoter dataset.\n",
    "\n",
    "---\n",
    "\n",
    "## 📚 Notebook Structure\n",
    "\n",
    "1. **Setup**: Load data and base model from HuggingFace\n",
    "2. **Model Architecture**: Add activity conditioning to NTv3\n",
    "3. **Training Components**: MDLM masking, loss functions, EMA\n",
    "4. **Training Loop**: Full training pipeline with LR scheduling\n",
    "5. **Dataset**: Prepare enhancer sequences with activity labels\n",
    "6. **Run Training**: Execute training on enhancer dataset\n",
    "\n",
    "---\n",
    "\n",
    "## 💻 Hardware Requirements\n",
    "`NTv3_tracks` has ≈700M parameters, and the whole model is fine-tuned, with an EMA copy of it kept alongside. At the batch size of 32 used here, training peaks at **44.8 GiB of GPU memory**, measured on an H100: it needs an **80 GB GPU** (H100, or A100 80 GB) and does **not** fit a 40 GB card. Lower `batch_size` for a smaller GPU.\n",
    "The hyperparameters are set for a full training run: at ≈0.76 s per step on an H100, 500,000 steps is about **four days**. For a quick demo, lower `num_steps` (and `warmup_steps` and `log_every` with it).\n",
    "\n",
    "Because of that, the cells below ship **without the outputs of a full run**. Only the cheap ones are filled in: the dataset size and the parameter counts. Run the notebook to get the rest.\n",
    "\n",
    "> 📝 **Note for Google Colab users:** This notebook is compatible with Colab! Enable a GPU: Runtime → Change runtime type → GPU.\n",
    "> The Colab T4 is far too small for this: see the memory figure above."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5df997b",
   "metadata": {},
   "source": [
    "## 0) 📦 Setup\n",
    "\n",
    "First, we install and import the required libraries and authenticate with HuggingFace to access the pre-trained model and dataset.\n",
    "\n",
    "### Dependencies\n",
    "\n",
    "| Library | Purpose |\n",
    "|---------|---------|\n",
    "| `torch` | Deep learning framework |\n",
    "| `transformers` | Load the NTv3_tracks model and tokenizer |\n",
    "| `datasets` | Load the enhancer generation dataset from HuggingFace |\n",
    "| `huggingface_hub` | Authentication for gated models |"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51af2d66",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Colab ships torch; reinstalling it there breaks the runtime. Elsewhere, also install torch\n",
    "!pip -q install --upgrade 'transformers>=5.17,<6' datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58abcba5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The NTv3 models are gated: request access once on any model page, then log in.\n",
    "from huggingface_hub import get_token, notebook_login\n",
    "\n",
    "if get_token() is None:\n",
    "    notebook_login()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c504c424",
   "metadata": {},
   "outputs": [],
   "source": [
    "import copy\n",
    "from typing import Any\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from datasets import load_dataset\n",
    "from torch.optim import AdamW\n",
    "from torch.optim.lr_scheduler import LambdaLR\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "from tqdm.auto import tqdm\n",
    "from transformers import AutoConfig, AutoModel, AutoTokenizer, PreTrainedTokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "46692f2e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset size: 804592 samples\n"
     ]
    }
   ],
   "source": [
    "# Load training dataset and tokenizer\n",
    "dataset = load_dataset(\"InstaDeepAI/NTv3_enhancer_generation\", split=\"train\")\n",
    "print(f\"Dataset size: {len(dataset)} samples\")\n",
    "\n",
    "model_name = \"InstaDeepAI/NTv3_tracks_1mb\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8797d6f6",
   "metadata": {},
   "source": [
    "### 📊 Understanding the Training Data\n",
    "\n",
    "The dataset `InstaDeepAI/NTv3_enhancer_generation` contains enhancer sequences with measured activity levels from STARR-seq experiments:\n",
    "\n",
    "| Field | Description |\n",
    "|-------|-------------|\n",
    "| `full_sequence` | Complete 4,096 bp construct sequence (promoter + enhancer + reporter) |\n",
    "| `enhancer_sequence` | The 249 bp enhancer inserted in the construct |\n",
    "| `discrete_label` | Activity bin (0-4): 0 = low activity, 4 = high activity |\n",
    "| `activity` | Measured activity, before binning |\n",
    "| `promoter` | Promoter type used (DSCP or RpS12) |\n",
    "\n",
    "**Why Activity Bins?**\n",
    "Continuous activity values are discretized into 5 bins for conditioning. This makes the conditioning signal more robust and enables discrete classifier-free guidance during generation.\n",
    "\n",
    "**Enhancer Region Masking**\n",
    "During training, we only mask the enhancer region (not the promoter/reporter), as this is the region we want the model to learn to generate. The promoter context is provided as input but not predicted.\n",
    "\n",
    "---\n",
    "\n",
    "## 1) 🏗️ Model Architecture\n",
    "\n",
    "### How NTv3 Conditioning Works\n",
    "\n",
    "NTv3 supports **discrete conditioning** through learned embedding tables. For each condition:\n",
    "1. An embedding table maps condition values (e.g., species ID, activity bin) to vectors\n",
    "2. These vectors modulate the model's hidden states through FiLM layers (a learned scale and shift) after its normalization layers\n",
    "3. A prediction head predicts the condition from the sequence\n",
    "\n",
    "`NTv3_tracks` is conditioned on species in its track-prediction head only: its trunk, which the language-model head reads, is not conditioned. To generate, the conditions have to reach that trunk.\n",
    "\n",
    "### Extending with New Conditions\n",
    "\n",
    "We extend the model with a **new activity condition**, and switch on conditioning in the trunk:\n",
    "- **Species** (index 0): Uses the `NTv3_tracks` embeddings → **frozen**\n",
    "- **Activity** (index 1): New embeddings initialized randomly → **trainable**\n",
    "- **Trunk FiLM layers**: New, one per condition, **initialized at zero**, so the extended model starts out predicting exactly what `NTv3_tracks` predicts\n",
    "- **Track head**: Dropped, generation only needs the language-model head\n",
    "\n",
    "All the sequences here are from *Drosophila melanogaster*, so the species condition is constant: its FiLM layers are frozen with it and stay at zero. It is the activity condition the model learns to follow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8fd164f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_generative_config(base_model_name: str, num_activity_bins: int = 5) -> AutoConfig:\n",
    "    \"\"\"NTv3_tracks config with an activity condition and FiLM conditioning in the trunk.\"\"\"\n",
    "    config = AutoConfig.from_pretrained(base_model_name, trust_remote_code=True)\n",
    "    ntv3 = dict(config.ntv3)\n",
    "    cond_dim = ntv3[\"species_embed_dim\"]\n",
    "    ntv3.update(\n",
    "        conditions_names=[\"species\", \"activity\"],\n",
    "        conditions_vocab_size=[ntv3[\"conditions_vocab_size\"][0], num_activity_bins],\n",
    "        # One FiLM input per condition, which switches trunk conditioning on\n",
    "        conditions_dims=[cond_dim, cond_dim],\n",
    "        enable_trunk_conditioning=True,\n",
    "        # Generation only needs the language-model head\n",
    "        number_bigwig_tracks=0,\n",
    "        bigwigs_per_file_assembly={},\n",
    "    )\n",
    "    config.ntv3 = ntv3\n",
    "    return config\n",
    "\n",
    "\n",
    "def zero_film_layers(model: nn.Module) -> None:\n",
    "    \"\"\"Zeroes every FiLM layer, so the new conditions start as a no-op.\"\"\"\n",
    "    for name, module in model.named_modules():\n",
    "        if name.endswith(\"modulation_layers\"):\n",
    "            for layer in module:\n",
    "                nn.init.zeros_(layer.weight)\n",
    "                nn.init.zeros_(layer.bias)\n",
    "\n",
    "\n",
    "def build_generative_model(base_model_name: str, num_activity_bins: int = 5) -> nn.Module:\n",
    "    \"\"\"NTv3 with activity conditioning, initialised from the NTv3_tracks weights.\"\"\"\n",
    "    config = build_generative_config(base_model_name, num_activity_bins)\n",
    "    model = AutoModel.from_config(config, trust_remote_code=True)\n",
    "    base_model = AutoModel.from_pretrained(base_model_name, trust_remote_code=True)\n",
    "    missing, unexpected = model.load_state_dict(base_model.state_dict(), strict=False)\n",
    "    del base_model\n",
    "\n",
    "    # Only the new conditioning layers are left to initialise, and only the track head is dropped\n",
    "    new_layers = (\"modulation_layers\", \"conditions_embed_layers.1.\", \"conditions_heads.1.\")\n",
    "    assert all(any(key in k for key in new_layers) for k in missing), missing\n",
    "    assert all(k.startswith(\"model.bigwig_head.\") for k in unexpected), unexpected\n",
    "    zero_film_layers(model)\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6f5143f5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Trainable parameters: 709,345,544 / 715,370,033 (99.2%)\n"
     ]
    }
   ],
   "source": [
    "# Build the generative model and freeze the species conditioning\n",
    "generation_model = build_generative_model(model_name, num_activity_bins=5)\n",
    "generation_config = generation_model.config\n",
    "\n",
    "# Freeze species conditioning layers (index 0) - only train activity (index 1)\n",
    "species_params = (\"modulation_layers.0.\", \"conditions_embed_layers.0.\", \"conditions_heads.0.\")\n",
    "for name, param in generation_model.named_parameters():\n",
    "    if any(x in name for x in species_params):\n",
    "        param.requires_grad = False\n",
    "\n",
    "trainable_params = sum(p.numel() for p in generation_model.parameters() if p.requires_grad)\n",
    "total_params = sum(p.numel() for p in generation_model.parameters())\n",
    "print(f\"Trainable parameters: {trainable_params:,} / {total_params:,} ({100*trainable_params/total_params:.1f}%)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0bc4e00d",
   "metadata": {},
   "source": [
    "## 2) 🔧 Training Components\n",
    "\n",
    "This section defines the core training components for MDLM training.\n",
    "\n",
    "---\n",
    "\n",
    "### MDLM Training Strategy\n",
    "\n",
    "The model is trained to predict masked tokens using:\n",
    "\n",
    "| Component | Strategy |\n",
    "|-----------|----------|\n",
    "| **Sequence loss** | MDLM-style masking with linear noise schedule, weighted CE loss |\n",
    "| **Activity loss** | Dropout-style masking (10%), CE loss on masked samples only |\n",
    "| **Species** | Frozen (input conditioning only, no loss) |\n",
    "\n",
    "> 💡 **Key insight:** Activity is masked during training so the model learns to predict it from the sequence, also enabling classifier-free guidance at inference. The species logits are not trained here and should not be used.\n",
    "\n",
    "---\n",
    "\n",
    "### The Noise Schedule\n",
    "\n",
    "In MDLM, the **noise level** $t \\in (0, 1]$ determines what fraction of tokens to mask:\n",
    "- $t = 0$: No masking (original sequence)\n",
    "- $t = 1$: All tokens masked\n",
    "\n",
    "During training, we sample $t$ uniformly for each sequence. The loss is weighted by $1/t$ (MDLM weighting) to account for the varying difficulty of different noise levels.\n",
    "\n",
    "### Classifier-Free Guidance (CFG)\n",
    "\n",
    "To enable CFG at inference time, we need the model to work both **with** and **without** the activity condition:\n",
    "- During training: Randomly drop the activity condition with 10% probability\n",
    "- The model learns: $P(\\text{sequence})$ and $P(\\text{sequence} | \\text{activity})$\n",
    "- At inference: Use both predictions to amplify the conditioning signal\n",
    "\n",
    "---\n",
    "\n",
    "### Masking Function\n",
    "\n",
    "The `mask_by_noise_level` function implements the forward diffusion process:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d17ec297",
   "metadata": {},
   "outputs": [],
   "source": [
    "def mask_by_noise_level(\n",
    "    tokens: torch.Tensor,\n",
    "    noise_level: torch.Tensor,\n",
    "    mask_token_id: int,\n",
    "    pad_token_id: int,\n",
    "    mask_vector: torch.Tensor | None = None,\n",
    ") -> tuple[torch.Tensor, torch.Tensor]:\n",
    "    \"\"\"\n",
    "    Mask tokens with probability = noise_level (MDLM diffusion-style).\n",
    "\n",
    "    Returns:\n",
    "        masked_tokens: (B, L) with some tokens replaced by mask_token_id\n",
    "        targets: (B, L) original tokens where masked, pad_token_id elsewhere\n",
    "    \"\"\"\n",
    "    B, L = tokens.shape\n",
    "\n",
    "    # Mask each position with probability = noise_level\n",
    "    mask = torch.rand(B, L, device=tokens.device) < noise_level.view(-1, 1)\n",
    "    mask = mask & (tokens != pad_token_id)  # Don't mask padding\n",
    "\n",
    "    # Only mask positions specified by mask_vector (if provided)\n",
    "    if mask_vector is not None:\n",
    "        mask = mask & mask_vector\n",
    "\n",
    "    masked_tokens = torch.where(mask, mask_token_id, tokens)\n",
    "    targets = torch.where(mask, tokens, pad_token_id)\n",
    "    return masked_tokens, targets\n",
    "\n",
    "\n",
    "def compute_sequence_loss(\n",
    "    logits: torch.Tensor,\n",
    "    targets: torch.Tensor,\n",
    "    weights: torch.Tensor,\n",
    "    pad_token_id: int,\n",
    "    input_tokens: torch.Tensor,\n",
    ") -> torch.Tensor:\n",
    "    \"\"\"MDLM weighted cross-entropy loss, normalized by all non-padding tokens.\"\"\"\n",
    "    B, L, V = logits.shape\n",
    "    sequence_mask = targets != pad_token_id\n",
    "    loss_denominator_mask = input_tokens != pad_token_id\n",
    "\n",
    "    ce = F.cross_entropy(\n",
    "        logits.reshape(-1, V),\n",
    "        targets.reshape(-1),\n",
    "        reduction=\"none\",\n",
    "        ignore_index=pad_token_id,\n",
    "    ).view(B, L)\n",
    "\n",
    "    weighted_ce = ce * weights.unsqueeze(1)\n",
    "    running_sum_loss = (weighted_ce * sequence_mask).sum(dim=1)\n",
    "    per_sample_loss = running_sum_loss / loss_denominator_mask.sum(dim=1).clamp(min=1)\n",
    "    return per_sample_loss.mean()\n",
    "\n",
    "\n",
    "def compute_activity_loss(\n",
    "    logits: torch.Tensor,\n",
    "    targets: torch.Tensor,\n",
    "    mask: torch.Tensor,\n",
    ") -> torch.Tensor:\n",
    "    \"\"\"Cross-entropy loss for activity prediction (only on masked samples).\"\"\"\n",
    "    if not mask.any():\n",
    "        return torch.tensor(0.0, device=logits.device)\n",
    "    return F.cross_entropy(logits[mask], targets[mask])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c54481cb",
   "metadata": {},
   "source": [
    "### Exponential Moving Average (EMA)\n",
    "\n",
    "EMA maintains a smoothed version of model weights throughout training:\n",
    "\n",
    "$$\\theta_{\\text{EMA}} = \\beta \\cdot \\theta_{\\text{EMA}} + (1 - \\beta) \\cdot \\theta_{\\text{current}}$$\n",
    "\n",
    "where $\\beta = 0.999$ (decay rate).\n",
    "\n",
    "**Why use EMA for generation?**\n",
    "- Training weights can be noisy due to stochastic gradient descent\n",
    "- EMA provides an \"averaged\" model that's more stable for inference\n",
    "- Particularly important for generative models where output quality matters\n",
    "- Adapted from existing literature\n",
    "\n",
    "At inference time, we use the **EMA model** instead of the training model. This has a large impact on the sequence generation quality.\n",
    "\n",
    "NTv3's normalization layers also keep running statistics (buffers, not parameters), updated in training mode. The EMA model copies them from the training model at each step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d982575",
   "metadata": {},
   "outputs": [],
   "source": [
    "class EMAModel:\n",
    "    \"\"\"\n",
    "    Exponential Moving Average of model parameters.\n",
    "    Formula: ema_param = decay * ema_param + (1 - decay) * current_param\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, model: nn.Module, decay: float = 0.999):\n",
    "        self.decay = decay\n",
    "        self.ema_model = copy.deepcopy(model)\n",
    "        self.ema_model.eval()\n",
    "        for param in self.ema_model.parameters():\n",
    "            param.requires_grad = False\n",
    "\n",
    "    @torch.no_grad()\n",
    "    def update(self, model: nn.Module):\n",
    "        \"\"\"Update EMA parameters after each training step.\"\"\"\n",
    "        for ema_param, model_param in zip(self.ema_model.parameters(), model.parameters()):\n",
    "            ema_param.data.mul_(self.decay).add_(model_param.data, alpha=1.0 - self.decay)\n",
    "        # Running normalization statistics are copied, not averaged\n",
    "        for ema_buffer, model_buffer in zip(self.ema_model.buffers(), model.buffers()):\n",
    "            ema_buffer.copy_(model_buffer)\n",
    "\n",
    "    def get_model(self) -> nn.Module:\n",
    "        return self.ema_model\n",
    "\n",
    "    def state_dict(self) -> dict[str, Any]:\n",
    "        return self.ema_model.state_dict()\n",
    "\n",
    "    def load_state_dict(self, state_dict: dict[str, Any]):\n",
    "        self.ema_model.load_state_dict(state_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e84a3c8c",
   "metadata": {},
   "source": [
    "### Training Step\n",
    "\n",
    "The `train_step` function implements a single MDLM training iteration:\n",
    "\n",
    "1. **Sample noise level**: Random $t \\sim \\text{Uniform}(\\epsilon, 1)$ for each sequence\n",
    "2. **Mask tokens**: Replace tokens with `<mask>` with probability $t$\n",
    "3. **Apply condition dropout**: Mask activity with 10% probability (for CFG)\n",
    "4. **Forward pass**: Model predicts original tokens and activity\n",
    "5. **Compute losses**: Weighted sequence loss + activity prediction loss\n",
    "6. **Track metrics**: Accuracy for both sequence and activity predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61e1b48d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_step(\n",
    "    model: nn.Module,\n",
    "    tokens: torch.Tensor,\n",
    "    activity: torch.Tensor,\n",
    "    species_ids: torch.Tensor,\n",
    "    mask_vector: torch.Tensor | None,\n",
    "    mask_token_id: int,\n",
    "    pad_token_id: int,\n",
    "    activity_masking_prob: float = 0.1,\n",
    "    activity_loss_weight: float = 1.0,\n",
    ") -> dict[str, Any]:\n",
    "    \"\"\"Single MDLM training step with sequence and activity losses.\"\"\"\n",
    "    B = tokens.shape[0]\n",
    "    device = tokens.device\n",
    "\n",
    "    # 1. Sample noise level (linear schedule, avoid t=0 for stability)\n",
    "    eps = 0.001\n",
    "    t = torch.rand(B, device=device) * (1 - eps) + eps\n",
    "    seq_weights = (1.0 / t).clamp(max=100.0)  # MDLM weight = 1/t\n",
    "\n",
    "    # 2. Mask sequence tokens\n",
    "    masked_tokens, seq_targets = mask_by_noise_level(\n",
    "        tokens, t, mask_token_id, pad_token_id, mask_vector\n",
    "    )\n",
    "\n",
    "    # 3. Condition masking (dropout for classifier-free guidance)\n",
    "    species_mask = torch.ones(B, dtype=torch.bool, device=device)  # Always use species\n",
    "    activity_mask = torch.rand(B, device=device) >= activity_masking_prob  # 10% dropout\n",
    "\n",
    "    # 4. Forward pass\n",
    "    outputs = model(\n",
    "        masked_tokens,\n",
    "        conditions_tokens=[species_ids, activity],\n",
    "        conditions_masks=[species_mask, activity_mask],\n",
    "    )\n",
    "\n",
    "    # 5. Compute losses\n",
    "    seq_loss = compute_sequence_loss(\n",
    "        outputs[\"logits\"], seq_targets, seq_weights, pad_token_id, tokens\n",
    "    )\n",
    "    activity_loss = compute_activity_loss(\n",
    "        outputs[\"activity_logits\"], activity, ~activity_mask\n",
    "    )\n",
    "    total_loss = seq_loss + activity_loss_weight * activity_loss\n",
    "\n",
    "    # 6. Compute metrics\n",
    "    seq_mask = seq_targets != pad_token_id\n",
    "    masked_tokens_count = seq_mask.sum().item()\n",
    "    total_tokens_count = (tokens != pad_token_id).sum().item()\n",
    "    masked_activity_count = (~activity_mask).sum().item()\n",
    "\n",
    "    with torch.no_grad():\n",
    "        seq_preds = outputs[\"logits\"].argmax(dim=-1)\n",
    "        seq_correct = (seq_preds == seq_targets) & seq_mask\n",
    "        seq_accuracy = seq_correct.sum().item() / max(1, masked_tokens_count)\n",
    "\n",
    "        if masked_activity_count > 0:\n",
    "            act_preds = outputs[\"activity_logits\"].argmax(dim=-1)\n",
    "            act_correct = (act_preds == activity) & (~activity_mask)\n",
    "            act_accuracy = act_correct.sum().item() / masked_activity_count\n",
    "        else:\n",
    "            act_accuracy = 0.0\n",
    "\n",
    "    return {\n",
    "        \"loss\": total_loss,\n",
    "        \"seq_loss\": seq_loss.item(),\n",
    "        \"activity_loss\": activity_loss.item(),\n",
    "        \"seq_accuracy\": seq_accuracy,\n",
    "        \"act_accuracy\": act_accuracy,\n",
    "        \"masked_tokens\": masked_tokens_count,\n",
    "        \"total_tokens\": total_tokens_count,\n",
    "        \"masked_activity\": masked_activity_count,\n",
    "        \"batch_size\": B,\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b48b24a7",
   "metadata": {},
   "source": [
    "## 3) 🔄 Training Loop\n",
    "\n",
    "This section implements the main training loop with learning rate scheduling.\n",
    "\n",
    "---\n",
    "\n",
    "### Learning Rate Schedule\n",
    "\n",
    "We use a **modified square-root decay** schedule with linear warmup:\n",
    "\n",
    "1. **Warmup phase**: Linear increase from `initial_lr` to `peak_lr`\n",
    "2. **Decay phase**: Power-law decay: $\\text{LR}(t) = \\text{peak\\_lr} \\times \\left(\\frac{\\text{warmup\\_steps}}{t}\\right)^\\alpha$\n",
    "\n",
    "This schedule provides:\n",
    "- Stable early training (warmup prevents large initial updates)\n",
    "- Gradual refinement (decay allows fine-tuning without overshooting)\n",
    "\n",
    "### Training Configuration\n",
    "\n",
    "| Parameter | Recommended Value | Description |\n",
    "|-----------|------------------|-------------|\n",
    "| `num_steps` | 500,000+ | Dataset dependent convergence time |\n",
    "| `lr` | 1e-4 | Peak learning rate |\n",
    "| `warmup_steps` | 25,000 | ≈5% of total steps |\n",
    "| `ema_decay` | 0.999 | Higher = smoother EMA |\n",
    "| `activity_masking_prob` | 0.1 | Activity conditioning dropout rate |"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67f9a841",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_modified_sqrt_decay_scheduler(\n",
    "    optimizer: torch.optim.Optimizer,\n",
    "    warmup_steps: int,\n",
    "    total_steps: int,\n",
    "    initial_lr: float,\n",
    "    peak_lr: float,\n",
    "    final_ratio: float = 0.5,\n",
    ") -> LambdaLR:\n",
    "    \"\"\"Linear warmup → power-law decay scheduler.\"\"\"\n",
    "    alpha = np.log(1 / final_ratio) / np.log(total_steps / max(1, warmup_steps))\n",
    "\n",
    "    def lr_lambda(step):\n",
    "        if step < warmup_steps:\n",
    "            return (initial_lr + (peak_lr - initial_lr) * step / max(1, warmup_steps)) / peak_lr\n",
    "        return (warmup_steps / (step + 1)) ** alpha\n",
    "\n",
    "    return LambdaLR(optimizer, lr_lambda)\n",
    "\n",
    "\n",
    "def train_mdlm(\n",
    "    model: nn.Module,\n",
    "    train_dataloader: DataLoader,\n",
    "    tokenizer: PreTrainedTokenizer,\n",
    "    num_steps: int = 10000,\n",
    "    lr: float = 5e-5,\n",
    "    initial_lr: float = 0.0,\n",
    "    device: str = \"cuda\",\n",
    "    activity_masking_prob: float = 0.1,\n",
    "    activity_loss_weight: float = 1.0,\n",
    "    warmup_steps: int = 0,\n",
    "    scheduler_type: str | None = None,\n",
    "    final_lr_ratio: float = 0.5,\n",
    "    log_every: int = 100,\n",
    "    ema_decay: float = 0.999,\n",
    "    use_ema: bool = True,\n",
    ") -> tuple[nn.Module, nn.Module]:\n",
    "    \"\"\"MDLM training loop. Returns (trained_model, ema_model).\"\"\"\n",
    "    model = model.to(device)\n",
    "    optimizer = AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=lr)\n",
    "\n",
    "    ema = EMAModel(model, decay=ema_decay) if use_ema else None\n",
    "    scheduler = get_modified_sqrt_decay_scheduler(\n",
    "        optimizer, warmup_steps, num_steps, initial_lr, lr, final_lr_ratio\n",
    "    ) if scheduler_type == \"modified_sqrt\" else None\n",
    "\n",
    "    mask_token_id = tokenizer.mask_token_id\n",
    "    pad_token_id = tokenizer.pad_token_id\n",
    "\n",
    "    model.train()\n",
    "    train_iter = iter(train_dataloader)\n",
    "    metrics = {\"loss\": 0, \"seq\": 0, \"act\": 0, \"seq_acc\": 0, \"act_acc\": 0,\n",
    "               \"masked\": 0, \"total\": 0, \"masked_act\": 0, \"batch\": 0}\n",
    "\n",
    "    pbar = tqdm(range(1, num_steps + 1), desc=\"Training\")\n",
    "    for step in pbar:\n",
    "        try:\n",
    "            batch = next(train_iter)\n",
    "        except StopIteration:\n",
    "            train_iter = iter(train_dataloader)\n",
    "            batch = next(train_iter)\n",
    "\n",
    "        tokens = batch[\"input_ids\"].to(device)\n",
    "        species_ids = batch[\"species_ids\"].to(device)\n",
    "        activity = batch[\"activity_labels\"].to(device)\n",
    "        mask_vector = batch[\"masking_vector\"].to(device)\n",
    "\n",
    "        losses = train_step(\n",
    "            model, tokens, activity, species_ids, mask_vector,\n",
    "            mask_token_id, pad_token_id, activity_masking_prob, activity_loss_weight\n",
    "        )\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        losses[\"loss\"].backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        if ema: ema.update(model)\n",
    "        if scheduler: scheduler.step()\n",
    "\n",
    "        # Track metrics\n",
    "        metrics[\"loss\"] += losses[\"loss\"].item()\n",
    "        metrics[\"seq\"] += losses[\"seq_loss\"]\n",
    "        metrics[\"act\"] += losses[\"activity_loss\"]\n",
    "        metrics[\"seq_acc\"] += losses[\"seq_accuracy\"]\n",
    "        metrics[\"act_acc\"] += losses[\"act_accuracy\"] * losses[\"masked_activity\"]\n",
    "        metrics[\"masked\"] += losses[\"masked_tokens\"]\n",
    "        metrics[\"total\"] += losses[\"total_tokens\"]\n",
    "        metrics[\"masked_act\"] += losses[\"masked_activity\"]\n",
    "        metrics[\"batch\"] += losses[\"batch_size\"]\n",
    "\n",
    "        if step % log_every == 0:\n",
    "            avg_act_acc = metrics[\"act_acc\"] / max(1, metrics[\"masked_act\"])\n",
    "            tqdm.write(\n",
    "                f\"Step {step:5d} | loss: {metrics['loss']/log_every:.4f} | \"\n",
    "                f\"seq_loss: {metrics['seq']/log_every:.4f} | seq_acc: {metrics['seq_acc']/log_every:.3f} | \"\n",
    "                f\"act_loss: {metrics['act']/log_every:.4f} | act_acc: {avg_act_acc:.3f} | \"\n",
    "                f\"lr: {optimizer.param_groups[0]['lr']:.2e}\"\n",
    "            )\n",
    "            metrics = {k: 0 for k in metrics}\n",
    "\n",
    "    return (model, ema.get_model()) if use_ema else (model, model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "263802a8",
   "metadata": {},
   "source": [
    "## 4) 📊 Dataset\n",
    "\n",
    "This section defines the PyTorch Dataset class for loading enhancer sequences.\n",
    "\n",
    "---\n",
    "\n",
    "### Dataset Structure\n",
    "\n",
    "The `EnhancerDataset` class handles:\n",
    "1. **Tokenization**: Convert DNA sequences to token IDs using the NTv3 tokenizer\n",
    "2. **Activity labels**: Extract discrete activity bins (0-4)\n",
    "3. **Masking vectors**: Create binary masks indicating which positions to train on\n",
    "\n",
    "### Enhancer Insertion Positions\n",
    "\n",
    "Different promoter types have the enhancer inserted at different positions:\n",
    "\n",
    "| Promoter | Enhancer Start | Enhancer End | Length |\n",
    "|----------|---------------|--------------|--------|\n",
    "| DSCP | 1018 | 1267 | 249 bp |\n",
    "| RpS12 | 968 | 1217 | 249 bp |\n",
    "\n",
    "The `masking_vector` ensures we only apply MDLM loss to the enhancer region, not the flanking promoter/reporter sequences.\n",
    "\n",
    "The constructs are 4,096 bp long, a multiple of 128 as NTv3 requires, and in upper case: `NTv3_tracks` reads soft-masked (lower-case) bases as different tokens."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95e60646",
   "metadata": {},
   "outputs": [],
   "source": [
    "class EnhancerDataset(Dataset):\n",
    "    \"\"\"Dataset for enhancer generation training.\"\"\"\n",
    "\n",
    "    # Enhancer insertion positions for each promoter type\n",
    "    PROMOTER_INSERTIONS = {\n",
    "        \"DSCP\": [1018, 1267],\n",
    "        \"RpS12\": [968, 1217],\n",
    "    }\n",
    "\n",
    "    def __init__(self, hf_dataset, tokenizer: PreTrainedTokenizer, species_id: int):\n",
    "        self.data = hf_dataset\n",
    "        self.tokenizer = tokenizer\n",
    "        self.species_id = species_id\n",
    "\n",
    "    def __len__(self) -> int:\n",
    "        return len(self.data)\n",
    "\n",
    "    def __getitem__(self, idx: int) -> dict[str, Any]:\n",
    "        item = self.data[idx]\n",
    "\n",
    "        # Tokenize sequence\n",
    "        tokens = self.tokenizer(\n",
    "            item[\"full_sequence\"], add_special_tokens=False, return_tensors=\"pt\"\n",
    "        )[\"input_ids\"].squeeze(0)\n",
    "\n",
    "        # Activity label (0-4)\n",
    "        activity = torch.tensor(item.get(\"discrete_label\", 0), dtype=torch.long)\n",
    "\n",
    "        # Create masking vector for enhancer region only\n",
    "        promoter_type = item.get(\"promoter\")\n",
    "        mask_interval = self.PROMOTER_INSERTIONS[promoter_type]\n",
    "        masking_vector = torch.zeros(len(tokens), dtype=torch.bool)\n",
    "        masking_vector[mask_interval[0]:mask_interval[1]] = True\n",
    "\n",
    "        return {\n",
    "            \"input_ids\": tokens,\n",
    "            \"species_ids\": self.species_id,\n",
    "            \"activity_labels\": activity,\n",
    "            \"masking_vector\": masking_vector,\n",
    "        }"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e119379f",
   "metadata": {},
   "source": [
    "## 5) 🚀 Run Training\n",
    "\n",
    "Now we instantiate the dataset and run the training loop.\n",
    "\n",
    "---\n",
    "\n",
    "### Training Parameters\n",
    "\n",
    "The configuration below is set for a **full training run**. For a quick demo, decrease `num_steps` and adjust the batch size to your GPU memory:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "41f2a2ea",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training on 804592 samples\n"
     ]
    }
   ],
   "source": [
    "# Create dataset and dataloader\n",
    "species_id = generation_config.track_metadata[\"species_vocabulary\"][\"drosophila\"]\n",
    "train_dataset = EnhancerDataset(dataset, tokenizer, species_id)\n",
    "train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=0)\n",
    "\n",
    "print(f\"Training on {len(train_dataset)} samples\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44497419",
   "metadata": {},
   "source": [
    "### Understanding Training Metrics\n",
    "\n",
    "During training, monitor these metrics:\n",
    "\n",
    "| Metric | What to Watch |\n",
    "|--------|---------------|\n",
    "| `loss` | Total loss (seq + activity) |\n",
    "| `seq_loss` | Sequence prediction loss |\n",
    "| `seq_acc` | Token prediction accuracy |\n",
    "| `act_loss` | Activity prediction loss |\n",
    "| `act_acc` | Activity classification accuracy |\n",
    "\n",
    "**Interpreting Results:**\n",
    "- **seq_acc** measures how well the model predicts masked DNA tokens\n",
    "- **act_acc** measures whether the model learns the activity-sequence relationship. Read it against the **majority-class baseline, not against zero**: bin 1 holds ≈52% of this dataset, so an `act_acc` at or below ≈0.52 means the head has not learned anything yet. Early in training that is what you should expect to see.\n",
    "- The activity conditioning is what makes generation steerable, so it is the metric to watch over a long run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6098e7ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train the model\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "print(f\"Training on: {device}\")\n",
    "\n",
    "trained_model, ema_model = train_mdlm(\n",
    "    model=generation_model,\n",
    "    train_dataloader=train_loader,\n",
    "    tokenizer=tokenizer,\n",
    "    num_steps=500_000,              # Decrease for a quick demo\n",
    "    lr=1e-4,                        # Peak LR\n",
    "    initial_lr=5e-5,                # Starting LR for warmup\n",
    "    device=device,\n",
    "    activity_masking_prob=0.1,      # 10% activity dropout for CFG\n",
    "    activity_loss_weight=0.1,\n",
    "    warmup_steps=25_000,            # Decrease with num_steps\n",
    "    scheduler_type=\"modified_sqrt\",\n",
    "    final_lr_ratio=0.5,\n",
    "    log_every=2000,                 # Decrease for a quick demo\n",
    "    ema_decay=0.999,\n",
    "    use_ema=True,\n",
    ")\n",
    "\n",
    "# Use EMA model for generation (smoother, more stable)\n",
    "generation_model = ema_model\n",
    "print(\"Training complete!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98df6da7",
   "metadata": {},
   "source": [
    "### Saving Checkpoints\n",
    "\n",
    "After training, save both the training model and the EMA model. **Use the EMA model for generation** as it produces smoother, more stable outputs.\n",
    "\n",
    "To generate with the model you just trained, take the sampler of [notebook 06](https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_tutorials/06_plant_promoter_generation.ipynb) and point it at `ntv3_generative_finetuned_ema.pt`, with this notebook's conditions and the enhancer construct in place of the plant promoter one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8e7ab96",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save checkpoints\n",
    "torch.save(trained_model.state_dict(), \"ntv3_generative_finetuned.pt\")\n",
    "torch.save(ema_model.state_dict(), \"ntv3_generative_finetuned_ema.pt\")\n",
    "print(\"Checkpoints saved!\")"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}