Replace the 2025 Space with the revised NTv3 release (#4)

- Replace the 2025 Space with the revised NTv3 release (a6af7a757e95fd80e6af34a772ff1e571bb11813)

.gitattributes

@@ -35,6 +35,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 assets/paper_summary.jpg filter=lfs diff=lfs merge=lfs -text
 assets/paper_summary.png filter=lfs diff=lfs merge=lfs -text
+assets/ntv3_overview.png filter=lfs diff=lfs merge=lfs -text
 assets/output_tracks.png filter=lfs diff=lfs merge=lfs -text
 notebooks_pipelines/bigwig_outputs/K562_DNAse.bw filter=lfs diff=lfs merge=lfs -text
 notebooks_pipelines/bigwig_outputs/K562_H3k4me3.bw filter=lfs diff=lfs merge=lfs -text
@@ -45,3 +46,7 @@ notebooks_pipelines/bigwig_outputs/HepG2_H3k4me3.bw filter=lfs diff=lfs merge=lf
 notebooks_pipelines/bigwig_outputs/HepG2_RNA_seq.bw filter=lfs diff=lfs merge=lfs -text
 notebooks_pipelines/bigwig_outputs/K562_CTCF.bw filter=lfs diff=lfs merge=lfs -text
 assets/saliency_example.png filter=lfs diff=lfs merge=lfs -text
+assets/skill_workflow_v6b.png filter=lfs diff=lfs merge=lfs -text
+assets/skill_report_hbe1_cagesal_top_hs.png filter=lfs diff=lfs merge=lfs -text
+assets/skill_report_hbe1_hs4_saliency_cage_r848.png filter=lfs diff=lfs merge=lfs -text
+assets/skill_report_hbe1_4kb_annotation.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,5 +1,5 @@
 ---
-title: NTv3 — Foundation Models for Long-Range Genomics
+title: NTv3 Foundation Models for Genomics
 emoji: 🧬
 colorFrom: indigo
 colorTo: blue
@@ -7,26 +7,71 @@ sdk: static
 pinned: false
 ---
 
-# 🧬 NTv3 — Foundation Models for Long-Range Genomics
+# 🧬 NTv3: Foundation Models for Genomics
 
-This Space is the companion hub for NTv3 checkpoints on the Hugging Face Hub. It provides PyTorch notebooks and minimal examples for inference, sequence-to-function prediction (functional tracks), genome annotation, fine-tuning, model interpretation and sequence generation.
+A genomic foundation model that reads up to 1 Mb of DNA at single-base resolution. One architecture, three
+capabilities: **functional-track prediction**, **gene annotation** and **regulatory sequence design**. Of the
+paper's design models, the plant promoter one is released as an example of the generative models. This Space is the
+companion hub, with runnable PyTorch notebooks.
+
+> **Release status.** This Space is being prepared and items still to be settled are marked `TBC` throughout.
+> Looking for the 2025 models? See [Initial release (2025)](#-initial-release-2025) below.
 
 ## 📖 About NTv3
 
-NTv3 is a multi-species genomic foundation model family that unifies representation learning, functional-track prediction, genome annotation, and controllable sequence generation within a single U-Net-style backbone. It models up to 1 Mb of DNA at single-base resolution, using a conv–Transformer–deconv architecture that efficiently captures both local motifs and long-range regulatory dependencies. NTv3 is first pretrained on ~9T base pairs from the OpenGenome2 corpus spanning >128k species using masked language modeling, and then post-trained with a joint objective on ~16k functional tracks and annotation labels across 24 animal and plant species, enabling state-of-the-art cross-species functional prediction and base-resolution genome annotation.
+NTv3 reads DNA at single-base resolution across up to **1 Mb** of context in one pass. One architecture supports three
+capabilities: **functional-track prediction**, **gene annotation** and **regulatory sequence design**. The backbone is
+U-Net-style: a convolutional tower that reads the sequence, a Transformer tower that carries long-range context, and a
+deconvolutional tower that returns predictions at base resolution.
 
-Beyond prediction, NTv3 can be fine-tuned into a controllable generative model via masked-diffusion language modeling, allowing targeted design of regulatory sequences (for example, enhancers with specified activity and promoter selectivity) that have been validated experimentally.
+They are acquired through a staged recipe, each stage building on the one before:
 
-## 📓 Notebooks
+1. **Pre-training**, self-supervised. Masked-token prediction over unlabeled genomes from the OpenGenome2 corpus -> **NTv3-pretrained**.
+2. **Tracks post-training**, supervised multi-task. 16,262 functional tracks across 15 species -> **NTv3-tracks**.
+3. **Annotation fine-tuning**, supervised single-task. Reference gene annotations -> one **NTv3-annotate** model
+   per clade (mammals, plants, insects).
+4. **Generative post-training**, masked diffusion on MPRA activity datasets -> **NTv3-generative**. The plant
+   promoter model is released as an example; the plant enhancer and fly enhancer models are not.
+
+## 🧩 What it does
+
+| Capability | Detail |
+| --- | --- |
+| Reads any genome | 1 Mb context, nucleotide resolution, base-pair-resolution prediction |
+| Predicts functional tracks | 16,262 tracks, 15 species: 9,603 gene expression, 3,115 histone marks, 2,193 TF binding, 1,351 accessibility |
+| Annotates genes | DNA alone to per-base posteriors over intergenic, intron, CDS and splice states, decoded to gene models on both strands |
+| Designs regulatory DNA | Conditional generation of plant promoters, plant enhancers and fly enhancers, validated by measured activity |
+
+## 🤖 The NTv3 models
+
+Each capability is released as its own checkpoint, all sharing the same architecture and the same pretrained
+representations, at **700M parameters**.
 
-Notebooks live in `./notebooks/`:
+| Model | Stage | What it is for | Repository |
+| --- | --- | --- | --- |
+| NTv3-pretrained | Pretrained | Embeddings, masked-token prediction, your own fine-tuning | [`InstaDeepAI/NTv3_pretrained`](https://huggingface.co/InstaDeepAI/NTv3_pretrained) |
+| NTv3-tracks | Tracks post-training | Functional-track prediction from 16 kb to 1 Mb inputs, and the backbone for downstream fine-tuning | [`InstaDeepAI/NTv3_tracks`](https://huggingface.co/InstaDeepAI/NTv3_tracks) |
+| NTv3-tracks-1mb | Tracks post-training | The 1 Mb specialist behind the paper's results; 1 Mb inputs only | [`InstaDeepAI/NTv3_tracks_1mb`](https://huggingface.co/InstaDeepAI/NTv3_tracks_1mb) |
+| NTv3-annotate | Annotation fine-tuning | *Ab initio* gene prediction, one model per clade | [`_mammals`](https://huggingface.co/InstaDeepAI/NTv3_annot_mammals) / [`_plants`](https://huggingface.co/InstaDeepAI/NTv3_annot_plants) / [`_insects`](https://huggingface.co/InstaDeepAI/NTv3_annot_insects) |
+| NTv3-generative-plant-promoter | Generative post-training | Example generative model used in the paper: designs plant core promoters at a requested strength, and is what notebook 06 runs | [`InstaDeepAI/NTv3_generative_plant_promoter`](https://huggingface.co/InstaDeepAI/NTv3_generative_plant_promoter) |
 
-- 🚀 `00_quickstart_inference.ipynb` — load a checkpoint + run inference
-- 📊 `01_tracks_prediction.ipynb` — sequence → functional tracks (+ plotting)
-- 🏷️ `02_genome_annotation_segmentation.ipynb` — sequence → annotation
-- 🎯 `03_finetune_head.ipynb` — fine-tune on bigwig tracks
-- 🔍 `04_model_interpretation.ipynb` — interpretation of post-trained model
-- 🧪 `05_sequence_generation.ipynb` — fine-tune NTv3 to generate enhancer sequences
+Of the paper's design models, only the plant promoter one is released, as an example; the plant enhancer and fly
+enhancer models are not. It is the model notebook 06 runs, and a worked example of the recipe notebook 05 teaches.
+Read its model card before quoting anything from it: it is not species-aware.
+
+**700M** is the figure quoted throughout, counted from the published weights rather than the internal registry, which
+still calls the pretrained checkpoint 650M. Exactly: 703,349,146 for pretrained, 711,256,680 for tracks,
+703,446,690 for each annotate clade, and 721,379,955 for the plant promoter example, which carries three condition
+tables and conditioning layers in its trunk.
+
+The [NTv3 collection](https://huggingface.co/collections/InstaDeepAI/ntv3) is the canonical index
+of what is downloadable today.
+
+## 🔐 Access
+
+The models are gated: request access once on any of their model pages, and one approval, reviewed manually, covers
+all of them. Then log in with `hf auth login`, or `notebook_login()` in a notebook. The datasets and this Space are
+open.
 
 ## 📦 Install
 
@@ -34,68 +79,158 @@ Notebooks live in `./notebooks/`:
 pip install torch transformers accelerate safetensors huggingface_hub numpy
 ```
 
-## 🤖 Load a pre-trained model
+Works on **transformers 4.57.1 and 5.17.0**, which now agree exactly. Two faults that broke 5.x are fixed: the rotary
+frequencies are rebuilt per call rather than restored from a buffer that 5.x zeroed, and the annotation checkpoints
+store `residual_scale` as the scalar the model declares. Re-measured on chr19:6,700,000-6,716,384 across both
+versions, the predictions are identical to the bit. Older 5.x releases may still carry the zeroed-buffer behaviour,
+which fails silently, so pin `transformers>=4.57.1` and check against the fixtures in each repo if you stray.
+
+Models use custom code, so pass `trust_remote_code=True`. Input sequence length must be a multiple of 128.
+
+## 🚀 Quickstart with NTv3-pretrained
+
+Both checkpoints are registered under `AutoModel`, and the forward takes the token tensor positionally rather than
+the tokenizer dict, so `model(**batch)` will not work. The forward returns a plain `dict`, not a `ModelOutput`, so
+read results with `out["logits"]` rather than `out.logits`.
 
 ```python
-from transformers import AutoTokenizer, AutoModelForMaskedLM
+import torch
+from transformers import AutoModel, AutoTokenizer
 
-repo = "InstaDeepAI/NTv3_650M_pre"
-tok = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
-model = AutoModelForMaskedLM.from_pretrained(repo, trust_remote_code=True)
+model_name = "InstaDeepAI/NTv3_pretrained"
 
-batch = tok(["ATCGNATCG", "ACGT"], add_special_tokens=False, padding=True, pad_to_multiple_of=128, return_tensors="pt")
-out = model(**batch, output_hidden_states=True, output_attentions=True)
+tok = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+model = AutoModel.from_pretrained(model_name, trust_remote_code=True).eval()
 
-print(out.logits.shape)       # (B, L, V = 11)
-print(len(out.hidden_states)) # convs + transformers + deconvs
-print(len(out.attentions))    # equals transformer layers = 12
+# Batch only equal-length sequences: padding is not masked out
+batch = tok(["ACGT" * 32, "TTGA" * 32], add_special_tokens=False, return_tensors="pt")
+
+with torch.no_grad():
+    out = model(batch["input_ids"])
+
+print(out["logits"].shape)      # (B, L, V = 11)
+print(out["embedding"].shape)   # (B, L, 1024)
 ```
 
-## 💻 Pipelines
+Input length must be a multiple of 128, and **padded positions are not masked out of attention**. The same sequence
+run alone and run padded to a longer length gives different values at its real positions, so batch only sequences of
+equal length and run ragged ones separately.
+
+## 📊 NTv3-tracks: functional-track prediction
 
-Here is a quick example of how to use the post-trained NTv3 650M model on a human genomic window.
+The `ntv3-tracks` pipeline takes a genomic region and the tracks you want, fetches the sequence, conditions on the
+species, predicts and plots.
 
 ```python
-from transformers import AutoConfig
-
-model_name = "InstaDeepAI/NTv3_100M"
-
-# Load track prediction pipeline
-cfg = AutoConfig.from_pretrained(model_name, trust_remote_code=True, force_download=True)
-pipe = cfg.load_tracks_pipeline(model_name, device="auto")  # or "cpu"/"cuda"/"mps"
-
-# Run track prediction
-out = pipe(
-    {
-        "chrom": "chr19",
-        "start": 6_700_000,
-        "end": 6_831_072,
-        "species": "human"
-    }
+from transformers import pipeline
+
+ntv3_tracks = pipeline("ntv3-tracks", model="InstaDeepAI/NTv3_tracks", trust_remote_code=True)
+
+out = ntv3_tracks(
+    {"chrom": "chr19", "start": 6_700_000, "end": 6_732_768, "species": "human"},
+    plot=True,
+    tracks_to_plot={"K562 RNA-seq": "ENCSR056HPM", "K562 CTCF": "ENCSR000AKO"},
 )
 
-print(out.bigwig_tracks_logits.shape)   # functional track predictions
-print(out.bed_tracks_logits.shape)      # genome annotation predictions
-print(out.mlm_logits.shape)             # MLM logits: (B, L, V = 11)
+print(out.tracks.shape)        # (12288, 2), the central 37.5 % of the window
+print(out.coordinates)         # (6710240, 6722528), the span actually predicted
+print(out.track("ENCSR000AKO").shape)  # (12288,): one track
+```
+
+Naming the tracks computes only those channels. The head is as wide as the largest species panel (hg38, 7,362
+tracks), so a full forward at 1 Mb materialises about 11 GB, while a named subset is bit-identical and a few MB.
+Species, assemblies and track order all come from the model's own `config.track_metadata`.
+
+For the plain forward, which is where to start if you want to fine-tune rather than predict, `NTv3_tracks` takes a
+species token alongside the sequence, and its nucleotide vocabulary is 16 tokens rather than the 11 of the pretrained
+model, so it reads soft-masked (lowercase) sequence directly. Notebook 01 covers it.
+
+The 2025 `ntv3-interpret` pipeline is not being carried over. Notebook 04 shows saliency and attention with a plain
+`AutoModel` forward instead.
+
+## 🏷️ NTv3-annotate: genome annotation
+
+```python
+import torch
+from transformers import AutoModel, AutoTokenizer
+
+CLADE = "mammals"  # or "plants", or "insects"
+repo = f"InstaDeepAI/NTv3_annot_{CLADE}"
+
+model = AutoModel.from_pretrained(repo, trust_remote_code=True).eval()
+tok = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
+batch = tok("ACGT" * 256, add_special_tokens=False, return_tensors="pt")
+
+# Annotation is not conditioned on species, so all three clades take an empty list
+with torch.no_grad():
+    out = model(batch["input_ids"], conditions_tokens=[])
+
+print(out["bed_tracks_logits"].shape)   # (1, 768, 1, 15); 768 = 1024 * 0.75
 ```
 
-## 🤖 Checkpoints
+The output covers the central 75% of the window, the crop the model is served at. Logits alone are not an annotation:
+notebook 03b decodes them into gene models (GFF3) with the HMM whose specification `config.json` carries under
+`annotation`, and notebook 03c scores an annotation the way the benchmark does.
+
+## 📓 Notebooks
+
+The notebooks cover the workflows end to end, in `notebooks_tutorials/`:
+
+| # | Notebook | File | Status |
+| --- | --- | --- | --- |
+| 00 | 🚀 Quickstart inference | [`00_quickstart_inference.ipynb`](https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_tutorials/00_quickstart_inference.ipynb) | ✅ |
+| 01 | 📊 Tracks prediction | [`01_tracks_prediction.ipynb`](https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_tutorials/01_tracks_prediction.ipynb) | ✅ |
+| 02 | 🎯 Fine-tune NTv3 on bigwig tracks | [`02_fine_tuning_bigwig_tracks.ipynb`](https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_tutorials/02_fine_tuning_bigwig_tracks.ipynb) | ✅ |
+| 03a | 🏷️ Train a gene-annotation model | [`03a_genome_annotation_training.ipynb`](https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_tutorials/03a_genome_annotation_training.ipynb) | ✅ |
+| 03b | 🏷️ From per-base predictions to gene models | [`03b_genome_annotation_decoding.ipynb`](https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_tutorials/03b_genome_annotation_decoding.ipynb) | ✅ |
+| 03c | 🏷️ Evaluate a gene annotation against a reference | [`03c_genome_annotation_evaluation.ipynb`](https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_tutorials/03c_genome_annotation_evaluation.ipynb) | ✅ |
+| 04 | 🔍 Model interpretation | [`04_model_interpretation.ipynb`](https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_tutorials/04_model_interpretation.ipynb) | ✅ |
+| 05 | 🧪 Fine-tuning NTv3 into a diffusion model | [`05_diffusion_fine_tuning.ipynb`](https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_tutorials/05_diffusion_fine_tuning.ipynb) | ✅ |
+| 06 | 🌱 Generating plant promoter sequences | [`06_plant_promoter_generation.ipynb`](https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_tutorials/06_plant_promoter_generation.ipynb) | ✅ |
+
+All run against the published checkpoints.
+
+## 🗄️ Initial release (2025)
+
+The initial-release repositories are **retained unchanged** for reproducibility and backward compatibility. Existing
+model IDs, files and interfaces continue to work. If you are reproducing published results or already have a working
+pipeline, stay on these.
+
+**Pretrained:** `NTv3_8M_pre`, `NTv3_100M_pre`, `NTv3_650M_pre`
 
-**📦 Pre-trained:** `InstaDeepAI/NTv3_8M_pre`, `InstaDeepAI/NTv3_100M_pre`, `InstaDeepAI/NTv3_650M_pre`
+**Post-trained:** `NTv3_100M_post`, `NTv3_650M_post`
 
-**🎯 Post-trained:** `InstaDeepAI/NTv3_100M`, `InstaDeepAI/NTv3_650M`
+**Context-length variants:** `NTv3_8M_pre_8kb`, `NTv3_100M_pre_8kb`, `NTv3_650M_pre_8kb`, `NTv3_100M_post_131kb`,
+`NTv3_650M_post_131kb`
+
+**Architecture ablations:** `NTv3_5downsample_pre`, `NTv3_5downsample_pre_8kb`, `NTv3_5downsample_post`,
+`NTv3_5downsample_post_131kb`
+
+**Generative:** `NTv3_generative`
+
+All under the `InstaDeepAI/` namespace. Note that some earlier materials wrote the post-trained checkpoints as
+`NTv3_100M_pos` / `NTv3_650M_pos`; the correct identifiers end in `_post`.
+
+These models were trained with a *joint* post-training objective over tracks and annotation labels, so numbers quoted
+in their materials do not always match the current paper. Where they differ, the current paper is authoritative.
 
 ## 🔗 Links
 
-- **📄 Paper:** (add link)
-- **💻 JAX research code (GitHub):** [https://github.com/instadeepai/nucleotide-transformer](https://github.com/instadeepai/nucleotide-transformer)
-- **🏆 NTv3 benchmark leaderboard: (add link)**
+- 📄 **Paper:** preprint `TBC`
+- 💻 **JAX model code:** [github.com/instadeepai/nucleotide-transformer](https://github.com/instadeepai/nucleotide-transformer)
+- 🎯 **Model collection:** [NTv3 collection](https://huggingface.co/collections/InstaDeepAI/ntv3)
+- 📚 **Benchmark dataset:** [`InstaDeepAI/NTv3_benchmark_dataset_revised`](https://huggingface.co/datasets/InstaDeepAI/NTv3_benchmark_dataset_revised)
+- 🏷️ **Annotation dataset:** [`InstaDeepAI/NTv3_annot_dataset`](https://huggingface.co/datasets/InstaDeepAI/NTv3_annot_dataset)
+- 🏆 **Benchmark leaderboard:** [`InstaDeepAI/ntv3_benchmark_revised`](https://huggingface.co/spaces/InstaDeepAI/ntv3_benchmark_revised)
+- 🚀 **Interactive demo:** [`InstaDeepAI/ntv3_tracks`](https://huggingface.co/spaces/InstaDeepAI/ntv3_tracks)
 
 ## 📝 Citation
 
+`TBC`. To be added on publication.
+
 ```bibtex
 @article{ntv3,
-  title   = {A foundational model for joint sequence-function multi-species modeling at scale for long-range genomic prediction},
+  title   = {…},
   author  = {…},
   journal = {…},
   year    = {…}
@@ -104,6 +239,8 @@ print(out.mlm_logits.shape)             # MLM logits: (B, L, V = 11)
 
 ## 📜 License
 
-**Code & notebooks in this Space:** (choose and add, e.g., Apache-2.0)
+**Model weights:** released under the InstaDeep Open Model Licence, which permits non-commercial use only and also
+places conditions on the **outputs** of the models. See the `MODEL-LICENSE.md` shipped in each model repository for the
+authoritative terms.
 
-**Model weights:** see the license specified in each model repository
+**Code and notebooks in this Space:** `TBC`

index.html

@@ -3,8 +3,8 @@
 <head>
   <meta charset="utf-8" />
   <meta name="viewport" content="width=device-width,initial-scale=1" />
-  <title>NTv3 — Next-Gen Foundation Models for Genomics</title>
-  <meta name="description" content="NTv3 companion hub: PyTorch notebooks for inference, fine-tuning, interpretation, and sequence generation on NTv3 models hosted on Hugging Face." />
+  <title>NTv3: Foundation Models for Genomics</title>
+  <meta name="description" content="NTv3 reads up to 1 Mb of DNA at single-base resolution. One architecture, three capabilities: functional-track prediction, gene annotation and regulatory sequence design." />
   <link href="https://cdnjs.cloudflare.com/ajax/libs/prism/1.29.0/themes/prism-tomorrow.min.css" rel="stylesheet" />
   <style>
     :root {
@@ -198,6 +198,51 @@
       display: block;
       border-radius: 12px;
     }
+    /* Placeholder marker for content pending release.
+       Every unresolved item on this Space is tagged with one of these,
+       so they can be found by searching the page for "TBC". */
+    .todo {
+      display: inline-block;
+      font-family: var(--mono);
+      font-size: 11px;
+      font-weight: 600;
+      letter-spacing: 0.04em;
+      padding: 2px 7px;
+      margin-left: 6px;
+      border-radius: 6px;
+      border: 1px solid rgba(251, 191, 36, 0.45);
+      background: rgba(251, 191, 36, 0.12);
+      color: #fcd34d;
+      vertical-align: middle;
+    }
+    .notice {
+      margin-top: 18px;
+      padding: 16px 18px;
+      border: 1px solid rgba(251, 191, 36, 0.35);
+      background: rgba(251, 191, 36, 0.07);
+      border-radius: 14px;
+      color: var(--muted);
+      font-size: 14px;
+      line-height: 1.6;
+    }
+    .notice strong { color: var(--text); }
+    .legacy-banner {
+      margin-top: 18px;
+      padding: 14px 18px;
+      border: 1px solid var(--border);
+      background: rgba(255, 255, 255, 0.04);
+      border-radius: 14px;
+      color: var(--muted);
+      font-size: 14px;
+      line-height: 1.6;
+    }
+    .capability-grid { margin-top: 12px; }
+    /* Keep model names such as NTv3-pretrained on one line rather than
+       breaking at the hyphen. */
+    .capability-grid td:first-child {
+      color: var(--text);
+      white-space: nowrap;
+    }
     .footer { margin-top: 22px; color: var(--muted); font-size: 13px; }
     
     /* Tab navigation styles */
@@ -244,6 +289,7 @@
     
     @media (max-width: 860px) {
       .card { grid-column: span 12; }
+      .card-stack { grid-column: span 12; }
       h1 { font-size: 28px; }
       .tab-button {
         padding: 10px 16px;
@@ -256,16 +302,19 @@
 <body>
   <div class="wrap">
     <div class="hero">
-      <h1>🧬 NTv3 — Next-Gen Foundation Models for Genomics</h1>
+      <h1>🧬 NTv3: Foundation Models for Genomics</h1>
       <p>
-        This Space is the companion hub for <strong>NTv3</strong> models: runnable notebooks for inference, fine-tuning, interpretation, and sequence generation.
+        A genomic foundation model that reads up to 1 Mb of DNA at single-base resolution. One architecture, three
+        capabilities: <strong>functional-track prediction</strong>, <strong>gene annotation</strong> and
+        <strong>regulatory sequence design</strong>. This Space is the companion hub, with runnable PyTorch notebooks.
       </p>
 
       <div class="pillrow">
-        <span class="pill">🤖 Foundation Models</span>
-		<span class="pill">🧬 Long-context genomics</span>
-		<span class="pill">🌍 Multi-species</span>
-        <span class="pill">⚡ Inference • Fine-tune • Interpret • Generate</span>
+        <span class="pill">🧬 1 Mb context, single-base resolution</span>
+        <span class="pill">🌍 Multi-species</span>
+        <span class="pill">📊 Tracks</span>
+        <span class="pill">🏷️ Annotate</span>
+        <span class="pill">✏️ Design</span>
         <span class="pill">📓 Torch notebooks</span>
       </div>
     </div>
@@ -273,7 +322,8 @@
     <!-- Tab Navigation -->
     <div class="tabs">
       <button class="tab-button active" data-tab="home">🏠 Home</button>
-      <button class="tab-button" data-tab="demo">🚀 Live Demo</button>
+      <button class="tab-button" data-tab="skill">🤖 Claude Skill</button>
+      <button class="tab-button" data-tab="initial">🗄️ Initial release (2025)</button>
     </div>
 
     <!-- Home Tab (Content loaded from tabs/home.html) -->
@@ -281,18 +331,18 @@
       <!-- Content will be loaded dynamically -->
     </div>
 
-    <!-- Live Demo Tab (Content loaded from tabs/demo.html) -->
-    <div id="demo" class="tab-content">
+    <!-- Claude Skill Tab (Content loaded from tabs/skill.html) -->
+    <div id="skill" class="tab-content">
+      <!-- Content will be loaded dynamically -->
+    </div>
+
+    <!-- Initial-release Tab (Content loaded from tabs/initial_release.html) -->
+    <div id="initial" class="tab-content">
       <!-- Content will be loaded dynamically -->
     </div>
-          
-    <!-- <div class="paper-summary">
-		<h2>📄 A foundational model for joint sequence-function multi-species modeling at scale for long-range genomic prediction</h2>
-		<img src="assets/paper_summary.png" alt="NTv3 Paper Summary" />
-    </div> -->
 
     <p class="footer">
-      © instadeep-ai — NTv3 companion Space.
+      © instadeep-ai. NTv3 companion Space. Model weights are released under the InstaDeep Open Model Licence; see each model repository.
     </p>
   </div>
   <script src="https://cdnjs.cloudflare.com/ajax/libs/prism/1.29.0/components/prism-core.min.js"></script>
@@ -301,8 +351,8 @@
     // Tab content mapping
     const tabFiles = {
       'home': 'tabs/home.html',
-      'demo': 'tabs/demo.html',
-      'functional_tracks': 'tabs/functional_tracks.html',
+      'skill': 'tabs/skill.html',
+      'initial': 'tabs/initial_release.html',
     };
 
     // Cache for loaded tab content
@@ -371,6 +421,21 @@
         showTab(activeTab.id);
       }
 
+      // In-page links that switch tabs, e.g. <a data-tab-link="initial">.
+      // Delegated, because tab content is injected after load.
+      document.addEventListener('click', (event) => {
+        const link = event.target.closest('[data-tab-link]');
+        if (!link) return;
+        event.preventDefault();
+        const targetButton = document.querySelector(
+          `.tab-button[data-tab="${link.getAttribute('data-tab-link')}"]`
+        );
+        if (targetButton) {
+          targetButton.click();
+          window.scrollTo({ top: 0, behavior: 'smooth' });
+        }
+      });
+
       tabButtons.forEach(button => {
         button.addEventListener('click', async () => {
           const targetTab = button.getAttribute('data-tab');

notebooks_tutorials/04_model_interpretation.ipynb

@@ -0,0 +1,1202 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "BXtGC0FwDK7i"
+   },
+   "source": [
+    "# 🧬 NTv3 Multi-Scale Interpretation Tutorial\n",
+    "\n",
+    "This notebook uses [`InstaDeepAI/NTv3_tracks`](https://huggingface.co/InstaDeepAI/NTv3_tracks), the NTv3 model post-trained on functional tracks, for multi-scale interpretation of regulatory logic, enabling analysis from individual nucleotides to enhancer elements and full gene loci.\n",
+    "\n",
+    "> 📝 **Note for Google Colab users:** This notebook is compatible with Colab! For faster inference, make sure to enable GPU: Runtime → Change runtime type → GPU (T4 or better recommended).\n",
+    "\n",
+    "## 📋 Overview\n",
+    "\n",
+    "Given a genomic window from the **human genome (hg38)**, we demonstrate three complementary interpretability modalities:\n",
+    "\n",
+    "1. **Predicted functional tracks** - Regulatory activity across thousands of assays\n",
+    "2. **Transformer attention maps** - How the model integrates information across long genomic distances\n",
+    "3. **Gradient-based saliency maps** - Sequence positions most influential for predictions\n",
+    "\n",
+    "> 💡 **Note:** Functional tracks are predicted only for the center 37.5% of the input sequence.\n",
+    "\n",
+    "`NTv3_tracks` predicts functional tracks, not gene annotations. Gene annotation is served by the separate NTv3-annotate checkpoints.\n",
+    "\n",
+    "## 📚 Notebook Structure\n",
+    "\n",
+    "1. **Setup**: Install dependencies and configure the analysis\n",
+    "2. **Data Loading**: Fetch the genomic sequence from UCSC\n",
+    "3. **Model Loading**: Load the NTv3 tracks model, its tokenizer and track metadata\n",
+    "4. **Inference**: Run the model to generate predictions and extract attention\n",
+    "5. **Functional Tracks**: Visualize predicted regulatory tracks\n",
+    "6. **Attention Maps**: Analyze long-range information\n",
+    "7. **Saliency Maps**: Compute gradient-based sequence attributions\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "MS_mIORoDK7k"
+   },
+   "source": [
+    "## 0) 📦 Imports + Setup\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "vKTwnGhmVYFu"
+   },
+   "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()\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "JRc8918_DK7l"
+   },
+   "outputs": [],
+   "source": [
+    "# Colab ships torch and numpy. Upgrading them breaks the preinstalled scipy.\n",
+    "!pip -q install --upgrade 'transformers>=5.17,<6'\n",
+    "!pip -q install fast-logomaker\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "qdXtceUiDK7l"
+   },
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import math\n",
+    "import warnings\n",
+    "import logging\n",
+    "\n",
+    "import torch\n",
+    "import torch.nn.functional as F\n",
+    "from transformers import AutoModel, AutoTokenizer\n",
+    "\n",
+    "import requests\n",
+    "import matplotlib.pyplot as plt\n",
+    "import numpy as np\n",
+    "import seaborn as sns\n",
+    "\n",
+    "# Device setup\n",
+    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
+    "if device == \"cuda\" and torch.cuda.get_device_capability(0)[0] >= 8:\n",
+    "    dtype = torch.bfloat16\n",
+    "else:\n",
+    "    dtype = torch.float16\n",
+    "print(f\"Device: {device}, dtype: {dtype}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "gOCdTnRADK7m"
+   },
+   "source": [
+    "## 1) 📦 Configuration\n",
+    "\n",
+    "Set your NTv3 model, species and genomic window here. Coordinates are in the assembly of the chosen species (hg38 for human, the default), and the tracks you can pick come from that assembly's track list.\n",
+    "\n",
+    "> 💡 **Note:** For a T4 GPU, we recommend using an input sequence length of 32kb or shorter.\n",
+    "\n",
+    "> The knobs are `species`, `assembly`, `chrom`, `center_position` and `sequence_length` in the cell below. `assembly` must be the one `SPECIES_TO_ASSEMBLY` gives for `species`, or empty to use it; a mismatch raises an error. `start` and `end` are derived from `center_position` and `sequence_length`. The tracks to plot are chosen later, in `tracks_to_plot` (section 5) and `saliency_track_id` (section 7).\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "iBmudDzIDK7m"
+   },
+   "outputs": [],
+   "source": [
+    "# -----------------------------\n",
+    "# User Configuration\n",
+    "# -----------------------------\n",
+    "model_name = \"InstaDeepAI/NTv3_tracks\"\n",
+    "\n",
+    "# Species and its assembly; leave assembly empty to use the species' assembly\n",
+    "species = \"human\"  #@param [\"human\", \"mouse\", \"drosophila\", \"arabidopsis\", \"maize\", \"rice\", \"soybean\", \"caenorhabditis_elegans\", \"zebrafish\", \"chicken\", \"horse\", \"sheep\"]\n",
+    "assembly = \"hg38\"  #@param [\"\", \"hg38\", \"mm10\", \"dm6\", \"TAIR10\", \"Zm-B73-REFERENCE-NAM-5.0\", \"IRGSP-1.0\", \"Glycine_max_v2.1\", \"WBcel235\", \"GRCz11\", \"GRCg6a\", \"EquCab3.0\", \"Oar_v3.1\"]\n",
+    "\n",
+    "# Genomic window (must be multiple of 128 due to model downsampling)\n",
+    "chrom = \"chr11\"  #@param {type:\"string\"}\n",
+    "# Center of region of interest (HBE1 locus)\n",
+    "center_position = 5269945  #@param {type:\"integer\"}\n",
+    "# use 131kb window on A100, 32_768 on T4 GPU\n",
+    "sequence_length = 32768  #@param [16384, 32768, 65536, 131072] {type:\"raw\"}\n",
+    "assert sequence_length % 128 == 0, \"sequence_length must be a multiple of 128\"\n",
+    "\n",
+    "# Calculate start/end from center\n",
+    "start = center_position - sequence_length // 2\n",
+    "end = center_position + sequence_length // 2\n",
+    "\n",
+    "# Analysis parameters\n",
+    "token_resolution = 128  # bp per token\n",
+    "\n",
+    "# Track lists are keyed by assembly, so each species maps to one assembly.\n",
+    "# The names match metadata[\"bigwigs_per_file_assembly\"], printed once the model loads.\n",
+    "SPECIES_TO_ASSEMBLY = {\n",
+    "    \"human\": \"hg38\",\n",
+    "    \"mouse\": \"mm10\",\n",
+    "    \"drosophila\": \"dm6\",\n",
+    "    \"arabidopsis\": \"TAIR10\",\n",
+    "    \"maize\": \"Zm-B73-REFERENCE-NAM-5.0\",\n",
+    "    \"rice\": \"IRGSP-1.0\",\n",
+    "    \"soybean\": \"Glycine_max_v2.1\",\n",
+    "    \"caenorhabditis_elegans\": \"WBcel235\",\n",
+    "    \"zebrafish\": \"GRCz11\",\n",
+    "    \"chicken\": \"GRCg6a\",\n",
+    "    \"horse\": \"EquCab3.0\",\n",
+    "    \"sheep\": \"Oar_v3.1\",\n",
+    "}\n",
+    "if species not in SPECIES_TO_ASSEMBLY:\n",
+    "    raise ValueError(f\"Unknown species {species!r}. Choose one of {list(SPECIES_TO_ASSEMBLY)}\")\n",
+    "if not assembly:\n",
+    "    assembly = SPECIES_TO_ASSEMBLY[species]\n",
+    "elif assembly != SPECIES_TO_ASSEMBLY[species]:\n",
+    "    raise ValueError(\n",
+    "        f\"assembly={assembly!r} does not match species={species!r}, \"\n",
+    "        f\"whose assembly is {SPECIES_TO_ASSEMBLY[species]!r}\"\n",
+    "    )\n",
+    "\n",
+    "print(f\"Model: {model_name}\")\n",
+    "print(f\"Region: {chrom}:{start:,}-{end:,} ({sequence_length:,} bp)\")\n",
+    "print(f\"Center: {center_position:,}\")\n",
+    "print(f\"Species: {species} -> Assembly: {assembly}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "yytbImqlDK7m"
+   },
+   "source": [
+    "## 2) 📥 Fetch Chromosome Sequence\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "cqXIi7mYDK7m"
+   },
+   "outputs": [],
+   "source": [
+    "# Get the sequence from the UCSC API\n",
+    "url = (\n",
+    "    f\"https://api.genome.ucsc.edu/getData/sequence?\"\n",
+    "    f\"genome={assembly};chrom={chrom};start={start};end={end}\"\n",
+    ")\n",
+    "seq = requests.get(url, timeout=60).json()[\"dna\"]\n",
+    "print(f\"Original sequence length: {len(seq):,} bp\")\n",
+    "\n",
+    "# Crop to multiple of 128 (ensures clean tokenization)\n",
+    "seq = seq[:int(len(seq) // 128) * 128]\n",
+    "n_tokens = len(seq) / 128\n",
+    "print(\n",
+    "    f\"Cropped sequence length: {len(seq):,} bp \"\n",
+    "    f\"({n_tokens:.0f} transformer tokens)\"\n",
+    ")\n",
+    "\n",
+    "# UCSC returns soft-masked sequence (repeats in lowercase). NTv3_tracks was trained\n",
+    "# on soft-masked genomes, so we keep the case as it is.\n",
+    "soft_masked_fraction = sum(b.islower() for b in seq) / len(seq)\n",
+    "print(f\"Soft-masked (lowercase) fraction: {soft_masked_fraction:.1%}\")\n",
+    "print(f\"First 100 bp: {seq[:100]}\")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "jtJMea-1DK7m"
+   },
+   "source": [
+    "## 3) 🤖 Load NTv3 Model + Tokenizer\n",
+    "\n",
+    "The model uses custom code, so we pass `trust_remote_code=True`. The species vocabulary and the track order ship with the model config, under `track_metadata`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "L8RFTmRaDK7m"
+   },
+   "outputs": [],
+   "source": [
+    "# Load model\n",
+    "print(\"Loading model...\")\n",
+    "model = AutoModel.from_pretrained(\n",
+    "    model_name,\n",
+    "    trust_remote_code=True,\n",
+    ").to(device)\n",
+    "model.eval()\n",
+    "\n",
+    "# Load tokenizer\n",
+    "print(\"Loading tokenizer...\")\n",
+    "tokenizer = AutoTokenizer.from_pretrained(\n",
+    "    model_name, trust_remote_code=True\n",
+    ")\n",
+    "\n",
+    "# Species ids and track order ship with the model config\n",
+    "metadata = model.config.track_metadata\n",
+    "species_vocab = metadata[\"species_vocabulary\"]\n",
+    "\n",
+    "print(f\"\\n✅ Model loaded successfully\")\n",
+    "supported_species = [s for s in species_vocab if not s.startswith(\"<\")]\n",
+    "print(f\"   Supported species: {supported_species}\")\n",
+    "assemblies = list(metadata[\"bigwigs_per_file_assembly\"])\n",
+    "print(f\"   Available assemblies: {assemblies}\")\n",
+    "if species not in supported_species:\n",
+    "    raise ValueError(f\"Species {species!r} is not supported by {model_name}\")\n",
+    "if assembly not in assemblies:\n",
+    "    raise ValueError(f\"Assembly {assembly!r} has no tracks in {model_name}\")\n",
+    "n_bigwigs = len(metadata[\"bigwigs_per_file_assembly\"][assembly])\n",
+    "print(f\"   Number of {assembly} BigWig tracks: {n_bigwigs}\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "_XZ3xXIMDK7n"
+   },
+   "outputs": [],
+   "source": [
+    "# Inspect available outputs: channel i of the output is track i of this list\n",
+    "bigwig_names = metadata[\"bigwigs_per_file_assembly\"][assembly]\n",
+    "print(\n",
+    "    f\"{len(bigwig_names)} functional tracks for {assembly}. \"\n",
+    "    f\"First 10: {bigwig_names[:10]}\"\n",
+    ")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "rFdiD_xZDK7n"
+   },
+   "source": [
+    "## 4) ⚡ Run Inference\n",
+    "\n",
+    "Following the manuscript methods: The functional-track module produces per-base predictions for all species-matched regulatory assays. Supervision is applied only to the central 37.5% of the input sequence.\n",
+    "\n",
+    "We pass:\n",
+    "- the tokenized DNA window, positionally: the forward takes the token tensor, not the tokenizer dict\n",
+    "- `conditions_tokens`: the species id, as a list with one `(batch,)` tensor\n",
+    "\n",
+    "`bigwig_tracks_logits` comes back on the forward output. The trunk uses scaled-dot-product attention and does not return weights, so the maps below are computed from each attention block.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "bAidYEvWDK7n"
+   },
+   "outputs": [],
+   "source": [
+    "# Tokenize DNA sequence (its length is already a multiple of 128)\n",
+    "batch = tokenizer([seq], add_special_tokens=False, return_tensors=\"pt\")\n",
+    "input_ids = batch[\"input_ids\"].to(device)\n",
+    "# Characters outside the alphabet (e.g. IUPAC codes) become <unk> without changing the length\n",
+    "assert tokenizer.unk_token_id not in input_ids, \"the sequence contains characters outside the alphabet\"\n",
+    "\n",
+    "# Encode species\n",
+    "species_ids = torch.tensor([species_vocab[species]], device=device)\n",
+    "\n",
+    "print(f\"Input IDs shape: {input_ids.shape}\")\n",
+    "print(f\"Species IDs shape: {species_ids.shape}\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "yWDD1Z0zDK7n"
+   },
+   "outputs": [],
+   "source": [
+    "# Run inference. Attention weights are not a forward argument on this model.\n",
+    "import importlib\n",
+    "\n",
+    "print(\"Running forward pass...\")\n",
+    "blocks = model.model.mha_blocks\n",
+    "modeling = importlib.import_module(type(blocks[0]).__module__)\n",
+    "captured = []\n",
+    "originals = [block.forward for block in blocks]\n",
+    "\n",
+    "def _capturing_forward(block, original):\n",
+    "    def forward(x, attention_mask=None, conditions=None, conditions_masks=None):\n",
+    "        batch, seq_len, _ = x.shape\n",
+    "        if block.conditioned:\n",
+    "            hidden = block.norm_in(x, conditions, conditions_masks)\n",
+    "        else:\n",
+    "            hidden = block.norm_in(x)\n",
+    "        query = block.norm_q(\n",
+    "            block.q_layer(hidden).view(batch, seq_len, block.attention_heads, -1)\n",
+    "        )\n",
+    "        key = block.norm_k(\n",
+    "            block.k_layer(hidden).view(batch, seq_len, block.num_kv_heads, -1)\n",
+    "        )\n",
+    "        inv_freq = modeling.build_rope_inv_freq(block.key_size, device=x.device)\n",
+    "        query = modeling.apply_rope(query, inv_freq)\n",
+    "        key = modeling.apply_rope(key, inv_freq)\n",
+    "        group = block.attention_heads // block.num_kv_heads\n",
+    "        query = query.transpose(1, 2)\n",
+    "        key = key.transpose(1, 2).repeat_interleave(group, dim=1)\n",
+    "        scale = 1.0 / math.sqrt(block.key_size)\n",
+    "        logits = torch.matmul(query.float(), key.float().transpose(-2, -1)) * scale\n",
+    "        if attention_mask is not None:\n",
+    "            logits = logits.masked_fill(~attention_mask, float(\"-inf\"))\n",
+    "        weights = torch.softmax(logits, dim=-1)\n",
+    "        captured.append(weights[0].cpu().numpy())\n",
+    "        return original(\n",
+    "            x,\n",
+    "            attention_mask=attention_mask,\n",
+    "            conditions=conditions,\n",
+    "            conditions_masks=conditions_masks,\n",
+    "        )\n",
+    "\n",
+    "    return forward\n",
+    "\n",
+    "for block, original in zip(blocks, originals):\n",
+    "    block.forward = _capturing_forward(block, original)\n",
+    "try:\n",
+    "    with torch.no_grad():\n",
+    "        outputs = model(input_ids, conditions_tokens=[species_ids])\n",
+    "finally:\n",
+    "    for block, original in zip(blocks, originals):\n",
+    "        block.forward = original\n",
+    "\n",
+    "# [batch, 0.375 * seq_len, num_tracks]\n",
+    "bigwig_logits = outputs[\"bigwig_tracks_logits\"]\n",
+    "# (layers, heads, blocks, blocks), one 128 bp block per position\n",
+    "attention_maps = np.stack(captured, axis=0)\n",
+    "\n",
+    "print(\"\\n✅ Forward pass completed\")\n",
+    "print(f\"   BigWig logits: {tuple(bigwig_logits.shape)}\")\n",
+    "print(f\"   Attention maps: {attention_maps.shape} (layers, heads, seq, seq)\")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "WpKxn1juDK7n"
+   },
+   "source": [
+    "## 5) 📊 Visualize Functional Tracks\n",
+    "\n",
+    "Following the manuscript methods: Functional-track predictions are obtained directly from the model output tensor. Predictions are center-cropped to the supervised region (49,152 bp for 131 kb inputs).\n",
+    "\n",
+    "Metadata for all available tracks can be downloaded [here](https://huggingface.co/spaces/InstaDeepAI/ntv3_tracks/blob/main/data/functional_tracks_metadata.csv)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "u3pzwP7VDK7n"
+   },
+   "outputs": [],
+   "source": [
+    "# Select functional tracks to plot\n",
+    "tracks_to_plot = {\n",
+    "    \"K562 RNA-seq\": \"ENCSR056HPM\",\n",
+    "    \"K562 CAGE\": \"CNhs11250_P\",\n",
+    "    \"K562 DNAse\": \"ENCSR921NMD\",\n",
+    "    \"K562 GATA1 ChIP-seq\": \"ENCSR000EFT\"\n",
+    "}\n",
+    "\n",
+    "# Every track must exist for this assembly\n",
+    "missing_tracks = [v for v in tracks_to_plot.values() if v not in bigwig_names]\n",
+    "if missing_tracks:\n",
+    "    raise ValueError(f\"Tracks not available for {assembly}: {missing_tracks}\")\n",
+    "\n",
+    "print(f\"Plotting {len(tracks_to_plot)} functional tracks\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "gtF2KGoQDK7n"
+   },
+   "outputs": [],
+   "source": [
+    "def plot_tracks(tracks, start_pos, end_pos, chrom_name, assembly_name,\n",
+    "                height=1.0):\n",
+    "    \"\"\"Plot multiple tracks in a stacked view.\n",
+    "\n",
+    "    Args:\n",
+    "        tracks: Dictionary mapping track names to 1D arrays of values.\n",
+    "        start_pos: Genomic start position for x-axis.\n",
+    "        end_pos: Genomic end position for x-axis.\n",
+    "        chrom_name: Chromosome name for labeling.\n",
+    "        assembly_name: Assembly name for labeling.\n",
+    "        height: Height multiplier for each track subplot.\n",
+    "\n",
+    "    Returns:\n",
+    "        matplotlib.figure.Figure: The created figure object.\n",
+    "    \"\"\"\n",
+    "    n_tracks = len(tracks)\n",
+    "    fig, axes = plt.subplots(\n",
+    "        n_tracks, 1, figsize=(16, height * n_tracks), sharex=True\n",
+    "    )\n",
+    "    if n_tracks == 1:\n",
+    "        axes = [axes]\n",
+    "\n",
+    "    for ax, (title, y) in zip(axes, tracks.items()):\n",
+    "        x_positions = np.linspace(start_pos, end_pos, num=len(y))\n",
+    "        ax.fill_between(x_positions, y, alpha=0.7)\n",
+    "        ax.set_title(title, fontsize=8)\n",
+    "        ax.yaxis.set_label_coords(-0.01, 0.5)\n",
+    "        sns.despine(ax=ax, top=True, right=True, bottom=True)\n",
+    "        ax.tick_params(axis='y', labelsize=7)\n",
+    "\n",
+    "    xlabel = f\"{chrom_name}:{start_pos:,}-{end_pos:,} ({assembly_name})\"\n",
+    "    axes[-1].set_xlabel(xlabel, fontsize=9)\n",
+    "    plt.tight_layout()\n",
+    "    return fig\n",
+    "\n",
+    "# Extract and plot functional tracks\n",
+    "bigwig = bigwig_logits.detach().float().cpu().numpy()[0]\n",
+    "bigwig_tracks = {}\n",
+    "for track_name, track_id in tracks_to_plot.items():\n",
+    "    if bigwig_names and track_id in bigwig_names:\n",
+    "        track_idx = bigwig_names.index(track_id)\n",
+    "        bigwig_tracks[track_name] = bigwig[:, track_idx]\n",
+    "\n",
+    "if not bigwig_tracks:\n",
+    "    print(\"⚠️ No tracks to plot (track names may not be available)\")\n",
+    "else:\n",
+    "    # Calculate prediction region (center 37.5%)\n",
+    "    window_len = end - start\n",
+    "    prediction_start = start + int(window_len * 0.3125)\n",
+    "    prediction_end = prediction_start + int(window_len * 0.375)\n",
+    "\n",
+    "    fig = plot_tracks(\n",
+    "        bigwig_tracks, prediction_start, prediction_end, chrom, assembly\n",
+    "    )\n",
+    "    plt.suptitle(\"Predicted Functional Tracks\", fontsize=11, y=1.02)\n",
+    "    plt.show()\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "> **How to read this plot.** The x axis is the genomic coordinate on `chrom` in `assembly`. It covers the predicted span, the center 37.5% of the input window, not the full input. Each panel is one track at 1 bp resolution, on its own y scale. A high value means the model predicts a strong signal for that assay at that position.\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "iqyiyLPzDK7o"
+   },
+   "source": [
+    "## 6) 🔗 Analyze Attention Maps\n",
+    "\n",
+    "Following the manuscript methods: Transformer self-attention reveals how NTv3 integrates information across distal regions. We extract multi-head self-attention weights and average them across heads. Matrices are symmetrized as (A + A^T)/2.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "StvwdIcUDK7o"
+   },
+   "outputs": [],
+   "source": [
+    "# Extract attention from last layer\n",
+    "layer_index = -1  # Last layer\n",
+    "# [heads, seq_len, seq_len]\n",
+    "attention_last_layer = attention_maps[layer_index]\n",
+    "\n",
+    "# Average over heads and symmetrize\n",
+    "# [seq_len, seq_len]\n",
+    "attention_mean = attention_last_layer.mean(axis=0)\n",
+    "attention_mean = 0.5 * (attention_mean + attention_mean.T)  # Symmetrize\n",
+    "\n",
+    "print(f\"Attention map shape: {attention_mean.shape}\")\n",
+    "attn_min = attention_mean.min()\n",
+    "attn_max = attention_mean.max()\n",
+    "print(f\"Attention value range: [{attn_min:.6f}, {attn_max:.6f}]\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "qRxVY-6VDK7o"
+   },
+   "outputs": [],
+   "source": [
+    "# ============================================================\n",
+    "# Pyramid-style Attention Map Plotting (45-degree rotated triangular)\n",
+    "# ============================================================\n",
+    "from matplotlib.colors import TwoSlopeNorm\n",
+    "\n",
+    "# Global settings\n",
+    "LINE_WIDTH = 0.5\n",
+    "\n",
+    "def _get_45deg_mesh(mat):\n",
+    "    \"\"\"Create 45-degree rotated mesh for triangular attention visualization.\n",
+    "\n",
+    "    Args:\n",
+    "        mat: Input matrix to create mesh for.\n",
+    "\n",
+    "    Returns:\n",
+    "        tuple: (X, Y) mesh coordinates rotated 45 degrees.\n",
+    "    \"\"\"\n",
+    "    theta = -np.pi / 4\n",
+    "    R = np.array([\n",
+    "        [np.cos(theta), -np.sin(theta)],\n",
+    "        [np.sin(theta), np.cos(theta)]\n",
+    "    ])\n",
+    "    K = len(mat) + 1\n",
+    "    g = np.arange(0, K) - 0.5\n",
+    "    X = np.tile(g[:, None], (1, K))\n",
+    "    Y = np.tile(g[None, :], (K, 1))\n",
+    "    xy = np.array([X.ravel(), Y.ravel()])\n",
+    "    xy_rot = R @ xy\n",
+    "    return xy_rot[0, :].reshape(K, K), xy_rot[1, :].reshape(K, K).T\n",
+    "\n",
+    "def plot_attention_panel(mat, filename=None, *, cmap=\"Blues\", vmin=0.0001,\n",
+    "                         vmax=0.005, colorbar=True, dpi=600, figsize=(10, 10),\n",
+    "                         variant_id=\"\", show_titles=True, show_xlabel=True,\n",
+    "                         show_xticks=True, show_yticks=False,\n",
+    "                         positions='TOKEN', token_resolution=128, show=True):\n",
+    "    \"\"\"Plot triangular attention heatmap with 45-degree rotation (pyramid style).\n",
+    "\n",
+    "    Args:\n",
+    "        mat: Attention matrix to plot.\n",
+    "        filename: Output filename. If None and show=True, displays plot.\n",
+    "        cmap: Colormap name. Defaults to \"Blues\".\n",
+    "        vmin: Minimum value for color scale. Defaults to 0.0001.\n",
+    "        vmax: Maximum value for color scale. Defaults to 0.005.\n",
+    "        colorbar: Whether to show colorbar. Defaults to True.\n",
+    "        dpi: Figure DPI. Defaults to 600.\n",
+    "        figsize: Figure size tuple. Defaults to (10, 10).\n",
+    "        variant_id: Title for the plot. Defaults to \"\".\n",
+    "        show_titles: Whether to show title. Defaults to True.\n",
+    "        show_xlabel: Whether to show x-axis label. Defaults to True.\n",
+    "        show_xticks: Whether to show x-axis ticks. Defaults to True.\n",
+    "        show_yticks: Whether to show y-axis ticks. Defaults to False.\n",
+    "        positions: 'TOKEN' or 'BP' for position units. Defaults to 'TOKEN'.\n",
+    "        token_resolution: Base pairs per token. Defaults to 128.\n",
+    "        show: Whether to display plot. Defaults to True.\n",
+    "\n",
+    "    Returns:\n",
+    "        matplotlib.figure.Figure: The created figure object.\n",
+    "    \"\"\"\n",
+    "    # Symmetrize matrix\n",
+    "    mat = 0.5 * (mat + mat.T)\n",
+    "\n",
+    "    # Mask lower triangle\n",
+    "    mat[np.tril_indices_from(mat, k=-1)] = np.nan\n",
+    "    n = mat.shape[0]\n",
+    "    X, Y = _get_45deg_mesh(mat)\n",
+    "\n",
+    "    # Coordinate normalization\n",
+    "    C = 1\n",
+    "    half_pixel_diag = 1 / (2*C)\n",
+    "    pixel_side = 1 / (C * np.sqrt(2))\n",
+    "    X = X * pixel_side + half_pixel_diag\n",
+    "    Y = Y * pixel_side\n",
+    "    Y = -Y  # Flip Y\n",
+    "\n",
+    "    # Clean up coordinates\n",
+    "    X = np.nan_to_num(X, nan=0.0, posinf=1.0, neginf=-1.0)\n",
+    "    Y = np.nan_to_num(Y, nan=0.0, posinf=1.0, neginf=-1.0)\n",
+    "\n",
+    "    # Set up normalization\n",
+    "    vcenter = (vmin + vmax) / 2\n",
+    "    norm = TwoSlopeNorm(vmin=vmin, vcenter=vcenter, vmax=vmax)\n",
+    "\n",
+    "    fig, ax = plt.subplots(figsize=figsize, dpi=dpi)\n",
+    "    im = ax.pcolormesh(X, Y, mat, cmap=cmap, norm=norm, shading=\"auto\")\n",
+    "\n",
+    "    # Set axis properties\n",
+    "    ax.set_aspect(\"equal\")\n",
+    "\n",
+    "    # Calculate coordinate extents\n",
+    "    finite_mask = np.isfinite(mat)\n",
+    "    if finite_mask.any():\n",
+    "        vertex_mask = np.zeros_like(X, dtype=bool)\n",
+    "        vertex_mask[:-1, :-1] |= finite_mask\n",
+    "        vertex_mask[:-1, 1:] |= finite_mask\n",
+    "        vertex_mask[1:, :-1] |= finite_mask\n",
+    "        vertex_mask[1:, 1:] |= finite_mask\n",
+    "\n",
+    "        valid_x = X[vertex_mask]\n",
+    "        valid_y = Y[vertex_mask]\n",
+    "        x_min, x_max = np.nanmin(valid_x), np.nanmax(valid_x)\n",
+    "        y_min, y_max = np.nanmin(valid_y), np.nanmax(valid_y)\n",
+    "    else:\n",
+    "        x_min, x_max = np.nanmin(X), np.nanmax(X)\n",
+    "        y_min, y_max = np.nanmin(Y), np.nanmax(Y)\n",
+    "\n",
+    "    triangle_bottom_y = y_min\n",
+    "    triangle_top_y = y_max\n",
+    "\n",
+    "    xlim_pad = 0.1\n",
+    "    ylim_pad = 0.1\n",
+    "    y_range = (y_max - y_min) if y_max > y_min else 1.0\n",
+    "    bottom_pad = 0.001 * y_range\n",
+    "\n",
+    "    ax.set_xlim(x_min - xlim_pad, x_max + xlim_pad)\n",
+    "    ax.set_ylim(y_min - bottom_pad, y_max + ylim_pad)\n",
+    "\n",
+    "    # Add colorbar\n",
+    "    if colorbar:\n",
+    "        pyramid_y_min = y_min\n",
+    "        pyramid_y_max = y_max\n",
+    "\n",
+    "        cb = plt.colorbar(im, ax=ax, shrink=1.0, aspect=20, pad=0.1)\n",
+    "        cb.ax.tick_params(width=LINE_WIDTH*0.8, length=3, labelsize=7, direction=\"in\")\n",
+    "        cb.outline.set_linewidth(LINE_WIDTH*0.8)\n",
+    "\n",
+    "        ax_pos = ax.get_position()\n",
+    "        axes_y_min_data = y_min - bottom_pad\n",
+    "        axes_y_max_data = y_max + ylim_pad\n",
+    "        axes_height_data = axes_y_max_data - axes_y_min_data\n",
+    "\n",
+    "        if axes_height_data > 0:\n",
+    "            pyramid_bottom_frac = (\n",
+    "                (pyramid_y_min - axes_y_min_data) / axes_height_data\n",
+    "            )\n",
+    "            pyramid_top_frac = (\n",
+    "                (pyramid_y_max - axes_y_min_data) / axes_height_data\n",
+    "            )\n",
+    "        else:\n",
+    "            pyramid_bottom_frac = 0\n",
+    "            pyramid_top_frac = 1\n",
+    "\n",
+    "        cb_pos = cb.ax.get_position()\n",
+    "        new_y0 = ax_pos.y0 + ax_pos.height * pyramid_bottom_frac\n",
+    "        new_height = ax_pos.height * (pyramid_top_frac - pyramid_bottom_frac)\n",
+    "        cb.ax.set_position(\n",
+    "            [cb_pos.x0, new_y0, cb_pos.width, new_height]\n",
+    "        )\n",
+    "\n",
+    "    # Format axes\n",
+    "    use_bp_display = positions.upper() == 'BP'\n",
+    "\n",
+    "    ax.tick_params(width=LINE_WIDTH, length=3, labelsize=7)\n",
+    "    for sp in (\"top\", \"right\", \"left\"):\n",
+    "        ax.spines[sp].set_visible(False)\n",
+    "    ax.spines[\"bottom\"].set_linewidth(LINE_WIDTH)\n",
+    "    ax.spines[\"bottom\"].set_position(('data', triangle_bottom_y))\n",
+    "    ax.xaxis.set_ticks_position('bottom')\n",
+    "    ax.xaxis.set_label_position('bottom')\n",
+    "\n",
+    "    if show_xticks and mat.size > 0:\n",
+    "        n_tokens = mat.shape[0]\n",
+    "        tick_positions = [0, n_tokens - 1]\n",
+    "        if use_bp_display:\n",
+    "            max_bp = (n_tokens - 1) * token_resolution\n",
+    "            tick_labels = ['0', str(max_bp)]\n",
+    "        else:\n",
+    "            tick_labels = ['0', str(n_tokens - 1)]\n",
+    "        ax.set_xticks(tick_positions)\n",
+    "        ax.set_xticklabels(tick_labels)\n",
+    "    else:\n",
+    "        ax.set_xticks([])\n",
+    "\n",
+    "    if not show_yticks:\n",
+    "        ax.set_yticks([])\n",
+    "\n",
+    "    if use_bp_display:\n",
+    "        xlabel_text = \"Position (base pairs)\"\n",
+    "    else:\n",
+    "        xlabel_text = \"Token position\"\n",
+    "    if show_xlabel:\n",
+    "        ax.set_xlabel(xlabel_text, fontsize=8)\n",
+    "\n",
+    "    if show_titles and variant_id:\n",
+    "        ax.set_title(variant_id, fontsize=9, pad=10)\n",
+    "\n",
+    "    ax.set_facecolor(\"white\")\n",
+    "\n",
+    "    if filename:\n",
+    "        dirname = os.path.dirname(filename)\n",
+    "        if dirname:\n",
+    "            os.makedirs(dirname, exist_ok=True)\n",
+    "        plt.savefig(\n",
+    "            filename, dpi=600, bbox_inches=\"tight\", facecolor='white'\n",
+    "        )\n",
+    "        print(f\"✅ Saved attention plot: {filename}\")\n",
+    "        if not show:\n",
+    "            plt.close(fig)\n",
+    "    elif show:\n",
+    "        plt.show()\n",
+    "    else:\n",
+    "        plt.close(fig)\n",
+    "\n",
+    "    return fig\n",
+    "\n",
+    "print(\"✅ Pyramid-style attention plotting function defined\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "qeziBpBmDK7o"
+   },
+   "outputs": [],
+   "source": [
+    "# Plot pyramid-style attention map (cropped to center window)\n",
+    "LAYER_INDEX = layer_index\n",
+    "TOKEN_RESOLUTION = token_resolution\n",
+    "window_size = min(1024, len(attention_mean))  # tokens to show\n",
+    "\n",
+    "# Crop to center window\n",
+    "n_tokens_full = len(attention_mean)\n",
+    "start_crop = (n_tokens_full - window_size) // 2\n",
+    "end_crop = start_crop + window_size\n",
+    "attention_mean_cropped = attention_mean[\n",
+    "    start_crop:end_crop, start_crop:end_crop\n",
+    "].copy()\n",
+    "\n",
+    "print(f\"Plotting attention map:\")\n",
+    "print(f\"  Layer: {LAYER_INDEX}\")\n",
+    "print(\n",
+    "    f\"  Window: {window_size} tokens \"\n",
+    "    f\"(cropped from {n_tokens_full} total)\"\n",
+    ")\n",
+    "attn_crop_min = attention_mean_cropped.min()\n",
+    "attn_crop_max = attention_mean_cropped.max()\n",
+    "print(f\"  Value range: [{attn_crop_min:.6f}, {attn_crop_max:.6f}]\")\n",
+    "\n",
+    "# Plot with pyramid style using shared utility\n",
+    "title = f\"Attention Map - Layer {LAYER_INDEX} (Head Average)\"\n",
+    "fig = plot_attention_panel(\n",
+    "    attention_mean_cropped,\n",
+    "    cmap=\"Blues\",\n",
+    "    vmin=0.0001,\n",
+    "    vmax=0.005,\n",
+    "    colorbar=True,\n",
+    "    figsize=(10, 10),\n",
+    "    dpi=150,\n",
+    "    show_xlabel=True,\n",
+    "    show_xticks=True,\n",
+    "    positions='TOKEN',\n",
+    "    token_resolution=TOKEN_RESOLUTION,\n",
+    "    show=True,\n",
+    ")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "> **How to read this plot.** The x axis is token position across the input window, one token per `token_resolution` bp (128 bp), up to 1,024 tokens around the center, so it covers the full input rather than only the predicted span. Each pixel is the last-layer attention between two tokens, averaged over heads and symmetrized. The map is the upper triangle rotated by 45°: the diagonal runs along the base and holds each token with itself and its neighbours, the local context. A pixel sits above the midpoint of its two tokens, higher the farther apart they are, so dark blobs high above the base are long-range interactions. Darker means more attention, and the colour scale is clipped at `vmax`, so the strongest pairs saturate."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "J7jJHp0uDK7o"
+   },
+   "source": [
+    "## 7) 🎯 Compute Saliency Maps\n",
+    "\n",
+    "Following the manuscript methods: Gradient-based attribution provides a first-order, nucleotide-resolved estimate of how NTv3 predictions depend on the underlying DNA sequence. We differentiate model outputs with respect to the one-hot encoded DNA input using a differentiable embedding override: a forward hook replaces the token embedding lookup with `one_hot @ embedding.weight`.\n",
+    "\n",
+    "For the selected functional track(s), we aggregate logits over sequence positions in a window and apply a log transform.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "FhAzknM4DK7o"
+   },
+   "outputs": [],
+   "source": [
+    "def find_token_embedding(model, vocab_size):\n",
+    "    \"\"\"Return the embedding layer that maps DNA tokens to vectors.\n",
+    "\n",
+    "    Args:\n",
+    "        model: The model to search.\n",
+    "        vocab_size: Size of the DNA token vocabulary.\n",
+    "\n",
+    "    Returns:\n",
+    "        torch.nn.Embedding: The token embedding layer.\n",
+    "    \"\"\"\n",
+    "    try:\n",
+    "        embedding = model.get_input_embeddings()\n",
+    "    except (AttributeError, NotImplementedError):\n",
+    "        embedding = None\n",
+    "    if isinstance(embedding, torch.nn.Embedding):\n",
+    "        return embedding\n",
+    "\n",
+    "    # Fall back to the only embedding table sized to the DNA vocabulary\n",
+    "    candidates = [\n",
+    "        module for module in model.modules()\n",
+    "        if isinstance(module, torch.nn.Embedding)\n",
+    "        and module.num_embeddings == vocab_size\n",
+    "    ]\n",
+    "    if len(candidates) != 1:\n",
+    "        raise ValueError(\n",
+    "            f\"Expected one token embedding of size {vocab_size}, \"\n",
+    "            f\"found {len(candidates)}. \"\n",
+    "            \"Saliency computation requires access to embedding weights.\"\n",
+    "        )\n",
+    "    return candidates[0]\n",
+    "\n",
+    "\n",
+    "class SaliencyComputer:\n",
+    "    \"\"\"Compute gradient-based saliency maps for NTv3_tracks.\"\"\"\n",
+    "\n",
+    "    def __init__(self, model, tokenizer, track_indices, species=\"human\",\n",
+    "                 device=None, target_window_bp=None):\n",
+    "        \"\"\"Initialize saliency computer.\n",
+    "\n",
+    "        Args:\n",
+    "            model: The model to compute saliency for.\n",
+    "            tokenizer: Tokenizer for input sequences.\n",
+    "            track_indices: List of track indices to compute saliency for.\n",
+    "            species: Species name. Defaults to \"human\".\n",
+    "            device: Device to use. If None, auto-detects. Defaults to None.\n",
+    "            target_window_bp: Optional (start, end) tuple in bp along the\n",
+    "                output axis to aggregate over. Indices are relative to the\n",
+    "                output (0 = first predicted position). If None (default),\n",
+    "                sums over all output positions. The output is at 1 bp\n",
+    "                resolution.\n",
+    "        \"\"\"\n",
+    "        # Auto-detect device if not provided\n",
+    "        if device is None:\n",
+    "            device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
+    "\n",
+    "        self.model = model\n",
+    "        self.tokenizer = tokenizer\n",
+    "        # Only these channels of the track head are computed\n",
+    "        self.track_indices = torch.tensor(track_indices, device=device)\n",
+    "        self.species = species\n",
+    "        self.device = device\n",
+    "        self.target_window_bp = target_window_bp\n",
+    "        self.embedding_layer = find_token_embedding(model, len(tokenizer))\n",
+    "\n",
+    "        # Species id from the vocabulary shipped with the model config\n",
+    "        species_vocab = model.config.track_metadata[\"species_vocabulary\"]\n",
+    "        self.species_ids = torch.tensor([species_vocab[species]], device=device)\n",
+    "\n",
+    "    def compute_saliency(self, sequence):\n",
+    "        \"\"\"Compute saliency map for given sequence.\n",
+    "\n",
+    "        Args:\n",
+    "            sequence: DNA sequence string to compute saliency for. Its\n",
+    "                length must be a multiple of 128.\n",
+    "\n",
+    "        Returns:\n",
+    "            tuple: (gradients, one_hot_sequence) where gradients is the\n",
+    "                saliency scores and one_hot_sequence is the one-hot encoded\n",
+    "                input sequence, both (positions, vocab_size).\n",
+    "        \"\"\"\n",
+    "        # Tokenize\n",
+    "        batch = self.tokenizer(\n",
+    "            [sequence], add_special_tokens=False, return_tensors=\"pt\"\n",
+    "        )\n",
+    "        token_ids = batch[\"input_ids\"].to(self.device)\n",
+    "\n",
+    "        # Create one-hot encoding\n",
+    "        one_hot = F.one_hot(\n",
+    "            token_ids[0], num_classes=self.embedding_layer.num_embeddings\n",
+    "        ).float()\n",
+    "        one_hot.requires_grad_(True)\n",
+    "\n",
+    "        # Replace the embedding lookup by one_hot @ weight, which is differentiable\n",
+    "        def embed_one_hot(module, inputs, output):\n",
+    "            return (one_hot @ module.weight).unsqueeze(0).to(output.dtype)\n",
+    "\n",
+    "        # Forward pass with gradient tracking; the token tensor is passed positionally\n",
+    "        handle = self.embedding_layer.register_forward_hook(embed_one_hot)\n",
+    "        try:\n",
+    "            outputs = self.model(\n",
+    "                token_ids,\n",
+    "                conditions_tokens=[self.species_ids],\n",
+    "                track_indices=self.track_indices,\n",
+    "                skip_lm_head=True,\n",
+    "            )\n",
+    "        finally:\n",
+    "            handle.remove()\n",
+    "\n",
+    "        # [batch, positions, len(track_indices)], read from the returned dict\n",
+    "        logits = outputs[\"bigwig_tracks_logits\"]\n",
+    "\n",
+    "        # Optionally slice output positions (output is at 1 bp resolution)\n",
+    "        if self.target_window_bp is not None:\n",
+    "            start, end = self.target_window_bp\n",
+    "            seq_len = logits.shape[1]\n",
+    "            start = max(0, min(start, seq_len))\n",
+    "            end = max(start, min(end, seq_len))\n",
+    "            logits = logits[:, start:end, :]\n",
+    "\n",
+    "        score = torch.log(logits.sum() + 1e-6)\n",
+    "\n",
+    "        # Backward pass\n",
+    "        score.backward()\n",
+    "\n",
+    "        gradients = one_hot.grad.cpu().detach().numpy()\n",
+    "        one_hot_np = one_hot.detach().cpu().numpy()\n",
+    "\n",
+    "        return gradients, one_hot_np\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "_YwKWAj-DK7o"
+   },
+   "outputs": [],
+   "source": [
+    "# Select track for saliency computation\n",
+    "# The track id must exist for this assembly\n",
+    "saliency_track_id = \"ENCSR000EFT\"  # K562 GATA1 ChIP-seq\n",
+    "\n",
+    "# Resolve track index\n",
+    "if saliency_track_id not in bigwig_names:\n",
+    "    raise ValueError(f\"Track not available for {assembly}: {saliency_track_id!r}\")\n",
+    "saliency_track_idx = bigwig_names.index(saliency_track_id)\n",
+    "print(\n",
+    "    f\"Computing saliency for track: {saliency_track_id} \"\n",
+    "    f\"(index {saliency_track_idx})\"\n",
+    ")\n",
+    "\n",
+    "# Use shorter sequence for saliency (faster computation)\n",
+    "saliency_seq_length = len(seq) // 2\n",
+    "seq_start = (len(seq) - saliency_seq_length) // 2\n",
+    "seq_end = seq_start + saliency_seq_length\n",
+    "saliency_seq = seq[seq_start:seq_end]\n",
+    "\n",
+    "print(\n",
+    "    f\"Saliency sequence: {len(saliency_seq):,} bp \"\n",
+    "    f\"(center crop of full sequence)\"\n",
+    ")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "G44_fw_bDK7o"
+   },
+   "outputs": [],
+   "source": [
+    "# Define target window: center +/- 1024 bp of the output axis\n",
+    "promoter_window_bp = 1024\n",
+    "output_len = int(len(saliency_seq) * 0.375)  # predicted center span at 1 bp resolution\n",
+    "center = output_len // 2\n",
+    "target_window_bp = (center - promoter_window_bp, center + promoter_window_bp + 1)\n",
+    "\n",
+    "# Initialize saliency computer\n",
+    "saliency_computer = SaliencyComputer(\n",
+    "    model=model,\n",
+    "    tokenizer=tokenizer,\n",
+    "    track_indices=[saliency_track_idx],\n",
+    "    species=species,\n",
+    "    device=device,\n",
+    "    target_window_bp=target_window_bp,\n",
+    ")\n",
+    "\n",
+    "# Compute saliency\n",
+    "print(\"Computing saliency (this may take a moment)...\")\n",
+    "saliency_scores, one_hot_sequence = saliency_computer.compute_saliency(\n",
+    "    saliency_seq\n",
+    ")\n",
+    "\n",
+    "print(f\"\\n✅ Saliency computed\")\n",
+    "print(f\"   Shape: {saliency_scores.shape}\")\n",
+    "sal_min = saliency_scores.min()\n",
+    "sal_max = saliency_scores.max()\n",
+    "print(f\"   Value range: [{sal_min:.6f}, {sal_max:.6f}]\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "A-eR1tK9DK7o"
+   },
+   "outputs": [],
+   "source": [
+    "# Nucleotide channels in logo order (A, T, C, G), upper case and soft-masked\n",
+    "upper_ids = tokenizer(\"ATCG\", add_special_tokens=False)[\"input_ids\"]\n",
+    "lower_ids = tokenizer(\"atcg\", add_special_tokens=False)[\"input_ids\"]\n",
+    "assert tokenizer.unk_token_id not in upper_ids + lower_ids\n",
+    "\n",
+    "# Read each position from the channels of its own case: (positions, 4)\n",
+    "soft_masked = np.array([base.islower() for base in saliency_seq])\n",
+    "channels = np.where(soft_masked[:, None], lower_ids, upper_ids)\n",
+    "gradients_acgt = np.take_along_axis(saliency_scores, channels, axis=1)\n",
+    "one_hot_acgt = np.take_along_axis(one_hot_sequence, channels, axis=1)\n",
+    "\n",
+    "# Apply gradient correction (subtract mean over the four nucleotides at each position)\n",
+    "saliency_corrected = gradients_acgt - np.mean(\n",
+    "    gradients_acgt, axis=1, keepdims=True\n",
+    ")\n",
+    "\n",
+    "# Calculate genomic positions\n",
+    "saliency_start = start + seq_start\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "mU9TdySBDK7o"
+   },
+   "outputs": [],
+   "source": [
+    "# Plot saliency sequence logo\n",
+    "from fast_logomaker import FastLogo\n",
+    "\n",
+    "# Nucleotide color scheme\n",
+    "NUCLEOTIDE_COLORS = {\n",
+    "    'A': '#00b050',  # green\n",
+    "    'C': '#3973fc',  # blue\n",
+    "    'G': '#fdb932',  # yellow/gold\n",
+    "    'T': '#ff554d'   # red\n",
+    "}\n",
+    "\n",
+    "# True = contribution mode (multiply by one-hot), False = hypothetical mode\n",
+    "CONTRIBUTION_MODE = True\n",
+    "LOGO_WINDOW_SIZE = 128  # Number of positions to show in logo\n",
+    "\n",
+    "# Saliency is at nucleotide resolution. Channels are already A, T, C, G.\n",
+    "n_saliency = len(saliency_corrected)\n",
+    "cropped_start = (n_saliency - LOGO_WINDOW_SIZE) // 2\n",
+    "cropped_end = cropped_start + LOGO_WINDOW_SIZE\n",
+    "window_start_bp = saliency_start + cropped_start\n",
+    "window_end_bp = saliency_start + cropped_end\n",
+    "\n",
+    "print(f\"Plotting saliency sequence logo...\")\n",
+    "print(f\"  Contribution mode: {CONTRIBUTION_MODE}\")\n",
+    "print(f\"  Window: {window_start_bp:,} - {window_end_bp:,} bp\")\n",
+    "\n",
+    "saliency_acgt = saliency_corrected\n",
+    "if CONTRIBUTION_MODE:\n",
+    "    saliency_acgt = saliency_acgt * one_hot_acgt\n",
+    "\n",
+    "window_saliency = saliency_acgt[cropped_start:cropped_end]\n",
+    "print(f\"  Window size: {len(window_saliency)} positions\")\n",
+    "\n",
+    "logo_processor = FastLogo(\n",
+    "    window_saliency[np.newaxis, :, :],\n",
+    "    alphabet=['A', 'T', 'C', 'G'],\n",
+    "    font_name='DejaVu Sans',\n",
+    "    font_weight='bold',\n",
+    "    color_scheme=NUCLEOTIDE_COLORS,\n",
+    "    fade_below=0.0,\n",
+    "    shade_below=0.0,\n",
+    "    width=0.9,\n",
+    "    figsize=(14, 3),\n",
+    "    center_values=not CONTRIBUTION_MODE,\n",
+    "    border=False,\n",
+    "    fixed_ylim=True,\n",
+    "    batch_size=1,\n",
+    ")\n",
+    "logo_processor.process_all()\n",
+    "fig, ax = logo_processor.draw_single(0, border=False)\n",
+    "logo_title = f'{chrom}:{window_start_bp:,}-{window_end_bp:,}'\n",
+    "ax.set_title(logo_title, fontsize=11)\n",
+    "\n",
+    "n_positions = len(window_saliency)\n",
+    "tick_positions = [0, n_positions // 4, n_positions // 2,\n",
+    "                  3 * n_positions // 4, n_positions - 1]\n",
+    "tick_labels = [f'{window_start_bp + pos:,}' for pos in tick_positions]\n",
+    "ax.set_xticks(tick_positions)\n",
+    "ax.set_xticklabels(tick_labels, fontsize=8)\n",
+    "ax.set_xlabel('Genomic Position', fontsize=9)\n",
+    "\n",
+    "plt.yticks([])\n",
+    "plt.tight_layout()\n",
+    "plt.show()\n",
+    "print(\"\\n✅ Saliency logo plotted\")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "> **How to read this plot.** The x axis is the genomic coordinate on `chrom`, for the central logo window of the saliency sequence. The score is the log of the predicted `saliency_track_id` signal summed over `target_window_bp`, ±`promoter_window_bp` bp (1,024 bp) around the center of the predicted span, not over the whole span. The logo window sits inside it. In contribution mode, each letter is the reference base and its height is the gradient after subtracting the mean over the four bases. A tall letter above the axis means that base raises the predicted signal. A letter below the axis means it lowers it. Saliency ranks positions. It does not predict the effect of a mutation.\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "JF7UOTNHDK7o"
+   },
+   "source": [
+    "## 📝 Summary\n",
+    "\n",
+    "This notebook demonstrated three complementary interpretability modalities of the NTv3 tracks model:\n",
+    "\n",
+    "- ✅ **Functional tracks**: Predicted regulatory activity across thousands of assays\n",
+    "- ✅ **Attention maps**: Long-range information integration patterns\n",
+    "- ✅ **Saliency maps**: Gradient-based sequence attributions\n",
+    "\n",
+    "### 💡 Next Steps\n",
+    "\n",
+    "- Try different genomic regions by changing `chrom`, `center_position`, `sequence_length`\n",
+    "- Try another species by changing `species`; `assembly` must be its entry in `SPECIES_TO_ASSEMBLY`\n",
+    "- Explore different tracks by modifying `tracks_to_plot` and `saliency_track_id`; the ids must be in that assembly's `bigwig_names`\n",
+    "- Compare attention patterns across different layers\n",
+    "- Compute saliency for different track combinations\n",
+    "\n",
+    "### 📚 Additional Resources\n",
+    "\n",
+    "- For track prediction tutorials, see: [01_tracks_prediction.ipynb](https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_tutorials/01_tracks_prediction.ipynb)\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "accelerator": "GPU",
+  "colab": {
+   "gpuType": "T4",
+   "machine_shape": "hm",
+   "provenance": []
+  },
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.13.15"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}

notebooks_tutorials/05_diffusion_fine_tuning.ipynb

@@ -0,0 +1,941 @@
+{
+ "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
+}

src/genome_annotation/__init__.py

@@ -0,0 +1 @@
+"""Helpers of the NTv3 genome-annotation tutorial notebook."""

src/genome_annotation/analysis.py

@@ -0,0 +1,594 @@
+"""Error analyses of a predicted annotation, as in the NTv3 annotation paper.
+
+Everything works on ``evaluation.ScoringSets``, i.e. after the benchmark filters.
+A locus is one reference gene; a prediction overlaps it when the two share coding
+sequence on the same strand (a prediction inside an intron does not).
+"""
+
+import bisect
+from collections import defaultdict
+from collections.abc import Iterable
+from dataclasses import dataclass
+
+from .evaluation import CdsKey, Gene
+
+LOCUS_CLASSES = ("correct", "boundary", "missed", "fragmented")
+EXON_CLASSES = ("exact", "shifted_one", "shifted_both", "missed")
+POSITIONS = ("first", "internal", "last", "single")
+
+# Inclusive bins unless noted.
+EXON_COUNT_BINS = ((1, 1), (2, 3), (4, 6), (7, 10), (11, 20), (21, 10**9))
+SPAN_BINS = (  # half-open, bp of CDS envelope
+    (0, 5_000),
+    (5_000, 20_000),
+    (20_000, 50_000),
+    (50_000, 131_072),
+    (131_072, 500_000),
+    (500_000, 10**12),
+)
+EXON_LENGTH_BINS = (
+    (1, 50),
+    (51, 100),
+    (101, 200),
+    (201, 500),
+    (501, 1_000),
+    (1_001, 10**9),
+)
+WRONG_EXON_BUCKETS = ((0, 0), (1, 1), (2, 2), (3, 3), (4, 5), (6, 10**9))
+OFFSET_BINS = (
+    (1, 3),
+    (4, 6),
+    (7, 12),
+    (13, 25),
+    (26, 100),
+    (101, 1_000),
+    (1_001, 10**9),
+)
+MERGE_SIZE_BINS = ((2, 2), (3, 3), (4, 5), (6, 10**9))
+
+
+def bin_label(lo: int, hi: int) -> str:
+    """``"4-6"``, ``"1"``, or ``"21+"`` for the open-ended bin."""
+    if lo == hi:
+        return str(lo)
+    return f"{lo}+" if hi >= 10**9 else f"{lo}-{hi}"
+
+
+def _bin_index(value: int, bins: tuple[tuple[int, int], ...]) -> int:
+    return next(
+        (i for i, (lo, hi) in enumerate(bins) if lo <= value <= hi), len(bins) - 1
+    )
+
+
+def _half_open_bin_index(value: int, bins: tuple[tuple[int, int], ...]) -> int:
+    return next(
+        (i for i, (lo, hi) in enumerate(bins) if lo <= value < hi), len(bins) - 1
+    )
+
+
+def cds_span(gene: Gene) -> tuple[int, int]:
+    """CDS envelope; the gene row would include UTRs the predictions lack."""
+    return min(k[1] for k in gene.cds_keys), max(k[2] for k in gene.cds_keys)
+
+
+def _boundary(key: CdsKey) -> tuple[str, int, int, str]:
+    return key[:4]
+
+
+def _ordered_cds(gene: Gene) -> list[CdsKey]:
+    """CDS in transcription order, 5' to 3'."""
+    return sorted(gene.cds_keys, key=lambda k: k[1], reverse=gene.strand == "-")
+
+
+def _position(index: int, total: int) -> str:
+    if total == 1:
+        return "single"
+    if index == 0:
+        return "first"
+    return "last" if index == total - 1 else "internal"
+
+
+class CdsIndex:
+    """Same-strand overlap lookup over individual CDS, returning gene IDs."""
+
+    def __init__(self, genes: dict[str, Gene]) -> None:
+        buckets: dict[tuple[str, str], list[tuple[int, int, str]]] = defaultdict(list)
+        for gid, gene in genes.items():
+            for _, start, end, strand, _ in gene.cds_keys:
+                buckets[(gene.chrom, strand)].append((start, end, gid))
+        self._items: dict[tuple[str, str], list[tuple[int, int, str]]] = {}
+        self._starts: dict[tuple[str, str], list[int]] = {}
+        self._max_end: dict[tuple[str, str], list[int]] = {}
+        for key, items in buckets.items():
+            items.sort()
+            self._items[key] = items
+            self._starts[key] = [item[0] for item in items]
+            running, best = [], 0
+            for _, end, _ in items:
+                best = max(best, end)
+                running.append(best)
+            self._max_end[key] = running
+
+    def overlapping(self, chrom: str, strand: str, start: int, end: int) -> set[str]:
+        """IDs of genes with a CDS overlapping [start, end] on this strand."""
+        items = self._items.get((chrom, strand))
+        if items is None:
+            return set()
+        max_end = self._max_end[(chrom, strand)]
+        hits: set[str] = set()
+        for i in range(
+            bisect.bisect_right(self._starts[(chrom, strand)], end) - 1, -1, -1
+        ):
+            if max_end[i] < start:
+                break
+            if items[i][1] >= start:
+                hits.add(items[i][2])
+        return hits
+
+    def overlapping_gene(self, gene: Gene) -> set[str]:
+        """IDs of indexed genes sharing coding sequence with ``gene``."""
+        hits: set[str] = set()
+        for _, start, end, strand, _ in gene.cds_keys:
+            hits |= self.overlapping(gene.chrom, strand, start, end)
+        return hits
+
+
+@dataclass(frozen=True)
+class LocusOutcome:
+    """One reference gene, classified against the predictions."""
+
+    gene_id: str
+    chrom: str
+    strand: str
+    start: int
+    end: int
+    n_cds: int
+    cds_length: int
+    outcome: str
+    n_overlapping: int
+
+
+def classify_loci(
+    ref_genes: dict[str, Gene], pred_genes: dict[str, Gene]
+) -> tuple[list[LocusOutcome], list[str]]:
+    """Classifies every reference locus; also returns the invented predictions.
+
+    ``correct``: its exact CDS set is predicted. Otherwise by the number of
+    overlapping predictions: 0 ``missed``, 1 ``boundary``, 2+ ``fragmented``.
+    Invented: a prediction overlapping no reference gene.
+    """
+    pred_structures = {frozenset(g.cds_keys) for g in pred_genes.values()}
+    pred_index, ref_index = CdsIndex(pred_genes), CdsIndex(ref_genes)
+    outcomes = []
+    for gid, gene in ref_genes.items():
+        if not gene.cds_keys:
+            continue
+        hits = pred_index.overlapping_gene(gene)
+        if frozenset(gene.cds_keys) in pred_structures:
+            outcome = "correct"
+        else:
+            outcome = (
+                ("missed", "boundary")[len(hits)] if len(hits) < 2 else "fragmented"
+            )
+        outcomes.append(
+            LocusOutcome(
+                gid,
+                gene.chrom,
+                gene.strand,
+                *cds_span(gene),
+                len(gene.cds_keys),
+                sum(k[2] - k[1] + 1 for k in gene.cds_keys),
+                outcome,
+                len(hits),
+            )
+        )
+    invented = [
+        gid
+        for gid, gene in pred_genes.items()
+        if gene.cds_keys and not ref_index.overlapping_gene(gene)
+    ]
+    return outcomes, invented
+
+
+def summarise_loci(
+    outcomes: list[LocusOutcome], invented: list[str], n_pred: int
+) -> dict[str, float]:
+    """Counts per class; precision here is correct loci over predictions."""
+    counts = dict.fromkeys(LOCUS_CLASSES, 0)
+    for outcome in outcomes:
+        counts[outcome.outcome] += 1
+    n_loci = len(outcomes)
+    recall = 100.0 * counts["correct"] / n_loci if n_loci else 0.0
+    precision = 100.0 * counts["correct"] / n_pred if n_pred else 0.0
+    total = recall + precision
+    return {
+        **counts,
+        "n_loci": n_loci,
+        "n_pred": n_pred,
+        "invented": len(invented),
+        "recall": recall,
+        "precision": precision,
+        "gene_f1": 2 * recall * precision / total if total else 0.0,
+    }
+
+
+def recovery_by_covariate(outcomes: list[LocusOutcome]) -> dict[str, object]:
+    """Loci and outcomes per exon-count and CDS-span bin."""
+    exon_total = [0] * len(EXON_COUNT_BINS)
+    exon_correct = [0] * len(EXON_COUNT_BINS)
+    span_total = [0] * len(SPAN_BINS)
+    span_correct = [0] * len(SPAN_BINS)
+    span_outcome = {cls: [0] * len(SPAN_BINS) for cls in LOCUS_CLASSES}
+    for outcome in outcomes:
+        i = _bin_index(outcome.n_cds, EXON_COUNT_BINS)
+        j = _half_open_bin_index(outcome.end - outcome.start + 1, SPAN_BINS)
+        exon_total[i] += 1
+        span_total[j] += 1
+        span_outcome[outcome.outcome][j] += 1
+        if outcome.outcome == "correct":
+            exon_correct[i] += 1
+            span_correct[j] += 1
+    return {
+        "exon_count_bins": [bin_label(lo, hi) for lo, hi in EXON_COUNT_BINS],
+        "exon_total": exon_total,
+        "exon_correct": exon_correct,
+        "span_bins": [f"{lo}-{hi}" for lo, hi in SPAN_BINS],
+        "span_total": span_total,
+        "span_correct": span_correct,
+        "span_outcome": span_outcome,
+    }
+
+
+def _overlapping_cds(
+    pred_genes: dict[str, Gene],
+    pred_index: CdsIndex,
+    chrom: str,
+    strand: str,
+    start: int,
+    end: int,
+) -> list[CdsKey]:
+    return [
+        k
+        for gid in sorted(pred_index.overlapping(chrom, strand, start, end))
+        for k in pred_genes[gid].cds_keys
+        if k[0] == chrom and k[3] == strand and k[1] <= end and k[2] >= start
+    ]
+
+
+def _best_overlap(candidates: list[CdsKey], start: int, end: int) -> CdsKey:
+    """Largest overlap; ties broken by coordinate so the result is deterministic."""
+    return max(
+        candidates, key=lambda k: (min(end, k[2]) - max(start, k[1]), -k[1], -k[2])
+    )
+
+
+def exon_taxonomy(
+    ref_genes: dict[str, Gene],
+    pred_genes: dict[str, Gene],
+    locus_outcome: dict[str, str],
+) -> dict[str, object]:
+    """Classifies every reference CDS (phase ignored), by position and length.
+
+    ``exact``: some predicted CDS has both boundaries. Otherwise the most-overlapping
+    same-strand predicted CDS has one (``shifted_one``) or both (``shifted_both``)
+    boundaries wrong, or nothing overlaps (``missed``).
+    """
+    pred_index = CdsIndex(pred_genes)
+    pred_boundaries = {_boundary(k) for g in pred_genes.values() for k in g.cds_keys}
+    by_position = {pos: dict.fromkeys(EXON_CLASSES, 0) for pos in POSITIONS}
+    by_length = [dict.fromkeys(EXON_CLASSES, 0) for _ in EXON_LENGTH_BINS]
+    ends = dict.fromkeys(("five_prime", "three_prime", "both"), 0)
+    wrong_per_failed_gene = [0] * len(WRONG_EXON_BUCKETS)
+    n_ref_cds = 0
+
+    for gid, gene in ref_genes.items():
+        ordered = _ordered_cds(gene)
+        wrong = 0
+        for index, key in enumerate(ordered):
+            chrom, start, end, strand, _ = key
+            n_ref_cds += 1
+            position = _position(index, len(ordered))
+            length_bin = by_length[_bin_index(end - start + 1, EXON_LENGTH_BINS)]
+            if _boundary(key) in pred_boundaries:
+                klass = "exact"
+            else:
+                wrong += 1
+                candidates = _overlapping_cds(
+                    pred_genes, pred_index, chrom, strand, start, end
+                )
+                if not candidates:
+                    klass = "missed"
+                else:
+                    best = _best_overlap(candidates, start, end)
+                    start_wrong, end_wrong = best[1] != start, best[2] != end
+                    klass = (
+                        "shifted_both" if start_wrong and end_wrong else "shifted_one"
+                    )
+                    # Transcript orientation: a minus-strand start is the 3' end.
+                    if start_wrong and end_wrong:
+                        ends["both"] += 1
+                    elif start_wrong:
+                        ends["five_prime" if strand == "+" else "three_prime"] += 1
+                    else:
+                        ends["three_prime" if strand == "+" else "five_prime"] += 1
+            by_position[position][klass] += 1
+            length_bin[klass] += 1
+        if locus_outcome.get(gid) != "correct":
+            wrong_per_failed_gene[_bin_index(wrong, WRONG_EXON_BUCKETS)] += 1
+
+    ref_boundaries = {_boundary(k) for g in ref_genes.values() for k in g.cds_keys}
+    ref_index = CdsIndex(ref_genes)
+    pred_keys = [k for g in pred_genes.values() for k in g.cds_keys]
+    invented = sum(
+        1
+        for k in pred_keys
+        if _boundary(k) not in ref_boundaries
+        and not ref_index.overlapping(k[0], k[3], k[1], k[2])
+    )
+    totals = {
+        klass: sum(c[klass] for c in by_position.values()) for klass in EXON_CLASSES
+    }
+    return {
+        "n_ref_cds": n_ref_cds,
+        "n_pred_cds": len(pred_keys),
+        "totals": totals,
+        "by_position": by_position,
+        "exon_length_bins": [bin_label(lo, hi) for lo, hi in EXON_LENGTH_BINS],
+        "by_length": by_length,
+        "ends": ends,
+        "invented": invented,
+        "wrong_exon_buckets": [bin_label(lo, hi) for lo, hi in WRONG_EXON_BUCKETS],
+        "wrong_exons_per_failed_gene": wrong_per_failed_gene,
+    }
+
+
+def boundary_offsets(
+    ref_genes: dict[str, Gene], pred_genes: dict[str, Gene]
+) -> dict[str, object]:
+    """How far off shifted boundaries are, over reference CDS a prediction overlaps."""
+    pred_index = CdsIndex(pred_genes)
+    aligned = dict.fromkeys(POSITIONS, 0)
+    shifted = dict.fromkeys(POSITIONS, 0)
+    offsets = [0] * len(OFFSET_BINS)
+    length_delta_mod3 = {0: 0, 1: 0, 2: 0}
+    for gene in ref_genes.values():
+        ordered = _ordered_cds(gene)
+        for index, (chrom, start, end, strand, _) in enumerate(ordered):
+            candidates = _overlapping_cds(
+                pred_genes, pred_index, chrom, strand, start, end
+            )
+            if not candidates:
+                continue
+            best = _best_overlap(candidates, start, end)
+            position = _position(index, len(ordered))
+            aligned[position] += 1
+            if best[1] == start and best[2] == end:
+                continue
+            shifted[position] += 1
+            for offset in (abs(best[1] - start), abs(best[2] - end)):
+                if offset:
+                    offsets[_bin_index(offset, OFFSET_BINS)] += 1
+            length_delta_mod3[((best[2] - best[1]) - (end - start)) % 3] += 1
+    return {
+        "aligned": aligned,
+        "shifted": shifted,
+        "offset_bins": [bin_label(lo, hi) for lo, hi in OFFSET_BINS],
+        "offsets": offsets,
+        "length_delta_mod3": length_delta_mod3,
+    }
+
+
+def failure_anatomy(
+    ref_genes: dict[str, Gene],
+    pred_genes: dict[str, Gene],
+    locus_outcome: dict[str, str],
+) -> dict[str, object]:
+    """Wrong exons per failed gene, by gene size, and where a lone wrong exon sits."""
+    boundaries = {_boundary(k) for g in pred_genes.values() for k in g.cds_keys}
+    grid = [[0] * len(WRONG_EXON_BUCKETS) for _ in EXON_COUNT_BINS]
+    observed = dict.fromkeys(POSITIONS, 0)
+    expected = dict.fromkeys(POSITIONS, 0.0)
+    convertible = [0] * len(WRONG_EXON_BUCKETS)
+    for gid, gene in ref_genes.items():
+        if locus_outcome.get(gid) == "correct":
+            continue
+        ordered = _ordered_cds(gene)
+        n = len(ordered)
+        wrong = [
+            _position(i, n)
+            for i, k in enumerate(ordered)
+            if _boundary(k) not in boundaries
+        ]
+        bucket = _bin_index(len(wrong), WRONG_EXON_BUCKETS)
+        grid[_bin_index(n, EXON_COUNT_BINS)][bucket] += 1
+        for i in range(bucket, len(convertible)):
+            convertible[i] += 1
+        if len(wrong) == 1:
+            observed[wrong[0]] += 1
+            if n == 1:
+                expected["single"] += 1.0
+            else:  # Uniform over the gene's exons.
+                expected["first"] += 1.0 / n
+                expected["last"] += 1.0 / n
+                expected["internal"] += (n - 2) / n
+    return {
+        "exon_count_bins": [bin_label(lo, hi) for lo, hi in EXON_COUNT_BINS],
+        "wrong_exon_buckets": [bin_label(lo, hi) for lo, hi in WRONG_EXON_BUCKETS],
+        "wrong_by_gene_size": grid,
+        "cumulative_convertible": convertible,
+        "single_wrong_observed": observed,
+        "single_wrong_expected": {k: round(v, 2) for k, v in expected.items()},
+    }
+
+
+def independence_check(
+    ref_genes: dict[str, Gene],
+    pred_genes: dict[str, Gene],
+    locus_outcome: dict[str, str],
+) -> dict[str, object]:
+    """Observed gene recovery vs. independent exon errors, per exon-count bin.
+
+    Predicted recovery of an n-exon gene is P(first) * P(internal)^(n-2) * P(last),
+    with per-position exact rates measured on the same predictions.
+    """
+    boundaries = {_boundary(k) for g in pred_genes.values() for k in g.cds_keys}
+    hit: dict[str, int] = defaultdict(int)
+    total: dict[str, int] = defaultdict(int)
+    for gene in ref_genes.values():
+        ordered = _ordered_cds(gene)
+        for index, key in enumerate(ordered):
+            position = _position(index, len(ordered))
+            total[position] += 1
+            hit[position] += _boundary(key) in boundaries
+    rate = {p: hit[p] / total[p] for p in total}
+
+    def predicted(n: int) -> float:
+        if n == 1:
+            return rate.get("single", 0.0)
+        return (
+            rate.get("first", 0.0)
+            * rate.get("last", 0.0)
+            * rate.get("internal", 0.0) ** (n - 2)
+        )
+
+    observed = [0] * len(EXON_COUNT_BINS)
+    n_genes = [0] * len(EXON_COUNT_BINS)
+    expected = [0.0] * len(EXON_COUNT_BINS)
+    for gid, gene in ref_genes.items():
+        n = len(gene.cds_keys)
+        i = _bin_index(n, EXON_COUNT_BINS)
+        n_genes[i] += 1
+        expected[i] += predicted(n)
+        observed[i] += locus_outcome.get(gid) == "correct"
+    return {
+        "exon_count_bins": [bin_label(lo, hi) for lo, hi in EXON_COUNT_BINS],
+        "n_genes": n_genes,
+        "observed": observed,
+        "predicted": [round(v, 3) for v in expected],
+        "position_rate": {p: round(v, 6) for p, v in sorted(rate.items())},
+    }
+
+
+def intron_chain(gene: Gene) -> tuple[object, ...]:
+    """Internal CDS boundaries (gffcompare's intron chain); phase ignored."""
+    ordered = sorted(gene.cds_keys, key=lambda k: (k[1], k[2]))
+    introns = tuple((ordered[i][2], ordered[i + 1][1]) for i in range(len(ordered) - 1))
+    return (gene.chrom, gene.strand, introns)
+
+
+def _rates(
+    tp_ref: int, n_ref: int, tp_pred: int, n_pred: int, prefix: str
+) -> dict[str, float]:
+    recall = 100.0 * tp_ref / n_ref if n_ref else 0.0
+    precision = 100.0 * tp_pred / n_pred if n_pred else 0.0
+    total = recall + precision
+    return {
+        f"{prefix}_recall": recall,
+        f"{prefix}_precision": precision,
+        f"{prefix}_f1": 2 * recall * precision / total if total else 0.0,
+    }
+
+
+def _chain_counts(
+    refs: list[Gene], preds: list[Gene], chains: bool
+) -> dict[str, int]:
+    ref_exact = {frozenset(g.cds_keys) for g in refs}
+    pred_exact = {frozenset(g.cds_keys) for g in preds}
+    counts = {
+        "n_ref": len(refs),
+        "n_pred": len(preds),
+        "exact_tp_ref": sum(frozenset(g.cds_keys) in pred_exact for g in refs),
+        "exact_tp_pred": sum(frozenset(g.cds_keys) in ref_exact for g in preds),
+    }
+    if chains:
+        ref_chains = {intron_chain(g) for g in refs}
+        pred_chains = {intron_chain(g) for g in preds}
+        counts["chain_tp_ref"] = sum(intron_chain(g) in pred_chains for g in refs)
+        counts["chain_tp_pred"] = sum(intron_chain(g) in ref_chains for g in preds)
+    else:  # Single-exon genes have no chain; keep the exact tally so rollups add up.
+        counts["chain_tp_ref"] = counts["exact_tp_ref"]
+        counts["chain_tp_pred"] = counts["exact_tp_pred"]
+    return counts
+
+
+def _with_rates(counts: dict[str, int]) -> dict[str, float]:
+    return {
+        **counts,
+        **_rates(
+            counts["exact_tp_ref"],
+            counts["n_ref"],
+            counts["exact_tp_pred"],
+            counts["n_pred"],
+            "exact",
+        ),
+        **_rates(
+            counts["chain_tp_ref"],
+            counts["n_ref"],
+            counts["chain_tp_pred"],
+            counts["n_pred"],
+            "chain",
+        ),
+    }
+
+
+def measure_chains(
+    ref_genes: dict[str, Gene], pred_genes: dict[str, Gene]
+) -> dict[str, object]:
+    """Exact-gene and intron-chain matching, for multi- and single-exon genes."""
+
+    def split(genes: Iterable[Gene]) -> tuple[list[Gene], list[Gene]]:
+        genes = [g for g in genes if g.cds_keys]
+        return [g for g in genes if len(g.cds_keys) > 1], [
+            g for g in genes if len(g.cds_keys) == 1
+        ]
+
+    ref_multi, ref_single = split(ref_genes.values())
+    pred_multi, pred_single = split(pred_genes.values())
+    multi = _chain_counts(ref_multi, pred_multi, chains=True)
+    single = _chain_counts(ref_single, pred_single, chains=False)
+    combined = {key: multi[key] + single[key] for key in multi}
+    return {
+        "multi_exon": _with_rates(multi),
+        "single_exon": _with_rates(single),
+        "all": _with_rates(combined),
+    }
+
+
+def measure_merges(
+    ref_genes: dict[str, Gene],
+    pred_genes: dict[str, Gene],
+    outcomes: list[LocusOutcome],
+) -> dict[str, object]:
+    """Predictions overlapping 2+ reference loci, and what they absorb."""
+    ref_index = CdsIndex(ref_genes)
+    outcome_of = {o.gene_id: o.outcome for o in outcomes}
+    sizes = dict.fromkeys((bin_label(lo, hi) for lo, hi in MERGE_SIZE_BINS), 0)
+    merging, largest = 0, 0
+    absorbed: set[str] = set()
+    for gene in pred_genes.values():
+        if not gene.cds_keys:
+            continue
+        hits = ref_index.overlapping_gene(gene)
+        if len(hits) < 2:
+            continue
+        merging += 1
+        lo, hi = MERGE_SIZE_BINS[_bin_index(len(hits), MERGE_SIZE_BINS)]
+        sizes[bin_label(lo, hi)] += 1
+        largest = max(largest, len(hits))
+        absorbed |= hits
+    absorbed_outcomes: dict[str, int] = defaultdict(int)
+    for gid in absorbed:
+        absorbed_outcomes[outcome_of.get(gid, "unscored")] += 1
+    return {
+        "n_pred": sum(1 for g in pred_genes.values() if g.cds_keys),
+        "n_ref": len(outcomes),
+        "merging_predictions": merging,
+        "absorbed_loci": len(absorbed),
+        "merge_sizes": sizes,
+        "largest_merge": largest,
+        "boundary_loci": sum(o.outcome == "boundary" for o in outcomes),
+        "boundary_from_merge": sum(
+            o.outcome == "boundary" and o.gene_id in absorbed for o in outcomes
+        ),
+        "absorbed_outcomes": dict(sorted(absorbed_outcomes.items())),
+    }

src/genome_annotation/benchmark.py

@@ -0,0 +1,194 @@
+"""The NTv3 annotation benchmark: test species, released files and published scores."""
+
+from dataclasses import dataclass
+
+DATASET_REPO = "InstaDeepAI/NTv3_annot_dataset"
+
+
+@dataclass(frozen=True)
+class TestSpecies:
+    """One held-out species of the benchmark."""
+
+    name: str
+    clade: str
+    assembly: str
+    released_gff: str
+    gene_f1: float
+    cds_f1: float
+
+    @property
+    def model_repo(self) -> str:
+        """Hugging Face repo of the clade model, which also holds its predictions."""
+        return f"InstaDeepAI/NTv3_annot_{self.clade}"
+
+    @property
+    def dataset_dir(self) -> str:
+        """Folder of the species in ``DATASET_REPO``."""
+        return f"{self.clade}/test_species/{self.assembly}"
+
+    @property
+    def reference_gff(self) -> str:
+        """Benchmark reference in ``DATASET_REPO``."""
+        return f"{self.dataset_dir}/reference_longest_isoform.gff3.gz"
+
+    @property
+    def genome_fasta(self) -> str:
+        """Soft-masked genome in ``DATASET_REPO`` (bgzipped, indexed)."""
+        return f"{self.dataset_dir}/genome.fa.gz"
+
+
+def _species(
+    name: str, clade: str, assembly: str, tag: str, gene_f1: float, cds_f1: float
+) -> TestSpecies:
+    return TestSpecies(
+        name,
+        clade,
+        assembly,
+        f"annot_results/ntv3_{name}_{tag}.gff3.gz",
+        gene_f1,
+        cds_f1,
+    )
+
+
+# Gene / CDS F1 (%) of the released NTv3 GFF3s, whole genome, as in the paper.
+TEST_SPECIES = {
+    s.name: s
+    for s in (
+        _species(
+            "bos_taurus",
+            "mammals",
+            "GCF_002263795.3_ARS-UCD2.0",
+            "GCF_002263795.3",
+            58.06,
+            90.60,
+        ),
+        _species(
+            "delphinapterus_leucas",
+            "mammals",
+            "GCF_002288925.2_ASM228892v3",
+            "GCF_002288925.2",
+            54.66,
+            89.68,
+        ),
+        _species(
+            "brachypodium_distachyon",
+            "plants",
+            "GCF_000005505.3_Brachypodium_distachyon_v3.0",
+            "GCF_000005505.3",
+            71.89,
+            89.68,
+        ),
+        _species(
+            "eschscholzia_californica",
+            "plants",
+            "EcalifornicavarOrangeKingHAP1v1.1",
+            "OrangeKingHAP1v1.1",
+            73.17,
+            91.07,
+        ),
+        _species(
+            "freycinetia_multiflora",
+            "plants",
+            "FmultifloravarABG20040013HAP1v1.1",
+            "ABG20040013HAP1v1.1",
+            68.36,
+            88.73,
+        ),
+        _species(
+            "medicago_truncatula",
+            "plants",
+            "GCF_003473485.1_MtrunA17r5.0-ANR",
+            "GCF_003473485.1",
+            72.63,
+            89.33,
+        ),
+        _species(
+            "mimulus_guttatus",
+            "plants",
+            "MguttatusvarIM62v3.1",
+            "IM62v3.1",
+            80.73,
+            93.54,
+        ),
+        _species(
+            "bombyx_mori",
+            "insects",
+            "GCF_030269925.1_ASM3026992v2",
+            "GCF_030269925.1",
+            56.53,
+            85.19,
+        ),
+        _species(
+            "cataglyphis_hispanica",
+            "insects",
+            "GCF_021464435.1_ULB_Chis1_1.0",
+            "GCF_021464435.1",
+            63.94,
+            87.85,
+        ),
+        _species(
+            "colias_croceus",
+            "insects",
+            "GCF_905220415.1_ilColCroc2.1",
+            "GCF_905220415.1",
+            62.65,
+            88.13,
+        ),
+        _species(
+            "danaus_plexippus",
+            "insects",
+            "GCF_018135715.1_MEX_DaPlex",
+            "GCF_018135715.1",
+            65.68,
+            89.38,
+        ),
+        _species(
+            "leptidea_sinapis",
+            "insects",
+            "GCF_905404315.1_ilLepSina1.1",
+            "GCF_905404315.1",
+            55.00,
+            81.99,
+        ),
+        _species(
+            "nymphalis_io",
+            "insects",
+            "GCF_905147045.1_ilAglIoxx1.1",
+            "GCF_905147045.1",
+            66.02,
+            89.19,
+        ),
+        _species(
+            "osmia_bicornis",
+            "insects",
+            "GCF_907164935.1_iOsmBic2.1",
+            "GCF_907164935.1",
+            58.75,
+            86.26,
+        ),
+        _species(
+            "tribolium_castaneum",
+            "insects",
+            "GCF_031307605.1_icTriCast1.1",
+            "GCF_031307605.1",
+            53.05,
+            80.60,
+        ),
+        _species(
+            "vanessa_cardui",
+            "insects",
+            "GCF_905220365.1_ilVanCard2.1",
+            "GCF_905220365.1",
+            64.37,
+            88.77,
+        ),
+        _species(
+            "zerene_cesonia",
+            "insects",
+            "GCF_012273895.1_Zerene_cesonia_1.1",
+            "GCF_012273895.1",
+            61.31,
+            85.14,
+        ),
+    )
+}

src/genome_annotation/data.py

@@ -0,0 +1,165 @@
+"""Training windows of (soft-masked tokens, 15-class labels) from FASTA + label BED."""
+
+from collections.abc import Iterator
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import pyfaidx
+import torch
+
+from genome_annotation.decoding import CLASS_INDEX
+
+IGNORE_INDEX = -1
+_MASK_LABELS = ("mask_short_intron", "mask_frame_error")
+_LABEL_TO_CLASS = CLASS_INDEX | dict.fromkeys(_MASK_LABELS, IGNORE_INDEX)
+# Where rows overlap, the label later in this order wins (masks last).
+_PAINT_ORDER = {label: i for i, label in enumerate(_LABEL_TO_CLASS)}
+_CDS_P0 = CLASS_INDEX["CDS_p0"]
+_COMPLEMENT = str.maketrans("ACGTNacgtn", "TGCANtgcan")
+
+
+def reverse_complement(sequence: str) -> str:
+    """Reverse-complements a DNA string, keeping soft-masking."""
+    return sequence.translate(_COMPLEMENT)[::-1]
+
+
+class Tokenizer:
+    """Maps DNA strings to NTv3 token ids, one token per base.
+
+    A lookup-table equivalent of the model's own (character-level) tokenizer, fast
+    enough for long windows: ACGTN and their soft-masked lowercase get their own
+    ids, anything else (IUPAC codes, ...) is ``<unk>``.
+    """
+
+    def __init__(self, vocab: dict[str, int]) -> None:
+        self.lut = np.full(256, vocab["<unk>"], dtype=np.int64)
+        for base in "ACGTNacgtn":
+            self.lut[ord(base)] = vocab[base]
+
+    def __call__(self, sequence: str) -> np.ndarray:
+        return self.lut[np.frombuffer(sequence.encode("ascii"), dtype=np.uint8)]
+
+
+class LabelTrack:
+    """Per-base 15-class labels of one assembly, read from its training BED.
+
+    Each BED row labels an interval of one strand. A CDS row carries the codon
+    position of its 5' base, which cycles 0, 1, 2 along the transcript.
+    """
+
+    def __init__(self, bed_path: str | Path) -> None:
+        bed = pd.read_csv(
+            bed_path,
+            sep="\t",
+            header=None,
+            names=["chrom", "start", "end", "label", "score", "strand"],
+            dtype={"chrom": str},
+        )
+        bed["cls"] = bed["label"].map(_LABEL_TO_CLASS).astype(np.int64)
+        bed["order"] = bed["label"].map(_PAINT_ORDER).astype(np.int64)
+        self.rows: dict[tuple[str, str], tuple[np.ndarray, ...]] = {}
+        for (chrom, strand), group in bed.groupby(["chrom", "strand"]):
+            group = group.sort_values("start")
+            self.rows[(chrom, strand)] = tuple(
+                group[col].to_numpy() for col in ("start", "end", "cls", "order")
+            )
+
+    def strand_labels(
+        self, chrom: str, strand: str, start: int, end: int
+    ) -> np.ndarray:
+        """(end - start,) labels of one strand, in genomic order."""
+        labels = np.zeros(end - start, dtype=np.int64)
+        if (chrom, strand) not in self.rows:
+            return labels
+        starts, ends, classes, order = self.rows[(chrom, strand)]
+        # Rows are sorted by start; a row overlaps the window iff it starts before
+        # its end and ends after its start (ends are not sorted where rows overlap).
+        hi = np.searchsorted(starts, end, side="left")
+        candidates = np.flatnonzero(ends[:hi] > start)
+        candidates = candidates[np.argsort(order[candidates], kind="stable")]
+        for row_start, row_end, cls in zip(
+            starts[candidates], ends[candidates], classes[candidates], strict=True
+        ):
+            s, e = max(start, row_start), min(end, row_end)
+            if _CDS_P0 <= cls < _CDS_P0 + 3:
+                frame = cls - _CDS_P0
+                if strand == "+":
+                    offsets = np.arange(s - row_start, e - row_start)
+                else:
+                    offsets = np.arange(row_end - 1 - s, row_end - 1 - e, -1)
+                labels[s - start : e - start] = _CDS_P0 + (frame + offsets) % 3
+            else:
+                labels[s - start : e - start] = cls
+        return labels
+
+    def window(self, chrom: str, start: int, end: int, reverse: bool) -> np.ndarray:
+        """Labels of the strand read 5'->3' in the window's orientation."""
+        if reverse:
+            return self.strand_labels(chrom, "-", start, end)[::-1].copy()
+        return self.strand_labels(chrom, "+", start, end)
+
+
+class Assembly:
+    """A genome and its labels, stored as ``<dir>/genome.fa.gz`` and ``labels.bed.gz``."""
+
+    def __init__(self, directory: str | Path) -> None:
+        self.directory = Path(directory)
+        self.name = self.directory.name
+        self.labels = LabelTrack(self.directory / "labels.bed.gz")
+        self._fasta: pyfaidx.Fasta | None = None
+        with pyfaidx.Fasta(str(self.directory / "genome.fa.gz")) as fasta:
+            self.lengths = {name: len(record) for name, record in fasta.items()}
+
+    @property
+    def fasta(self) -> pyfaidx.Fasta:
+        # Opened lazily so that each dataloader worker gets its own file handle.
+        if self._fasta is None:
+            self._fasta = pyfaidx.Fasta(str(self.directory / "genome.fa.gz"))
+        return self._fasta
+
+    def sequence(self, chrom: str, start: int, end: int) -> str:
+        return str(self.fasta[chrom][start:end])
+
+
+class TrainingWindows(torch.utils.data.IterableDataset):
+    """Endless random windows: species uniform, position uniform within the species.
+
+    Each window is reverse-complemented with probability 1/2, together with its
+    labels, so the single-strand head learns genes on both strands.
+    """
+
+    def __init__(
+        self,
+        assemblies: list[Assembly],
+        tokenizer: Tokenizer,
+        window: int,
+        seed: int = 0,
+    ) -> None:
+        self.assemblies = assemblies
+        self.tokenizer = tokenizer
+        self.window = window
+        self.seed = seed
+
+    def __iter__(self) -> Iterator[tuple[torch.Tensor, torch.Tensor]]:
+        info = torch.utils.data.get_worker_info()
+        rng = np.random.default_rng([self.seed, info.id if info else 0])
+        contigs = [
+            [(c, n) for c, n in a.lengths.items() if n >= self.window]
+            for a in self.assemblies
+        ]
+        while True:
+            index = rng.integers(len(self.assemblies))
+            assembly = self.assemblies[index]
+            sizes = np.array([n - self.window + 1 for _, n in contigs[index]])
+            flat = rng.integers(sizes.sum())
+            contig = int(np.searchsorted(np.cumsum(sizes), flat, side="right"))
+            chrom = contigs[index][contig][0]
+            start = int(flat - (np.cumsum(sizes)[contig] - sizes[contig]))
+            end = start + self.window
+            reverse = bool(rng.integers(2))
+            sequence = assembly.sequence(chrom, start, end)
+            if reverse:
+                sequence = reverse_complement(sequence)
+            labels = assembly.labels.window(chrom, start, end, reverse)
+            yield torch.from_numpy(self.tokenizer(sequence)), torch.from_numpy(labels)

src/genome_annotation/decoding.py

@@ -0,0 +1,338 @@
+"""Constrained Viterbi decoding of 15-class annotation logits into GFF3 genes.
+
+A simplified port of the NTv3 genome-annotation decoder: the same 15-state
+Tiberius-style HMM and start / splice / stop motif constraints, without the U12
+(AT..AC) intron lane.
+"""
+
+import gzip
+import itertools
+from collections.abc import Callable
+from pathlib import Path
+
+import numba
+import numpy as np
+
+CLASSES: tuple[str, ...] = (
+    "intergenic",
+    "intron_p0",
+    "intron_p1",
+    "intron_p2",
+    "CDS_p0",
+    "CDS_p1",
+    "CDS_p2",
+    "splice_acceptor_p0",
+    "splice_acceptor_p1",
+    "splice_acceptor_p2",
+    "splice_donor_p0",
+    "splice_donor_p1",
+    "splice_donor_p2",
+    "start_codon",
+    "stop_codon",
+)
+NUM_CLASSES = len(CLASSES)
+CLASS_INDEX: dict[str, int] = {name: i for i, name in enumerate(CLASSES)}
+
+# Tiberius mean lengths: intergenic ~10 kb, intron ~4.5 kb, exon ~200 bp.
+_P_IR = 1 / 10_000
+_P_I = 1 / 4_500
+_P_E = 1 / 200
+
+# Allowed transitions per state; the self-loop takes the remaining mass. Start,
+# stop and splice states last exactly 1 bp.
+TRANSITIONS: dict[str, dict[str, float]] = {
+    "intergenic": {"start_codon": _P_IR},
+    "intron_p0": {"splice_acceptor_p1": _P_I},
+    "intron_p1": {"splice_acceptor_p2": _P_I},
+    "intron_p2": {"splice_acceptor_p0": _P_I},
+    "CDS_p0": {"CDS_p1": 1 - _P_E, "splice_donor_p1": _P_E},
+    "CDS_p1": {"CDS_p2": 1 - 2 * _P_E, "splice_donor_p2": _P_E, "stop_codon": _P_E},
+    "CDS_p2": {"CDS_p0": 1 - _P_E, "splice_donor_p0": _P_E},
+    "splice_acceptor_p0": {"CDS_p1": 1.0},
+    "splice_acceptor_p1": {"CDS_p2": 1.0},
+    "splice_acceptor_p2": {"CDS_p0": 1.0},
+    "splice_donor_p0": {"intron_p0": 1.0},
+    "splice_donor_p1": {"intron_p1": 1.0},
+    "splice_donor_p2": {"intron_p2": 1.0},
+    "start_codon": {"CDS_p1": 1.0},
+    "stop_codon": {"intergenic": 1.0},
+}
+
+
+def build_log_transitions() -> np.ndarray:
+    """Returns the (15, 15) log transition matrix of the HMM."""
+    matrix = np.zeros((NUM_CLASSES, NUM_CLASSES), dtype=np.float64)
+    for src, dests in TRANSITIONS.items():
+        i = CLASS_INDEX[src]
+        matrix[i, i] = max(0.0, 1.0 - sum(dests.values()))
+        for dst, prob in dests.items():
+            matrix[i, CLASS_INDEX[dst]] = prob
+    with np.errstate(divide="ignore"):
+        return np.log(matrix.astype(np.float32))
+
+
+LOG_TRANSITIONS = build_log_transitions()
+
+# Base codes: A=0, C=1, G=2, T=3, anything else (N, IUPAC) = 4.
+A, C, G, T, N = 0, 1, 2, 3, 4
+_BASE_CODE_LUT = np.full(256, N, dtype=np.uint8)
+for _base, _code in zip("ACGT", (A, C, G, T), strict=True):
+    _BASE_CODE_LUT[ord(_base)] = _code
+    _BASE_CODE_LUT[ord(_base.lower())] = _code
+
+
+def encode_bases(sequence: str) -> np.ndarray:
+    """Maps a DNA string to case-insensitive base codes."""
+    return _BASE_CODE_LUT[np.frombuffer(sequence.encode("ascii"), dtype=np.uint8)]
+
+
+def reverse_complement_codes(codes: np.ndarray) -> np.ndarray:
+    """Reverse-complements an array of base codes."""
+    return np.where(codes == N, N, T - codes)[::-1].astype(np.uint8)
+
+
+# (offsets relative to the state position, bases, log-score) per accepted motif.
+# Scores follow Tiberius: a slight TAG preference, and GT >> GC donors.
+_MOTIFS: dict[str, list[tuple[tuple[int, ...], tuple[int, ...], float]]] = {
+    "start": [((0, 1, 2), (A, T, G), 0.0)],
+    "stop": [
+        ((-2, -1, 0), (T, A, G), float(np.log(np.float32(0.34)))),
+        ((-2, -1, 0), (T, A, A), float(np.log(np.float32(0.33)))),
+        ((-2, -1, 0), (T, G, A), float(np.log(np.float32(0.33)))),
+    ],
+    "donor": [
+        ((1, 2), (G, T), float(np.log(np.float32(0.99)))),
+        ((1, 2), (G, C), float(np.log(np.float32(0.01)))),
+    ],
+    "acceptor": [((-2, -1), (A, G), 0.0)],
+}
+_MOTIF_STATES: dict[str, tuple[str, ...]] = {
+    "start": ("start_codon",),
+    "stop": ("stop_codon",),
+    "donor": ("splice_donor_p0", "splice_donor_p1", "splice_donor_p2"),
+    "acceptor": ("splice_acceptor_p0", "splice_acceptor_p1", "splice_acceptor_p2"),
+}
+
+
+def _motif_log_mask(
+    codes: np.ndarray, rules: list[tuple[tuple[int, ...], tuple[int, ...], float]]
+) -> np.ndarray:
+    """Per-position log-score of one motif: -inf if no rule matches, 0 near an N."""
+    length = len(codes)
+    padded = np.pad(codes, (2, 2), constant_values=N)
+    offsets = {off for rule_offsets, _, _ in rules for off in rule_offsets}
+
+    def shifted(off: int) -> np.ndarray:
+        return padded[2 + off : 2 + off + length]
+
+    has_n = np.zeros(length, dtype=bool)
+    for off in offsets:
+        has_n |= shifted(off) == N
+    score = np.full(length, -np.inf, dtype=np.float32)
+    for rule_offsets, bases, log_score in rules:
+        match = np.ones(length, dtype=bool)
+        for off, base in zip(rule_offsets, bases, strict=True):
+            match &= shifted(off) == base
+        score[match] = log_score
+    score[has_n] = 0.0
+    # Motifs whose context falls outside the window are forbidden.
+    pos = np.arange(length)
+    score[(pos < -min(offsets)) | (pos >= length - max(offsets))] = -np.inf
+    return score
+
+
+def constraint_log_mask(codes: np.ndarray) -> np.ndarray:
+    """(L, 15) additive log-mask vetoing start/stop/splice states off their motif."""
+    mask = np.zeros((len(codes), NUM_CLASSES), dtype=np.float32)
+    for motif, rules in _MOTIFS.items():
+        column = _motif_log_mask(codes, rules)
+        for state in _MOTIF_STATES[motif]:
+            mask[:, CLASS_INDEX[state]] = column
+    return mask
+
+
+@numba.njit(cache=True)
+def _viterbi(log_emissions: np.ndarray, log_transitions: np.ndarray) -> np.ndarray:
+    length, num_states = log_emissions.shape
+    backpointers = np.zeros((length, num_states), dtype=np.int8)
+    scores = log_emissions[0].copy()
+    stepped = np.empty(num_states, dtype=np.float32)
+    for t in range(1, length):
+        for j in range(num_states):
+            best = -np.inf
+            best_i = 0
+            for i in range(num_states):
+                value = scores[i] + log_transitions[i, j]
+                if value > best:
+                    best = value
+                    best_i = i
+            stepped[j] = best + log_emissions[t, j]
+            backpointers[t, j] = best_i
+        # Rescale by the running max so scores stay bounded; argmaxes are unchanged.
+        offset = max(stepped.max(), np.float32(-1e30))
+        for j in range(num_states):
+            scores[j] = stepped[j] - offset
+    path = np.empty(length, dtype=np.int8)
+    path[-1] = np.argmax(scores)
+    for t in range(length - 1, 0, -1):
+        path[t - 1] = backpointers[t, path[t]]
+    return path
+
+
+def viterbi(
+    log_emissions: np.ndarray, constrain_codes: np.ndarray | None
+) -> np.ndarray:
+    """Most likely state path of one strand, read 5'->3'.
+
+    Args:
+        log_emissions: (L, 15) log-probabilities in transcription direction.
+        constrain_codes: Base codes in the same direction, or None to decode
+            without motif constraints.
+
+    Returns:
+        (L,) int8 class indices.
+    """
+    log_emissions = np.asarray(log_emissions, dtype=np.float32)
+    if constrain_codes is not None:
+        log_emissions = log_emissions + constraint_log_mask(constrain_codes)
+    return _viterbi(np.ascontiguousarray(log_emissions), LOG_TRANSITIONS)
+
+
+def tile_windows(length: int, window: int) -> list[tuple[int, int]]:
+    """(window start, first bp the window owns) covering [0, length) end to end.
+
+    The last window is pulled back to end on the final base; the bp it shares
+    with the previous window stay with the previous one.
+    """
+    if length < window:
+        raise ValueError(f"Sequence of {length} bp is shorter than one window.")
+    starts = list(range(0, length - window + 1, window))
+    tiles = [(s, 0) for s in starts]
+    if starts[-1] + window < length:
+        last = length - window
+        tiles.append((last, starts[-1] + window - last))
+    return tiles
+
+
+def annotate_sequence(
+    sequence: str,
+    predict: Callable[[str], tuple[np.ndarray, np.ndarray]],
+    window: int,
+    constrained: bool = True,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Decodes both strands of a chromosome window by window.
+
+    Args:
+        sequence: Chromosome sequence (soft-masking allowed).
+        predict: Maps a window's sequence to (L, 15) log-probabilities for the
+            window itself and for its reverse complement.
+        window: Window length in bp.
+        constrained: Whether to apply the motif constraints.
+
+    Returns:
+        Plus- and minus-strand state paths in genomic order, each (len(sequence),).
+    """
+    codes = encode_bases(sequence)
+    plus = np.zeros(len(sequence), dtype=np.int8)
+    minus = np.zeros(len(sequence), dtype=np.int8)
+    for start, keep in tile_windows(len(sequence), window):
+        end = start + window
+        log_probs_plus, log_probs_rc = predict(sequence[start:end])
+        window_codes = codes[start:end]
+        rc_codes = reverse_complement_codes(window_codes)
+        states_plus = viterbi(log_probs_plus, window_codes if constrained else None)
+        states_rc = viterbi(log_probs_rc, rc_codes if constrained else None)
+        plus[start + keep : end] = states_plus[keep:]
+        minus[start + keep : end] = states_rc[::-1][keep:]
+    return plus, minus
+
+
+_INTRON_STATES = {CLASS_INDEX[f"intron_p{p}"] for p in range(3)}
+# Codon position at the 5' end of a CDS run, which sets its GFF3 phase.
+_CODON_POSITION: dict[int, int] = {
+    CLASS_INDEX[f"{kind}_p{p}"]: p
+    for kind in ("CDS", "splice_acceptor", "splice_donor")
+    for p in range(3)
+} | {CLASS_INDEX["start_codon"]: 0, CLASS_INDEX["stop_codon"]: 0}
+_STOP_CODONS = {(T, A, A), (T, A, G), (T, G, A)}
+
+
+def _has_internal_stop(
+    gene_states: np.ndarray, gene_codes: np.ndarray, strand: str
+) -> bool:
+    """Whether the spliced CDS holds an in-frame stop before its last codon."""
+    is_intron = np.isin(gene_states, list(_INTRON_STATES))
+    cds = gene_codes[~is_intron]
+    if strand == "-":
+        cds = reverse_complement_codes(cds)
+    codons = cds[: len(cds) // 3 * 3].reshape(-1, 3)[:-1]
+    return any(tuple(int(b) for b in codon) in _STOP_CODONS for codon in codons)
+
+
+def find_genes(
+    states: np.ndarray,
+    strand: str,
+    chrom: str,
+    codes: np.ndarray | None = None,
+    first_gene_id: int = 0,
+) -> list[str]:
+    """Turns one strand's state path into GFF3 gene/CDS/intron lines.
+
+    A gene is a maximal run of non-intergenic states. Genes with an internal
+    in-frame stop codon are dropped when ``codes`` is given.
+
+    Args:
+        states: (L,) class indices of one strand, in genomic order.
+        strand: "+" or "-".
+        chrom: Sequence name for the GFF3 seqid column.
+        codes: Optional base codes of the chromosome, in genomic order.
+        first_gene_id: Genes are numbered from first_gene_id + 1.
+
+    Returns:
+        GFF3 lines, each gene followed by its sub-features.
+    """
+    is_genic = np.concatenate([[False], states != CLASS_INDEX["intergenic"], [False]])
+    edges = np.flatnonzero(is_genic[1:] != is_genic[:-1])
+    feature_type = np.where(
+        np.isin(states, list(_INTRON_STATES)), 2, np.where(states == 0, 0, 1)
+    )
+    type_changes = np.flatnonzero(feature_type[1:] != feature_type[:-1]) + 1
+    lines: list[str] = []
+    gene_id = first_gene_id
+    for gene_start, gene_end in zip(edges[::2], edges[1::2], strict=True):
+        gene_states = states[gene_start:gene_end]
+        if codes is not None and _has_internal_stop(
+            gene_states, codes[gene_start:gene_end], strand
+        ):
+            continue
+        gene_id += 1
+        gid = f"gene_{gene_id}_{'p' if strand == '+' else 'm'}"
+        lines.append(
+            f"{chrom}\tNTv3\tgene\t{gene_start + 1}\t{gene_end}\t.\t{strand}\t.\tID={gid}"
+        )
+        lo = np.searchsorted(type_changes, gene_start, "right")
+        hi = np.searchsorted(type_changes, gene_end, "left")
+        bounds = [gene_start, *type_changes[lo:hi].tolist(), gene_end]
+        for feat_start, feat_end in itertools.pairwise(bounds):
+            is_cds = feature_type[feat_start] == 1
+            phase = "."
+            if is_cds:
+                five_prime = feat_start if strand == "+" else feat_end - 1
+                phase = str((3 - _CODON_POSITION[int(states[five_prime])]) % 3)
+            kind = "CDS" if is_cds else "intron"
+            lines.append(
+                f"{chrom}\tNTv3\t{kind}\t{feat_start + 1}\t{feat_end}\t.\t{strand}"
+                f"\t{phase}\tID={gid}_{kind}_{feat_start + 1};Parent={gid}"
+            )
+    return lines
+
+
+def write_gff3(path: str | Path, lines: list[str]) -> None:
+    """Writes GFF3 lines sorted by (seqid, start), gzipped if path ends in .gz."""
+    ordered = sorted(
+        lines, key=lambda line: (line.split("\t")[0], int(line.split("\t")[3]))
+    )
+    text = "##gff-version 3\n" + "\n".join(ordered) + "\n"
+    opener = gzip.open if str(path).endswith(".gz") else open
+    with opener(path, "wt") as handle:
+        handle.write(text)

src/genome_annotation/evaluation.py

@@ -0,0 +1,672 @@
+"""Scores a predicted GFF3 against a reference the way the NTv3 benchmark did.
+
+Tiberius / BRAKER-style (Gabriel et al. 2024): a CDS is correct if chromosome,
+start, end, strand and phase all match; a gene is correct if its whole set of CDS
+matches a reference gene. Reference genes are restricted to those the NTv3
+training labels keep, predicted genes below 201 nt of CDS are dropped, and
+predicted CDS in pseudogene regions away from reference genes are ignored.
+
+Standard library only, so it runs anywhere the notebook does.
+"""
+
+import bisect
+import gzip
+import re
+from collections import defaultdict
+from collections.abc import Callable, Iterable, Iterator, Mapping
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import IO, NamedTuple
+
+CdsKey = tuple[str, int, int, str, str]  # (chrom, start, end, strand, phase), 1-based
+Intervals = dict[str, list[tuple[int, int]]]
+
+# GENCODE keeps pseudogenes as generic rows with a pseudogene ``gene_type``.
+_PSEUDOGENE_GENE_TYPE = re.compile(r"(?:^|;)gene_type=[^;]*pseudogene")
+
+LEVELS = {
+    "CDS": "Exon",
+    "Gene": "Gene",
+    "CDS_no_phase": "Exon (no phase)",
+    "Gene_no_phase": "Gene (no phase)",
+}
+
+
+# ---------------------------------------------------------------------------
+# GFF reading
+# ---------------------------------------------------------------------------
+
+
+class GffRow(NamedTuple):
+    """One GFF3 feature row (1-based closed coordinates)."""
+
+    seqid: str
+    type: str
+    start: int
+    end: int
+    strand: str
+    phase: str
+    attribute: str
+
+    def attributes(self) -> dict[str, str]:
+        """Column 9 as ``{tag: value}``."""
+        out = {}
+        for item in self.attribute.split(";"):
+            if "=" in item:
+                key, value = item.strip().split("=", 1)
+                out[key.strip()] = value.strip()
+        return out
+
+    def get_id(self) -> str | None:
+        """First ``ID=`` value."""
+        for item in self.attribute.split(";"):
+            item = item.strip()
+            if item.startswith("ID="):
+                return item[3:].split(",")[0].strip() or None
+        return None
+
+    def get_parents(self) -> list[str]:
+        """All ``Parent=`` values."""
+        for item in self.attribute.split(";"):
+            item = item.strip()
+            if item.startswith("Parent="):
+                return [v.strip() for v in item[7:].split(",") if v.strip()]
+        return []
+
+
+def open_text(path: str | Path) -> IO[str]:
+    """Opens a plain or gzipped text file."""
+    if str(path).endswith(".gz"):
+        return gzip.open(path, "rt", encoding="utf-8")
+    return open(path, encoding="utf-8")
+
+
+def iter_gff(
+    path: str | Path, chromosomes: Iterable[str] | None = None
+) -> Iterator[GffRow]:
+    """Yields feature rows, skipping comments and malformed lines."""
+    keep = set(chromosomes) if chromosomes is not None else None
+    with open_text(path) as handle:
+        for line in handle:
+            if not line.strip() or line.startswith("#"):
+                continue
+            cols = line.rstrip("\n").split("\t")
+            if len(cols) < 9 or (keep is not None and cols[0] not in keep):
+                continue
+            try:
+                yield GffRow(cols[0], cols[2], int(cols[3]), int(cols[4]), *cols[6:9])
+            except ValueError:
+                continue
+
+
+def merge_intervals(intervals: Iterable[tuple[int, int]]) -> list[tuple[int, int]]:
+    """Merges overlapping or adjacent 1-based closed intervals."""
+    merged: list[tuple[int, int]] = []
+    for start, end in sorted(intervals):
+        if merged and start <= merged[-1][1] + 1:
+            merged[-1] = (merged[-1][0], max(merged[-1][1], end))
+        else:
+            merged.append((start, end))
+    return merged
+
+
+# ---------------------------------------------------------------------------
+# Reference preparation
+# ---------------------------------------------------------------------------
+
+
+def write_longest_isoform(
+    source: str | Path, output: str | Path, keep_pseudogenes: bool = False
+) -> None:
+    """Keeps each protein-coding gene with its longest-CDS transcript.
+
+    Turns a raw multi-isoform annotation (e.g. an NCBI ``*_genomic.gff.gz``) into
+    the single-isoform reference the benchmark scores against. Rows must be in
+    the usual gene -> transcript -> exon/CDS order.
+
+    Args:
+        source: Raw GFF3, plain or gzipped.
+        output: Destination GFF3, gzipped if it ends in ``.gz``.
+        keep_pseudogenes: Also copy ``pseudogene`` rows, which turns on
+            pseudogene filtering. The benchmark references carried none outside
+            hg38, so leave off to reproduce it.
+    """
+    opener = gzip.open if str(output).endswith(".gz") else open
+    transcripts: dict[str, list[str]] = {}
+    lengths: dict[str, int] = {}
+    gene_line: str | None = None
+    gene_id: str | None = None
+    transcript_id: str | None = None
+    keep_gene = False
+
+    def flush(out: IO[str]) -> None:
+        if transcripts and gene_line is not None:
+            longest = max(transcripts, key=lambda tid: lengths[tid])
+            out.write(gene_line)
+            out.writelines(transcripts[longest])
+
+    with open_text(source) as src, opener(output, "wt", encoding="utf-8") as out:
+        out.write("##gff-version 3\n")
+        for line in src:
+            if not line.strip() or line.startswith("#"):
+                continue
+            cols = line.rstrip("\n").split("\t")
+            if len(cols) < 9:
+                continue
+            row = GffRow(cols[0], cols[2], int(cols[3]), int(cols[4]), *cols[6:9])
+            if row.type in ("gene", "pseudogene"):
+                flush(out)
+                if row.type == "pseudogene" and keep_pseudogenes:
+                    out.write(line)
+                attrs = dict(
+                    item.split("=", 1)
+                    for item in row.attribute.split(";")
+                    if "=" in item
+                )
+                biotype = attrs.get("gene_type") or attrs.get("gene_biotype")
+                keep_gene = biotype == "protein_coding" or not biotype
+                transcripts, lengths = {}, {}
+                gene_id, gene_line = row.get_id(), line
+            if not keep_gene:
+                continue
+            if row.type in ("mRNA", "transcript"):
+                parents = row.get_parents()
+                if parents and gene_id is not None and parents[0] == gene_id:
+                    transcript_id = row.get_id()
+                    if transcript_id is not None:
+                        transcripts[transcript_id] = [line]
+                        lengths[transcript_id] = 0
+            elif (
+                row.type in ("exon", "CDS")
+                and transcript_id is not None
+                and transcript_id in transcripts
+            ):
+                # Children join the last transcript seen, as in the pipeline.
+                transcripts[transcript_id].append(line)
+                if row.type == "CDS":
+                    lengths[transcript_id] += row.end - row.start + 1
+        flush(out)
+
+
+@dataclass
+class _CodingGene:
+    gene_id: str
+    name: str
+    chrom: str
+    strand: str
+    cds: list[tuple[int, int]] = field(default_factory=list)
+
+
+def _parse_coding_genes(
+    path: str | Path, chromosomes: Iterable[str] | None = None
+) -> dict[str, _CodingGene]:
+    genes: dict[str, _CodingGene] = {}
+    transcript_to_gene: dict[str, str] = {}
+    cds_by_transcript: dict[str, list[tuple[int, int]]] = defaultdict(list)
+    for row in iter_gff(path, chromosomes):
+        if row.type == "gene":
+            attrs = row.attributes()
+            biotype = attrs.get("gene_type") or attrs.get("gene_biotype")
+            gene_id = row.get_id()
+            if (biotype == "protein_coding" or not biotype) and gene_id:
+                if row.strand in ("+", "-"):
+                    name = attrs.get("Name") or gene_id.removeprefix("gene-")
+                    genes[gene_id] = _CodingGene(gene_id, name, row.seqid, row.strand)
+        elif row.type in ("mRNA", "transcript"):
+            parents, tid = row.get_parents(), row.get_id()
+            if parents and tid:
+                transcript_to_gene[tid] = parents[0]
+        elif row.type == "CDS" and row.phase != ".":
+            parents = row.get_parents()
+            if parents:
+                cds_by_transcript[parents[0]].append((row.start, row.end))
+    for tid, gid in transcript_to_gene.items():
+        if gid in genes:
+            genes[gid].cds.extend(cds_by_transcript.get(tid, []))
+    for gene in genes.values():
+        gene.cds.sort()
+    return genes
+
+
+def _intervals_overlap(a: list[tuple[int, int]], b: list[tuple[int, int]]) -> bool:
+    i = j = 0
+    while i < len(a) and j < len(b):
+        if a[i][0] <= b[j][1] and b[j][0] <= a[i][1]:
+            return True
+        if a[i][1] < b[j][1]:
+            i += 1
+        else:
+            j += 1
+    return False
+
+
+def _cluster_by_span(
+    gene_ids: list[str], genes: dict[str, _CodingGene]
+) -> list[list[str]]:
+    ordered = sorted(gene_ids, key=lambda g: (genes[g].cds[0][0], genes[g].cds[-1][1]))
+    clusters: list[list[str]] = []
+    end = -1
+    for gid in ordered:
+        if clusters and genes[gid].cds[0][0] <= end:
+            clusters[-1].append(gid)
+            end = max(end, genes[gid].cds[-1][1])
+        else:
+            clusters.append([gid])
+            end = genes[gid].cds[-1][1]
+    return clusters
+
+
+def _is_readthrough(name: str, other_names: set[str]) -> bool:
+    parts = name.split("-")
+    return any(
+        "-".join(parts[:i]) in other_names and "-".join(parts[i:]) in other_names
+        for i in range(1, len(parts))
+    )
+
+
+def _cds_length(gene: _CodingGene) -> int:
+    return sum(end - start + 1 for start, end in gene.cds)
+
+
+def _conflicting_gene_ids(genes: dict[str, _CodingGene]) -> set[str]:
+    """Readthrough fusions, then the shorter of two genes sharing CDS bases."""
+    by_strand: dict[tuple[str, str], list[str]] = defaultdict(list)
+    for gid, gene in genes.items():
+        by_strand[(gene.chrom, gene.strand)].append(gid)
+    conflicts: set[str] = set()
+    for gids in by_strand.values():
+        for cluster in _cluster_by_span(gids, genes):
+            if len(cluster) < 2:
+                continue
+            names = {gid: genes[gid].name for gid in cluster}
+            all_names = set(names.values())
+            remaining = []
+            for gid in cluster:
+                if _is_readthrough(names[gid], all_names - {names[gid]}):
+                    conflicts.add(gid)
+                else:
+                    remaining.append(gid)
+            kept: list[str] = []
+            for gid in remaining:
+                clash = next(
+                    (
+                        other
+                        for other in kept
+                        if genes[gid].cds[0][0] <= genes[other].cds[-1][1]
+                        and _intervals_overlap(genes[other].cds, genes[gid].cds)
+                    ),
+                    None,
+                )
+                if clash is None:
+                    kept.append(gid)
+                elif _cds_length(genes[gid]) > _cds_length(genes[clash]):
+                    conflicts.add(clash)
+                    kept.remove(clash)
+                    kept.append(gid)
+                else:
+                    conflicts.add(gid)
+    return conflicts
+
+
+def scorable_reference_genes(
+    path: str | Path, chromosomes: Iterable[str] | None = None
+) -> set[str]:
+    """IDs of the reference genes the benchmark scores.
+
+    Drops genes with no CDS, overlapping CDS, readthrough fusions and CDS shared
+    with a longer gene, then keeps the longest gene of every same-strand cluster
+    of overlapping CDS spans (so a gene nested in another's intron goes too).
+    """
+    genes = {
+        gid: gene
+        for gid, gene in _parse_coding_genes(path, chromosomes).items()
+        if gene.cds
+        and all(b[0] > a[1] for a, b in zip(gene.cds, gene.cds[1:], strict=False))
+    }
+    conflicts = _conflicting_gene_ids(genes)
+    genes = {gid: gene for gid, gene in genes.items() if gid not in conflicts}
+    by_strand: dict[tuple[str, str], list[str]] = defaultdict(list)
+    for gid, gene in genes.items():
+        by_strand[(gene.chrom, gene.strand)].append(gid)
+    return {
+        max(cluster, key=lambda g: (_cds_length(genes[g]), g))
+        for gids in by_strand.values()
+        for cluster in _cluster_by_span(gids, genes)
+    }
+
+
+# ---------------------------------------------------------------------------
+# Scoring
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class Gene:
+    """A gene reduced to what the scorer compares."""
+
+    gene_id: str
+    chrom: str
+    strand: str
+    cds_keys: list[CdsKey] = field(default_factory=list)
+
+
+def parse_cds_and_genes(
+    path: str | Path, chromosomes: Iterable[str] | None = None
+) -> tuple[set[CdsKey], dict[str, Gene]]:
+    """All CDS keys, and each gene's CDS keys (CDS reach genes via transcripts)."""
+    cds: set[CdsKey] = set()
+    genes: dict[str, Gene] = {}
+    transcript_to_gene: dict[str, str] = {}
+    for row in iter_gff(path, chromosomes):
+        kind = row.type.lower()
+        if kind in ("gene", "pseudogene"):
+            gid = row.get_id()
+            if gid and gid not in genes:
+                genes[gid] = Gene(gid, row.seqid, row.strand)
+        elif kind in ("mrna", "transcript"):
+            tid, parents = row.get_id(), row.get_parents()
+            if tid and parents:
+                transcript_to_gene[tid] = parents[0]
+        elif kind == "cds":
+            parents = row.get_parents()
+            parent = parents[0] if parents else ""
+            gid = transcript_to_gene.get(parent, parent)
+            key = (row.seqid, row.start, row.end, row.strand, row.phase)
+            cds.add(key)
+            if gid:
+                genes.setdefault(gid, Gene(gid, row.seqid, row.strand))
+                genes[gid].cds_keys.append(key)
+    return cds, genes
+
+
+def parse_pseudogene_regions(
+    path: str | Path, chromosomes: Iterable[str] | None = None
+) -> list[tuple[str, int, int]]:
+    """``pseudogene`` rows, or rows whose ``gene_type`` is a pseudogene (GENCODE)."""
+    return [
+        (row.seqid, row.start, row.end)
+        for row in iter_gff(path, chromosomes)
+        if row.type == "pseudogene" or _PSEUDOGENE_GENE_TYPE.search(row.attribute)
+    ]
+
+
+def _interval_index(regions: Iterable[tuple[str, int, int]]) -> Intervals:
+    by_chrom: dict[str, list[tuple[int, int]]] = defaultdict(list)
+    for chrom, start, end in regions:
+        by_chrom[chrom].append((start, end))
+    return {chrom: merge_intervals(ivs) for chrom, ivs in by_chrom.items()}
+
+
+def _overlaps_any(chrom: str, start: int, end: int, index: Intervals) -> bool:
+    intervals = index.get(chrom)
+    if not intervals:
+        return False
+    i = bisect.bisect_left(intervals, (end + 1,)) - 1
+    return i >= 0 and intervals[i][1] >= start
+
+
+@dataclass
+class ScoringSets:
+    """Reference and prediction after every benchmark filter."""
+
+    ref_cds: set[CdsKey]
+    ref_genes: dict[str, Gene]
+    pred_cds: set[CdsKey]
+    pred_genes: dict[str, Gene]
+
+
+def prepare_scoring_sets(
+    reference: str | Path,
+    prediction: str | Path,
+    pseudogenes: str | Path | None = None,
+    chromosomes: Iterable[str] | None = None,
+    min_pred_gene_cds_length: int = 201,
+    full_reference: bool = False,
+) -> ScoringSets:
+    """Parses both files and applies the benchmark filters.
+
+    Args:
+        reference: Single-isoform reference GFF3.
+        prediction: Predicted GFF3.
+        pseudogenes: GFF3 to read pseudogene regions from; the benchmark passed
+            the reference itself.
+        chromosomes: Seqids to score; None scores the whole reference, so a
+            contig the prediction skips counts as missed.
+        min_pred_gene_cds_length: Predicted genes with less CDS (nt) are dropped.
+        full_reference: Score every reference gene with a CDS instead of the
+            ones NTv3 training labels keep.
+
+    Returns:
+        The filtered CDS and gene sets.
+    """
+    chroms = sorted(set(chromosomes)) if chromosomes is not None else None
+    ref_cds, ref_genes = parse_cds_and_genes(reference, chroms)
+    if full_reference:
+        ref_genes = {gid: gene for gid, gene in ref_genes.items() if gene.cds_keys}
+    else:
+        keep = scorable_reference_genes(reference, chroms)
+        ref_genes = {gid: gene for gid, gene in ref_genes.items() if gid in keep}
+    ref_cds = {key for gene in ref_genes.values() for key in gene.cds_keys}
+
+    pred_cds, pred_genes = parse_cds_and_genes(prediction, chroms)
+    if min_pred_gene_cds_length > 0:
+        short = {
+            gid
+            for gid, gene in pred_genes.items()
+            if sum(end - start + 1 for _, start, end, _, _ in gene.cds_keys)
+            < min_pred_gene_cds_length
+        }
+        pred_cds -= {key for gid in short for key in pred_genes[gid].cds_keys}
+        pred_genes = {gid: g for gid, g in pred_genes.items() if gid not in short}
+
+    if pseudogenes is not None:
+        ref_index = _interval_index((c, s, e) for c, s, e, _, _ in ref_cds)
+        away_from_genes = _interval_index(
+            region
+            for region in parse_pseudogene_regions(pseudogenes, chroms)
+            if not _overlaps_any(*region, ref_index)
+        )
+        pred_cds = {k for k in pred_cds if not _overlaps_any(*k[:3], away_from_genes)}
+        pred_genes = {
+            gid: Gene(gid, g.chrom, g.strand, kept)
+            for gid, g in pred_genes.items()
+            if (kept := [k for k in g.cds_keys if k in pred_cds])
+        }
+    return ScoringSets(ref_cds, ref_genes, pred_cds, pred_genes)
+
+
+def restrict_to_region(
+    sets: ScoringSets, chrom: str, start: int, end: int
+) -> ScoringSets:
+    """Keeps the genes whose CDS lie inside [start, end) (1-based, closed ends).
+
+    Genes cut by the region edges are dropped on both sides, so neither the
+    reference nor the prediction is penalised for what lies outside.
+    """
+
+    def inside(gene: Gene) -> bool:
+        return gene.chrom == chrom and all(
+            start <= k[1] and k[2] <= end for k in gene.cds_keys
+        )
+
+    ref_genes = {g: x for g, x in sets.ref_genes.items() if inside(x)}
+    pred_genes = {g: x for g, x in sets.pred_genes.items() if inside(x)}
+    return ScoringSets(
+        {k for x in ref_genes.values() for k in x.cds_keys},
+        ref_genes,
+        {k for x in pred_genes.values() for k in x.cds_keys},
+        pred_genes,
+    )
+
+
+def _prf(ref_tp: int, n_ref: int, pred_tp: int, n_pred: int) -> dict[str, float]:
+    recall = ref_tp / n_ref if n_ref else 0.0
+    precision = pred_tp / n_pred if n_pred else 0.0
+    f1 = 2 * recall * precision / (recall + precision) if recall + precision else 0.0
+    return {
+        "Ref_TP": ref_tp,
+        "Ref_Total": n_ref,
+        "Pred_TP": pred_tp,
+        "Pred_Total": n_pred,
+        "Recall": recall,
+        "Precision": precision,
+        "F1": f1,
+    }
+
+
+def _no_phase(keys: Iterable[CdsKey]) -> frozenset[tuple[str, int, int, str]]:
+    return frozenset(key[:4] for key in keys)
+
+
+def _cds_scores(ref: set, pred: set) -> dict[str, float]:
+    tp = len(ref & pred)
+    scores = _prf(tp, len(ref), tp, len(pred))
+    return {
+        "TP": tp,
+        "FN": len(ref) - tp,
+        "FP": len(pred) - tp,
+        **{k: scores[k] for k in ("Recall", "Precision", "F1")},
+    }
+
+
+def _gene_scores(
+    ref_genes: Iterable[Gene],
+    pred_genes: Iterable[Gene],
+    structure: Callable[[Gene], frozenset],
+    all_ref: set[frozenset],
+    all_pred: set[frozenset],
+) -> dict[str, float]:
+    ref = [structure(g) for g in ref_genes]
+    pred = [structure(g) for g in pred_genes]
+    return _prf(
+        sum(s in all_pred for s in ref),
+        len(ref),
+        sum(s in all_ref for s in pred),
+        len(pred),
+    )
+
+
+def score_sets(sets: ScoringSets) -> dict[str, dict[str, float]]:
+    """CDS- and gene-level scores, with and without phase (fractions, not %)."""
+    return _score(sets, sets.ref_cds, sets.pred_cds, sets.ref_genes, sets.pred_genes)
+
+
+def _score(
+    sets: ScoringSets,
+    ref_cds: set[CdsKey],
+    pred_cds: set[CdsKey],
+    ref_genes: dict[str, Gene],
+    pred_genes: dict[str, Gene],
+) -> dict[str, dict[str, float]]:
+    with_phase: Callable[[Gene], frozenset] = lambda g: frozenset(g.cds_keys)  # noqa: E731
+    without_phase: Callable[[Gene], frozenset] = lambda g: _no_phase(g.cds_keys)  # noqa: E731
+    results = {"CDS": _cds_scores(ref_cds, pred_cds)}
+    results["Gene"] = _gene_scores(
+        ref_genes.values(),
+        pred_genes.values(),
+        with_phase,
+        {with_phase(g) for g in sets.ref_genes.values()},
+        {with_phase(g) for g in sets.pred_genes.values()},
+    )
+    results["CDS_no_phase"] = _cds_scores(
+        {key[:4] for key in ref_cds}, {key[:4] for key in pred_cds}
+    )
+    results["Gene_no_phase"] = _gene_scores(
+        ref_genes.values(),
+        pred_genes.values(),
+        without_phase,
+        {without_phase(g) for g in sets.ref_genes.values()},
+        {without_phase(g) for g in sets.pred_genes.values()},
+    )
+    return results
+
+
+def evaluate(
+    reference: str | Path,
+    prediction: str | Path,
+    pseudogenes: str | Path | None = None,
+    chromosomes: Iterable[str] | None = None,
+    min_pred_gene_cds_length: int = 201,
+    full_reference: bool = False,
+) -> dict[str, dict[str, float]]:
+    """Benchmark scores of one prediction; see ``prepare_scoring_sets`` for args."""
+    return score_sets(
+        prepare_scoring_sets(
+            reference,
+            prediction,
+            pseudogenes,
+            chromosomes,
+            min_pred_gene_cds_length,
+            full_reference,
+        )
+    )
+
+
+def score_per_contig(sets: ScoringSets) -> dict[str, dict[str, dict[str, float]]]:
+    """Scores per seqid, matched against the whole genome (counts add up)."""
+    chroms = sorted(
+        {g.chrom for g in sets.ref_genes.values()}
+        | {g.chrom for g in sets.pred_genes.values()}
+    )
+    return {
+        chrom: _score(
+            sets,
+            {k for k in sets.ref_cds if k[0] == chrom},
+            {k for k in sets.pred_cds if k[0] == chrom},
+            {gid: g for gid, g in sets.ref_genes.items() if g.chrom == chrom},
+            {gid: g for gid, g in sets.pred_genes.items() if g.chrom == chrom},
+        )
+        for chrom in chroms
+    }
+
+
+def compare_seqids(reference: str | Path, prediction: str | Path) -> dict[str, list]:
+    """Seqids with CDS in each file, to catch naming mismatches (``chr1`` vs ``1``)."""
+    ref = {key[0] for key in parse_cds_and_genes(reference)[0]}
+    pred = {key[0] for key in parse_cds_and_genes(prediction)[0]}
+    return {
+        "shared": sorted(ref & pred),
+        "reference_only": sorted(ref - pred),
+        "prediction_only": sorted(pred - ref),
+    }
+
+
+def rename_seqids(
+    source: str | Path,
+    output: str | Path,
+    rename: Mapping[str, str] | Callable[[str], str],
+) -> None:
+    """Writes a copy of a GFF3 with column 1 (and ``##sequence-region``) renamed.
+
+    Args:
+        source: GFF3 to rename, plain or gzipped.
+        output: Destination, gzipped if it ends in ``.gz``.
+        rename: Old to new seqid (unlisted ones are kept), or a function.
+    """
+    if callable(rename):
+        new = rename
+    else:
+        mapping = dict(rename)
+
+        def new(seqid: str) -> str:
+            return mapping.get(seqid, seqid)
+
+    opener = gzip.open if str(output).endswith(".gz") else open
+    with open_text(source) as src, opener(output, "wt", encoding="utf-8") as out:
+        for line in src:
+            if line.startswith("##sequence-region"):
+                tag, seqid, *rest = line.split()
+                line = " ".join([tag, str(new(seqid)), *rest]) + "\n"
+            elif line.strip() and not line.startswith("#"):
+                seqid, rest = line.split("\t", 1)
+                line = f"{new(seqid)}\t{rest}"
+            out.write(line)
+
+
+def format_summary(results: dict[str, dict[str, float]]) -> list[str]:
+    """One line per level, in %."""
+    return [
+        f"{LEVELS.get(level, level):>16s}:  Recall {100 * m['Recall']:6.2f}%  "
+        f"Precision {100 * m['Precision']:6.2f}%  F1 {100 * m['F1']:6.2f}%"
+        for level, m in results.items()
+    ]

src/genome_annotation/inference.py

@@ -0,0 +1,255 @@
+"""Annotates a genomic region with a released NTv3 model, tiled as the pipeline does.
+
+Each window of ``sequence_length`` bp predicts its central ``keep`` fraction; the
+kept regions form a grid anchored at the contig's first base, the first window
+reaching ``crop`` bp ahead of the contig (read as N) and the last one shifted back
+onto its final base. Each kept region is decoded on its own and the state paths
+are stitched, the overlap of the shifted last region going to the earlier window.
+"""
+
+from collections import defaultdict
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+from transformers import AutoModel
+
+from . import decoding
+from .data import Tokenizer, reverse_complement
+from .model import CODE_REVISION, predict_log_probs
+
+DECODERS = ("argmax", "hmm", "hmm_constrained")
+
+
+def load_released(clade: str, device: torch.device) -> torch.nn.Module:
+    """``InstaDeepAI/NTv3_annot_<clade>`` (mammals, plants or insects), in eval mode."""
+    model = AutoModel.from_pretrained(
+        f"InstaDeepAI/NTv3_annot_{clade}",
+        trust_remote_code=True,
+        code_revision=CODE_REVISION,
+    )
+    return model.to(device).eval()
+
+
+@dataclass(frozen=True)
+class Window:
+    """One model input and the contig bp its decode is kept for (0-based, half-open)."""
+
+    start: int
+    keep_start: int
+    keep_stop: int
+
+
+def crop_of(sequence_length: int, keep: float) -> int:
+    """Context bp the head does not predict, on each side of a window."""
+    return (sequence_length - int(sequence_length * keep)) // 2
+
+
+def plan_windows(contig_length: int, sequence_length: int, keep: float) -> list[Window]:
+    """The pipeline's windows over a whole contig, kept regions disjoint and covering it."""
+    crop = crop_of(sequence_length, keep)
+    stride = sequence_length - 2 * crop
+    last_start = max(contig_length - stride, 0)
+    windows: list[Window] = []
+    covered = 0
+    for grid in range(0, contig_length, stride):
+        keep_start = min(grid, last_start)
+        keep_stop = min(keep_start + stride, contig_length)
+        windows.append(Window(keep_start - crop, max(keep_start, covered), keep_stop))
+        covered = keep_stop
+    return windows
+
+
+def windows_for_region(
+    contig_length: int, start: int, end: int, sequence_length: int, keep: float
+) -> list[Window]:
+    """The pipeline windows whose kept region overlaps [start, end)."""
+    return [
+        w
+        for w in plan_windows(contig_length, sequence_length, keep)
+        if w.keep_start < end and w.keep_stop > start
+    ]
+
+
+def window_sequence(contig: str, window: Window, sequence_length: int) -> str:
+    """The window's bases, N-padded where it reaches past either contig end."""
+    lo, hi = max(window.start, 0), min(window.start + sequence_length, len(contig))
+    return (
+        "N" * (lo - window.start)
+        + contig[lo:hi]
+        + "N" * (window.start + sequence_length - hi)
+    )
+
+
+def predict_window(
+    model: torch.nn.Module,
+    tokenizer: Tokenizer,
+    sequence: str,
+    device: torch.device,
+    dtype: torch.dtype,
+    one_strand_at_a_time: bool = False,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Log-probabilities of the window's kept region, for it and its reverse complement.
+
+    ``one_strand_at_a_time`` halves peak memory (about 31 GB instead of 57 GB for a
+    1 Mb window in bfloat16).
+
+    Returns:
+        (kept, 15) for the plus strand in genomic order, and (kept, 15) for the
+        minus strand in reverse-complement order (5' to 3' of the minus strand).
+    """
+    tokens = np.stack([tokenizer(sequence), tokenizer(reverse_complement(sequence))])
+    if one_strand_at_a_time:
+        plus, rc = (predict_log_probs(model, t[None], device, dtype)[0] for t in tokens)
+        return plus, rc
+    log_probs = predict_log_probs(model, tokens, device, dtype)
+    return log_probs[0], log_probs[1]
+
+
+def decode(log_probs: np.ndarray, codes: np.ndarray, decoder: str) -> np.ndarray:
+    """State path of one strand's kept region (both in the strand's own orientation)."""
+    if decoder == "argmax":
+        return log_probs.argmax(-1).astype(np.int8)
+    return decoding.viterbi(log_probs, codes if decoder == "hmm_constrained" else None)
+
+
+@dataclass
+class RegionPrediction:
+    """Model outputs of every window covering a region.
+
+    ``per_window`` holds, for each window, the log-probabilities of its whole kept
+    region: plus strand in genomic order, minus strand in its own 5'->3' order.
+    """
+
+    chrom: str
+    start: int
+    end: int
+    windows: list[Window]
+    per_window: list[tuple[np.ndarray, np.ndarray]]
+
+    def stitched(self, strand: str) -> np.ndarray:
+        """(end - start, 15) log-probabilities of one strand ("+"/"-"), genomic order."""
+        out = np.zeros((self.end - self.start, decoding.NUM_CLASSES), np.float32)
+        for window, (plus, rc) in zip(self.windows, self.per_window, strict=True):
+            values = plus if strand == "+" else rc[::-1]
+            skip = window.keep_start - (window.keep_stop - len(plus))
+            out[window.keep_start - self.start : window.keep_stop - self.start] = (
+                values[skip:]
+            )
+        return out
+
+
+def predict_region(
+    model: torch.nn.Module,
+    tokenizer: Tokenizer,
+    contig: str,
+    chrom: str,
+    start: int,
+    end: int,
+    sequence_length: int,
+    keep: float,
+    device: torch.device,
+    dtype: torch.dtype,
+    one_strand_at_a_time: bool = False,
+) -> RegionPrediction:
+    """Runs the model on every pipeline window covering [start, end) of a contig.
+
+    The region is widened to whole kept regions, so decoding it reproduces the
+    pipeline's decode of those bp exactly.
+    """
+    windows = windows_for_region(len(contig), start, end, sequence_length, keep)
+    crop = crop_of(sequence_length, keep)
+    per_window = []
+    for window in windows:
+        plus, rc = predict_window(
+            model,
+            tokenizer,
+            window_sequence(contig, window, sequence_length),
+            device,
+            dtype,
+            one_strand_at_a_time,
+        )
+        n = min(window.start + crop + len(plus), len(contig)) - (window.start + crop)
+        per_window.append((plus[:n], rc[len(rc) - n :]))
+    return RegionPrediction(
+        chrom, windows[0].keep_start, windows[-1].keep_stop, windows, per_window
+    )
+
+
+def decode_region(
+    prediction: RegionPrediction, contig: str, decoder: str
+) -> tuple[np.ndarray, np.ndarray]:
+    """Plus- and minus-strand state paths over the region, in genomic order.
+
+    Each window is decoded over its whole kept region, as in the pipeline; the
+    shifted last window then drops the bp its predecessor already decoded.
+    """
+    codes = decoding.encode_bases(contig)
+    n = prediction.end - prediction.start
+    plus, minus = np.zeros(n, np.int8), np.zeros(n, np.int8)
+    for window, (lp_plus, lp_rc) in zip(
+        prediction.windows, prediction.per_window, strict=True
+    ):
+        full_start = window.keep_stop - len(lp_plus)
+        window_codes = codes[full_start : window.keep_stop]
+        states_plus = decode(lp_plus, window_codes, decoder)
+        states_minus = decode(
+            lp_rc, decoding.reverse_complement_codes(window_codes), decoder
+        )[::-1]
+        skip = window.keep_start - full_start
+        out = slice(
+            window.keep_start - prediction.start, window.keep_stop - prediction.start
+        )
+        plus[out] = states_plus[skip:]
+        minus[out] = states_minus[skip:]
+    return plus, minus
+
+
+def region_genes(
+    plus: np.ndarray,
+    minus: np.ndarray,
+    contig: str,
+    chrom: str,
+    start: int,
+    filters: bool = True,
+    min_cds_length: int = 201,
+) -> list[str]:
+    """GFF3 lines of both strands, in contig coordinates.
+
+    With ``filters``, genes with an internal in-frame stop or less than
+    ``min_cds_length`` bp of CDS are dropped, as in the released annotations.
+    """
+    codes = (
+        decoding.encode_bases(contig[start : start + len(plus)]) if filters else None
+    )
+    lines = decoding.find_genes(plus, "+", chrom, codes)
+    lines += decoding.find_genes(minus, "-", chrom, codes)
+    lines = [_shift(line, start) for line in lines]
+    return drop_short_genes(lines, min_cds_length) if filters else lines
+
+
+def _shift(line: str, offset: int) -> str:
+    cols = line.split("\t")
+    cols[3], cols[4] = str(int(cols[3]) + offset), str(int(cols[4]) + offset)
+    if "Parent=" in cols[8]:  # sub-feature IDs carry their start
+        parent = cols[8].split("Parent=")[1]
+        cols[8] = f"ID={parent}_{cols[2]}_{cols[3]};Parent={parent}"
+    return "\t".join(cols)
+
+
+def drop_short_genes(lines: list[str], min_cds_length: int) -> list[str]:
+    """Drops ``find_genes`` genes with less than ``min_cds_length`` bp of CDS."""
+    cds_length: dict[str, int] = defaultdict(int)
+    for line in lines:
+        cols = line.split("\t")
+        if cols[2] == "CDS":
+            parent = cols[8].split("Parent=")[1]
+            cds_length[parent] += int(cols[4]) - int(cols[3]) + 1
+
+    def gene_of(line: str) -> str:
+        attributes = line.split("\t")[8]
+        if "Parent=" in attributes:
+            return attributes.split("Parent=")[1]
+        return attributes.split("ID=")[1].split(";")[0]
+
+    return [line for line in lines if cds_length[gene_of(line)] >= min_cds_length]

src/genome_annotation/metrics.py

@@ -0,0 +1,143 @@
+"""Gene- and CDS-level accuracy of a predicted GFF3 (Tiberius / BRAKER style).
+
+A CDS is correct if its chromosome, start, end, strand and phase all match the
+reference; a gene is correct if its whole set of CDS matches a reference gene.
+"""
+
+import bisect
+import gzip
+from collections import defaultdict
+from pathlib import Path
+
+import numpy as np
+
+CdsKey = tuple[str, int, int, str, str]  # (chrom, start, end, strand, phase), 1-based
+
+
+def _open_text(path: str | Path):
+    return gzip.open(path, "rt") if str(path).endswith(".gz") else open(path)
+
+
+def _attribute(attributes: str, key: str) -> str | None:
+    for field in attributes.split(";"):
+        if field.startswith(f"{key}="):
+            return field[len(key) + 1 :]
+    return None
+
+
+def parse_genes(
+    path: str | Path, chromosomes: set[str] | None = None
+) -> dict[str, list[CdsKey]]:
+    """Maps each gene ID to its CDS keys; CDS reach genes directly or via mRNAs."""
+    transcript_to_gene: dict[str, str] = {}
+    genes: dict[str, list[CdsKey]] = defaultdict(list)
+    with _open_text(path) as handle:
+        for line in handle:
+            if line.startswith("#"):
+                continue
+            cols = line.rstrip("\n").split("\t")
+            if len(cols) < 9 or (chromosomes and cols[0] not in chromosomes):
+                continue
+            kind = cols[2].lower()
+            if kind in ("mrna", "transcript"):
+                tid, parent = _attribute(cols[8], "ID"), _attribute(cols[8], "Parent")
+                if tid and parent:
+                    transcript_to_gene[tid] = parent.split(",")[0]
+            elif kind == "cds":
+                parent = (_attribute(cols[8], "Parent") or "").split(",")[0]
+                gene = transcript_to_gene.get(parent, parent)
+                genes[gene].append(
+                    (cols[0], int(cols[3]), int(cols[4]), cols[6], cols[7])
+                )
+    return dict(genes)
+
+
+def read_regions(path: str | Path) -> dict[str, list[tuple[int, int]]]:
+    """Reads a BED of regions into sorted, merged 1-based closed intervals."""
+    regions: dict[str, list[tuple[int, int]]] = defaultdict(list)
+    with _open_text(path) as handle:
+        for line in handle:
+            chrom, start, end = line.split("\t")[:3]
+            regions[chrom].append((int(start) + 1, int(end)))
+    merged: dict[str, list[tuple[int, int]]] = {}
+    for chrom, intervals in regions.items():
+        out: list[tuple[int, int]] = []
+        for start, end in sorted(intervals):
+            if out and start <= out[-1][1] + 1:
+                out[-1] = (out[-1][0], max(out[-1][1], end))
+            else:
+                out.append((start, end))
+        merged[chrom] = out
+    return merged
+
+
+def _overlaps(regions: dict[str, list[tuple[int, int]]], key: CdsKey) -> bool:
+    intervals = regions.get(key[0])
+    if not intervals:
+        return False
+    i = bisect.bisect_left(intervals, (key[2] + 1,)) - 1
+    return i >= 0 and intervals[i][1] >= key[1]
+
+
+def _prf(
+    tp_recall: int, n_ref: int, tp_precision: int, n_pred: int
+) -> dict[str, float]:
+    recall = tp_recall / n_ref if n_ref else 0.0
+    precision = tp_precision / n_pred if n_pred else 0.0
+    f1 = 2 * recall * precision / (recall + precision) if recall + precision else 0.0
+    return {"Precision": 100 * precision, "Recall": 100 * recall, "F1": 100 * f1}
+
+
+def evaluate_gff(
+    reference: dict[str, list[CdsKey]],
+    prediction: dict[str, list[CdsKey]],
+    pseudogenes: dict[str, list[tuple[int, int]]] | None = None,
+    min_cds_length: int = 201,
+) -> dict[str, dict[str, float]]:
+    """CDS- and gene-level precision, recall and F1 (in %).
+
+    Args:
+        reference: Reference genes, as returned by ``parse_genes``.
+        prediction: Predicted genes, restricted to the reference's chromosomes.
+        pseudogenes: Pseudogene regions away from reference CDS; predicted CDS
+            overlapping them are ignored.
+        min_cds_length: Predicted genes with less CDS than this (bp) are ignored.
+
+    Returns:
+        {"CDS": {...}, "Gene": {...}} with Precision, Recall and F1.
+    """
+    ref_cds = {key for keys in reference.values() for key in keys}
+    kept = {
+        gid: keys
+        for gid, keys in prediction.items()
+        if sum(end - start + 1 for _, start, end, _, _ in keys) >= min_cds_length
+    }
+    dropped = {
+        key for gid, keys in prediction.items() if gid not in kept for key in keys
+    }
+    pred_cds = {key for keys in prediction.values() for key in keys} - dropped
+    if pseudogenes:
+        pred_cds = {key for key in pred_cds if not _overlaps(pseudogenes, key)}
+        kept = {gid: [k for k in keys if k in pred_cds] for gid, keys in kept.items()}
+        kept = {gid: keys for gid, keys in kept.items() if keys}
+
+    tp = len(ref_cds & pred_cds)
+    ref_structures = [frozenset(keys) for keys in reference.values()]
+    pred_structures = [frozenset(keys) for keys in kept.values()]
+    ref_set, pred_set = set(ref_structures), set(pred_structures)
+    return {
+        "CDS": _prf(tp, len(ref_cds), tp, len(pred_cds)),
+        "Gene": _prf(
+            sum(s in pred_set for s in ref_structures),
+            len(ref_structures),
+            sum(s in ref_set for s in pred_structures),
+            len(pred_structures),
+        ),
+    }
+
+
+def per_class_f1(confusion: np.ndarray) -> np.ndarray:
+    """Per-class F1 from a (true, predicted) confusion matrix."""
+    tp = np.diag(confusion).astype(np.float64)
+    denominator = confusion.sum(axis=0) + confusion.sum(axis=1)
+    return np.divide(2 * tp, denominator, out=np.zeros_like(tp), where=denominator > 0)

src/genome_annotation/model.py

@@ -0,0 +1,118 @@
+"""NTv3 annotation model warm-started from NTv3-tracks, its loss and its inference."""
+
+import functools
+
+import numpy as np
+import torch
+from torch.utils.checkpoint import checkpoint
+from transformers import AutoConfig, AutoModel
+
+ANNOTATION_REPO = "InstaDeepAI/NTv3_annot_plants"
+TRACKS_REPO = "InstaDeepAI/NTv3_tracks_1mb"
+# Revision of the modelling code (hosted in InstaDeepAI/NTv3_pretrained), pinned
+# because batch-norm running statistics depend on it.
+CODE_REVISION = "ea80a2df95bb605f33e1077b652da3c36c1dad56"
+
+# Which parts of the network are trained; the rest stays at the NTv3-tracks weights.
+TRAINABLE_PARTS = {
+    "all": ("",),
+    "decoder": ("model.decoder_blocks.", "model.bed_head."),
+    "head": ("model.bed_head.",),
+}
+
+
+def set_training(model: torch.nn.Module, training: bool) -> torch.nn.Module:
+    """Switches the trained batch norms between batch statistics and running ones.
+
+    In training mode a batch norm normalises with the window's own statistics and
+    updates its running ones. Everything else stays in eval mode: NTv3 has no
+    dropout, and batch norms of frozen blocks keep their NTv3-tracks statistics.
+    """
+    model.eval()
+    if training and model.batch_norm == "batch":
+        for module in model.modules():
+            if hasattr(module, "var_ema") and module.scale.requires_grad:
+                module.train()
+    return model
+
+
+def load_model(
+    trainable: str = "all",
+    gradient_checkpointing: bool = False,
+    batch_norm: str = "frozen",
+) -> torch.nn.Module:
+    """NTv3-tracks trunk under a new 15-class, single-strand annotation head.
+
+    The architecture (and head size) comes from the released plant annotation
+    model's config; every trunk weight comes from NTv3-tracks.
+
+    Args:
+        trainable: "all", "decoder" (U-Net decoder + head) or "head".
+        gradient_checkpointing: Recompute block activations in the backward pass.
+        batch_norm: "frozen" keeps the NTv3-tracks running statistics; "batch"
+            normalises trained blocks with per-window statistics during training
+            and updates their running statistics, as our training pipeline does.
+    """
+    if batch_norm == "batch" and gradient_checkpointing:
+        # Recomputing a block in the backward pass would update its statistics twice.
+        raise ValueError("batch_norm='batch' does not support gradient checkpointing.")
+    config = AutoConfig.from_pretrained(
+        ANNOTATION_REPO, trust_remote_code=True, code_revision=CODE_REVISION
+    )
+    config.ntv3["keep_target_center_fraction"] = 1.0
+    model = AutoModel.from_pretrained(
+        TRACKS_REPO, config=config, trust_remote_code=True, code_revision=CODE_REVISION
+    )
+    prefixes = TRAINABLE_PARTS[trainable]
+    for name, param in model.named_parameters():
+        param.requires_grad = name.startswith(prefixes)
+    if gradient_checkpointing:
+        trunk = model.model
+        for block in [
+            *trunk.encoder_blocks,
+            *trunk.mha_blocks,
+            *trunk.mlp_blocks,
+            *trunk.decoder_blocks,
+        ]:
+            block.forward = functools.partial(
+                checkpoint, block.forward, use_reentrant=False
+            )
+    model.batch_norm = batch_norm
+    return model.eval()
+
+
+def logits(model: torch.nn.Module, tokens: torch.Tensor) -> torch.Tensor:
+    """(batch, length) tokens -> (batch, length, 15) annotation logits."""
+    out = model(tokens, conditions_tokens=[], skip_lm_head=True)
+    return out["bed_tracks_logits"][:, :, 0, :]
+
+
+def focal_loss(
+    logits: torch.Tensor,
+    labels: torch.Tensor,
+    class_weights: torch.Tensor,
+    gamma: float = 2.0,
+) -> torch.Tensor:
+    """Class-weighted focal loss averaged over labelled positions.
+
+    Each position is weighted by the larger of its true and predicted class
+    weights, so false positives of rare classes cost as much as misses.
+    """
+    log_probs = torch.log_softmax(logits.float(), dim=-1)
+    valid = labels >= 0
+    target = labels.clamp(min=0)
+    log_p = log_probs.gather(-1, target[..., None])[..., 0]
+    loss = -((1 - log_p.exp()) ** gamma) * log_p
+    weight = torch.maximum(class_weights[target], class_weights[log_probs.argmax(-1)])
+    return (loss * weight * valid).sum() / valid.sum().clamp(min=1)
+
+
+@torch.no_grad()
+def predict_log_probs(
+    model: torch.nn.Module, tokens: np.ndarray, device: torch.device, dtype: torch.dtype
+) -> np.ndarray:
+    """(batch, length) tokens -> (batch, length, 15) float32 log-probabilities."""
+    batch = torch.from_numpy(tokens).to(device)
+    with torch.autocast(device.type, dtype=dtype, enabled=dtype != torch.float32):
+        out = logits(model, batch)
+    return torch.log_softmax(out.float(), dim=-1).cpu().numpy()

src/genome_annotation/plotting.py

@@ -0,0 +1,710 @@
+"""Figures of the annotation notebooks (matplotlib, one call per figure)."""
+
+import itertools
+from collections.abc import Mapping, Sequence
+from typing import Any
+
+import matplotlib.pyplot as plt
+import numpy as np
+from matplotlib.axes import Axes
+from matplotlib.colors import to_rgb
+from matplotlib.figure import Figure
+from matplotlib.patches import Rectangle
+
+from . import decoding
+from .analysis import cds_span
+from .benchmark import TestSpecies
+from .evaluation import Gene
+
+# Paper hues; the gold is darkened to sit in the readable lightness band. `correct`
+# and `invented` are deliberate neutrals so the three error classes carry the colour.
+CLASS_COLORS = {
+    "correct": "#b3b3b3",
+    "boundary": "#d99a00",
+    "missed": "#b428a0",
+    "fragmented": "#ff554d",
+    "invented": "#4b5563",
+}
+EXON_COLORS = {
+    "exact": CLASS_COLORS["correct"],
+    "shifted_one": CLASS_COLORS["boundary"],
+    "shifted_both": CLASS_COLORS["fragmented"],
+    "missed": CLASS_COLORS["missed"],
+}
+NTV3_COLOR = "#3973fc"
+OTHER_COLOR = "#808080"
+INK, MUTED = "#1f2328", "#6b7280"
+BAR_HEIGHT = 0.6
+
+
+def _style(ax: Axes) -> None:
+    for side in ("top", "right"):
+        ax.spines[side].set_visible(False)
+    for side in ("left", "bottom"):
+        ax.spines[side].set_color(MUTED)
+    ax.tick_params(colors=MUTED, labelcolor=INK)
+    ax.grid(axis="x", color="#e5e7eb", linewidth=0.8)
+    ax.set_axisbelow(True)
+
+
+def _text_color(background: str) -> str:
+    r, g, b = to_rgb(background)
+    return "white" if 0.2126 * r + 0.7152 * g + 0.0722 * b < 0.45 else INK
+
+
+def _stacked_rows(
+    ax: Axes,
+    rows: Sequence[str],
+    shares: Sequence[Mapping[str, float]],
+    colors: Mapping[str, str],
+) -> None:
+    """Horizontal 100% bars, one per row, labelled inside segments wide enough."""
+    for y, share in enumerate(shares):
+        left = 0.0
+        for klass, color in colors.items():
+            width = share.get(klass, 0.0)
+            ax.barh(
+                y,
+                width,
+                left=left,
+                height=BAR_HEIGHT,
+                color=color,
+                edgecolor="white",
+                linewidth=2,
+            )
+            if width >= 6:
+                ax.text(
+                    left + width / 2,
+                    y,
+                    f"{width:.0f}%",
+                    ha="center",
+                    va="center",
+                    fontsize=8,
+                    color=_text_color(color),
+                )
+            left += width
+    ax.set_yticks(range(len(rows)), rows)
+    ax.invert_yaxis()
+    ax.set_xlim(0, 100)
+    ax.set_xlabel("share (%)")
+    handles = [Rectangle((0, 0), 1, 1, color=c) for c in colors.values()]
+    labels = [k.replace("_", " ") for k in colors]
+    ax.legend(
+        handles,
+        labels,
+        ncol=len(colors),
+        frameon=False,
+        loc="lower left",
+        bbox_to_anchor=(0, 1.0),
+        fontsize=8,
+    )
+    _style(ax)
+
+
+def _percent(counts: Mapping[str, int], keys: Sequence[str]) -> dict[str, float]:
+    total = sum(counts[k] for k in keys)
+    return {k: 100.0 * counts[k] / total if total else 0.0 for k in keys}
+
+
+def plot_scores(
+    results: Mapping[str, Mapping[str, Mapping[str, float]]],
+    species: TestSpecies | None = None,
+) -> Figure:
+    """Gene and exon F1 of each annotation, with the published NTv3 score as a tick.
+
+    Args:
+        results: Annotation label to ``evaluation.evaluate`` output.
+        species: Benchmark species, to mark the paper's NTv3 numbers.
+    """
+    fig, axes = plt.subplots(1, 2, figsize=(9, 0.6 + 0.5 * len(results)), sharey=True)
+    labels = list(results)
+    for ax, (level, title) in zip(
+        axes, (("Gene", "Gene F1 (%)"), ("CDS", "Exon (CDS) F1 (%)")), strict=True
+    ):
+        values = [100 * results[label][level]["F1"] for label in labels]
+        colors = [NTV3_COLOR if "NTv3" in label else OTHER_COLOR for label in labels]
+        ax.barh(labels, values, height=BAR_HEIGHT, color=colors)
+        for y, value in enumerate(values):
+            ax.text(value + 1, y, f"{value:.2f}", va="center", fontsize=9, color=INK)
+        if species is not None:
+            paper = species.gene_f1 if level == "Gene" else species.cds_f1
+            ax.axvline(paper, color=INK, linestyle=":", linewidth=1, label="paper")
+            ax.legend(
+                frameon=False, fontsize=8, loc="lower right", bbox_to_anchor=(1, 1)
+            )
+        ax.set_xlim(0, 110)
+        ax.set_title(title, loc="left", fontsize=10, color=INK)
+        _style(ax)
+    axes[0].invert_yaxis()
+    fig.tight_layout()
+    return fig
+
+
+def plot_locus_composition(summaries: Mapping[str, Mapping[str, float]]) -> Figure:
+    """Where every reference gene ends up, and the share of invented predictions."""
+    fig, (ax, side) = plt.subplots(
+        1,
+        2,
+        figsize=(10, 1.2 + 0.6 * len(summaries)),
+        gridspec_kw={"width_ratios": (5, 1)},
+        sharey=True,
+    )
+    classes = ("correct", "boundary", "missed", "fragmented")
+    shares = [
+        _percent({k: int(s[k]) for k in classes}, classes) for s in summaries.values()
+    ]
+    _stacked_rows(ax, list(summaries), shares, {k: CLASS_COLORS[k] for k in classes})
+    ax.set_xlabel("reference genes (%)")
+    for y, summary in enumerate(summaries.values()):
+        share = 100 * summary["invented"] / summary["n_pred"]
+        side.barh(y, share, height=BAR_HEIGHT, color=CLASS_COLORS["invented"])
+        side.text(share, y, f" {share:.1f}%", va="center", fontsize=8, color=INK)
+    side.set_xlim(0, max(10, 1.4 * side.get_xlim()[1]))
+    side.set_xlabel("invented\n(% of predictions)")
+    _style(side)
+    fig.tight_layout()
+    return fig
+
+
+def plot_recovery_by_size(recovery: Mapping[str, Any]) -> Figure:
+    """Locus outcome by CDS span, and exact recovery by exon count."""
+    fig, (left, right) = plt.subplots(1, 2, figsize=(11, 3.6))
+    span_bins = [_short_span(label) for label in recovery["span_bins"]]
+    totals = recovery["span_total"]
+    outcome = recovery["span_outcome"]
+    shares = [
+        _percent({k: outcome[k][i] for k in outcome}, list(outcome))
+        for i in range(len(span_bins))
+    ]
+    rows = [f"{b} (n={n:,})" for b, n in zip(span_bins, totals, strict=True)]
+    _stacked_rows(left, rows, shares, {k: CLASS_COLORS[k] for k in outcome})
+    left.set_title("by CDS span", loc="left", fontsize=10, color=INK, pad=22)
+
+    bins = recovery["exon_count_bins"]
+    total = recovery["exon_total"]
+    correct = recovery["exon_correct"]
+    rate = [100 * c / t if t else 0.0 for c, t in zip(correct, total, strict=True)]
+    right.bar(bins, rate, color=NTV3_COLOR, width=0.6)
+    for x, (value, n) in enumerate(zip(rate, total, strict=True)):
+        right.text(x, value + 1, f"{value:.0f}%\nn={n:,}", ha="center", fontsize=8)
+    right.set_ylim(0, 110)
+    right.set_xlabel("coding exons per gene")
+    right.set_ylabel("genes recovered exactly (%)")
+    right.set_title("by exon count", loc="left", fontsize=10, color=INK)
+    _style(right)
+    right.grid(axis="x", visible=False)
+    right.grid(axis="y", color="#e5e7eb", linewidth=0.8)
+    fig.tight_layout()
+    return fig
+
+
+def _short_span(label: str) -> str:
+    lo, hi = (int(x) for x in label.split("-"))
+    fmt = lambda v: f"{v / 1e6:.0f} Mb" if v >= 1e6 else f"{v / 1e3:.0f} kb"  # noqa: E731
+    return (
+        f">{fmt(lo)}" if hi >= 10**12 else f"{fmt(lo)}-{fmt(hi)}".replace("0 kb-", "0-")
+    )
+
+
+def plot_exon_taxonomy(exons: Mapping[str, Any]) -> Figure:
+    """Reference exons: exact, one or both boundaries wrong, or missed."""
+    by_position = exons["by_position"]
+    by_length = exons["by_length"]
+    classes = list(EXON_COLORS)
+    fig, (top, bottom) = plt.subplots(2, 1, figsize=(9, 6.6))
+    positions = [
+        p
+        for p in ("first", "internal", "last", "single")
+        if sum(by_position[p].values())
+    ]
+    totals = exons["totals"]
+    _stacked_rows(
+        top,
+        [f"all exons (n={sum(totals.values()):,})"]
+        + [f"{p} (n={sum(by_position[p].values()):,})" for p in positions],
+        [_percent(totals, classes)]
+        + [_percent(by_position[p], classes) for p in positions],
+        EXON_COLORS,
+    )
+    top.set_title("by position in the gene", loc="left", fontsize=10, color=INK, pad=22)
+    lengths = exons["exon_length_bins"]
+    _stacked_rows(
+        bottom,
+        [
+            f"{b.replace('+', '+ ')} bp (n={sum(c.values()):,})"
+            for b, c in zip(lengths, by_length, strict=True)
+        ],
+        [_percent(c, classes) for c in by_length],
+        EXON_COLORS,
+    )
+    bottom.set_title("by exon length", loc="left", fontsize=10, color=INK, pad=22)
+    fig.tight_layout()
+    return fig
+
+
+def plot_failure_anatomy(anatomy: Mapping[str, Any]) -> Figure:
+    """How many exons are wrong in each failed gene, and which one when it is one."""
+    grid = anatomy["wrong_by_gene_size"]
+    sizes = anatomy["exon_count_bins"]
+    buckets = anatomy["wrong_exon_buckets"]
+    fig, (left, right) = plt.subplots(1, 2, figsize=(11, 3.6))
+    ramp = ["#c3d5fc", "#7ea4fc", "#3973fc", "#1d56d6", "#004697", "#002a5c"]
+    shares = [
+        {
+            b: 100 * c / sum(row) if sum(row) else 0.0
+            for b, c in zip(buckets, row, strict=True)
+        }
+        for row in grid
+    ]
+    rows = [
+        f"{s} exon{'' if s == '1' else 's'} (n={sum(r):,})"
+        for s, r in zip(sizes, grid, strict=True)
+    ]
+    _stacked_rows(left, rows, shares, dict(zip(buckets, ramp, strict=False)))
+    left.set_title(
+        "wrong exons per failed gene", loc="left", fontsize=10, color=INK, pad=22
+    )
+
+    observed = anatomy["single_wrong_observed"]
+    expected = anatomy["single_wrong_expected"]
+    positions = list(observed)
+    x = range(len(positions))
+    right.bar(
+        [i - 0.2 for i in x],
+        [observed[p] for p in positions],
+        width=0.4,
+        edgecolor="white",
+        linewidth=2,
+        color=NTV3_COLOR,
+        label="observed",
+    )
+    right.bar(
+        [i + 0.2 for i in x],
+        [expected[p] for p in positions],
+        width=0.4,
+        edgecolor="white",
+        linewidth=2,
+        color=OTHER_COLOR,
+        label="if uniform",
+    )
+    right.set_xticks(list(x), positions)
+    right.set_ylabel("failed genes")
+    right.set_title(
+        "position of the single wrong exon", loc="left", fontsize=10, color=INK
+    )
+    right.legend(frameon=False, fontsize=8)
+    _style(right)
+    right.grid(axis="x", visible=False)
+    fig.tight_layout()
+    return fig
+
+
+def plot_independence(check: Mapping[str, Any]) -> Figure:
+    """Observed gene recovery vs. what independent exon errors would give."""
+    bins = check["exon_count_bins"]
+    n = check["n_genes"]
+    observed = [
+        100 * o / t if t else 0 for o, t in zip(check["observed"], n, strict=True)
+    ]
+    predicted = [
+        100 * p / t if t else 0 for p, t in zip(check["predicted"], n, strict=True)
+    ]
+    fig, ax = plt.subplots(figsize=(6, 3.4))
+    ax.plot(bins, observed, marker="o", color=NTV3_COLOR, linewidth=2, label="observed")
+    ax.plot(
+        bins,
+        predicted,
+        marker="o",
+        color=OTHER_COLOR,
+        linewidth=2,
+        linestyle="--",
+        label="independent exons",
+    )
+    ax.set_ylim(0, 100)
+    ax.set_xlabel("coding exons per gene")
+    ax.set_ylabel("genes recovered exactly (%)")
+    ax.legend(frameon=False, fontsize=8)
+    _style(ax)
+    ax.grid(axis="y", color="#e5e7eb", linewidth=0.8)
+    fig.tight_layout()
+    return fig
+
+
+def plot_boundary_offsets(offsets: Mapping[str, Any]) -> Figure:
+    """How far a wrong exon boundary is from the reference one."""
+    bins = offsets["offset_bins"]
+    counts = offsets["offsets"]
+    total = sum(counts)
+    fig, ax = plt.subplots(figsize=(7, 3.2))
+    shares = [100 * c / total if total else 0 for c in counts]
+    ax.bar([f"{b} bp" for b in bins], shares, color=NTV3_COLOR, width=0.6)
+    for x, value in enumerate(shares):
+        ax.text(x, value + 0.5, f"{value:.0f}%", ha="center", fontsize=8)
+    ax.set_ylabel("shifted boundaries (%)")
+    ax.set_xlabel("distance to the reference boundary")
+    _style(ax)
+    ax.grid(axis="x", visible=False)
+    fig.tight_layout()
+    return fig
+
+
+def plot_structure(chains: Mapping[str, Mapping[str, float]]) -> Figure:
+    """Exact gene F1 vs. intron-chain F1 (outer ends forgiven), by gene type."""
+    groups = [
+        ("multi_exon", "multi-exon"),
+        ("single_exon", "single-exon"),
+        ("all", "all"),
+    ]
+    fig, ax = plt.subplots(figsize=(6.5, 3.2))
+    for i, (key, _) in enumerate(groups):
+        for offset, (metric, color, label) in zip(
+            (-0.2, 0.2),
+            (
+                ("exact_f1", NTV3_COLOR, "exact structure"),
+                ("chain_f1", "#004697", "intron chain"),
+            ),
+            strict=True,
+        ):
+            value = chains[key][metric]
+            ax.bar(
+                i + offset,
+                value,
+                width=0.4,
+                edgecolor="white",
+                linewidth=2,
+                color=color,
+                label=label if i == 0 else None,
+            )
+            ax.text(i + offset, value + 1, f"{value:.1f}", ha="center", fontsize=8)
+    ax.set_xticks(range(len(groups)), [label for _, label in groups])
+    ax.set_ylim(0, 110)
+    ax.set_ylabel("gene F1 (%)")
+    ax.legend(frameon=False, fontsize=8, ncol=2, loc="upper left")
+    _style(ax)
+    ax.grid(axis="x", visible=False)
+    fig.tight_layout()
+    return fig
+
+
+def plot_proteins(counts: Mapping[str, Any]) -> Figure:
+    """Identity of each reference protein's best predicted counterpart."""
+    histogram = counts["identity_histogram"]
+    total = sum(histogram.values())
+    fig, ax = plt.subplots(figsize=(6.5, 3.2))
+    labels = list(histogram)
+    shares = [100 * histogram[k] / total if total else 0 for k in labels]
+    ax.bar([f"{k}%" for k in labels], shares, color=NTV3_COLOR, width=0.6)
+    for x, value in enumerate(shares):
+        ax.text(x, value + 1, f"{value:.0f}%", ha="center", fontsize=8)
+    ax.set_ylabel("reference proteins (%)")
+    ax.set_xlabel("identity to the best overlapping prediction")
+    _style(ax)
+    ax.grid(axis="x", visible=False)
+    fig.tight_layout()
+    return fig
+
+
+def plot_failed_gene_repair(
+    exons: Mapping[str, Any], anatomy: Mapping[str, Any], summary: Mapping[str, Any]
+) -> Figure:
+    """Wrong exons per failed gene, and the genes a fix of at most k exons would add."""
+    buckets = exons["wrong_exon_buckets"]
+    per_gene = exons["wrong_exons_per_failed_gene"]
+    fig, (left, right) = plt.subplots(1, 2, figsize=(11, 3.4))
+    failed = sum(per_gene)
+    shares = [100 * n / failed if failed else 0 for n in per_gene]
+    left.bar(buckets, shares, color=NTV3_COLOR, width=0.6)
+    for x, value in enumerate(shares):
+        left.text(x, value + 1, f"{value:.0f}%", ha="center", fontsize=8)
+    left.set_xlabel("wrong exons in the gene")
+    left.set_ylabel("failed genes (%)")
+    left.set_title("how wrong a failed gene is", loc="left", fontsize=10, color=INK)
+    n_loci = summary["n_loci"]
+    recall = 100 * summary["correct"] / n_loci
+    gained = [100 * c / n_loci for c in anatomy["cumulative_convertible"]]
+    right.bar(
+        buckets,
+        [recall] * len(buckets),
+        color=OTHER_COLOR,
+        width=0.6,
+        label="recovered",
+    )
+    right.bar(
+        buckets,
+        gained,
+        bottom=recall,
+        color=NTV3_COLOR,
+        width=0.6,
+        label="after the fix",
+    )
+    for x, value in enumerate(gained):
+        right.text(
+            x, recall + value + 1, f"{recall + value:.0f}%", ha="center", fontsize=8
+        )
+    right.set_ylim(0, 105)
+    right.set_xlabel("fix at most this many exons per failed gene")
+    right.set_ylabel("reference genes exact (%)")
+    right.set_title(
+        "what fixing a few exons would buy", loc="left", fontsize=10, color=INK
+    )
+    right.legend(frameon=False, fontsize=8, loc="lower right")
+    for ax in (left, right):
+        _style(ax)
+        ax.grid(axis="x", visible=False)
+    fig.tight_layout()
+    return fig
+
+
+def plot_structure_and_merges(
+    chains: Mapping[str, Mapping[str, float]], merges: Mapping[str, Any]
+) -> Figure:
+    """Exact vs. intron-chain gene F1, and how often one prediction merges loci."""
+    fig, (left, right) = plt.subplots(1, 2, figsize=(10, 3.2))
+    values = [chains["all"]["exact_f1"], chains["all"]["chain_f1"]]
+    left.bar(
+        ["exact structure", "intron chain"],
+        values,
+        color=[NTV3_COLOR, "#004697"],
+        width=0.5,
+    )
+    for x, value in enumerate(values):
+        left.text(x, value + 1, f"{value:.1f}", ha="center", fontsize=9)
+    left.set_ylim(0, 105)
+    left.set_ylabel("gene F1 (%)")
+    shares = [
+        100 * merges["absorbed_loci"] / merges["n_ref"],
+        100 * merges["merging_predictions"] / merges["n_pred"],
+    ]
+    right.barh(
+        ["reference genes\nabsorbed in a merge", "predictions that\nmerge 2+ genes"],
+        shares,
+        color=CLASS_COLORS["fragmented"],
+        height=0.5,
+    )
+    for y, value in enumerate(shares):
+        right.text(value, y, f" {value:.1f}%", va="center", fontsize=9)
+    right.set_xlim(0, max(5, 1.3 * max(shares)))
+    right.set_xlabel("%")
+    right.invert_yaxis()
+    for ax in (left, right):
+        _style(ax)
+    left.grid(axis="x", visible=False)
+    fig.tight_layout()
+    return fig
+
+
+def plot_protein_recovery(counts: Mapping[str, Any]) -> Figure:
+    """Reference proteins recovered at each identity, and well-formed predicted ORFs."""
+    fig, (left, right) = plt.subplots(
+        1, 2, figsize=(10, 3.2), gridspec_kw={"width_ratios": (3, 1)}
+    )
+    n_ref = counts["n_reference"]
+    thresholds = list(counts["recovered"])
+    shares = [100 * counts["recovered"][t] / n_ref for t in thresholds]
+    left.bar([f">= {t}%" for t in thresholds], shares, color=NTV3_COLOR, width=0.6)
+    for x, value in enumerate(shares):
+        left.text(x, value + 1, f"{value:.1f}%", ha="center", fontsize=8)
+    left.set_ylim(0, 105)
+    left.set_xlabel("identity to the best overlapping prediction")
+    left.set_ylabel("reference proteins (%)")
+    well = 100 * counts["well_formed"] / counts["n_predicted"]
+    right.bar(["well-formed\nORFs"], [well], color=NTV3_COLOR, width=0.5)
+    right.text(0, well + 1, f"{well:.2f}%", ha="center", fontsize=8)
+    right.set_ylim(0, 105)
+    right.set_ylabel("predicted proteins (%)")
+    for ax in (left, right):
+        _style(ax)
+        ax.grid(axis="x", visible=False)
+    fig.tight_layout()
+    return fig
+
+
+def plot_locus(
+    reference: Gene,
+    predictions: Sequence[Gene],
+    title: str = "",
+    pad_fraction: float = 0.04,
+) -> Figure:
+    """Gene models of one reference locus and the predictions overlapping it."""
+    genes = [reference, *predictions]
+    spans = [cds_span(g) for g in genes if g.cds_keys]
+    start, end = min(s for s, _ in spans), max(e for _, e in spans)
+    pad = max(1, int((end - start) * pad_fraction))
+    fig, ax = plt.subplots(figsize=(10, 0.8 + 0.5 * len(genes)))
+    ref_boundaries = {k[:4] for k in reference.cds_keys}
+    for y, gene in enumerate(genes):
+        lo, hi = cds_span(gene)
+        ax.plot([lo, hi], [y, y], color=MUTED, linewidth=1)
+        for key in gene.cds_keys:
+            good = y == 0 or key[:4] in ref_boundaries
+            color = (
+                (OTHER_COLOR if y == 0 else NTV3_COLOR)
+                if good
+                else CLASS_COLORS["fragmented"]
+            )
+            ax.add_patch(
+                Rectangle((key[1], y - 0.25), key[2] - key[1] + 1, 0.5, color=color)
+            )
+    ax.set_yticks(range(len(genes)), ["reference", *[g.gene_id for g in predictions]])
+    ax.invert_yaxis()
+    ax.set_xlim(start - pad, end + pad)
+    ax.xaxis.set_major_formatter(lambda x, _: f"{x:,.0f}")
+    strand = "forward" if reference.strand == "+" else "reverse"
+    ax.set_xlabel(f"{reference.chrom} ({strand} strand), bp")
+    ax.set_title(title or reference.gene_id, loc="left", fontsize=10, color=INK)
+    _style(ax)
+    fig.tight_layout()
+    return fig
+
+
+# --- Decoding walkthrough (03b) ---------------------------------------------------
+
+_INTRON = np.array([c.startswith("intron") for c in decoding.CLASSES])
+_EXONIC = ~_INTRON & (np.arange(decoding.NUM_CLASSES) != 0)
+GROUP_COLORS = {"intergenic": "#e5e7eb", "exon": NTV3_COLOR, "intron": "#c3d5fc"}
+
+
+def plot_transitions() -> Figure:
+    """The HMM's allowed transitions and their probabilities (log scale)."""
+    probs = np.exp(decoding.LOG_TRANSITIONS)
+    fig, ax = plt.subplots(figsize=(7, 6))
+    shown = np.where(probs > 0, np.log10(np.maximum(probs, 1e-12)), np.nan)
+    image = ax.imshow(shown, cmap="Blues", vmin=-5, vmax=0)
+    ax.set_xticks(
+        range(decoding.NUM_CLASSES), decoding.CLASSES, rotation=90, fontsize=8
+    )
+    ax.set_yticks(range(decoding.NUM_CLASSES), decoding.CLASSES, fontsize=8)
+    ax.set_xlabel("to")
+    ax.set_ylabel("from")
+    fig.colorbar(image, ax=ax, label="log10 probability (blank: forbidden)")
+    fig.tight_layout()
+    return fig
+
+
+def _feature_runs(
+    states: np.ndarray,
+) -> tuple[list[tuple[int, int]], list[tuple[int, int]]]:
+    """(exon runs, intron runs) of a state path, as half-open index ranges."""
+    kind = np.where(_INTRON[states], 2, np.where(_EXONIC[states], 1, 0))
+    edges = np.flatnonzero(np.diff(np.concatenate([[-1], kind, [-1]])))
+    runs = list(itertools.pairwise(edges))
+    return (
+        [(a, b) for a, b in runs if kind[a] == 1],
+        [(a, b) for a, b in runs if kind[a] == 2],
+    )
+
+
+def plot_tracks(
+    tracks: Mapping[str, np.ndarray],
+    region_start: int,
+    zoom: tuple[int, int],
+    reference: Sequence[Gene] = (),
+    title: str = "",
+    mark_errors: bool = False,
+) -> Figure:
+    """Gene structure of each state path over ``zoom`` (contig bp), reference on top.
+
+    Args:
+        tracks: Label to one strand's state path over the region, genomic order.
+        region_start: Contig bp of the paths' first position.
+        zoom: Contig bp range to draw.
+        reference: Reference genes of that strand to draw on the first row.
+        title: Axes title.
+        mark_errors: Draw predicted exons whose boundaries match no reference exon
+            in red.
+    """
+    ref_exons = {(k[1], k[2]) for gene in reference for k in gene.cds_keys}
+    rows = ["reference", *tracks] if reference else list(tracks)
+    fig, ax = plt.subplots(figsize=(11, 0.9 + 0.45 * len(rows)))
+    lo, hi = zoom
+    y = 0
+    if reference:
+        for gene in reference:
+            s, e = cds_span(gene)
+            if e < lo or s > hi:
+                continue
+            ax.plot([s, e], [y, y], color=MUTED, linewidth=1)
+            for key in gene.cds_keys:
+                ax.add_patch(
+                    Rectangle(
+                        (key[1], y - 0.25), key[2] - key[1] + 1, 0.5, color=OTHER_COLOR
+                    )
+                )
+        y += 1
+    for states in tracks.values():
+        window = states[lo - region_start : hi - region_start]
+        exons, introns = _feature_runs(window)
+        for a, b in introns:
+            ax.plot([lo + a, lo + b], [y, y], color=MUTED, linewidth=1)
+        for a, b in exons:
+            wrong = mark_errors and (lo + a + 1, lo + b) not in ref_exons
+            color = CLASS_COLORS["fragmented"] if wrong else NTV3_COLOR
+            ax.add_patch(Rectangle((lo + a, y - 0.25), b - a, 0.5, color=color))
+        y += 1
+    ax.set_yticks(range(len(rows)), rows)
+    ax.set_ylim(len(rows) - 0.5, -0.5)
+    ax.set_xlim(lo, hi)
+    ax.xaxis.set_major_formatter(lambda x, _: f"{x:,.0f}")
+    ax.set_title(title, loc="left", fontsize=10, color=INK)
+    _style(ax)
+    fig.tight_layout()
+    return fig
+
+
+def plot_class_probabilities(
+    log_probs: np.ndarray, region_start: int, zoom: tuple[int, int], title: str = ""
+) -> Figure:
+    """Probability of intergenic / exon / intron along ``zoom``, stacked to 1."""
+    lo, hi = zoom
+    probs = np.exp(log_probs[lo - region_start : hi - region_start])
+    groups = {
+        "intergenic": probs[:, 0],
+        "exon": probs[:, _EXONIC].sum(1),
+        "intron": probs[:, _INTRON].sum(1),
+    }
+    fig, ax = plt.subplots(figsize=(11, 2.4))
+    x = np.arange(lo, hi)
+    ax.stackplot(
+        x, *groups.values(), colors=list(GROUP_COLORS.values()), labels=list(groups)
+    )
+    ax.set_xlim(lo, hi)
+    ax.set_ylim(0, 1)
+    ax.set_ylabel("probability")
+    ax.xaxis.set_major_formatter(lambda v, _: f"{v:,.0f}")
+    ax.legend(
+        frameon=False, fontsize=8, ncol=3, loc="lower left", bbox_to_anchor=(0, 1)
+    )
+    ax.set_title(title, loc="right", fontsize=10, color=INK)
+    _style(ax)
+    fig.tight_layout()
+    return fig
+
+
+def plot_edge_drift(
+    distance_kb: np.ndarray, drift: np.ndarray, crop_kb: float
+) -> Figure:
+    """How far predictions move from their full-context values near a window edge.
+
+    Args:
+        distance_kb: Bin centres, kb from the window edge.
+        drift: Mean total-variation distance to the full-context prediction.
+        crop_kb: Context margin the pipeline discards, drawn as a line.
+    """
+    fig, ax = plt.subplots(figsize=(7.5, 3.2))
+    ax.axvspan(
+        distance_kb.min() / 2, crop_kb, color="#eef2ff", label="discarded margin"
+    )
+    ax.plot(distance_kb, drift, color=NTV3_COLOR, linewidth=2, marker="o", markersize=5)
+    ax.set_xscale("log")
+    ax.set_xlim(distance_kb.min() / 2, distance_kb.max() * 1.5)
+    ticks = [1, 4, 16, 64, round(crop_kb)]
+    ax.set_xticks(ticks, [f"{t:g}" for t in ticks])
+    ax.minorticks_off()
+    ax.set_ylim(0, None)
+    ax.set_xlabel("distance to the window edge (kb)")
+    ax.set_ylabel("change in predicted\nprobabilities (TV)")
+    ax.legend(frameon=False, fontsize=8)
+    _style(ax)
+    ax.grid(axis="y", color="#e5e7eb", linewidth=0.8)
+    fig.tight_layout()
+    return fig

src/genome_annotation/proteins.py

@@ -0,0 +1,150 @@
+"""Protein-level checks of a predicted annotation (needs the genome FASTA).
+
+Translates every gene, checks predicted ORFs (start with M, no internal stop), and
+scores each reference protein by the identity of its best overlapping prediction.
+"""
+
+import dataclasses
+from collections.abc import Callable, Iterable, Mapping
+from dataclasses import dataclass, field
+
+from Bio.Align import PairwiseAligner
+from Bio.Seq import Seq
+
+from .analysis import CdsIndex
+from .evaluation import Gene
+
+IDENTITY_THRESHOLDS = (100.0, 95.0, 90.0, 50.0)
+IDENTITY_BINS = ("100", "95-100", "90-95", "50-90", "<50")
+# Shorter proteins are fragments, not scored for ORF shape.
+MIN_PROTEIN_LENGTH = 20
+
+
+@dataclass
+class ProteinCounts:
+    """Protein-level tallies; ORF fields cover predictions of 20+ residues."""
+
+    n_reference: int = 0
+    recovered: dict[str, int] = field(
+        default_factory=lambda: {f"{t:g}": 0 for t in IDENTITY_THRESHOLDS}
+    )
+    no_counterpart: int = 0
+    n_predicted: int = 0
+    starts_with_methionine: int = 0
+    no_internal_stop: int = 0
+    well_formed: int = 0
+    identity_histogram: dict[str, int] = field(
+        default_factory=lambda: dict.fromkeys(IDENTITY_BINS, 0)
+    )
+
+
+def translate_gene(gene: Gene, sequence: str) -> str:
+    """Protein encoded by a gene's CDS on its (uppercase) chromosome sequence.
+
+    CDS are joined in ascending genomic order and reverse-complemented once on the
+    minus strand; the phase of the transcript's first CDS is honoured, a trailing
+    partial codon and one terminal stop are dropped.
+    """
+    ascending = sorted(gene.cds_keys, key=lambda k: k[1])
+    coding = "".join(sequence[start - 1 : end] for _, start, end, _, _ in ascending)
+    if gene.strand == "-":
+        coding = str(Seq(coding).reverse_complement())
+    phase = (ascending[-1] if gene.strand == "-" else ascending[0])[4]
+    coding = coding[int(phase) if phase.isdigit() else 0 :]
+    coding = coding[: len(coding) - len(coding) % 3]
+    if not coding:
+        return ""
+    protein = str(Seq(coding).translate())
+    return protein[:-1] if protein.endswith("*") else protein
+
+
+def identity(a: str, b: str, aligner: PairwiseAligner) -> float:
+    """Percent identity of a global alignment, over the longer protein."""
+    if a == b:
+        return 100.0
+    if not a or not b:
+        return 0.0
+    alignment = aligner.align(a, b)[0]
+    matches = sum(
+        x == y and x != "-" for x, y in zip(alignment[0], alignment[1], strict=True)
+    )
+    return 100.0 * matches / max(len(a), len(b))
+
+
+def _identity_bin(value: float) -> str:
+    for threshold, label in zip(IDENTITY_THRESHOLDS, IDENTITY_BINS, strict=False):
+        if value >= threshold:
+            return label
+    return "<50"
+
+
+def measure_proteins(
+    ref_genes: dict[str, Gene],
+    pred_genes: dict[str, Gene],
+    genome: Mapping[str, str] | Callable[[str], str],
+    chromosomes: Iterable[str] | None = None,
+) -> dict[str, object]:
+    """ORF shape of the predictions and identity recovery of the reference.
+
+    Args:
+        ref_genes: Reference genes, e.g. ``ScoringSets.ref_genes``.
+        pred_genes: Predicted genes.
+        genome: Chromosome name to sequence, or a function returning it (e.g. over a
+            ``pyfaidx.Fasta``); read one chromosome at a time.
+        chromosomes: Chromosomes to measure; defaults to every one with a gene.
+
+    Returns:
+        ``ProteinCounts`` as a dict.
+    """
+    fetch = genome if callable(genome) else genome.__getitem__
+    if chromosomes is None:
+        chromosomes = sorted(
+            {g.chrom for g in ref_genes.values()}
+            | {g.chrom for g in pred_genes.values()}
+        )
+    aligner = PairwiseAligner(scoring="blastp")
+    aligner.mode = "global"
+    counts = ProteinCounts()
+
+    for chrom in chromosomes:
+        chrom_ref = {
+            g: x for g, x in ref_genes.items() if x.chrom == chrom and x.cds_keys
+        }
+        chrom_pred = {
+            g: x for g, x in pred_genes.items() if x.chrom == chrom and x.cds_keys
+        }
+        if not chrom_ref and not chrom_pred:
+            continue
+        sequence = str(fetch(chrom)).upper()
+        ref_proteins = {g: translate_gene(x, sequence) for g, x in chrom_ref.items()}
+        pred_proteins = {g: translate_gene(x, sequence) for g, x in chrom_pred.items()}
+
+        for protein in pred_proteins.values():
+            if len(protein) < MIN_PROTEIN_LENGTH:
+                continue
+            starts, clean = protein.startswith("M"), "*" not in protein
+            counts.n_predicted += 1
+            counts.starts_with_methionine += starts
+            counts.no_internal_stop += clean
+            counts.well_formed += starts and clean
+
+        index = CdsIndex(chrom_pred)
+        for gid, gene in chrom_ref.items():
+            reference = ref_proteins[gid]
+            if not reference:
+                continue
+            counts.n_reference += 1
+            candidates = [
+                pred_proteins[other]
+                for other in index.overlapping_gene(gene)
+                if pred_proteins.get(other)
+            ]
+            if not candidates:
+                counts.no_counterpart += 1
+                counts.identity_histogram["<50"] += 1
+                continue
+            best = max(identity(reference, c, aligner) for c in candidates)
+            for threshold in IDENTITY_THRESHOLDS:
+                counts.recovered[f"{threshold:g}"] += best >= threshold
+            counts.identity_histogram[_identity_bin(best)] += 1
+    return dataclasses.asdict(counts)

tabs/demo.html

@@ -1,33 +0,0 @@
-<div class="grid">
-  <div class="card" style="grid-column: span 12; text-align: center; padding: 60px 40px;">
-    <h2 style="margin-bottom: 20px;">🚀 NTv3 Live Demo</h2>
-    <p style="margin-bottom: 30px; color: var(--muted); line-height: 1.7;">
-      Try out the NTv3 model interactively in this live demo. Predict functional tracks and genome annotation directly in your browser. Click the button below to open the interactive demo in a new window.
-    </p>
-    <a 
-      href="https://huggingface.co/spaces/InstaDeepAI/ntv3_tracks" 
-      target="_blank" 
-      rel="noopener noreferrer"
-      style="
-        display: inline-block;
-        padding: 16px 32px;
-        background: var(--link);
-        color: white;
-        text-decoration: none;
-        border-radius: 12px;
-        font-weight: 600;
-        font-size: 16px;
-        transition: all 0.2s ease;
-        box-shadow: 0 4px 12px rgba(125, 211, 252, 0.3);
-      "
-      onmouseover="this.style.background='#60b8e8'; this.style.transform='translateY(-2px)'; this.style.boxShadow='0 6px 16px rgba(125, 211, 252, 0.4)';"
-      onmouseout="this.style.background='var(--link)'; this.style.transform='translateY(0)'; this.style.boxShadow='0 4px 12px rgba(125, 211, 252, 0.3)';"
-    >
-      🚀 Open Live Demo →
-    </a>
-    <p style="margin-top: 30px; font-size: 13px; color: var(--muted);">
-      The demo will open in a new tab. Make sure pop-ups are enabled for this site.
-    </p>
-  </div>
-</div>
-

tabs/home.html

@@ -1,78 +1,101 @@
+<div class="legacy-banner">
+  🗄️ Looking for the 2025 NTv3 models? They are listed under the
+  <a href="#" data-tab-link="initial">Initial release (2025)</a> tab, with their notebooks.
+</div>
+
 <div class="summary">
   <h2>📖 About NTv3</h2>
   <p>
-    NTv3 is a multi-species genomic foundation model family that unifies representation learning, functional-track prediction, genome annotation, and controllable sequence generation within a single U-Net-style backbone. It models up to 1 Mb of DNA at single-base resolution, using a conv–Transformer–deconv architecture that efficiently captures both local motifs and long-range regulatory dependencies. NTv3 is first pretrained on ~9T base pairs from the OpenGenome2 corpus spanning >128k species using masked language modeling, and then post-trained with a joint objective on ~16k functional tracks and annotation labels across 24 animal and plant species, enabling state-of-the-art cross-species functional prediction and base-resolution genome annotation.
+    NTv3 is a genomic foundation model that reads DNA at single-base resolution across up to <strong>1 Mb</strong> of
+    context in one pass. One architecture supports three capabilities: <strong>functional-track prediction</strong>,
+    <strong>gene annotation</strong> and <strong>regulatory sequence design</strong>. The backbone is U-Net-style: a
+    convolutional tower that reads the sequence, a Transformer tower that carries long-range context, and a
+    deconvolutional tower that returns predictions at base resolution.
+  </p>
+  <p>
+    The three capabilities are acquired in stages, each building on the one before. NTv3 is first
+    <strong>pretrained</strong> by masked-token prediction on unlabeled genomes from the OpenGenome2 corpus. It is then
+    <strong>post-trained on functional tracks</strong>, a supervised multi-task stage over 16,262 tracks from 15 species,
+    which gives the first capability. The other two are built from those same functional representations:
+    <strong>annotation fine-tuning</strong> on reference gene annotations, and <strong>generative post-training</strong>
+    by masked diffusion on MPRA activity datasets. The checkpoints below are the pretrained backbone, tracks and
+    annotation, plus the plant promoter model of the paper as an example of the generative ones.
   </p>
   <p>
-    NTv3 also acts as a controllable generative model via masked-diffusion language modeling, allowing targeted design of regulatory sequences (for example, enhancers with specified activity and promoter selectivity) that have been validated experimentally.
+    Because it is one backbone and one set of learned representations throughout, what NTv3 learns about sequence
+    function carries across tasks and across species rather than being relearned for each.
   </p>
 </div>
 
 <div class="paper-summary">
-  <!-- <h2>📄 A foundational model for joint sequence-function multi-species modeling at scale for long-range genomic prediction</h2> -->
-  <img src="assets/paper_summary.png" alt="NTv3 Paper Summary" />
+  <img src="assets/ntv3_overview.png" alt="NTv3 overview: pre-training, tracks post-training, annotation fine-tuning and generative post-training, with the four resulting capabilities" />
 </div>
 
 <div class="why-ntv3">
-  <h2>✨ Why NTv3?</h2>
+  <h2>✨ What NTv3 does</h2>
   <ul>
-    <li>📏 <strong>1 Mb long context at nucleotide resolution</strong> — ~100× longer than typical genomics models.</li>
-    <li>🏗️ <strong>Unified architecture</strong> for: masked language modeling, functional-track prediction, genome annotation, and sequence generation.</li>
-    <li>🌍 <strong>Cross-species generalization</strong> across 24 animals + plants with a shared conditioned representation space.</li>
-    <li>⚡ <strong>U-Net–style architecture</strong> improves stability and GPU efficiency on very long sequences.</li>
-    <li>🎯 <strong>Controllable generative modeling</strong>, enabling targeted enhancer/promoter engineering validated by experimental assays.</li>
+    <li>📏 <strong>Reads any genome.</strong> 1 Mb of context at nucleotide resolution, with a U-Net design that keeps very long sequences tractable on GPU.</li>
+    <li>📊 <strong>Predicts functional tracks.</strong> 16,262 tracks across 15 species: 9,603 gene expression, 3,115 histone marks, 2,193 TF binding, 1,351 accessibility.</li>
+    <li>🏷️ <strong>Annotates genes from DNA alone.</strong> Per-base posteriors over intergenic, intron, CDS and splice states, decoded into gene models on both strands. One model per clade: mammals, plants and insects.</li>
+    <li>✏️ <strong>Designs regulatory DNA.</strong> Conditional generation of plant promoters, plant enhancers and fly enhancers at requested activity levels, with activity measured experimentally.</li>
+    <li>🌍 <strong>Generalizes across species</strong>, including to genomes held out of post-training entirely.</li>
   </ul>
 </div>
 
 <div class="grid">
-  <div class="card">
-    <h2>🤖 Models (see <a href="https://huggingface.co/collections/InstaDeepAI/nucleotide-transformer-v3" target="_blank" rel="noopener noreferrer">collection</a>)</h2>
-    <ul>
-      <li>📦 Pretrained checkpoints:
-        <div style="margin-top: 8px; margin-left: 0;">
-          <div><a href="https://huggingface.co/InstaDeepAI/NTv3_8M_pre" target="_blank" rel="noopener noreferrer"><code>InstaDeepAI/NTv3_8M_pre</code></a></div>
-          <div><a href="https://huggingface.co/InstaDeepAI/NTv3_100M_pre" target="_blank" rel="noopener noreferrer"><code>InstaDeepAI/NTv3_100M_pre</code></a></div>
-          <div><a href="https://huggingface.co/InstaDeepAI/NTv3_650M_pre" target="_blank" rel="noopener noreferrer"><code>InstaDeepAI/NTv3_650M_pre</code></a></div>
-        </div>
-      </li>
-      <li>🎯 Post-trained checkpoints:
-        <div style="margin-top: 8px; margin-left: 0;">
-          <div><a href="https://huggingface.co/InstaDeepAI/NTv3_100M_pos" target="_blank" rel="noopener noreferrer"><code>InstaDeepAI/NTv3_100M_pos</code></a></div>
-          <div><a href="https://huggingface.co/InstaDeepAI/NTv3_650M_pos" target="_blank" rel="noopener noreferrer"><code>InstaDeepAI/NTv3_650M_pos</code></a></div>
-        </div>
-      </li>
-    </ul>
-    <table>
+  <div class="card" style="grid-column: span 12;">
+    <h2>🤖 The NTv3 models</h2>
+    <p>
+      NTv3 is a family rather than a single model. One <strong>700M-parameter</strong> backbone is trained once on
+      unlabeled genomes, then specialised into a checkpoint per capability, so each one inherits what the others
+      learned about sequence rather than starting over.
+    </p>
+    <p>
+      Start from <strong>NTv3-pretrained</strong> for embeddings, masked-token prediction, or your own fine-tuning on
+      a task we have not covered. Reach for <strong>NTv3-tracks</strong> to predict functional signal, chromatin
+      accessibility, histone marks, transcription factor binding and expression, anywhere in a supported genome. Use
+      <strong>NTv3-annotate</strong> to call genes from sequence alone, choosing the model for your clade.
+    </p>
+    <table class="capability-grid">
       <thead>
         <tr>
           <th>Model</th>
-          <th>Size</th>
-          <th>Pre-training</th>
-          <th>Post-training</th>
-          <th>Usage</th>
+          <th>Stage</th>
+          <th>What it is for</th>
+          <th>Repository</th>
         </tr>
       </thead>
       <tbody>
         <tr>
-          <td><strong>NTv3-8M</strong></td>
-          <td>8M params</td>
-          <td><span class="checkmark">✅</span></td>
-          <td>❌</td>
-          <td>Embeddings, light inference</td>
+          <td><strong>NTv3-pretrained</strong></td>
+          <td>Pretrained</td>
+          <td>Embeddings, masked-token prediction, starting point for your own fine-tuning</td>
+          <td><a href="https://huggingface.co/InstaDeepAI/NTv3_pretrained" target="_blank" rel="noopener noreferrer"><code>InstaDeepAI/NTv3_pretrained</code></a></td>
         </tr>
         <tr>
-          <td><strong>NTv3-100M</strong></td>
-          <td>100M params</td>
-          <td><span class="checkmark">✅</span></td>
-          <td><span class="checkmark">✅</span></td>
-          <td>Embeddings, tracks, annotation</td>
+          <td><strong>NTv3-tracks</strong></td>
+          <td>Tracks post-training</td>
+          <td>Functional-track prediction, and the backbone for downstream fine-tuning. <code>NTv3_tracks_1mb</code> is the 1 Mb specialist; 1 Mb inputs only</td>
+          <td>
+            <a href="https://huggingface.co/InstaDeepAI/NTv3_tracks" target="_blank" rel="noopener noreferrer"><code>InstaDeepAI/NTv3_tracks</code></a><br />
+            <a href="https://huggingface.co/InstaDeepAI/NTv3_tracks_1mb" target="_blank" rel="noopener noreferrer"><code>InstaDeepAI/NTv3_tracks_1mb</code></a>
+          </td>
         </tr>
         <tr>
-          <td><strong>NTv3-650M</strong></td>
-          <td>650M params</td>
-          <td><span class="checkmark">✅</span></td>
-          <td><span class="checkmark">✅</span></td>
-          <td>Embeddings, tracks, annotation, best accuracy</td>
+          <td><strong>NTv3-annotate</strong></td>
+          <td>Annotation fine-tuning</td>
+          <td><em>Ab initio</em> gene prediction, one model per clade</td>
+          <td>
+            <a href="https://huggingface.co/InstaDeepAI/NTv3_annot_mammals" target="_blank" rel="noopener noreferrer"><code>InstaDeepAI/NTv3_annot_mammals</code></a><br />
+            <a href="https://huggingface.co/InstaDeepAI/NTv3_annot_plants" target="_blank" rel="noopener noreferrer"><code>InstaDeepAI/NTv3_annot_plants</code></a><br />
+            <a href="https://huggingface.co/InstaDeepAI/NTv3_annot_insects" target="_blank" rel="noopener noreferrer"><code>InstaDeepAI/NTv3_annot_insects</code></a>
+          </td>
+        </tr>
+        <tr>
+          <td><strong>NTv3-generative-plant-promoter</strong></td>
+          <td>Generative post-training</td>
+          <td>Example generative model used in the paper: designs plant core promoters at a requested strength, and is what notebook 06 runs</td>
+          <td><a href="https://huggingface.co/InstaDeepAI/NTv3_generative_plant_promoter" target="_blank" rel="noopener noreferrer"><code>InstaDeepAI/NTv3_generative_plant_promoter</code></a></td>
         </tr>
       </tbody>
     </table>
@@ -80,152 +103,147 @@
 
   <div class="card-stack">
     <div class="card">
-      <h2>📓 Tutorial notebooks (browse <a href="https://huggingface.co/spaces/InstaDeepAI/ntv3/tree/main/notebooks_tutorials" target="_blank" rel="noopener noreferrer">folder</a>)</h2>
+      <h2>📓 Notebooks</h2>
+      <p>
+        The notebooks cover the workflows end to end. All run against the published checkpoints and are linked.
+      </p>
       <ul>
-        <li><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_tutorials/00_quickstart_inference.ipynb" target="_blank" rel="noopener noreferrer">🚀 00 — Quickstart inference</a></li>
-        <li><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_tutorials/01_tracks_prediction.ipynb" target="_blank" rel="noopener noreferrer">📊 01 — Tracks prediction</a></li>
-        <li><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_tutorials/02_fine_tuning_pretrained_model_biwig.ipynb" target="_blank" rel="noopener noreferrer">🎯 02 — Fine-tune a pre-trained model on bigwig tracks</a></li>
-        <li><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_tutorials/03_fine_tuning_posttrained_model_biwig.ipynb" target="_blank" rel="noopener noreferrer">🎯 03 — Fine-tune a post-trained model on bigwig tracks</a></li>
-        <li><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_tutorials/04_fine_tuning_pretrained_model_annotation.ipynb" target="_blank" rel="noopener noreferrer">🏷️ 04 — Fine-tune a pre-trained model on annotations</a></li>
-        <li><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_tutorials/05_model_interpretation.ipynb" target="_blank" rel="noopener noreferrer">🔍 05 — Model interpretation</a></li>
-        <li><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_tutorials/06_NTv3_generative_training.ipynb" target="_blank" rel="noopener noreferrer">🧪 06 — Fine-tuning NTv3 into a diffusion model</a></li>
-        <li><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_tutorials/07_enhancer_generation.ipynb" target="_blank" rel="noopener noreferrer">🪰 07 — Generating enhancer sequences</a></li>
-      </ul>
-    </div>
-    <div class="card">
-      <h2>📓 Pipeline notebooks (browse <a href="https://huggingface.co/spaces/InstaDeepAI/ntv3/tree/main/notebooks_pipelines" target="_blank" rel="noopener noreferrer">folder</a>)</h2>
-      <ul>
-        <li><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_pipelines/01_functional_track_prediction.ipynb" target="_blank" rel="noopener noreferrer">🎯 01 — Generate bigwig predictions for certain tracks</a></li>
-        <li><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_pipelines/02_model_interpretation.ipynb" target="_blank" rel="noopener noreferrer">🔍 02 — Interpret a given genomic region</a></li>
+        <li>🚀 <strong>00 · Quickstart inference</strong><br /><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_tutorials/00_quickstart_inference.ipynb" target="_blank" rel="noopener noreferrer"><code>00_quickstart_inference.ipynb</code></a></li>
+        <li>📊 <strong>01 · Tracks prediction</strong><br /><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_tutorials/01_tracks_prediction.ipynb" target="_blank" rel="noopener noreferrer"><code>01_tracks_prediction.ipynb</code></a></li>
+        <li>🎯 <strong>02 · Fine-tune NTv3 on bigwig tracks</strong><br /><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_tutorials/02_fine_tuning_bigwig_tracks.ipynb" target="_blank" rel="noopener noreferrer"><code>02_fine_tuning_bigwig_tracks.ipynb</code></a></li>
+        <li>🏷️ <strong>03 · Genome annotation with NTv3</strong>
+          <ul>
+            <li><strong>03a · Train a gene-annotation model</strong><br /><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_tutorials/03a_genome_annotation_training.ipynb" target="_blank" rel="noopener noreferrer"><code>03a_genome_annotation_training.ipynb</code></a></li>
+            <li><strong>03b · From per-base predictions to gene models</strong><br /><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_tutorials/03b_genome_annotation_decoding.ipynb" target="_blank" rel="noopener noreferrer"><code>03b_genome_annotation_decoding.ipynb</code></a></li>
+            <li><strong>03c · Evaluate a gene annotation against a reference</strong><br /><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_tutorials/03c_genome_annotation_evaluation.ipynb" target="_blank" rel="noopener noreferrer"><code>03c_genome_annotation_evaluation.ipynb</code></a></li>
+          </ul>
+        </li>
+        <li>🔍 <strong>04 · Model interpretation</strong><br /><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_tutorials/04_model_interpretation.ipynb" target="_blank" rel="noopener noreferrer"><code>04_model_interpretation.ipynb</code></a></li>
+        <li>🧪 <strong>05 · Fine-tuning NTv3 into a diffusion model</strong><br /><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_tutorials/05_diffusion_fine_tuning.ipynb" target="_blank" rel="noopener noreferrer"><code>05_diffusion_fine_tuning.ipynb</code></a></li>
+        <li>🌱 <strong>06 · Generating plant promoter sequences</strong><br /><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/main/notebooks_tutorials/06_plant_promoter_generation.ipynb" target="_blank" rel="noopener noreferrer"><code>06_plant_promoter_generation.ipynb</code></a></li>
       </ul>
     </div>
     <div class="card">
       <h2>🔗 Links</h2>
       <ul>
-        <li><a href="https://www.biorxiv.org/content/10.64898/2025.12.22.695963v1" target="_blank" rel="noopener noreferrer">📄 Paper</a></li>
+        <li>📄 Paper <span class="todo">link TBC</span></li>
         <li><a href="https://github.com/instadeepai/nucleotide-transformer" target="_blank" rel="noopener noreferrer">💻 JAX model code (GitHub)</a></li>
-        <li><a href="https://huggingface.co/collections/InstaDeepAI/nucleotide-transformer-v3" target="_blank" rel="noopener noreferrer">🎯 HF Model Collection (all NTv3 models)</a></li>
-        <li><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3/tree/main/notebooks_tutorials" target="_blank" rel="noopener noreferrer">📚 Tutorial </a> and <a href="https://huggingface.co/spaces/InstaDeepAI/ntv3/tree/main/notebooks_pipelines" target="_blank" rel="noopener noreferrer">🔧 Pipeline</a> notebooks</li>
-        <li><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3_benchmark" target="_blank" rel="noopener noreferrer">🏆 NTv3 benchmark leaderboard</a></li>
+        <li><a href="https://huggingface.co/collections/InstaDeepAI/ntv3" target="_blank" rel="noopener noreferrer">🎯 HF model collection</a></li>
+        <li><a href="https://huggingface.co/datasets/InstaDeepAI/NTv3_benchmark_dataset_revised" target="_blank" rel="noopener noreferrer">📚 NTv3 benchmark dataset</a></li>
+        <li><a href="https://huggingface.co/datasets/InstaDeepAI/NTv3_annot_dataset" target="_blank" rel="noopener noreferrer">🏷️ NTv3 annotation dataset</a></li>
+        <li><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3_benchmark_revised" target="_blank" rel="noopener noreferrer">🏆 NTv3 benchmark leaderboard</a></li>
       </ul>
     </div>
   </div>
 
   <div class="card-stack">
     <div class="card">
-      <h2>🤖 Load a pre-trained model</h2>
-      <p>Here is an example of how to load and use a pre-trained NTv3 model.</p>
-      <div class="code"><pre><code class="language-python">from transformers import AutoTokenizer, AutoModelForMaskedLM
+      <h2>🤖 NTv3-pretrained</h2>
+      <p>Embeddings and masked-token prediction, the starting point for your own fine-tuning.</p>
+      <div class="code"><pre><code class="language-python">import torch
+from transformers import AutoModel, AutoTokenizer
 
-model_name = "InstaDeepAI/NTv3_650M_pre"
+model_name = "InstaDeepAI/NTv3_pretrained"
 
-# Load model and tokenizer
-model = AutoModelForMaskedLM.from_pretrained(model_name, trust_remote_code=True)
+model = AutoModel.from_pretrained(model_name, trust_remote_code=True).eval()
 tok = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
 
-# Tokenize input sequences
-batch = tok(["ATCGNATCG", "ACGT"], add_special_tokens=False, padding=True, pad_to_multiple_of=128, return_tensors="pt")
+# Input length must be a multiple of 128. Batch only equal-length sequences:
+# padding is NOT masked out, so a padded sequence gets different values than
+# the same sequence run on its own.
+batch = tok(["ACGT" * 32, "TTGA" * 32], add_special_tokens=False, return_tensors="pt")
 
-# Run model
-out = model(**batch)
+# The forward takes the token tensor positionally, not the tokenizer dict
+with torch.no_grad():
+    out = model(batch["input_ids"])
 
-# Print output shapes
-print(out.logits.shape)       # (B, L, V = 11)
+print(out["logits"].shape)      # (B, L, V = 11)
+print(out["embedding"].shape)   # (B, L, 1024)
 </code></pre></div>
-      <p>Model embeddings can be used for fine-tuning on downstream tasks.</p>
     </div>
     <div class="card">
-      <h2>🔍 Model interpretation</h2>
-      <p>Here is an example of how to use the interpretation pipeline on the NTv3 post-trained model for multi-scale analysis of DNA sequences:</p>
-      <div class="code"><pre><code class="language-python">from transformers import pipeline
-import torch
-import matplotlib.pyplot as plt
-
-model_name = "InstaDeepAI/NTv3_650M_post"
-
-# Build interpretation pipeline
-ntv3_interpret = pipeline(
-    "ntv3-interpret",
-    model=model_name,
-    trust_remote_code=True,
-    device=0 if torch.cuda.is_available() else -1,
-)
-
-# Run interpretation on a given genomic region with tracks, annotations, attention, and saliency
-result = ntv3_interpret(
-    {"chrom": "chr11", "start": 5_253_561, "end": 5_286_329, "species": "human"},
-    output_attention=True,
-    output_saliency=True,
-    saliency_track_id="ENCSR000EFT",  # K562 GATA1 ChIP-seq
-    plot=True,  # plot predictons on tracks and annotations
-    tracks_to_plot={"K562 RNA-seq": "ENCSR056HPM", "K562 GATA1": "ENCSR000EFT"},
-    elements_to_plot=["exon", "promoter_Tissue_specific"],
-)
-
-# Access attention map results
-result.plot_attention()  # attention map (last layer)
-plt.show()
-
-# Access saliency scores results
-result.plot_saliency(window_size=128)
-plt.show()
-</code></pre></div>
-      <img src="assets/saliency_example.png" alt="Output tracks visualization" style="max-width: 100%; margin-top: 20px;" />
+      <h2>🔍 NTv3-tracks: model interpretation</h2>
+      <p>Multi-scale analysis of a genomic region: predicted tracks, attention and saliency.</p>
+      <img src="assets/saliency_example.png" alt="Saliency example" style="max-width: 100%; margin-top: 20px;" />
     </div>
   </div>
-  
-  <div class="card">
-    <h2>💻 Use a post-trained model</h2>
-    <p>Here is a quick example of how to use the post-trained NTv3 650M model to predict tracks for a human genomic window.</p>
-    <div class="code"><pre><code class="language-python">from transformers import pipeline
-import torch
 
-model_name = "InstaDeepAI/NTv3_650M_pos"
+  <div class="card" style="grid-column: span 12;">
+    <h2>📊 NTv3-tracks: predict functional tracks</h2>
+    <p>
+      Give the pipeline a genomic region and the tracks you want, and it fetches the sequence, conditions on the
+      species, predicts and plots.
+    </p>
+    <div class="code"><pre><code class="language-python">from transformers import pipeline
 
-ntv3_tracks = pipeline(
-    "ntv3-tracks",
-    model=model_name,
-    trust_remote_code=True,
-    device=0 if torch.cuda.is_available() else -1,
-)
+ntv3_tracks = pipeline("ntv3-tracks", model="InstaDeepAI/NTv3_tracks", trust_remote_code=True)
 
-# Run track prediction
 out = ntv3_tracks(
-  {
-    "chrom": "chr19",
-    "start": 6_700_000,
-    "end": 6_831_072,
-    "species": "human"
-  }
+    {"chrom": "chr19", "start": 6_700_000, "end": 6_732_768, "species": "human"},
+    plot=True,
+    tracks_to_plot={
+        "K562 RNA-seq": "ENCSR056HPM",
+        "K562 DNAse": "ENCSR921NMD",
+        "K562 H3k4me3": "ENCSR000DWD",
+        "K562 CTCF": "ENCSR000AKO",
+    },
 )
 
-# Print output shapes
-# 7k human tracks over 37.5 % center region of the input sequence
-print("bigwig_tracks_logits:", tuple(out.bigwig_tracks_logits.shape))
-# Location of 21 genomic elements over 37.5 % center region of the input sequence
-print("bed_tracks_logits:", tuple(out.bed_tracks_logits.shape))
-# Language model logits for whole sequence over vocabulary
-print("language model logits:", tuple(out.mlm_logits.shape))</code></pre></div>
-    <p>Predictions can also be plotted for a subset of functional tracks and genomic elements:</p>
-    <div class="code"><pre><code class="language-python">tracks_to_plot = {
-    "K562 RNA-seq": "ENCSR056HPM",
-    "K562 DNAse": "ENCSR921NMD",
-    "K562 H3k4me3": "ENCSR000DWD",
-    "K562 CTCF": "ENCSR000AKO",
-    "HepG2 RNA-seq": "ENCSR561FEE_P",
-    "HepG2 DNAse": "ENCSR000EJV",
-    "HepG2 H3k4me3": "ENCSR000AMP",
-    "HepG2 CTCF": "ENCSR000BIE",
-}
-elements_to_plot = ["protein_coding_gene", "exon", "intron", "splice_donor", "splice_acceptor"]
+print(out.tracks.shape)        # (12288, 4): the central 37.5 % of the window
+print(out.coordinates)         # (6710240, 6722528), the span actually predicted
+print(out.track("ENCSR000AKO").shape)  # (12288,): one track</code></pre></div>
+    <img src="assets/output_tracks.png" alt="Predicted tracks over a human genomic window" style="max-width: 100%; margin-top: 20px;" />
+  </div>
 
-out = ntv3_tracks(
-    {"chrom": "chr19", "start": 6_700_000, "end": 6_831_072, "species": "human"},
-    plot=True,
-    tracks_to_plot=tracks_to_plot,
-    elements_to_plot=elements_to_plot,
-)</code></pre></div>
-    <img src="assets/output_tracks.png" alt="Output tracks visualization" style="max-width: 100%; margin-top: 20px;" />
+  <div class="card" style="grid-column: span 12;">
+    <h2>🏷️ NTv3-annotate: end-to-end FASTA annotation pipeline</h2>
+    <p>
+      Turns a genome FASTA into its predicted gene annotation (GFF3): the production pipeline tiles every chromosome
+      into 1 Mb windows, runs the clade model and decodes the gene models. Here is a single-chromosome example on a
+      test species.
+    </p>
+    <div class="code"><pre><code class="language-bash"># 1. Get the code
+git clone git@github.com:instadeepai/nucleotide-transformer-private.git
+cd nucleotide-transformer-private
+
+# 2. Install with GPU support, then check that JAX sees the GPU
+python3 -m venv .venv && source .venv/bin/activate
+pip install '.[gpu]'
+python -c "import jax; print(jax.devices())"   # should list a CudaDevice
+
+# 3. Get chromosome 31 of the monarch butterfly (2.8 Mb, soft-masked)
+NCBI=https://ftp.ncbi.nlm.nih.gov/genomes/all/GCF/018/135/715/GCF_018135715.1_MEX_DaPlex
+curl -sSfLO $NCBI/GCF_018135715.1_MEX_DaPlex_genomic.fna.gz
+gunzip GCF_018135715.1_MEX_DaPlex_genomic.fna.gz
+faidx GCF_018135715.1_MEX_DaPlex_genomic.fna NC_083558.1 > danaus_chr31.fa
+
+# 4. Annotate it with the insect model (weights download on the first run)
+export HF_TOKEN=hf_...   # needed while the model repos are private
+ntv3-annotate --genome danaus_chr31.fa --model_cfg insects --out danaus_chr31.gff3
+
+# 5. Look at one predicted gene: a gene line, then its exons (CDS) and introns
+awk '$4 >= 2334916 && $5 <= 2337940' danaus_chr31.gff3 | cut -f 1,3-5,7
+# NC_083558.1  gene    2334916  2337940  -
+# NC_083558.1  CDS     2334916  2334969  -
+# NC_083558.1  intron  2334970  2335582  -
+# NC_083558.1  CDS     2335583  2335753  -
+# NC_083558.1  intron  2335754  2336418  -
+# NC_083558.1  CDS     2336419  2336565  -
+# NC_083558.1  intron  2336566  2336820  -
+# NC_083558.1  CDS     2336821  2336932  -
+# NC_083558.1  intron  2336933  2337785  -
+# NC_083558.1  CDS     2337786  2337940  -</code></pre></div>
+    <img src="assets/output_annotation.png" alt="Predicted gene models on monarch butterfly chromosome 31 against the RefSeq reference" style="max-width: 100%; margin-top: 20px;" />
+    <p>
+      The predicted genes (blue) over 34 kb of chromosome 31, against the RefSeq reference (grey). Four genes match the
+      reference exon for exon, including the one printed above; the rightmost reference gene is missed.
+    </p>
+    <div class="notice">
+      Needs an <strong>A100-40GB or larger</strong> GPU. For the other clades, use
+      <code style="white-space: nowrap;">--model_cfg mammals</code> or
+      <code style="white-space: nowrap;">--model_cfg plants</code> with a genome of that clade. Notebook
+      <strong>03b</strong> explains how the gene models are decoded, and notebook <strong>03c</strong> scores a GFF3
+      against the benchmark reference.
+    </div>
   </div>
 </div>
-

tabs/initial_release.html

@@ -0,0 +1,263 @@
+<div class="notice">
+  🗄️ <strong>This page describes the initial NTv3 release, from 2025.</strong> Those repositories are retained unchanged for
+  reproducibility and backward compatibility. Existing model IDs, files and interfaces continue to work exactly as they
+  did, and nothing on this page is being withdrawn. If you are starting something new, prefer the current models on the
+  <strong>Home</strong> tab; if you are reproducing published results or already have a working pipeline, stay here.
+</div>
+
+<div class="summary">
+  <h2>📖 About the initial release (2025)</h2>
+  <p>
+    The initial release described NTv3 as a single U-Net-style backbone pretrained on the OpenGenome2 corpus and then
+    post-trained with a <em>joint</em> objective over functional tracks and annotation labels. It shipped three model sizes
+    (8M, 100M, 650M), each in pretrained and post-trained form, plus context-length and architecture ablation variants.
+  </p>
+  <p>
+    The current release changes that structure: tracks post-training is its own stage, and annotation becomes a separate
+    per-clade fine-tune rather than a head trained jointly. Numbers quoted in these materials therefore do not always
+    match the current paper. Where the two differ, the current paper is authoritative.
+  </p>
+</div>
+
+<div class="paper-summary">
+  <img src="assets/paper_summary.png" alt="NTv3 initial-release overview figure" />
+</div>
+
+<div class="grid">
+  <div class="card" style="grid-column: span 12;">
+    <h2>🤖 Initial-release checkpoints (2025)</h2>
+    <p>
+      All of the following are public on the Hub and unchanged. Sizes are the parameter counts as published.
+    </p>
+    <table>
+      <thead>
+        <tr>
+          <th>Model</th>
+          <th>Size</th>
+          <th>Pretrained</th>
+          <th>Post-trained</th>
+          <th>Usage</th>
+        </tr>
+      </thead>
+      <tbody>
+        <tr>
+          <td><strong>NTv3-8M</strong></td>
+          <td>8M params</td>
+          <td><span class="checkmark">✅</span></td>
+          <td>❌</td>
+          <td>Embeddings, light inference</td>
+        </tr>
+        <tr>
+          <td><strong>NTv3-100M</strong></td>
+          <td>100M params</td>
+          <td><span class="checkmark">✅</span></td>
+          <td><span class="checkmark">✅</span></td>
+          <td>Embeddings, tracks, annotation</td>
+        </tr>
+        <tr>
+          <td><strong>NTv3-650M</strong></td>
+          <td>650M params</td>
+          <td><span class="checkmark">✅</span></td>
+          <td><span class="checkmark">✅</span></td>
+          <td>Embeddings, tracks, annotation, best accuracy</td>
+        </tr>
+      </tbody>
+    </table>
+
+    <h2 style="margin-top: 22px;">📦 Pretrained</h2>
+    <ul>
+      <li><a href="https://huggingface.co/InstaDeepAI/NTv3_8M_pre" target="_blank" rel="noopener noreferrer"><code>InstaDeepAI/NTv3_8M_pre</code></a></li>
+      <li><a href="https://huggingface.co/InstaDeepAI/NTv3_100M_pre" target="_blank" rel="noopener noreferrer"><code>InstaDeepAI/NTv3_100M_pre</code></a></li>
+      <li><a href="https://huggingface.co/InstaDeepAI/NTv3_650M_pre" target="_blank" rel="noopener noreferrer"><code>InstaDeepAI/NTv3_650M_pre</code></a></li>
+    </ul>
+
+    <h2 style="margin-top: 18px;">🎯 Post-trained</h2>
+    <ul>
+      <li><a href="https://huggingface.co/InstaDeepAI/NTv3_100M_post" target="_blank" rel="noopener noreferrer"><code>InstaDeepAI/NTv3_100M_post</code></a></li>
+      <li><a href="https://huggingface.co/InstaDeepAI/NTv3_650M_post" target="_blank" rel="noopener noreferrer"><code>InstaDeepAI/NTv3_650M_post</code></a></li>
+    </ul>
+
+    <h2 style="margin-top: 18px;">📏 Context-length variants</h2>
+    <ul>
+      <li><a href="https://huggingface.co/InstaDeepAI/NTv3_8M_pre_8kb" target="_blank" rel="noopener noreferrer"><code>InstaDeepAI/NTv3_8M_pre_8kb</code></a></li>
+      <li><a href="https://huggingface.co/InstaDeepAI/NTv3_100M_pre_8kb" target="_blank" rel="noopener noreferrer"><code>InstaDeepAI/NTv3_100M_pre_8kb</code></a></li>
+      <li><a href="https://huggingface.co/InstaDeepAI/NTv3_650M_pre_8kb" target="_blank" rel="noopener noreferrer"><code>InstaDeepAI/NTv3_650M_pre_8kb</code></a></li>
+      <li><a href="https://huggingface.co/InstaDeepAI/NTv3_100M_post_131kb" target="_blank" rel="noopener noreferrer"><code>InstaDeepAI/NTv3_100M_post_131kb</code></a></li>
+      <li><a href="https://huggingface.co/InstaDeepAI/NTv3_650M_post_131kb" target="_blank" rel="noopener noreferrer"><code>InstaDeepAI/NTv3_650M_post_131kb</code></a></li>
+    </ul>
+
+    <h2 style="margin-top: 18px;">🔬 Architecture ablations (five-downsample variants)</h2>
+    <ul>
+      <li><a href="https://huggingface.co/InstaDeepAI/NTv3_5downsample_pre" target="_blank" rel="noopener noreferrer"><code>InstaDeepAI/NTv3_5downsample_pre</code></a></li>
+      <li><a href="https://huggingface.co/InstaDeepAI/NTv3_5downsample_pre_8kb" target="_blank" rel="noopener noreferrer"><code>InstaDeepAI/NTv3_5downsample_pre_8kb</code></a></li>
+      <li><a href="https://huggingface.co/InstaDeepAI/NTv3_5downsample_post" target="_blank" rel="noopener noreferrer"><code>InstaDeepAI/NTv3_5downsample_post</code></a></li>
+      <li><a href="https://huggingface.co/InstaDeepAI/NTv3_5downsample_post_131kb" target="_blank" rel="noopener noreferrer"><code>InstaDeepAI/NTv3_5downsample_post_131kb</code></a></li>
+    </ul>
+
+    <h2 style="margin-top: 18px;">🧪 Generative</h2>
+    <ul>
+      <li><a href="https://huggingface.co/InstaDeepAI/NTv3_generative" target="_blank" rel="noopener noreferrer"><code>InstaDeepAI/NTv3_generative</code></a></li>
+    </ul>
+
+    <div class="notice">
+      <strong>Naming note.</strong> Some earlier materials referred to the post-trained checkpoints as
+      <code>NTv3_100M_pos</code> and <code>NTv3_650M_pos</code>. The correct identifiers end in <code>_post</code>. If a
+      script of yours 404s, that is the likely cause.
+    </div>
+  </div>
+
+  <div class="card-stack">
+    <div class="card">
+      <h2>📓 Notebooks for the initial release (2025)</h2>
+      <p>These notebooks were written and run against the initial-release checkpoints.</p>
+      <ul>
+        <li><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/initial-release/notebooks_tutorials/00_quickstart_inference.ipynb" target="_blank" rel="noopener noreferrer">🚀 00 · Quickstart inference</a></li>
+        <li><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/initial-release/notebooks_tutorials/01_tracks_prediction.ipynb" target="_blank" rel="noopener noreferrer">📊 01 · Tracks prediction</a></li>
+        <li><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/initial-release/notebooks_tutorials/02_fine_tuning_pretrained_model_biwig.ipynb" target="_blank" rel="noopener noreferrer">🎯 02 · Fine-tune a pretrained model on bigwig tracks</a></li>
+        <li><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/initial-release/notebooks_tutorials/03_fine_tuning_posttrained_model_biwig.ipynb" target="_blank" rel="noopener noreferrer">🎯 03 · Fine-tune a post-trained model on bigwig tracks</a></li>
+        <li><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/initial-release/notebooks_tutorials/04_fine_tuning_pretrained_model_annotation.ipynb" target="_blank" rel="noopener noreferrer">🏷️ 04 · Fine-tune a pretrained model on annotations</a></li>
+        <li><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/initial-release/notebooks_tutorials/05_model_interpretation.ipynb" target="_blank" rel="noopener noreferrer">🔍 05 · Model interpretation</a></li>
+        <li><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/initial-release/notebooks_tutorials/06_NTv3_generative_training.ipynb" target="_blank" rel="noopener noreferrer">🧪 06 · Fine-tuning NTv3 into a diffusion model</a></li>
+        <li><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/initial-release/notebooks_tutorials/07_enhancer_generation.ipynb" target="_blank" rel="noopener noreferrer">🪰 07 · Generating enhancer sequences</a></li>
+      </ul>
+      <p style="margin-top: 12px;">Pipeline notebooks:</p>
+      <ul>
+        <li><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/initial-release/notebooks_pipelines/01_functional_track_prediction.ipynb" target="_blank" rel="noopener noreferrer">🎯 01 · Generate bigwig predictions for given tracks</a></li>
+        <li><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3/blob/initial-release/notebooks_pipelines/02_model_interpretation.ipynb" target="_blank" rel="noopener noreferrer">🔍 02 · Interpret a given genomic region</a></li>
+      </ul>
+    </div>
+    <div class="card">
+      <h2>🔗 Initial-release links (2025)</h2>
+      <ul>
+        <li><a href="https://www.biorxiv.org/content/10.64898/2025.12.22.695963v1" target="_blank" rel="noopener noreferrer">📄 Initial preprint</a></li>
+        <li><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3/tree/initial-release" target="_blank" rel="noopener noreferrer">🏠 Original NTv3 Space (files at tag <code>initial-release</code>)</a></li>
+        <li><a href="https://github.com/instadeepai/nucleotide-transformer/tree/ntv3-initial-release" target="_blank" rel="noopener noreferrer">💻 JAX model code (GitHub, tag <code>ntv3-initial-release</code>)</a></li>
+        <li><a href="https://huggingface.co/collections/InstaDeepAI/nucleotide-transformer-v3" target="_blank" rel="noopener noreferrer">🎯 HF model collection (2025)</a></li>
+        <li><a href="https://huggingface.co/datasets/InstaDeepAI/NTv3_benchmark_dataset" target="_blank" rel="noopener noreferrer">📚 NTv3 benchmark dataset</a></li>
+        <li><a href="https://huggingface.co/datasets/InstaDeepAI/NTv3_enhancer_generation" target="_blank" rel="noopener noreferrer">📚 NTv3 enhancer generation dataset</a></li>
+        <li><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3_benchmark" target="_blank" rel="noopener noreferrer">🏆 NTv3 benchmark leaderboard (2025)</a></li>
+        <li><a href="https://huggingface.co/spaces/InstaDeepAI/ntv3_tracks" target="_blank" rel="noopener noreferrer">🚀 Interactive tracks demo</a></li>
+      </ul>
+    </div>
+  </div>
+
+  <div class="card">
+    <h2>🤖 Load a pre-trained model</h2>
+    <p>Here is an example of how to load and use a pre-trained NTv3 model.</p>
+    <div class="code"><pre><code class="language-python">from transformers import AutoTokenizer, AutoModelForMaskedLM
+
+model_name = "InstaDeepAI/NTv3_650M_pre"
+
+model = AutoModelForMaskedLM.from_pretrained(model_name, trust_remote_code=True)
+tok = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+
+batch = tok(["ATCGNATCG", "ACGT"], add_special_tokens=False, padding=True,
+            pad_to_multiple_of=128, return_tensors="pt")
+
+out = model(**batch)
+print(out.logits.shape)       # (B, L, V = 11)
+</code></pre></div>
+    <p>
+      To pin a specific version rather than tracking the branch, pass a tag:
+      <code>from_pretrained(model_name, revision="v1.0", trust_remote_code=True)</code>
+    </p>
+  </div>
+
+  <div class="card">
+    <h2>🔍 Model interpretation</h2>
+    <p>
+      Here is an example of how to use the interpretation pipeline on the NTv3 post-trained model for multi-scale
+      analysis of DNA sequences:
+    </p>
+    <div class="code"><pre><code class="language-python">from transformers import pipeline
+import torch
+import matplotlib.pyplot as plt
+
+model_name = "InstaDeepAI/NTv3_650M_post"
+
+# Build interpretation pipeline
+ntv3_interpret = pipeline(
+    "ntv3-interpret",
+    model=model_name,
+    trust_remote_code=True,
+    device=0 if torch.cuda.is_available() else -1,
+)
+
+# Run interpretation on a given genomic region with tracks, annotations, attention, and saliency
+result = ntv3_interpret(
+    {"chrom": "chr11", "start": 5_253_561, "end": 5_286_329, "species": "human"},
+    output_attention=True,
+    output_saliency=True,
+    saliency_track_id="ENCSR000EFT",  # K562 GATA1 ChIP-seq
+    plot=True,  # plot predictions on tracks and annotations
+    tracks_to_plot={"K562 RNA-seq": "ENCSR056HPM", "K562 GATA1": "ENCSR000EFT"},
+    elements_to_plot=["exon", "promoter_Tissue_specific"],
+)
+
+# Access attention map results
+result.plot_attention()  # attention map (last layer)
+plt.show()
+
+# Access saliency scores results
+result.plot_saliency(window_size=128)
+plt.show()</code></pre></div>
+    <img src="assets/saliency_example.png" alt="Saliency example" style="max-width: 100%; margin-top: 20px;" />
+  </div>
+
+  <div class="card" style="grid-column: span 12;">
+    <h2>💻 Use a post-trained model</h2>
+    <p>
+      Here is a quick example of how to use the post-trained NTv3 650M model to predict tracks for a human genomic
+      window.
+    </p>
+    <div class="code"><pre><code class="language-python">from transformers import pipeline
+import torch
+
+model_name = "InstaDeepAI/NTv3_650M_post"
+
+ntv3_tracks = pipeline(
+    "ntv3-tracks",
+    model=model_name,
+    trust_remote_code=True,
+    device=0 if torch.cuda.is_available() else -1,
+)
+
+# Run track prediction
+out = ntv3_tracks(
+  {
+    "chrom": "chr19",
+    "start": 6_700_000,
+    "end": 6_831_072,
+    "species": "human"
+  }
+)
+
+# Print output shapes
+# 7k human tracks over 37.5 % center region of the input sequence
+print("bigwig_tracks_logits:", tuple(out.bigwig_tracks_logits.shape))
+# Location of 21 genomic elements over 37.5 % center region of the input sequence
+print("bed_tracks_logits:", tuple(out.bed_tracks_logits.shape))
+# Language model logits for whole sequence over vocabulary
+print("language model logits:", tuple(out.mlm_logits.shape))</code></pre></div>
+    <p>Predictions can be plotted for a subset of tracks and genomic elements:</p>
+    <div class="code"><pre><code class="language-python">tracks_to_plot = {
+    "K562 RNA-seq": "ENCSR056HPM",
+    "K562 DNAse": "ENCSR921NMD",
+    "K562 H3k4me3": "ENCSR000DWD",
+    "K562 CTCF": "ENCSR000AKO",
+    "HepG2 RNA-seq": "ENCSR561FEE_P",
+    "HepG2 DNAse": "ENCSR000EJV",
+    "HepG2 H3k4me3": "ENCSR000AMP",
+    "HepG2 CTCF": "ENCSR000BIE",
+}
+elements_to_plot = ["protein_coding_gene", "exon", "intron", "splice_donor", "splice_acceptor"]
+
+out = ntv3_tracks(
+    {"chrom": "chr19", "start": 6_700_000, "end": 6_831_072, "species": "human"},
+    plot=True,
+    tracks_to_plot=tracks_to_plot,
+    elements_to_plot=elements_to_plot,
+)</code></pre></div>
+    <img src="assets/output_tracks.png" alt="Predicted tracks over a human genomic window" style="max-width: 100%; margin-top: 20px;" />
+  </div>
+</div>

tabs/skill.html

@@ -0,0 +1,148 @@
+<div class="summary">
+<h2>🤖 NTv3 in Claude Code</h2>
+<p>The NTv3 Claude Skill provides a natural-language interface to NTv3 models and genomic analysis tools. It combines model outputs with external genomic data and annotations to support analysis and interpretation of genomic regions.</p><p>Current analyses include functional-track prediction, saliency (gradient-based attribution), single-nucleotide variant comparison and gene annotation, with further interpretation methods and tools planned as the skill develops.</p>
+<p>Give the agent a gene or region and the analysis you want to run. It asks for any missing details, prepares an analysis plan and waits for your approval before running it.</p>
+</div>
+<div class="grid">
+<div class="card" style="grid-column: span 12;">
+<h2>🧭 How it works</h2>
+<p>The agent retrieves the reference sequence, runs the selected analyses and gathers relevant external data. Results are assembled into a locus-level report with an interactive genome browser, delivered as a Claude artifact. A Python script supports rerunning the analysis, and an optional notebook provides access to the saved results.</p><p style="margin-top: 12px;">An evidence audit checks references to analysis files and literature and screens report claims for predefined forms of overstatement. Its decisions are recorded alongside the results.</p>
+<a href="assets/skill_workflow_v6b.png" rel="noopener noreferrer" target="_blank"><img alt="NTv3 skill workflow: the natural-language request 'Predict K562 CAGE and DNase at HBE1 and show which bases drive the prediction', the LLM agent that clarifies inputs and resolves a plan, user approval, then NTv3 track prediction, gradient attribution and single-variant comparison with public ENCODE and UCSC context, an evidence audit, and a locus-level report published as a Claude artifact" src="assets/skill_workflow_v6b.png" style="max-width: 100%; margin-top: 20px; border-radius: 12px;"/></a>
+</div>
+<div class="card" style="grid-column: span 12;">
+<h2>🖥️ Example: the HBE1 locus in K562</h2>
+<style>
+      /* Skill-run transcript (matches supp figure v10a, panel b). Scoped to .skill-chat. */
+      .skill-chat {
+        margin-top: 12px;
+        padding: 14px 16px;
+        border: 1px solid var(--border);
+        border-radius: 14px;
+        background: rgba(125, 211, 252, 0.06);
+        font-size: 14px;
+        line-height: 1.55;
+        color: var(--text);
+        display: grid;
+        grid-template-columns: 4.6em minmax(0, 1fr);
+        column-gap: 12px;
+        row-gap: 10px;
+        overflow-wrap: anywhere;
+      }
+      .skill-chat .who { color: var(--muted); font-weight: 600; }
+      .skill-chat .turn { min-width: 0; }
+      .skill-chat .pair { display: inline-block; max-width: 100%; margin-right: 0.35em; }
+      .skill-chat .sep { color: var(--muted); margin-left: 0.15em; }
+      .skill-chat .q { font-weight: 600; }
+      .skill-chat .key { color: var(--muted); }
+      .skill-chat .plan {
+        padding: 8px 10px;
+        border: 1px solid rgba(125, 211, 252, 0.35);
+        border-radius: 10px;
+        background: rgba(0, 0, 0, 0.18);
+      }
+      .skill-chat .arrow { color: var(--muted); margin: 0 0.35em; }
+      .skill-prompt pre, .skill-prompt code[class*="language-"] { white-space: pre-wrap; overflow-wrap: anywhere; }
+      @media (max-width: 520px) {
+        .skill-chat { font-size: 13px; padding: 12px; column-gap: 8px; grid-template-columns: 3.9em minmax(0, 1fr); }
+      }
+    </style>
+<div aria-label="Skill run transcript" class="skill-chat" role="group">
+<div class="who">User</div>
+<div class="turn">Predict K562 CAGE and DNase at <em>HBE1</em> and show which bases drive the prediction.</div>
+<div class="who">Agent</div>
+<div class="turn">
+<span class="pair"><span class="q">Species?</span> human <span class="sep">·</span></span>
+<span class="pair"><span class="q">Assembly?</span> hg38 <span class="sep">·</span></span>
+<span class="pair"><span class="q">Tracks?</span> CAGE, DNase <span class="sep">·</span></span>
+<span class="pair"><span class="q">Window length?</span> 131,072 bp <span class="sep">·</span></span>
+<span class="pair"><span class="q">Browser lanes?</span> GENCODE (knownGene), ClinVar (clinvarMain)</span>
+</div>
+<div class="who">Plan</div>
+<div class="turn plan">
+<span class="pair"><span class="key">Window</span> chr11:5,204,408–5,335,480 (131,072 bp, GPU, centered on the <em>HBE1</em> TSS)</span><br/>
+<span class="pair"><span class="key">Prediction region</span> 49,152 bp <span class="sep">·</span></span>
+<span class="pair"><span class="key">Tracks</span> CAGE CNhs11250_P; DNase ENCSR921NMD</span><br/>
+<span class="pair"><span class="key">Measured signal</span> ENCODE ENCSR921NMD (DNase) <span class="sep">·</span></span>
+<span class="pair"><span class="key">Attribution</span> CAGE <span class="sep">·</span></span>
+<span class="pair"><span class="key">Evidence audit</span> on</span>
+</div>
+<div class="who">User</div>
+<div class="turn">
+<strong>Approve</strong><span class="arrow">→</span><span class="who">Agent</span> 
+        lookup → predict + saliency → evidence audit → locus-level report (Claude artifact)
+      </div>
+</div>
+<p style="margin-top: 14px;">Here, the skill uses NTv3_tracks to predict CAGE and DNase signal in K562 cells around <em>HBE1</em>. The 131,072-bp input is centered on the transcription start site; predictions cover the central 49,152 bp, including the β-globin locus control region.</p>
+<p style="margin-top: 10px;">The report brings together predicted CAGE and DNase tracks, measured DNase signal from ENCODE (ENCSR921NMD), and GENCODE and ClinVar annotations. CAGE attributions are shown below the signal tracks. Each lane has its own scale, so compare peak shapes rather than heights. The <em>HBE1</em> TSS line and the HS1–HS5 labels of the locus control region were added for this figure.</p>
+<a href="assets/skill_report_hbe1_cagesal_top_hs.png" rel="noopener noreferrer" target="_blank"><img alt="NTv3 skill report browser for HBE1 in K562 over the 49,152 bp predicted crop, with its controls row, range line and ideogram: GENCODE and ClinVar lanes, the NTv3 predicted CAGE track, the measured and predicted DNase tracks, and the CAGE saliency, with the HBE1 TSS and HS1 to HS5 marked" src="assets/skill_report_hbe1_cagesal_top_hs.png" style="max-width: 100%; margin-top: 20px; border-radius: 12px;"/></a>
+<p style="margin-top: 16px;">Zooming in on HS4 shows the CAGE attributions at base resolution. Letter heights show the gradient × input of each base for the mean predicted CAGE over the 49,152-bp prediction region. The GATA- and AP-1-like boxes were added for this figure.</p>
+<a href="assets/skill_report_hbe1_hs4_saliency_cage_r848.png" rel="noopener noreferrer" target="_blank"><img alt="HS4, chr11:5,288,225-5,288,345 (120 bp), Saliency · CAGE lane: the saliency letters show a GATA site pair and an AP-1 site" src="assets/skill_report_hbe1_hs4_saliency_cage_r848.png" style="max-width: 100%; margin-top: 20px; border-radius: 12px; background: #fff;"/></a>
+<p style="margin-top: 6px; font-size: 13px;">HS4 · chr11:5,288,225–5,288,345 (hg38) · CAGE attribution</p>
+</div>
+<div class="card" style="grid-column: span 12;">
+<h2>🎯 Current capabilities</h2>
+<table class="capability-grid">
+<thead>
+<tr>
+<th>Analysis</th>
+<th>What the skill does</th>
+</tr>
+</thead>
+<tbody>
+<tr>
+<td><strong>Predict</strong></td>
+<td>Predict functional tracks for a gene or genomic region, with the species and tracks selected for the analysis.</td>
+</tr>
+<tr>
+<td><strong>Saliency</strong></td>
+<td>Show which bases drive a selected prediction, as gradient × input at base resolution.</td>
+</tr>
+<tr>
+<td><strong>Compare</strong></td>
+<td>Predict each allele of a single-nucleotide variant and report the change in each selected track (alternative − reference).</td>
+</tr>
+<tr>
+<td><strong>Annotate</strong></td>
+<td>Predict gene structures for a FASTA, or for the 1 Mb around a Predict run, with the clade-specific NTv3 annotation models.</td>
+</tr>
+<tr>
+<td><strong>Explain</strong></td>
+<td>Show which model, window, crop and head produced each number, read from the run's saved files.</td>
+</tr>
+</tbody>
+</table>
+<p style="margin-top: 12px;"><strong>Scope:</strong> human (hg38) and mouse (mm10) for Predict, Saliency and Compare; any mammal, plant or insect FASTA for Annotate.</p>
+<div class="notice">The <strong>Home</strong> tab describes the models and provides notebooks for running the analyses directly.</div>
+</div>
+<div class="card" style="grid-column: span 12;">
+<h2>🚀 Run the skill</h2>
+<p>The skill requires Claude Code, Python 3.11 with uv, and network access to UCSC, ENCODE and Hugging Face. Install it from the <a href="https://github.com/instadeepai/nucleotide-transformer" rel="noopener noreferrer" target="_blank">NTv3 repository</a> (its README has the details):</p>
+<div class="code"><pre><code class="language-bash">git clone https://github.com/instadeepai/nucleotide-transformer.git
+cd nucleotide-transformer
+"$(uv python find 3.11)" -m venv .venv
+.venv/bin/pip install -e .      # on a GPU host: .venv/bin/pip install -e '.[gpu]'
+uv sync --project ntv3_skills --python 3.11 --frozen
+ntv3_skills/.venv/bin/ntv3 doctor</code></pre></div>
+<p style="margin-top: 12px;"><code>ntv3 doctor</code> finds the Python that runs the model and says whether it sees a GPU; the skill picks the device from it.</p>
+<div class="notice"><strong>GPU or CPU:</strong> with the <code>gpu</code> install on a GPU host, the skill uses the GPU for everything: windows up to 1,048,576 bp, such as the 131,072-bp example above, and Annotate (about 25 GB of GPU memory per 1 Mb window). On a laptop CPU, Predict, Saliency and Compare also run; from 16,384 bp the plan warns that it is slow and needs a lot of memory. On either device, Saliency above 131,072 bp comes with a warning that it may run out of memory.</div>
+<p style="margin-top: 12px;">Open Claude Code in the repository directory and try a request that runs on any laptop:</p>
+<div class="code skill-prompt"><pre><code class="language-text">/ntv3 HBE1 at chr11 5266144 5270240 on hg38 in K562. Predict K562 CAGE with a 4,096 bp window and the GENCODE and ClinVar lanes.</code></pre></div>
+<p style="margin-top: 12px;">The agent asks for any missing details and presents the analysis plan for approval. Once approved, NTv3_tracks predicts K562 CAGE over the central 1,536 bp of the window, at the 3′ end of <em>HBE1</em>, next to the measured ENCODE CAGE signal. Running Annotate on the 1 Mb around this run adds the NTv3 gene model, which matches the three coding exons of <em>HBE1</em> in GENCODE:</p>
+<a href="assets/skill_report_hbe1_4kb_annotation.png" rel="noopener noreferrer" target="_blank"><img alt="NTv3 skill report browser for the 4,096 bp HBE1 window in K562: GENCODE V50, the NTv3 annotation lane with a three-exon minus-strand gene model matching HBE1, ClinVar, the measured CAGE signal and the NTv3 predicted CAGE track" src="assets/skill_report_hbe1_4kb_annotation.png" style="max-width: 100%; margin-top: 16px; border-radius: 12px; background: #fff;"/></a>
+<p style="margin-top: 6px; font-size: 13px;">Report browser for the request above, after Annotate with the mammals model · chr11:5,266,144–5,270,240 (hg38). Thin GENCODE blocks are UTRs; NTv3 predicts coding exons only.</p>
+<p style="margin-top: 16px;">More requests to try:</p>
+<div class="code skill-prompt"><pre><code class="language-text"># Saliency
+/ntv3 Predict K562 CAGE at HBE1 with a 4,096 bp window and show which bases drive it.
+# Compare (rs334, the HBB sickle-cell variant)
+/ntv3 Compare chr11 5227002 T A on hg38 in K562: CAGE with a 4,096 bp window.
+# The 131 kb example above (GPU)
+/ntv3 Predict K562 CAGE and DNase at HBE1 with a 131,072 bp window and show which bases drive the prediction.
+# Annotate a genome (GPU)
+/ntv3 Annotate /path/to/genome.fa with the mammals model.
+# Annotate the 1 Mb around a Predict run (GPU)
+/ntv3 Annotate the 1 Mb around my HBE1 run.</code></pre></div>
+<p style="margin-top: 12px;">The skill's tools write only to <code>results/</code> in the repository and upload nothing; the report reaches you as a Claude artifact, and every model download is listed in the plan you approve.</p>
+<h2 style="margin-top: 20px;">🔗 Links</h2>
+<ul><li>Analysis notebooks: <strong>Home</strong> tab</li><li><a href="https://github.com/instadeepai/nucleotide-transformer" rel="noopener noreferrer" target="_blank">NTv3 model code (GitHub)</a></li><li><a href="https://huggingface.co/collections/InstaDeepAI/ntv3" rel="noopener noreferrer" target="_blank">NTv3 model collection</a></li></ul>
+</div>
+</div>